Skip to content

Commit b80d23f

Browse files
tlopextlopex
authored
[Docs] Add tvm.s_tir.tensor_intrin API reference and remove empty legacy tvm/tir directory (#19386)
as per title
1 parent 0c43ab2 commit b80d23f

7 files changed

Lines changed: 123 additions & 42 deletions

File tree

‎docs/reference/api/python/index.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@ Python API
7676
s_tir/transform
7777
s_tir/dlight
7878
s_tir/backend
79+
s_tir/tensor_intrin
7980

8081
.. toctree::
8182
:maxdepth: 1

‎docs/reference/api/python/s_tir/tensor_intrin.rst‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
.. Licensed to the Apache Software Foundation (ASF) under one
2+
or more contributor license agreements. See the NOTICE file
3+
distributed with this work for additional information
4+
regarding copyright ownership. The ASF licenses this file
5+
to you under the Apache License, Version 2.0 (the
6+
"License"); you may not use this file except in compliance
7+
with the License. You may obtain a copy of the License at
8+
9+
.. http://www.apache.org/licenses/LICENSE-2.0
10+
11+
.. Unless required by applicable law or agreed to in writing,
12+
software distributed under the License is distributed on an
13+
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
KIND, either express or implied. See the License for the
15+
specific language governing permissions and limitations
16+
under the License.
17+
18+
tvm.s_tir.tensor_intrin
19+
-----------------------
20+
21+
tvm.s_tir.tensor_intrin
22+
***********************
23+
.. automodule:: tvm.s_tir.tensor_intrin
24+
:members:
25+
:imported-members:
26+
27+
tvm.s_tir.tensor_intrin.cuda
28+
****************************
29+
.. automodule:: tvm.s_tir.tensor_intrin.cuda
30+
:members:
31+
:noindex:
32+
33+
tvm.s_tir.tensor_intrin.arm_cpu
34+
*******************************
35+
.. automodule:: tvm.s_tir.tensor_intrin.arm_cpu
36+
:members:
37+
:noindex:
38+
39+
tvm.s_tir.tensor_intrin.x86
40+
****************************
41+
.. automodule:: tvm.s_tir.tensor_intrin.x86
42+
:members:
43+
:noindex:
44+
45+
tvm.s_tir.tensor_intrin.rocm
46+
*****************************
47+
.. automodule:: tvm.s_tir.tensor_intrin.rocm
48+
:members:
49+
:noindex:
50+
51+
tvm.s_tir.tensor_intrin.metal
52+
*****************************
53+
.. automodule:: tvm.s_tir.tensor_intrin.metal
54+
:members:
55+
:noindex:
56+
57+
tvm.s_tir.tensor_intrin.hexagon
58+
*******************************
59+
.. automodule:: tvm.s_tir.tensor_intrin.hexagon
60+
:members:
61+
:noindex:
62+
63+
tvm.s_tir.tensor_intrin.riscv_cpu
64+
*********************************
65+
.. automodule:: tvm.s_tir.tensor_intrin.riscv_cpu
66+
:members:
67+
:noindex:
68+
69+
tvm.s_tir.tensor_intrin.dot_product_common
70+
******************************************
71+
.. automodule:: tvm.s_tir.tensor_intrin.dot_product_common
72+
:members:
73+
:noindex:

‎python/tvm/arith/bound.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,16 +24,16 @@ def deduce_bound(var, cond, hint_map, relax_map):
2424
2525
Parameters
2626
----------
27-
var : tvm.tir.Var
27+
var : tvm.tirx.Var
2828
The target variable to be deduced.
2929
3030
cond : PrimExpr
3131
The condition
3232
33-
hint_map : Map[tvm.tir.Var, IntSet]
33+
hint_map : Map[tvm.tirx.Var, IntSet]
3434
Domain of variables used to help deduction.
3535
36-
relax_map : Map[tvm.tir.Var, IntSet]
36+
relax_map : Map[tvm.tirx.Var, IntSet]
3737
The fomain of the variables to be relaxed
3838
using the provided domain.
3939
"""

‎python/tvm/arith/int_set.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ def estimate_region_lower_bound(region, var_dom, predicate):
102102
region : List[Range]
103103
The region to be analyzed.
104104
105-
var_dom : Dict[tvm.tir.Var, Range]
105+
var_dom : Dict[tvm.tirx.Var, Range]
106106
The ranges of the variables
107107
108108
predicate : PrimExpr
@@ -125,7 +125,7 @@ def estimate_region_strict_bound(region, var_dom, predicate):
125125
region : List[Range]
126126
The region to be analyzed.
127127
128-
var_dom : Dict[tvm.tir.Var, Range]
128+
var_dom : Dict[tvm.tirx.Var, Range]
129129
The ranges of the variables
130130
131131
predicate : PrimExpr
@@ -149,7 +149,7 @@ def estimate_region_upper_bound(region, var_dom, predicate):
149149
region : List[Range]
150150
The region to be analyzed.
151151
152-
var_dom : Dict[tvm.tir.Var, Range]
152+
var_dom : Dict[tvm.tirx.Var, Range]
153153
The ranges of the variables
154154
155155
predicate : PrimExpr
@@ -168,7 +168,7 @@ def pos_inf():
168168
169169
Returns
170170
----------
171-
pos_inf : tvm.tir.Var
171+
pos_inf : tvm.tirx.Var
172172
A symbolic var that indicates positive infinity
173173
"""
174174
return _ffi_api.PosInf()
@@ -179,7 +179,7 @@ def neg_inf():
179179
180180
Returns
181181
----------
182-
neg_inf : tvm.tir.Var
182+
neg_inf : tvm.tirx.Var
183183
A symbolic var that indicates positive infinity
184184
"""
185185
return _ffi_api.NegInf()

‎python/tvm/arith/iter_affine_map.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ def detect_iter_map(
137137
indices : List[PrimExpr]
138138
The input indices
139139
140-
input_iters : Map[tvm.tir.Var, Range]
140+
input_iters : Map[tvm.tirx.Var, Range]
141141
The domain of each input iterators.
142142
143143
predicate : PrimExpr
@@ -178,7 +178,7 @@ def normalize_to_iter_sum(index, input_iters):
178178
index : PrimExpr
179179
The input index
180180
181-
input_iters : Map[tvm.tir.Var, Range]
181+
input_iters : Map[tvm.tirx.Var, Range]
182182
The domain of each input iterators.
183183
184184
Returns
@@ -211,7 +211,7 @@ def iter_map_simplify(
211211
indices : List[PrimExpr]
212212
The input indices
213213
214-
input_iters : Map[tvm.tir.Var, Range]
214+
input_iters : Map[tvm.tirx.Var, Range]
215215
The domain of each input iterators.
216216
217217
predicate : PrimExpr
@@ -289,10 +289,10 @@ def subspace_divide(
289289
bindings : List[PrimExpr]
290290
The input bindings
291291
292-
input_iters : Map[tvm.tir.Var, Range]
292+
input_iters : Map[tvm.tirx.Var, Range]
293293
The domain of input iterator, which is the basis of the whole space
294294
295-
sub_iters : Array[tvm.tir.Var]
295+
sub_iters : Array[tvm.tirx.Var]
296296
The subset of input_iters, which is the basis of the subspace
297297
298298
predicate : PrimExpr
@@ -344,7 +344,7 @@ def inverse_affine_iter_map(iter_map, outputs):
344344
345345
Returns
346346
-------
347-
results : Map[tvm.tir.Var, PrimExpr]
347+
results : Map[tvm.tirx.Var, PrimExpr]
348348
The map from the input to the transformed result.
349349
"""
350350
return _ffi_api.InverseAffineIterMap(iter_map, outputs)

‎python/tvm/s_tir/tensor_intrin/arm_cpu.py‎

Lines changed: 33 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -234,16 +234,18 @@ def get_sme_transpose_interleave_2svlx2svl_fp32_intrin(cols, rows):
234234
the contents of sub-tile 1 and 2 are stored in opposite locations - see the diagram
235235
below.
236236
237-
A: Accumulator tile: A_t:
238-
2SVL 2SVL 2SVL
239-
+----------------+ +-----------------+ +-------------------+
240-
| --0a-- --1a-- | | | | | | | | |
241-
| --0b-- --1b-- | | 0 1 | | 0a 0b .. 2a 2b .. |
242-
| ... ... | ld1w.horiz | | st1w.vert | | | | | |
243-
2SVL | --2a-- --3a-- | ====> 2SVL | | ====> 2SVL | | | | | |
244-
| --2a-- --3b-- | | 2 3 | | 1a 1b .. 3a 3b .. |
245-
| ... ... | | | | | | | | |
246-
+----------------+ +-----------------+ +-------------------+
237+
::
238+
239+
A: Accumulator tile: A_t:
240+
2SVL 2SVL 2SVL
241+
+----------------+ +-----------------+ +-------------------+
242+
| --0a-- --1a-- | | | | | | | | |
243+
| --0b-- --1b-- | | 0 1 | | 0a 0b .. 2a 2b .. |
244+
| ... ... | ld1w.horiz | | st1w.vert | | | | | |
245+
2SVL | --2a-- --3a-- | ====> 2SVL | | ====> 2SVL | | | | | |
246+
| --2a-- --3b-- | | 2 3 | | 1a 1b .. 3a 3b .. |
247+
| ... ... | | | | | | | | |
248+
+----------------+ +-----------------+ +-------------------+
247249
248250
Returns
249251
-------
@@ -521,24 +523,26 @@ def get_sme_gemm_interleaved_mopa_2svlx2svl_intrin(M, K, in_dtype):
521523
Diagram showing outer-product performed on each of the accumulator sub-tiles
522524
for the fp32 datatype:
523525
524-
SVL SVL
525-
+----------------------------+
526-
| l | h | K
527-
K +----------------------------+
528-
+---+ +----------------------------+
529-
| | | 0: 1: |-+
530-
| | | mopa(l, l) mopa(l, h) | |-+
531-
l | | | | | |
532-
| | | | | |
533-
|---| | | | |
534-
| | | 2: 3: | | |
535-
h | | | mopa(h, l) mopa(h, h) | | |
536-
| | | | | |
537-
| | | | | |
538-
+---+ +----------------------------+ | |
539-
+----------------------------+ |
540-
+---------------------------+
541-
(accumulate K times)
526+
::
527+
528+
SVL SVL
529+
+----------------------------+
530+
| l | h | K
531+
K +----------------------------+
532+
+---+ +----------------------------+
533+
| | | 0: 1: |-+
534+
| | | mopa(l, l) mopa(l, h) | |-+
535+
l | | | | | |
536+
| | | | | |
537+
|---| | | | |
538+
| | | 2: 3: | | |
539+
h | | | mopa(h, l) mopa(h, h) | | |
540+
| | | | | |
541+
| | | | | |
542+
+---+ +----------------------------+ | |
543+
+----------------------------+ |
544+
+---------------------------+
545+
(accumulate K times)
542546
543547
Pseudo code computing 2SVL x 2SVL GEMM for fp32 inputs:
544548
@@ -572,6 +576,7 @@ def get_sme_gemm_interleaved_mopa_2svlx2svl_intrin(M, K, in_dtype):
572576
}
573577
574578
Notes:
579+
575580
- Recall that A has been transposed beforehand such that each column is now accessed
576581
by row.
577582
- 'sme.zero' resets the accumulator tile to contain all zero's.

‎python/tvm/s_tir/tensor_intrin/riscv_cpu.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,9 @@ def rvv_vec_dot_product_kernels(
6161
with C[LANES].
6262
6363
The pseudo code is as follows:
64+
6465
.. code-block:: c
66+
6567
void vec_dot_prod(A[ELEMS], B[LANES][ELEMS], C[LANES]){
6668
for (j = 0; j < LANES; j++) {
6769
for (k = 0; k < ELEMS; k++) {

0 commit comments

Comments
 (0)