This is an automated email from the ASF dual-hosted git repository.
ruihangl pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new b6f67b06db [Docs] Add Python API reference for tvm submodule docs
(#19379)
b6f67b06db is described below
commit b6f67b06db4dcfdefd0ec55d9bee38698248222f
Author: Shushi Hong <[email protected]>
AuthorDate: Fri Apr 10 14:52:49 2026 -0400
[Docs] Add Python API reference for tvm submodule docs (#19379)
as per title
---
docs/reference/api/python/{ir.rst => arith.rst} | 7 +--
docs/reference/api/python/contrib.rst | 14 +++++
docs/reference/api/python/{ir.rst => exec.rst} | 7 +--
docs/reference/api/python/index.rst | 4 ++
docs/reference/api/python/instrument.rst | 1 +
docs/reference/api/python/ir.rst | 6 ++
docs/reference/api/python/target.rst | 6 ++
docs/reference/api/python/{ir.rst => testing.rst} | 7 +--
.../python/{instrument.rst => tirx/backend.rst} | 15 ++++-
docs/reference/api/python/topi.rst | 26 +++++++++
python/tvm/arith/bound.py | 6 +-
python/tvm/arith/int_set.py | 10 ++--
python/tvm/arith/iter_affine_map.py | 55 ++++++++++--------
python/tvm/contrib/hexagon/tools.py | 4 +-
python/tvm/testing/utils.py | 14 +++--
python/tvm/topi/testing/strided_slice_python.py | 9 +--
python/tvm/topi/vision/nms.py | 18 +++---
python/tvm/topi/vision/nms_util.py | 67 ++++++++++++----------
18 files changed, 179 insertions(+), 97 deletions(-)
diff --git a/docs/reference/api/python/ir.rst
b/docs/reference/api/python/arith.rst
similarity index 94%
copy from docs/reference/api/python/ir.rst
copy to docs/reference/api/python/arith.rst
index 1f0dc0c5e2..742fbf5aef 100644
--- a/docs/reference/api/python/ir.rst
+++ b/docs/reference/api/python/arith.rst
@@ -15,9 +15,8 @@
specific language governing permissions and limitations
under the License.
-tvm.ir
-------
-.. automodule:: tvm.ir
+tvm.arith
+---------
+.. automodule:: tvm.arith
:members:
:imported-members:
- :autosummary:
diff --git a/docs/reference/api/python/contrib.rst
b/docs/reference/api/python/contrib.rst
index ebe96910af..f16da9ecb9 100644
--- a/docs/reference/api/python/contrib.rst
+++ b/docs/reference/api/python/contrib.rst
@@ -113,3 +113,17 @@ tvm.contrib.xcode
~~~~~~~~~~~~~~~~~
.. automodule:: tvm.contrib.xcode
:members:
+
+
+tvm.contrib.cutlass
+~~~~~~~~~~~~~~~~~~~
+.. automodule:: tvm.contrib.cutlass
+ :members:
+ :imported-members:
+
+
+tvm.contrib.hexagon
+~~~~~~~~~~~~~~~~~~~
+.. automodule:: tvm.contrib.hexagon
+ :members:
+ :imported-members:
diff --git a/docs/reference/api/python/ir.rst
b/docs/reference/api/python/exec.rst
similarity index 94%
copy from docs/reference/api/python/ir.rst
copy to docs/reference/api/python/exec.rst
index 1f0dc0c5e2..42976fc484 100644
--- a/docs/reference/api/python/ir.rst
+++ b/docs/reference/api/python/exec.rst
@@ -15,9 +15,8 @@
specific language governing permissions and limitations
under the License.
-tvm.ir
-------
-.. automodule:: tvm.ir
+tvm.exec
+--------
+.. automodule:: tvm.exec
:members:
:imported-members:
- :autosummary:
diff --git a/docs/reference/api/python/index.rst
b/docs/reference/api/python/index.rst
index c826a87e89..21d2d8bb32 100644
--- a/docs/reference/api/python/index.rst
+++ b/docs/reference/api/python/index.rst
@@ -24,10 +24,13 @@ Python API
error
ir
+ arith
instrument
transform
target
driver
+ testing
+ exec
.. toctree::
:maxdepth: 1
@@ -62,6 +65,7 @@ Python API
tirx/analysis
tirx/stmt_functor
tirx/transform
+ tirx/backend
.. toctree::
:maxdepth: 1
diff --git a/docs/reference/api/python/instrument.rst
b/docs/reference/api/python/instrument.rst
index 56897b2cfa..7008dec70f 100644
--- a/docs/reference/api/python/instrument.rst
+++ b/docs/reference/api/python/instrument.rst
@@ -20,3 +20,4 @@ tvm.ir.instrument
.. automodule:: tvm.ir.instrument
:members:
:imported-members:
+ :exclude-members: Path
diff --git a/docs/reference/api/python/ir.rst b/docs/reference/api/python/ir.rst
index 1f0dc0c5e2..a8d39e7b01 100644
--- a/docs/reference/api/python/ir.rst
+++ b/docs/reference/api/python/ir.rst
@@ -21,3 +21,9 @@ tvm.ir
:members:
:imported-members:
:autosummary:
+
+tvm.ir.diagnostics
+~~~~~~~~~~~~~~~~~~
+.. automodule:: tvm.ir.diagnostics
+ :members:
+ :imported-members:
diff --git a/docs/reference/api/python/target.rst
b/docs/reference/api/python/target.rst
index b3e763379a..9ccfb248d8 100644
--- a/docs/reference/api/python/target.rst
+++ b/docs/reference/api/python/target.rst
@@ -21,3 +21,9 @@ tvm.target
:members:
:imported-members:
:autosummary:
+
+tvm.target.tag
+~~~~~~~~~~~~~~
+.. automodule:: tvm.target.tag
+ :members:
+ :imported-members:
diff --git a/docs/reference/api/python/ir.rst
b/docs/reference/api/python/testing.rst
similarity index 94%
copy from docs/reference/api/python/ir.rst
copy to docs/reference/api/python/testing.rst
index 1f0dc0c5e2..5ef7181fd7 100644
--- a/docs/reference/api/python/ir.rst
+++ b/docs/reference/api/python/testing.rst
@@ -15,9 +15,8 @@
specific language governing permissions and limitations
under the License.
-tvm.ir
-------
-.. automodule:: tvm.ir
+tvm.testing
+-----------
+.. automodule:: tvm.testing
:members:
:imported-members:
- :autosummary:
diff --git a/docs/reference/api/python/instrument.rst
b/docs/reference/api/python/tirx/backend.rst
similarity index 78%
copy from docs/reference/api/python/instrument.rst
copy to docs/reference/api/python/tirx/backend.rst
index 56897b2cfa..c4e6689bfe 100644
--- a/docs/reference/api/python/instrument.rst
+++ b/docs/reference/api/python/tirx/backend.rst
@@ -15,8 +15,17 @@
specific language governing permissions and limitations
under the License.
-tvm.ir.instrument
-------------------
-.. automodule:: tvm.ir.instrument
+tvm.tirx.backend
+================
+
+tvm.tirx.backend
+*****************
+.. automodule:: tvm.tirx.backend
+ :members:
+ :imported-members:
+
+tvm.tirx.backend.adreno
+***********************
+.. automodule:: tvm.tirx.backend.adreno
:members:
:imported-members:
diff --git a/docs/reference/api/python/topi.rst
b/docs/reference/api/python/topi.rst
index e8ba9edd5f..65bc1a7e17 100644
--- a/docs/reference/api/python/topi.rst
+++ b/docs/reference/api/python/topi.rst
@@ -37,3 +37,29 @@ tvm.topi.image
:members:
:imported-members:
:autosummary:
+
+tvm.topi.vision
+~~~~~~~~~~~~~~~
+.. automodule:: tvm.topi.vision
+ :members:
+ :imported-members:
+ :noindex:
+
+tvm.topi.gpu
+~~~~~~~~~~~~
+.. automodule:: tvm.topi.gpu
+ :members:
+ :imported-members:
+ :noindex:
+
+tvm.topi.cpp
+~~~~~~~~~~~~
+.. automodule:: tvm.topi.cpp
+ :members:
+ :imported-members:
+
+tvm.topi.testing
+~~~~~~~~~~~~~~~~
+.. automodule:: tvm.topi.testing
+ :members:
+ :imported-members:
diff --git a/python/tvm/arith/bound.py b/python/tvm/arith/bound.py
index 9eb53d22b5..2abe386e1d 100644
--- a/python/tvm/arith/bound.py
+++ b/python/tvm/arith/bound.py
@@ -24,16 +24,16 @@ def deduce_bound(var, cond, hint_map, relax_map):
Parameters
----------
- var : Var
+ var : tvm.tir.Var
The target variable to be deduced.
cond : PrimExpr
The condition
- hint_map : Map[Var, IntSet]
+ hint_map : Map[tvm.tir.Var, IntSet]
Domain of variables used to help deduction.
- relax_map : Map[Var, IntSet]
+ relax_map : Map[tvm.tir.Var, IntSet]
The fomain of the variables to be relaxed
using the provided domain.
"""
diff --git a/python/tvm/arith/int_set.py b/python/tvm/arith/int_set.py
index 28694edfb7..746390187b 100644
--- a/python/tvm/arith/int_set.py
+++ b/python/tvm/arith/int_set.py
@@ -102,7 +102,7 @@ def estimate_region_lower_bound(region, var_dom, predicate):
region : List[Range]
The region to be analyzed.
- var_dom : Dict[Var, Range]
+ var_dom : Dict[tvm.tir.Var, Range]
The ranges of the variables
predicate : PrimExpr
@@ -125,7 +125,7 @@ def estimate_region_strict_bound(region, var_dom,
predicate):
region : List[Range]
The region to be analyzed.
- var_dom : Dict[Var, Range]
+ var_dom : Dict[tvm.tir.Var, Range]
The ranges of the variables
predicate : PrimExpr
@@ -149,7 +149,7 @@ def estimate_region_upper_bound(region, var_dom, predicate):
region : List[Range]
The region to be analyzed.
- var_dom : Dict[Var, Range]
+ var_dom : Dict[tvm.tir.Var, Range]
The ranges of the variables
predicate : PrimExpr
@@ -168,7 +168,7 @@ def pos_inf():
Returns
----------
- pos_inf : Var
+ pos_inf : tvm.tir.Var
A symbolic var that indicates positive infinity
"""
return _ffi_api.PosInf()
@@ -179,7 +179,7 @@ def neg_inf():
Returns
----------
- neg_inf : Var
+ neg_inf : tvm.tir.Var
A symbolic var that indicates positive infinity
"""
return _ffi_api.NegInf()
diff --git a/python/tvm/arith/iter_affine_map.py
b/python/tvm/arith/iter_affine_map.py
index b2548fc669..8371d368ee 100644
--- a/python/tvm/arith/iter_affine_map.py
+++ b/python/tvm/arith/iter_affine_map.py
@@ -137,7 +137,7 @@ def detect_iter_map(
indices : List[PrimExpr]
The input indices
- input_iters : Map[Var, Range]
+ input_iters : Map[tvm.tir.Var, Range]
The domain of each input iterators.
predicate : PrimExpr
@@ -178,7 +178,7 @@ def normalize_to_iter_sum(index, input_iters):
index : PrimExpr
The input index
- input_iters : Map[Var, Range]
+ input_iters : Map[tvm.tir.Var, Range]
The domain of each input iterators.
Returns
@@ -211,7 +211,7 @@ def iter_map_simplify(
indices : List[PrimExpr]
The input indices
- input_iters : Map[Var, Range]
+ input_iters : Map[tvm.tir.Var, Range]
The domain of each input iterators.
predicate : PrimExpr
@@ -265,28 +265,34 @@ def subspace_divide(
simplify_trivial_iterators=True,
):
"""Detect if bindings can be written as
- [a_0*e_0 + b_0 + c_0, a_1*e_1 + b_1, ..., a_n*e_n + b_n]
- where a = some-quasi-affine-iter-map(input_iters set_minus sub_iters)
- b = some-quasi-affine-iter-map(sub_iters)
- c is constant symbols
- e is the extent of b
- For example, z*12 + y*3 + x + c = (z*4+y)*3 + x
- bindings = [z*12 + y*3 + x + c]
- input_iters = [z, y, x]
- sub_iter = [x]
- Then the result will be [a, b] where
- a = [z*4 + y]
- b = [x]
+ ``[a_0*e_0 + b_0 + c_0, a_1*e_1 + b_1, ..., a_n*e_n + b_n]``
+
+ where::
+
+ a = some-quasi-affine-iter-map(input_iters set_minus sub_iters)
+ b = some-quasi-affine-iter-map(sub_iters)
+ c is constant symbols
+ e is the extent of b
+
+ For example::
+
+ z*12 + y*3 + x + c = (z*4+y)*3 + x
+ bindings = [z*12 + y*3 + x + c]
+ input_iters = [z, y, x]
+ sub_iter = [x]
+ Then the result will be [a, b] where
+ a = [z*4 + y]
+ b = [x]
Parameters
----------
bindings : List[PrimExpr]
The input bindings
- input_iters : Map[Var, Range]
+ input_iters : Map[tvm.tir.Var, Range]
The domain of input iterator, which is the basis of the whole space
- sub_iters : Array[Var]
+ sub_iters : Array[tvm.tir.Var]
The subset of input_iters, which is the basis of the subspace
predicate : PrimExpr
@@ -302,12 +308,13 @@ def subspace_divide(
Returns
-------
results : List[List[PrimExpr]]
- The result list has length len(bindings) + 1
- [0, len(bindings)): The iter map matching result. The inner list is of
length 2.
- The first expr is the basis of the quotient space.
- The second expr is the basis of the subspace.
- len(bindings): the predicate of outer space and inner space
- Empty array if no match can be found.
+ The result list has length ``len(bindings) + 1``.
+
+ - ``[0, len(bindings))``: The iter map matching result.
+ The inner list is of length 2. The first expr is the basis
+ of the quotient space. The second expr is the basis of the subspace.
+ - ``len(bindings)``: the predicate of outer space and inner space.
+ - Empty array if no match can be found.
"""
if isinstance(check_level, str):
check_level = IterMapLevel.from_str(check_level)
@@ -337,7 +344,7 @@ def inverse_affine_iter_map(iter_map, outputs):
Returns
-------
- results : Map[Var, PrimExpr]
+ results : Map[tvm.tir.Var, PrimExpr]
The map from the input to the transformed result.
"""
return _ffi_api.InverseAffineIterMap(iter_map, outputs)
diff --git a/python/tvm/contrib/hexagon/tools.py
b/python/tvm/contrib/hexagon/tools.py
index bf7913cf62..632149391f 100644
--- a/python/tvm/contrib/hexagon/tools.py
+++ b/python/tvm/contrib/hexagon/tools.py
@@ -115,7 +115,7 @@ def link_shared(so_name, objs, extra_args=None):
----------
so_name : str
Name of the shared library file.
- objs : list[str,StringImm]
+ objs : list[str, tvm.tirx.StringImm]
extra_args : dict (str->str) or Map<String,String>
Additional arguments:
'hex_arch' - Hexagon architecture, e.g. v68
@@ -187,7 +187,7 @@ def link_shared_macos(so_name, objs, extra_args=None):
----------
so_name : str
Name of the shared library file.
- objs : list[str,StringImm]
+ objs : list[str, tvm.tirx.StringImm]
extra_args : dict (str->str) or Map<String,String>
Additional arguments:
'hex_arch' - Hexagon architecture, e.g. v68
diff --git a/python/tvm/testing/utils.py b/python/tvm/testing/utils.py
index 125fa9586b..f15768858b 100644
--- a/python/tvm/testing/utils.py
+++ b/python/tvm/testing/utils.py
@@ -281,12 +281,14 @@ def check_bool_expr_is_true(bool_expr, vranges,
cond=None):
"""Check that bool_expr holds given the condition cond
for every value of free variables from vranges.
- for example, 2x > 4y solves to x > 2y given x in (0, 10) and y in (0, 10)
- here bool_expr is x > 2y, vranges is {x: (0, 10), y: (0, 10)}, cond is 2x
> 4y
- We creates iterations to check,
- for x in range(10):
- for y in range(10):
- assert !(2x > 4y) || (x > 2y)
+ For example, ``2x > 4y`` solves to ``x > 2y`` given ``x in (0, 10)``
+ and ``y in (0, 10)``. Here bool_expr is ``x > 2y``,
+ vranges is ``{x: (0, 10), y: (0, 10)}``, cond is ``2x > 4y``.
+ We create iterations to check::
+
+ for x in range(10):
+ for y in range(10):
+ assert !(2x > 4y) || (x > 2y)
Parameters
----------
diff --git a/python/tvm/topi/testing/strided_slice_python.py
b/python/tvm/topi/testing/strided_slice_python.py
index 3843d09967..54e041bf03 100644
--- a/python/tvm/topi/testing/strided_slice_python.py
+++ b/python/tvm/topi/testing/strided_slice_python.py
@@ -36,10 +36,11 @@ def strided_slice_python(data, begin, end, strides,
slice_mode="end", axes=None)
slice_mode : str, optional
The slice mode [end, size].
- end: The default slice mode, ending indices for the slice.
- size: The input strides will be ignored, input end in this mode
indicates
- the sizeof a slice starting at the location specified by begin.
If end[i] is -1,
- all remaining elements in that dimension are included in the
slice.
+
+ - ``"end"``: The default slice mode, ending indices for the slice.
+ - ``"size"``: The input strides will be ignored, input end in this
mode indicates
+ the size of a slice starting at the location specified by begin. If
end[i] is -1,
+ all remaining elements in that dimension are included in the slice.
axes : list, optional
Axes along which slicing is applied
diff --git a/python/tvm/topi/vision/nms.py b/python/tvm/topi/vision/nms.py
index b69e9c2aa1..a602527fcc 100644
--- a/python/tvm/topi/vision/nms.py
+++ b/python/tvm/topi/vision/nms.py
@@ -787,20 +787,22 @@ def all_class_non_max_suppression(
"""Non-maximum suppression operator for object detection, corresponding to
ONNX
NonMaxSuppression and TensorFlow combined_non_max_suppression.
NMS is performed for each class separately.
+
Parameters
----------
boxes : tvm.te.Tensor
3-D tensor with shape (batch_size, num_boxes, 4)
- scores: tvm.te.Tensor
+ scores : tvm.te.Tensor
3-D tensor with shape (batch_size, num_classes, num_boxes)
max_output_boxes_per_class : int or tvm.te.Tensor, optional
The maxinum number of output selected boxes per class
- iou_threshold : float or tvm.te.Tensor, optionaIl
+ iou_threshold : float or tvm.te.Tensor, optional
IoU test threshold
score_threshold : float or tvm.te.Tensor, optional
Score threshold to filter out low score boxes early
output_format : str, optional
"onnx" or "tensorflow", see below.
+
Returns
-------
out : list of tvm.te.Tensor
@@ -819,12 +821,12 @@ def all_class_non_max_suppression(
may contain garbage values. When comparing with ONNX Runtime or
other implementations
that output dynamic shapes, you should only compare the first
`num_total_detection` rows.
- Example:
- ```python
- selected_indices, valid_count = nms_output
- actual_count = int(valid_count.numpy()[0])
- valid_indices = selected_indices.numpy()[:actual_count, :]
- ```
+ Example::
+
+ selected_indices, valid_count = nms_output
+ actual_count = int(valid_count.numpy()[0])
+ valid_indices = selected_indices.numpy()[:actual_count, :]
+
If `output_format` is "tensorflow", the output is three tensors, the
first
is `indices` of size `(batch_size, num_class * num_boxes , 2)`, the
second is `scores` of
size `(batch_size, num_class * num_boxes)`, and the third is
`num_total_detection` of size
diff --git a/python/tvm/topi/vision/nms_util.py
b/python/tvm/topi/vision/nms_util.py
index a4b4c78363..f9bb460bc8 100644
--- a/python/tvm/topi/vision/nms_util.py
+++ b/python/tvm/topi/vision/nms_util.py
@@ -134,21 +134,23 @@ def collect_selected_indices(
num_total_detections=None,
input_image_size=None,
):
- """Collect selected indices from the core NMS loop into one linear output
+ """Collect selected indices from the core NMS loop into one linear output.
+
Parameters
----------
num_class : int
- selected_indices: tvm.te.Tensor
+ selected_indices : tvm.te.Tensor
2-D tensor with shape (batch_size * num_classes, num_boxes),
representing the indices
of selected boxes by the core NMS loop.
- num_detections tvm.te.Tensor
+ num_detections : tvm.te.Tensor
1-D tensor with shape (batch_size * num_classes,), representing
- the number of boxes selected by the core NMS loop, per batch and class
- row_offsets tvm.te.Tensor
+ the number of boxes selected by the core NMS loop, per batch and class.
+ row_offsets : tvm.te.Tensor
1-D tensor with shape (batch_size * num_classes,), this should be the
exclusive scan
- of num_detections
+ of num_detections.
ir : function
- A function to generate IR for CPU or GPU, see its usage in
vision/nms.py and cuda/nms.py
+ A function to generate IR for CPU or GPU, see its usage in
vision/nms.py and cuda/nms.py.
+
Returns
-------
out : tvm.te.Tensor
@@ -243,24 +245,27 @@ def collect_selected_indices(
def collect_selected_indices_and_scores(
selected_indices, selected_scores, num_detections, row_offsets,
num_total_detections, ir
):
- """Collect selected indices and scores from the core NMS loop into one
linear output
+ """Collect selected indices and scores from the core NMS loop into one
linear output.
+
Parameters
----------
- num_class : int
- selected_indices: tvm.te.Tensor
+ selected_indices : tvm.te.Tensor
2-D tensor with shape (batch_size * num_classes, num_boxes),
representing the indices
of selected boxes by the core NMS loop.
- selected_indices: tvm.te.Tensor
+ selected_scores : tvm.te.Tensor
2-D tensor with shape (batch_size * num_classes, num_boxes),
representing the scores
of selected boxes by the core NMS loop.
- num_detections tvm.te.Tensor
+ num_detections : tvm.te.Tensor
2-D tensor with shape (batch_size, num_classes), representing
- the number of boxes selected by the core NMS loop, per batch and class
- row_offsets tvm.te.Tensor
+ the number of boxes selected by the core NMS loop, per batch and class.
+ row_offsets : tvm.te.Tensor
2-D tensor with shape (batch_size, num_classes), this should be the
exclusive scan
- of num_detections along axis 1
+ of num_detections along axis 1.
+ num_total_detections : tvm.te.Tensor
+ Total number of detections.
ir : function
- A function to generate IR for CPU or GPU, see its usage in
vision/nms.py and cuda/nms.py
+ A function to generate IR for CPU or GPU, see its usage in
vision/nms.py and cuda/nms.py.
+
Returns
-------
out : [tvm.te.Tensor, tvm.te.Tensor]
@@ -379,28 +384,30 @@ def run_all_class_nms(
return_scores=False,
score_threshold=None,
):
- """The core all class NMS routine
+ """The core all class NMS routine.
+
Parameters
----------
boxes : tvm.te.Tensor
3-D tensor with shape (batch_size, num_boxes, 4)
- sorted_scores: tvm.te.Tensor
- 2-D tensor with shape (batch_size * num_classes, num_boxes)
- One of the outputs from argsort
- sorted_indices: tvm.te.Tensor
- 2-D tensor with shape (batch_size * num_classes, num_boxes)
- The other output from argsort
- valid_count: tvm.te.Tensor
+ sorted_scores : tvm.te.Tensor
+ 2-D tensor with shape (batch_size * num_classes, num_boxes).
+ One of the outputs from argsort.
+ sorted_indices : tvm.te.Tensor
+ 2-D tensor with shape (batch_size * num_classes, num_boxes).
+ The other output from argsort.
+ valid_count : tvm.te.Tensor
1-D tensor with shape (batch_size * num_classes,), representing
- the number of boxes whose score is above score_threshold, per batch
and class
- max_output_boxes_per_class : int or tvm.te.Tensor, optional
- The maxinum number of output selected boxes per class
- iou_threshold : float or tvm.te.Tensor, optionaIl
- IoU test threshold
+ the number of boxes whose score is above score_threshold, per batch
and class.
+ max_output_size_per_class : int or tvm.te.Tensor, optional
+ The maxinum number of output selected boxes per class.
+ iou_threshold : float or tvm.te.Tensor, optional
+ IoU test threshold.
nms_loop : function
- A core NMS loop, see its usage in vision/nms.py and cuda/nms.py
+ A core NMS loop, see its usage in vision/nms.py and cuda/nms.py.
return_scores : bool, optional
Whether or not to return selected scores, needed by the tensorflow
output format.
+
Returns
-------
out : a list of tvm.te.Tensor
