You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/08 13:58:36 UTC
(tvm) branch main updated: [REFACTOR][IR] Align native and Python operator family ownership (#20603)
This is an automated email from the ASF dual-hosted git repository.
tqchen 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 02be634cdf [REFACTOR][IR] Align native and Python operator family ownership (#20603)
02be634cdf is described below
commit 02be634cdf40657a0780e39fb35c05e769007973
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 8 09:58:28 2026 -0400
[REFACTOR][IR] Align native and Python operator family ownership (#20603)
Give native and Python operator families consistent owning modules while
preserving registry identities and call behavior. Existing C++ Op
getters use cached `const Op& xxx_op()` handles in their owning
namespaces; expression-building APIs retain their spelling.
Keep Relax attrs in dedicated `include/tvm/relax/op/*.h` family headers
and tile nodes and getters in `include/tvm/tirx/tile_op.h`. Other native
attrs and getters live in their owning `op.h`, with reflection
definitions beside registration. Ordinary operations precede region
operations, retaining nearby documentation and attrs.
Python operator packages own the corresponding attrs and primitive/tile
builders, and package/script exports reuse those definitions. Direct
consumers migrate from the former attrs/builtin headers and getter
names; Python call spelling and public argument order remain unchanged.
---
docs/reference/api/python/relax/op.rst | 8 +-
docs/reference/api/python/script/script.rst | 18 +-
docs/tirx/arch/tile_dispatch.rst | 2 +-
include/tvm/backend/cuda/op.h | 163 ++++
include/tvm/ir/prim/builtin.h | 58 --
include/tvm/ir/prim/op.h | 38 +-
include/tvm/relax/attrs/image.h | 185 -----
include/tvm/relax/attrs/manipulate.h | 324 --------
include/tvm/relax/attrs/nn.h | 874 ---------------------
include/tvm/relax/attrs/sorting.h | 109 ---
include/tvm/relax/attrs/vision.h | 197 -----
include/tvm/relax/{attrs => op}/ccl.h | 45 +-
include/tvm/relax/{attrs => op}/create.h | 22 +-
include/tvm/relax/{attrs => op}/datatype.h | 19 +-
include/tvm/relax/{attrs => op}/distributed.h | 17 +-
include/tvm/relax/op/image.h | 86 ++
include/tvm/relax/{attrs => op}/index.h | 33 +-
include/tvm/relax/{attrs => op}/linear_algebra.h | 21 +-
include/tvm/relax/op/manipulate.h | 200 +++++
include/tvm/relax/op/nn.h | 345 ++++++++
include/tvm/relax/{attrs => op}/op.h | 45 +-
include/tvm/relax/{attrs => op}/qdq.h | 18 +-
include/tvm/relax/{attrs => op}/sampling.h | 15 +-
include/tvm/relax/{attrs => op}/search.h | 31 +-
include/tvm/relax/op/sorting.h | 67 ++
include/tvm/relax/{attrs => op}/statistical.h | 36 +-
include/tvm/relax/op/vision.h | 112 +++
include/tvm/s_tir/transform.h | 2 +-
include/tvm/tirx/attrs.h | 73 --
include/tvm/tirx/builtin.h | 757 ------------------
include/tvm/tirx/op.h | 753 +++++++++++++++++-
include/tvm/tirx/script/ir_builder/ir.h | 2 +-
include/tvm/tirx/stmt_functor.h | 2 +-
include/tvm/tirx/{tile_primitive.h => tile_op.h} | 83 +-
include/tvm/topi/detail/extern.h | 15 +-
include/tvm/topi/elemwise.h | 3 +-
python/tvm/backend/cuda/op.py | 14 +-
python/tvm/backend/cuda/script.py | 8 +
.../cuda/tile_primitive/elementwise/__init__.py | 5 +-
python/tvm/backend/metal/op.py | 33 +-
python/tvm/backend/metal/script.py | 27 +-
.../trn/tile_primitive/compose_op/binary_chain.py | 2 +-
.../trn/tile_primitive/compose_op/binary_reduce.py | 2 +-
.../trn/tile_primitive/compose_op/reduce_negate.py | 2 +-
.../trn/tile_primitive/compose_op/unary_reduce.py | 2 +-
.../backend/trn/tile_primitive/private_alloc.py | 2 +-
.../backend/trn/tile_primitive/select/default.py | 2 +-
python/tvm/ir/prim/__init__.py | 30 +-
python/tvm/ir/prim/op.py | 483 ++++++++++++
python/tvm/relax/op/__init__.py | 101 ++-
python/tvm/relax/op/base.py | 27 +-
python/tvm/relax/op/ccl/__init__.py | 10 +-
python/tvm/relax/op/ccl/ccl.py | 18 +
python/tvm/relax/op/create.py | 14 +-
python/tvm/relax/op/datatype.py | 14 +-
python/tvm/relax/op/image/__init__.py | 8 +-
python/tvm/relax/op/image/image.py | 9 +-
python/tvm/relax/op/index.py | 13 +
python/tvm/relax/op/linear_algebra.py | 13 +
python/tvm/relax/op/manipulate.py | 99 ++-
python/tvm/relax/op/nn/__init__.py | 26 +
python/tvm/relax/op/nn/nn.py | 133 ++++
python/tvm/relax/op/op_attrs.py | 401 ----------
python/tvm/relax/op/qdq.py | 8 +
python/tvm/relax/op/sampling.py | 8 +
python/tvm/relax/op/search.py | 8 +
python/tvm/relax/op/sorting.py | 19 +-
python/tvm/relax/op/statistical.py | 14 +-
.../tvm/relax/op/vision/multibox_transform_loc.py | 8 +
python/tvm/relax/op/vision/nms.py | 18 +
python/tvm/relax/op/vision/roi_align.py | 8 +
python/tvm/relax/op/vision/roi_pool.py | 8 +
python/tvm/relax/script/ir_builder/op.py | 80 ++
.../tvm/relax/transform/legalize_ops/inspect_op.py | 2 +-
python/tvm/tirx/{op.py => op/__init__.py} | 528 +------------
.../{operator/tile_primitive/ops.py => op/tile.py} | 0
.../tvm/tirx/operator/tile_primitive/__init__.py | 4 +-
python/tvm/tirx/script/ir_builder/op.py | 3 +
.../tvm/tirx/script/ir_builder/parser_protocol.py | 20 +-
python/tvm/tirx/tile_primitive.py | 2 +-
python/tvm/tirx/transform/transform.py | 2 +-
src/backend/cuda/codegen/codegen_cuda.cc | 32 +-
src/backend/cuda/codegen/codegen_cuda_host.cc | 13 +-
src/backend/cuda/codegen/intrin_rule_cuda.cc | 12 +-
src/backend/cuda/codegen/llvm/codegen_nvptx.cc | 12 +-
src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc | 5 +-
src/backend/cuda/op/attrs.h | 108 ---
src/backend/cuda/op/iket.cc | 37 +
src/backend/cuda/op/target_builtin.cc | 59 +-
src/backend/cuda/transforms/lower_iket.cc | 83 +-
.../hexagon/codegen/llvm/intrin_rule_hexagon.cc | 4 +-
src/backend/metal/codegen/codegen_metal.cc | 14 +-
src/backend/metal/codegen/intrin_rule_metal.cc | 8 +-
src/backend/opencl/codegen/codegen_opencl.cc | 17 +-
src/backend/opencl/codegen/intrin_rule_opencl.cc | 4 +-
src/backend/rocm/codegen/llvm/codegen_amdgpu.cc | 4 +-
src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc | 17 +-
src/backend/trn/codegen/codegen_trn.cc | 6 +-
src/backend/trn/op/target_builtin.cc | 3 +-
src/backend/trn/transform/lower_trainium_layout.cc | 14 +-
src/backend/vulkan/codegen/codegen_spirv.cc | 35 +-
src/backend/vulkan/codegen/intrin_rule_spirv.cc | 5 +-
src/backend/webgpu/codegen/codegen_webgpu.cc | 26 +-
src/backend/webgpu/codegen/intrin_rule_webgpu.cc | 6 +-
src/ir/prim/builtin.cc | 16 +-
src/ir/prim/op.cc | 23 +-
src/ir/prim/vector_expr.cc | 7 +-
.../backend/adreno/annotate_custom_storage.cc | 2 +-
.../backend/adreno/fold_vdevice_scope_change.cc | 2 +-
src/relax/backend/contrib/cutlass/codegen.cc | 2 +-
src/relax/backend/contrib/nnapi/codegen.cc | 2 +-
src/relax/backend/contrib/tensorrt/codegen.cc | 6 +-
src/relax/backend/vm/lower_runtime_builtin.cc | 2 +-
src/relax/distributed/axis_group_graph.cc | 10 +-
.../distributed/transform/legalize_redistribute.cc | 4 +-
src/relax/distributed/transform/lower_distir.cc | 2 +-
.../transform/lower_global_view_to_local_view.cc | 2 +-
.../distributed/transform/propagate_sharding.cc | 8 +-
src/relax/op/ccl/ccl.cc | 37 +
src/relax/op/ccl/ccl.h | 2 +-
src/relax/op/distributed/ccl.cc | 2 +-
src/relax/op/distributed/distributed.cc | 11 +-
src/relax/op/distributed/distributed.h | 2 +-
src/relax/op/image/resize.cc | 107 +++
src/relax/op/image/resize.h | 2 +-
src/relax/op/nn/attention.cc | 15 +
src/relax/op/nn/attention.h | 2 +-
src/relax/op/nn/convolution.cc | 250 ++++++
src/relax/op/nn/convolution.h | 2 +-
src/relax/op/nn/nn.cc | 146 ++++
src/relax/op/nn/nn.h | 2 +-
src/relax/op/nn/pooling.cc | 159 ++++
src/relax/op/nn/pooling.h | 2 +-
src/relax/op/op.cc | 38 +-
src/relax/op/tensor/create.cc | 14 +
src/relax/op/tensor/create.h | 2 +-
src/relax/op/tensor/datatype.cc | 11 +
src/relax/op/tensor/datatype.h | 2 +-
src/relax/op/tensor/grad.cc | 7 +
src/relax/op/tensor/grad.h | 4 +-
src/relax/op/tensor/index.cc | 18 +
src/relax/op/tensor/index.h | 2 +-
src/relax/op/tensor/inspect.cc | 47 +-
src/relax/op/tensor/linear_algebra.cc | 13 +
src/relax/op/tensor/linear_algebra.h | 2 +-
src/relax/op/tensor/manipulate.cc | 157 ++++
src/relax/op/tensor/manipulate.h | 2 +-
src/relax/op/tensor/qdq.cc | 10 +
src/relax/op/tensor/qdq.h | 2 +-
src/relax/op/tensor/sampling.cc | 7 +
src/relax/op/tensor/sampling.h | 2 +-
src/relax/op/tensor/search.cc | 23 +
src/relax/op/tensor/search.h | 2 +-
src/relax/op/tensor/sorting.cc | 48 ++
src/relax/op/tensor/sorting.h | 2 +-
src/relax/op/tensor/statistical.cc | 28 +
src/relax/op/tensor/statistical.h | 2 +-
src/relax/op/vision/multibox_transform_loc.cc | 19 +
src/relax/op/vision/multibox_transform_loc.h | 2 +-
src/relax/op/vision/nms.cc | 51 +-
src/relax/op/vision/nms.h | 2 +-
src/relax/op/vision/roi_align.cc | 18 +
src/relax/op/vision/roi_align.h | 2 +-
src/relax/op/vision/roi_pool.cc | 10 +
src/relax/op/vision/roi_pool.h | 2 +-
src/relax/script/ir_builder/distributed.cc | 2 +-
src/relax/script/printer/call.cc | 2 +-
src/relax/transform/adjust_matmul_order.cc | 4 +-
src/relax/transform/alter_op_impl.cc | 2 +-
src/relax/transform/call_tir_rewrite.cc | 2 +-
src/relax/transform/canonicalize_bindings.cc | 2 +-
src/relax/transform/compute_prim_value.cc | 4 +-
src/relax/transform/dataflow_inplace.cc | 2 +-
src/relax/transform/decompose_ops.cc | 2 +-
src/relax/transform/fuse_tir.cc | 2 +-
src/relax/transform/gradient.cc | 2 +-
src/relax/transform/gradient_simplifier.cc | 2 +-
src/relax/transform/infer_amp_utils.h | 2 +-
src/relax/transform/infer_layout_utils.h | 14 +-
src/relax/transform/realize_vdevice.cc | 2 +-
.../specialize_function_based_on_callsite.cc | 2 +-
src/relax/transform/split_call_tir_by_pattern.cc | 3 +-
src/relax/utils.cc | 2 +-
src/s_tir/analysis/calculate_allocated_memory.cc | 4 +-
src/s_tir/analysis/conditional_bounds.cc | 4 +-
src/s_tir/analysis/is_pure_function.cc | 4 +-
.../analysis/sblock_access_region_detector.cc | 14 +-
.../analysis/sblock_buffer_access_lca_detector.cc | 4 +-
src/s_tir/analysis/verify_gpu_code.cc | 6 +-
src/s_tir/backend/adreno/inject_texture_alloc.cc | 12 +-
src/s_tir/backend/adreno/texture_flatten.cc | 10 +-
src/s_tir/schedule/analysis/reducer.cc | 4 +-
src/s_tir/schedule/ir_comparator.cc | 10 +-
src/s_tir/schedule/primitive/cache_read_write.cc | 6 +-
src/s_tir/schedule/primitive/decompose_padding.cc | 4 +-
.../schedule/primitive/layout_transformation.cc | 5 +-
src/s_tir/schedule/transform.cc | 4 +-
src/s_tir/transform/annotate_irregular_loop.cc | 4 +-
src/s_tir/transform/compact_buffer_region.cc | 15 +-
src/s_tir/transform/hoist_expression.cc | 10 +-
src/s_tir/transform/inject_permuted_layout.cc | 4 +-
src/s_tir/transform/inject_ptx_async_copy.cc | 7 +-
src/s_tir/transform/inject_software_pipeline.cc | 10 +-
src/s_tir/transform/inject_virtual_thread.cc | 36 +-
src/s_tir/transform/ir_utils.cc | 3 +-
src/s_tir/transform/loop_partition.cc | 24 +-
.../transform/lower_cross_thread_reduction.cc | 2 +-
src/s_tir/transform/lower_match_buffer.cc | 5 +-
src/s_tir/transform/lower_opaque_block.cc | 2 +-
src/s_tir/transform/lower_thread_binding.cc | 7 +-
src/s_tir/transform/lower_vtcm_alloc.cc | 25 +-
.../transform/memhammer_tensorcore_rewrite.cc | 10 +-
.../transform/merge_shared_memory_allocations.cc | 31 +-
src/s_tir/transform/remove_store_undef.cc | 15 +-
src/s_tir/transform/rewrite_unsafe_select.cc | 10 +-
src/s_tir/transform/storage_access.cc | 22 +-
src/s_tir/transform/tensorcore_infer_fragment.cc | 6 +-
src/s_tir/transform/thread_storage_sync.cc | 11 +-
.../transform/using_assume_to_reduce_branches.cc | 9 +-
src/script/printer/script_printer.cc | 2 +-
src/sym/const_int_bound.cc | 6 +-
src/sym/constraint_helpers.h | 2 +-
src/sym/modular_set.cc | 2 +-
src/sym/pattern_match.h | 3 +-
src/sym/rewrite_simplify.cc | 17 +-
src/sym/simplify_base.cc | 3 +-
src/sym/z3_prover.cc | 4 +-
src/target/intrin_rule.h | 6 +-
src/target/llvm/codegen_cpu.cc | 60 +-
src/target/llvm/codegen_llvm.cc | 57 +-
src/target/llvm/codegen_llvm.h | 10 +-
src/target/llvm/codegen_x86_64.cc | 5 +-
src/target/llvm/intrin_rule_llvm.cc | 2 +-
src/target/llvm/intrin_rule_llvm.h | 8 +-
src/target/source/codegen_c.cc | 89 ++-
src/target/source/codegen_c.h | 8 +-
src/target/source/codegen_c_host.cc | 16 +-
src/te/operation/compute_op.cc | 4 +-
src/tirx/analysis/filter_canonical.cc | 4 +-
src/tirx/analysis/var_use_def_analysis.cc | 4 +-
src/tirx/analysis/verify_memory.cc | 10 +-
src/tirx/analysis/verify_tirx_well_formed.cc | 6 +-
src/tirx/analysis/verify_well_formed.h | 2 +-
src/tirx/ir/buffer.cc | 9 +-
src/tirx/ir/data_type_rewriter.cc | 13 +-
src/tirx/ir/exec_scope.cc | 5 +-
src/tirx/ir/ir_mutator_with_analyzer.cc | 3 +-
src/tirx/ir/ir_visitor_with_analyzer.cc | 7 +-
src/tirx/ir/specialize.cc | 7 +-
src/tirx/ir/stmt.cc | 11 +-
src/tirx/ir/stmt_functor.cc | 2 +-
src/tirx/ir/tir_visitor_with_path.h | 2 +-
src/tirx/ir/tirx_stmt.cc | 2 +-
src/tirx/op/builtin.cc | 222 +++---
src/tirx/op/op.cc | 25 +-
src/tirx/op/tirx.cc | 76 +-
src/tirx/script/ir_builder/frame.cc | 3 +-
src/tirx/script/ir_builder/ir.cc | 15 +-
src/tirx/script/printer/buffer.cc | 6 +-
src/tirx/script/printer/expr.cc | 9 +-
src/tirx/script/printer/stmt.cc | 20 +-
src/tirx/transform/bind_target.cc | 8 +-
src/tirx/transform/common_subexpr_elim.cc | 6 +-
src/tirx/transform/dtype_conversion.h | 3 +-
src/tirx/transform/flatten_buffer.cc | 19 +-
src/tirx/transform/force_narrow_index_to_i32.h | 2 +-
src/tirx/transform/inline_private_functions.cc | 3 +-
src/tirx/transform/ir_utils.h | 20 +-
src/tirx/transform/lower_intrin.cc | 17 +-
src/tirx/transform/lower_thread_allreduce.h | 49 +-
src/tirx/transform/lower_tirx_cleanup.cc | 19 +-
src/tirx/transform/lower_tirx_dedup_tensormap.cc | 8 +-
src/tirx/transform/lower_tirx_opaque.cc | 7 +-
src/tirx/transform/lower_tvm_builtin.cc | 121 ++-
src/tirx/transform/lower_warp_memory.cc | 24 +-
src/tirx/transform/make_packed_api.cc | 37 +-
src/tirx/transform/narrow_datatype.cc | 5 +-
src/tirx/transform/remap_thread_axis.cc | 8 +-
src/tirx/transform/remove_assume.cc | 9 +-
src/tirx/transform/split_host_device.cc | 18 +-
src/tirx/transform/stmt_simplify.cc | 7 +-
src/tirx/transform/storage_rewrite.cc | 70 +-
src/tirx/transform/tile_primitive_dispatch.cc | 47 +-
src/tirx/transform/tvm_ffi_binder.cc | 74 +-
src/tirx/transform/unsupported_dtype_legalize.cc | 31 +-
src/tirx/transform/update_pointer_storage_scope.cc | 7 +-
src/tirx/transform/vectorize_loop.cc | 31 +-
tests/cpp/ir_functor_test.cc | 21 +-
tests/cpp/pattern_match_test.cc | 2 +-
tests/cpp/tir_analysis_side_effect.cc | 10 +-
tests/cpp/tir_scalable_datatype.cc | 6 +-
tests/python/relax/test_analysis_well_formed.py | 2 +-
tests/python/script/test_constructor_contracts.py | 14 +-
.../tile_primitive/cuda/copy_async/test_dsmem.py | 2 +-
.../tile_primitive/cuda/copy_async/test_tma.py | 2 +-
.../cuda/gemm_async/test_gemm_async.py | 4 +-
tests/python/tirx/script/test_tirx_script_gpu.py | 2 +-
297 files changed, 5599 insertions(+), 5215 deletions(-)
diff --git a/docs/reference/api/python/relax/op.rst b/docs/reference/api/python/relax/op.rst
index f6a115d1fd..dbc283f674 100644
--- a/docs/reference/api/python/relax/op.rst
+++ b/docs/reference/api/python/relax/op.rst
@@ -79,7 +79,7 @@ tvm.relax.op.vision
.. automodule:: tvm.relax.op.vision
:members:
:imported-members:
- :exclude-members: Expr
+ :exclude-members: Expr, Attrs
tvm.relax.op.vm
***************
@@ -87,9 +87,3 @@ tvm.relax.op.vm
:members:
:imported-members:
:exclude-members: Expr
-
-tvm.relax.op.op_attrs
-*********************
-.. automodule:: tvm.relax.op.op_attrs
- :members:
- :exclude-members: Attrs
diff --git a/docs/reference/api/python/script/script.rst b/docs/reference/api/python/script/script.rst
index 5c77edebf8..907dd77563 100644
--- a/docs/reference/api/python/script/script.rst
+++ b/docs/reference/api/python/script/script.rst
@@ -27,9 +27,20 @@ tvm.script
tvm.script.relax
****************
+
+Operator attributes re-exported by this namespace are documented in :mod:`tvm.relax.op`.
+
.. automodule:: tvm.script.relax
:members:
- :exclude-members: ExternFunc, ShapeExpr, TupleGetItem, Range, function, macro
+ :exclude-members: ExternFunc, ShapeExpr, TupleGetItem, Range, function, macro,
+ ArgmaxArgminAttrs, ArgsortAttrs, AstypeAttrs, CallInplacePackedAttrs, CallTIRInplaceAttrs,
+ CallTIRWithGradAttrs, ConcatAttrs, EinsumAttrs, ExpandDimsAttrs, FlipAttrs,
+ GatherElementsAttrs, GatherNDAttrs, HintOnDeviceAttrs, IndexPutAttrs, InitAttrs,
+ LayoutTransformAttrs, MatmulAttrs, MeshgridAttrs, MultinomialFromUniformAttrs,
+ OneHotAttrs, PermuteDimsAttrs, QuantizeAttrs, RepeatAttrs, ReverseSequenceAttrs,
+ ScanopAttrs, ScatterElementsAttrs, ScatterNDAttrs, SliceScatterAttrs, SortAttrs,
+ SplitAttrs, SqueezeAttrs, StackAttrs, StatisticalAttrs, StridedSliceAttrs, TakeAttrs,
+ TileAttrs, TopKAttrs, ToVDeviceAttrs, TriluAttrs, WrapParamAttrs
.. autofunction:: tvm.script.relax.function
@@ -37,9 +48,12 @@ tvm.script.relax
tvm.script.tirx
***************
+
+Operator attributes re-exported by this namespace are documented in :mod:`tvm.tirx`.
+
.. automodule:: tvm.script.tirx
:members:
- :exclude-members: Range, meta_var, Var, Call, CommReducer, Reduce, SMEMPool, TMEMPool, FloatImm, IntImm, Cast, Add, Sub, Mul, Div, Mod, FloorDiv, FloorMod, LShift, RShift, BitwiseAnd, BitwiseOr, BitwiseXor, BitwiseNot, Min, Max, EQ, NE, LT, LE, GT, GE, And, Or, Not, Select, Ramp, Broadcast, Shuffle, CallEffectKind, IterVar, ComposeLayout, DtypeConstructor, ExecScope, Iter, Layout, LetAnnotation, LocalVectorAnnotation, ScopeIdDef, TileLayout, function, jit, inline, macro
+ :exclude-members: Range, meta_var, Var, Call, CommReducer, Reduce, SMEMPool, TMEMPool, FloatImm, IntImm, Cast, Add, Sub, Mul, Div, Mod, FloorDiv, FloorMod, LShift, RShift, BitwiseAnd, BitwiseOr, BitwiseXor, BitwiseNot, Min, Max, EQ, NE, LT, LE, GT, GE, And, Or, Not, Select, Ramp, Broadcast, Shuffle, CallEffectKind, IterVar, ComposeLayout, DtypeConstructor, ExecScope, Iter, Layout, LetAnnotation, LocalVectorAnnotation, ScopeIdDef, TileLayout, CallFFIKernelAttr, TensorMapEncodeTiledAttr [...]
.. autofunction:: tvm.script.tirx.function
diff --git a/docs/tirx/arch/tile_dispatch.rst b/docs/tirx/arch/tile_dispatch.rst
index b0d8d7d591..d5ddf8f9b9 100644
--- a/docs/tirx/arch/tile_dispatch.rst
+++ b/docs/tirx/arch/tile_dispatch.rst
@@ -30,7 +30,7 @@ Implementation surfaces
The authoritative operation list is the C++ registry
(``src/tirx/op/tirx.cc``, with operations named ``tirx.tile.<name>``). IR
wrapper classes live in
-``python/tvm/tirx/operator/tile_primitive/ops.py``. Raw
+``python/tvm/tirx/op/tile.py``. Raw
``TilePrimitiveCall`` constructors live in
``python/tvm/tirx/script/ir_builder/tirx.py``, while the validated authoring facade
is in ``python/tvm/tirx/script/tile.py``. Both Python construction surfaces
diff --git a/include/tvm/backend/cuda/op.h b/include/tvm/backend/cuda/op.h
new file mode 100644
index 0000000000..cedf4e69ae
--- /dev/null
+++ b/include/tvm/backend/cuda/op.h
@@ -0,0 +1,163 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*!
+ * \file tvm/backend/cuda/op.h
+ * \brief CUDA-owned operators.
+ */
+#ifndef TVM_BACKEND_CUDA_OP_H_
+#define TVM_BACKEND_CUDA_OP_H_
+
+#include <tvm/ir/attrs.h>
+#include <tvm/ir/op.h>
+
+namespace tvm {
+namespace backend {
+namespace cuda {
+
+// tcgen05 instruction descriptor attributes.
+
+/*! \brief Static options for the dense tcgen05 instruction descriptor. */
+struct TCGen05InstrDescriptorAttrs : public AttrsNode {
+ ffi::String d_dtype;
+ ffi::String a_dtype;
+ ffi::String b_dtype;
+ int64_t M;
+ int64_t N;
+ int64_t K;
+ bool trans_a;
+ bool trans_b;
+ int64_t n_cta_groups = 1;
+ bool neg_a = false;
+ bool neg_b = false;
+ bool sat_d = false;
+ bool is_sparse = false;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.cuda.TCGen05InstrDescriptorAttrs",
+ TCGen05InstrDescriptorAttrs, AttrsNode);
+};
+
+/*! \brief Static options for the block-scaled tcgen05 instruction descriptor. */
+struct TCGen05InstrDescriptorBlockScaledAttrs : public AttrsNode {
+ ffi::String d_dtype;
+ ffi::String a_dtype;
+ ffi::String b_dtype;
+ ffi::String sfa_dtype;
+ ffi::String sfb_dtype;
+ int64_t M;
+ int64_t N;
+ int64_t K;
+ bool trans_a;
+ bool trans_b;
+ int64_t n_cta_groups = 1;
+ bool neg_a = false;
+ bool neg_b = false;
+ bool is_sparse = false;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.cuda.TCGen05InstrDescriptorBlockScaledAttrs",
+ TCGen05InstrDescriptorBlockScaledAttrs, AttrsNode);
+};
+
+/*!
+ * \name Frontend-only NVIDIA IKET annotations
+ *
+ * These CUDA device intrinsics have opaque call effects and no registered attrs
+ * schema or attrs defaults. Marker, range-start, range-end and range-push
+ * annotations have a registered variadic operand tail; the Python convenience
+ * builders and IKET lowering use it for one optional scalar payload. Below, C denotes
+ * tvm.backend.cuda.op and T is the TIRx script namespace.
+ *
+ * Match a structurally valid marker call and inspect its operands:
+ * \code
+ * if (const auto* call = expr.as<tvm::CallNode>();
+ * call && call->op.same_as(tvm::backend::cuda::iket_mark_op())) {
+ * tvm::Expr name = call->args[0];
+ * if (call->args.size() > 1) {
+ * tvm::Expr payload = call->args[1];
+ * }
+ * tvm::Type result_type = call->ty;
+ * }
+ * \endcode
+ * The other getters use the same identity-matching pattern with the operand
+ * layouts documented below.
+ * \{
+ */
+
+/*!
+ * \brief Get the NVIDIA IKET marker operator, tirx.cuda.iket_mark.
+ * Python: C.cuda_iket_mark(name, payload=None).
+ * Script: T.cuda.iket.mark(name, payload=None).
+ * Operands: [0] name, the marker name; [1...] args, the optional payload tail.
+ * Result: void. Effects: opaque marker annotation.
+ */
+TVM_DLL const Op& iket_mark_op();
+
+/*!
+ * \brief Get the NVIDIA IKET token-range start operator, tirx.cuda.iket_range_start.
+ * Python: C.cuda_iket_range_start(name, payload=None).
+ * Script: T.cuda.iket.range_start(name, payload=None).
+ * Operands: [0] name, the range name; [1...] args, the optional payload tail.
+ * Result: uint32 range token. Effects: opaque range-start annotation.
+ */
+TVM_DLL const Op& iket_range_start_op();
+
+/*!
+ * \brief Get the NVIDIA IKET token-range end operator, tirx.cuda.iket_range_end.
+ * Python: C.cuda_iket_range_end(token, payload=None).
+ * Script: T.cuda.iket.range_end(token, payload=None).
+ * Operands: [0] token, an IntExpr identifying the range to end;
+ * [1...] args, the optional payload tail. IKET lowering requires a uint32 token.
+ * Result: void. Effects: opaque range-end annotation.
+ */
+TVM_DLL const Op& iket_range_end_op();
+
+/*!
+ * \brief Get the NVIDIA IKET stack-range push operator, tirx.cuda.iket_range_push.
+ * Python: C.cuda_iket_range_push(name, payload=None).
+ * Script: T.cuda.iket.range_push(name, payload=None).
+ * Operands: [0] name, the range name; [1...] args, the optional payload tail.
+ * Result: void. Effects: opaque stack-range push annotation.
+ */
+TVM_DLL const Op& iket_range_push_op();
+
+/*!
+ * \brief Get the NVIDIA IKET stack-range pop operator, tirx.cuda.iket_range_pop.
+ * Python: C.cuda_iket_range_pop(). Script: T.cuda.iket.range_pop().
+ * Operands: none; IKET lowering enforces the empty operand list.
+ * Result: void. Effects: opaque stack-range pop annotation.
+ */
+TVM_DLL const Op& iket_range_pop_op();
+
+/*!
+ * \brief Get the NVIDIA IKET sentinel-token operator, tirx.cuda.iket_sentinel_token.
+ * Python: C.cuda_iket_sentinel_token(name). Script: T.cuda.iket.sentinel_token(name).
+ * Operands: [0] name, the event name.
+ * Result: uint32 sentinel token. Registered effects: opaque. The sentinel carries
+ * token-flow identity; IKET lowering emits no runtime event for it.
+ */
+TVM_DLL const Op& iket_sentinel_token_op();
+
+/*! \} */
+
+} // namespace cuda
+} // namespace backend
+} // namespace tvm
+
+#endif // TVM_BACKEND_CUDA_OP_H_
diff --git a/include/tvm/ir/prim/builtin.h b/include/tvm/ir/prim/builtin.h
deleted file mode 100644
index 0e423e2803..0000000000
--- a/include/tvm/ir/prim/builtin.h
+++ /dev/null
@@ -1,58 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-#ifndef TVM_IR_PRIM_BUILTIN_H_
-#define TVM_IR_PRIM_BUILTIN_H_
-
-#include <tvm/ir/op.h>
-
-namespace tvm {
-namespace prim {
-namespace builtin {
-
-/*!
- * \brief Get the target's vscale value. It will be lowered to llvm.vscale intrinsic
- * (https://llvm.org/docs/LangRef.html#llvm-vscale-intrinsic)
- */
-TVM_DLL const Op& vscale();
-
-/*! \brief Round up to the nearest integral value. */
-TVM_DLL const Op& ceil();
-
-/*! \brief Base-two logarithm. */
-TVM_DLL const Op& log2();
-
-/*! \brief Count leading zero bits. */
-TVM_DLL const Op& clz();
-
-/*!
- * \brief Same as select, used for unsafe memory access.
- *
- * Type tvm_if_then_else(cond, a, b) {
- * return cond ? a : b;
- * }
- */
-TVM_DLL const Op& if_then_else();
-
-/*! \brief Marks a condition is likely going to happen. */
-TVM_DLL const Op& likely();
-
-} // namespace builtin
-} // namespace prim
-} // namespace tvm
-#endif // TVM_IR_PRIM_BUILTIN_H_
diff --git a/include/tvm/ir/prim/op.h b/include/tvm/ir/prim/op.h
index 095f60b186..06b3740d2e 100644
--- a/include/tvm/ir/prim/op.h
+++ b/include/tvm/ir/prim/op.h
@@ -22,7 +22,7 @@
*/
#ifndef TVM_IR_PRIM_OP_H_
#define TVM_IR_PRIM_OP_H_
-#include <tvm/ir/prim/builtin.h>
+#include <tvm/ir/op.h>
#include <tvm/ir/prim/expr.h>
#include <algorithm>
@@ -30,6 +30,37 @@
#include <type_traits>
#include <utility>
namespace tvm {
+namespace prim {
+
+/*!
+ * \brief Get the target's vscale value. It will be lowered to llvm.vscale intrinsic
+ * (https://llvm.org/docs/LangRef.html#llvm-vscale-intrinsic)
+ */
+TVM_DLL const Op& vscale_op();
+
+/*! \brief Round up to the nearest integral value. */
+TVM_DLL const Op& ceil_op();
+
+/*! \brief Base-two logarithm. */
+TVM_DLL const Op& log2_op();
+
+/*! \brief Count leading zero bits. */
+TVM_DLL const Op& clz_op();
+
+/*!
+ * \brief Same as select, used for unsafe memory access.
+ *
+ * Type tvm_if_then_else(cond, a, b) {
+ * return cond ? a : b;
+ * }
+ */
+TVM_DLL const Op& if_then_else_op();
+
+/*! \brief Marks a condition is likely going to happen. */
+TVM_DLL const Op& likely_op();
+
+} // namespace prim
+
/*!
* \brief add operator
*
@@ -527,9 +558,8 @@ inline PrimExpr MakeConst(PrimType dtype, ValueType value, Span span) {
if (dtype.IsFixedLengthVector()) {
return prim::Broadcast(MakeConstScalar(elem_ty, value, span), dtype.lanes(), span);
}
- PrimExpr lanes =
- prim::Mul(Call(PrimType::Int(32), prim::builtin::vscale(), {}).as_or_throw<PrimExpr>(),
- dtype.VScaleFactor());
+ PrimExpr lanes = prim::Mul(Call(PrimType::Int(32), prim::vscale_op(), {}).as_or_throw<PrimExpr>(),
+ dtype.VScaleFactor());
return prim::Broadcast(MakeConstScalar(elem_ty, value, span), lanes, span);
}
diff --git a/include/tvm/relax/attrs/image.h b/include/tvm/relax/attrs/image.h
deleted file mode 100644
index 7e3f8c27d1..0000000000
--- a/include/tvm/relax/attrs/image.h
+++ /dev/null
@@ -1,185 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file tvm/relax/attrs/image.h
- * \brief Attributes for image operators.
- */
-#ifndef TVM_RELAX_ATTRS_IMAGE_H_
-#define TVM_RELAX_ATTRS_IMAGE_H_
-
-#include <tvm/relax/expr.h>
-
-namespace tvm {
-namespace relax {
-
-/*! \brief Attributes used in image resize2d operator */
-struct Resize2DAttrs : public AttrsNode {
- ffi::Array<FloatImm> roi;
- ffi::String layout;
- ffi::String method;
- ffi::String coordinate_transformation_mode;
- ffi::String rounding_method;
- double cubic_alpha;
- int cubic_exclude;
- double extrapolation_value;
- ffi::Optional<DLDataType> out_dtype;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<Resize2DAttrs>()
- .def_ro("roi", &Resize2DAttrs::roi,
- "Region of Interest for coordinate transformation mode 'tf_crop_and_resize'")
- .def_ro("layout", &Resize2DAttrs::layout,
- "Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
- "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
- "dimensions respectively. Resize is applied on the 'H' and"
- "'W' dimensions.",
- refl::DefaultValue(ffi::String("NCHW")))
- .def_ro("method", &Resize2DAttrs::method,
- "Specify the mode to use for scaling."
- "nearest_neighbor - Nearest Neighbor"
- "linear - Bilinear Interpolation"
- "cubic - Bicubic Interpolation",
- refl::DefaultValue(ffi::String("linear")))
- .def_ro("coordinate_transformation_mode", &Resize2DAttrs::coordinate_transformation_mode,
- "Describes how to transform the coordinate in the resized tensor"
- "to the coordinate in the original tensor."
- "Refer to the ONNX Resize operator specification for details"
- "Available options are half_pixel, align_corners and asymmetric",
- refl::DefaultValue(ffi::String("half_pixel")))
- .def_ro("rounding_method", &Resize2DAttrs::rounding_method,
- "indicates how to find the \"nearest\" pixel in nearest_neighbor method"
- "Available options are round, floor, and ceil.",
- refl::DefaultValue(ffi::String("round")))
- .def_ro("cubic_alpha", &Resize2DAttrs::cubic_alpha,
- "Spline Coefficient for Bicubic Interpolation", refl::DefaultValue(-0.75))
- .def_ro("cubic_exclude", &Resize2DAttrs::cubic_exclude,
- "Flag to exclude exterior of the image during bicubic interpolation",
- refl::DefaultValue(0))
- .def_ro("extrapolation_value", &Resize2DAttrs::extrapolation_value,
- "Value to return when roi is outside of the image", refl::DefaultValue(0.0))
- .def_ro(
- "out_dtype", &Resize2DAttrs::out_dtype,
- "The dtype of the output tensor. It it is not specified, the output will have the same "
- "dtype as input if not specified.",
- refl::DefaultValue(ffi::Optional<DLDataType>{}));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Resize2DAttrs", Resize2DAttrs, AttrsNode);
-}; // struct Resize2dAttrs
-
-/*! \brief Attributes used in image resize3d operator */
-struct Resize3DAttrs : public AttrsNode {
- ffi::Array<FloatImm> roi;
- ffi::String layout;
- ffi::String method;
- ffi::String coordinate_transformation_mode;
- ffi::String rounding_method;
- double cubic_alpha;
- int cubic_exclude;
- double extrapolation_value;
- ffi::Optional<DLDataType> out_dtype;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<Resize3DAttrs>()
- .def_ro("roi", &Resize3DAttrs::roi,
- "Region of Interest for coordinate transformation mode 'tf_crop_and_resize'")
- .def_ro("layout", &Resize3DAttrs::layout,
- "Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
- "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
- "dimensions respectively. Resize is applied on the 'D', 'H' and"
- "'W' dimensions.",
- refl::DefaultValue(ffi::String("NCDHW")))
- .def_ro("method", &Resize3DAttrs::method,
- "Specify the mode to use for scaling."
- "nearest_neighbor - Nearest Neighbor"
- "linear - Trilinear Interpolation"
- "cubic - Tricubic Interpolation",
- refl::DefaultValue(ffi::String("linear")))
- .def_ro("coordinate_transformation_mode", &Resize3DAttrs::coordinate_transformation_mode,
- "Describes how to transform the coordinate in the resized tensor"
- "to the coordinate in the original tensor."
- "Refer to the ONNX Resize operator specification for details"
- "Available options are half_pixel, align_corners and asymmetric",
- refl::DefaultValue(ffi::String("half_pixel")))
- .def_ro("rounding_method", &Resize3DAttrs::rounding_method,
- "indicates how to find the \"nearest\" pixel in nearest_neighbor method"
- "Available options are round, floor, and ceil.",
- refl::DefaultValue(ffi::String("")))
- .def_ro("cubic_alpha", &Resize3DAttrs::cubic_alpha,
- "Spline Coefficient for Tricubic Interpolation", refl::DefaultValue(-0.75))
- .def_ro("cubic_exclude", &Resize3DAttrs::cubic_exclude,
- "Flag to exclude exterior of the image during tricubic interpolation",
- refl::DefaultValue(0))
- .def_ro("extrapolation_value", &Resize3DAttrs::extrapolation_value,
- "Value to return when roi is outside of the image", refl::DefaultValue(0.0))
- .def_ro(
- "out_dtype", &Resize3DAttrs::out_dtype,
- "The dtype of the output tensor. It it is not specified, the output will have the same "
- "dtype as input if not specified.",
- refl::DefaultValue(ffi::Optional<DLDataType>{}));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Resize3DAttrs", Resize3DAttrs, AttrsNode);
-}; // struct Resize3DAttrs
-
-/*! \brief Attributes used in image grid_sample operator */
-struct GridSampleAttrs : public AttrsNode {
- ffi::String method;
- ffi::String layout;
- ffi::String padding_mode;
- bool align_corners;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<GridSampleAttrs>()
- .def_ro("method", &GridSampleAttrs::method,
- "Interpolation method. Can be 'nearest', 'bilinear', or 'bicubic'.",
- refl::DefaultValue(ffi::String("bilinear")))
- .def_ro("layout", &GridSampleAttrs::layout,
- "Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc.",
- refl::DefaultValue(ffi::String("NCHW")))
- .def_ro("padding_mode", &GridSampleAttrs::padding_mode,
- "Padding mode for outside grid values. Can be 'zeros', 'border', or 'reflection'.",
- refl::DefaultValue(ffi::String("zeros")))
- .def_ro("align_corners", &GridSampleAttrs::align_corners,
- "If True, the corner pixels of the input and output tensors are aligned.",
- refl::DefaultValue(false));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GridSampleAttrs", GridSampleAttrs, AttrsNode);
-}; // struct GridSampleAttrs
-
-/*! \brief Attributes used in image affine_grid operator */
-struct AffineGridAttrs : public AttrsNode {
- bool align_corners;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<AffineGridAttrs>().def_ro(
- "align_corners", &AffineGridAttrs::align_corners,
- "If True, normalized grid coordinates map to corner pixels; otherwise to pixel centers.",
- refl::DefaultValue(true));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AffineGridAttrs", AffineGridAttrs, AttrsNode);
-}; // struct AffineGridAttrs
-
-} // namespace relax
-} // namespace tvm
-
-#endif // TVM_RELAX_ATTRS_IMAGE_H_
diff --git a/include/tvm/relax/attrs/manipulate.h b/include/tvm/relax/attrs/manipulate.h
deleted file mode 100644
index 1affa55b98..0000000000
--- a/include/tvm/relax/attrs/manipulate.h
+++ /dev/null
@@ -1,324 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file tvm/relax/attrs/manipulate.h
- * \brief Attributes for tensor manipulation operators.
- */
-#ifndef TVM_RELAX_ATTRS_MANIPULATE_H_
-#define TVM_RELAX_ATTRS_MANIPULATE_H_
-
-#include <tvm/relax/expr.h>
-#include <tvm/tirx/index_map.h>
-
-namespace tvm {
-namespace relax {
-
-/*! \brief Attributes used in concat operators */
-struct ConcatAttrs : public AttrsNode {
- ffi::Optional<int64_t> axis;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ConcatAttrs>().def_ro("axis", &ConcatAttrs::axis,
- "The axis at which the input arrays are concatenated."
- "Should lie in range `[-ndim, ndim)`.",
- refl::DefaultValue(ffi::Optional<int64_t>{0}));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ConcatAttrs", ConcatAttrs, AttrsNode);
-}; // struct ConcatAttrs
-
-/*! \brief Attributes used in expand_dims operators */
-struct ExpandDimsAttrs : public AttrsNode {
- ffi::Array<int64_t> axis;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ExpandDimsAttrs>().def_ro(
- "axis", &ExpandDimsAttrs::axis,
- "The axes at which the input array are expanded. "
- "All values are required to lie in range `[-data.ndim - 1, data.ndim]`, "
- "with the convention of negative indexing.");
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ExpandDimsAttrs", ExpandDimsAttrs, AttrsNode);
-}; // struct ExpandDimsAttrs
-
-/*! \brief Attributes used in layout_transform operator */
-struct LayoutTransformAttrs : public AttrsNode {
- tirx::IndexMap index_map;
- // pad_value is chosen to be of PrimExpr type, as it represents constant TIR POD expression. This
- // needs to be revisited in case PrimExpr is evolved to represent symbolic expression in future.
- ffi::Optional<PrimExpr> pad_value;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<LayoutTransformAttrs>()
- .def_ro("index_map", &LayoutTransformAttrs::index_map,
- "The layout transformation to apply.")
- .def_ro(
- "pad_value", &LayoutTransformAttrs::pad_value,
- "The specific value to be used to pad if the layout transform would result in implicit "
- "padding. If not specified, the compiler is free to choose any value.",
- refl::DefaultValue(ffi::Optional<PrimExpr>{}));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LayoutTransformAttrs", LayoutTransformAttrs,
- AttrsNode);
-}; // struct LayoutTransformAttrs
-
-/*! \brief Attributes used in permute_dims operator */
-struct PermuteDimsAttrs : public AttrsNode {
- ffi::Optional<ffi::Array<int64_t>> axes;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<PermuteDimsAttrs>().def_ro(
- "axes", &PermuteDimsAttrs::axes, "The target axes order, reverse order if not specified.",
- refl::DefaultValue(ffi::Optional<ffi::Array<int64_t>>{}));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PermuteDimsAttrs", PermuteDimsAttrs, AttrsNode);
-}; // struct PermuteDimsAttrs
-
-/*! \brief Attributes used in split operator */
-struct SplitAttrs : public AttrsNode {
- ffi::ObjectRef indices_or_sections;
- int axis;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<SplitAttrs>()
- .def_ro("indices_or_sections", &SplitAttrs::indices_or_sections,
- "The input array of indices or the number of split sections.")
- .def_ro("axis", &SplitAttrs::axis, "The axis to be splitted", refl::DefaultValue(0));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SplitAttrs", SplitAttrs, AttrsNode);
-}; // struct SplitAttrs
-
-/*! \brief Attributes used in squeeze operators */
-struct SqueezeAttrs : public AttrsNode {
- ffi::Optional<ffi::Array<int64_t>> axis;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<SqueezeAttrs>().def_ro(
- "axis", &SqueezeAttrs::axis,
- "The axis to squeeze in the input tensor."
- "If `axis = None`, all axis of dimension 1 get squeezed;"
- "Else, the dimension in axes get squeezed."
- "It is an error if an axis does not has dimension 1.",
- refl::DefaultValue(ffi::Optional<ffi::Array<int64_t>>{}));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SqueezeAttrs", SqueezeAttrs, AttrsNode);
-}; // struct SqueezeAttrs
-
-/*! \brief Attributes used in stack operators */
-struct StackAttrs : public AttrsNode {
- ffi::Optional<int64_t> axis;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<StackAttrs>().def_ro(
- "axis", &StackAttrs::axis,
- "The axis along which to stack the input tensors. "
- "The axis will be inserted at this position in the output, "
- "so it must be in range [-ndim-1, ndim] where ndim is the "
- "number of dimensions of the input tensors.",
- refl::DefaultValue(ffi::Optional<int64_t>{0}));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.StackAttrs", StackAttrs, AttrsNode);
-}; // struct StackAttrs
-
-/*! \brief Attributes used in repeat operators */
-struct RepeatAttrs : public AttrsNode {
- int repeats;
- ffi::Optional<int64_t> axis;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<RepeatAttrs>()
- .def_ro("repeats", &RepeatAttrs::repeats, "The number of repetitions.")
- .def_ro("axis", &RepeatAttrs::axis,
- "The axis along which to repeat values. The negative numbers are interpreted "
- "counting from the backward. By default, use the flattened input array, and "
- "return a flat output array.",
- refl::DefaultValue(ffi::Optional<int64_t>{}));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.RepeatAttrs", RepeatAttrs, AttrsNode);
-}; // struct RepeatAttrs
-
-/*! \brief Attributes used in tile operators */
-struct TileAttrs : public AttrsNode {
- ffi::Array<int64_t> repeats;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<TileAttrs>().def_ro("repeats", &TileAttrs::repeats,
- "The number of repetitions of data along each axis.");
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TileAttrs", TileAttrs, AttrsNode);
-}; // struct TileAttrs
-
-/*! \brief Attributes used in flip operators */
-struct FlipAttrs : public AttrsNode {
- int64_t axis;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<FlipAttrs>().def_ro("axis", &FlipAttrs::axis,
- "The axis along which to flip over.");
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.FlipAttrs", FlipAttrs, AttrsNode);
-}; // struct FlipAttrs
-
-/*! \brief Attributes used in reverse_sequence operators */
-struct ReverseSequenceAttrs : public AttrsNode {
- int64_t seq_axis;
- int64_t batch_axis;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ReverseSequenceAttrs>()
- .def_ro("seq_axis", &ReverseSequenceAttrs::seq_axis,
- "The axis along which to reverse variable length slices.", refl::DefaultValue(1))
- .def_ro("batch_axis", &ReverseSequenceAttrs::batch_axis, "The axis that indexes the batch.",
- refl::DefaultValue(0));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ReverseSequenceAttrs", ReverseSequenceAttrs,
- AttrsNode);
-}; // struct ReverseSequenceAttrs
-
-/*! \brief Attributes used in gather_elements operators */
-struct GatherElementsAttrs : public AttrsNode {
- int64_t axis;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<GatherElementsAttrs>().def_ro("axis", &GatherElementsAttrs::axis,
- "The axis along which to index.",
- refl::DefaultValue(0));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GatherElementsAttrs", GatherElementsAttrs,
- AttrsNode);
-}; // struct GatherElementsAttrs
-
-/*! \brief Attributes used in gather_nd operators */
-struct GatherNDAttrs : public AttrsNode {
- int64_t batch_dims;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<GatherNDAttrs>().def_ro("batch_dims", &GatherNDAttrs::batch_dims,
- "The number of batch dims.", refl::DefaultValue(0));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GatherNDAttrs", GatherNDAttrs, AttrsNode);
-}; // struct GatherNDAttrs
-
-/*! \brief Attributes used in index_put operator */
-struct IndexPutAttrs : public AttrsNode {
- bool accumulate;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<IndexPutAttrs>().def_ro(
- "accumulate", &IndexPutAttrs::accumulate,
- "Whether to accumulate (add) values rather than replace. "
- "If true, performs tensor[indices] += values, "
- "otherwise performs tensor[indices] = values.",
- refl::DefaultValue(false));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.IndexPutAttrs", IndexPutAttrs, AttrsNode);
-}; // struct IndexPutAttrs
-
-/*! \brief Attribute used in meshgrid operator */
-struct MeshgridAttrs : public AttrsNode {
- ffi::Optional<ffi::String> indexing;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<MeshgridAttrs>().def_ro("indexing", &MeshgridAttrs::indexing,
- "Specifies how the grid dimensions are ordered.",
- refl::DefaultValue(ffi::String("ij")));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MeshgridAttrs", MeshgridAttrs, AttrsNode);
-};
-
-/*! \brief Attributes used in scatter_elements operators */
-struct ScatterElementsAttrs : public AttrsNode {
- int64_t axis;
- ffi::String reduction;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ScatterElementsAttrs>()
- .def_ro("axis", &ScatterElementsAttrs::axis, "The axis over which to select values.",
- refl::DefaultValue(0))
- .def_ro("reduction", &ScatterElementsAttrs::reduction,
- "Reduction mode of the scatter elements, "
- "either \"update\", \"add\", \"mul\", \"mean\", \"min\" or \"max\".",
- refl::DefaultValue("update"));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScatterElementsAttrs", ScatterElementsAttrs,
- AttrsNode);
-}; // struct ScatterElementsAttrs
-
-/*! \brief Attributes used in scatter_nd operators */
-struct ScatterNDAttrs : public AttrsNode {
- ffi::String reduction;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ScatterNDAttrs>().def_ro(
- "reduction", &ScatterNDAttrs::reduction,
- "Accumulation mode of the ScatterND, "
- "either \"update\", \"add\", \"mul\", \"min\" or \"max\".",
- refl::DefaultValue("update"));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScatterNDAttrs", ScatterNDAttrs, AttrsNode);
-}; // struct ScatterNDAttrs
-
-/*! \brief Attributes used in slice_scatter operator */
-struct SliceScatterAttrs : public AttrsNode {
- int axis;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<SliceScatterAttrs>().def_ro("axis", &SliceScatterAttrs::axis,
- "the dimension to insert the slice into ",
- refl::DefaultValue(0));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SliceScatterAttrs", SliceScatterAttrs, AttrsNode);
-}; // struct SliceScatterAttrs
-
-/*! \brief Attributes used in one_hot operator */
-struct OneHotAttrs : public AttrsNode {
- int depth;
- int axis;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<OneHotAttrs>()
- .def_ro("depth", &OneHotAttrs::depth, "Depth of the one hot dimension.")
- .def_ro("axis", &OneHotAttrs::axis, "Axis to fill.", refl::DefaultValue(-1));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.OneHotAttrs", OneHotAttrs, AttrsNode);
-}; // struct OneHotAttrs
-
-} // namespace relax
-} // namespace tvm
-
-#endif // TVM_RELAX_ATTRS_MANIPULATE_H_
diff --git a/include/tvm/relax/attrs/nn.h b/include/tvm/relax/attrs/nn.h
deleted file mode 100644
index e28a3c80d8..0000000000
--- a/include/tvm/relax/attrs/nn.h
+++ /dev/null
@@ -1,874 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file tvm/relax/attrs/nn.h
- * \brief Attributes for neural network operators.
- */
-#ifndef TVM_RELAX_ATTRS_NN_H_
-#define TVM_RELAX_ATTRS_NN_H_
-
-#include <tvm/relax/expr.h>
-
-namespace tvm {
-namespace relax {
-
-/*! \brief Attributes used in Conv1d operator */
-struct Conv1DAttrs : public AttrsNode {
- ffi::Array<int64_t> strides;
- ffi::Array<int64_t> padding;
- ffi::Array<int64_t> dilation;
- int groups;
- ffi::String data_layout;
- ffi::String kernel_layout;
- ffi::String out_layout;
- ffi::Optional<DLDataType> out_dtype;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<Conv1DAttrs>()
- .def_ro("strides", &Conv1DAttrs::strides, "Specifies the strides of the convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1}))
- .def_ro("padding", &Conv1DAttrs::padding,
- "If padding is non-zero, then the input is implicitly zero-padded"
- "Padding support both symmetric and asymmetric as"
- "one int : same padding used on both sides"
- "two int : padding width in the order of (left, right)",
- refl::DefaultValue(ffi::Array<int64_t>{0, 0}))
- .def_ro("dilation", &Conv1DAttrs::dilation,
- "Specifies the dilation rate to use for dilated convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1}))
- .def_ro("groups", &Conv1DAttrs::groups,
- "Number of groups to split the input into for grouped convolution. The number of "
- "input and "
- "output channels should be divisible by the number of groups.",
- refl::DefaultValue(1))
- .def_ro("data_layout", &Conv1DAttrs::data_layout,
- "Dimension ordering of input data. Can be 'NCW', 'NWC', etc."
- "'N', 'C', 'W' stands for batch, channel, width"
- "dimensions respectively. Convolution is applied on the 'W' dimensions.",
- refl::DefaultValue(ffi::String("NCW")))
- .def_ro("kernel_layout", &Conv1DAttrs::kernel_layout,
- "Dimension ordering of weight. Can be 'OIW', 'IOW', etc."
- "'O', 'I', 'W' stands for num_filter, input_channel, and width"
- "dimensions respectively.",
- refl::DefaultValue(ffi::String("OIW")))
- .def_ro("out_layout", &Conv1DAttrs::out_layout,
- "Dimension ordering of output. Can be 'NCW', 'NWC', etc."
- "'N', 'C', 'W' stands for batch, channel, and width"
- "dimensions respectively. Default to be same as input layout.")
- .def_ro("out_dtype", &Conv1DAttrs::out_dtype,
- "Output data type, set to explicit type under mixed precision setting",
- refl::DefaultValue(ffi::Optional<DLDataType>{}));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv1DAttrs", Conv1DAttrs, AttrsNode);
-}; // struct Conv1dAttrs
-
-/*! \brief Attributes used in Conv2d operator */
-struct Conv2DAttrs : public AttrsNode {
- ffi::Array<int64_t> strides;
- ffi::Array<int64_t> padding;
- ffi::Array<int64_t> dilation;
- int groups;
- ffi::String data_layout;
- ffi::String kernel_layout;
- ffi::String out_layout;
- ffi::Optional<DLDataType> out_dtype;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<Conv2DAttrs>()
- .def_ro("strides", &Conv2DAttrs::strides, "Specifies the strides of the convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1, 1}))
- .def_ro("padding", &Conv2DAttrs::padding,
- "If padding is non-zero, then the input is implicitly zero-padded"
- "Padding support both symmetric and asymmetric as"
- "one int : same padding used on all sides"
- "two int : bottom, right will use same padding as top, left"
- "four int : padding width in the order of (top, left, bottom, right)",
- refl::DefaultValue(ffi::Array<int64_t>{0, 0, 0, 0}))
- .def_ro("dilation", &Conv2DAttrs::dilation,
- "Specifies the dilation rate to use for dilated convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1, 1}))
- .def_ro("groups", &Conv2DAttrs::groups,
- "Number of groups to split the input into for grouped convolution. The number of "
- "input and "
- "output channels should be divisible by the number of groups.",
- refl::DefaultValue(1))
- .def_ro("data_layout", &Conv2DAttrs::data_layout,
- "Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
- "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
- "dimensions respectively. Convolution is applied on the 'H' and"
- "'W' dimensions.",
- refl::DefaultValue(ffi::String("NCHW")))
- .def_ro("kernel_layout", &Conv2DAttrs::kernel_layout,
- "Dimension ordering of weight. Can be 'OIHW', 'OIHW16o16i', etc."
- "'O', 'I', 'H', 'W' stands for num_filter, input_channel, height, and width"
- "dimensions respectively.",
- refl::DefaultValue(ffi::String("OIHW")))
- .def_ro("out_layout", &Conv2DAttrs::out_layout,
- "Dimension ordering of output. Can be 'NCHW', 'NHWC', etc."
- "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
- "dimensions respectively. Default to be same as input layout.")
- .def_ro("out_dtype", &Conv2DAttrs::out_dtype,
- "Output data type, set to explicit type under mixed precision setting",
- refl::DefaultValue(ffi::Optional<DLDataType>{}));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv2DAttrs", Conv2DAttrs, AttrsNode);
-}; // struct Conv2dAttrs
-
-/*! \brief Attributes used in Conv3d operator */
-struct Conv3DAttrs : public AttrsNode {
- ffi::Array<int64_t> strides;
- ffi::Array<int64_t> padding;
- ffi::Array<int64_t> dilation;
- int groups;
- ffi::String data_layout;
- ffi::String kernel_layout;
- ffi::String out_layout;
- ffi::Optional<DLDataType> out_dtype;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<Conv3DAttrs>()
- .def_ro("strides", &Conv3DAttrs::strides, "Specifies the strides of the convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1, 1, 1}))
- .def_ro(
- "padding", &Conv3DAttrs::padding,
- "If padding is non-zero, then the input is implicitly zero-padded"
- "Padding support both symmetric and asymmetric as"
- "one int : same padding used on all sides"
- "two int : bottom, right will use same padding as top, left"
- "four int : padding width in the order of (forward, back, top, left, bottom, right)",
- refl::DefaultValue(ffi::Array<int64_t>{0, 0, 0, 0, 0, 0}))
- .def_ro("dilation", &Conv3DAttrs::dilation,
- "Specifies the dilation rate to use for dilated convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1, 1, 1}))
- .def_ro("groups", &Conv3DAttrs::groups,
- "Number of groups to split the input into for grouped convolution. The number of "
- "input and "
- "output channels should be divisible by the number of groups.",
- refl::DefaultValue(1))
- .def_ro("data_layout", &Conv3DAttrs::data_layout,
- "Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
- "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
- "dimensions respectively. Convolution is applied on the 'D', 'H', and"
- "'W' dimensions.",
- refl::DefaultValue(ffi::String("NCDHW")))
- .def_ro(
- "kernel_layout", &Conv3DAttrs::kernel_layout,
- "Dimension ordering of weight. Can be 'OIDHW', 'OIDHW16o16i', etc."
- "'O', 'I', 'D', 'H', 'W' stands for num_filter, input_channel, depth, height, and width"
- "dimensions respectively.",
- refl::DefaultValue(ffi::String("OIDHW")))
- .def_ro("out_layout", &Conv3DAttrs::out_layout,
- "Dimension ordering of output. Can be 'NCDHW', 'NDHWC', etc."
- "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
- "dimensions respectively. Default to be same as input layout.")
- .def_ro("out_dtype", &Conv3DAttrs::out_dtype,
- "Output data type, set to explicit type under mixed precision setting",
- refl::DefaultValue(ffi::Optional<DLDataType>{}));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv3DAttrs", Conv3DAttrs, AttrsNode);
-}; // struct Conv3dAttrs
-
-/*! \brief Attributes used in Conv1DTranspose operator */
-struct Conv1DTransposeAttrs : public AttrsNode {
- ffi::Array<int64_t> strides;
- ffi::Array<int64_t> padding;
- ffi::Array<int64_t> output_padding;
- ffi::Array<int64_t> dilation;
- int groups;
- ffi::String data_layout;
- ffi::String kernel_layout;
- ffi::String out_layout;
- ffi::Optional<DLDataType> out_dtype;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<Conv1DTransposeAttrs>()
- .def_ro("strides", &Conv1DTransposeAttrs::strides,
- "Specifies the strides of the convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1}))
- .def_ro("padding", &Conv1DTransposeAttrs::padding,
- "If padding is non-zero, then the input is implicitly zero-padded"
- "Padding support both symmetric and asymmetric as"
- "one int : same padding used on both sides"
- "two int : padding width in the order of (left, right)",
- refl::DefaultValue(ffi::Array<int64_t>{0, 0}))
- .def_ro("output_padding", &Conv1DTransposeAttrs::output_padding,
- "Used to disambiguate the output shape.",
- refl::DefaultValue(ffi::Array<int64_t>{0}))
- .def_ro("dilation", &Conv1DTransposeAttrs::dilation,
- "Specifies the dilation rate to use for dilated convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1}))
- .def_ro("groups", &Conv1DTransposeAttrs::groups,
- "Number of groups to split the input into for grouped convolution. The number of "
- "input and "
- "output channels should be divisible by the number of groups.",
- refl::DefaultValue(1))
- .def_ro("data_layout", &Conv1DTransposeAttrs::data_layout,
- "Dimension ordering of input data. Can be 'NCW', 'NWC', etc."
- "'N', 'C', 'W' stands for batch, channel, width"
- "dimensions respectively. Convolution is applied on the 'W' dimensions.",
- refl::DefaultValue(ffi::String("NCW")))
- .def_ro("kernel_layout", &Conv1DTransposeAttrs::kernel_layout,
- "Dimension ordering of weight. Can be 'OIW', 'IOW', etc."
- "'O', 'I', 'W' stands for num_filter, input_channel, and width"
- "dimensions respectively.",
- refl::DefaultValue(ffi::String("IOW")))
- .def_ro("out_layout", &Conv1DTransposeAttrs::out_layout,
- "Dimension ordering of output. Can be 'NCW', 'NWC', etc."
- "'N', 'C', 'W' stands for batch, channel, and width"
- "dimensions respectively. Default to be same as input layout.")
- .def_ro("out_dtype", &Conv1DTransposeAttrs::out_dtype,
- "Output data type, set to explicit type under mixed precision setting",
- refl::DefaultValue(ffi::Optional<DLDataType>{}));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv1DTransposeAttrs", Conv1DTransposeAttrs,
- AttrsNode);
-}; // struct Conv1DTransposeAttrs
-
-/*! \brief Attributes used in Conv2d operator */
-struct Conv2DTransposeAttrs : public AttrsNode {
- ffi::Array<int64_t> strides;
- ffi::Array<int64_t> padding;
- ffi::Array<int64_t> output_padding;
- ffi::Array<int64_t> dilation;
- int groups;
- ffi::String data_layout;
- ffi::String kernel_layout;
- ffi::String out_layout;
- ffi::Optional<DLDataType> out_dtype;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<Conv2DTransposeAttrs>()
- .def_ro("strides", &Conv2DTransposeAttrs::strides,
- "Specifies the strides of the convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1, 1}))
- .def_ro("padding", &Conv2DTransposeAttrs::padding,
- "If padding is non-zero, then the input is implicitly zero-padded"
- "Padding support both symmetric and asymmetric as"
- "one int : same padding used on all sides"
- "two int : bottom, right will use same padding as top, left"
- "four int : padding width in the order of (top, left, bottom, right)",
- refl::DefaultValue(ffi::Array<int64_t>{0, 0, 0, 0}))
- .def_ro("output_padding", &Conv2DTransposeAttrs::output_padding,
- "Used to disambiguate the output shape.",
- refl::DefaultValue(ffi::Array<int64_t>{0, 0}))
- .def_ro("dilation", &Conv2DTransposeAttrs::dilation,
- "Specifies the dilation rate to use for dilated convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1, 1}))
- .def_ro("groups", &Conv2DTransposeAttrs::groups,
- "Number of groups to split the input into for grouped convolution. The number of "
- "input and "
- "output channels should be divisible by the number of groups.",
- refl::DefaultValue(1))
- .def_ro("data_layout", &Conv2DTransposeAttrs::data_layout,
- "Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
- "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
- "dimensions respectively. Convolution is applied on the 'H' and"
- "'W' dimensions.",
- refl::DefaultValue(ffi::String("NCHW")))
- .def_ro("kernel_layout", &Conv2DTransposeAttrs::kernel_layout,
- "Dimension ordering of weight. Can be 'OIHW', 'OIHW16o16i', etc."
- "'O', 'I', 'H', 'W' stands for num_filter, input_channel, height, and width"
- "dimensions respectively.",
- refl::DefaultValue(ffi::String("IOHW")))
- .def_ro("out_layout", &Conv2DTransposeAttrs::out_layout,
- "Dimension ordering of output. Can be 'NCHW', 'NHWC', etc."
- "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
- "dimensions respectively. Default to be same as input layout.")
- .def_ro("out_dtype", &Conv2DTransposeAttrs::out_dtype,
- "Output data type, set to explicit type under mixed precision setting",
- refl::DefaultValue(ffi::Optional<DLDataType>{}));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv2DTransposeAttrs", Conv2DTransposeAttrs,
- AttrsNode);
-}; // struct Conv2DTransposeAttrs
-
-/*! \brief Attributes used in Conv3dTranspose operator */
-struct Conv3DTransposeAttrs : public AttrsNode {
- ffi::Array<int64_t> strides;
- ffi::Array<int64_t> padding;
- ffi::Array<int64_t> output_padding;
- ffi::Array<int64_t> dilation;
- int groups;
- ffi::String data_layout;
- ffi::String kernel_layout;
- ffi::String out_layout;
- ffi::Optional<DLDataType> out_dtype;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<Conv3DTransposeAttrs>()
- .def_ro("strides", &Conv3DTransposeAttrs::strides,
- "Specifies the strides of the convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1, 1, 1}))
- .def_ro("padding", &Conv3DTransposeAttrs::padding,
- "If padding is non-zero, then the input is implicitly zero-padded"
- "Padding support both symmetric and asymmetric as"
- "one int : same padding used on all sides"
- "three int : back/bottom/right will use same padding as front/top/left"
- "six int : padding width in the order of (front, top, left, back, bottom, right)",
- refl::DefaultValue(ffi::Array<int64_t>{0, 0, 0, 0, 0, 0}))
- .def_ro("output_padding", &Conv3DTransposeAttrs::output_padding,
- "Used to disambiguate the output shape.",
- refl::DefaultValue(ffi::Array<int64_t>{0, 0, 0}))
- .def_ro("dilation", &Conv3DTransposeAttrs::dilation,
- "Specifies the dilation rate to use for dilated convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1, 1, 1}))
- .def_ro("groups", &Conv3DTransposeAttrs::groups,
- "Number of groups to split the input into for grouped convolution. The number of "
- "input and "
- "output channels should be divisible by the number of groups.",
- refl::DefaultValue(1))
- .def_ro("data_layout", &Conv3DTransposeAttrs::data_layout,
- "Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
- "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
- "dimensions respectively. Convolution is applied on the 'D', 'H', and"
- "'W' dimensions.",
- refl::DefaultValue(ffi::String("NCDHW")))
- .def_ro(
- "kernel_layout", &Conv3DTransposeAttrs::kernel_layout,
- "Dimension ordering of weight. Can be 'IODHW', etc."
- "'I', 'O', 'D', 'H', 'W' stands for input_channel, output_channel, depth, height, and "
- "width"
- "dimensions respectively.",
- refl::DefaultValue(ffi::String("IODHW")))
- .def_ro("out_layout", &Conv3DTransposeAttrs::out_layout,
- "Dimension ordering of output. Can be 'NCDHW', 'NDHWC', etc."
- "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
- "dimensions respectively. Default to be same as input layout.")
- .def_ro("out_dtype", &Conv3DTransposeAttrs::out_dtype,
- "Output data type, set to explicit type under mixed precision setting",
- refl::DefaultValue(ffi::Optional<DLDataType>{}));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv3DTransposeAttrs", Conv3DTransposeAttrs,
- AttrsNode);
-}; // struct Conv3DTransposeAttrs
-
-/*! \brief Attributes used in max_pool1d and avg_pool1d operator */
-struct Pool1DAttrs : public AttrsNode {
- ffi::Array<int64_t> pool_size;
- ffi::Array<int64_t> strides;
- ffi::Array<int64_t> padding;
- ffi::Array<int64_t> dilation;
- bool ceil_mode;
- bool count_include_pad;
- ffi::String layout;
- ffi::String out_layout;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<Pool1DAttrs>()
- .def_ro("pool_size", &Pool1DAttrs::pool_size, "Size of the pooling windows.",
- refl::DefaultValue(ffi::Array<int64_t>{1}))
- .def_ro("strides", &Pool1DAttrs::strides, "Specifies the strides of the convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1}))
- .def_ro("dilation", &Pool1DAttrs::dilation, "Specifies the dilation of the convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1}))
- .def_ro("padding", &Pool1DAttrs::padding,
- "If padding is non-zero, then the input is implicitly zero-padded"
- "Padding support both symmetric and asymmetric as"
- "one int : same padding used on all sides"
- "two int : padding width in the order of (left, right)",
- refl::DefaultValue(ffi::Array<int64_t>{0, 0}))
- .def_ro(
- "ceil_mode", &Pool1DAttrs::ceil_mode,
- "A boolean indicating if use ceil or floor to compute the output shape. By using ceil, "
- "every element in the input tensor will be covered by a sliding window.",
- refl::DefaultValue(false))
- .def_ro("count_include_pad", &Pool1DAttrs::count_include_pad,
- "When true, will include padding to compute the average")
- .def_ro("layout", &Pool1DAttrs::layout,
- "Dimension ordering of input data. Can be 'NCW', 'NWC', etc."
- "'N', 'C', 'W' stands for batch, channel, and width"
- "dimensions respectively. Pooling is applied on the 'W' dimensions.",
- refl::DefaultValue("NCW"))
- .def_ro("out_layout", &Pool1DAttrs::out_layout,
- "Dimension ordering of output data. Can be 'NCW', 'NWC', etc."
- "'N', 'C', 'W' stands for batch, channel, and width"
- "dimensions respectively. Pooling is applied on the 'W' dimensions.");
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool1DAttrs", Pool1DAttrs, AttrsNode);
-}; // struct Pool1dAttrs
-
-/*! \brief Attributes used in max_pool2d and avg_pool2d operator */
-struct Pool2DAttrs : public AttrsNode {
- ffi::Array<int64_t> pool_size;
- ffi::Array<int64_t> strides;
- ffi::Array<int64_t> padding;
- ffi::Array<int64_t> dilation;
- bool ceil_mode;
- bool count_include_pad;
- ffi::String layout;
- ffi::String out_layout;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<Pool2DAttrs>()
- .def_ro("pool_size", &Pool2DAttrs::pool_size, "Size of the pooling windows.",
- refl::DefaultValue(ffi::Array<int64_t>{1, 1}))
- .def_ro("strides", &Pool2DAttrs::strides, "Specifies the strides of the convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1, 1}))
- .def_ro("dilation", &Pool2DAttrs::dilation, "Specifies the dilation of the convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1, 1}))
- .def_ro("padding", &Pool2DAttrs::padding,
- "If padding is non-zero, then the input is implicitly zero-padded"
- "Padding support both symmetric and asymmetric as"
- "one int : same padding used on all sides"
- "two int : bottom, right will use same padding as top, left"
- "four int : padding width in the order of (top, left, bottom, right)",
- refl::DefaultValue(ffi::Array<int64_t>{0, 0, 0, 0}))
- .def_ro(
- "ceil_mode", &Pool2DAttrs::ceil_mode,
- "A boolean indicating if use ceil or floor to compute the output shape. By using ceil, "
- "every element in the input tensor will be covered by a sliding window.",
- refl::DefaultValue(false))
- .def_ro("count_include_pad", &Pool2DAttrs::count_include_pad,
- "When true, will include padding to compute the average", refl::DefaultValue(false))
- .def_ro("layout", &Pool2DAttrs::layout,
- "Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
- "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
- "dimensions respectively. Pooling is applied on the 'H' and"
- "'W' dimensions.",
- refl::DefaultValue(ffi::String("NCHW")))
- .def_ro("out_layout", &Pool2DAttrs::out_layout,
- "Dimension ordering of output data. Can be 'NCHW', 'NHWC', etc."
- "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
- "dimensions respectively. Pooling is applied on the 'H' and"
- "'W' dimensions.");
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool2DAttrs", Pool2DAttrs, AttrsNode);
-}; // struct Pool2dAttrs
-
-/*! \brief Attributes used in max_pool3d and avg_pool3d operator */
-struct Pool3DAttrs : public AttrsNode {
- ffi::Array<int64_t> pool_size;
- ffi::Array<int64_t> strides;
- ffi::Array<int64_t> padding;
- ffi::Array<int64_t> dilation;
- bool ceil_mode;
- bool count_include_pad;
- ffi::String layout;
- ffi::String out_layout;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<Pool3DAttrs>()
- .def_ro("pool_size", &Pool3DAttrs::pool_size, "Size of the pooling windows.",
- refl::DefaultValue(ffi::Array<int64_t>{1, 1, 1}))
- .def_ro("strides", &Pool3DAttrs::strides, "Specifies the strides of the convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1, 1, 1}))
- .def_ro("dilation", &Pool3DAttrs::dilation, "Specifies the dilation of the convolution.",
- refl::DefaultValue(ffi::Array<int64_t>{1, 1, 1}))
- .def_ro("padding", &Pool3DAttrs::padding,
- "If padding is non-zero, then the input is implicitly zero-padded"
- "Padding support both symmetric and asymmetric as"
- "one int : same padding used on all sides"
- "three int : back, bottom, right will use same padding as front, top, left"
- "four int : padding width in the order of (front, top, left, back, bottom, right)",
- refl::DefaultValue(ffi::Array<int64_t>{0, 0, 0, 0, 0, 0}))
- .def_ro(
- "ceil_mode", &Pool3DAttrs::ceil_mode,
- "A boolean indicating if use ceil or floor to compute the output shape. By using ceil, "
- "every element in the input tensor will be covered by a sliding window.",
- refl::DefaultValue(false))
- .def_ro("count_include_pad", &Pool3DAttrs::count_include_pad,
- "When true, will include padding to compute the average")
- .def_ro("layout", &Pool3DAttrs::layout,
- "Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
- "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
- "dimensions respectively. Pooling is applied on the 'D', 'H' and"
- "'W' dimensions.",
- refl::DefaultValue(ffi::String("NCDHW")))
- .def_ro("out_layout", &Pool3DAttrs::out_layout,
- "Dimension ordering of output data. Can be 'NCDHW', 'NDHWC', etc."
- "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
- "dimensions respectively. Pooling is applied on the 'D', 'H' and"
- "'W' dimensions.");
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool3DAttrs", Pool3DAttrs, AttrsNode);
-}; // struct Pool3dAttrs
-
-/*! \brief Attributes for 1d adaptive pool operator */
-struct AdaptivePool1DAttrs : public AttrsNode {
- ffi::Optional<ffi::Array<int64_t>> output_size;
- ffi::String layout;
- ffi::String out_layout;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<AdaptivePool1DAttrs>()
- .def_ro("output_size", &AdaptivePool1DAttrs::output_size, "Output width.",
- refl::DefaultValue(ffi::Optional<ffi::Array<int64_t>>{}))
- .def_ro("layout", &AdaptivePool1DAttrs::layout,
- "Dimension ordering of input data. Can be 'NCW', 'NWC', etc."
- "'N', 'C', 'W' stands for batch, channel and width"
- "dimensions respectively. Pooling is applied on the"
- "'W' dimensions.",
- refl::DefaultValue(ffi::String("NCW")))
- .def_ro("out_layout", &AdaptivePool1DAttrs::out_layout,
- "Dimension ordering of output data. Can be 'NCW', 'NWC', etc."
- "'N', 'C', 'W' stands for batch, channel and width"
- "dimensions respectively. Pooling is applied on the"
- "'W' dimensions.");
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AdaptivePool1DAttrs", AdaptivePool1DAttrs,
- AttrsNode);
-}; // struct AdaptivePool1DAttrs
-
-/*! \brief Attributes for 2d adaptive pool operator */
-struct AdaptivePool2DAttrs : public AttrsNode {
- ffi::Optional<ffi::Array<int64_t>> output_size;
- ffi::String layout;
- ffi::String out_layout;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<AdaptivePool2DAttrs>()
- .def_ro("output_size", &AdaptivePool2DAttrs::output_size, "Output height and width.",
- refl::DefaultValue(ffi::Optional<ffi::Array<int64_t>>{}))
- .def_ro("layout", &AdaptivePool2DAttrs::layout,
- "Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
- "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
- "dimensions respectively. Pooling is applied on the 'H' and"
- "'W' dimensions.",
- refl::DefaultValue(ffi::String("NCHW")))
- .def_ro("out_layout", &AdaptivePool2DAttrs::out_layout,
- "Dimension ordering of output data. Can be 'NCHW', 'NHWC', etc."
- "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
- "dimensions respectively. Pooling is applied on the 'H' and"
- "'W' dimensions.");
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AdaptivePool2DAttrs", AdaptivePool2DAttrs,
- AttrsNode);
-}; // struct AdaptivePool2DAttrs
-
-/*! \brief Attributes for 3d adaptive pool operator */
-struct AdaptivePool3DAttrs : public AttrsNode {
- ffi::Optional<ffi::Array<int64_t>> output_size;
- ffi::String layout;
- ffi::String out_layout;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<AdaptivePool3DAttrs>()
- .def_ro("output_size", &AdaptivePool3DAttrs::output_size, "Output depth, height and width.",
- refl::DefaultValue(ffi::Optional<ffi::Array<int64_t>>{}))
- .def_ro("layout", &AdaptivePool3DAttrs::layout,
- "Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
- "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
- "dimensions respectively. Pooling is applied on 'D', 'H' and"
- "'W' dimensions.",
- refl::DefaultValue(ffi::String("NCDHW")))
- .def_ro("out_layout", &AdaptivePool3DAttrs::out_layout,
- "Dimension ordering of output data. Can be 'NCDHW', 'NDHWC', etc."
- "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
- "dimensions respectively. Pooling is applied on 'D', 'H' and"
- "'W' dimensions.");
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AdaptivePool3DAttrs", AdaptivePool3DAttrs,
- AttrsNode);
-}; // struct AdaptivePool3DAttrs
-
-/*! \brief Attributes used in softmax operators */
-struct SoftmaxAttrs : public AttrsNode {
- int axis;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<SoftmaxAttrs>().def_ro("axis", &SoftmaxAttrs::axis,
- "The axis to sum over when computing softmax.",
- refl::DefaultValue(-1));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SoftmaxAttrs", SoftmaxAttrs, AttrsNode);
-};
-
-/*! \brief Attributes used in softmax operators */
-struct LeakyReluAttrs : public AttrsNode {
- double alpha;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<LeakyReluAttrs>().def_ro("alpha", &LeakyReluAttrs::alpha,
- "The slope of the negative part.",
- refl::DefaultValue(0.01));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LeakyReluAttrs", LeakyReluAttrs, AttrsNode);
-};
-
-/*! \brief Attributes used in softplus operators */
-struct SoftplusAttrs : public AttrsNode {
- double beta;
- double threshold;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<SoftplusAttrs>()
- .def_ro("beta", &SoftplusAttrs::beta,
- "Scaling factor controlling the sharpness of the Softplus transition.",
- refl::DefaultValue(1.0))
- .def_ro("threshold", &SoftplusAttrs::threshold,
- "Value determining when to use linear approximation for numerical stability.",
- refl::DefaultValue(20.0));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SoftplusAttrs", SoftplusAttrs, AttrsNode);
-};
-
-/*! \brief Attributes used in PReLU operator */
-struct PReluAttrs : public AttrsNode {
- int axis;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<PReluAttrs>().def_ro("axis", &PReluAttrs::axis,
- "The axis along which the alpha values are applied.",
- refl::DefaultValue(1));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PReluAttrs", PReluAttrs, AttrsNode);
-};
-
-/*! \brief Attributes used in batch_norm operator */
-struct BatchNormAttrs : public AttrsNode {
- int axis;
- double epsilon;
- bool center;
- bool scale;
- double momentum;
- bool training;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<BatchNormAttrs>()
- .def_ro("axis", &BatchNormAttrs::axis, "The axis along which the normalization is applied.")
- .def_ro("epsilon", &BatchNormAttrs::epsilon,
- "Small float added to variance to avoid dividing by zero",
- refl::DefaultValue(1e-05))
- .def_ro("center", &BatchNormAttrs::center,
- "Indicating if the beta offset will be added to the normalized tensor.",
- refl::DefaultValue(true))
- .def_ro("scale", &BatchNormAttrs::scale,
- "Indicating if the gamma scale will be multiplied.", refl::DefaultValue(true))
- .def_ro("momentum", &BatchNormAttrs::momentum,
- "The value used for the moving_mean and moving_var update.",
- refl::DefaultValue(0.1))
- .def_ro("training", &BatchNormAttrs::training,
- "Whether we are training (i.e., not in eval mode).", refl::DefaultValue(true));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.BatchNormAttrs", BatchNormAttrs, AttrsNode);
-}; // struct BatchNormAttrs
-
-/*! \brief Attributes used in layer_norm operator */
-struct LayerNormAttrs : public AttrsNode {
- ffi::Array<int64_t> axes;
- double epsilon;
- bool center;
- bool scale;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<LayerNormAttrs>()
- .def_ro("axes", &LayerNormAttrs::axes,
- "The axes that along which the normalization is applied.")
- .def_ro("epsilon", &LayerNormAttrs::epsilon,
- "Small float added to variance to avoid dividing by zero",
- refl::DefaultValue(1e-05))
- .def_ro("center", &LayerNormAttrs::center,
- "Indicating if the beta offset will be added to the normalized tensor.",
- refl::DefaultValue(true))
- .def_ro("scale", &LayerNormAttrs::scale,
- "Indicating if the gamma scale will be multiplied.", refl::DefaultValue(true));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LayerNormAttrs", LayerNormAttrs, AttrsNode);
-}; // struct LayerNormAttrs
-
-/*! \brief Attributes used in group_norm operator */
-struct GroupNormAttrs : public AttrsNode {
- int num_groups;
- int channel_axis;
- ffi::Array<int64_t> axes;
- double epsilon;
- bool center;
- bool scale;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<GroupNormAttrs>()
- .def_ro("num_groups", &GroupNormAttrs::num_groups,
- "The number of groups to separate the channels into.")
- .def_ro("channel_axis", &GroupNormAttrs::channel_axis,
- "The axis that represents the channel.")
- .def_ro(
- "axes", &GroupNormAttrs::axes,
- "The axes that along which the normalization is applied (excluding the channel axis).")
- .def_ro("epsilon", &GroupNormAttrs::epsilon,
- "Small float added to variance to avoid dividing by zero",
- refl::DefaultValue(1e-05))
- .def_ro("center", &GroupNormAttrs::center,
- "Indicating if the beta offset will be added to the normalized tensor.",
- refl::DefaultValue(true))
- .def_ro("scale", &GroupNormAttrs::scale,
- "Indicating if the gamma scale will be multiplied.", refl::DefaultValue(true));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GroupNormAttrs", GroupNormAttrs, AttrsNode);
-}; // struct GroupNormAttrs
-
-/*! \brief Attributes used in instance_norm operator */
-struct InstanceNormAttrs : public AttrsNode {
- int channel_axis;
- ffi::Array<int64_t> axes;
- double epsilon;
- bool center;
- bool scale;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<InstanceNormAttrs>()
- .def_ro("channel_axis", &InstanceNormAttrs::channel_axis,
- "The axis that represents the channel.")
- .def_ro("axes", &InstanceNormAttrs::axes,
- "The axes that along which the normalization is applied.")
- .def_ro("epsilon", &InstanceNormAttrs::epsilon,
- "Small float added to variance to avoid dividing by zero",
- refl::DefaultValue(1e-05))
- .def_ro("center", &InstanceNormAttrs::center,
- "Indicating if the beta offset will be added to the normalized tensor.",
- refl::DefaultValue(true))
- .def_ro("scale", &InstanceNormAttrs::scale,
- "Indicating if the gamma scale will be multiplied.", refl::DefaultValue(true));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.InstanceNormAttrs", InstanceNormAttrs, AttrsNode);
-}; // struct InstanceNormAttrs
-
-/*! \brief Attributes used in rms_norm operator */
-struct RMSNormAttrs : public AttrsNode {
- ffi::Array<int64_t> axes;
- double epsilon;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<RMSNormAttrs>()
- .def_ro("axes", &RMSNormAttrs::axes,
- "The axes that along which the normalization is applied.",
- refl::DefaultValue(ffi::Array<int64_t>{-1}))
- .def_ro("epsilon", &RMSNormAttrs::epsilon,
- "Small float added to variance to avoid dividing by zero",
- refl::DefaultValue(1e-05));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.RMSNormAttrs", RMSNormAttrs, AttrsNode);
-}; // struct RMSNormAttrs
-
-/*! \brief Attributes used in nll_loss operator */
-struct NLLLossAttrs : public AttrsNode {
- ffi::String reduction;
- int ignore_index;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<NLLLossAttrs>()
- .def_ro("reduction", &NLLLossAttrs::reduction,
- "The reduction method to apply to the output. Can be"
- "'none', 'mean' or 'sum'.",
- refl::DefaultValue("mean"))
- .def_ro("ignore_index", &NLLLossAttrs::ignore_index, "The target value to ignore.",
- refl::DefaultValue(-100));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.NLLLossAttrs", NLLLossAttrs, AttrsNode);
-}; // struct NLLLossAttrs
-
-/*! \brief Attributes used in dropout operator */
-struct DropoutAttrs : public AttrsNode {
- double rate;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<DropoutAttrs>().def_ro(
- "rate", &DropoutAttrs::rate,
- "Fraction of the input that gets dropped out during training time",
- refl::DefaultValue(0.5));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.DropoutAttrs", DropoutAttrs, AttrsNode);
-}; // struct DropoutAttrs
-
-/*! \brief Attributes used in Attention operator */
-struct AttentionAttrs : public AttrsNode {
- ffi::Optional<FloatImm> scale;
- ffi::Optional<ffi::String> causal_mask;
- ffi::Optional<IntImm> window_size;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<AttentionAttrs>()
- .def_ro(
- "scale", &AttentionAttrs::scale,
- "The custom scale applied before the softmax. The default value is 1 / sqrt(head_dim).",
- refl::DefaultValue(ffi::Optional<FloatImm>{}))
- .def_ro("causal_mask", &AttentionAttrs::causal_mask,
- "The type of the causal mask, i.e. 'TopLeft' and 'BottomRight'.",
- refl::DefaultValue(ffi::Optional<ffi::String>{}))
- .def_ro("window_size", &AttentionAttrs::window_size,
- "The size of the window for sliding-window attention.",
- refl::DefaultValue(ffi::Optional<IntImm>{}));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AttentionAttrs", AttentionAttrs, AttrsNode);
-}; // struct AttentionAttrs
-
-/*! \brief Attributes used for the padding operator */
-struct PadAttrs : public AttrsNode {
- ffi::Array<int64_t> pad_width;
- double pad_value = 0.0;
- tvm::ffi::String pad_mode;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<PadAttrs>()
- .def_ro("pad_width", &PadAttrs::pad_width,
- "Number of values padded to the edges of each axis, "
- "in the format of (before_1, after_1, ..., before_N, after_N)")
- .def_ro("pad_value", &PadAttrs::pad_value, "The value to fill in padded area with",
- refl::DefaultValue(0.0))
- .def_ro("pad_mode", &PadAttrs::pad_mode,
- "Padding type to use. \"constant\" pads with constant_value, "
- "\"edge\" pads using the edge values of the input array, "
- "\"reflect\" pads by reflecting values with respect to the edges.",
- refl::DefaultValue("constant"));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PadAttrs", PadAttrs, AttrsNode);
-};
-
-/*! \brief Attributes used for the pixel shuffle operator */
-struct PixelShuffleAttrs : public AttrsNode {
- int upscale_factor;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<PixelShuffleAttrs>().def_ro("upscale_factor",
- &PixelShuffleAttrs::upscale_factor,
- "Scale factor for spatial upsampling.");
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PixelShuffleAttrs", PixelShuffleAttrs, AttrsNode);
-};
-
-} // namespace relax
-} // namespace tvm
-
-#endif // TVM_RELAX_ATTRS_NN_H_
diff --git a/include/tvm/relax/attrs/sorting.h b/include/tvm/relax/attrs/sorting.h
deleted file mode 100644
index 606a5e3c85..0000000000
--- a/include/tvm/relax/attrs/sorting.h
+++ /dev/null
@@ -1,109 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file tvm/relax/attrs/sorting.h
- * \brief Attributes for sorting operators.
- */
-#ifndef TVM_RELAX_ATTRS_SORTING_H_
-#define TVM_RELAX_ATTRS_SORTING_H_
-
-#include <tvm/relax/expr.h>
-#include <tvm/tirx/index_map.h>
-
-namespace tvm {
-namespace relax {
-
-/*! \brief Attributes used in sort operator */
-struct SortAttrs : public AttrsNode {
- int axis;
- bool descending;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<SortAttrs>()
- .def_ro("axis", &SortAttrs::axis,
- "Axis along which the sort is computed."
- "The default the last axis is used.",
- refl::DefaultValue(-1))
- .def_ro("descending", &SortAttrs::descending,
- "Whether to sort in descending order."
- "If it is not specified, it defaults to the ascending order.",
- refl::DefaultValue(false));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SortAttrs", SortAttrs, AttrsNode);
-}; // struct SortAttrs
-
-/*! \brief Attributes used in argsort operator */
-struct ArgsortAttrs : public AttrsNode {
- int axis;
- bool descending;
- ffi::Optional<DLDataType> dtype;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ArgsortAttrs>()
- .def_ro("axis", &ArgsortAttrs::axis,
- "Axis along which the argsort is computed."
- "The default the last axis is used.",
- refl::DefaultValue(-1))
- .def_ro("descending", &ArgsortAttrs::descending,
- "Whether to argsort in descending order."
- "If it is not specified, it defaults to the ascending order.",
- refl::DefaultValue(false))
- .def_ro("dtype", &ArgsortAttrs::dtype, "DType of the output indices.",
- refl::DefaultValue((DLDataType{kDLInt, 32, 1})));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ArgsortAttrs", ArgsortAttrs, AttrsNode);
-}; // struct ArgsortAttrs
-
-/*! \brief Attributes used in topk operator */
-struct TopKAttrs : public AttrsNode {
- int k;
- int axis;
- bool largest;
- ffi::String ret_type;
- ffi::Optional<DLDataType> dtype;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<TopKAttrs>()
- .def_ro("k", &TopKAttrs::k, "Number of top elements to select", refl::DefaultValue(1))
- .def_ro("axis", &TopKAttrs::axis, "Axis along which to sort the input tensor.",
- refl::DefaultValue(-1))
- .def_ro("ret_type", &TopKAttrs::ret_type,
- "The return type [both, values, indices]."
- "both - return both top k data and indices."
- "values - return top k data only."
- "indices - return top k indices only.",
- refl::DefaultValue("both"))
- .def_ro("largest", &TopKAttrs::largest,
- "Whether to return largest or smallest elements."
- "By default, return the largest k elements.",
- refl::DefaultValue(true))
- .def_ro("dtype", &TopKAttrs::dtype, "Data type of the output indices.",
- refl::DefaultValue((DLDataType{kDLInt, 32, 1})));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TopKAttrs", TopKAttrs, AttrsNode);
-}; // struct TopKAttrs
-
-} // namespace relax
-} // namespace tvm
-
-#endif // TVM_RELAX_ATTRS_SORTING_H_
diff --git a/include/tvm/relax/attrs/vision.h b/include/tvm/relax/attrs/vision.h
deleted file mode 100644
index 8091ad1c32..0000000000
--- a/include/tvm/relax/attrs/vision.h
+++ /dev/null
@@ -1,197 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-/*!
- * \file tvm/relax/attrs/vision.h
- * \brief Auxiliary attributes for vision operators.
- */
-#ifndef TVM_RELAX_ATTRS_VISION_H_
-#define TVM_RELAX_ATTRS_VISION_H_
-
-#include <tvm/ffi/string.h>
-#include <tvm/ir/attrs.h>
-#include <tvm/ir/type.h>
-#include <tvm/relax/expr.h>
-
-namespace tvm {
-namespace relax {
-
-/*! \brief Attributes used in AllClassNonMaximumSuppression operator */
-struct AllClassNonMaximumSuppressionAttrs : public AttrsNode {
- ffi::String output_format;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<AllClassNonMaximumSuppressionAttrs>().def_ro(
- "output_format", &AllClassNonMaximumSuppressionAttrs::output_format,
- "Output format, onnx or tensorflow. Returns outputs in a way that can be easily "
- "consumed by each frontend.",
- refl::DefaultValue(ffi::String("onnx")));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AllClassNonMaximumSuppressionAttrs",
- AllClassNonMaximumSuppressionAttrs, AttrsNode);
-}; // struct AllClassNonMaximumSuppressionAttrs
-
-/*! \brief Attributes used in ROIAlign operator */
-struct ROIAlignAttrs : public AttrsNode {
- ffi::Array<int64_t> pooled_size;
- double spatial_scale;
- int sample_ratio;
- bool aligned;
- ffi::String layout;
- ffi::String mode;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ROIAlignAttrs>()
- .def_ro("pooled_size", &ROIAlignAttrs::pooled_size, "Output size of roi align.")
- .def_ro("spatial_scale", &ROIAlignAttrs::spatial_scale,
- "Ratio of input feature map height (or width) to raw image height (or width).")
- .def_ro("sample_ratio", &ROIAlignAttrs::sample_ratio,
- "Optional sampling ratio of ROI align, using adaptive size by default.",
- refl::DefaultValue(-1))
- .def_ro("aligned", &ROIAlignAttrs::aligned,
- "Whether to use the aligned ROIAlign semantics without the legacy 1-pixel clamp.",
- refl::DefaultValue(false))
- .def_ro("layout", &ROIAlignAttrs::layout, "Dimension ordering of the input data.",
- refl::DefaultValue(ffi::String("NCHW")))
- .def_ro("mode", &ROIAlignAttrs::mode, "Mode for ROI Align. Can be 'avg' or 'max'.",
- refl::DefaultValue(ffi::String("avg")));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ROIAlignAttrs", ROIAlignAttrs, AttrsNode);
-}; // struct ROIAlignAttrs
-
-/*! \brief Attributes used in ROIPool operator */
-struct ROIPoolAttrs : public AttrsNode {
- ffi::Array<int64_t> pooled_size;
- double spatial_scale;
- ffi::String layout;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ROIPoolAttrs>()
- .def_ro("pooled_size", &ROIPoolAttrs::pooled_size, "Output size of roi pool.")
- .def_ro("spatial_scale", &ROIPoolAttrs::spatial_scale,
- "Ratio of input feature map height (or width) to raw image height (or width).")
- .def_ro("layout", &ROIPoolAttrs::layout, "Dimension ordering of the input data.",
- refl::DefaultValue(ffi::String("NCHW")));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ROIPoolAttrs", ROIPoolAttrs, AttrsNode);
-}; // struct ROIPoolAttrs
-
-/*! \brief Attributes used in GetValidCounts operator */
-struct GetValidCountsAttrs : public AttrsNode {
- double score_threshold;
- int id_index;
- int score_index;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<GetValidCountsAttrs>()
- .def_ro("score_threshold", &GetValidCountsAttrs::score_threshold,
- "Lower limit of score for valid bounding boxes.", refl::DefaultValue(0))
- .def_ro("id_index", &GetValidCountsAttrs::id_index,
- "Index of the class categories, -1 to disable.", refl::DefaultValue(0))
- .def_ro("score_index", &GetValidCountsAttrs::score_index,
- "Index of the scores/confidence of boxes.", refl::DefaultValue(1));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GetValidCountsAttrs", GetValidCountsAttrs,
- AttrsNode);
-}; // struct GetValidCountsAttrs
-
-/*! \brief Attributes used in NonMaximumSuppression operator */
-struct NonMaximumSuppressionAttrs : public AttrsNode {
- int max_output_size;
- double iou_threshold;
- bool force_suppress;
- int top_k;
- int coord_start;
- int score_index;
- int id_index;
- bool return_indices;
- bool invalid_to_bottom;
- double soft_nms_sigma;
- double score_threshold;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<NonMaximumSuppressionAttrs>()
- .def_ro("max_output_size", &NonMaximumSuppressionAttrs::max_output_size,
- "Max number of output valid boxes, -1 for no limit.", refl::DefaultValue(-1))
- .def_ro("iou_threshold", &NonMaximumSuppressionAttrs::iou_threshold,
- "Non-maximum suppression IoU threshold.", refl::DefaultValue(0.5))
- .def_ro("force_suppress", &NonMaximumSuppressionAttrs::force_suppress,
- "Whether to suppress all detections regardless of class_id.",
- refl::DefaultValue(false))
- .def_ro("top_k", &NonMaximumSuppressionAttrs::top_k,
- "Keep maximum top k detections before nms, -1 for no limit.",
- refl::DefaultValue(-1))
- .def_ro("coord_start", &NonMaximumSuppressionAttrs::coord_start,
- "Start index of the consecutive 4 coordinates.", refl::DefaultValue(2))
- .def_ro("score_index", &NonMaximumSuppressionAttrs::score_index,
- "Index of the scores/confidence of boxes.", refl::DefaultValue(1))
- .def_ro("id_index", &NonMaximumSuppressionAttrs::id_index,
- "Index of the class categories, -1 to disable.", refl::DefaultValue(0))
- .def_ro("return_indices", &NonMaximumSuppressionAttrs::return_indices,
- "Whether to return box indices in input data.", refl::DefaultValue(true))
- .def_ro("invalid_to_bottom", &NonMaximumSuppressionAttrs::invalid_to_bottom,
- "Whether to move all valid bounding boxes to the top.", refl::DefaultValue(false))
- .def_ro("soft_nms_sigma", &NonMaximumSuppressionAttrs::soft_nms_sigma,
- "Sigma for soft-NMS; 0.0 means standard hard NMS.", refl::DefaultValue(0.0))
- .def_ro("score_threshold", &NonMaximumSuppressionAttrs::score_threshold,
- "Score threshold for soft-NMS validity check; 0.0 when unused.",
- refl::DefaultValue(0.0));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.NonMaximumSuppressionAttrs",
- NonMaximumSuppressionAttrs, AttrsNode);
-}; // struct NonMaximumSuppressionAttrs
-
-/*! \brief Attributes for multibox_transform_loc (SSD / TFLite-style box decode). */
-struct MultiboxTransformLocAttrs : public AttrsNode {
- bool clip;
- double threshold;
- ffi::Array<double> variances;
- bool keep_background;
- bool apply_softmax;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<MultiboxTransformLocAttrs>()
- .def_ro("clip", &MultiboxTransformLocAttrs::clip,
- "Clip decoded ymin,xmin,ymax,xmax to [0,1].", refl::DefaultValue(false))
- .def_ro("threshold", &MultiboxTransformLocAttrs::threshold,
- "After softmax, zero scores strictly below this value.", refl::DefaultValue(0.0))
- .def_ro("variances", &MultiboxTransformLocAttrs::variances,
- "(x,y,w,h) scales = TFLite 1/x_scale,1/y_scale,1/w_scale,1/h_scale on "
- "encodings. Very large w/h scales can overflow exp in decode.",
- refl::DefaultValue(ffi::Array<double>{1.0, 1.0, 1.0, 1.0}))
- .def_ro("keep_background", &MultiboxTransformLocAttrs::keep_background,
- "If false, force output scores[:,0,:] to 0 (background class).",
- refl::DefaultValue(true))
- .def_ro("apply_softmax", &MultiboxTransformLocAttrs::apply_softmax,
- "Whether to apply softmax to class predictions before thresholding.",
- refl::DefaultValue(true));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MultiboxTransformLocAttrs",
- MultiboxTransformLocAttrs, AttrsNode);
-}; // struct MultiboxTransformLocAttrs
-
-} // namespace relax
-} // namespace tvm
-
-#endif // TVM_RELAX_ATTRS_VISION_H_
diff --git a/include/tvm/relax/attrs/ccl.h b/include/tvm/relax/op/ccl.h
similarity index 50%
rename from include/tvm/relax/attrs/ccl.h
rename to include/tvm/relax/op/ccl.h
index 2d6a1d183d..a0c986e237 100644
--- a/include/tvm/relax/attrs/ccl.h
+++ b/include/tvm/relax/op/ccl.h
@@ -18,11 +18,11 @@
*/
/*!
- * \file tvm/relax/attrs/ccl.h
+ * \file tvm/relax/op/ccl.h
* \brief Attributes for ccl operators.
*/
-#ifndef TVM_RELAX_ATTRS_CCL_H_
-#define TVM_RELAX_ATTRS_CCL_H_
+#ifndef TVM_RELAX_OP_CCL_H_
+#define TVM_RELAX_OP_CCL_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/expr.h>
@@ -35,18 +35,7 @@ struct AllReduceAttrs : public tvm::AttrsNode {
ffi::String op_type;
bool in_group;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<AllReduceAttrs>()
- .def_ro("op_type", &AllReduceAttrs::op_type,
- "The type of reduction operation to be applied to the input data. Now only sum is "
- "supported.",
- refl::DefaultValue(ffi::String("sum")))
- .def_ro("in_group", &AllReduceAttrs::in_group,
- "Whether the reduction operation performs in group or globally or in group as "
- "default.",
- refl::DefaultValue(true));
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AllReduceAttrs", AllReduceAttrs, AttrsNode);
}; // struct AllReduceAttrs
@@ -55,17 +44,7 @@ struct AllGatherAttrs : public tvm::AttrsNode {
int num_workers;
bool in_group;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<AllGatherAttrs>()
- .def_ro("num_workers", &AllGatherAttrs::num_workers,
- "The number of workers, also the number of parts the given buffer should be "
- "chunked into.")
- .def_ro("in_group", &AllGatherAttrs::in_group,
- "Whether the allgather operation performs in group or globally or in group as "
- "default.",
- refl::DefaultValue(true));
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AllGatherAttrs", AllGatherAttrs, AttrsNode);
}; // struct AllGatherAttrs
@@ -74,17 +53,7 @@ struct ScatterCollectiveAttrs : public tvm::AttrsNode {
int num_workers;
int axis;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ScatterCollectiveAttrs>()
- .def_ro("num_workers", &ScatterCollectiveAttrs::num_workers,
- "The number of workers, also the number of parts the given buffer should be "
- "chunked into.")
- .def_ro("axis", &ScatterCollectiveAttrs::axis,
- "The axis of the tensor to be scattered. The tensor will be chunked along "
- "this axis.",
- refl::DefaultValue(0));
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScatterCollectiveAttrs", ScatterCollectiveAttrs,
AttrsNode);
}; // struct ScatterCollectiveAttrs
@@ -92,4 +61,4 @@ struct ScatterCollectiveAttrs : public tvm::AttrsNode {
} // namespace relax
} // namespace tvm
-#endif // TVM_RELAX_ATTRS_CCL_H_
+#endif // TVM_RELAX_OP_CCL_H_
diff --git a/include/tvm/relax/attrs/create.h b/include/tvm/relax/op/create.h
similarity index 67%
rename from include/tvm/relax/attrs/create.h
rename to include/tvm/relax/op/create.h
index f2ab653443..9b4d0a5062 100644
--- a/include/tvm/relax/attrs/create.h
+++ b/include/tvm/relax/op/create.h
@@ -18,11 +18,11 @@
*/
/*!
- * \file tvm/relax/attrs/create.h
+ * \file tvm/relax/op/create.h
* \brief Attributes for tensor creation operators.
*/
-#ifndef TVM_RELAX_ATTRS_CREATE_H_
-#define TVM_RELAX_ATTRS_CREATE_H_
+#ifndef TVM_RELAX_OP_CREATE_H_
+#define TVM_RELAX_OP_CREATE_H_
#include <tvm/relax/expr.h>
@@ -33,12 +33,7 @@ namespace relax {
struct InitAttrs : public AttrsNode {
ffi::Optional<DLDataType> dtype;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<InitAttrs>().def_ro("dtype", &InitAttrs::dtype,
- "The data type of the created tensor.",
- refl::DefaultValue(ffi::Optional<DLDataType>{}));
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.InitAttrs", InitAttrs, AttrsNode);
}; // struct InitAttrs
@@ -46,16 +41,11 @@ struct InitAttrs : public AttrsNode {
struct TriluAttrs : public AttrsNode {
int k;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<TriluAttrs>().def_ro(
- "k", &TriluAttrs::k,
- "The number of diagonals above or below the main diagonal to exclude or include.");
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TriluAttrs", TriluAttrs, AttrsNode);
}; // struct TriluAttrs
} // namespace relax
} // namespace tvm
-#endif // TVM_RELAX_ATTRS_CREATE_H_
+#endif // TVM_RELAX_OP_CREATE_H_
diff --git a/include/tvm/relax/attrs/datatype.h b/include/tvm/relax/op/datatype.h
similarity index 70%
rename from include/tvm/relax/attrs/datatype.h
rename to include/tvm/relax/op/datatype.h
index 5ee6e8af33..9c04904acd 100644
--- a/include/tvm/relax/attrs/datatype.h
+++ b/include/tvm/relax/op/datatype.h
@@ -18,11 +18,11 @@
*/
/*!
- * \file tvm/relax/attrs/datatype.h
+ * \file tvm/relax/op/datatype.h
* \brief Attributes for datatype operators.
*/
-#ifndef TVM_RELAX_ATTRS_DATATYPE_H_
-#define TVM_RELAX_ATTRS_DATATYPE_H_
+#ifndef TVM_RELAX_OP_DATATYPE_H_
+#define TVM_RELAX_OP_DATATYPE_H_
#include <tvm/relax/expr.h>
@@ -33,10 +33,7 @@ namespace relax {
struct AstypeAttrs : public AttrsNode {
DLDataType dtype;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<AstypeAttrs>().def_ro("dtype", &AstypeAttrs::dtype, "Target data type");
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AstypeAttrs", AstypeAttrs, AttrsNode);
}; // struct AstypeAttrs.
@@ -44,15 +41,11 @@ struct AstypeAttrs : public AttrsNode {
struct WrapParamAttrs : public AttrsNode {
DLDataType dtype;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<WrapParamAttrs>().def_ro("dtype", &WrapParamAttrs::dtype, "Target data type",
- refl::DefaultValue((DLDataType{kDLFloat, 32, 1})));
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.WrapParamAttrs", WrapParamAttrs, AttrsNode);
}; // struct WrapParamAttrs.
} // namespace relax
} // namespace tvm
-#endif // TVM_RELAX_ATTRS_DATATYPE_H_
+#endif // TVM_RELAX_OP_DATATYPE_H_
diff --git a/include/tvm/relax/attrs/distributed.h b/include/tvm/relax/op/distributed.h
similarity index 72%
rename from include/tvm/relax/attrs/distributed.h
rename to include/tvm/relax/op/distributed.h
index 78e63b3306..42049ad84f 100644
--- a/include/tvm/relax/attrs/distributed.h
+++ b/include/tvm/relax/op/distributed.h
@@ -18,11 +18,11 @@
*/
/*!
- * \file tvm/relax/attrs/distributed.h
+ * \file tvm/relax/op/distributed.h
* \brief Attributes for redistribute and annotate_sharding operators.
*/
-#ifndef TVM_RELAX_ATTRS_DISTRIBUTED_H_
-#define TVM_RELAX_ATTRS_DISTRIBUTED_H_
+#ifndef TVM_RELAX_OP_DISTRIBUTED_H_
+#define TVM_RELAX_OP_DISTRIBUTED_H_
#include <tvm/relax/distributed/global_info.h>
#include <tvm/relax/distributed/type.h>
@@ -36,18 +36,11 @@ struct DistributionAttrs : public AttrsNode {
distributed::DeviceMesh device_mesh;
distributed::Placement placement;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<DistributionAttrs>()
- .def_ro("device_mesh", &DistributionAttrs::device_mesh,
- "The device mesh of a tensor's distribution plan")
- .def_ro("placement", &DistributionAttrs::placement,
- "The placement of a tensor's distribution plan");
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.DistributionAttrs", DistributionAttrs, AttrsNode);
}; // struct DistributionAttrs
} // namespace relax
} // namespace tvm
-#endif // TVM_RELAX_ATTRS_DISTRIBUTED_H_
+#endif // TVM_RELAX_OP_DISTRIBUTED_H_
diff --git a/include/tvm/relax/op/image.h b/include/tvm/relax/op/image.h
new file mode 100644
index 0000000000..9aa1beb98b
--- /dev/null
+++ b/include/tvm/relax/op/image.h
@@ -0,0 +1,86 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm/relax/op/image.h
+ * \brief Attributes for image operators.
+ */
+#ifndef TVM_RELAX_OP_IMAGE_H_
+#define TVM_RELAX_OP_IMAGE_H_
+
+#include <tvm/relax/expr.h>
+
+namespace tvm {
+namespace relax {
+
+/*! \brief Attributes used in image resize2d operator */
+struct Resize2DAttrs : public AttrsNode {
+ ffi::Array<FloatImm> roi;
+ ffi::String layout;
+ ffi::String method;
+ ffi::String coordinate_transformation_mode;
+ ffi::String rounding_method;
+ double cubic_alpha;
+ int cubic_exclude;
+ double extrapolation_value;
+ ffi::Optional<DLDataType> out_dtype;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Resize2DAttrs", Resize2DAttrs, AttrsNode);
+}; // struct Resize2dAttrs
+
+/*! \brief Attributes used in image resize3d operator */
+struct Resize3DAttrs : public AttrsNode {
+ ffi::Array<FloatImm> roi;
+ ffi::String layout;
+ ffi::String method;
+ ffi::String coordinate_transformation_mode;
+ ffi::String rounding_method;
+ double cubic_alpha;
+ int cubic_exclude;
+ double extrapolation_value;
+ ffi::Optional<DLDataType> out_dtype;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Resize3DAttrs", Resize3DAttrs, AttrsNode);
+}; // struct Resize3DAttrs
+
+/*! \brief Attributes used in image grid_sample operator */
+struct GridSampleAttrs : public AttrsNode {
+ ffi::String method;
+ ffi::String layout;
+ ffi::String padding_mode;
+ bool align_corners;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GridSampleAttrs", GridSampleAttrs, AttrsNode);
+}; // struct GridSampleAttrs
+
+/*! \brief Attributes used in image affine_grid operator */
+struct AffineGridAttrs : public AttrsNode {
+ bool align_corners;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AffineGridAttrs", AffineGridAttrs, AttrsNode);
+}; // struct AffineGridAttrs
+
+} // namespace relax
+} // namespace tvm
+
+#endif // TVM_RELAX_OP_IMAGE_H_
diff --git a/include/tvm/relax/attrs/index.h b/include/tvm/relax/op/index.h
similarity index 57%
rename from include/tvm/relax/attrs/index.h
rename to include/tvm/relax/op/index.h
index f275f67bb2..adff401a8b 100644
--- a/include/tvm/relax/attrs/index.h
+++ b/include/tvm/relax/op/index.h
@@ -18,11 +18,11 @@
*/
/*!
- * \file tvm/relax/attrs/index.h
+ * \file tvm/relax/op/index.h
* \brief Attributes for indexing operators.
*/
-#ifndef TVM_RELAX_ATTRS_INDEX_H_
-#define TVM_RELAX_ATTRS_INDEX_H_
+#ifndef TVM_RELAX_OP_INDEX_H_
+#define TVM_RELAX_OP_INDEX_H_
#include <tvm/relax/expr.h>
@@ -34,14 +34,7 @@ struct TakeAttrs : public AttrsNode {
ffi::Optional<int64_t> axis;
ffi::String mode;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<TakeAttrs>()
- .def_ro("axis", &TakeAttrs::axis, "The axis over which to select values.",
- refl::DefaultValue(ffi::Optional<int64_t>{}))
- .def_ro("mode", &TakeAttrs::mode, "The mode for handling out-of-bounds indices.",
- refl::DefaultValue("fast"));
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TakeAttrs", TakeAttrs, AttrsNode);
}; // struct TakeAttrs
@@ -49,12 +42,7 @@ struct TakeAttrs : public AttrsNode {
struct TakeBackwardAttrs : public AttrsNode {
ffi::Optional<int64_t> axis;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<TakeBackwardAttrs>().def_ro("axis", &TakeBackwardAttrs::axis,
- "The axis over which values were selected.",
- refl::DefaultValue(ffi::Optional<int64_t>{}));
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TakeBackwardAttrs", TakeBackwardAttrs, AttrsNode);
};
@@ -62,18 +50,11 @@ struct TakeBackwardAttrs : public AttrsNode {
struct StridedSliceAttrs : public AttrsNode {
bool assume_inbound;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<StridedSliceAttrs>().def_ro(
- "assume_inbound", &StridedSliceAttrs::assume_inbound,
- "Whether to assume the indices are in bound. If it is set to false, "
- "out of bound indices will be clipped to the bound.",
- refl::DefaultValue(true));
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.StridedSliceAttrs", StridedSliceAttrs, AttrsNode);
}; // struct StridedSliceAttrs
} // namespace relax
} // namespace tvm
-#endif // TVM_RELAX_ATTRS_INDEX_H_
+#endif // TVM_RELAX_OP_INDEX_H_
diff --git a/include/tvm/relax/attrs/linear_algebra.h b/include/tvm/relax/op/linear_algebra.h
similarity index 65%
rename from include/tvm/relax/attrs/linear_algebra.h
rename to include/tvm/relax/op/linear_algebra.h
index 856bee34e7..75415b2f1a 100644
--- a/include/tvm/relax/attrs/linear_algebra.h
+++ b/include/tvm/relax/op/linear_algebra.h
@@ -18,11 +18,11 @@
*/
/*!
- * \file tvm/relax/attrs/linear_algebra.h
+ * \file tvm/relax/op/linear_algebra.h
* \brief Attributes for linear algebra operators.
*/
-#ifndef TVM_RELAX_ATTRS_LINEAR_ALGEBRA_H_
-#define TVM_RELAX_ATTRS_LINEAR_ALGEBRA_H_
+#ifndef TVM_RELAX_OP_LINEAR_ALGEBRA_H_
+#define TVM_RELAX_OP_LINEAR_ALGEBRA_H_
#include <tvm/relax/expr.h>
@@ -33,12 +33,7 @@ namespace relax {
struct MatmulAttrs : public AttrsNode {
ffi::Optional<DLDataType> out_dtype;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<MatmulAttrs>().def_ro("out_dtype", &MatmulAttrs::out_dtype,
- "The data type of the output tensor",
- refl::DefaultValue(ffi::Optional<DLDataType>{}));
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MatmulAttrs", MatmulAttrs, AttrsNode);
}; // struct MatmulAttrs
@@ -46,15 +41,11 @@ struct MatmulAttrs : public AttrsNode {
struct EinsumAttrs : public AttrsNode {
ffi::String subscripts;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<EinsumAttrs>().def_ro("subscripts", &EinsumAttrs::subscripts,
- "The einsum expression string");
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.EinsumAttrs", EinsumAttrs, AttrsNode);
}; // struct EinsumAttrs
} // namespace relax
} // namespace tvm
-#endif // TVM_RELAX_ATTRS_LINEAR_ALGEBRA_H_
+#endif // TVM_RELAX_OP_LINEAR_ALGEBRA_H_
diff --git a/include/tvm/relax/op/manipulate.h b/include/tvm/relax/op/manipulate.h
new file mode 100644
index 0000000000..ffae61bc2c
--- /dev/null
+++ b/include/tvm/relax/op/manipulate.h
@@ -0,0 +1,200 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm/relax/op/manipulate.h
+ * \brief Attributes for tensor manipulation operators.
+ */
+#ifndef TVM_RELAX_OP_MANIPULATE_H_
+#define TVM_RELAX_OP_MANIPULATE_H_
+
+#include <tvm/relax/expr.h>
+#include <tvm/tirx/index_map.h>
+
+namespace tvm {
+namespace relax {
+
+/*! \brief Attributes used in concat operators */
+struct ConcatAttrs : public AttrsNode {
+ ffi::Optional<int64_t> axis;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ConcatAttrs", ConcatAttrs, AttrsNode);
+}; // struct ConcatAttrs
+
+/*! \brief Attributes used in expand_dims operators */
+struct ExpandDimsAttrs : public AttrsNode {
+ ffi::Array<int64_t> axis;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ExpandDimsAttrs", ExpandDimsAttrs, AttrsNode);
+}; // struct ExpandDimsAttrs
+
+/*! \brief Attributes used in layout_transform operator */
+struct LayoutTransformAttrs : public AttrsNode {
+ tirx::IndexMap index_map;
+ // pad_value is chosen to be of PrimExpr type, as it represents constant TIR POD expression. This
+ // needs to be revisited in case PrimExpr is evolved to represent symbolic expression in future.
+ ffi::Optional<PrimExpr> pad_value;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LayoutTransformAttrs", LayoutTransformAttrs,
+ AttrsNode);
+}; // struct LayoutTransformAttrs
+
+/*! \brief Attributes used in permute_dims operator */
+struct PermuteDimsAttrs : public AttrsNode {
+ ffi::Optional<ffi::Array<int64_t>> axes;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PermuteDimsAttrs", PermuteDimsAttrs, AttrsNode);
+}; // struct PermuteDimsAttrs
+
+/*! \brief Attributes used in split operator */
+struct SplitAttrs : public AttrsNode {
+ ffi::ObjectRef indices_or_sections;
+ int axis;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SplitAttrs", SplitAttrs, AttrsNode);
+}; // struct SplitAttrs
+
+/*! \brief Attributes used in squeeze operators */
+struct SqueezeAttrs : public AttrsNode {
+ ffi::Optional<ffi::Array<int64_t>> axis;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SqueezeAttrs", SqueezeAttrs, AttrsNode);
+}; // struct SqueezeAttrs
+
+/*! \brief Attributes used in stack operators */
+struct StackAttrs : public AttrsNode {
+ ffi::Optional<int64_t> axis;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.StackAttrs", StackAttrs, AttrsNode);
+}; // struct StackAttrs
+
+/*! \brief Attributes used in repeat operators */
+struct RepeatAttrs : public AttrsNode {
+ int repeats;
+ ffi::Optional<int64_t> axis;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.RepeatAttrs", RepeatAttrs, AttrsNode);
+}; // struct RepeatAttrs
+
+/*! \brief Attributes used in tile operators */
+struct TileAttrs : public AttrsNode {
+ ffi::Array<int64_t> repeats;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TileAttrs", TileAttrs, AttrsNode);
+}; // struct TileAttrs
+
+/*! \brief Attributes used in flip operators */
+struct FlipAttrs : public AttrsNode {
+ int64_t axis;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.FlipAttrs", FlipAttrs, AttrsNode);
+}; // struct FlipAttrs
+
+/*! \brief Attributes used in reverse_sequence operators */
+struct ReverseSequenceAttrs : public AttrsNode {
+ int64_t seq_axis;
+ int64_t batch_axis;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ReverseSequenceAttrs", ReverseSequenceAttrs,
+ AttrsNode);
+}; // struct ReverseSequenceAttrs
+
+/*! \brief Attributes used in gather_elements operators */
+struct GatherElementsAttrs : public AttrsNode {
+ int64_t axis;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GatherElementsAttrs", GatherElementsAttrs,
+ AttrsNode);
+}; // struct GatherElementsAttrs
+
+/*! \brief Attributes used in gather_nd operators */
+struct GatherNDAttrs : public AttrsNode {
+ int64_t batch_dims;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GatherNDAttrs", GatherNDAttrs, AttrsNode);
+}; // struct GatherNDAttrs
+
+/*! \brief Attributes used in index_put operator */
+struct IndexPutAttrs : public AttrsNode {
+ bool accumulate;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.IndexPutAttrs", IndexPutAttrs, AttrsNode);
+}; // struct IndexPutAttrs
+
+/*! \brief Attribute used in meshgrid operator */
+struct MeshgridAttrs : public AttrsNode {
+ ffi::Optional<ffi::String> indexing;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MeshgridAttrs", MeshgridAttrs, AttrsNode);
+};
+
+/*! \brief Attributes used in scatter_elements operators */
+struct ScatterElementsAttrs : public AttrsNode {
+ int64_t axis;
+ ffi::String reduction;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScatterElementsAttrs", ScatterElementsAttrs,
+ AttrsNode);
+}; // struct ScatterElementsAttrs
+
+/*! \brief Attributes used in scatter_nd operators */
+struct ScatterNDAttrs : public AttrsNode {
+ ffi::String reduction;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScatterNDAttrs", ScatterNDAttrs, AttrsNode);
+}; // struct ScatterNDAttrs
+
+/*! \brief Attributes used in slice_scatter operator */
+struct SliceScatterAttrs : public AttrsNode {
+ int axis;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SliceScatterAttrs", SliceScatterAttrs, AttrsNode);
+}; // struct SliceScatterAttrs
+
+/*! \brief Attributes used in one_hot operator */
+struct OneHotAttrs : public AttrsNode {
+ int depth;
+ int axis;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.OneHotAttrs", OneHotAttrs, AttrsNode);
+}; // struct OneHotAttrs
+
+} // namespace relax
+} // namespace tvm
+
+#endif // TVM_RELAX_OP_MANIPULATE_H_
diff --git a/include/tvm/relax/op/nn.h b/include/tvm/relax/op/nn.h
new file mode 100644
index 0000000000..3a147ac54a
--- /dev/null
+++ b/include/tvm/relax/op/nn.h
@@ -0,0 +1,345 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm/relax/op/nn.h
+ * \brief Attributes for neural network operators.
+ */
+#ifndef TVM_RELAX_OP_NN_H_
+#define TVM_RELAX_OP_NN_H_
+
+#include <tvm/relax/expr.h>
+
+namespace tvm {
+namespace relax {
+
+/*! \brief Attributes used in Conv1d operator */
+struct Conv1DAttrs : public AttrsNode {
+ ffi::Array<int64_t> strides;
+ ffi::Array<int64_t> padding;
+ ffi::Array<int64_t> dilation;
+ int groups;
+ ffi::String data_layout;
+ ffi::String kernel_layout;
+ ffi::String out_layout;
+ ffi::Optional<DLDataType> out_dtype;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv1DAttrs", Conv1DAttrs, AttrsNode);
+}; // struct Conv1dAttrs
+
+/*! \brief Attributes used in Conv2d operator */
+struct Conv2DAttrs : public AttrsNode {
+ ffi::Array<int64_t> strides;
+ ffi::Array<int64_t> padding;
+ ffi::Array<int64_t> dilation;
+ int groups;
+ ffi::String data_layout;
+ ffi::String kernel_layout;
+ ffi::String out_layout;
+ ffi::Optional<DLDataType> out_dtype;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv2DAttrs", Conv2DAttrs, AttrsNode);
+}; // struct Conv2dAttrs
+
+/*! \brief Attributes used in Conv3d operator */
+struct Conv3DAttrs : public AttrsNode {
+ ffi::Array<int64_t> strides;
+ ffi::Array<int64_t> padding;
+ ffi::Array<int64_t> dilation;
+ int groups;
+ ffi::String data_layout;
+ ffi::String kernel_layout;
+ ffi::String out_layout;
+ ffi::Optional<DLDataType> out_dtype;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv3DAttrs", Conv3DAttrs, AttrsNode);
+}; // struct Conv3dAttrs
+
+/*! \brief Attributes used in Conv1DTranspose operator */
+struct Conv1DTransposeAttrs : public AttrsNode {
+ ffi::Array<int64_t> strides;
+ ffi::Array<int64_t> padding;
+ ffi::Array<int64_t> output_padding;
+ ffi::Array<int64_t> dilation;
+ int groups;
+ ffi::String data_layout;
+ ffi::String kernel_layout;
+ ffi::String out_layout;
+ ffi::Optional<DLDataType> out_dtype;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv1DTransposeAttrs", Conv1DTransposeAttrs,
+ AttrsNode);
+}; // struct Conv1DTransposeAttrs
+
+/*! \brief Attributes used in Conv2d operator */
+struct Conv2DTransposeAttrs : public AttrsNode {
+ ffi::Array<int64_t> strides;
+ ffi::Array<int64_t> padding;
+ ffi::Array<int64_t> output_padding;
+ ffi::Array<int64_t> dilation;
+ int groups;
+ ffi::String data_layout;
+ ffi::String kernel_layout;
+ ffi::String out_layout;
+ ffi::Optional<DLDataType> out_dtype;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv2DTransposeAttrs", Conv2DTransposeAttrs,
+ AttrsNode);
+}; // struct Conv2DTransposeAttrs
+
+/*! \brief Attributes used in Conv3dTranspose operator */
+struct Conv3DTransposeAttrs : public AttrsNode {
+ ffi::Array<int64_t> strides;
+ ffi::Array<int64_t> padding;
+ ffi::Array<int64_t> output_padding;
+ ffi::Array<int64_t> dilation;
+ int groups;
+ ffi::String data_layout;
+ ffi::String kernel_layout;
+ ffi::String out_layout;
+ ffi::Optional<DLDataType> out_dtype;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv3DTransposeAttrs", Conv3DTransposeAttrs,
+ AttrsNode);
+}; // struct Conv3DTransposeAttrs
+
+/*! \brief Attributes used in max_pool1d and avg_pool1d operator */
+struct Pool1DAttrs : public AttrsNode {
+ ffi::Array<int64_t> pool_size;
+ ffi::Array<int64_t> strides;
+ ffi::Array<int64_t> padding;
+ ffi::Array<int64_t> dilation;
+ bool ceil_mode;
+ bool count_include_pad;
+ ffi::String layout;
+ ffi::String out_layout;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool1DAttrs", Pool1DAttrs, AttrsNode);
+}; // struct Pool1dAttrs
+
+/*! \brief Attributes used in max_pool2d and avg_pool2d operator */
+struct Pool2DAttrs : public AttrsNode {
+ ffi::Array<int64_t> pool_size;
+ ffi::Array<int64_t> strides;
+ ffi::Array<int64_t> padding;
+ ffi::Array<int64_t> dilation;
+ bool ceil_mode;
+ bool count_include_pad;
+ ffi::String layout;
+ ffi::String out_layout;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool2DAttrs", Pool2DAttrs, AttrsNode);
+}; // struct Pool2dAttrs
+
+/*! \brief Attributes used in max_pool3d and avg_pool3d operator */
+struct Pool3DAttrs : public AttrsNode {
+ ffi::Array<int64_t> pool_size;
+ ffi::Array<int64_t> strides;
+ ffi::Array<int64_t> padding;
+ ffi::Array<int64_t> dilation;
+ bool ceil_mode;
+ bool count_include_pad;
+ ffi::String layout;
+ ffi::String out_layout;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool3DAttrs", Pool3DAttrs, AttrsNode);
+}; // struct Pool3dAttrs
+
+/*! \brief Attributes for 1d adaptive pool operator */
+struct AdaptivePool1DAttrs : public AttrsNode {
+ ffi::Optional<ffi::Array<int64_t>> output_size;
+ ffi::String layout;
+ ffi::String out_layout;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AdaptivePool1DAttrs", AdaptivePool1DAttrs,
+ AttrsNode);
+}; // struct AdaptivePool1DAttrs
+
+/*! \brief Attributes for 2d adaptive pool operator */
+struct AdaptivePool2DAttrs : public AttrsNode {
+ ffi::Optional<ffi::Array<int64_t>> output_size;
+ ffi::String layout;
+ ffi::String out_layout;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AdaptivePool2DAttrs", AdaptivePool2DAttrs,
+ AttrsNode);
+}; // struct AdaptivePool2DAttrs
+
+/*! \brief Attributes for 3d adaptive pool operator */
+struct AdaptivePool3DAttrs : public AttrsNode {
+ ffi::Optional<ffi::Array<int64_t>> output_size;
+ ffi::String layout;
+ ffi::String out_layout;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AdaptivePool3DAttrs", AdaptivePool3DAttrs,
+ AttrsNode);
+}; // struct AdaptivePool3DAttrs
+
+/*! \brief Attributes used in softmax operators */
+struct SoftmaxAttrs : public AttrsNode {
+ int axis;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SoftmaxAttrs", SoftmaxAttrs, AttrsNode);
+};
+
+/*! \brief Attributes used in softmax operators */
+struct LeakyReluAttrs : public AttrsNode {
+ double alpha;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LeakyReluAttrs", LeakyReluAttrs, AttrsNode);
+};
+
+/*! \brief Attributes used in softplus operators */
+struct SoftplusAttrs : public AttrsNode {
+ double beta;
+ double threshold;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SoftplusAttrs", SoftplusAttrs, AttrsNode);
+};
+
+/*! \brief Attributes used in PReLU operator */
+struct PReluAttrs : public AttrsNode {
+ int axis;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PReluAttrs", PReluAttrs, AttrsNode);
+};
+
+/*! \brief Attributes used in batch_norm operator */
+struct BatchNormAttrs : public AttrsNode {
+ int axis;
+ double epsilon;
+ bool center;
+ bool scale;
+ double momentum;
+ bool training;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.BatchNormAttrs", BatchNormAttrs, AttrsNode);
+}; // struct BatchNormAttrs
+
+/*! \brief Attributes used in layer_norm operator */
+struct LayerNormAttrs : public AttrsNode {
+ ffi::Array<int64_t> axes;
+ double epsilon;
+ bool center;
+ bool scale;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LayerNormAttrs", LayerNormAttrs, AttrsNode);
+}; // struct LayerNormAttrs
+
+/*! \brief Attributes used in group_norm operator */
+struct GroupNormAttrs : public AttrsNode {
+ int num_groups;
+ int channel_axis;
+ ffi::Array<int64_t> axes;
+ double epsilon;
+ bool center;
+ bool scale;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GroupNormAttrs", GroupNormAttrs, AttrsNode);
+}; // struct GroupNormAttrs
+
+/*! \brief Attributes used in instance_norm operator */
+struct InstanceNormAttrs : public AttrsNode {
+ int channel_axis;
+ ffi::Array<int64_t> axes;
+ double epsilon;
+ bool center;
+ bool scale;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.InstanceNormAttrs", InstanceNormAttrs, AttrsNode);
+}; // struct InstanceNormAttrs
+
+/*! \brief Attributes used in rms_norm operator */
+struct RMSNormAttrs : public AttrsNode {
+ ffi::Array<int64_t> axes;
+ double epsilon;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.RMSNormAttrs", RMSNormAttrs, AttrsNode);
+}; // struct RMSNormAttrs
+
+/*! \brief Attributes used in nll_loss operator */
+struct NLLLossAttrs : public AttrsNode {
+ ffi::String reduction;
+ int ignore_index;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.NLLLossAttrs", NLLLossAttrs, AttrsNode);
+}; // struct NLLLossAttrs
+
+/*! \brief Attributes used in dropout operator */
+struct DropoutAttrs : public AttrsNode {
+ double rate;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.DropoutAttrs", DropoutAttrs, AttrsNode);
+}; // struct DropoutAttrs
+
+/*! \brief Attributes used in Attention operator */
+struct AttentionAttrs : public AttrsNode {
+ ffi::Optional<FloatImm> scale;
+ ffi::Optional<ffi::String> causal_mask;
+ ffi::Optional<IntImm> window_size;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AttentionAttrs", AttentionAttrs, AttrsNode);
+}; // struct AttentionAttrs
+
+/*! \brief Attributes used for the padding operator */
+struct PadAttrs : public AttrsNode {
+ ffi::Array<int64_t> pad_width;
+ double pad_value = 0.0;
+ tvm::ffi::String pad_mode;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PadAttrs", PadAttrs, AttrsNode);
+};
+
+/*! \brief Attributes used for the pixel shuffle operator */
+struct PixelShuffleAttrs : public AttrsNode {
+ int upscale_factor;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PixelShuffleAttrs", PixelShuffleAttrs, AttrsNode);
+};
+
+} // namespace relax
+} // namespace tvm
+
+#endif // TVM_RELAX_OP_NN_H_
diff --git a/include/tvm/relax/attrs/op.h b/include/tvm/relax/op/op.h
similarity index 63%
rename from include/tvm/relax/attrs/op.h
rename to include/tvm/relax/op/op.h
index 54d6bba916..b5bc9ae5bb 100644
--- a/include/tvm/relax/attrs/op.h
+++ b/include/tvm/relax/op/op.h
@@ -18,11 +18,11 @@
*/
/*!
- * \file tvm/relax/attrs/op.h
+ * \file tvm/relax/op/op.h
* \brief Attributes for relax specific operators.
*/
-#ifndef TVM_RELAX_ATTRS_OP_H_
-#define TVM_RELAX_ATTRS_OP_H_
+#ifndef TVM_RELAX_OP_OP_H_
+#define TVM_RELAX_OP_OP_H_
#include <tvm/relax/expr.h>
#include <tvm/relax/global_info.h>
@@ -35,15 +35,7 @@ struct CallTIRWithGradAttrs : public AttrsNode {
ffi::String te_grad_name;
ffi::Map<ffi::String, Any> te_grad_kwargs;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<CallTIRWithGradAttrs>()
- .def_ro(
- "te_grad_name", &CallTIRWithGradAttrs::te_grad_name,
- "The name of the te gradient function associated with this call_tir_with_grad node.")
- .def_ro("te_grad_kwargs", &CallTIRWithGradAttrs::te_grad_kwargs,
- "The keyword arguments passed to the te gradient function.");
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.CallTIRWithGradAttrs", CallTIRWithGradAttrs,
AttrsNode);
}; // struct CallTIRAttrs
@@ -59,11 +51,7 @@ struct CallTIRInplaceAttrs : public AttrsNode {
*/
ffi::Array<int64_t> inplace_indices;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<CallTIRInplaceAttrs>().def_ro("inplace_indices",
- &CallTIRInplaceAttrs::inplace_indices);
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.CallTIRInplaceAttrs", CallTIRInplaceAttrs,
AttrsNode);
}; // struct CallTIRInplaceAttrs
@@ -79,11 +67,7 @@ struct CallInplacePackedAttrs : public AttrsNode {
*/
ffi::Array<int64_t> inplace_indices;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<CallInplacePackedAttrs>().def_ro("inplace_indices",
- &CallInplacePackedAttrs::inplace_indices);
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.CallInplacePackedAttrs", CallInplacePackedAttrs,
AttrsNode);
}; // struct CallInplacePackedAttrs
@@ -92,11 +76,7 @@ struct CallInplacePackedAttrs : public AttrsNode {
struct ToVDeviceAttrs : public AttrsNode {
VDevice dst_vdevice;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ToVDeviceAttrs>().def_ro("dst_vdevice", &ToVDeviceAttrs::dst_vdevice,
- "The destination device where the data is copied to.");
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ToVDeviceAttrs", ToVDeviceAttrs, AttrsNode);
}; // struct ToVDeviceAttrs
@@ -106,18 +86,11 @@ struct HintOnDeviceAttrs : public AttrsNode {
int32_t index;
MemoryScope memory_scope;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<HintOnDeviceAttrs>()
- .def_ro("device_type", &HintOnDeviceAttrs::device_type,
- "The device type where the data is supposed to be executed.")
- .def_ro("index", &HintOnDeviceAttrs::index, "The device id.")
- .def_ro("memory_scope", &HintOnDeviceAttrs::memory_scope, "The device memory scope.");
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.HintOnDeviceAttrs", HintOnDeviceAttrs, AttrsNode);
}; // struct HintOnDeviceAttrs
} // namespace relax
} // namespace tvm
-#endif // TVM_RELAX_ATTRS_OP_H_
+#endif // TVM_RELAX_OP_OP_H_
diff --git a/include/tvm/relax/attrs/qdq.h b/include/tvm/relax/op/qdq.h
similarity index 68%
rename from include/tvm/relax/attrs/qdq.h
rename to include/tvm/relax/op/qdq.h
index be95b9e7b8..978028c916 100644
--- a/include/tvm/relax/attrs/qdq.h
+++ b/include/tvm/relax/op/qdq.h
@@ -18,11 +18,11 @@
*/
/*!
- * \file include/tvm/relax/attrs/qdq.h
+ * \file include/tvm/relax/op/qdq.h
* \brief Attributes for quantize/dequantize operators.
*/
-#ifndef TVM_RELAX_ATTRS_QDQ_H_
-#define TVM_RELAX_ATTRS_QDQ_H_
+#ifndef TVM_RELAX_OP_QDQ_H_
+#define TVM_RELAX_OP_QDQ_H_
#include <tvm/relax/expr.h>
@@ -34,19 +34,11 @@ struct QuantizeAttrs : public AttrsNode {
DLDataType out_dtype;
int axis;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<QuantizeAttrs>()
- .def_ro("out_dtype", &QuantizeAttrs::out_dtype, "Output data type.")
- .def_ro("axis", &QuantizeAttrs::axis,
- "The output channel axis for channel wise quantization/dequantization. "
- "Default value is -1, which corresponds to the last axis.",
- refl::DefaultValue(-1));
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.QuantizeAttrs", QuantizeAttrs, AttrsNode);
}; // QuantizeAttrs
} // namespace relax
} // namespace tvm
-#endif // TVM_RELAX_ATTRS_QDQ_H_
+#endif // TVM_RELAX_OP_QDQ_H_
diff --git a/include/tvm/relax/attrs/sampling.h b/include/tvm/relax/op/sampling.h
similarity index 74%
rename from include/tvm/relax/attrs/sampling.h
rename to include/tvm/relax/op/sampling.h
index 07b7de25e5..771e75cce2 100644
--- a/include/tvm/relax/attrs/sampling.h
+++ b/include/tvm/relax/op/sampling.h
@@ -18,11 +18,11 @@
*/
/*!
- * \file tvm/relax/attrs/sampling.h
+ * \file tvm/relax/op/sampling.h
* \brief Attributes for sampling operators.
*/
-#ifndef TVM_RELAX_ATTRS_SAMPLING_H_
-#define TVM_RELAX_ATTRS_SAMPLING_H_
+#ifndef TVM_RELAX_OP_SAMPLING_H_
+#define TVM_RELAX_OP_SAMPLING_H_
#include <tvm/relax/expr.h>
@@ -33,12 +33,7 @@ namespace relax {
struct MultinomialFromUniformAttrs : public AttrsNode {
DLDataType dtype;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<MultinomialFromUniformAttrs>().def_ro(
- "dtype", &MultinomialFromUniformAttrs::dtype, "Data type of the output indices.",
- refl::DefaultValue((DLDataType{kDLInt, 64, 1})));
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MultinomialFromUniformAttrs",
MultinomialFromUniformAttrs, AttrsNode);
}; // struct MultinomialFromUniformAttrs
@@ -46,4 +41,4 @@ struct MultinomialFromUniformAttrs : public AttrsNode {
} // namespace relax
} // namespace tvm
-#endif // TVM_RELAX_ATTRS_SAMPLING_H_
+#endif // TVM_RELAX_OP_SAMPLING_H_
diff --git a/include/tvm/relax/attrs/search.h b/include/tvm/relax/op/search.h
similarity index 56%
rename from include/tvm/relax/attrs/search.h
rename to include/tvm/relax/op/search.h
index c750093368..6bfe61ad5f 100644
--- a/include/tvm/relax/attrs/search.h
+++ b/include/tvm/relax/op/search.h
@@ -18,11 +18,11 @@
*/
/*!
- * \file tvm/relax/attrs/search.h
+ * \file tvm/relax/op/search.h
* \brief Attributes for search operators.
*/
-#ifndef TVM_RELAX_ATTRS_SEARCH_H_
-#define TVM_RELAX_ATTRS_SEARCH_H_
+#ifndef TVM_RELAX_OP_SEARCH_H_
+#define TVM_RELAX_OP_SEARCH_H_
#include <tvm/relax/expr.h>
@@ -34,18 +34,7 @@ struct ArgmaxArgminAttrs : public AttrsNode {
ffi::Optional<int64_t> axis;
bool keepdims;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ArgmaxArgminAttrs>()
- .def_ro("axis", &ArgmaxArgminAttrs::axis,
- "The axis along which to perform the argmin/argmax.",
- refl::DefaultValue(ffi::Optional<int64_t>{}))
- .def_ro("keepdims", &ArgmaxArgminAttrs::keepdims,
- "If this is set to `True`, the reduced axis is left in the result as dimension "
- "with size "
- "one.",
- refl::DefaultValue(false));
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ArgmaxArgminAttrs", ArgmaxArgminAttrs, AttrsNode);
}; // struct ArgmaxArgminAttrs
@@ -54,19 +43,11 @@ struct BucketizeAttrs : public tvm::AttrsNode {
bool out_int32;
bool right;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<BucketizeAttrs>()
- .def_ro("out_int32", &BucketizeAttrs::out_int32,
- "Indicate the output datatype, int32 if True, int64 otherwise.",
- refl::DefaultValue(false))
- .def_ro("right", &BucketizeAttrs::right, "Determines the behavior for values in boundaries",
- refl::DefaultValue(false));
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.BucketizeAttrs", BucketizeAttrs, AttrsNode);
}; // struct BucketizeAttrs
} // namespace relax
} // namespace tvm
-#endif // TVM_RELAX_ATTRS_SEARCH_H_
+#endif // TVM_RELAX_OP_SEARCH_H_
diff --git a/include/tvm/relax/op/sorting.h b/include/tvm/relax/op/sorting.h
new file mode 100644
index 0000000000..4be091cd36
--- /dev/null
+++ b/include/tvm/relax/op/sorting.h
@@ -0,0 +1,67 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm/relax/op/sorting.h
+ * \brief Attributes for sorting operators.
+ */
+#ifndef TVM_RELAX_OP_SORTING_H_
+#define TVM_RELAX_OP_SORTING_H_
+
+#include <tvm/relax/expr.h>
+#include <tvm/tirx/index_map.h>
+
+namespace tvm {
+namespace relax {
+
+/*! \brief Attributes used in sort operator */
+struct SortAttrs : public AttrsNode {
+ int axis;
+ bool descending;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SortAttrs", SortAttrs, AttrsNode);
+}; // struct SortAttrs
+
+/*! \brief Attributes used in argsort operator */
+struct ArgsortAttrs : public AttrsNode {
+ int axis;
+ bool descending;
+ ffi::Optional<DLDataType> dtype;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ArgsortAttrs", ArgsortAttrs, AttrsNode);
+}; // struct ArgsortAttrs
+
+/*! \brief Attributes used in topk operator */
+struct TopKAttrs : public AttrsNode {
+ int k;
+ int axis;
+ bool largest;
+ ffi::String ret_type;
+ ffi::Optional<DLDataType> dtype;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TopKAttrs", TopKAttrs, AttrsNode);
+}; // struct TopKAttrs
+
+} // namespace relax
+} // namespace tvm
+
+#endif // TVM_RELAX_OP_SORTING_H_
diff --git a/include/tvm/relax/attrs/statistical.h b/include/tvm/relax/op/statistical.h
similarity index 51%
rename from include/tvm/relax/attrs/statistical.h
rename to include/tvm/relax/op/statistical.h
index 53e4481557..2d7ec0c733 100644
--- a/include/tvm/relax/attrs/statistical.h
+++ b/include/tvm/relax/op/statistical.h
@@ -18,11 +18,11 @@
*/
/*!
- * \file tvm/relax/attrs/statistical.h
+ * \file tvm/relax/op/statistical.h
* \brief Attributes for statistical operators.
*/
-#ifndef TVM_RELAX_ATTRS_STATISTICAL_H_
-#define TVM_RELAX_ATTRS_STATISTICAL_H_
+#ifndef TVM_RELAX_OP_STATISTICAL_H_
+#define TVM_RELAX_OP_STATISTICAL_H_
#include <tvm/relax/expr.h>
@@ -34,18 +34,7 @@ struct StatisticalAttrs : public AttrsNode {
ffi::Optional<ffi::Array<int64_t>> axis;
bool keepdims;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<StatisticalAttrs>()
- .def_ro("axis", &StatisticalAttrs::axis,
- "The axis or axes along which to perform the reduction.",
- refl::DefaultValue(ffi::Optional<ffi::Array<int64_t>>{}))
- .def_ro("keepdims", &StatisticalAttrs::keepdims,
- "If this is set to `True`, the reduced axes are left in the result as dimension "
- "with size "
- "one.",
- refl::DefaultValue(false));
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.StatisticalAttrs", StatisticalAttrs, AttrsNode);
}; // struct StatisticalAttrs
@@ -55,24 +44,11 @@ struct ScanopAttrs : public AttrsNode {
ffi::Optional<DLDataType> dtype;
bool exclusive = false;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ScanopAttrs>()
- .def_ro("axis", &ScanopAttrs::axis,
- "The axis along which to perform the scan computation."
- "The default (None) is to compute over the flattened array.",
- refl::DefaultValue(ffi::Optional<int64_t>{}))
- .def_ro("dtype", &ScanopAttrs::dtype,
- "The output data type."
- "If dtype is not specified, it defaults to the dtype of input data.",
- refl::DefaultValue(ffi::Optional<DLDataType>{}))
- .def_ro("exclusive", &ScanopAttrs::exclusive, "The first element is not included",
- refl::DefaultValue(false));
- }
+ static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScanopAttrs", ScanopAttrs, AttrsNode);
}; // struct ScanopAttrs
} // namespace relax
} // namespace tvm
-#endif // TVM_RELAX_ATTRS_STATISTICAL_H_
+#endif // TVM_RELAX_OP_STATISTICAL_H_
diff --git a/include/tvm/relax/op/vision.h b/include/tvm/relax/op/vision.h
new file mode 100644
index 0000000000..ee9f35ed16
--- /dev/null
+++ b/include/tvm/relax/op/vision.h
@@ -0,0 +1,112 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file tvm/relax/op/vision.h
+ * \brief Auxiliary attributes for vision operators.
+ */
+#ifndef TVM_RELAX_OP_VISION_H_
+#define TVM_RELAX_OP_VISION_H_
+
+#include <tvm/ffi/string.h>
+#include <tvm/ir/attrs.h>
+#include <tvm/ir/type.h>
+#include <tvm/relax/expr.h>
+
+namespace tvm {
+namespace relax {
+
+/*! \brief Attributes used in AllClassNonMaximumSuppression operator */
+struct AllClassNonMaximumSuppressionAttrs : public AttrsNode {
+ ffi::String output_format;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AllClassNonMaximumSuppressionAttrs",
+ AllClassNonMaximumSuppressionAttrs, AttrsNode);
+}; // struct AllClassNonMaximumSuppressionAttrs
+
+/*! \brief Attributes used in ROIAlign operator */
+struct ROIAlignAttrs : public AttrsNode {
+ ffi::Array<int64_t> pooled_size;
+ double spatial_scale;
+ int sample_ratio;
+ bool aligned;
+ ffi::String layout;
+ ffi::String mode;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ROIAlignAttrs", ROIAlignAttrs, AttrsNode);
+}; // struct ROIAlignAttrs
+
+/*! \brief Attributes used in ROIPool operator */
+struct ROIPoolAttrs : public AttrsNode {
+ ffi::Array<int64_t> pooled_size;
+ double spatial_scale;
+ ffi::String layout;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ROIPoolAttrs", ROIPoolAttrs, AttrsNode);
+}; // struct ROIPoolAttrs
+
+/*! \brief Attributes used in GetValidCounts operator */
+struct GetValidCountsAttrs : public AttrsNode {
+ double score_threshold;
+ int id_index;
+ int score_index;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GetValidCountsAttrs", GetValidCountsAttrs,
+ AttrsNode);
+}; // struct GetValidCountsAttrs
+
+/*! \brief Attributes used in NonMaximumSuppression operator */
+struct NonMaximumSuppressionAttrs : public AttrsNode {
+ int max_output_size;
+ double iou_threshold;
+ bool force_suppress;
+ int top_k;
+ int coord_start;
+ int score_index;
+ int id_index;
+ bool return_indices;
+ bool invalid_to_bottom;
+ double soft_nms_sigma;
+ double score_threshold;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.NonMaximumSuppressionAttrs",
+ NonMaximumSuppressionAttrs, AttrsNode);
+}; // struct NonMaximumSuppressionAttrs
+
+/*! \brief Attributes for multibox_transform_loc (SSD / TFLite-style box decode). */
+struct MultiboxTransformLocAttrs : public AttrsNode {
+ bool clip;
+ double threshold;
+ ffi::Array<double> variances;
+ bool keep_background;
+ bool apply_softmax;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MultiboxTransformLocAttrs",
+ MultiboxTransformLocAttrs, AttrsNode);
+}; // struct MultiboxTransformLocAttrs
+
+} // namespace relax
+} // namespace tvm
+
+#endif // TVM_RELAX_OP_VISION_H_
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index 2c89636f4f..d2aa71098d 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -330,7 +330,7 @@ TVM_DLL Pass DefaultGPUSchedule();
TVM_DLL Pass RemoveWeightLayoutRewriteBlock(bool skip_tensor_rewrite = false);
/*!
- * \brief Remove stores of tirx::builtin::undef.
+ * \brief Remove stores of tirx::undef_op.
* \return The pass.
*/
TVM_DLL Pass RemoveStoreUndef();
diff --git a/include/tvm/tirx/attrs.h b/include/tvm/tirx/attrs.h
deleted file mode 100644
index 1d173b3561..0000000000
--- a/include/tvm/tirx/attrs.h
+++ /dev/null
@@ -1,73 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-#ifndef TVM_TIRX_ATTRS_H_
-#define TVM_TIRX_ATTRS_H_
-
-#include <tvm/ffi/container/array.h>
-#include <tvm/ir/attrs.h>
-
-namespace tvm {
-namespace tirx {
-
-/*! \brief Launch metadata for call_ffi_kernel. */
-struct CallFFIKernelAttr : public AttrsNode {
- /*!
- * \brief Ordered launch tags describing the suffix of the call arguments.
- *
- * The first call argument is the kernel symbol, followed by kernel operands
- * and launch values. Flag-only tags consume no value; dynamic shared-memory
- * bytes, when present, are last. Runtime expressions remain in Call.args.
- */
- ffi::Array<ffi::String> launch_params;
-
- static void RegisterReflection() {
- ffi::reflection::ObjectDef<CallFFIKernelAttr>().def_ro("launch_params",
- &CallFFIKernelAttr::launch_params);
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.CallFFIKernelAttr", CallFFIKernelAttr, AttrsNode);
-};
-
-/*! \brief Fixed encoding options for tensormap_encode_tiled. */
-struct TensorMapEncodeTiledAttr : public AttrsNode {
- DLDataType descriptor_dtype;
- int64_t rank;
- int64_t interleave;
- int64_t swizzle;
- int64_t l2_promotion;
- int64_t oob_fill;
- int64_t force_cu_dtype;
-
- static void RegisterReflection() {
- ffi::reflection::ObjectDef<TensorMapEncodeTiledAttr>()
- .def_ro("descriptor_dtype", &TensorMapEncodeTiledAttr::descriptor_dtype)
- .def_ro("rank", &TensorMapEncodeTiledAttr::rank)
- .def_ro("interleave", &TensorMapEncodeTiledAttr::interleave)
- .def_ro("swizzle", &TensorMapEncodeTiledAttr::swizzle)
- .def_ro("l2_promotion", &TensorMapEncodeTiledAttr::l2_promotion)
- .def_ro("oob_fill", &TensorMapEncodeTiledAttr::oob_fill)
- .def_ro("force_cu_dtype", &TensorMapEncodeTiledAttr::force_cu_dtype);
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.TensorMapEncodeTiledAttr", TensorMapEncodeTiledAttr,
- AttrsNode);
-};
-
-} // namespace tirx
-} // namespace tvm
-
-#endif // TVM_TIRX_ATTRS_H_
diff --git a/include/tvm/tirx/builtin.h b/include/tvm/tirx/builtin.h
deleted file mode 100644
index 20d24fa89d..0000000000
--- a/include/tvm/tirx/builtin.h
+++ /dev/null
@@ -1,757 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file tvm/tirx/builtin.h
- * \brief TIRx builtin region and call operations.
- *
- * TIRx builtin operations are stored as tvm::Op.
- * They are processed in the same way as we process Ops.
- *
- * It is not necessary to create a function for every Op,
- * as we can obtain them through Op::Get.
- *
- * This file contains the most commonly used intrinsics or
- * those that have special semantics and need compiler support.
- */
-#ifndef TVM_TIR_BUILTIN_H_
-#define TVM_TIR_BUILTIN_H_
-
-#include <tvm/ir/op.h>
-#include <tvm/ir/prim/expr.h>
-
-namespace tvm {
-namespace tirx {
-
-/*! \brief Collection of builtin region and call operations as Ops. */
-namespace builtin {
-/*!
- * \name Region operations
- * \brief Operations used by RegionStmt to enclose a lexical body.
- *
- * Every region operation registers FRegionGetBodyParams, returning fresh typed
- * body parameters or an empty array for no parameters. Attribute presence alone
- * identifies region support; classification does not invoke the hook. Operands
- * and attributes are evaluated outside the body-parameter scope.
- * \{
- */
-/*!
- * \brief Thread launch region: operands are a nonempty StringImm tag and a
- * scalar signed or unsigned integer extent wider than one bit.
- * The sole body parameter is a fresh thread-index PrimVar matching the extent type.
- * Tags starting with vthread denote virtual threads. There are no attrs or results.
- */
-TVM_DLL const Op& launch_thread();
-
-/*!
- * \brief Mark a user-facing device entry containing device scope definitions.
- * Takes no operands, body parameters, attributes or results.
- */
-TVM_DLL const Op& device_entry();
-/*!
- * \brief Supply lexical device context for allocation and packed-call lowering.
- * Operands are integer device type and device ID; there are no body parameters,
- * attributes or results. The region does not change the active runtime device.
- */
-TVM_DLL const Op& device_context();
-/*!
- * \brief Outline the body as a CPU compute helper named by a StringImm operand.
- * There are no body parameters, attributes or results.
- */
-TVM_DLL const Op& compute_scope();
-/*!
- * \brief Launch a CPU worker team around parallel loops and team barriers.
- * Takes no operands, body parameters, attributes or results.
- */
-TVM_DLL const Op& parallel_launch();
-/*! \} */
-
-/*!
- * \name Call operations
- * \brief Operations invoked through Call with explicit operands and result types.
- * \{
- */
-/*!
- * \brief Allocate a buffer: alloc_tensor(shape, dtype, scope) -> TensorType.
- *
- * Arguments, in order:
- * - args[0]: shape, Tuple of integer extents (IntImm or symbolic integer expressions).
- * - args[1]: dtype, DataTypeImm with a DLDataType payload for the element type.
- * - args[2]: scope, StringImm naming the storage scope.
- *
- * DictAttrs directly holds the allocation annotations, defaulting to an empty dictionary.
- * The TensorType result agrees with the operands and retains buffer access/storage metadata.
- *
- * \code
- * // Example pattern match code for a given Binding:
- * if (const auto* call = binding->value.as<CallNode>();
- * call && call->op.same_as(builtin::alloc_tensor())) {
- * tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
- * DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
- * ffi::String scope = call->args[2].as_or_throw<StringImm>()->value;
- * DictAttrs annotations = call->attrs.as_or_throw<DictAttrs>();
- * }
- * \endcode
- */
-TVM_DLL const Op& alloc_tensor();
-/*!
- * \brief Declare a buffer view: decl_tensor(data, shape, dtype, scope) -> TensorType.
- *
- * Arguments, in order:
- * - args[0]: data, Expr for the existing physical pointer backing the buffer view.
- * - args[1]: shape, Tuple of integer extents (IntImm or symbolic integer expressions).
- * - args[2]: dtype, DataTypeImm with a DLDataType payload for the element type.
- * - args[3]: scope, StringImm naming the storage scope.
- *
- * There are no attributes. The TensorType result agrees with the operands and retains
- * buffer access/storage metadata. The operation binds a view without allocating memory.
- *
- * \code
- * // Example pattern match code for a given Binding:
- * if (const auto* call = binding->value.as<CallNode>();
- * call && call->op.same_as(builtin::decl_tensor())) {
- * Expr data = call->args[0];
- * tvm::Tuple shape = call->args[1].as_or_throw<tvm::Tuple>();
- * DLDataType dtype = call->args[2].as_or_throw<DataTypeImm>()->value;
- * ffi::String scope = call->args[3].as_or_throw<StringImm>()->value;
- * }
- * \endcode
- */
-TVM_DLL const Op& decl_tensor();
-/*!
- * \brief Return from a GPU thread without returning a function value.
- */
-TVM_DLL const Op& thread_return();
-/*!
- * \brief Reinterpret the value using the target type.
- */
-TVM_DLL const Op& reinterpret();
-
-/*!
- * \brief Thread-set filter predicate. Used as the condition of an IfThenElse
- * to narrow the active thread set A for the then-branch. Two forms:
- * filter(var, lo, hi) -- range form, true iff var in [lo, hi)
- * filter(var, cond) -- predicate form (e.g. var == k); true iff cond
- * `var` must be a ScopeIdDef-declared Var at parse time (Verifier Rule 2).
- */
-TVM_DLL const Op& filter();
-
-/*!
- * \brief Analysis-only active-thread selector.
- *
- * ``selector(var, pred)`` denotes the unique value of ``var`` in the current
- * active domain for which ``pred`` is true. It is used only inside
- * ExecContext/DispatchContext metadata, for predicates such as
- * ``ptx.elect_sync()`` whose selected lane cannot be inferred structurally.
- */
-TVM_DLL const Op& selector();
-
-/*!
- * \brief Execute a multiplication between two Q-numbers x and y
- * followed by a right shift s
- * The default rounding rule is to the nearest value, rounding half up
- * (i.e., round(x.1) = x and round (x.5) = x+1)
- */
-TVM_DLL const Op& q_multiply_shift();
-TVM_DLL const Op& q_multiply_shift_per_axis();
-
-/*!
- * \brief Returns the address of an element in the buffer (see pseudocode below).
- *
- * The number of indices should match the dimensionality of the buffer
- * being accessed. If this operation occurs after buffer flattening,
- * the number of indices must be supported by the target (i.e. N>1
- * only on targets that support non-flat memory buffers).
- *
- * Handle address_of(TensorLoad *op) {
- * return &op->buffer_var[op->indices[0], op->indices[1], ..., op->indices[N-1]];
- * }
- */
-TVM_DLL const Op& address_of();
-
-/*!
- * \brief See pesudo code
- *
- * bool isnullptr(void* handle) {
- * return handle == nullptr
- * }
- */
-TVM_DLL const Op& isnullptr();
-
-/*!
- * \brief Check if value is nan
- */
-TVM_DLL const Op& isnan();
-
-/*!
- * \brief Popcount
- */
-TVM_DLL const Op& popcount();
-
-/*!
- * \brief Fused multiply add
- *
- * Type fma(a, b, c) {
- * return a * b + c;
- * }
- */
-TVM_DLL const Op& fma();
-
-/*!
- * \brief Call an extern C function with given name
- * and signature from the types of args in the runtime environment.
- *
- * Type call_extern(name, args...) {
- * return dlsym(name)(args...);
- * }
- *
- * \note This intrinsic does not provide any type checking,
- * and is main used for backward compatibility reasons.
- * Always consider use pre-registered and typed tvm::Op first.
- */
-TVM_DLL const Op& call_extern();
-
-/*!
- * \brief Call an pure extern C function with given name
- * and signature from the types of args in the runtime environment.
- *
- * Type call_pure_extern(name, args...) {
- * return dlsym(name)(args...);
- * }
- *
- * \note This intrinsic does not provide any type checking,
- * and is main used for backward compatibility reasons.
- * Always consider use pre-registered and typed tvm::Op first.
- */
-TVM_DLL const Op& call_pure_extern();
-
-/*!
- * \brief Call an LLVM intrinsic with a given intrinsic id
- * and signature from the types of args in the runtime environment.
- *
- * Type call_llvm_pure_intrin(intrin_id, args...) {
- * return dlsym(name)(args...);
- * }
- *
- * \note This op does not provide any type checking.
- */
-TVM_DLL const Op& call_llvm_intrin();
-
-/*!
- * \brief Call an LLVM pure intrinsic with a given intrinsic id
- * and signature from the types of args in the runtime environment.
- *
- * Type call_llvm_pure_intrin(intrin_id, args...) {
- * return dlsym(name)(args...);
- * }
- *
- * \note This op does not provide any type checking.
- */
-TVM_DLL const Op& call_llvm_pure_intrin();
-
-/*!
- * \brief Call an SPIRV pure GLSL450 intrinsic.
- *
- * Type call_spirv_pure_glsl450(intrin_id, args...) {
- * return dlsym(name)(args...);
- * }
- *
- * \note This op does not provide any type checking.
- */
-TVM_DLL const Op& call_spirv_pure_glsl450();
-
-// TODO(tvm-team) revisit the builtins below
-// some of them can simply become ops with special codegen attr.
-/*!
- * \brief same signature as llvm.prefetch
- */
-TVM_DLL const Op& prefetch();
-
-/*!
- * \brief Get head access address with memory access pattern info.
- *
- * This operator also marks range of the memory access
- * The offset and extent are in unit of the DType(including vectorization factor).
- * rw_mask is a bit_mask setting whether the access is a read(1) or write(2).
- * The access is assume to happen in the current expression.
- *
- * PtrType tvm_access_ptr(Expr dtype, DType* data,
- * int offset, int extent,
- * int rw_mask) {
- * // DType == dtype.type();
- * return &data[offset];
- * }
- */
-TVM_DLL const Op& tvm_access_ptr();
-
-/*!
- * \brief Cast a handle to a typed pointer after adding a byte offset.
- *
- * DType* ptr_byte_offset(void* data, int byte_offset, Expr dtype) {
- * return reinterpret_cast<DType*>(reinterpret_cast<char*>(data) + byte_offset);
- * }
- */
-TVM_DLL const Op& ptr_byte_offset();
-
-/*!
- * \brief Create a function local static handle that iniitalizes to nullptr.
- * can be used to cache function local static resources.
- */
-TVM_DLL const Op& tvm_static_handle();
-
-/*!
- * \brief See pesudo code
- *
- * void* handle_add_byte_offset(void* handle, int offset) {
- * return reinterpret_cast<v*>(reinterpret_cast<char*>(handle) + offset);
- * }
- */
-TVM_DLL const Op& handle_add_byte_offset();
-
-/*!
- * \brief See pesudo code
- *
- * Type tvm_struct_get(StructType* arr, int index, int field_id) {
- * return arr[index]->field;
- * }
- * \sa TVMStructFieldKind
- */
-TVM_DLL const Op& tvm_struct_get();
-
-/*!
- * \brief See pesudo code
- *
- * Handle tvm_struct_set(StructType* arr, int index, int field_id, value) {
- * arr[index]->field = value;
- * }
- * \sa TVMStructFieldKind
- */
-TVM_DLL const Op& tvm_struct_set();
-
-/*!
- * \brief See pesudo code
- *
- * void tvm_throw_last_error() {
- * throw TVMGetLastError();
- * }
- */
-TVM_DLL const Op& tvm_throw_last_error();
-
-/*!
- * \brief See pesudo code
- *
- * dtype in {shape, array, arg_value, arg_tcode}
- *
- * Handle tvm_stack_alloca(string dtype, int num) {
- * return new on stack dtype[num];
- * }
- */
-TVM_DLL const Op& tvm_stack_alloca();
-
-/*!
- * \brief Allocate a shape tuple on stack, return the handle.
- *
- * Handle tvm_stack_make_shape(list args) {
- * ret = alloca stack int64_t[len(args)];
- * for i in range(len(args)):
- * ret[i] = args[i]
- * return &ret[0];
- * }
- */
-TVM_DLL const Op& tvm_stack_make_shape();
-
-/*!
- * \brief Allocate a Tensor(DLTensor) on stack, return the handle.
- *
- * Type tvm_stack_make_array(Expr data,
- * Expr shape,
- * Expr strides,
- * Expr ndim,
- * Expr dtype,
- * Expr elem_offset) {
- * ret = alloca stack DLTensor();
- * ret->data = data;
- * ret->shape = shape;
- * ret->strides = strides != 0 ? strides : nullptr;
- * ret->ndim = ndim;
- * ret->dtype = dtype.type();
- * ret->byte_offset = elem_offset * sizeof(dtype);
- * return ret;
- * }
- */
-TVM_DLL const Op& tvm_stack_make_array();
-
-/*!
- * \brief See pesudo code
- *
- * return_type tvm_call_packed(name, TVMFFIAny* args) {
- * TVMFFIAny result;
- * ModuleNode* env = GetCurrentEnv();
- * const ffi::Function* f = env->GetFuncFromEnv(name);
- * (*f)(args, args, len(args), &result);
- * // return type can be int, float, handle.
- * return cast(return_type, result);
- * }
- */
-TVM_DLL const Op& tvm_call_packed();
-
-/*!
- * \brief Launch a kernel using the packed-function argument convention.
- *
- * Arguments are the kernel symbol, kernel operands, then launch values.
- * CallFFIKernelAttr::launch_params describes the launch-value suffix.
- * Host backends may consume this call directly or lower it to tvm_call_packed.
- */
-TVM_DLL const Op& call_ffi_kernel();
-
-/*!
- * \brief Encode a tiled tensor map at invocation time.
- *
- * TensorMapEncodeTiledAttr stores the descriptor dtype, rank and fixed options.
- * Arguments are descriptor and data pointers, global dimensions (rank), byte
- * strides (rank - 1), box dimensions (rank), then element strides (rank).
- */
-TVM_DLL const Op& tensormap_encode_tiled();
-
-/*!
- * \brief See pesudo code
- *
- * return_type tvm_call_packed(fname, TVMFFIAny* args) {
- * TVMFFIAny result;
- * (*fname)(args, args, len(args), &result);
- * return cast(return_type, result);
- * }
- */
-TVM_DLL const Op& tvm_call_cpacked();
-
-/*!
- * \brief Mark a condition to be thread invariant.
- * This means the condition must be the same for all threads.
- */
-TVM_DLL const Op& tvm_thread_invariant();
-
-/*!
- * \brief Lowered version of call packed, the space of value and
- * type codes are explicitly allocated.
- *
- * return_type tvm_call_packed_lowered(name,
- * TVMFFIAny* args_stack,
- * int begin,
- * int end) {
- * ModuleNode* env = GetCurrentEnv();
- * const ffi::Function* f = env->GetFuncFromEnv(name);
- * f->CallPacked(ffi::PackedArgs(args_stack[begin:end]),
- * ffi::Any(args_stack + end));
- * // return type can be int, float, handle.
- * return cast(return_type, load_return_from(args_stack + end))
- * }
- */
-TVM_DLL const Op& tvm_call_packed_lowered();
-
-/*!
- * \brief Lowered version of call c-packed, the space of value and
- * type codes are explicitly allocated.
- *
- * int tvm_call_packed_lowered(fname,
- * TVMFFIAny* args_stack,
- * int begin,
- * int end,
- * void* self) {
- * fname(ffi::PackedArgs(value_stack[begin:end], tcode_stack[begin:end]),
- * ffi::Any(value_stack + end, tcode_stack + end));
- * }
- */
-TVM_DLL const Op& tvm_call_cpacked_lowered();
-
-/*!
- * \brief See pseudo code
- *
- * int tvm_storage_sync(std::string storage_scope) {
- * __sync(storage_scope);
- * return 0;
- * }
- */
-TVM_DLL const Op& tvm_storage_sync();
-
-/*!
- * \brief Synchronize all workers in the current CPU parallel launch.
- *
- * Every worker must reach this operation. It must be outside partitioned
- * parallel loops, whose iteration counts can differ between workers.
- */
-TVM_DLL const Op& cpu_parallel_barrier();
-
-/*!
- * \brief Marker where a transform should replace generated kernel initialization.
- */
-TVM_DLL const Op& tvm_kernel_replace_point();
-
-/*!
- * \brief See pseudo code
- *
- * Type tvm_warp_shuffle(mask, Type value, warp_id, width, warp_size) {
- * return (value passed in by warp indicated by this_warp_id);
- * }
- *
- * Type tvm_warp_shuffle_up(mask, Type value, offset, width, warp_size) {
- * return (value passed in by warp indicated by this_warp_id - offset);
- * }
- *
- * Type tvm_warp_shuffle_down(mask, Type value, offset, width, warp_size) {
- * return (value passed in by warp indicated by this_warp_id + offset);
- * }
- *
- * unsigned tvm_warp_activemask() {
- * return (32-bit mask of currently active threads in the calling warp);
- * }
- *
- * Parameter warp_id indicates the source thread ID in a warp.
- *
- * Parameter offset indicates the relative distance to this_warp_id.
- *
- * Parameter width indicates the number of threads involved in one
- * shuffle. See CUDA document for __shfl_sync, __shfl_up_sync,
- * __shfl_down_sync, __shfl_xor_sync and __activemask.
- *
- * Parameter warp_size is the size of a warp, which helps a backend
- * to determine whether the width parameter is legal.
- *
- */
-TVM_DLL const Op& tvm_warp_shuffle();
-TVM_DLL const Op& tvm_warp_shuffle_up();
-TVM_DLL const Op& tvm_warp_shuffle_down();
-TVM_DLL const Op& tvm_warp_shuffle_xor();
-TVM_DLL const Op& tvm_warp_activemask();
-
-/*!
- * \brief Cross-thread reduction with an explicit typed combiner and identities.
- *
- * void tvm_thread_allreduce(LambdaExpr combine, Expr identity, Expr values,
- * PrimExpr predicate, Expr destinations, Expr thread_axes);
- *
- * For N values, combine binds lhs[0:N] followed by rhs[0:N] and returns an
- * N-element Tuple, or a scalar when N is one. Identity, values, destinations
- * and thread_axes may each be a scalar or an explicit Tuple of fields.
- * Each value, identity, pair of parameters and result have
- * the same primitive type. Inactive inputs are replaced by their identities.
- * Destinations are N tensor loads at index zero of one-element result temporaries
- * (optionally cast for boolean storage). Each result temporary must be accessed
- * only at index zero. Thread axes are reduction thread variables or zero for
- * simplified unit axes.
- * Other thread indices remain fixed. The operation writes the reduced values
- * to the destination tensors and returns void.
- */
-TVM_DLL const Op& tvm_thread_allreduce();
-
-/*!
- * \brief View a scalar all-reduce operand/result as one field, or expose its Tuple fields.
- * \param value The scalar expression or explicit Tuple.
- * \return The fields without changing the expression's representation in IR.
- */
-inline ffi::Array<Expr> GetAllreduceFields(const Expr& value) {
- if (const auto* tuple = value.as<tvm::TupleNode>()) return tuple->fields;
- return {value};
-}
-
-// Metal cooperative_tensor intrinsics (MetalPerformancePrimitives / Metal 4)
-
-/*!
- * \brief Fill a cooperative_tensor with a given value.
- *
- * void cooperative_tensor_fill(Var d, PrimExpr index, PrimExpr value,
- * int rows, int cols);
- */
-TVM_DLL const Op& cooperative_tensor_fill();
-
-/*!
- * \brief Load data from device or threadgroup memory into a cooperative_tensor.
- *
- * void cooperative_tensor_load(Var d, PrimExpr index, PrimExpr ptr,
- * PrimExpr stride, int rows, int cols,
- * bool transpose_matrix,
- * int mma_M, int mma_N, int mma_K,
- * int operand_role);
- * operand_role: 0=left(A), 1=right(B), 2=destination(C)
- */
-TVM_DLL const Op& cooperative_tensor_load();
-
-/*!
- * \brief Store data from a cooperative_tensor to device or threadgroup memory.
- *
- * void cooperative_tensor_store(Var d, PrimExpr index, PrimExpr ptr,
- * PrimExpr stride, int rows, int cols,
- * bool transpose_matrix,
- * int mma_M, int mma_N, int mma_K,
- * int operand_role);
- * operand_role: 0=left(A), 1=right(B), 2=destination(C)
- */
-TVM_DLL const Op& cooperative_tensor_store();
-
-/*!
- * \brief Multiply and accumulate two matrices using cooperative_tensor
- * (MetalPerformancePrimitives matmul2d).
- *
- * void cooperative_tensor_multiply_accumulate(
- * Var d, PrimExpr index_d, Var a, PrimExpr index_a,
- * Var b, PrimExpr index_b, Var c, PrimExpr index_c,
- * int M, int N, int K, bool transpose_a, bool transpose_b);
- */
-TVM_DLL const Op& cooperative_tensor_multiply_accumulate();
-
-// TODO(tvm-team) replace the usage of the vector operations by Shuffle.
-/*!
- * \brief Get the high level half of the vector
- */
-TVM_DLL const Op& vectorhigh();
-
-/*!
- * \brief Get the low-level half of the vector
- */
-TVM_DLL const Op& vectorlow();
-
-/*!
- * \brief Concat two vectors.
- */
-TVM_DLL const Op& vectorcombine();
-
-/*!
- * \brief Dot product of two int8x4 vectors and add an optional accumulator
- */
-TVM_DLL const Op& dp4a();
-
-/*!
- * \brief atomic add instruction, corresponding e.g. to atomicAdd in CUDA
- */
-TVM_DLL const Op& atomic_add();
-/*!
- * \brief Create an Nd memory allocation with storage scope
- */
-TVM_DLL const Op& nd_mem_alloc_with_scope();
-
-/*!
- * \brief Store to texture 2d memory
- */
-TVM_DLL const Op& texture2d_store();
-
-/*!
- * \brief Load from texture 2d memory
- */
-TVM_DLL const Op& texture2d_load();
-
-/*!
- * \brief Provide a true statement that can be used for simplifications
- *
- * Compile-time representation of known constraints about function
- * inputs. This assumption is removed when lowering, and does not
- * occur in codegen.
- */
-TVM_DLL const Op& assume();
-
-/*!
- * \brief Assume a tensor's base address has the given constant byte alignment.
- *
- * This leaf operation carries a compiler fact, without checking or changing
- * the address. Alignment must be a power of two between 1 and 2^27 bytes.
- */
-TVM_DLL const Op& assume_aligned();
-
-/*!
- * \brief Returns an initialized but arbitrary value
- *
- * Compile-time representation of memory locations whose values may be
- * altered as a result of optimizations.
- */
-TVM_DLL const Op& undef();
-
-/*!
- * \brief Calculate a predicate mask given an upper bound (limit) and a current value (base).
- *
- * It will be lowered to the llvm.get.active.lane.mask intrinsic.
- * (https://llvm.org/docs/LangRef.html#llvm-get-active-lane-mask-intrinsics)
- */
-TVM_DLL const Op& get_active_lane_mask();
-
-/*!
- * \brief Masked buffer load.
- *
- * Arguments are the buffer variable, one or more indices, and a trailing boolean lane mask.
- * The result type is the vector type loaded from the selected lanes.
- */
-TVM_DLL const Op& masked_load();
-
-/*!
- * \brief Masked buffer store.
- *
- * Arguments are the buffer variable, value, one or more indices, and a trailing boolean lane
- * mask. The result type is void.
- */
-TVM_DLL const Op& masked_store();
-
-/*! \brief Annotate a predicate not be considered as target condition of loop partition. */
-TVM_DLL const Op& ignore_loop_partition();
-/*!
- * \brief Get the element offset of a buffer given logical indices.
-
- The offset is determined by the layout of the buffer.
- */
-TVM_DLL const Op& buffer_offset();
-
-/*!
- * \brief Project the physical pointer associated with a TensorVar definition.
- *
- * The result pointer type is derived from the TensorType dtype and storage
- * scope of the sole TensorVar argument. This operation is consumed by TIRx
- * lowering and code generation.
- */
-TVM_DLL const Op& buffer_data();
-
-/*! \brief The kind of structure field info used in intrinsic */
-enum TVMStructFieldKind : int {
- // DLTensor fields
- kDLTensorAddr,
- kDLTensorData,
- kDLTensorShape,
- kDLTensorStrides,
- kDLTensorNDim,
- kDLTensorTypeCode,
- kDLTensorTypeBits,
- kDLTensorTypeLanes,
- kDLTensorByteOffset,
- kDLTensorDeviceId,
- kDLTensorDeviceType,
- kDLTensorKindBound_,
- // TVMValue field
- kTVMValueContent,
- kTVMFFIAnyTypeIndex,
- kTVMFFIAnyZeroPadding,
- kTVMFFIAnyUnionValue,
- kTVMValueKindBound_,
- // Generic int64 array element access: ((int64_t*)buf)[index]
- kInt64ArrayElem,
-};
-
-/*!
- * \brief Print the content of a buffer during runtime.
- */
-TVM_DLL const Op& print_buffer();
-/*! \} */
-} // namespace builtin
-} // namespace tirx
-} // namespace tvm
-#endif // TVM_TIR_BUILTIN_H_
diff --git a/include/tvm/tirx/op.h b/include/tvm/tirx/op.h
index 8536bf5a70..c88bd88076 100644
--- a/include/tvm/tirx/op.h
+++ b/include/tvm/tirx/op.h
@@ -19,7 +19,7 @@
/*!
* \file tvm/tirx/op.h
- * \brief Common operators defined for Expr.
+ * \brief TIRx operator attributes, builtin Ops and expression helpers.
*
* \note Most of the operator defined here perform simple constant folding
* when the type is int32 or int64 for simplifying the index expressions.
@@ -28,13 +28,13 @@
#ifndef TVM_TIRX_OP_H_
#define TVM_TIRX_OP_H_
+#include <tvm/ffi/container/array.h>
+#include <tvm/ir/attrs.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/op.h>
-#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/prim/op.h>
#include <tvm/ir/type.h>
-#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op_attr_types.h>
#include <tvm/tirx/stmt.h>
@@ -43,6 +43,753 @@
#include <type_traits>
#include <utility>
+namespace tvm {
+namespace tirx {
+
+/*!
+ * \name Call operations
+ * \brief Operations invoked through Call with explicit operands and result types.
+ * \{
+ */
+/*!
+ * \brief Allocate a buffer: alloc_tensor(shape, dtype, scope) -> TensorType.
+ *
+ * Arguments, in order:
+ * - args[0]: shape, Tuple of integer extents (IntImm or symbolic integer expressions).
+ * - args[1]: dtype, DataTypeImm with a DLDataType payload for the element type.
+ * - args[2]: scope, StringImm naming the storage scope.
+ *
+ * DictAttrs directly holds the allocation annotations, defaulting to an empty dictionary.
+ * The TensorType result agrees with the operands and retains buffer access/storage metadata.
+ *
+ * \code
+ * // Example pattern match code for a given Binding:
+ * if (const auto* call = binding->value.as<CallNode>();
+ * call && call->op.same_as(tirx::alloc_tensor_op())) {
+ * tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
+ * DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
+ * ffi::String scope = call->args[2].as_or_throw<StringImm>()->value;
+ * DictAttrs annotations = call->attrs.as_or_throw<DictAttrs>();
+ * }
+ * \endcode
+ */
+TVM_DLL const Op& alloc_tensor_op();
+/*!
+ * \brief Declare a buffer view: decl_tensor(data, shape, dtype, scope) -> TensorType.
+ *
+ * Arguments, in order:
+ * - args[0]: data, Expr for the existing physical pointer backing the buffer view.
+ * - args[1]: shape, Tuple of integer extents (IntImm or symbolic integer expressions).
+ * - args[2]: dtype, DataTypeImm with a DLDataType payload for the element type.
+ * - args[3]: scope, StringImm naming the storage scope.
+ *
+ * There are no attributes. The TensorType result agrees with the operands and retains
+ * buffer access/storage metadata. The operation binds a view without allocating memory.
+ *
+ * \code
+ * // Example pattern match code for a given Binding:
+ * if (const auto* call = binding->value.as<CallNode>();
+ * call && call->op.same_as(tirx::decl_tensor_op())) {
+ * Expr data = call->args[0];
+ * tvm::Tuple shape = call->args[1].as_or_throw<tvm::Tuple>();
+ * DLDataType dtype = call->args[2].as_or_throw<DataTypeImm>()->value;
+ * ffi::String scope = call->args[3].as_or_throw<StringImm>()->value;
+ * }
+ * \endcode
+ */
+TVM_DLL const Op& decl_tensor_op();
+/*!
+ * \brief Return from a GPU thread without returning a function value.
+ */
+TVM_DLL const Op& thread_return_op();
+/*!
+ * \brief Reinterpret the value using the target type.
+ */
+TVM_DLL const Op& reinterpret_op();
+
+/*!
+ * \brief Thread-set filter predicate. Used as the condition of an IfThenElse
+ * to narrow the active thread set A for the then-branch. Two forms:
+ * filter(var, lo, hi) -- range form, true iff var in [lo, hi)
+ * filter(var, cond) -- predicate form (e.g. var == k); true iff cond
+ * `var` must be a ScopeIdDef-declared Var at parse time (Verifier Rule 2).
+ */
+TVM_DLL const Op& filter_op();
+
+/*!
+ * \brief Analysis-only active-thread selector.
+ *
+ * ``selector(var, pred)`` denotes the unique value of ``var`` in the current
+ * active domain for which ``pred`` is true. It is used only inside
+ * ExecContext/DispatchContext metadata, for predicates such as
+ * ``ptx.elect_sync()`` whose selected lane cannot be inferred structurally.
+ */
+TVM_DLL const Op& selector_op();
+
+/*!
+ * \brief Execute a multiplication between two Q-numbers x and y
+ * followed by a right shift s
+ * The default rounding rule is to the nearest value, rounding half up
+ * (i.e., round(x.1) = x and round (x.5) = x+1)
+ */
+TVM_DLL const Op& q_multiply_shift_op();
+TVM_DLL const Op& q_multiply_shift_per_axis_op();
+
+/*!
+ * \brief Returns the address of an element in the buffer (see pseudocode below).
+ *
+ * The number of indices should match the dimensionality of the buffer
+ * being accessed. If this operation occurs after buffer flattening,
+ * the number of indices must be supported by the target (i.e. N>1
+ * only on targets that support non-flat memory buffers).
+ *
+ * Handle address_of(TensorLoad *op) {
+ * return &op->buffer_var[op->indices[0], op->indices[1], ..., op->indices[N-1]];
+ * }
+ */
+TVM_DLL const Op& address_of_op();
+
+/*!
+ * \brief See pesudo code
+ *
+ * bool isnullptr(void* handle) {
+ * return handle == nullptr
+ * }
+ */
+TVM_DLL const Op& isnullptr_op();
+
+/*!
+ * \brief Check if value is nan
+ */
+TVM_DLL const Op& isnan_op();
+
+/*!
+ * \brief Popcount
+ */
+TVM_DLL const Op& popcount_op();
+
+/*!
+ * \brief Fused multiply add
+ *
+ * Type fma(a, b, c) {
+ * return a * b + c;
+ * }
+ */
+TVM_DLL const Op& fma_op();
+
+/*!
+ * \brief Call an extern C function with given name
+ * and signature from the types of args in the runtime environment.
+ *
+ * Type call_extern(name, args...) {
+ * return dlsym(name)(args...);
+ * }
+ *
+ * \note This intrinsic does not provide any type checking,
+ * and is main used for backward compatibility reasons.
+ * Always consider use pre-registered and typed tvm::Op first.
+ */
+TVM_DLL const Op& call_extern_op();
+
+/*!
+ * \brief Call an pure extern C function with given name
+ * and signature from the types of args in the runtime environment.
+ *
+ * Type call_pure_extern(name, args...) {
+ * return dlsym(name)(args...);
+ * }
+ *
+ * \note This intrinsic does not provide any type checking,
+ * and is main used for backward compatibility reasons.
+ * Always consider use pre-registered and typed tvm::Op first.
+ */
+TVM_DLL const Op& call_pure_extern_op();
+
+/*!
+ * \brief Call an LLVM intrinsic with a given intrinsic id
+ * and signature from the types of args in the runtime environment.
+ *
+ * Type call_llvm_pure_intrin(intrin_id, args...) {
+ * return dlsym(name)(args...);
+ * }
+ *
+ * \note This op does not provide any type checking.
+ */
+TVM_DLL const Op& call_llvm_intrin_op();
+
+/*!
+ * \brief Call an LLVM pure intrinsic with a given intrinsic id
+ * and signature from the types of args in the runtime environment.
+ *
+ * Type call_llvm_pure_intrin(intrin_id, args...) {
+ * return dlsym(name)(args...);
+ * }
+ *
+ * \note This op does not provide any type checking.
+ */
+TVM_DLL const Op& call_llvm_pure_intrin_op();
+
+/*!
+ * \brief Call an SPIRV pure GLSL450 intrinsic.
+ *
+ * Type call_spirv_pure_glsl450(intrin_id, args...) {
+ * return dlsym(name)(args...);
+ * }
+ *
+ * \note This op does not provide any type checking.
+ */
+TVM_DLL const Op& call_spirv_pure_glsl450_op();
+
+// TODO(tvm-team) revisit the builtins below
+// some of them can simply become ops with special codegen attr.
+/*!
+ * \brief same signature as llvm.prefetch
+ */
+TVM_DLL const Op& prefetch_op();
+
+/*!
+ * \brief Get head access address with memory access pattern info.
+ *
+ * This operator also marks range of the memory access
+ * The offset and extent are in unit of the DType(including vectorization factor).
+ * rw_mask is a bit_mask setting whether the access is a read(1) or write(2).
+ * The access is assume to happen in the current expression.
+ *
+ * PtrType tvm_access_ptr(Expr dtype, DType* data,
+ * int offset, int extent,
+ * int rw_mask) {
+ * // DType == dtype.type();
+ * return &data[offset];
+ * }
+ */
+TVM_DLL const Op& tvm_access_ptr_op();
+
+/*!
+ * \brief Cast a handle to a typed pointer after adding a byte offset.
+ *
+ * DType* ptr_byte_offset(void* data, int byte_offset, Expr dtype) {
+ * return reinterpret_cast<DType*>(reinterpret_cast<char*>(data) + byte_offset);
+ * }
+ */
+TVM_DLL const Op& ptr_byte_offset_op();
+
+/*!
+ * \brief Create a function local static handle that iniitalizes to nullptr.
+ * can be used to cache function local static resources.
+ */
+TVM_DLL const Op& tvm_static_handle_op();
+
+/*!
+ * \brief See pesudo code
+ *
+ * void* handle_add_byte_offset(void* handle, int offset) {
+ * return reinterpret_cast<v*>(reinterpret_cast<char*>(handle) + offset);
+ * }
+ */
+TVM_DLL const Op& handle_add_byte_offset_op();
+
+/*!
+ * \brief See pesudo code
+ *
+ * Type tvm_struct_get(StructType* arr, int index, int field_id) {
+ * return arr[index]->field;
+ * }
+ * \sa TVMStructFieldKind
+ */
+TVM_DLL const Op& tvm_struct_get_op();
+
+/*!
+ * \brief See pesudo code
+ *
+ * Handle tvm_struct_set(StructType* arr, int index, int field_id, value) {
+ * arr[index]->field = value;
+ * }
+ * \sa TVMStructFieldKind
+ */
+TVM_DLL const Op& tvm_struct_set_op();
+
+/*!
+ * \brief See pesudo code
+ *
+ * void tvm_throw_last_error() {
+ * throw TVMGetLastError();
+ * }
+ */
+TVM_DLL const Op& tvm_throw_last_error_op();
+
+/*!
+ * \brief See pesudo code
+ *
+ * dtype in {shape, array, arg_value, arg_tcode}
+ *
+ * Handle tvm_stack_alloca(string dtype, int num) {
+ * return new on stack dtype[num];
+ * }
+ */
+TVM_DLL const Op& tvm_stack_alloca_op();
+
+/*!
+ * \brief Allocate a shape tuple on stack, return the handle.
+ *
+ * Handle tvm_stack_make_shape(list args) {
+ * ret = alloca stack int64_t[len(args)];
+ * for i in range(len(args)):
+ * ret[i] = args[i]
+ * return &ret[0];
+ * }
+ */
+TVM_DLL const Op& tvm_stack_make_shape_op();
+
+/*!
+ * \brief Allocate a Tensor(DLTensor) on stack, return the handle.
+ *
+ * Type tvm_stack_make_array(Expr data,
+ * Expr shape,
+ * Expr strides,
+ * Expr ndim,
+ * Expr dtype,
+ * Expr elem_offset) {
+ * ret = alloca stack DLTensor();
+ * ret->data = data;
+ * ret->shape = shape;
+ * ret->strides = strides != 0 ? strides : nullptr;
+ * ret->ndim = ndim;
+ * ret->dtype = dtype.type();
+ * ret->byte_offset = elem_offset * sizeof(dtype);
+ * return ret;
+ * }
+ */
+TVM_DLL const Op& tvm_stack_make_array_op();
+
+/*!
+ * \brief See pesudo code
+ *
+ * return_type tvm_call_packed(name, TVMFFIAny* args) {
+ * TVMFFIAny result;
+ * ModuleNode* env = GetCurrentEnv();
+ * const ffi::Function* f = env->GetFuncFromEnv(name);
+ * (*f)(args, args, len(args), &result);
+ * // return type can be int, float, handle.
+ * return cast(return_type, result);
+ * }
+ */
+TVM_DLL const Op& tvm_call_packed_op();
+
+/*! \brief Launch metadata for call_ffi_kernel. */
+struct CallFFIKernelAttr : public AttrsNode {
+ /*!
+ * \brief Ordered launch tags describing the suffix of the call arguments.
+ *
+ * The first call argument is the kernel symbol, followed by kernel operands
+ * and launch values. Flag-only tags consume no value; dynamic shared-memory
+ * bytes, when present, are last. Runtime expressions remain in Call.args.
+ */
+ ffi::Array<ffi::String> launch_params;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.CallFFIKernelAttr", CallFFIKernelAttr, AttrsNode);
+};
+
+/*!
+ * \brief Launch a kernel using the packed-function argument convention.
+ *
+ * Arguments are the kernel symbol, kernel operands, then launch values.
+ * CallFFIKernelAttr::launch_params describes the launch-value suffix.
+ * Host backends may consume this call directly or lower it to tvm_call_packed.
+ */
+TVM_DLL const Op& call_ffi_kernel_op();
+
+/*! \brief Fixed encoding options for tensormap_encode_tiled. */
+struct TensorMapEncodeTiledAttr : public AttrsNode {
+ DLDataType descriptor_dtype;
+ int64_t rank;
+ int64_t interleave;
+ int64_t swizzle;
+ int64_t l2_promotion;
+ int64_t oob_fill;
+ int64_t force_cu_dtype;
+
+ static void RegisterReflection();
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.TensorMapEncodeTiledAttr", TensorMapEncodeTiledAttr,
+ AttrsNode);
+};
+
+/*!
+ * \brief Encode a tiled tensor map at invocation time.
+ *
+ * TensorMapEncodeTiledAttr stores the descriptor dtype, rank and fixed options.
+ * Arguments are descriptor and data pointers, global dimensions (rank), byte
+ * strides (rank - 1), box dimensions (rank), then element strides (rank).
+ */
+TVM_DLL const Op& tensormap_encode_tiled_op();
+
+/*!
+ * \brief See pesudo code
+ *
+ * return_type tvm_call_packed(fname, TVMFFIAny* args) {
+ * TVMFFIAny result;
+ * (*fname)(args, args, len(args), &result);
+ * return cast(return_type, result);
+ * }
+ */
+TVM_DLL const Op& tvm_call_cpacked_op();
+
+/*!
+ * \brief Mark a condition to be thread invariant.
+ * This means the condition must be the same for all threads.
+ */
+TVM_DLL const Op& tvm_thread_invariant_op();
+
+/*!
+ * \brief Lowered version of call packed, the space of value and
+ * type codes are explicitly allocated.
+ *
+ * return_type tvm_call_packed_lowered(name,
+ * TVMFFIAny* args_stack,
+ * int begin,
+ * int end) {
+ * ModuleNode* env = GetCurrentEnv();
+ * const ffi::Function* f = env->GetFuncFromEnv(name);
+ * f->CallPacked(ffi::PackedArgs(args_stack[begin:end]),
+ * ffi::Any(args_stack + end));
+ * // return type can be int, float, handle.
+ * return cast(return_type, load_return_from(args_stack + end))
+ * }
+ */
+TVM_DLL const Op& tvm_call_packed_lowered_op();
+
+/*!
+ * \brief Lowered version of call c-packed, the space of value and
+ * type codes are explicitly allocated.
+ *
+ * int tvm_call_packed_lowered(fname,
+ * TVMFFIAny* args_stack,
+ * int begin,
+ * int end,
+ * void* self) {
+ * fname(ffi::PackedArgs(value_stack[begin:end], tcode_stack[begin:end]),
+ * ffi::Any(value_stack + end, tcode_stack + end));
+ * }
+ */
+TVM_DLL const Op& tvm_call_cpacked_lowered_op();
+
+/*!
+ * \brief See pseudo code
+ *
+ * int tvm_storage_sync(std::string storage_scope) {
+ * __sync(storage_scope);
+ * return 0;
+ * }
+ */
+TVM_DLL const Op& tvm_storage_sync_op();
+
+/*!
+ * \brief Synchronize all workers in the current CPU parallel launch.
+ *
+ * Every worker must reach this operation. It must be outside partitioned
+ * parallel loops, whose iteration counts can differ between workers.
+ */
+TVM_DLL const Op& cpu_parallel_barrier_op();
+
+/*!
+ * \brief Marker where a transform should replace generated kernel initialization.
+ */
+TVM_DLL const Op& tvm_kernel_replace_point_op();
+
+/*!
+ * \brief See pseudo code
+ *
+ * Type tvm_warp_shuffle(mask, Type value, warp_id, width, warp_size) {
+ * return (value passed in by warp indicated by this_warp_id);
+ * }
+ *
+ * Type tvm_warp_shuffle_up(mask, Type value, offset, width, warp_size) {
+ * return (value passed in by warp indicated by this_warp_id - offset);
+ * }
+ *
+ * Type tvm_warp_shuffle_down(mask, Type value, offset, width, warp_size) {
+ * return (value passed in by warp indicated by this_warp_id + offset);
+ * }
+ *
+ * unsigned tvm_warp_activemask() {
+ * return (32-bit mask of currently active threads in the calling warp);
+ * }
+ *
+ * Parameter warp_id indicates the source thread ID in a warp.
+ *
+ * Parameter offset indicates the relative distance to this_warp_id.
+ *
+ * Parameter width indicates the number of threads involved in one
+ * shuffle. See CUDA document for __shfl_sync, __shfl_up_sync,
+ * __shfl_down_sync, __shfl_xor_sync and __activemask.
+ *
+ * Parameter warp_size is the size of a warp, which helps a backend
+ * to determine whether the width parameter is legal.
+ *
+ */
+TVM_DLL const Op& tvm_warp_shuffle_op();
+TVM_DLL const Op& tvm_warp_shuffle_up_op();
+TVM_DLL const Op& tvm_warp_shuffle_down_op();
+TVM_DLL const Op& tvm_warp_shuffle_xor_op();
+TVM_DLL const Op& tvm_warp_activemask_op();
+
+/*!
+ * \brief Cross-thread reduction with an explicit typed combiner and identities.
+ *
+ * void tvm_thread_allreduce(LambdaExpr combine, Expr identity, Expr values,
+ * PrimExpr predicate, Expr destinations, Expr thread_axes);
+ *
+ * For N values, combine binds lhs[0:N] followed by rhs[0:N] and returns an
+ * N-element Tuple, or a scalar when N is one. Identity, values, destinations
+ * and thread_axes may each be a scalar or an explicit Tuple of fields.
+ * Each value, identity, pair of parameters and result have
+ * the same primitive type. Inactive inputs are replaced by their identities.
+ * Destinations are N tensor loads at index zero of one-element result temporaries
+ * (optionally cast for boolean storage). Each result temporary must be accessed
+ * only at index zero. Thread axes are reduction thread variables or zero for
+ * simplified unit axes.
+ * Other thread indices remain fixed. The operation writes the reduced values
+ * to the destination tensors and returns void.
+ */
+TVM_DLL const Op& tvm_thread_allreduce_op();
+
+/*!
+ * \brief View a scalar all-reduce operand/result as one field, or expose its Tuple fields.
+ * \param value The scalar expression or explicit Tuple.
+ * \return The fields without changing the expression's representation in IR.
+ */
+inline ffi::Array<Expr> GetAllreduceFields(const Expr& value) {
+ if (const auto* tuple = value.as<tvm::TupleNode>()) return tuple->fields;
+ return {value};
+}
+
+// Metal cooperative_tensor intrinsics (MetalPerformancePrimitives / Metal 4)
+
+/*!
+ * \brief Fill a cooperative_tensor with a given value.
+ *
+ * void cooperative_tensor_fill(Var d, PrimExpr index, PrimExpr value,
+ * int rows, int cols);
+ */
+TVM_DLL const Op& cooperative_tensor_fill_op();
+
+/*!
+ * \brief Load data from device or threadgroup memory into a cooperative_tensor.
+ *
+ * void cooperative_tensor_load(Var d, PrimExpr index, PrimExpr ptr,
+ * PrimExpr stride, int rows, int cols,
+ * bool transpose_matrix,
+ * int mma_M, int mma_N, int mma_K,
+ * int operand_role);
+ * operand_role: 0=left(A), 1=right(B), 2=destination(C)
+ */
+TVM_DLL const Op& cooperative_tensor_load_op();
+
+/*!
+ * \brief Store data from a cooperative_tensor to device or threadgroup memory.
+ *
+ * void cooperative_tensor_store(Var d, PrimExpr index, PrimExpr ptr,
+ * PrimExpr stride, int rows, int cols,
+ * bool transpose_matrix,
+ * int mma_M, int mma_N, int mma_K,
+ * int operand_role);
+ * operand_role: 0=left(A), 1=right(B), 2=destination(C)
+ */
+TVM_DLL const Op& cooperative_tensor_store_op();
+
+/*!
+ * \brief Multiply and accumulate two matrices using cooperative_tensor
+ * (MetalPerformancePrimitives matmul2d).
+ *
+ * void cooperative_tensor_multiply_accumulate(
+ * Var d, PrimExpr index_d, Var a, PrimExpr index_a,
+ * Var b, PrimExpr index_b, Var c, PrimExpr index_c,
+ * int M, int N, int K, bool transpose_a, bool transpose_b);
+ */
+TVM_DLL const Op& cooperative_tensor_multiply_accumulate_op();
+
+// TODO(tvm-team) replace the usage of the vector operations by Shuffle.
+/*!
+ * \brief Get the high level half of the vector
+ */
+TVM_DLL const Op& vectorhigh_op();
+
+/*!
+ * \brief Get the low-level half of the vector
+ */
+TVM_DLL const Op& vectorlow_op();
+
+/*!
+ * \brief Concat two vectors.
+ */
+TVM_DLL const Op& vectorcombine_op();
+
+/*!
+ * \brief Dot product of two int8x4 vectors and add an optional accumulator
+ */
+TVM_DLL const Op& dp4a_op();
+
+/*!
+ * \brief atomic add instruction, corresponding e.g. to atomicAdd in CUDA
+ */
+TVM_DLL const Op& atomic_add_op();
+/*!
+ * \brief Create an Nd memory allocation with storage scope
+ */
+TVM_DLL const Op& nd_mem_alloc_with_scope_op();
+
+/*!
+ * \brief Store to texture 2d memory
+ */
+TVM_DLL const Op& texture2d_store_op();
+
+/*!
+ * \brief Load from texture 2d memory
+ */
+TVM_DLL const Op& texture2d_load_op();
+
+/*!
+ * \brief Provide a true statement that can be used for simplifications
+ *
+ * Compile-time representation of known constraints about function
+ * inputs. This assumption is removed when lowering, and does not
+ * occur in codegen.
+ */
+TVM_DLL const Op& assume_op();
+
+/*!
+ * \brief Assume a tensor's base address has the given constant byte alignment.
+ *
+ * This leaf operation carries a compiler fact, without checking or changing
+ * the address. Alignment must be a power of two between 1 and 2^27 bytes.
+ */
+TVM_DLL const Op& assume_aligned_op();
+
+/*!
+ * \brief Returns an initialized but arbitrary value
+ *
+ * Compile-time representation of memory locations whose values may be
+ * altered as a result of optimizations.
+ */
+TVM_DLL const Op& undef_op();
+
+/*!
+ * \brief Calculate a predicate mask given an upper bound (limit) and a current value (base).
+ *
+ * It will be lowered to the llvm.get.active.lane.mask intrinsic.
+ * (https://llvm.org/docs/LangRef.html#llvm-get-active-lane-mask-intrinsics)
+ */
+TVM_DLL const Op& get_active_lane_mask_op();
+
+/*!
+ * \brief Masked buffer load.
+ *
+ * Arguments are the buffer variable, one or more indices, and a trailing boolean lane mask.
+ * The result type is the vector type loaded from the selected lanes.
+ */
+TVM_DLL const Op& masked_load_op();
+
+/*!
+ * \brief Masked buffer store.
+ *
+ * Arguments are the buffer variable, value, one or more indices, and a trailing boolean lane
+ * mask. The result type is void.
+ */
+TVM_DLL const Op& masked_store_op();
+
+/*! \brief Annotate a predicate not be considered as target condition of loop partition. */
+TVM_DLL const Op& ignore_loop_partition_op();
+/*!
+ * \brief Get the element offset of a buffer given logical indices.
+
+ The offset is determined by the layout of the buffer.
+ */
+TVM_DLL const Op& buffer_offset_op();
+
+/*!
+ * \brief Project the physical pointer associated with a TensorVar definition.
+ *
+ * The result pointer type is derived from the TensorType dtype and storage
+ * scope of the sole TensorVar argument. This operation is consumed by TIRx
+ * lowering and code generation.
+ */
+TVM_DLL const Op& buffer_data_op();
+
+/*! \brief The kind of structure field info used in intrinsic */
+enum TVMStructFieldKind : int {
+ // DLTensor fields
+ kDLTensorAddr,
+ kDLTensorData,
+ kDLTensorShape,
+ kDLTensorStrides,
+ kDLTensorNDim,
+ kDLTensorTypeCode,
+ kDLTensorTypeBits,
+ kDLTensorTypeLanes,
+ kDLTensorByteOffset,
+ kDLTensorDeviceId,
+ kDLTensorDeviceType,
+ kDLTensorKindBound_,
+ // TVMValue field
+ kTVMValueContent,
+ kTVMFFIAnyTypeIndex,
+ kTVMFFIAnyZeroPadding,
+ kTVMFFIAnyUnionValue,
+ kTVMValueKindBound_,
+ // Generic int64 array element access: ((int64_t*)buf)[index]
+ kInt64ArrayElem,
+};
+
+/*!
+ * \brief Print the content of a buffer during runtime.
+ */
+TVM_DLL const Op& print_buffer_op();
+/*! \} */
+
+/*!
+ * \name Region operations
+ * \brief Operations used by RegionStmt to enclose a lexical body.
+ *
+ * Every region operation registers FRegionGetBodyParams, returning fresh typed
+ * body parameters or an empty array for no parameters. Attribute presence alone
+ * identifies region support; classification does not invoke the hook. Operands
+ * and attributes are evaluated outside the body-parameter scope.
+ * \{
+ */
+/*!
+ * \brief Thread launch region: operands are a nonempty StringImm tag and a
+ * scalar signed or unsigned integer extent wider than one bit.
+ * The sole body parameter is a fresh thread-index PrimVar matching the extent type.
+ * Tags starting with vthread denote virtual threads. There are no attrs or results.
+ */
+TVM_DLL const Op& launch_thread_op();
+
+/*!
+ * \brief Mark a user-facing device entry containing device scope definitions.
+ * Takes no operands, body parameters, attributes or results.
+ */
+TVM_DLL const Op& device_entry_op();
+/*!
+ * \brief Supply lexical device context for allocation and packed-call lowering.
+ * Operands are integer device type and device ID; there are no body parameters,
+ * attributes or results. The region does not change the active runtime device.
+ */
+TVM_DLL const Op& device_context_op();
+/*!
+ * \brief Outline the body as a CPU compute helper named by a StringImm operand.
+ * There are no body parameters, attributes or results.
+ */
+TVM_DLL const Op& compute_scope_op();
+/*!
+ * \brief Launch a CPU worker team around parallel loops and team barriers.
+ * Takes no operands, body parameters, attributes or results.
+ */
+TVM_DLL const Op& parallel_launch_op();
+/*! \} */
+
+} // namespace tirx
+} // namespace tvm
+
namespace tvm::prim {
// Shared primitive construction and constants are declared in ir/prim/op.h.
diff --git a/include/tvm/tirx/script/ir_builder/ir.h b/include/tvm/tirx/script/ir_builder/ir.h
index fcea336171..4a2e2134d7 100644
--- a/include/tvm/tirx/script/ir_builder/ir.h
+++ b/include/tvm/tirx/script/ir_builder/ir.h
@@ -27,7 +27,7 @@
#include <tvm/tirx/layout.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/script/ir_builder/frame.h>
-#include <tvm/tirx/tile_primitive.h>
+#include <tvm/tirx/tile_op.h>
#include <tvm/tirx/type.h>
namespace tvm {
diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h
index 2a3b679ce2..08dc71bd11 100644
--- a/include/tvm/tirx/stmt_functor.h
+++ b/include/tvm/tirx/stmt_functor.h
@@ -32,7 +32,7 @@
#include <tvm/tirx/expr_functor.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/tile_primitive.h>
+#include <tvm/tirx/tile_op.h>
#include <unordered_map>
#include <utility>
diff --git a/include/tvm/tirx/tile_primitive.h b/include/tvm/tirx/tile_op.h
similarity index 87%
rename from include/tvm/tirx/tile_primitive.h
rename to include/tvm/tirx/tile_op.h
index 0b10ff13b1..67b8acf929 100644
--- a/include/tvm/tirx/tile_primitive.h
+++ b/include/tvm/tirx/tile_op.h
@@ -17,11 +17,11 @@
* under the License.
*/
/*!
- * \file tvm/tirx/tile_primitive.h
+ * \file tvm/tirx/tile_op.h
* \brief TIRX tile primitive statements, operators, and reified lambda expressions.
*/
-#ifndef TVM_TIRX_TILE_PRIMITIVE_H_
-#define TVM_TIRX_TILE_PRIMITIVE_H_
+#ifndef TVM_TIRX_TILE_OP_H_
+#define TVM_TIRX_TILE_OP_H_
#include <tvm/ffi/object.h>
#include <tvm/ir/op.h>
@@ -209,40 +209,43 @@ class TilePrimitiveCall : public Stmt {
TVM_DEFINE_OBJECT_REF_COW_METHOD(TilePrimitiveCallNode);
};
+/*! \brief Tile primitive operator handles. */
+namespace tile {
+
/*!
* \brief See pesudo code below:
*
* Tx.cast(TensorRegion dst, TensorRegion src)
*/
-TVM_DLL const Op& cast();
+TVM_DLL const Op& cast_op();
/*!
* \brief See pesudo code below:
*
* Tx.copy(TensorRegion dst, TensorRegion src)
*/
-TVM_DLL const Op& copy();
+TVM_DLL const Op& copy_op();
/*!
* \brief See pesudo code below:
*
* Tx.Async.copy(TensorRegion dst, TensorRegion src)
*/
-TVM_DLL const Op& copy_async();
+TVM_DLL const Op& copy_async_op();
/*!
* \brief See pesudo code below:
*
* Tx.fill(TensorRegion dst, PrimExpr value)
*/
-TVM_DLL const Op& fill();
+TVM_DLL const Op& fill_op();
/*!
* \brief See pesudo code below:
*
* Tx.gemm(TensorVar A, TensorVar B, TensorVar C, TensorVar D, PrimExpr alpha, PrimExpr beta)
*/
-TVM_DLL const Op& gemm();
+TVM_DLL const Op& gemm_op();
/*!
* \brief See pesudo code below:
@@ -250,62 +253,64 @@ TVM_DLL const Op& gemm();
* Tx.gemm_async(TensorRegion C, TensorRegion A, TensorRegion B, bool transA, bool transB,
* bool accum)
*/
-TVM_DLL const Op& gemm_async();
+TVM_DLL const Op& gemm_async_op();
+
+TVM_DLL const Op& zero_op();
-TVM_DLL const Op& zero();
+TVM_DLL const Op& sqrt_op();
+TVM_DLL const Op& sqrt_with_scale_bias_op();
-TVM_DLL const Op& sqrt();
-TVM_DLL const Op& sqrt_with_scale_bias();
+TVM_DLL const Op& exp_op();
+TVM_DLL const Op& exp_with_scale_bias_op();
-TVM_DLL const Op& exp();
-TVM_DLL const Op& exp_with_scale_bias();
+TVM_DLL const Op& exp2_op();
+TVM_DLL const Op& exp2_with_scale_bias_op();
-TVM_DLL const Op& exp2();
-TVM_DLL const Op& exp2_with_scale_bias();
+TVM_DLL const Op& log2_op();
+TVM_DLL const Op& log2_with_scale_bias_op();
-TVM_DLL const Op& log2();
-TVM_DLL const Op& log2_with_scale_bias();
+TVM_DLL const Op& add_op();
-TVM_DLL const Op& add();
+TVM_DLL const Op& sub_op();
-TVM_DLL const Op& sub();
+TVM_DLL const Op& mul_op();
-TVM_DLL const Op& mul();
+TVM_DLL const Op& fdiv_op();
-TVM_DLL const Op& fdiv();
+TVM_DLL const Op& minimum_op();
-TVM_DLL const Op& minimum();
+TVM_DLL const Op& maximum_op();
-TVM_DLL const Op& maximum();
+TVM_DLL const Op& reciprocal_op();
-TVM_DLL const Op& reciprocal();
+TVM_DLL const Op& sum_op();
-TVM_DLL const Op& sum();
+TVM_DLL const Op& max_op();
-TVM_DLL const Op& max();
+TVM_DLL const Op& min_op();
-TVM_DLL const Op& min();
+TVM_DLL const Op& memset_op();
-TVM_DLL const Op& memset();
+TVM_DLL const Op& reduce_negate_op();
-TVM_DLL const Op& reduce_negate();
+TVM_DLL const Op& binary_reduce_op();
-TVM_DLL const Op& binary_reduce();
+TVM_DLL const Op& unary_reduce_op();
+TVM_DLL const Op& unary_reduce_with_scale_bias_op();
-TVM_DLL const Op& unary_reduce();
-TVM_DLL const Op& unary_reduce_with_scale_bias();
+TVM_DLL const Op& binary_chain_op();
-TVM_DLL const Op& binary_chain();
+TVM_DLL const Op& select_op();
-TVM_DLL const Op& select();
+TVM_DLL const Op& fma_op();
-TVM_DLL const Op& fma();
+TVM_DLL const Op& silu_op();
-TVM_DLL const Op& silu();
+TVM_DLL const Op& permute_layout_op();
-TVM_DLL const Op& permute_layout();
+} // namespace tile
} // namespace tirx
} // namespace tvm
-#endif // TVM_TIRX_TILE_PRIMITIVE_H_
+#endif // TVM_TIRX_TILE_OP_H_
diff --git a/include/tvm/topi/detail/extern.h b/include/tvm/topi/detail/extern.h
index a344938d64..7c147facfb 100644
--- a/include/tvm/topi/detail/extern.h
+++ b/include/tvm/topi/detail/extern.h
@@ -24,9 +24,9 @@
#ifndef TVM_TOPI_DETAIL_EXTERN_H_
#define TVM_TOPI_DETAIL_EXTERN_H_
-#include <tvm/ir/prim/builtin.h>
+#include <tvm/ir/prim/op.h>
#include <tvm/te/operation.h>
-#include <tvm/tirx/builtin.h>
+#include <tvm/tirx/op.h>
#include <string>
#include <utility>
@@ -103,11 +103,11 @@ inline Expr pack_buffer(TensorVar buf) {
using namespace tvm::prim;
TVM_FFI_ICHECK_GT(buf->shape.size(), 0) << "buf shape must have at least one element";
Expr shape =
- Call(PointerType(PrimType::Int(64)), tvm::tirx::builtin::tvm_stack_make_shape(), buf->shape);
+ Call(PointerType(PrimType::Int(64)), tvm::tirx::tvm_stack_make_shape_op(), buf->shape);
Expr strides;
if (buf->strides.size() > 0) {
- strides = Call(PointerType(PrimType::Int(64)), tvm::tirx::builtin::tvm_stack_make_shape(),
- buf->strides);
+ strides =
+ Call(PointerType(PrimType::Int(64)), tvm::tirx::tvm_stack_make_shape_op(), buf->strides);
} else {
strides = PrimExpr(0);
}
@@ -117,7 +117,7 @@ inline Expr pack_buffer(TensorVar buf) {
IntImm::Int32(static_cast<int64_t>(buf->shape.size())),
MakeConst(PrimType(buf->dtype), 0),
buf->elem_offset};
- return Call(PointerType::VoidPointerTy(), tvm::tirx::builtin::tvm_stack_make_array(), pack_args);
+ return Call(PointerType::VoidPointerTy(), tvm::tirx::tvm_stack_make_array_op(), pack_args);
}
/*!
@@ -130,8 +130,7 @@ inline Expr pack_buffer(TensorVar buf) {
* \return An expression representing the invocation
*/
inline PrimExpr call_packed(ffi::Array<Expr> args) {
- return Call(PrimType::Int(32), tvm::tirx::builtin::tvm_call_packed(), args)
- .as_or_throw<PrimExpr>();
+ return Call(PrimType::Int(32), tvm::tirx::tvm_call_packed_op(), args).as_or_throw<PrimExpr>();
}
} // namespace detail
diff --git a/include/tvm/topi/elemwise.h b/include/tvm/topi/elemwise.h
index 457fc4e3da..0cb7590dfd 100644
--- a/include/tvm/topi/elemwise.h
+++ b/include/tvm/topi/elemwise.h
@@ -24,9 +24,8 @@
#ifndef TVM_TOPI_ELEMWISE_H_
#define TVM_TOPI_ELEMWISE_H_
-#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/builtin.h>
+#include <tvm/ir/prim/op.h>
#include <tvm/tirx/op.h>
#include <tvm/topi/tags.h>
diff --git a/python/tvm/backend/cuda/op.py b/python/tvm/backend/cuda/op.py
index c437d04e07..971979eac8 100644
--- a/python/tvm/backend/cuda/op.py
+++ b/python/tvm/backend/cuda/op.py
@@ -21,8 +21,10 @@ from __future__ import annotations
from enum import Enum
+import tvm_ffi
+
from tvm import tirx
-from tvm.ir import Call, Op, StringImm
+from tvm.ir import Attrs, Call, Op, StringImm
from tvm.ir.op import _init_op_api, _make_op_api
from tvm.ir.type import PointerType, PrimType
from tvm.runtime import const
@@ -591,6 +593,16 @@ def ptx_legacy_ldmatrix(
)
+@tvm_ffi.register_object("tirx.cuda.TCGen05InstrDescriptorAttrs")
+class TCGen05InstrDescriptorAttrs(Attrs):
+ """Static options for the dense tcgen05 instruction descriptor."""
+
+
+@tvm_ffi.register_object("tirx.cuda.TCGen05InstrDescriptorBlockScaledAttrs")
+class TCGen05InstrDescriptorBlockScaledAttrs(Attrs):
+ """Static options for the block-scaled tcgen05 instruction descriptor."""
+
+
_encode_instr_descriptor = _make_op_api(
Op.get("tirx.cuda.tcgen05_encode_instr_descriptor"), __name__
)
diff --git a/python/tvm/backend/cuda/script.py b/python/tvm/backend/cuda/script.py
index cc95702241..c068dcbf95 100644
--- a/python/tvm/backend/cuda/script.py
+++ b/python/tvm/backend/cuda/script.py
@@ -102,6 +102,9 @@ class IketNamespace:
self.official_event = _cuda_op.cuda_iket_official_event
+TCGen05InstrDescriptorAttrs = _cuda_op.TCGen05InstrDescriptorAttrs
+TCGen05InstrDescriptorBlockScaledAttrs = _cuda_op.TCGen05InstrDescriptorBlockScaledAttrs
+
iket = IketNamespace()
wgmma = CudaWgmmaNamespace()
tcgen05 = CudaTcgen05Namespace()
@@ -128,6 +131,11 @@ ldg = _cuda_op.cuda_ldg
mov_sreg = _cuda_op.cuda_mov_sreg
wait_until = _cuda_op.cuda_wait_until
warp_reduce = _cuda_op.cuda_warp_reduce
+__shfl_sync = _cuda_op.__shfl_sync
+__shfl_up_sync = _cuda_op.__shfl_up_sync
+__shfl_down_sync = _cuda_op.__shfl_down_sync
+__shfl_xor_sync = _cuda_op.__shfl_xor_sync
+__activemask = _cuda_op.__activemask
_ir.op._init_op_api("tirx.cuda", __name__)
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/__init__.py b/python/tvm/backend/cuda/tile_primitive/elementwise/__init__.py
index c4dcfb63b8..4760792d96 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/__init__.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/__init__.py
@@ -32,7 +32,6 @@ from .register import *
# Suppress submodule-attribute leakage. Without an explicit ``__all__`` here,
# ``from tvm.backend.cuda.tile_primitive.elementwise import *`` (run by
# tile_primitive/__init__.py) re-exports the implicit submodule attributes
-# (``ops``, ``reg``, ``smem``, ``vec_emit``) — and ``ops`` in particular
-# shadows the top-level ``tile_primitive/ops.py`` (BinaryReduce / UnaryReduce
-# / ...) when downstream code does ``from tile_primitive import ops``.
+# (``ops``, ``reg``, ``smem``, ``vec_emit``) into the public dispatcher namespace.
+# Tile operator classes are owned by tvm.tirx.op.tile.
__all__: list[str] = []
diff --git a/python/tvm/backend/metal/op.py b/python/tvm/backend/metal/op.py
index d47f7cec02..a9e605eee5 100644
--- a/python/tvm/backend/metal/op.py
+++ b/python/tvm/backend/metal/op.py
@@ -18,7 +18,9 @@
from __future__ import annotations
-from tvm.ir import Call
+from tvm.ir import Call, Op
+from tvm.ir.op import _make_op_api
+from tvm.tirx import is_tensor_var
def make_filled_simdgroup_matrix(d, index, value, col=8, row=8, *, ty=None, span=None):
@@ -89,8 +91,37 @@ def simdgroup_multiply_accumulate(
)
+_simd_shuffle = _make_op_api(Op.get("tirx.metal.simd_shuffle"), __name__)
+_simd_shuffle_up = _make_op_api(Op.get("tirx.metal.simd_shuffle_up"), __name__)
+_simd_shuffle_down = _make_op_api(Op.get("tirx.metal.simd_shuffle_down"), __name__)
+
+
+def simd_shuffle(var, lane, *, ty=None, span=None):
+ """Shuffle a value from the selected SIMD lane."""
+ if is_tensor_var(var):
+ var = var[0]
+ return _simd_shuffle(var, lane, ty=ty, span=span)
+
+
+def simd_shuffle_up(var, delta, *, ty=None, span=None):
+ """Shuffle a value from a lower SIMD lane."""
+ if is_tensor_var(var):
+ var = var[0]
+ return _simd_shuffle_up(var, delta, ty=ty, span=span)
+
+
+def simd_shuffle_down(var, delta, *, ty=None, span=None):
+ """Shuffle a value from a higher SIMD lane."""
+ if is_tensor_var(var):
+ var = var[0]
+ return _simd_shuffle_down(var, delta, ty=ty, span=span)
+
+
__all__ = [
"make_filled_simdgroup_matrix",
+ "simd_shuffle",
+ "simd_shuffle_down",
+ "simd_shuffle_up",
"simdgroup_load",
"simdgroup_multiply_accumulate",
"simdgroup_store",
diff --git a/python/tvm/backend/metal/script.py b/python/tvm/backend/metal/script.py
index dd220af853..d0a0cd023b 100644
--- a/python/tvm/backend/metal/script.py
+++ b/python/tvm/backend/metal/script.py
@@ -19,13 +19,6 @@
from __future__ import annotations
from tvm.backend.metal import op as _metal_op
-from tvm.ir import Op
-from tvm.ir.op import _make_op_api
-from tvm.tirx import is_tensor_var
-
-_simd_shuffle = _make_op_api(Op.get("tirx.metal.simd_shuffle"), __name__)
-_simd_shuffle_up = _make_op_api(Op.get("tirx.metal.simd_shuffle_up"), __name__)
-_simd_shuffle_down = _make_op_api(Op.get("tirx.metal.simd_shuffle_down"), __name__)
class MetalNamespace:
@@ -37,23 +30,9 @@ class MetalNamespace:
self.simdgroup_store = _metal_op.simdgroup_store
self.simdgroup_multiply_accumulate = _metal_op.simdgroup_multiply_accumulate
- @staticmethod
- def simd_shuffle(var, lane, *, ty=None, span=None):
- if is_tensor_var(var):
- var = var[0]
- return _simd_shuffle(var, lane, ty=ty, span=span)
-
- @staticmethod
- def simd_shuffle_up(var, delta, *, ty=None, span=None):
- if is_tensor_var(var):
- var = var[0]
- return _simd_shuffle_up(var, delta, ty=ty, span=span)
-
- @staticmethod
- def simd_shuffle_down(var, delta, *, ty=None, span=None):
- if is_tensor_var(var):
- var = var[0]
- return _simd_shuffle_down(var, delta, ty=ty, span=span)
+ simd_shuffle = staticmethod(_metal_op.simd_shuffle)
+ simd_shuffle_up = staticmethod(_metal_op.simd_shuffle_up)
+ simd_shuffle_down = staticmethod(_metal_op.simd_shuffle_down)
__all__ = ["MetalNamespace"]
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py b/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
index 65c133da48..a30bedca3b 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
@@ -20,8 +20,8 @@
from tvm.ir import TensorRegion
from tvm.script import tirx as T
from tvm.tirx import Function, TilePrimitiveCall
+from tvm.tirx.op.tile import BinaryChain
from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
-from tvm.tirx.operator.tile_primitive.ops import BinaryChain
from ..binary.utils import InstType, try_find_inst_nary
from ..common import init_analyzer, nki_dim
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py b/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
index a994e79d5d..b7a89d47e6 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
@@ -20,8 +20,8 @@
from tvm.ir import TensorRegion
from tvm.script import tirx as T
from tvm.tirx import Function, TilePrimitiveCall
+from tvm.tirx.op.tile import BinaryReduce
from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
-from tvm.tirx.operator.tile_primitive.ops import BinaryReduce
from ..binary.utils import InstType, try_find_inst_nary
from ..common import init_analyzer, nki_dim
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/reduce_negate.py b/python/tvm/backend/trn/tile_primitive/compose_op/reduce_negate.py
index 5215c25f7e..bcdb5bb0ba 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/reduce_negate.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/reduce_negate.py
@@ -18,8 +18,8 @@
"""Implementation of ReduceNegate dispatch."""
from tvm.tirx import Function, TilePrimitiveCall
+from tvm.tirx.op.tile import ReduceNegate
from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
-from tvm.tirx.operator.tile_primitive.ops import ReduceNegate
from ..reduction.utils import reduction_trn
from .utils import optype_table
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py b/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
index b0ab572b0c..54f60ad6cd 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
@@ -20,8 +20,8 @@
from tvm.ir import TensorRegion
from tvm.script import tirx as T
from tvm.tirx import Function, TilePrimitiveCall
+from tvm.tirx.op.tile import UnaryReduce
from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
-from tvm.tirx.operator.tile_primitive.ops import UnaryReduce
from ..binary.utils import try_find_inst_nary
from ..common import init_analyzer, nki_dim
diff --git a/python/tvm/backend/trn/tile_primitive/private_alloc.py b/python/tvm/backend/trn/tile_primitive/private_alloc.py
index 16b737900a..88ef636c01 100644
--- a/python/tvm/backend/trn/tile_primitive/private_alloc.py
+++ b/python/tvm/backend/trn/tile_primitive/private_alloc.py
@@ -22,7 +22,7 @@ from tvm.backend.trn.tile_primitive.dim_utils import get_ewise_dim_map
from tvm.backend.trn.tile_primitive.instruction_generator import InstructionGenerator
from tvm.script import tirx as T
from tvm.tirx import FloatImm, IntImm, Stmt, Var
-from tvm.tirx.operator.tile_primitive.ops import (
+from tvm.tirx.op.tile import (
BinaryReduce,
Copy,
Exp,
diff --git a/python/tvm/backend/trn/tile_primitive/select/default.py b/python/tvm/backend/trn/tile_primitive/select/default.py
index 1d5dc5f46c..5e9e389b30 100644
--- a/python/tvm/backend/trn/tile_primitive/select/default.py
+++ b/python/tvm/backend/trn/tile_primitive/select/default.py
@@ -21,13 +21,13 @@ from tvm.backend.trn.layout import is_trainium_layout
from tvm.ir import TensorRegion
from tvm.script import tirx as T
from tvm.tirx import FloatImm, Function, TilePrimitiveCall
+from tvm.tirx.op.tile import Select
from tvm.tirx.operator.tile_primitive import (
DispatchContext,
fail,
predicate,
register_dispatch,
)
-from tvm.tirx.operator.tile_primitive.ops import Select
from ..common import init_analyzer, nki_dim
from ..dim_utils import get_ewise_dim_map
diff --git a/python/tvm/ir/prim/__init__.py b/python/tvm/ir/prim/__init__.py
index 3194de88c5..9cadc68f81 100644
--- a/python/tvm/ir/prim/__init__.py
+++ b/python/tvm/ir/prim/__init__.py
@@ -57,7 +57,35 @@ from .expr import (
Shuffle,
Sub,
)
-from .op import clz, convert, max_value, min_value
+from .op import (
+ all,
+ any,
+ bitwise_and,
+ bitwise_not,
+ bitwise_or,
+ bitwise_xor,
+ ceil,
+ ceildiv,
+ clz,
+ convert,
+ div,
+ floordiv,
+ floormod,
+ if_then_else,
+ indexdiv,
+ indexmod,
+ likely,
+ log2,
+ max,
+ max_value,
+ min,
+ min_value,
+ shift_left,
+ shift_right,
+ truncdiv,
+ truncmod,
+ vscale,
+)
def expr_deep_equal(lhs: Expr, rhs: Expr) -> bool:
diff --git a/python/tvm/ir/prim/op.py b/python/tvm/ir/prim/op.py
index da4d4b04e3..86d2598fa9 100644
--- a/python/tvm/ir/prim/op.py
+++ b/python/tvm/ir/prim/op.py
@@ -81,3 +81,486 @@ def clz(x):
The result.
"""
return Call("prim.clz", [x], ty="int32")
+
+
+def any(*args, span=None):
+ """Create a new experssion of the union of all conditions in the arguments
+
+ Parameters
+ ----------
+ args : list
+ List of symbolic boolean expressions
+
+ span : Optional[Span]
+ The location of this operator in the source code.
+
+ Returns
+ -------
+ expr: Expr
+ Expression
+ """
+ if not args:
+ raise ValueError("Any must take at least 1 argument")
+ if len(args) == 1:
+ return args[0]
+ val = _ffi_api._OpOr(args[0], args[1], span) # type: ignore
+ for i in range(2, len(args)):
+ val = _ffi_api._OpOr(val, args[i], span) # type: ignore
+ return val
+
+
+def all(*args, span=None):
+ """Create a new expression of the intersection of all conditions in the
+ arguments
+
+ Parameters
+ ----------
+ args : list
+ List of symbolic boolean expressions
+
+ span : Optional[Span]
+ The location of this operator in the source code.
+
+ Returns
+ -------
+ expr: Expr
+ Expression
+ """
+ if not args:
+ raise ValueError("Any must take at least 1 argument")
+ if len(args) == 1:
+ return args[0]
+ val = _ffi_api._OpAnd(args[0], args[1], span) # type: ignore
+ for i in range(2, len(args)):
+ val = _ffi_api._OpAnd(val, args[i], span) # type: ignore
+ return val
+
+
+def log2(x, *, ty=None, span=None):
+ """Take log2 of input x.
+
+ Parameters
+ ----------
+ x : Expr
+ Input argument.
+
+ Returns
+ -------
+ y : Expr
+ The result.
+ """
+ x = convert(x)
+ return Call("prim.log2", [x], ty=ty, span=span)
+
+
+def ceil(x, span=None):
+ """Take ceil of float input x.
+
+ Parameters
+ ----------
+ x : Expr
+ Input argument.
+
+ span : Optional[Span]
+ The location of this operator in the source code.
+
+ Returns
+ -------
+ y : Expr
+ The result.
+ """
+ return _ffi_api.ceil(x, span) # type: ignore
+
+
+def bitwise_and(x, y, span=None):
+ """Take bitwise and of two values
+
+ Parameters
+ ----------
+ x : Expr
+ Left operand
+
+ y : Expr
+ Right operand
+
+ span : Optional[Span]
+ The location of this operator in the source code.
+
+ Returns
+ -------
+ res : Expr
+ The result.
+ """
+ return _ffi_api.bitwise_and(x, y, span)
+
+
+def bitwise_not(x, span=None):
+ """Take bitwise not of input value
+
+ Parameters
+ ----------
+ x : Expr
+ Input operand
+
+ span : Optional[Span]
+ The location of this operator in the source code.
+
+ Returns
+ -------
+ res : Expr
+ The result.
+ """
+ return _ffi_api.bitwise_not(x, span)
+
+
+def bitwise_or(x, y, span=None):
+ """Take bitwise or of two values
+
+ Parameters
+ ----------
+ x : Expr
+ Left operand
+
+ y : Expr
+ Right operand
+
+ span : Optional[Span]
+ The location of this operator in the source code.
+
+ Returns
+ -------
+ res : Expr
+ The result.
+ """
+ return _ffi_api.bitwise_or(x, y, span)
+
+
+def bitwise_xor(x, y, span=None):
+ """Take bitwise xor of two values
+
+ Parameters
+ ----------
+ x : Expr
+ Left operand
+
+ y : Expr
+ Right operand
+
+ span : Optional[Span]
+ The location of this operator in the source code.
+
+ Returns
+ -------
+ res : Expr
+ The result.
+ """
+ return _ffi_api.bitwise_xor(x, y, span)
+
+
+def likely(cond, span=None):
+ """Mark condition as likely.
+
+ Parameters
+ ----------
+
+ cond : Expr
+ Input argument.
+
+ span : Optional[Span]
+ The location of this operator in the source code.
+
+ Returns
+ -------
+ y : Expr
+ The marked expression.
+ """
+ return _ffi_api.likely(cond, span) # type: ignore
+
+
+def shift_left(x, y, span=None):
+ """Return the result of x left shifted by y bits.
+
+ Parameters
+ ----------
+ x : Expr
+ Input argument.
+
+ y : Expr
+ Input argument.
+
+ Returns
+ -------
+ z : Expr
+ The result.
+ """
+ return _ffi_api.left_shift(x, y, span)
+
+
+def shift_right(x, y, span=None):
+ """Return the result of x right shifted by y bits.
+
+ Parameters
+ ----------
+ x : Expr
+ Input argument.
+
+ y : Expr
+ Input argument.
+
+ Returns
+ -------
+ z : Expr
+ The result.
+ """
+ return _ffi_api.right_shift(x, y, span)
+
+
+def if_then_else(cond, t, f, span=None):
+ """Conditional selection expression.
+
+ Parameters
+ ----------
+ cond : Expr
+ The condition
+
+ t : Expr
+ The result expression if cond is true.
+
+ f : Expr
+ The result expression if cond is false.
+
+ span : Optional[Span]
+ The location of this operator in the source.
+
+ Returns
+ -------
+ result : Node
+ The result of conditional expression.
+
+ Note
+ ----
+ Unlike Select, if_then_else will not execute
+ the branch that does not satisfy the condition.
+ You can use it to guard against out of bound access.
+ Unlike Select, if_then_else cannot be vectorized
+ if some lanes in the vector have different conditions.
+ """
+ return _ffi_api._OpIfThenElse(cond, t, f, span) # type: ignore
+
+
+def div(a, b, span=None):
+ """Compute a / b as in C/C++ semantics.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand, known to be non-negative.
+
+ b : Expr
+ The right hand operand, known to be non-negative.
+
+ span : Optional[Span]
+ The location of this operator in the source.
+
+ Returns
+ -------
+ res : Expr
+ The result expression.
+ Note
+ ----
+ When operands are integers, returns truncdiv(a, b, span).
+ """
+ return _ffi_api._OpDiv(a, b, span) # type: ignore
+
+
+def indexdiv(a, b, span=None):
+ """Compute floor(a / b) where a and b are non-negative.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand, known to be non-negative.
+
+ b : Expr
+ The right hand operand, known to be non-negative.
+
+ span : Optional[Span]
+ The location of this operator in the source.
+
+ Returns
+ -------
+ res : Expr
+ The result expression.
+
+ Note
+ ----
+ Use this function to split non-negative indices.
+ This function may take advantage of operands'
+ non-negativeness.
+ """
+ return _ffi_api._OpIndexDiv(a, b, span) # type: ignore
+
+
+def indexmod(a, b, span=None):
+ """Compute the remainder of indexdiv. a and b are non-negative.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand, known to be non-negative.
+
+ b : Expr
+ The right hand operand, known to be non-negative.
+
+ span : Optional[Span]
+ The location of this operator in the source.
+
+ Returns
+ -------
+ res : Expr
+ The result expression.
+
+ Note
+ ----
+ Use this function to split non-negative indices.
+ This function may take advantage of operands'
+ non-negativeness.
+ """
+ return _ffi_api._OpIndexMod(a, b, span) # type: ignore
+
+
+def truncdiv(a, b, span=None):
+ """Compute the truncdiv of two expressions.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand
+
+ b : Expr
+ The right hand operand
+
+ span : Optional[Span]
+ The location of this operator in the source.
+
+ Returns
+ -------
+ res : Expr
+ The result expression.
+
+ Note
+ ----
+ This is the default integer division behavior in C.
+ """
+ return _ffi_api._OpTruncDiv(a, b, span) # type: ignore
+
+
+def truncmod(a, b, span=None):
+ """Compute the truncmod of two expressions.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand
+
+ b : Expr
+ The right hand operand
+
+ span : Optional[Span]
+ The location of this operator in the source.
+
+ Returns
+ -------
+ res : Expr
+ The result expression.
+
+ Note
+ ----
+ This is the default integer division behavior in C.
+ """
+ return _ffi_api._OpTruncMod(a, b, span) # type: ignore
+
+
+def floordiv(a, b, span=None):
+ """Compute the floordiv of two expressions.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand
+
+ b : Expr
+ The right hand operand
+
+ span : Optional[Span]
+ The location of this operator in the source.
+
+ Returns
+ -------
+ res : Expr
+ The result expression.
+ """
+ return _ffi_api._OpFloorDiv(a, b, span) # type: ignore
+
+
+def floormod(a, b, span=None):
+ """Compute the floormod of two expressions.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand
+
+ b : Expr
+ The right hand operand
+
+ span : Optional[Span]
+ The location of this operator in the source.
+
+ Returns
+ -------
+ res : Expr
+ The result expression.
+ """
+ return _ffi_api._OpFloorMod(a, b, span) # type: ignore
+
+
+def ceildiv(lhs, rhs, span=None):
+ """Generic ceildiv operator.
+
+ Parameters
+ ----------
+ lhs : object
+ The left operand.
+ rhs : object
+ The right operand.
+ span : Optional[Span]
+ The location of this operator in the source.
+
+ Returns
+ -------
+ op : tvm.Expr
+ The result Expr of ceildiv operaton.
+ """
+ return _ffi_api._OpCeilDiv(lhs, rhs, span) # type: ignore
+
+
+def vscale(*, ty=None, span=None):
+ """Get the target's vscale value. It will be lowered to llvm.vscale intrinsic
+ (https://llvm.org/docs/LangRef.html#llvm-vscale-intrinsic)
+ Returns
+ -------
+ call : Expr
+ Call to the vscale intrinsic
+ """
+ return Call("prim.vscale", [], ty=ty, span=span)
+
+
+def min(a, b, span=None):
+ """Elementwise minimum of two primitive expressions."""
+ return _ffi_api._OpMin(a, b, span)
+
+
+def max(a, b, span=None):
+ """Elementwise maximum of two primitive expressions."""
+ return _ffi_api._OpMax(a, b, span)
diff --git a/python/tvm/relax/op/__init__.py b/python/tvm/relax/op/__init__.py
index 170c61e7bd..2740d97015 100644
--- a/python/tvm/relax/op/__init__.py
+++ b/python/tvm/relax/op/__init__.py
@@ -18,11 +18,16 @@
# pylint: disable= redefined-builtin
"""Relax core operators."""
-# Register operator gradient functions
-from . import _op_gradient, builtin, ccl, distributed, grad, image, memory, nn, op_attrs
+# Operator families register their attrs before exposure.
+from . import builtin, ccl, distributed, grad, image, memory, nn
# Operators
from .base import (
+ CallTIRWithGradAttrs,
+ CallInplacePackedAttrs,
+ CallTIRInplaceAttrs,
+ ToVDeviceAttrs,
+ HintOnDeviceAttrs,
assert_op,
call_builtin_with_ctx,
call_dps_packed,
@@ -75,6 +80,8 @@ from .binary import (
subtract,
)
from .create import (
+ InitAttrs,
+ TriluAttrs,
arange,
full,
full_like,
@@ -88,10 +95,47 @@ from .create import (
zeros,
zeros_like,
)
-from .datatype import astype, wrap_param
-from .index import dynamic_strided_slice, strided_slice, take
-from .linear_algebra import einsum, linear, matmul, outer
+from .datatype import (
+ AstypeAttrs,
+ WrapParamAttrs,
+ astype,
+ wrap_param,
+)
+from .index import (
+ TakeAttrs,
+ StridedSliceAttrs,
+ dynamic_strided_slice,
+ strided_slice,
+ take,
+)
+from .linear_algebra import (
+ MatmulAttrs,
+ EinsumAttrs,
+ einsum,
+ linear,
+ matmul,
+ outer,
+)
from .manipulate import (
+ ConcatAttrs,
+ ExpandDimsAttrs,
+ PermuteDimsAttrs,
+ SplitAttrs,
+ SqueezeAttrs,
+ StackAttrs,
+ IndexPutAttrs,
+ LayoutTransformAttrs,
+ RepeatAttrs,
+ TileAttrs,
+ FlipAttrs,
+ ReverseSequenceAttrs,
+ GatherElementsAttrs,
+ GatherNDAttrs,
+ MeshgridAttrs,
+ ScatterElementsAttrs,
+ ScatterNDAttrs,
+ SliceScatterAttrs,
+ OneHotAttrs,
broadcast_to,
collapse_sum_like,
collapse_sum_to,
@@ -119,12 +163,45 @@ from .manipulate import (
tile,
)
from .mask import masked_fill
-from .qdq import dequantize, quantize
-from .sampling import multinomial_from_uniform
-from .search import argmax, argmin, where, bucketize
+from .qdq import (
+ QuantizeAttrs,
+ dequantize,
+ quantize,
+)
+from .sampling import (
+ MultinomialFromUniformAttrs,
+ multinomial_from_uniform,
+)
+from .search import (
+ ArgmaxArgminAttrs,
+ argmax,
+ argmin,
+ where,
+ bucketize,
+)
from .set import nonzero, unique
-from .sorting import argsort, sort, topk
-from .statistical import cumprod, cumsum, max, mean, min, prod, std, sum, variance, median
+from .sorting import (
+ SortAttrs,
+ ArgsortAttrs,
+ TopKAttrs,
+ argsort,
+ sort,
+ topk,
+)
+from .statistical import (
+ StatisticalAttrs,
+ ScanopAttrs,
+ cumprod,
+ cumsum,
+ max,
+ mean,
+ min,
+ prod,
+ std,
+ sum,
+ variance,
+ median,
+)
from .ternary import ewise_fma
from .unary import (
abs,
@@ -170,6 +247,10 @@ from .vision import (
)
+# Register gradients after the operator families and their attrs are available.
+from . import _op_gradient
+
+
def _register_op_make():
# pylint: disable=import-outside-toplevel
from .. import expr
diff --git a/python/tvm/relax/op/base.py b/python/tvm/relax/op/base.py
index f67ef1ef03..270873d807 100644
--- a/python/tvm/relax/op/base.py
+++ b/python/tvm/relax/op/base.py
@@ -23,7 +23,7 @@ import tvm_ffi
import tvm
import tvm.runtime
-from tvm.ir import Call, Op, StringImm
+from tvm.ir import Attrs, Call, Op, StringImm
from tvm.ir.attrs import make_node as _make_attrs
from tvm.runtime import Object, ObjectConvertible
@@ -207,6 +207,11 @@ def call_tir_packed(gvar: GlobalVar, args: Expr, *, ty=None, span=None) -> Call:
return Call("relax.call_tir_packed", [gvar, args], ty=ty, span=span)
+@tvm_ffi.register_object("relax.attrs.CallTIRWithGradAttrs")
+class CallTIRWithGradAttrs(Attrs):
+ """Attributes used in call_tir_with_grad operator"""
+
+
def call_tir_with_grad(
gvar: GlobalVar,
args: Expr,
@@ -259,6 +264,11 @@ def call_tir_with_grad(
)
+@tvm_ffi.register_object("relax.attrs.CallTIRInplaceAttrs")
+class CallTIRInplaceAttrs(Attrs):
+ """Attributes used in call_tir_inplace operator"""
+
+
def call_tir_inplace(
gvar: GlobalVar,
args: Expr,
@@ -745,6 +755,11 @@ def shape_to_tensor(expr: Expr, *, ty=None, span=None) -> Expr:
return Call("relax.shape_to_tensor", [expr], ty=ty, span=span) # type: ignore # pylint: disable=no-member
+@tvm_ffi.register_object("relax.attrs.CallInplacePackedAttrs")
+class CallInplacePackedAttrs(Attrs):
+ """Attributes used in call_inplace_packed operator"""
+
+
def call_inplace_packed(
func: str | ExternFunc | GlobalVar,
*args: Expr,
@@ -925,6 +940,11 @@ def invoke_pure_closure(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.ToVDeviceAttrs")
+class ToVDeviceAttrs(Attrs):
+ """Attributes used in to_vdevice operator"""
+
+
def to_vdevice(data, dst_vdevice, *, ty=None, span=None) -> Expr:
"""Copy data to the destination device. This
operator helps data transferring between difference devices for
@@ -952,6 +972,11 @@ def to_vdevice(data, dst_vdevice, *, ty=None, span=None) -> Expr:
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.HintOnDeviceAttrs")
+class HintOnDeviceAttrs(Attrs):
+ """Attributes used in hint_on_device operator"""
+
+
def hint_on_device(data, dst_vdevice, memory_scope="global") -> Expr:
"""It provides a hint specifying the device on which the input data should be executed.
This hint is utilized by RealizeVDevice to propagate the virtual device."
diff --git a/python/tvm/relax/op/ccl/__init__.py b/python/tvm/relax/op/ccl/__init__.py
index 47006358f9..92be50094c 100644
--- a/python/tvm/relax/op/ccl/__init__.py
+++ b/python/tvm/relax/op/ccl/__init__.py
@@ -17,4 +17,12 @@
# under the License.
"""CCL related operators."""
-from .ccl import allgather, allreduce, broadcast_from_worker0, scatter_from_worker0
+from .ccl import (
+ ScatterCollectiveAttrs,
+ AllReduceAttrs,
+ AllGatherAttrs,
+ allgather,
+ allreduce,
+ broadcast_from_worker0,
+ scatter_from_worker0,
+)
diff --git a/python/tvm/relax/op/ccl/ccl.py b/python/tvm/relax/op/ccl/ccl.py
index da7b3fddee..d9d464e357 100644
--- a/python/tvm/relax/op/ccl/ccl.py
+++ b/python/tvm/relax/op/ccl/ccl.py
@@ -16,12 +16,20 @@
# under the License.
"""Relax Collective Communications Library (CCL) operators"""
+import tvm_ffi
+
+from tvm.ir import Attrs
from tvm.ir import Call as _Call
from tvm.ir.attrs import make_node as _make_attrs
from ...expr import Expr
+@tvm_ffi.register_object("relax.attrs.AllReduceAttrs")
+class AllReduceAttrs(Attrs):
+ """Attributes used in allreduce operator"""
+
+
def allreduce(x, op_type: str = "sum", in_group: bool = True, *, ty=None, span=None): # pylint: disable=invalid-name
"""Allreduce operator
@@ -56,6 +64,11 @@ def allreduce(x, op_type: str = "sum", in_group: bool = True, *, ty=None, span=N
) # type: ignore # pylint: disable=no-member
+@tvm_ffi.register_object("relax.attrs.AllGatherAttrs")
+class AllGatherAttrs(Attrs):
+ """Attributes used in allgather operator"""
+
+
def allgather(x, num_workers: int, in_group: bool = True, *, ty=None, span=None): # pylint: disable=invalid-name
"""AllGather operator
@@ -105,6 +118,11 @@ def broadcast_from_worker0(x: Expr, *, ty=None, span=None) -> Expr:
)
+@tvm_ffi.register_object("relax.attrs.ScatterCollectiveAttrs")
+class ScatterCollectiveAttrs(Attrs):
+ """Attributes used in scatter collective operators"""
+
+
def scatter_from_worker0(x: Expr, num_workers: int, axis: int = 0, *, ty=None, span=None) -> Expr:
"""Perform a scatter operation from worker-0, chunking the given buffer into equal parts.
diff --git a/python/tvm/relax/op/create.py b/python/tvm/relax/op/create.py
index 1e602c5a36..13dd04d6fa 100644
--- a/python/tvm/relax/op/create.py
+++ b/python/tvm/relax/op/create.py
@@ -16,9 +16,11 @@
# under the License.
"""Creation operators."""
+import tvm_ffi
+
from tvm import DataType, DataTypeCode
+from tvm.ir import Attrs, PrimType, is_prim_expr
from tvm.ir import Call as _Call
-from tvm.ir import PrimType, is_prim_expr
from tvm.ir.attrs import make_node as _make_attrs
from ..expr import Expr, ShapeExpr, prim_value
@@ -38,6 +40,11 @@ def _normalize_shape(shape):
return shape
+@tvm_ffi.register_object("relax.attrs.InitAttrs")
+class InitAttrs(Attrs):
+ """Attributes used in full/full_like, ones/ones_like, and zeros/zeros_like operator"""
+
+
def full(
shape: tuple[PrimExprLike] | Expr,
fill_value: Expr,
@@ -418,6 +425,11 @@ def hamming_window(window_size, periodic, alpha, beta, dtype, *, ty=None, span=N
)
+@tvm_ffi.register_object("relax.attrs.TriluAttrs")
+class TriluAttrs(Attrs):
+ """Attributes used in tril and triu operator"""
+
+
def tril(x: Expr, k: int | Expr = 0, *, ty=None, span=None) -> Expr:
"""Return the lower triangular part of a matrix or a batch of matrices.
diff --git a/python/tvm/relax/op/datatype.py b/python/tvm/relax/op/datatype.py
index 9e1c7186ca..6c58b340c7 100644
--- a/python/tvm/relax/op/datatype.py
+++ b/python/tvm/relax/op/datatype.py
@@ -16,9 +16,11 @@
# under the License.
"""Datatype operators."""
+import tvm_ffi
+
from tvm import DataType
+from tvm.ir import Attrs, PrimType
from tvm.ir import Call as _Call
-from tvm.ir import PrimType
from tvm.ir.attrs import make_node as _make_attrs
from ..expr import Expr
@@ -28,6 +30,11 @@ def _raw_dtype(dtype):
return dtype.dtype if isinstance(dtype, PrimType) else dtype
+@tvm_ffi.register_object("relax.attrs.AstypeAttrs")
+class AstypeAttrs(Attrs):
+ """Attributes used in astype operator"""
+
+
def astype(x: Expr, dtype: str | DataType | PrimType, *, ty=None, span=None) -> Expr:
"""Cast input tensor to the given data type.
@@ -53,6 +60,11 @@ def astype(x: Expr, dtype: str | DataType | PrimType, *, ty=None, span=None) ->
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.WrapParamAttrs")
+class WrapParamAttrs(Attrs):
+ """Attributes used in wrap_param operator"""
+
+
def wrap_param(
data: Expr,
dtype: str | DataType | PrimType = "float32",
diff --git a/python/tvm/relax/op/image/__init__.py b/python/tvm/relax/op/image/__init__.py
index dcc0d1f883..daa645b5e6 100644
--- a/python/tvm/relax/op/image/__init__.py
+++ b/python/tvm/relax/op/image/__init__.py
@@ -17,4 +17,10 @@
# under the License.
"""Image operators."""
-from .image import affine_grid, grid_sample, resize2d, resize3d
+from .image import (
+ Resize2DAttrs,
+ affine_grid,
+ grid_sample,
+ resize2d,
+ resize3d,
+)
diff --git a/python/tvm/relax/op/image/image.py b/python/tvm/relax/op/image/image.py
index 7f8ad29da6..a3c5e22f20 100644
--- a/python/tvm/relax/op/image/image.py
+++ b/python/tvm/relax/op/image/image.py
@@ -18,10 +18,12 @@
from typing import cast
+import tvm_ffi
+
from tvm import DataType
+from tvm.ir import Attrs, is_prim_expr
from tvm.ir import Call as _Call
from tvm.ir import PrimType as _PrimType
-from tvm.ir import is_prim_expr
from tvm.ir.attrs import make_node as _make_attrs
from ...expr import Expr, ShapeExpr
@@ -30,6 +32,11 @@ PrimExprLike = int | Expr
SizeLike = PrimExprLike | tuple[PrimExprLike, ...]
+@tvm_ffi.register_object("relax.attrs.Resize2DAttrs")
+class Resize2DAttrs(Attrs):
+ """Attributes used in image resize2d operator"""
+
+
def resize2d(
data: Expr,
size: SizeLike,
diff --git a/python/tvm/relax/op/index.py b/python/tvm/relax/op/index.py
index a7fda22b45..df7c8c891a 100644
--- a/python/tvm/relax/op/index.py
+++ b/python/tvm/relax/op/index.py
@@ -16,7 +16,10 @@
# under the License.
"""Indexing operators."""
+import tvm_ffi
+
from tvm.error import InternalError as _InternalError
+from tvm.ir import Attrs
from tvm.ir import Call as _Call
from tvm.ir.attrs import make_node as _make_attrs
from tvm.relax.expr import Tuple as _Tuple
@@ -27,6 +30,11 @@ from ..utils import convert_to_expr
PrimExprLike = int | Expr
+@tvm_ffi.register_object("relax.attrs.TakeAttrs")
+class TakeAttrs(Attrs):
+ """Attributes used in take operator"""
+
+
def take(
x: Expr,
indices: Expr,
@@ -74,6 +82,11 @@ def take(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.StridedSliceAttrs")
+class StridedSliceAttrs(Attrs):
+ """Attributes used in strided_slice operator"""
+
+
def strided_slice(
x: Expr,
axes: Expr,
diff --git a/python/tvm/relax/op/linear_algebra.py b/python/tvm/relax/op/linear_algebra.py
index 314470cabd..493092e748 100644
--- a/python/tvm/relax/op/linear_algebra.py
+++ b/python/tvm/relax/op/linear_algebra.py
@@ -17,7 +17,10 @@
# pylint: disable=invalid-name
"""Relax linear algebra operators"""
+import tvm_ffi
+
from tvm import DataType
+from tvm.ir import Attrs
from tvm.ir import Call as _Call
from tvm.ir import PrimType as _PrimType
from tvm.ir.attrs import make_node as _make_attrs
@@ -27,6 +30,11 @@ from ..expr import Tuple as RxTuple
from .manipulate import permute_dims
+@tvm_ffi.register_object("relax.attrs.MatmulAttrs")
+class MatmulAttrs(Attrs):
+ """Attributes for matmul operator"""
+
+
def matmul(
x1: Expr,
x2: Expr,
@@ -108,6 +116,11 @@ def linear(
return x + bias if bias is not None else x
+@tvm_ffi.register_object("relax.attrs.EinsumAttrs")
+class EinsumAttrs(Attrs):
+ """Attributes for einsum operator"""
+
+
def einsum(operands, subscripts, *, ty=None, span=None):
"""Evaluates the Einstein summation convention on data
diff --git a/python/tvm/relax/op/manipulate.py b/python/tvm/relax/op/manipulate.py
index 58c58c7a63..b02253978a 100644
--- a/python/tvm/relax/op/manipulate.py
+++ b/python/tvm/relax/op/manipulate.py
@@ -18,9 +18,11 @@
from collections.abc import Callable
+import tvm_ffi
+
from tvm.error import InternalError as _InternalError
+from tvm.ir import Attrs, is_prim_expr
from tvm.ir import Call as _Call
-from tvm.ir import is_prim_expr
from tvm.ir.attrs import make_node as _make_attrs
from tvm.runtime import DataTypeCode
from tvm.tirx import FloatImm, IndexMap, IntImm
@@ -53,6 +55,11 @@ def broadcast_to(x: Expr, shape: tuple[PrimExprLike] | Expr, *, ty=None, span=No
return _Call("relax.broadcast_to", [x, shape], ty=ty, span=span) # type: ignore
+@tvm_ffi.register_object("relax.attrs.ConcatAttrs")
+class ConcatAttrs(Attrs):
+ """Attributes for concat operator"""
+
+
def concat(
tensors: Expr | list[Expr],
axis: int | None = 0,
@@ -88,6 +95,11 @@ def concat(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.ExpandDimsAttrs")
+class ExpandDimsAttrs(Attrs):
+ """Attributes for expand_dims operator"""
+
+
def expand_dims(x: Expr, axis: int | list[int], *, ty=None, span=None) -> Expr:
"""Insert new axes at the positions given by `axis`.
@@ -133,6 +145,11 @@ def flatten(x: Expr, *, ty=None, span=None) -> Expr:
return _Call("relax.flatten", [x], ty=ty, span=span) # type: ignore
+@tvm_ffi.register_object("relax.attrs.LayoutTransformAttrs")
+class LayoutTransformAttrs(Attrs):
+ """Attributes used in layout_transform operator"""
+
+
def layout_transform(
x: Expr,
index_map: Callable | IndexMap,
@@ -190,6 +207,11 @@ def layout_transform(
)
+@tvm_ffi.register_object("relax.attrs.PermuteDimsAttrs")
+class PermuteDimsAttrs(Attrs):
+ """Attributes for permute_dims operator"""
+
+
def permute_dims(x: Expr, axes: list[int] | None = None, *, ty=None, span=None) -> Expr:
"""Permutes the dimensions of an array.
@@ -253,6 +275,11 @@ def reshape(x: Expr, shape: tuple[PrimExprLike] | Expr, *, ty=None, span=None) -
return _Call("relax.reshape", [x, shape], ty=ty, span=span) # type: ignore
+@tvm_ffi.register_object("relax.attrs.SplitAttrs")
+class SplitAttrs(Attrs):
+ """Attributes used in split operator"""
+
+
def split(
x: Expr,
indices_or_sections: int | list[PrimExprLike],
@@ -306,6 +333,11 @@ def split(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.SqueezeAttrs")
+class SqueezeAttrs(Attrs):
+ """Attributes for squeeze operator"""
+
+
def squeeze(x: Expr, axis: int | list[int] | None = None, *, ty=None, span=None) -> Expr:
"""Squeeze axes in the array.
@@ -335,6 +367,11 @@ def squeeze(x: Expr, axis: int | list[int] | None = None, *, ty=None, span=None)
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.StackAttrs")
+class StackAttrs(Attrs):
+ """Attributes for concat operator"""
+
+
def stack(tensors: Expr | list[Expr], axis: int = 0, *, ty=None, span=None) -> Expr:
"""Stack the input tensors along a new axis.
@@ -428,6 +465,11 @@ def collapse_sum_to(data: Expr, shape: tuple[PrimExprLike] | Expr, *, ty=None, s
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.RepeatAttrs")
+class RepeatAttrs(Attrs):
+ """Attributes for repeat operator"""
+
+
def repeat(
data: Expr,
repeats: int,
@@ -474,6 +516,11 @@ def repeat(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.TileAttrs")
+class TileAttrs(Attrs):
+ """Attributes for tile operator"""
+
+
def tile(
data: Expr,
repeats: int | tuple[int] | list[int],
@@ -528,6 +575,11 @@ def tile(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.FlipAttrs")
+class FlipAttrs(Attrs):
+ """Attributes for flip operator"""
+
+
def flip(data, axis, *, ty=None, span=None):
"""Reverses the order of elements along given axis while preserving array shape.
@@ -562,6 +614,11 @@ def flip(data, axis, *, ty=None, span=None):
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.ReverseSequenceAttrs")
+class ReverseSequenceAttrs(Attrs):
+ """Attributes for reverse_sequence operator"""
+
+
def reverse_sequence(
data: Expr,
seq_lengths: Expr,
@@ -603,6 +660,11 @@ def reverse_sequence(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.GatherElementsAttrs")
+class GatherElementsAttrs(Attrs):
+ """Attributes for gather_elements operator"""
+
+
def gather_elements(data: Expr, indices: Expr, axis: int = 0, *, ty=None, span=None) -> Expr:
"""Gather elements from data according to indices along the specified axis.
@@ -645,6 +707,11 @@ def gather_elements(data: Expr, indices: Expr, axis: int = 0, *, ty=None, span=N
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.GatherNDAttrs")
+class GatherNDAttrs(Attrs):
+ """Attributes for gather_nd operator"""
+
+
def gather_nd(data: Expr, indices: Expr, batch_dims: int = 0, *, ty=None, span=None) -> Expr:
"""Update data at positions defined by indices with values in updates.
@@ -756,6 +823,11 @@ def index_tensor(data: Expr, indices: Expr | list[Expr], *, ty=None, span=None)
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.IndexPutAttrs")
+class IndexPutAttrs(Attrs):
+ """Attributes for index_put operator"""
+
+
def index_put(
data: Expr,
indices: Expr | tuple[Expr],
@@ -817,6 +889,11 @@ def index_put(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.MeshgridAttrs")
+class MeshgridAttrs(Attrs):
+ """Attributes for meshgrid operator"""
+
+
def meshgrid(
tensors: Expr | list[Expr],
indexing: str | None = "ij",
@@ -852,6 +929,11 @@ def meshgrid(
)
+@tvm_ffi.register_object("relax.attrs.ScatterElementsAttrs")
+class ScatterElementsAttrs(Attrs):
+ """Attributes for scatter_elements operator"""
+
+
def scatter_elements(
data: Expr,
indices: Expr,
@@ -943,6 +1025,11 @@ def scatter_elements(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.ScatterNDAttrs")
+class ScatterNDAttrs(Attrs):
+ """Attributes for scatter_nd operator"""
+
+
def scatter_nd(
data: Expr,
indices: Expr,
@@ -996,6 +1083,11 @@ def scatter_nd(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.SliceScatterAttrs")
+class SliceScatterAttrs(Attrs):
+ """Attributes for slice_scatter operator"""
+
+
def slice_scatter(
input_tensor: Expr,
src: Expr,
@@ -1050,6 +1142,11 @@ def slice_scatter(
)
+@tvm_ffi.register_object("relax.attrs.OneHotAttrs")
+class OneHotAttrs(Attrs):
+ """Attributes for one_hot operator"""
+
+
def one_hot(
indices: Expr,
on_value: int | float | Expr,
diff --git a/python/tvm/relax/op/nn/__init__.py b/python/tvm/relax/op/nn/__init__.py
index 57128282b9..9304b1dcbb 100644
--- a/python/tvm/relax/op/nn/__init__.py
+++ b/python/tvm/relax/op/nn/__init__.py
@@ -18,6 +18,32 @@
"""Neural network related operators."""
from .nn import (
+ Conv2DAttrs,
+ Conv3DAttrs,
+ Conv2DTransposeAttrs,
+ Conv3DTransposeAttrs,
+ Pool2DAttrs,
+ AdaptivePool2DAttrs,
+ SoftmaxAttrs,
+ BatchNormAttrs,
+ LayerNormAttrs,
+ InstanceNormAttrs,
+ DropoutAttrs,
+ PadAttrs,
+ AttentionAttrs,
+ Conv1DAttrs,
+ Conv1DTransposeAttrs,
+ Pool1DAttrs,
+ Pool3DAttrs,
+ AdaptivePool1DAttrs,
+ AdaptivePool3DAttrs,
+ LeakyReluAttrs,
+ SoftplusAttrs,
+ PReluAttrs,
+ PixelShuffleAttrs,
+ GroupNormAttrs,
+ RMSNormAttrs,
+ NLLLossAttrs,
adaptive_avg_pool1d,
adaptive_avg_pool2d,
adaptive_avg_pool3d,
diff --git a/python/tvm/relax/op/nn/nn.py b/python/tvm/relax/op/nn/nn.py
index 22b4d48d88..4800748f56 100644
--- a/python/tvm/relax/op/nn/nn.py
+++ b/python/tvm/relax/op/nn/nn.py
@@ -16,8 +16,11 @@
# under the License.
"""Relax Neural Network (NN) operators"""
+import tvm_ffi
+
from tvm import DataType, relax
from tvm.error import InternalError as _InternalError
+from tvm.ir import Attrs
from tvm.ir import Call as _Call
from tvm.ir import PrimType as _PrimType
from tvm.ir.attrs import make_node as _make_attrs
@@ -26,6 +29,11 @@ from tvm.tirx import FloatImm
from ...expr import Expr
+@tvm_ffi.register_object("relax.attrs.Conv1DAttrs")
+class Conv1DAttrs(Attrs):
+ """Attributes for nn.conv1d"""
+
+
def conv1d(
data: Expr,
weight: Expr,
@@ -144,6 +152,11 @@ def conv1d(
)
+@tvm_ffi.register_object("relax.attrs.Conv2DAttrs")
+class Conv2DAttrs(Attrs):
+ """Attributes for nn.conv2d"""
+
+
def conv2d(
data: Expr,
weight: Expr,
@@ -266,6 +279,11 @@ def conv2d(
)
+@tvm_ffi.register_object("relax.attrs.Conv3DAttrs")
+class Conv3DAttrs(Attrs):
+ """Attributes for nn.conv3d"""
+
+
def conv3d(
data: Expr,
weight: Expr,
@@ -394,6 +412,11 @@ def conv3d(
)
+@tvm_ffi.register_object("relax.attrs.Conv1DTransposeAttrs")
+class Conv1DTransposeAttrs(Attrs):
+ """Attributes for nn.conv1d_transpose"""
+
+
def conv1d_transpose(
data: Expr,
weight: Expr,
@@ -507,6 +530,11 @@ def conv1d_transpose(
)
+@tvm_ffi.register_object("relax.attrs.Conv2DTransposeAttrs")
+class Conv2DTransposeAttrs(Attrs):
+ """Attributes for nn.conv2d_transpose"""
+
+
def conv2d_transpose(
data: Expr,
weight: Expr,
@@ -637,6 +665,11 @@ def conv2d_transpose(
)
+@tvm_ffi.register_object("relax.attrs.Conv3DTransposeAttrs")
+class Conv3DTransposeAttrs(Attrs):
+ """Attributes for nn.conv3d_transpose"""
+
+
def conv3d_transpose(
data: Expr,
weight: Expr,
@@ -768,6 +801,11 @@ def conv3d_transpose(
)
+@tvm_ffi.register_object("relax.attrs.PadAttrs")
+class PadAttrs(Attrs):
+ """Attributes used in pad operator"""
+
+
def pad(
data: Expr,
pad_width: list[int] | tuple[int, ...],
@@ -815,6 +853,11 @@ def pad(
)
+@tvm_ffi.register_object("relax.attrs.PixelShuffleAttrs")
+class PixelShuffleAttrs(Attrs):
+ """Attributes used in pixel_shuffle operator"""
+
+
def pixel_shuffle(data: Expr, upscale_factor: int, *, ty=None, span=None):
r"""
Pixel Shuffle Operator
@@ -854,6 +897,11 @@ def pixel_shuffle(data: Expr, upscale_factor: int, *, ty=None, span=None):
)
+@tvm_ffi.register_object("relax.attrs.Pool1DAttrs")
+class Pool1DAttrs(Attrs):
+ """Attributes for nn.max_pool1d and nn.avg_pool1d"""
+
+
def max_pool1d(
data: Expr,
pool_size: int | tuple[int, int] = (1,),
@@ -954,6 +1002,11 @@ def max_pool1d(
)
+@tvm_ffi.register_object("relax.attrs.Pool2DAttrs")
+class Pool2DAttrs(Attrs):
+ """Attributes for nn.max_pool2d"""
+
+
def max_pool2d(
data: Expr,
pool_size: int | tuple[int, int] = (1, 1),
@@ -1069,6 +1122,11 @@ def max_pool2d(
)
+@tvm_ffi.register_object("relax.attrs.Pool3DAttrs")
+class Pool3DAttrs(Attrs):
+ """Attributes for nn.max_pool3d and nn.avg_pool3d"""
+
+
def max_pool3d(
data: Expr,
pool_size: int | tuple[int, int] = (1, 1, 1),
@@ -1495,6 +1553,11 @@ def avg_pool3d(
)
+@tvm_ffi.register_object("relax.attrs.AdaptivePool1DAttrs")
+class AdaptivePool1DAttrs(Attrs):
+ """Attributes for 1d adaptive pool operator"""
+
+
def adaptive_avg_pool1d(
data: Expr,
output_size: int | tuple[int] | None = None,
@@ -1564,6 +1627,11 @@ def adaptive_avg_pool1d(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.AdaptivePool2DAttrs")
+class AdaptivePool2DAttrs(Attrs):
+ """Attributes for 2d adaptive pool operator"""
+
+
def adaptive_avg_pool2d(
data: Expr,
output_size: int | tuple[int, int] | None = None,
@@ -1638,6 +1706,11 @@ def adaptive_avg_pool2d(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.AdaptivePool3DAttrs")
+class AdaptivePool3DAttrs(Attrs):
+ """Attributes for 3d adaptive pool operator"""
+
+
def adaptive_avg_pool3d(
data: Expr,
output_size: int | tuple[int, int] | None = None,
@@ -1748,6 +1821,11 @@ def relu6(data: Expr) -> Expr:
return relax.op.clip(data, 0, 6)
+@tvm_ffi.register_object("relax.attrs.LeakyReluAttrs")
+class LeakyReluAttrs(Attrs):
+ """Attributes used in leaky_relu operator"""
+
+
def leakyrelu(data: Expr, alpha: float = 0.01, *, ty_args=None, ty=None, span=None) -> Expr:
"""Rectified linear unit.
@@ -1873,6 +1951,11 @@ def silu(data: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
return _Call("relax.nn.silu", [data], ty_args=ty_args, ty=ty, span=span) # type: ignore
+@tvm_ffi.register_object("relax.attrs.SoftmaxAttrs")
+class SoftmaxAttrs(Attrs):
+ """Attributes for nn.softmax"""
+
+
def softmax(data: Expr, axis: int = -1, *, ty=None, span=None) -> Expr:
r"""Computes softmax.
@@ -1906,6 +1989,11 @@ def softmax(data: Expr, axis: int = -1, *, ty=None, span=None) -> Expr:
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.SoftplusAttrs")
+class SoftplusAttrs(Attrs):
+ """Attributes used in softplus operator"""
+
+
def softplus(
data: Expr,
beta: float = 1.0,
@@ -1980,6 +2068,11 @@ def log_softmax(data: Expr, axis: int = -1, *, ty=None, span=None) -> Expr:
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.PReluAttrs")
+class PReluAttrs(Attrs):
+ """Attributes used in prelu operator"""
+
+
def prelu(data: Expr, alpha: Expr, axis: int = 1, *, ty=None, span=None) -> Expr:
r"""Parametric Rectified Linear Unit (PReLU).
@@ -2012,6 +2105,11 @@ def prelu(data: Expr, alpha: Expr, axis: int = 1, *, ty=None, span=None) -> Expr
)
+@tvm_ffi.register_object("relax.attrs.BatchNormAttrs")
+class BatchNormAttrs(Attrs):
+ """Attributes used in batch_norm operator"""
+
+
def batch_norm(
data: Expr,
gamma: Expr,
@@ -2142,6 +2240,11 @@ def batch_norm(
)
+@tvm_ffi.register_object("relax.attrs.LayerNormAttrs")
+class LayerNormAttrs(Attrs):
+ """Attributes used in layer_norm operator"""
+
+
def layer_norm(
data: Expr,
gamma: Expr,
@@ -2214,6 +2317,11 @@ def layer_norm(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.GroupNormAttrs")
+class GroupNormAttrs(Attrs):
+ """Attributes used in group_norm operator"""
+
+
def group_norm(
data: Expr,
gamma: Expr,
@@ -2287,6 +2395,11 @@ def group_norm(
)
+@tvm_ffi.register_object("relax.attrs.InstanceNormAttrs")
+class InstanceNormAttrs(Attrs):
+ """Attributes used in instance_norm operator"""
+
+
def instance_norm(
data: Expr,
gamma: Expr,
@@ -2349,6 +2462,11 @@ def instance_norm(
)
+@tvm_ffi.register_object("relax.attrs.RMSNormAttrs")
+class RMSNormAttrs(Attrs):
+ """Attributes used in rms_norm operator"""
+
+
def rms_norm(
data: Expr,
weight: Expr,
@@ -2398,6 +2516,11 @@ def rms_norm(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.DropoutAttrs")
+class DropoutAttrs(Attrs):
+ """Attributes for dropout operator"""
+
+
def dropout(data: Expr, rate: float = 0.5, *, ty=None, span=None) -> Expr:
"""Applies the dropout operation to the input tensor.
@@ -2464,6 +2587,11 @@ def cross_entropy_with_logits(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.NLLLossAttrs")
+class NLLLossAttrs(Attrs):
+ """Attributes used in nll_loss operator"""
+
+
def nll_loss(
predictions: Expr,
targets: Expr,
@@ -2524,6 +2652,11 @@ def nll_loss(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.AttentionAttrs")
+class AttentionAttrs(Attrs):
+ """Attributes used in attention operator"""
+
+
def attention(
query: Expr,
key: Expr,
diff --git a/python/tvm/relax/op/op_attrs.py b/python/tvm/relax/op/op_attrs.py
deleted file mode 100644
index b879a19b46..0000000000
--- a/python/tvm/relax/op/op_attrs.py
+++ /dev/null
@@ -1,401 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""The attributes node used for Relax operators"""
-
-import tvm_ffi
-
-from tvm.ir import Attrs
-
-
-@tvm_ffi.register_object("relax.attrs.CallTIRWithGradAttrs")
-class CallTIRWithGradAttrs(Attrs):
- """Attributes used in call_tir_with_grad operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.InitAttrs")
-class InitAttrs(Attrs):
- """Attributes used in full/full_like, ones/ones_like, and zeros/zeros_like operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.TriluAttrs")
-class TriluAttrs(Attrs):
- """Attributes used in tril and triu operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.AstypeAttrs")
-class AstypeAttrs(Attrs):
- """Attributes used in astype operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.TakeAttrs")
-class TakeAttrs(Attrs):
- """Attributes used in take operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.StridedSliceAttrs")
-class StridedSliceAttrs(Attrs):
- """Attributes used in strided_slice operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.MatmulAttrs")
-class MatmulAttrs(Attrs):
- """Attributes for matmul operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.Conv2DAttrs")
-class Conv2DAttrs(Attrs):
- """Attributes for nn.conv2d"""
-
-
-@tvm_ffi.register_object("relax.attrs.Conv3DAttrs")
-class Conv3DAttrs(Attrs):
- """Attributes for nn.conv3d"""
-
-
-@tvm_ffi.register_object("relax.attrs.Conv2DTransposeAttrs")
-class Conv2DTransposeAttrs(Attrs):
- """Attributes for nn.conv2d_transpose"""
-
-
-@tvm_ffi.register_object("relax.attrs.Conv3DTransposeAttrs")
-class Conv3DTransposeAttrs(Attrs):
- """Attributes for nn.conv3d_transpose"""
-
-
-@tvm_ffi.register_object("relax.attrs.Pool2DAttrs")
-class Pool2DAttrs(Attrs):
- """Attributes for nn.max_pool2d"""
-
-
-@tvm_ffi.register_object("relax.attrs.AdaptivePool2DAttrs")
-class AdaptivePool2DAttrs(Attrs):
- """Attributes for 2d adaptive pool operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.SoftmaxAttrs")
-class SoftmaxAttrs(Attrs):
- """Attributes for nn.softmax"""
-
-
-@tvm_ffi.register_object("relax.attrs.BatchNormAttrs")
-class BatchNormAttrs(Attrs):
- """Attributes used in batch_norm operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.LayerNormAttrs")
-class LayerNormAttrs(Attrs):
- """Attributes used in layer_norm operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.InstanceNormAttrs")
-class InstanceNormAttrs(Attrs):
- """Attributes used in instance_norm operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.DropoutAttrs")
-class DropoutAttrs(Attrs):
- """Attributes for dropout operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.StatisticalAttrs")
-class StatisticalAttrs(Attrs):
- """Attributes used in statistical operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.ConcatAttrs")
-class ConcatAttrs(Attrs):
- """Attributes for concat operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.ExpandDimsAttrs")
-class ExpandDimsAttrs(Attrs):
- """Attributes for expand_dims operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.PermuteDimsAttrs")
-class PermuteDimsAttrs(Attrs):
- """Attributes for permute_dims operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.SortAttrs")
-class SortAttrs(Attrs):
- """Attributes for sort operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.ArgsortAttrs")
-class ArgsortAttrs(Attrs):
- """Attributes for argsort operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.SplitAttrs")
-class SplitAttrs(Attrs):
- """Attributes used in split operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.SqueezeAttrs")
-class SqueezeAttrs(Attrs):
- """Attributes for squeeze operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.StackAttrs")
-class StackAttrs(Attrs):
- """Attributes for concat operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.IndexPutAttrs")
-class IndexPutAttrs(Attrs):
- """Attributes for index_put operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.LayoutTransformAttrs")
-class LayoutTransformAttrs(Attrs):
- """Attributes used in layout_transform operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.Resize2DAttrs")
-class Resize2DAttrs(Attrs):
- """Attributes used in image resize2d operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.ArgmaxArgminAttrs")
-class ArgmaxArgminAttrs(Attrs):
- """Attributes for argmax/argmin operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.RepeatAttrs")
-class RepeatAttrs(Attrs):
- """Attributes for repeat operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.TileAttrs")
-class TileAttrs(Attrs):
- """Attributes for tile operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.ScanopAttrs")
-class ScanopAttrs(Attrs):
- """Attributes for scan operators"""
-
-
-@tvm_ffi.register_object("relax.attrs.TopKAttrs")
-class TopKAttrs(Attrs):
- """Attributes for topk operators"""
-
-
-@tvm_ffi.register_object("relax.attrs.EinsumAttrs")
-class EinsumAttrs(Attrs):
- """Attributes for einsum operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.FlipAttrs")
-class FlipAttrs(Attrs):
- """Attributes for flip operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.ReverseSequenceAttrs")
-class ReverseSequenceAttrs(Attrs):
- """Attributes for reverse_sequence operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.PadAttrs")
-class PadAttrs(Attrs):
- """Attributes used in pad operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.MultinomialFromUniformAttrs")
-class MultinomialFromUniformAttrs(Attrs):
- """Attributes for multinomial_from_uniform operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.CallInplacePackedAttrs")
-class CallInplacePackedAttrs(Attrs):
- """Attributes used in call_inplace_packed operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.CallTIRInplaceAttrs")
-class CallTIRInplaceAttrs(Attrs):
- """Attributes used in call_tir_inplace operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.ToVDeviceAttrs")
-class ToVDeviceAttrs(Attrs):
- """Attributes used in to_vdevice operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.HintOnDeviceAttrs")
-class HintOnDeviceAttrs(Attrs):
- """Attributes used in hint_on_device operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.ScatterCollectiveAttrs")
-class ScatterCollectiveAttrs(Attrs):
- """Attributes used in scatter collective operators"""
-
-
-@tvm_ffi.register_object("relax.attrs.AttentionAttrs")
-class AttentionAttrs(Attrs):
- """Attributes used in attention operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.AllClassNonMaximumSuppressionAttrs")
-class AllClassNonMaximumSuppressionAttrs(Attrs):
- """Attributes for vision.all_class_non_max_suppression"""
-
-
-@tvm_ffi.register_object("relax.attrs.GetValidCountsAttrs")
-class GetValidCountsAttrs(Attrs):
- """Attributes for vision.get_valid_counts"""
-
-
-@tvm_ffi.register_object("relax.attrs.NonMaximumSuppressionAttrs")
-class NonMaximumSuppressionAttrs(Attrs):
- """Attributes for vision.non_max_suppression"""
-
-
-@tvm_ffi.register_object("relax.attrs.ROIAlignAttrs")
-class ROIAlignAttrs(Attrs):
- """Attributes for vision.roi_align"""
-
-
-@tvm_ffi.register_object("relax.attrs.ROIPoolAttrs")
-class ROIPoolAttrs(Attrs):
- """Attributes for vision.roi_pool"""
-
-
-@tvm_ffi.register_object("relax.attrs.MultiboxTransformLocAttrs")
-class MultiboxTransformLocAttrs(Attrs):
- """Attributes for vision.multibox_transform_loc"""
-
-
-@tvm_ffi.register_object("relax.attrs.Conv1DAttrs")
-class Conv1DAttrs(Attrs):
- """Attributes for nn.conv1d"""
-
-
-@tvm_ffi.register_object("relax.attrs.Conv1DTransposeAttrs")
-class Conv1DTransposeAttrs(Attrs):
- """Attributes for nn.conv1d_transpose"""
-
-
-@tvm_ffi.register_object("relax.attrs.Pool1DAttrs")
-class Pool1DAttrs(Attrs):
- """Attributes for nn.max_pool1d and nn.avg_pool1d"""
-
-
-@tvm_ffi.register_object("relax.attrs.Pool3DAttrs")
-class Pool3DAttrs(Attrs):
- """Attributes for nn.max_pool3d and nn.avg_pool3d"""
-
-
-@tvm_ffi.register_object("relax.attrs.AdaptivePool1DAttrs")
-class AdaptivePool1DAttrs(Attrs):
- """Attributes for 1d adaptive pool operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.AdaptivePool3DAttrs")
-class AdaptivePool3DAttrs(Attrs):
- """Attributes for 3d adaptive pool operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.LeakyReluAttrs")
-class LeakyReluAttrs(Attrs):
- """Attributes used in leaky_relu operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.SoftplusAttrs")
-class SoftplusAttrs(Attrs):
- """Attributes used in softplus operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.PReluAttrs")
-class PReluAttrs(Attrs):
- """Attributes used in prelu operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.PixelShuffleAttrs")
-class PixelShuffleAttrs(Attrs):
- """Attributes used in pixel_shuffle operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.GroupNormAttrs")
-class GroupNormAttrs(Attrs):
- """Attributes used in group_norm operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.RMSNormAttrs")
-class RMSNormAttrs(Attrs):
- """Attributes used in rms_norm operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.NLLLossAttrs")
-class NLLLossAttrs(Attrs):
- """Attributes used in nll_loss operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.AllReduceAttrs")
-class AllReduceAttrs(Attrs):
- """Attributes used in allreduce operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.AllGatherAttrs")
-class AllGatherAttrs(Attrs):
- """Attributes used in allgather operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.WrapParamAttrs")
-class WrapParamAttrs(Attrs):
- """Attributes used in wrap_param operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.QuantizeAttrs")
-class QuantizeAttrs(Attrs):
- """Attributes used in quantize/dequantize operators"""
-
-
-@tvm_ffi.register_object("relax.attrs.GatherElementsAttrs")
-class GatherElementsAttrs(Attrs):
- """Attributes for gather_elements operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.GatherNDAttrs")
-class GatherNDAttrs(Attrs):
- """Attributes for gather_nd operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.MeshgridAttrs")
-class MeshgridAttrs(Attrs):
- """Attributes for meshgrid operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.ScatterElementsAttrs")
-class ScatterElementsAttrs(Attrs):
- """Attributes for scatter_elements operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.ScatterNDAttrs")
-class ScatterNDAttrs(Attrs):
- """Attributes for scatter_nd operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.SliceScatterAttrs")
-class SliceScatterAttrs(Attrs):
- """Attributes for slice_scatter operator"""
-
-
-@tvm_ffi.register_object("relax.attrs.OneHotAttrs")
-class OneHotAttrs(Attrs):
- """Attributes for one_hot operator"""
diff --git a/python/tvm/relax/op/qdq.py b/python/tvm/relax/op/qdq.py
index a883ec52d7..9381b05504 100644
--- a/python/tvm/relax/op/qdq.py
+++ b/python/tvm/relax/op/qdq.py
@@ -16,6 +16,9 @@
# under the License.
"""Relax quantize/dequantize operators"""
+import tvm_ffi
+
+from tvm.ir import Attrs
from tvm.ir import Call as _Call
from tvm.ir import PrimType as _PrimType
from tvm.ir.attrs import make_node as _make_attrs
@@ -23,6 +26,11 @@ from tvm.ir.attrs import make_node as _make_attrs
from ..expr import Expr
+@tvm_ffi.register_object("relax.attrs.QuantizeAttrs")
+class QuantizeAttrs(Attrs):
+ """Attributes used in quantize/dequantize operators"""
+
+
def quantize(
data: Expr,
scale: Expr,
diff --git a/python/tvm/relax/op/sampling.py b/python/tvm/relax/op/sampling.py
index 4f973676ae..31cf4c4aad 100644
--- a/python/tvm/relax/op/sampling.py
+++ b/python/tvm/relax/op/sampling.py
@@ -16,12 +16,20 @@
# under the License.
"""Sampling operators."""
+import tvm_ffi
+
+from tvm.ir import Attrs
from tvm.ir import Call as _Call
from tvm.ir.attrs import make_node as _make_attrs
from ..expr import Expr
+@tvm_ffi.register_object("relax.attrs.MultinomialFromUniformAttrs")
+class MultinomialFromUniformAttrs(Attrs):
+ """Attributes for multinomial_from_uniform operator"""
+
+
def multinomial_from_uniform(
prob: Expr,
uniform_sample: Expr,
diff --git a/python/tvm/relax/op/search.py b/python/tvm/relax/op/search.py
index 85ce91db2d..f1e418fb09 100644
--- a/python/tvm/relax/op/search.py
+++ b/python/tvm/relax/op/search.py
@@ -17,6 +17,9 @@
# pylint: disable=invalid-name
"""Search operators."""
+import tvm_ffi
+
+from tvm.ir import Attrs
from tvm.ir import Call as _Call
from tvm.ir.attrs import make_node as _make_attrs
@@ -53,6 +56,11 @@ def where(condition: Expr, x1: Expr, x2: Expr, *, ty=None, span=None) -> Expr:
return _Call("relax.where", [condition, x1, x2], ty=ty, span=span) # type: ignore
+@tvm_ffi.register_object("relax.attrs.ArgmaxArgminAttrs")
+class ArgmaxArgminAttrs(Attrs):
+ """Attributes for argmax/argmin operator"""
+
+
def argmax(
x: Expr,
axis: int | None = None,
diff --git a/python/tvm/relax/op/sorting.py b/python/tvm/relax/op/sorting.py
index f892a43947..d543552907 100644
--- a/python/tvm/relax/op/sorting.py
+++ b/python/tvm/relax/op/sorting.py
@@ -16,13 +16,20 @@
# under the License.
"""Sortings operators."""
+import tvm_ffi
+
+from tvm.ir import Attrs, GenericConst
from tvm.ir import Call as _Call
-from tvm.ir import GenericConst
from tvm.ir.attrs import make_node as _make_attrs
from ..expr import Expr
+@tvm_ffi.register_object("relax.attrs.SortAttrs")
+class SortAttrs(Attrs):
+ """Attributes for sort operator"""
+
+
def sort(
x: Expr,
axis: int = -1,
@@ -61,6 +68,11 @@ def sort(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.ArgsortAttrs")
+class ArgsortAttrs(Attrs):
+ """Attributes for argsort operator"""
+
+
def argsort(
data: Expr,
axis: int = -1,
@@ -103,6 +115,11 @@ def argsort(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.TopKAttrs")
+class TopKAttrs(Attrs):
+ """Attributes for topk operators"""
+
+
def topk(
data: Expr,
k: int = 1,
diff --git a/python/tvm/relax/op/statistical.py b/python/tvm/relax/op/statistical.py
index 80fa0ea115..1a7a6b5b82 100644
--- a/python/tvm/relax/op/statistical.py
+++ b/python/tvm/relax/op/statistical.py
@@ -17,9 +17,11 @@
# pylint: disable=redefined-builtin
"""Statistical operators."""
+import tvm_ffi
+
from tvm import DataType
+from tvm.ir import Attrs, PrimType
from tvm.ir import Call as _Call
-from tvm.ir import PrimType
from tvm.ir.attrs import make_node as _make_attrs
from ..expr import Expr
@@ -29,6 +31,11 @@ def _raw_dtype(dtype):
return dtype.dtype if isinstance(dtype, PrimType) else dtype
+@tvm_ffi.register_object("relax.attrs.StatisticalAttrs")
+class StatisticalAttrs(Attrs):
+ """Attributes used in statistical operator"""
+
+
def max(
x: Expr,
axis: int | list[int] | None = None,
@@ -275,6 +282,11 @@ def sum(
) # type: ignore
+@tvm_ffi.register_object("relax.attrs.ScanopAttrs")
+class ScanopAttrs(Attrs):
+ """Attributes for scan operators"""
+
+
def cumprod(
data: Expr,
axis: int | None = None,
diff --git a/python/tvm/relax/op/vision/multibox_transform_loc.py b/python/tvm/relax/op/vision/multibox_transform_loc.py
index dcbc37d900..65d47c0b72 100644
--- a/python/tvm/relax/op/vision/multibox_transform_loc.py
+++ b/python/tvm/relax/op/vision/multibox_transform_loc.py
@@ -16,10 +16,18 @@
# under the License.
"""Multibox location transform for object detection."""
+import tvm_ffi
+
+from tvm.ir import Attrs
from tvm.ir import Call as _Call
from tvm.ir.attrs import make_node as _make_attrs
+@tvm_ffi.register_object("relax.attrs.MultiboxTransformLocAttrs")
+class MultiboxTransformLocAttrs(Attrs):
+ """Attributes for vision.multibox_transform_loc"""
+
+
def multibox_transform_loc(
cls_pred,
loc_pred,
diff --git a/python/tvm/relax/op/vision/nms.py b/python/tvm/relax/op/vision/nms.py
index 2715d20bc7..6fe45663b7 100644
--- a/python/tvm/relax/op/vision/nms.py
+++ b/python/tvm/relax/op/vision/nms.py
@@ -16,10 +16,18 @@
# under the License.
"""Non-maximum suppression operators."""
+import tvm_ffi
+
+from tvm.ir import Attrs
from tvm.ir import Call as _Call
from tvm.ir.attrs import make_node as _make_attrs
+@tvm_ffi.register_object("relax.attrs.AllClassNonMaximumSuppressionAttrs")
+class AllClassNonMaximumSuppressionAttrs(Attrs):
+ """Attributes for vision.all_class_non_max_suppression"""
+
+
def all_class_non_max_suppression(
boxes,
scores,
@@ -83,6 +91,11 @@ def all_class_non_max_suppression(
)
+@tvm_ffi.register_object("relax.attrs.GetValidCountsAttrs")
+class GetValidCountsAttrs(Attrs):
+ """Attributes for vision.get_valid_counts"""
+
+
def get_valid_counts(
data,
score_threshold=0,
@@ -131,6 +144,11 @@ def get_valid_counts(
)
+@tvm_ffi.register_object("relax.attrs.NonMaximumSuppressionAttrs")
+class NonMaximumSuppressionAttrs(Attrs):
+ """Attributes for vision.non_max_suppression"""
+
+
def non_max_suppression(
data,
valid_count,
diff --git a/python/tvm/relax/op/vision/roi_align.py b/python/tvm/relax/op/vision/roi_align.py
index 36075342d8..d9c8ebe7a0 100644
--- a/python/tvm/relax/op/vision/roi_align.py
+++ b/python/tvm/relax/op/vision/roi_align.py
@@ -16,12 +16,20 @@
# under the License.
"""ROI Align operator"""
+import tvm_ffi
+
+from tvm.ir import Attrs
from tvm.ir import Call as _Call
from tvm.ir.attrs import make_node as _make_attrs
from ..base import Expr
+@tvm_ffi.register_object("relax.attrs.ROIAlignAttrs")
+class ROIAlignAttrs(Attrs):
+ """Attributes for vision.roi_align"""
+
+
def roi_align(
data: Expr,
rois: Expr,
diff --git a/python/tvm/relax/op/vision/roi_pool.py b/python/tvm/relax/op/vision/roi_pool.py
index 50f97412dc..fe39e0d9c0 100644
--- a/python/tvm/relax/op/vision/roi_pool.py
+++ b/python/tvm/relax/op/vision/roi_pool.py
@@ -16,12 +16,20 @@
# under the License.
"""ROI Pool operator"""
+import tvm_ffi
+
+from tvm.ir import Attrs
from tvm.ir import Call as _Call
from tvm.ir.attrs import make_node as _make_attrs
from ..base import Expr
+@tvm_ffi.register_object("relax.attrs.ROIPoolAttrs")
+class ROIPoolAttrs(Attrs):
+ """Attributes for vision.roi_pool"""
+
+
def roi_pool(
data: Expr,
rois: Expr,
diff --git a/python/tvm/relax/script/ir_builder/op.py b/python/tvm/relax/script/ir_builder/op.py
index 8813d25933..30a6bcf4e0 100644
--- a/python/tvm/relax/script/ir_builder/op.py
+++ b/python/tvm/relax/script/ir_builder/op.py
@@ -36,6 +36,46 @@ from tvm.ir.prim import _ffi_api as _prim_ffi
from tvm.relax import Call, Expr, ExternFunc
from tvm.relax.global_info import VDevice
from tvm.relax.op import (
+ ArgmaxArgminAttrs,
+ ArgsortAttrs,
+ AstypeAttrs,
+ CallInplacePackedAttrs,
+ CallTIRInplaceAttrs,
+ CallTIRWithGradAttrs,
+ ConcatAttrs,
+ EinsumAttrs,
+ ExpandDimsAttrs,
+ FlipAttrs,
+ GatherElementsAttrs,
+ GatherNDAttrs,
+ HintOnDeviceAttrs,
+ IndexPutAttrs,
+ InitAttrs,
+ LayoutTransformAttrs,
+ MatmulAttrs,
+ MeshgridAttrs,
+ MultinomialFromUniformAttrs,
+ OneHotAttrs,
+ PermuteDimsAttrs,
+ QuantizeAttrs,
+ RepeatAttrs,
+ ReverseSequenceAttrs,
+ ScanopAttrs,
+ ScatterElementsAttrs,
+ ScatterNDAttrs,
+ SliceScatterAttrs,
+ SortAttrs,
+ SplitAttrs,
+ SqueezeAttrs,
+ StackAttrs,
+ StatisticalAttrs,
+ StridedSliceAttrs,
+ TakeAttrs,
+ TileAttrs,
+ TopKAttrs,
+ ToVDeviceAttrs,
+ TriluAttrs,
+ WrapParamAttrs,
abs,
acos,
acosh,
@@ -574,6 +614,46 @@ inspect = _Namespace()
_init_op_api("relax", __name__)
__all__ = [
+ "ArgmaxArgminAttrs",
+ "ArgsortAttrs",
+ "AstypeAttrs",
+ "CallInplacePackedAttrs",
+ "CallTIRInplaceAttrs",
+ "CallTIRWithGradAttrs",
+ "ConcatAttrs",
+ "EinsumAttrs",
+ "ExpandDimsAttrs",
+ "FlipAttrs",
+ "GatherElementsAttrs",
+ "GatherNDAttrs",
+ "HintOnDeviceAttrs",
+ "IndexPutAttrs",
+ "InitAttrs",
+ "LayoutTransformAttrs",
+ "MatmulAttrs",
+ "MeshgridAttrs",
+ "MultinomialFromUniformAttrs",
+ "OneHotAttrs",
+ "PermuteDimsAttrs",
+ "QuantizeAttrs",
+ "RepeatAttrs",
+ "ReverseSequenceAttrs",
+ "ScanopAttrs",
+ "ScatterElementsAttrs",
+ "ScatterNDAttrs",
+ "SliceScatterAttrs",
+ "SortAttrs",
+ "SplitAttrs",
+ "SqueezeAttrs",
+ "StackAttrs",
+ "StatisticalAttrs",
+ "StridedSliceAttrs",
+ "TakeAttrs",
+ "TileAttrs",
+ "ToVDeviceAttrs",
+ "TopKAttrs",
+ "TriluAttrs",
+ "WrapParamAttrs",
"abs",
"acos",
"acosh",
diff --git a/python/tvm/relax/transform/legalize_ops/inspect_op.py b/python/tvm/relax/transform/legalize_ops/inspect_op.py
index 78fe31d1e2..e924edbbf7 100644
--- a/python/tvm/relax/transform/legalize_ops/inspect_op.py
+++ b/python/tvm/relax/transform/legalize_ops/inspect_op.py
@@ -30,7 +30,7 @@ from .common import register_legalize
class TVMStructFieldKind(enum.IntEnum):
- """Equivalent to tvm::tirx::builtin::TVMStructFieldKind
+ """Equivalent to tvm::tirx::TVMStructFieldKind
This does not use `enum.auto()` to define the values, because
`enum.auto()` starts from 1, and this must match the C++
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op/__init__.py
similarity index 85%
rename from python/tvm/tirx/op.py
rename to python/tvm/tirx/op/__init__.py
index 7c7f81e03b..710b2edb92 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op/__init__.py
@@ -27,7 +27,6 @@ import tvm_ffi
from tvm_ffi import Array
import tvm
-import tvm.ir.prim._ffi_api as _prim_ffi_api
from tvm import tirx
from tvm.ir import (
Call,
@@ -43,13 +42,36 @@ from tvm.ir.base import Span
from tvm.ir.prim import clz as clz
from tvm.ir.prim import max_value as max_value
from tvm.ir.prim import min_value as min_value
+from tvm.ir.prim.op import all as all
+from tvm.ir.prim.op import any as any
+from tvm.ir.prim.op import bitwise_and as bitwise_and
+from tvm.ir.prim.op import bitwise_not as bitwise_not
+from tvm.ir.prim.op import bitwise_or as bitwise_or
+from tvm.ir.prim.op import bitwise_xor as bitwise_xor
+from tvm.ir.prim.op import ceil as ceil
+from tvm.ir.prim.op import ceildiv as ceildiv
+from tvm.ir.prim.op import div as div
+from tvm.ir.prim.op import floordiv as floordiv
+from tvm.ir.prim.op import floormod as floormod
+from tvm.ir.prim.op import if_then_else as if_then_else
+from tvm.ir.prim.op import indexdiv as indexdiv
+from tvm.ir.prim.op import indexmod as indexmod
+from tvm.ir.prim.op import likely as likely
+from tvm.ir.prim.op import log2 as log2
+from tvm.ir.prim.op import max as max
+from tvm.ir.prim.op import min as min
+from tvm.ir.prim.op import shift_left as shift_left
+from tvm.ir.prim.op import shift_right as shift_right
+from tvm.ir.prim.op import truncdiv as truncdiv
+from tvm.ir.prim.op import truncmod as truncmod
+from tvm.ir.prim.op import vscale as vscale
from tvm.runtime import const
-from . import _ffi_api
-from .buffer import buffer_data, is_tensor_var
-from .expr import ExprOp, IntImm
-from .expr import TensorLoad as _make_tensor_load
-from .type import TensorMapType
+from .. import _ffi_api
+from ..buffer import buffer_data, is_tensor_var
+from ..expr import ExprOp, IntImm
+from ..expr import TensorLoad as _make_tensor_load
+from ..type import TensorMapType
tir = tirx # alias for backward compat with upstream tir.convert() calls
@@ -1419,59 +1441,6 @@ def dp4a(vec1, vec2, acc=0, **kwargs):
return Call("tirx.dp4a", [vec1, vec2, acc], **kwargs)
-def any(*args, span=None):
- """Create a new experssion of the union of all conditions in the arguments
-
- Parameters
- ----------
- args : list
- List of symbolic boolean expressions
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- expr: Expr
- Expression
- """
- if not args:
- raise ValueError("Any must take at least 1 argument")
- if len(args) == 1:
- return args[0]
- val = _prim_ffi_api._OpOr(args[0], args[1], span) # type: ignore
- for i in range(2, len(args)):
- val = _prim_ffi_api._OpOr(val, args[i], span) # type: ignore
- return val
-
-
-def all(*args, span=None):
- """Create a new expression of the intersection of all conditions in the
- arguments
-
- Parameters
- ----------
- args : list
- List of symbolic boolean expressions
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- expr: Expr
- Expression
- """
- if not args:
- raise ValueError("Any must take at least 1 argument")
- if len(args) == 1:
- return args[0]
- val = _prim_ffi_api._OpAnd(args[0], args[1], span) # type: ignore
- for i in range(2, len(args)):
- val = _prim_ffi_api._OpAnd(val, args[i], span) # type: ignore
- return val
-
-
def infinity(dtype: str, span: Span | None = None) -> Any:
"""infinity value of dtype
@@ -1661,23 +1630,6 @@ def log(x, *, ty=None, span=None):
return call_intrin(ty, "tirx.log", x, span=span)
-def log2(x, *, ty=None, span=None):
- """Take log2 of input x.
-
- Parameters
- ----------
- x : Expr
- Input argument.
-
- Returns
- -------
- y : Expr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(ty, "prim.log2", x, span=span)
-
-
def log10(x, *, ty=None, span=None):
"""Take log10 of input x.
@@ -1973,25 +1925,6 @@ def floor(x: ExprWithOp, span=None):
return _ffi_api.floor(x, span) # type: ignore
-def ceil(x, span=None):
- """Take ceil of float input x.
-
- Parameters
- ----------
- x : Expr
- Input argument.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- y : Expr
- The result.
- """
- return _prim_ffi_api.ceil(x, span) # type: ignore
-
-
def trunc(x, span=None):
"""Get truncated value of the input.
@@ -2033,91 +1966,6 @@ def abs(x, span=None):
return _ffi_api.abs(x, span) # type: ignore
-def bitwise_and(x, y, span=None):
- """Take bitwise and of two values
-
- Parameters
- ----------
- x : Expr
- Left operand
-
- y : Expr
- Right operand
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- res : Expr
- The result.
- """
- return _prim_ffi_api.bitwise_and(x, y, span)
-
-
-def bitwise_not(x, span=None):
- """Take bitwise not of input value
-
- Parameters
- ----------
- x : Expr
- Input operand
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- res : Expr
- The result.
- """
- return _prim_ffi_api.bitwise_not(x, span)
-
-
-def bitwise_or(x, y, span=None):
- """Take bitwise or of two values
-
- Parameters
- ----------
- x : Expr
- Left operand
-
- y : Expr
- Right operand
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- res : Expr
- The result.
- """
- return _prim_ffi_api.bitwise_or(x, y, span)
-
-
-def bitwise_xor(x, y, span=None):
- """Take bitwise xor of two values
-
- Parameters
- ----------
- x : Expr
- Left operand
-
- y : Expr
- Right operand
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- res : Expr
- The result.
- """
- return _prim_ffi_api.bitwise_xor(x, y, span)
-
-
def round(x, span=None):
"""Round elements of the array to the nearest integer.
@@ -2247,26 +2095,6 @@ def ldexp(x1, x2, *, ty=None, span=None):
return call_intrin(ty, "tirx.ldexp", x1, x2, span=span) # type: ignore
-def likely(cond, span=None):
- """Mark condition as likely.
-
- Parameters
- ----------
-
- cond : Expr
- Input argument.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- y : Expr
- The marked expression.
- """
- return _prim_ffi_api.likely(cond, span) # type: ignore
-
-
def filter(var, pred, *, span=None, ty=None): # pylint: disable=redefined-builtin
"""Thread-set filter escape hatch.
@@ -2515,44 +2343,6 @@ def q_multiply_shift_per_axis(
)
-def shift_left(x, y, span=None):
- """Return the result of x left shifted by y bits.
-
- Parameters
- ----------
- x : Expr
- Input argument.
-
- y : Expr
- Input argument.
-
- Returns
- -------
- z : Expr
- The result.
- """
- return _prim_ffi_api.left_shift(x, y, span)
-
-
-def shift_right(x, y, span=None):
- """Return the result of x right shifted by y bits.
-
- Parameters
- ----------
- x : Expr
- Input argument.
-
- y : Expr
- Input argument.
-
- Returns
- -------
- z : Expr
- The result.
- """
- return _prim_ffi_api.right_shift(x, y, span)
-
-
def fmod(x, y, *, ty=None, span=None):
"""Return the remainder of x divided by y with the same sign as x.
@@ -2573,194 +2363,6 @@ def fmod(x, y, *, ty=None, span=None):
return call_intrin(ty, "tirx.fmod", x, y, span=span)
-def if_then_else(cond, t, f, span=None):
- """Conditional selection expression.
-
- Parameters
- ----------
- cond : Expr
- The condition
-
- t : Expr
- The result expression if cond is true.
-
- f : Expr
- The result expression if cond is false.
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- result : Node
- The result of conditional expression.
-
- Note
- ----
- Unlike Select, if_then_else will not execute
- the branch that does not satisfy the condition.
- You can use it to guard against out of bound access.
- Unlike Select, if_then_else cannot be vectorized
- if some lanes in the vector have different conditions.
- """
- return _prim_ffi_api._OpIfThenElse(cond, t, f, span) # type: ignore
-
-
-def div(a, b, span=None):
- """Compute a / b as in C/C++ semantics.
-
- Parameters
- ----------
- a : Expr
- The left hand operand, known to be non-negative.
-
- b : Expr
- The right hand operand, known to be non-negative.
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- res : Expr
- The result expression.
- Note
- ----
- When operands are integers, returns truncdiv(a, b, span).
- """
- return _prim_ffi_api._OpDiv(a, b, span) # type: ignore
-
-
-def indexdiv(a, b, span=None):
- """Compute floor(a / b) where a and b are non-negative.
-
- Parameters
- ----------
- a : Expr
- The left hand operand, known to be non-negative.
-
- b : Expr
- The right hand operand, known to be non-negative.
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- res : Expr
- The result expression.
-
- Note
- ----
- Use this function to split non-negative indices.
- This function may take advantage of operands'
- non-negativeness.
- """
- return _prim_ffi_api._OpIndexDiv(a, b, span) # type: ignore
-
-
-def indexmod(a, b, span=None):
- """Compute the remainder of indexdiv. a and b are non-negative.
-
- Parameters
- ----------
- a : Expr
- The left hand operand, known to be non-negative.
-
- b : Expr
- The right hand operand, known to be non-negative.
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- res : Expr
- The result expression.
-
- Note
- ----
- Use this function to split non-negative indices.
- This function may take advantage of operands'
- non-negativeness.
- """
- return _prim_ffi_api._OpIndexMod(a, b, span) # type: ignore
-
-
-def truncdiv(a, b, span=None):
- """Compute the truncdiv of two expressions.
-
- Parameters
- ----------
- a : Expr
- The left hand operand
-
- b : Expr
- The right hand operand
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- res : Expr
- The result expression.
-
- Note
- ----
- This is the default integer division behavior in C.
- """
- return _prim_ffi_api._OpTruncDiv(a, b, span) # type: ignore
-
-
-def truncmod(a, b, span=None):
- """Compute the truncmod of two expressions.
-
- Parameters
- ----------
- a : Expr
- The left hand operand
-
- b : Expr
- The right hand operand
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- res : Expr
- The result expression.
-
- Note
- ----
- This is the default integer division behavior in C.
- """
- return _prim_ffi_api._OpTruncMod(a, b, span) # type: ignore
-
-
-def floordiv(a, b, span=None):
- """Compute the floordiv of two expressions.
-
- Parameters
- ----------
- a : Expr
- The left hand operand
-
- b : Expr
- The right hand operand
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- res : Expr
- The result expression.
- """
- return _prim_ffi_api._OpFloorDiv(a, b, span) # type: ignore
-
-
def logaddexp(a, b, span=None):
"""Compute the logaddexp of two expressions.
@@ -2783,48 +2385,6 @@ def logaddexp(a, b, span=None):
return _ffi_api._OpLogAddExp(a, b, span) # type: ignore
-def floormod(a, b, span=None):
- """Compute the floormod of two expressions.
-
- Parameters
- ----------
- a : Expr
- The left hand operand
-
- b : Expr
- The right hand operand
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- res : Expr
- The result expression.
- """
- return _prim_ffi_api._OpFloorMod(a, b, span) # type: ignore
-
-
-def ceildiv(lhs, rhs, span=None):
- """Generic ceildiv operator.
-
- Parameters
- ----------
- lhs : object
- The left operand.
- rhs : object
- The right operand.
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- op : tvm.Expr
- The result Expr of ceildiv operaton.
- """
- return _prim_ffi_api._OpCeilDiv(lhs, rhs, span) # type: ignore
-
-
def TVMBackendAllocWorkspace(
device_type,
device_id,
@@ -2900,17 +2460,6 @@ def TVMBackendFreeWorkspace(device_type, device_id, ptr, *, ty=None, span=None):
)
-def vscale(*, ty=None, span=None):
- """Get the target's vscale value. It will be lowered to llvm.vscale intrinsic
- (https://llvm.org/docs/LangRef.html#llvm-vscale-intrinsic)
- Returns
- -------
- call : Expr
- Call to the vscale intrinsic
- """
- return call_intrin(ty, "prim.vscale", span=span)
-
-
def get_active_lane_mask(dtype, base, limit):
"""
Calculate a predicate mask given an upper bound (limit) and a current value (base).
@@ -3013,17 +2562,6 @@ def ignore_loop_partition(predicate, *, ty=None, span=None) -> Expr:
return call_intrin(ty, "tirx.ignore_loop_partition", predicate, span=span)
-# pylint: disable=unnecessary-lambda
-def min(a, b, span=None):
- """Elementwise minimum of two primitive expressions."""
- return _prim_ffi_api._OpMin(a, b, span)
-
-
-def max(a, b, span=None):
- """Elementwise maximum of two primitive expressions."""
- return _prim_ffi_api._OpMax(a, b, span)
-
-
def tvm_load_matrix_sync(
fragment,
m,
@@ -3312,3 +2850,13 @@ def tvm_store_matrix_sync(
def thread_return(*, ty=None, span=None):
"""Return from the current GPU thread without a function value."""
return call_intrin(ty, "tirx.thread_return", span=span)
+
+
+def __getattr__(name):
+ # Tile classes resolve registered Ops at import time. Load the family only
+ # when requested, after the core TIRx types are available.
+ if name == "tile":
+ from importlib import import_module # pylint: disable=import-outside-toplevel
+
+ return import_module(".tile", __name__)
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/python/tvm/tirx/operator/tile_primitive/ops.py b/python/tvm/tirx/op/tile.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/ops.py
rename to python/tvm/tirx/op/tile.py
diff --git a/python/tvm/tirx/operator/tile_primitive/__init__.py b/python/tvm/tirx/operator/tile_primitive/__init__.py
index e8e21c3c52..3b1f2778da 100644
--- a/python/tvm/tirx/operator/tile_primitive/__init__.py
+++ b/python/tvm/tirx/operator/tile_primitive/__init__.py
@@ -18,9 +18,9 @@
# ruff: noqa: I001
# Op class declarations (Add, Sub, Gemm, ...) — must run first so their
-# `op = Op.get("tirx.<name>")` registrations execute before any dispatch
+# `op = Op.get("tirx.tile.<name>")` registrations execute before any dispatch
# code refers to the same ops.
-from .ops import *
+from ...op.tile import *
# Dispatch infrastructure. Per-backend schedule registrations are loaded via
# ``tvm.backend.load(<name>)``.
diff --git a/python/tvm/tirx/script/ir_builder/op.py b/python/tvm/tirx/script/ir_builder/op.py
index 4325298115..25fdbe6e68 100644
--- a/python/tvm/tirx/script/ir_builder/op.py
+++ b/python/tvm/tirx/script/ir_builder/op.py
@@ -78,6 +78,7 @@ from tvm.tirx.expr import (
Shuffle,
Sub,
)
+from tvm.tirx.op import CallFFIKernelAttr, TensorMapEncodeTiledAttr
from . import _ffi_api
from .external_kernel import call_kernel
@@ -806,6 +807,7 @@ __all__ = [
"Broadcast",
"Call",
"CallEffectKind",
+ "CallFFIKernelAttr",
"Cast",
"CommReducer",
"Div",
@@ -827,6 +829,7 @@ __all__ = [
"Sub",
"TVMBackendAllocWorkspace",
"TVMBackendFreeWorkspace",
+ "TensorMapEncodeTiledAttr",
"abs",
"acos",
"acosh",
diff --git a/python/tvm/tirx/script/ir_builder/parser_protocol.py b/python/tvm/tirx/script/ir_builder/parser_protocol.py
index abe4d4b1c9..66e90f6132 100644
--- a/python/tvm/tirx/script/ir_builder/parser_protocol.py
+++ b/python/tvm/tirx/script/ir_builder/parser_protocol.py
@@ -148,16 +148,6 @@ def func_attr(attrs: dict[str, Any]) -> None:
_ffi_api.FuncAttrs(attrs) # type: ignore[attr-defined] # pylint: disable=no-member
-def device_entry() -> frame.RegionFrame:
- """Mark a device-entry region containing scope definitions.
-
- Use a flat ``T.device_entry()`` to scope the remaining statements in the
- enclosing body, or ``with T.device_entry():`` for an explicit boundary.
- Statements before the region remain host code.
- """
- return region("tirx.device_entry", [])
-
-
def check_well_formed_(function: _tir.Function) -> None:
"""Validate a completed TIRx function."""
try:
@@ -936,6 +926,16 @@ def region(
return _ffi_api.Region(op, args, body_params, attrs)
+def device_entry() -> frame.RegionFrame:
+ """Mark a device-entry region containing scope definitions.
+
+ Use a flat ``T.device_entry()`` to scope the remaining statements in the
+ enclosing body, or ``with T.device_entry():`` for an explicit boundary.
+ Statements before the region remain host code.
+ """
+ return region("tirx.device_entry", [])
+
+
def device_context(device_type: Expr, device_id: Expr) -> frame.RegionFrame:
"""Supply lexical device context for allocation and packed-call lowering.
diff --git a/python/tvm/tirx/tile_primitive.py b/python/tvm/tirx/tile_primitive.py
index de321acac8..13685dd6da 100644
--- a/python/tvm/tirx/tile_primitive.py
+++ b/python/tvm/tirx/tile_primitive.py
@@ -16,7 +16,7 @@
# under the License.
"""TIRx tile primitive IR nodes: DispatchContext, TilePrimitiveCall.
-Mirrors the C++ header ``include/tvm/tirx/tile_primitive.h``.
+Mirrors the C++ header ``include/tvm/tirx/tile_op.h``.
"""
# pylint: disable=no-member
diff --git a/python/tvm/tirx/transform/transform.py b/python/tvm/tirx/transform/transform.py
index 4332ca3db8..ff7e536a1c 100644
--- a/python/tvm/tirx/transform/transform.py
+++ b/python/tvm/tirx/transform/transform.py
@@ -146,7 +146,7 @@ def RemoveNoOp():
def RemoveAssume():
- """Remove all instances of builtin::assume
+ """Remove all instances of tirx::assume_op
Returns
-------
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc
index fbfe880efd..9322281bba 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -27,8 +27,8 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/sym/analyzer.h>
-#include <tvm/tirx/builtin.h>
#include <tvm/tirx/index_map.h>
+#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt_functor.h>
#include <cmath>
@@ -202,7 +202,7 @@ class ThreadIdxExtractor : public tirx::StmtExprVisitor {
private:
ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op) final {
- if (op->op.same_as(tirx::builtin::launch_thread())) {
+ if (op->op.same_as(tirx::launch_thread_op())) {
TVM_FFI_CHECK(
std::string(op->args[0].as_or_throw<StringImm>()->value).rfind("vthread", 0) != 0,
ValueError)
@@ -1313,7 +1313,7 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) {
os << "for (int i = 0; i < " << num_elem << "; ++i) {\n";
os << dst << "[" << dst_offset << " + i] = 0.0;";
os << "}\n";
- } else if (op->op.same_as(tirx::builtin::reinterpret())) {
+ } else if (op->op.same_as(tirx::reinterpret_op())) {
// Compile-time pointer reinterpret of a literal (e.g. a tcgen05 descriptor
// template encoded at address 0): emit C++-style reinterpret_cast<T*>(...)
// to match the encoded template. Runtime pointer reinterprets fall through
@@ -1379,8 +1379,7 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) {
if (IsFloat4(tgt_ty)) {
// We view the source as an uint16, and then extract bits of two fp4 numbers,
// and finally reinterpret the result as fp4x2.
- value =
- Call(PrimType::UInt(16), tirx::builtin::reinterpret(), {value}).as_or_throw<PrimExpr>();
+ value = Call(PrimType::UInt(16), tirx::reinterpret_op(), {value}).as_or_throw<PrimExpr>();
PrimVar temp_var("temp_var", PrimType::UInt(16));
value = prim::Let(temp_var, value,
prim::Cast(PrimType::UInt(8),
@@ -1389,19 +1388,18 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) {
} else {
value = prim::Cast(
PrimType::UInt(16),
- Call(PrimType::UInt(8), tirx::builtin::reinterpret(), {value}).as_or_throw<PrimExpr>());
+ Call(PrimType::UInt(8), tirx::reinterpret_op(), {value}).as_or_throw<PrimExpr>());
PrimVar temp_var("temp_var", PrimType::UInt(16));
value = prim::Let(temp_var, value,
(temp_var & IntImm(PrimType::UInt(16), 0xF)) |
((temp_var & IntImm(PrimType::UInt(16), 0xF0)) << 4));
}
- os << PrintExpr(Call(tgt_ty, tirx::builtin::reinterpret(), {value}).as_or_throw<PrimExpr>());
+ os << PrintExpr(Call(tgt_ty, tirx::reinterpret_op(), {value}).as_or_throw<PrimExpr>());
} else if (lanes == 4) {
if (IsFloat4(tgt_ty)) {
// We view the source as an uint32, and then extract bits of four fp4 numbers,
// and finally reinterpret the result as fp4x4.
- value =
- Call(PrimType::UInt(32), tirx::builtin::reinterpret(), {value}).as_or_throw<PrimExpr>();
+ value = Call(PrimType::UInt(32), tirx::reinterpret_op(), {value}).as_or_throw<PrimExpr>();
PrimVar temp_var("temp_var", PrimType::UInt(32));
value = prim::Let(temp_var, value,
prim::Cast(PrimType::UInt(16),
@@ -1410,9 +1408,9 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) {
((temp_var >> 8) & IntImm(PrimType::UInt(32), 0xF00)) |
((temp_var >> 12) & IntImm(PrimType::UInt(32), 0xF000))));
} else {
- value = prim::Cast(PrimType::UInt(32),
- Call(PrimType::UInt(16), tirx::builtin::reinterpret(), {value})
- .as_or_throw<PrimExpr>());
+ value = prim::Cast(
+ PrimType::UInt(32),
+ Call(PrimType::UInt(16), tirx::reinterpret_op(), {value}).as_or_throw<PrimExpr>());
PrimVar temp_var("temp_var", PrimType::UInt(32));
value = prim::Let(temp_var, value,
(temp_var & IntImm(PrimType::UInt(32), 0xF)) |
@@ -1420,19 +1418,19 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) {
((temp_var & IntImm(PrimType::UInt(32), 0xF00)) << 8) |
((temp_var & IntImm(PrimType::UInt(32), 0xF000)) << 12));
}
- os << PrintExpr(Call(tgt_ty, tirx::builtin::reinterpret(), {value}).as_or_throw<PrimExpr>());
+ os << PrintExpr(Call(tgt_ty, tirx::reinterpret_op(), {value}).as_or_throw<PrimExpr>());
} else {
TVM_FFI_THROW(InternalError)
<< "Invalid number of lanes for float4_e2m1fn reinterpret: " << lanes;
}
EndScope(ssa_scope);
- } else if (op->op.same_as(tirx::builtin::print_buffer())) {
+ } else if (op->op.same_as(tirx::print_buffer_op())) {
TVM_FFI_ICHECK_GE(op->args.size(), 5U) << "Print operation expects at least 5 arguments";
Expr arg = op->args[0];
const auto* var_node = arg.as<VarNode>();
if (const auto* call = arg.as<CallNode>();
- call && call->op.same_as(tirx::builtin::buffer_data()) && call->args.size() == 1) {
+ call && call->op.same_as(tirx::buffer_data_op()) && call->args.size() == 1) {
var_node = call->args[0].as<VarNode>();
TVM_FFI_ICHECK(var_node && var_node->ty.as<tirx::TensorTypeNode>())
<< "print_buffer expects buffer_data to project a TensorVar";
@@ -1576,7 +1574,7 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) {
} else if (op->op.same_as(cuda_func_call_op) ||
(op->op.as<Op>() && op->op.as<Op>().value()->name == "tirx.cuda.func_call")) {
print_cuda_func_call(op, os);
- } else if (op->op.same_as(tirx::builtin::thread_return())) {
+ } else if (op->op.same_as(tirx::thread_return_op())) {
os << "return";
} else {
CodeGenC::Dispatch_(op, os);
@@ -1585,7 +1583,7 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) {
void CodeGenCUDA::Dispatch_(const BindNode* op) {
if (const auto* call = op->value.as<CallNode>(); call) {
- if (call->op.same_as(tirx::builtin::alloc_tensor())) return DispatchAllocTensor(op, call);
+ if (call->op.same_as(tirx::alloc_tensor_op())) return DispatchAllocTensor(op, call);
}
CodeGenC::Dispatch_(op);
}
diff --git a/src/backend/cuda/codegen/codegen_cuda_host.cc b/src/backend/cuda/codegen/codegen_cuda_host.cc
index d3d91934f6..6fcbdf3393 100644
--- a/src/backend/cuda/codegen/codegen_cuda_host.cc
+++ b/src/backend/cuda/codegen/codegen_cuda_host.cc
@@ -21,7 +21,7 @@
* \brief C host wrappers with direct CUDA kernel launches.
*/
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/tirx/attrs.h>
+#include <tvm/tirx/op.h>
#include <tvm/tirx/type.h>
#include <algorithm>
@@ -131,7 +131,7 @@ class CodeGenCUDAHost : public CodeGenCHost {
}
void Dispatch_(const CallNode* op, std::ostream& os) override {
- if (op->op.same_as(tirx::builtin::tvm_stack_alloca()) &&
+ if (op->op.same_as(tirx::tvm_stack_alloca_op()) &&
op->args[0].as_or_throw<StringImm>()->value == "tensormap") {
auto count = op->args[1].as_or_throw<IntImm>()->value.as<int64_t>();
TVM_FFI_CHECK(count.has_value() && *count > 0, ValueError)
@@ -142,7 +142,7 @@ class CodeGenCUDAHost : public CodeGenCHost {
os << name;
return;
}
- if (op->op.same_as(tirx::builtin::tvm_call_packed_lowered())) {
+ if (op->op.same_as(tirx::tvm_call_packed_lowered_op())) {
const auto& name = op->args[0].as_or_throw<StringImm>()->value;
if (name == "runtime.cuTensorMapEncodeTiled" || name == "runtime.cuTensorMapInit") {
TVM_FFI_THROW(ValueError)
@@ -175,18 +175,17 @@ class CodeGenCUDAHost : public CodeGenCHost {
<< "cuda_host does not provide TVM runtime operation: " << symbol;
}
}
- if (op->op.same_as(tirx::builtin::call_extern()) ||
- op->op.same_as(tirx::builtin::call_pure_extern())) {
+ if (op->op.same_as(tirx::call_extern_op()) || op->op.same_as(tirx::call_pure_extern_op())) {
const auto& symbol = op->args[0].as_or_throw<StringImm>()->value;
TVM_FFI_CHECK(std::string(symbol).rfind("TVMBackend", 0) != 0, ValueError)
<< "cuda_host does not provide TVM runtime operation: " << symbol;
}
- if (op->op.same_as(tirx::builtin::tensormap_encode_tiled())) {
+ if (op->op.same_as(tirx::tensormap_encode_tiled_op())) {
PrintTensorMapEncode(op);
os << "0";
return;
}
- if (!op->op.same_as(tirx::builtin::call_ffi_kernel())) {
+ if (!op->op.same_as(tirx::call_ffi_kernel_op())) {
CodeGenCHost::Dispatch_(op, os);
return;
}
diff --git a/src/backend/cuda/codegen/intrin_rule_cuda.cc b/src/backend/cuda/codegen/intrin_rule_cuda.cc
index ba01ae7b22..2f1c4b3ebe 100644
--- a/src/backend/cuda/codegen/intrin_rule_cuda.cc
+++ b/src/backend/cuda/codegen/intrin_rule_cuda.cc
@@ -22,8 +22,8 @@
* \brief CUDA intrinsic rules.
*/
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/prim/builtin.h>
-#include <tvm/tirx/builtin.h>
+#include <tvm/ir/prim/op.h>
+#include <tvm/tirx/op.h>
#include <tvm/tirx/op_attr_types.h>
#include "../../../target/intrin_rule.h"
@@ -135,17 +135,17 @@ struct CUDAPopcount {
struct CUDAWarpIntrinsic {
const Op operator()(const PrimType& ty, const Op& orig_op) const {
- if (orig_op.same_as(builtin::tvm_warp_shuffle())) {
+ if (orig_op.same_as(tirx::tvm_warp_shuffle_op())) {
static const Op cuda_shfl_sync_op = Op::Get("tirx.cuda.__shfl_sync");
return cuda_shfl_sync_op;
- } else if (orig_op.same_as(builtin::tvm_warp_shuffle_up())) {
+ } else if (orig_op.same_as(tirx::tvm_warp_shuffle_up_op())) {
static const Op cuda_shfl_up_sync_op = Op::Get("tirx.cuda.__shfl_up_sync");
return cuda_shfl_up_sync_op;
- } else if (orig_op.same_as(builtin::tvm_warp_shuffle_down())) {
+ } else if (orig_op.same_as(tirx::tvm_warp_shuffle_down_op())) {
static const Op cuda_shfl_down_sync_op = Op::Get("tirx.cuda.__shfl_down_sync");
return cuda_shfl_down_sync_op;
} else {
- TVM_FFI_ICHECK(orig_op.same_as(builtin::tvm_warp_shuffle_xor()));
+ TVM_FFI_ICHECK(orig_op.same_as(tirx::tvm_warp_shuffle_xor_op()));
static const Op cuda_shfl_xor_sync_op = Op::Get("tirx.cuda.__shfl_xor_sync");
return cuda_shfl_xor_sync_op;
}
diff --git a/src/backend/cuda/codegen/llvm/codegen_nvptx.cc b/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
index 4e0e853a87..c6f0763f67 100644
--- a/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
+++ b/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
@@ -80,7 +80,7 @@ class CodeGenNVPTX : public CodeGenLLVM {
void Dispatch_(const BindNode* op) final {
if (const auto* call = op->value.as<CallNode>();
- call && call->op.same_as(tirx::builtin::alloc_tensor())) {
+ call && call->op.same_as(tirx::alloc_tensor_op())) {
return DispatchAllocTensor(op, call);
}
CodeGenLLVM::Dispatch_(op);
@@ -260,11 +260,11 @@ static bool GetWarpShuffleIntrinsic(const CallNode* op, llvm::Intrinsic::ID* id)
llvm::Intrinsic::nvvm_shfl_down_i32, llvm::Intrinsic::nvvm_shfl_down_f32};
int offset = 0;
- if (op->op.same_as(builtin::tvm_warp_shuffle())) {
+ if (op->op.same_as(tirx::tvm_warp_shuffle_op())) {
offset = 0;
- } else if (op->op.same_as(builtin::tvm_warp_shuffle_up())) {
+ } else if (op->op.same_as(tirx::tvm_warp_shuffle_up_op())) {
offset = 2;
- } else if (op->op.same_as(builtin::tvm_warp_shuffle_down())) {
+ } else if (op->op.same_as(tirx::tvm_warp_shuffle_down_op())) {
offset = 4;
} else {
return false;
@@ -291,13 +291,13 @@ llvm::Value* CodeGenNVPTX::CreateIntrinsic(const CallNode* op) {
llvm::Type* return_type = arg_type[0];
llvm::Function* func = GetIntrinsicDecl(id, return_type, arg_type);
return builder_->CreateCall(func, arg_value);
- } else if (op->op.same_as(builtin::tvm_warp_activemask())) {
+ } else if (op->op.same_as(tirx::tvm_warp_activemask_op())) {
// Only nvptx target may keep this intrinsic at this point.
// PTX assembly: asm "activemask.b32 r1;"
auto fty = llvm::FunctionType::get(t_int32_, false);
auto val = llvm::InlineAsm::get(fty, "activemask.b32 %0", "=r", true);
return builder_->CreateCall(val);
- } else if (op->op.same_as(builtin::atomic_add())) {
+ } else if (op->op.same_as(tirx::atomic_add_op())) {
PrimExpr value = args[1].as_or_throw<PrimExpr>();
PrimType value_ty = value.ty();
TVM_FFI_ICHECK(value_ty.bits() == 32) << "Only supports 32 bit atomic for now";
diff --git a/src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc b/src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc
index bf456d1980..459fae4d39 100644
--- a/src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc
+++ b/src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc
@@ -24,9 +24,8 @@
#include <tvm/ffi/function.h>
#include <tvm/ir/op.h>
-#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/builtin.h>
+#include <tvm/ir/prim/op.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/op_attr_types.h>
@@ -55,7 +54,7 @@ inline PrimExpr DispatchPureExternLibDevice(const PrimExpr& e) {
ffi::Array<Expr> new_args = {StringImm(intrinsic_name.str())};
new_args.insert(new_args.end(), call->args.begin(), call->args.end());
- return Call(call_ty, builtin::call_pure_extern(), new_args).as_or_throw<PrimExpr>();
+ return Call(call_ty, tirx::call_pure_extern_op(), new_args).as_or_throw<PrimExpr>();
}
namespace llvm {
diff --git a/src/backend/cuda/op/attrs.h b/src/backend/cuda/op/attrs.h
deleted file mode 100644
index 93fa8d3fa3..0000000000
--- a/src/backend/cuda/op/attrs.h
+++ /dev/null
@@ -1,108 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-#ifndef TVM_BACKEND_CUDA_OP_ATTRS_H_
-#define TVM_BACKEND_CUDA_OP_ATTRS_H_
-
-#include <tvm/ir/attrs.h>
-
-namespace tvm {
-namespace tirx {
-
-/*! \brief Static options for the dense tcgen05 instruction descriptor. */
-struct TCGen05InstrDescriptorAttrs : public AttrsNode {
- ffi::String d_dtype;
- ffi::String a_dtype;
- ffi::String b_dtype;
- int64_t M;
- int64_t N;
- int64_t K;
- bool trans_a;
- bool trans_b;
- int64_t n_cta_groups = 1;
- bool neg_a = false;
- bool neg_b = false;
- bool sat_d = false;
- bool is_sparse = false;
-
- static void RegisterReflection() {
- namespace refl = ffi::reflection;
- refl::ObjectDef<TCGen05InstrDescriptorAttrs>()
- .def_ro("d_dtype", &TCGen05InstrDescriptorAttrs::d_dtype)
- .def_ro("a_dtype", &TCGen05InstrDescriptorAttrs::a_dtype)
- .def_ro("b_dtype", &TCGen05InstrDescriptorAttrs::b_dtype)
- .def_ro("M", &TCGen05InstrDescriptorAttrs::M)
- .def_ro("N", &TCGen05InstrDescriptorAttrs::N)
- .def_ro("K", &TCGen05InstrDescriptorAttrs::K)
- .def_ro("trans_a", &TCGen05InstrDescriptorAttrs::trans_a)
- .def_ro("trans_b", &TCGen05InstrDescriptorAttrs::trans_b)
- .def_ro("n_cta_groups", &TCGen05InstrDescriptorAttrs::n_cta_groups, refl::DefaultValue(1))
- .def_ro("neg_a", &TCGen05InstrDescriptorAttrs::neg_a, refl::DefaultValue(false))
- .def_ro("neg_b", &TCGen05InstrDescriptorAttrs::neg_b, refl::DefaultValue(false))
- .def_ro("sat_d", &TCGen05InstrDescriptorAttrs::sat_d, refl::DefaultValue(false))
- .def_ro("is_sparse", &TCGen05InstrDescriptorAttrs::is_sparse, refl::DefaultValue(false));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.cuda.TCGen05InstrDescriptorAttrs",
- TCGen05InstrDescriptorAttrs, AttrsNode);
-};
-
-/*! \brief Static options for the block-scaled tcgen05 instruction descriptor. */
-struct TCGen05InstrDescriptorBlockScaledAttrs : public AttrsNode {
- ffi::String d_dtype;
- ffi::String a_dtype;
- ffi::String b_dtype;
- ffi::String sfa_dtype;
- ffi::String sfb_dtype;
- int64_t M;
- int64_t N;
- int64_t K;
- bool trans_a;
- bool trans_b;
- int64_t n_cta_groups = 1;
- bool neg_a = false;
- bool neg_b = false;
- bool is_sparse = false;
-
- static void RegisterReflection() {
- namespace refl = ffi::reflection;
- refl::ObjectDef<TCGen05InstrDescriptorBlockScaledAttrs>()
- .def_ro("d_dtype", &TCGen05InstrDescriptorBlockScaledAttrs::d_dtype)
- .def_ro("a_dtype", &TCGen05InstrDescriptorBlockScaledAttrs::a_dtype)
- .def_ro("b_dtype", &TCGen05InstrDescriptorBlockScaledAttrs::b_dtype)
- .def_ro("sfa_dtype", &TCGen05InstrDescriptorBlockScaledAttrs::sfa_dtype)
- .def_ro("sfb_dtype", &TCGen05InstrDescriptorBlockScaledAttrs::sfb_dtype)
- .def_ro("M", &TCGen05InstrDescriptorBlockScaledAttrs::M)
- .def_ro("N", &TCGen05InstrDescriptorBlockScaledAttrs::N)
- .def_ro("K", &TCGen05InstrDescriptorBlockScaledAttrs::K)
- .def_ro("trans_a", &TCGen05InstrDescriptorBlockScaledAttrs::trans_a)
- .def_ro("trans_b", &TCGen05InstrDescriptorBlockScaledAttrs::trans_b)
- .def_ro("n_cta_groups", &TCGen05InstrDescriptorBlockScaledAttrs::n_cta_groups,
- refl::DefaultValue(1))
- .def_ro("neg_a", &TCGen05InstrDescriptorBlockScaledAttrs::neg_a, refl::DefaultValue(false))
- .def_ro("neg_b", &TCGen05InstrDescriptorBlockScaledAttrs::neg_b, refl::DefaultValue(false))
- .def_ro("is_sparse", &TCGen05InstrDescriptorBlockScaledAttrs::is_sparse,
- refl::DefaultValue(false));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.cuda.TCGen05InstrDescriptorBlockScaledAttrs",
- TCGen05InstrDescriptorBlockScaledAttrs, AttrsNode);
-};
-
-} // namespace tirx
-} // namespace tvm
-
-#endif // TVM_BACKEND_CUDA_OP_ATTRS_H_
diff --git a/src/backend/cuda/op/iket.cc b/src/backend/cuda/op/iket.cc
index fbd427a842..fa66ae734f 100644
--- a/src/backend/cuda/op/iket.cc
+++ b/src/backend/cuda/op/iket.cc
@@ -21,6 +21,7 @@
* \brief Frontend-only NVIDIA IKET annotation operators for CUDA.
*/
+#include <tvm/backend/cuda/op.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/op_attr_types.h>
@@ -79,4 +80,40 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
} // namespace tirx
+
+namespace backend {
+namespace cuda {
+
+const Op& iket_mark_op() {
+ static const Op op = Op::Get("tirx.cuda.iket_mark");
+ return op;
+}
+
+const Op& iket_range_start_op() {
+ static const Op op = Op::Get("tirx.cuda.iket_range_start");
+ return op;
+}
+
+const Op& iket_range_end_op() {
+ static const Op op = Op::Get("tirx.cuda.iket_range_end");
+ return op;
+}
+
+const Op& iket_range_push_op() {
+ static const Op op = Op::Get("tirx.cuda.iket_range_push");
+ return op;
+}
+
+const Op& iket_range_pop_op() {
+ static const Op op = Op::Get("tirx.cuda.iket_range_pop");
+ return op;
+}
+
+const Op& iket_sentinel_token_op() {
+ static const Op op = Op::Get("tirx.cuda.iket_sentinel_token");
+ return op;
+}
+
+} // namespace cuda
+} // namespace backend
} // namespace tvm
diff --git a/src/backend/cuda/op/target_builtin.cc b/src/backend/cuda/op/target_builtin.cc
index 0c9e9c4590..3a1efc1ead 100644
--- a/src/backend/cuda/op/target_builtin.cc
+++ b/src/backend/cuda/op/target_builtin.cc
@@ -22,18 +22,61 @@
*
* builtin intrinsic operators specific to CUDA target.
*/
+#include <tvm/backend/cuda/op.h>
#include <tvm/ffi/function.h>
-#include <tvm/ir/prim/builtin.h>
+#include <tvm/ir/prim/op.h>
#include <tvm/runtime/base.h>
-#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/op_attr_types.h>
#include <string>
-#include "attrs.h"
-
namespace tvm {
+namespace backend {
+namespace cuda {
+
+void TCGen05InstrDescriptorAttrs::RegisterReflection() {
+ namespace refl = ffi::reflection;
+ refl::ObjectDef<TCGen05InstrDescriptorAttrs>()
+ .def_ro("d_dtype", &TCGen05InstrDescriptorAttrs::d_dtype)
+ .def_ro("a_dtype", &TCGen05InstrDescriptorAttrs::a_dtype)
+ .def_ro("b_dtype", &TCGen05InstrDescriptorAttrs::b_dtype)
+ .def_ro("M", &TCGen05InstrDescriptorAttrs::M)
+ .def_ro("N", &TCGen05InstrDescriptorAttrs::N)
+ .def_ro("K", &TCGen05InstrDescriptorAttrs::K)
+ .def_ro("trans_a", &TCGen05InstrDescriptorAttrs::trans_a)
+ .def_ro("trans_b", &TCGen05InstrDescriptorAttrs::trans_b)
+ .def_ro("n_cta_groups", &TCGen05InstrDescriptorAttrs::n_cta_groups, refl::DefaultValue(1))
+ .def_ro("neg_a", &TCGen05InstrDescriptorAttrs::neg_a, refl::DefaultValue(false))
+ .def_ro("neg_b", &TCGen05InstrDescriptorAttrs::neg_b, refl::DefaultValue(false))
+ .def_ro("sat_d", &TCGen05InstrDescriptorAttrs::sat_d, refl::DefaultValue(false))
+ .def_ro("is_sparse", &TCGen05InstrDescriptorAttrs::is_sparse, refl::DefaultValue(false));
+}
+
+void TCGen05InstrDescriptorBlockScaledAttrs::RegisterReflection() {
+ namespace refl = ffi::reflection;
+ refl::ObjectDef<TCGen05InstrDescriptorBlockScaledAttrs>()
+ .def_ro("d_dtype", &TCGen05InstrDescriptorBlockScaledAttrs::d_dtype)
+ .def_ro("a_dtype", &TCGen05InstrDescriptorBlockScaledAttrs::a_dtype)
+ .def_ro("b_dtype", &TCGen05InstrDescriptorBlockScaledAttrs::b_dtype)
+ .def_ro("sfa_dtype", &TCGen05InstrDescriptorBlockScaledAttrs::sfa_dtype)
+ .def_ro("sfb_dtype", &TCGen05InstrDescriptorBlockScaledAttrs::sfb_dtype)
+ .def_ro("M", &TCGen05InstrDescriptorBlockScaledAttrs::M)
+ .def_ro("N", &TCGen05InstrDescriptorBlockScaledAttrs::N)
+ .def_ro("K", &TCGen05InstrDescriptorBlockScaledAttrs::K)
+ .def_ro("trans_a", &TCGen05InstrDescriptorBlockScaledAttrs::trans_a)
+ .def_ro("trans_b", &TCGen05InstrDescriptorBlockScaledAttrs::trans_b)
+ .def_ro("n_cta_groups", &TCGen05InstrDescriptorBlockScaledAttrs::n_cta_groups,
+ refl::DefaultValue(1))
+ .def_ro("neg_a", &TCGen05InstrDescriptorBlockScaledAttrs::neg_a, refl::DefaultValue(false))
+ .def_ro("neg_b", &TCGen05InstrDescriptorBlockScaledAttrs::neg_b, refl::DefaultValue(false))
+ .def_ro("is_sparse", &TCGen05InstrDescriptorBlockScaledAttrs::is_sparse,
+ refl::DefaultValue(false));
+}
+
+} // namespace cuda
+} // namespace backend
+
namespace tirx {
namespace builtin {
@@ -76,8 +119,8 @@ void RegisterCudaTargetBuiltins() {
if (registered) return;
registered = true;
- TCGen05InstrDescriptorAttrs::RegisterReflection();
- TCGen05InstrDescriptorBlockScaledAttrs::RegisterReflection();
+ backend::cuda::TCGen05InstrDescriptorAttrs::RegisterReflection();
+ backend::cuda::TCGen05InstrDescriptorBlockScaledAttrs::RegisterReflection();
OpDef("tirx.tvm_load_matrix_sync")
.set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void())
@@ -430,11 +473,11 @@ void RegisterDeviceIntrinsics() {
.set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Int(64));
RegisterDeviceIntrinsic(OpDef("tirx.cuda.tcgen05_encode_instr_descriptor"), "cuda",
CallEffectKind::kOpaque, sig::arg("desc"),
- sig::call_attrs<TCGen05InstrDescriptorAttrs>())
+ sig::call_attrs<backend::cuda::TCGen05InstrDescriptorAttrs>())
.set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void());
RegisterDeviceIntrinsic(OpDef("tirx.cuda.tcgen05_encode_instr_descriptor_block_scaled"), "cuda",
CallEffectKind::kOpaque, sig::arg("desc"),
- sig::call_attrs<TCGen05InstrDescriptorBlockScaledAttrs>())
+ sig::call_attrs<backend::cuda::TCGen05InstrDescriptorBlockScaledAttrs>())
.set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void());
RegisterDeviceIntrinsic(OpDef("tirx.cuda.tcgen05_encode_matrix_descriptor"), "cuda",
CallEffectKind::kOpaque, sig::arg("desc"), sig::arg("addr"),
diff --git a/src/backend/cuda/transforms/lower_iket.cc b/src/backend/cuda/transforms/lower_iket.cc
index 59b87b825f..2fb5ac464c 100644
--- a/src/backend/cuda/transforms/lower_iket.cc
+++ b/src/backend/cuda/transforms/lower_iket.cc
@@ -22,6 +22,7 @@
* \brief Lower frontend-only TIRx IKET annotations to CUDA tracing helpers.
*/
+#include <tvm/backend/cuda/op.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/type.h>
#include <tvm/runtime/logging.h>
@@ -86,46 +87,20 @@ enum class PayloadType : uint32_t {
kUI64 = 16,
};
-const Op& IketMarkOp() {
- static const Op op = Op::Get("tirx.cuda.iket_mark");
- return op;
-}
-
-const Op& IketRangeStartOp() {
- static const Op op = Op::Get("tirx.cuda.iket_range_start");
- return op;
-}
-
-const Op& IketRangeEndOp() {
- static const Op op = Op::Get("tirx.cuda.iket_range_end");
- return op;
-}
-
-const Op& IketRangePushOp() {
- static const Op op = Op::Get("tirx.cuda.iket_range_push");
- return op;
-}
-
-const Op& IketRangePopOp() {
- static const Op op = Op::Get("tirx.cuda.iket_range_pop");
- return op;
-}
-
-const Op& IketSentinelOp() {
- static const Op op = Op::Get("tirx.cuda.iket_sentinel_token");
- return op;
-}
-
bool IsIketOp(const ffi::ObjectRef& value) {
auto op = value.as<Op>();
if (!op.has_value()) return false;
- return op.value().same_as(IketMarkOp()) || op.value().same_as(IketRangeStartOp()) ||
- op.value().same_as(IketRangeEndOp()) || op.value().same_as(IketRangePushOp()) ||
- op.value().same_as(IketRangePopOp()) || op.value().same_as(IketSentinelOp());
+ return op.value().same_as(backend::cuda::iket_mark_op()) ||
+ op.value().same_as(backend::cuda::iket_range_start_op()) ||
+ op.value().same_as(backend::cuda::iket_range_end_op()) ||
+ op.value().same_as(backend::cuda::iket_range_push_op()) ||
+ op.value().same_as(backend::cuda::iket_range_pop_op()) ||
+ op.value().same_as(backend::cuda::iket_sentinel_token_op());
}
bool IsTokenProducer(const CallNode* call) {
- return call->op.same_as(IketRangeStartOp()) || call->op.same_as(IketSentinelOp());
+ return call->op.same_as(backend::cuda::iket_range_start_op()) ||
+ call->op.same_as(backend::cuda::iket_sentinel_token_op());
}
bool IsValidUTF8(const std::string& value) {
@@ -336,15 +311,15 @@ class AnnotationCollector : public StmtExprVisitor {
return StmtExprVisitor::Visit_(call);
}
has_annotations = true;
- if (call->op.same_as(IketMarkOp())) {
+ if (call->op.same_as(backend::cuda::iket_mark_op())) {
AddDeclaration(DeclarationKind::kMark, call);
- } else if (call->op.same_as(IketRangeStartOp())) {
+ } else if (call->op.same_as(backend::cuda::iket_range_start_op())) {
TVM_FFI_CHECK(call->ty.as_or_throw<PrimType>()->dtype.code == kDLUInt &&
call->ty.as_or_throw<PrimType>()->dtype.bits == 32,
TypeError)
<< "IKET range_start must return uint32";
AddDeclaration(DeclarationKind::kRange, call);
- } else if (call->op.same_as(IketSentinelOp())) {
+ } else if (call->op.same_as(backend::cuda::iket_sentinel_token_op())) {
TVM_FFI_CHECK_EQ(call->args.size(), 1, TypeError)
<< "IKET sentinel_token expects one event name";
TVM_FFI_CHECK(call->ty.as_or_throw<PrimType>()->dtype.code == kDLUInt &&
@@ -354,9 +329,9 @@ class AnnotationCollector : public StmtExprVisitor {
// A sentinel carries only token-flow identity. It emits no runtime
// event and therefore must not create metadata or consume an event ID.
GetName(call);
- } else if (call->op.same_as(IketRangePushOp())) {
+ } else if (call->op.same_as(backend::cuda::iket_range_push_op())) {
AddDeclaration(DeclarationKind::kPush, call);
- } else if (call->op.same_as(IketRangeEndOp())) {
+ } else if (call->op.same_as(backend::cuda::iket_range_end_op())) {
TVM_FFI_CHECK(call->args.size() == 1 || call->args.size() == 2, TypeError)
<< "IKET range_end expects a token and optional payload";
DLDataType token_dtype = call->args[0].as_or_throw<PrimExpr>().ty()->dtype;
@@ -367,7 +342,7 @@ class AnnotationCollector : public StmtExprVisitor {
has_payload_calls = true;
ValidatePayload(call->args[1]);
}
- } else if (call->op.same_as(IketRangePopOp())) {
+ } else if (call->op.same_as(backend::cuda::iket_range_pop_op())) {
TVM_FFI_CHECK_EQ(call->args.size(), 0, TypeError) << "IKET range_pop takes no arguments";
}
@@ -460,7 +435,7 @@ class RangeEndSchemaVerifier : public StmtExprVisitor {
private:
ffi::Optional<VisitInterrupt> Visit_(const CallNode* call) final {
- if (!call->op.same_as(IketRangeEndOp())) {
+ if (!call->op.same_as(backend::cuda::iket_range_end_op())) {
return StmtExprVisitor::Visit_(call);
}
const auto* token = call->args[0].as<TensorLoadNode>();
@@ -567,7 +542,7 @@ class TokenVerifier : public StmtExprVisitor {
allow_producer_ = old_allow_producer;
return std::nullopt;
}
- if (call->op.same_as(IketRangeEndOp())) {
+ if (call->op.same_as(backend::cuda::iket_range_end_op())) {
TVM_FFI_CHECK_GE(call->args.size(), 1, TypeError) << "range_end requires a RangeToken";
const auto* token = call->args[0].as<TensorLoadNode>();
TVM_FFI_CHECK(
@@ -596,9 +571,9 @@ class StripIket : public StmtExprMutator {
private:
UnchangedOr<Stmt> Mutate_(const BindNode* alloc, InplaceMode inplace_mode) final {
- if (const auto* call = alloc->value.as<CallNode>();
- call && call->op.same_as(tirx::builtin::alloc_tensor()) &&
- token_buffers_.count(alloc->var.get())) {
+ if (const auto* call = alloc->value.as<CallNode>(); call &&
+ call->op.same_as(tirx::alloc_tensor_op()) &&
+ token_buffers_.count(alloc->var.get())) {
return Evaluate(0);
}
return StmtExprMutator::Mutate_(alloc, inplace_mode);
@@ -617,7 +592,8 @@ class StripIket : public StmtExprMutator {
}
UnchangedOr<Expr> Mutate_(const CallNode* call, InplaceMode inplace_mode) final {
- if (call->op.same_as(IketRangeStartOp()) || call->op.same_as(IketSentinelOp())) {
+ if (call->op.same_as(backend::cuda::iket_range_start_op()) ||
+ call->op.same_as(backend::cuda::iket_sentinel_token_op())) {
return IntImm(PrimType::UInt(32), 0);
}
return StmtExprMutator::Mutate_(call, inplace_mode);
@@ -1127,7 +1103,7 @@ class InstrumentOfficialKernel : public StmtExprMutator {
UnchangedOr<Stmt> Mutate_(const EvaluateNode* evaluate, InplaceMode inplace_mode) final {
if (const auto* call = evaluate->value.as<CallNode>();
- call && call->op.same_as(IketRangeEndOp())) {
+ call && call->op.same_as(backend::cuda::iket_range_end_op())) {
PrimExpr token =
Mutate(call->args[0]).ValueOrUnchanged(call->args[0]).as_or_throw<PrimExpr>();
if (call->args.size() == 2) {
@@ -1143,7 +1119,7 @@ class InstrumentOfficialKernel : public StmtExprMutator {
}
UnchangedOr<Expr> Mutate_(const CallNode* call, InplaceMode inplace_mode) final {
- if (call->op.same_as(IketRangeStartOp())) {
+ if (call->op.same_as(backend::cuda::iket_range_start_op())) {
const Declaration& declaration = Lookup(DeclarationKind::kRange, call);
PrimExpr event_id = IntImm(PrimType::UInt(32), declaration.event_id);
if (declaration.has_payload) {
@@ -1155,8 +1131,9 @@ class InstrumentOfficialKernel : public StmtExprMutator {
}
return Event(event_id);
}
- if (call->op.same_as(IketSentinelOp())) return IntImm(PrimType::UInt(32), 0);
- if (call->op.same_as(IketMarkOp())) {
+ if (call->op.same_as(backend::cuda::iket_sentinel_token_op()))
+ return IntImm(PrimType::UInt(32), 0);
+ if (call->op.same_as(backend::cuda::iket_mark_op())) {
const Declaration& declaration = Lookup(DeclarationKind::kMark, call);
PrimExpr event_id = IntImm(PrimType::UInt(32), declaration.event_id);
if (declaration.has_payload) {
@@ -1168,7 +1145,7 @@ class InstrumentOfficialKernel : public StmtExprMutator {
}
return Event(event_id);
}
- if (call->op.same_as(IketRangePushOp())) {
+ if (call->op.same_as(backend::cuda::iket_range_push_op())) {
const Declaration& declaration = Lookup(DeclarationKind::kPush, call);
PrimExpr event_id = IntImm(PrimType::UInt(32), declaration.event_id);
if (declaration.has_payload) {
@@ -1180,10 +1157,10 @@ class InstrumentOfficialKernel : public StmtExprMutator {
}
return Event(event_id);
}
- if (call->op.same_as(IketRangePopOp())) {
+ if (call->op.same_as(backend::cuda::iket_range_pop_op())) {
... 9435 lines suppressed ...