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 f66d50df63 refactor(op-dispatch)!: dispatch tile ops through Op
attributes (#20671)
f66d50df63 is described below
commit f66d50df63f18ca8d17886c8ef8475abf8d47fbd
Author: Bohan Hou <[email protected]>
AuthorDate: Mon Oct 12 08:15:08 2026 -0400
refactor(op-dispatch)!: dispatch tile ops through Op attributes (#20671)
TileDispatch currently identifies tile Calls through `TIRxOpCategory`
and invokes fixed Python dispatcher and scope globals. Register
`FTileOpDispatch(Call, DispatchContext) -> Function` on each Op and use
that attribute for both tile detection and lowering. Resolve execution
scopes from reflected Call attributes so native C++ callbacks can run
without loading Python.
Remove the category definitions and registrations, migrate Trainium
workspace callers, and update documentation and regression coverage for
custom Ops, live callback replacement, recursive expansion, residual
Calls, and invalid results.
Downstream compatibility PR:
https://github.com/mlc-ai/TIRx-kernels/pull/10
Validation:
- Incremental C++ build and pre-commit passed.
- 54 targeted dispatch, namespace, scope, and Trainium
private-allocation tests passed.
- Standalone C++ registration and TileDispatch passed without Python,
including an initially absent dispatch attribute column.
- Extended CUDA coverage: 389 passed, 2 skipped, 24 vectorization
failures. All 24 failures reproduced on the original revision with the
same `lanes > 1` error.
- Op tests: 12 passed, 3 JSON round-trip failures. All 3 failures
reproduced on the original revision.
---
docs/tirx/api/tile_dispatch.rst | 11 +-
docs/tirx/arch/tile_dispatch.rst | 18 ++-
include/tvm/tirx/op_attr_types.h | 11 --
include/tvm/tirx/tile/dispatch.h | 16 ++-
python/tvm/backend/cuda/ptx/engine.py | 2 -
python/tvm/backend/trn/tile/private_alloc.py | 6 +-
python/tvm/tirx/tile/dispatch.py | 6 -
python/tvm/tirx/tile/instruction.py | 13 +-
src/backend/cuda/codegen/intrin_rule_cuda.cc | 5 -
src/backend/cuda/op/iket.cc | 7 --
src/backend/cuda/op/target_builtin.cc | 14 +--
src/backend/cuda/op/tcgen05.cc | 3 -
src/backend/cuda/op/tensormap.cc | 1 -
src/backend/metal/codegen/intrin_rule_metal.cc | 3 -
src/backend/metal/op/cooperative_tensor.cc | 4 -
src/backend/metal/op/target_builtin.cc | 4 -
src/backend/opencl/op/memory.cc | 3 -
src/backend/trn/op/target_builtin.cc | 2 -
src/backend/webgpu/codegen/intrin_rule_webgpu.cc | 3 -
src/tirx/op/abi.cc | 18 ---
src/tirx/op/annotation.cc | 2 -
src/tirx/op/cpu.cc | 1 -
src/tirx/op/gpu.cc | 16 ---
src/tirx/op/memory.cc | 15 ---
src/tirx/op/prim_metadata.cc | 132 +++++++--------------
src/tirx/op/region.cc | 15 +--
src/tirx/op/vector.cc | 4 -
src/tirx/tile/dispatch.cc | 11 +-
src/tirx/tile/lower.cc | 29 +++--
src/tirx/transform/split_host_device.cc | 3 +-
tests/python/tirx/iket/test_iket_profiler.py | 2 +-
tests/python/tirx/test_op.py | 5 +-
tests/python/tirx/test_op_namespace_cleanup.py | 35 ++----
.../tirx/tile/cuda/gemm/test_gemm_mma_m16n8k_.py | 2 +-
tests/python/tirx/tile/test_dispatcher.py | 85 ++++++++++++-
35 files changed, 218 insertions(+), 289 deletions(-)
diff --git a/docs/tirx/api/tile_dispatch.rst b/docs/tirx/api/tile_dispatch.rst
index 0f74e61eea..6e3b4d3ef2 100644
--- a/docs/tirx/api/tile_dispatch.rst
+++ b/docs/tirx/api/tile_dispatch.rst
@@ -24,7 +24,7 @@ The shared implementation is ``tvm.tirx.tile.instruction``;
CUDA and
Trainium's ``tile/instructions.py`` files contain the concrete contracts.
Registration installs the fixed argument count, void return type, opaque
-side effects, category, printer name, and native Call validator on the Op.
+side effects, ``FTileOpDispatch``, printer name, and native Call validator on
the Op.
The validator also runs for directly constructed Calls and printer entry.
Runtime expressions must be operands. New static qualifiers require a typed
attribute field rather than an arbitrary dictionary.
@@ -34,6 +34,15 @@ Lowerers receive a semantic view and ``DispatchContext`` and
return a
layouts for the selected core instruction. They must report unsupported
contracts rather than select a different instruction family.
+``FTileOpDispatch`` has the signature
+``(tvm.ir.Call, DispatchContext) -> tvm.tirx.Function``. Its presence
identifies
+a tile Op; no separate category registration is required. ``Instruction``
+registers an adapter that decodes the semantic view and invokes its lowerer.
+Custom implementations can register the callback directly using
+``tvm.ir.register_op_attr(op_name, "FTileOpDispatch", dispatch)`` or C++
+``OpDef::set_attr<tile::FTileOpDispatch>``. The Call must carry reflected
+attributes with a string ``scope`` field in the target's execution hierarchy.
+
Named workspace operands have fixed optional slots. Declare a workspace policy
only for instructions that own such scratch storage; do not use it to hide
caller-owned algorithms. See :doc:`../arch/tile_dispatch` for callbacks and
diff --git a/docs/tirx/arch/tile_dispatch.rst b/docs/tirx/arch/tile_dispatch.rst
index 4c963f44af..fbde2a6eec 100644
--- a/docs/tirx/arch/tile_dispatch.rst
+++ b/docs/tirx/arch/tile_dispatch.rst
@@ -27,7 +27,7 @@ types and the native Call validator are in
``src/tirx/tile/op.cc``.
The public C++ tile interfaces are grouped under ``include/tvm/tirx/tile/``:
- ``op.h`` declares the instruction registration and validation contract.
-- ``dispatch.h`` declares the dispatcher entry point.
+- ``dispatch.h`` declares ``FTileOpDispatch`` and the dispatcher entry point.
- ``exec_scope.h`` defines backend execution hierarchy nodes used by layouts
and tile calls. The shared ``tirx::ExecScope`` type remains unchanged.
- ``context.h`` defines the backend-neutral ``tile::DispatchContext`` and its
@@ -39,9 +39,10 @@ active-thread execution contexts, and predicate analysis
remain private to
``include/tvm/tirx/layout.h``.
``tirx.TileDispatch`` remains the first phase of ``LowerTIRx``.
-The pass recognizes ``Evaluate(Call)`` by the operator's ``TIRxOpCategory``:
-``tile_primitive`` or ``tile_composite``. It validates the call, resolves the
-scope against the active thread set, constructs a ``DispatchContext``, and
+The pass recognizes ``Evaluate(Call)`` by the presence of the operator's
+``FTileOpDispatch`` attribute. It validates the call, reads its reflected
+``scope`` attribute, resolves the scope against the active thread set,
+constructs a ``DispatchContext``, and
invokes the instruction's registered lowerer. It replaces the Evaluate with
the returned function body and lowers any nested tensor calls.
@@ -49,8 +50,13 @@ This preserves launch parameters, inter/intra scope maps,
value ranges,
descriptor caches, host initialization, and allocation/initialization
callbacks. Lowering occurs at a statement boundary because one instruction
may require address preparation, layout expansion, and multiple occurrences
-of the same core instruction. A residual-call verifier rejects either
-category after expansion.
+of the same core instruction. A residual-call verifier rejects any remaining
+tile instruction after expansion using the same dispatch attribute.
+
+The dispatch table is the native Op attribute registry. Python registers
+callbacks on individual Ops; C++ implementations can register the same typed
+attribute directly. The pass does not call a global Python dispatcher or
+consult a Python instruction table to identify tile Ops or resolve their scope.
There is one lowerer per instruction identity. Unsupported layouts and scopes
produce an error that includes the instruction, target, scope, and original
diff --git a/include/tvm/tirx/op_attr_types.h b/include/tvm/tirx/op_attr_types.h
index a1cd0767ad..729c62a3c1 100644
--- a/include/tvm/tirx/op_attr_types.h
+++ b/include/tvm/tirx/op_attr_types.h
@@ -56,16 +56,6 @@ using FLowerIntrinsic =
ffi::TypedFunction<PrimExpr(PrimExpr)>;
*/
using FLegalize = ffi::TypedFunction<PrimExpr(PrimExpr)>;
-/*!
- * \brief Broad TIRx op category.
- *
- * Expected values:
- * - "builtin"
- * - "tile_primitive"
- * - "device_intrin"
- */
-using TIRxOpCategory = ffi::String;
-
/*!
* \brief Device intrinsic namespace.
*
@@ -76,7 +66,6 @@ using TDeviceIntrinsicNamespace = ffi::String;
namespace op_attr {
inline constexpr const char* kGlobalSymbol = "TGlobalSymbol";
inline constexpr const char* kVectorizable = "TVectorizable";
-inline constexpr const char* kOpCategory = "TIRxOpCategory";
inline constexpr const char* kDeviceIntrinsicNamespace =
"TDeviceIntrinsicNamespace";
} // namespace op_attr
diff --git a/include/tvm/tirx/tile/dispatch.h b/include/tvm/tirx/tile/dispatch.h
index d912655664..d2672ec3f2 100644
--- a/include/tvm/tirx/tile/dispatch.h
+++ b/include/tvm/tirx/tile/dispatch.h
@@ -20,13 +20,27 @@
#ifndef TVM_TIRX_TILE_DISPATCH_H_
#define TVM_TIRX_TILE_DISPATCH_H_
+#include <tvm/ffi/function.h>
+#include <tvm/ir/op.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/tile/context.h>
namespace tvm::tirx::tile {
/*!
- * \brief Invoke the registered tile dispatcher.
+ * \brief Lower a tile instruction Call using its execution context.
+ *
+ * Registering this attribute identifies an Op as a tile instruction. The
+ * dispatcher can be implemented in C++ or registered from Python.
+ */
+using FTileOpDispatch = ffi::TypedFunction<Function(const Call&, const
DispatchContext&)>;
+
+namespace op_attr {
+inline constexpr const char* kTileOpDispatch = "FTileOpDispatch";
+} // namespace op_attr
+
+/*!
+ * \brief Invoke the Call's FTileOpDispatch attribute.
* \param call The tile instruction to lower.
* \param context The operator's execution context and accumulated callbacks.
* \return The instruction implementation, whose body is expanded by
TileDispatch.
diff --git a/python/tvm/backend/cuda/ptx/engine.py
b/python/tvm/backend/cuda/ptx/engine.py
index 5e869f671e..d86a8e18f4 100644
--- a/python/tvm/backend/cuda/ptx/engine.py
+++ b/python/tvm/backend/cuda/ptx/engine.py
@@ -107,7 +107,6 @@ def register_table(table: dict[str, InstructionEntry]) ->
None:
# other family, and `gen_stubs` already spells the attribute this way.
family = escape_token(entry.family)
register_op_attr(entry.op_name, "TScriptPrinterName",
f"tirx.ptx.{family}")
- register_op_attr(entry.op_name, "TIRxOpCategory", "device_intrin")
register_op_attr(entry.op_name, "TDeviceIntrinsicNamespace", "ptx")
register_op_attr(entry.op_name, "TFixedReturnType", PrimType("void"))
if Op.get(entry.op_name).get_attr("__tvm_doc_translate_op_call__") is
None:
@@ -126,7 +125,6 @@ def register_addr() -> None:
op.add_arg("base", "The base pointer.")
op.add_arg("byte_offset", "The offset in bytes.")
register_op_attr(_ADDR_OP_NAME, "TScriptPrinterName", "tirx.ptx.addr")
- register_op_attr(_ADDR_OP_NAME, "TIRxOpCategory", "device_intrin")
register_op_attr(_ADDR_OP_NAME, "TDeviceIntrinsicNamespace", "ptx")
register_op_attr(_ADDR_OP_NAME, "FInferType", _infer_addr_type)
if op.get_attr("__tvm_doc_translate_op_call__") is None:
diff --git a/python/tvm/backend/trn/tile/private_alloc.py
b/python/tvm/backend/trn/tile/private_alloc.py
index 3914d68337..9f0cce70a5 100644
--- a/python/tvm/backend/trn/tile/private_alloc.py
+++ b/python/tvm/backend/trn/tile/private_alloc.py
@@ -22,7 +22,7 @@ from tvm.ir import Bind, Call, DictAttrs, Stmt, Tuple
from tvm.script import tirx as T
from tvm.tirx import FloatImm, IntImm, Var
from tvm.tirx.tile import DispatchContext
-from tvm.tirx.tile.dispatch import f_op_dispatcher
+from tvm.tirx.tile.dispatch import run_dispatch
from tvm.tirx.tile.instruction import TensorCall
from tvm.tirx.tile.views import (
BinaryReduce,
@@ -105,7 +105,7 @@ def alloc_partial_reduce_trn(
) -> dict[str, Any]:
if "partial_reduce" in op.workspaces:
return {}
- f_op_dispatcher(op, sctx)
+ run_dispatch(op, sctx)
partial_reduce_tensor = None
if DispatchContext.kPrivateAlloc not in sctx.callbacks:
return {}
@@ -145,7 +145,7 @@ def alloc_unary_reduce_trn(
else:
if "const_bias" in op.workspaces and "partial_reduce" in op.workspaces:
return {}
- f_op_dispatcher(op, sctx)
+ run_dispatch(op, sctx)
partial_reduce_tensor = None
const_bias_tensor = None
if DispatchContext.kPrivateAlloc not in sctx.callbacks:
diff --git a/python/tvm/tirx/tile/dispatch.py b/python/tvm/tirx/tile/dispatch.py
index 5afbf8a1fb..12477fb3e8 100644
--- a/python/tvm/tirx/tile/dispatch.py
+++ b/python/tvm/tirx/tile/dispatch.py
@@ -20,7 +20,6 @@ Mirrors the C++ header ``include/tvm/tirx/tile/context.h``.
"""
# pylint: disable=no-member
-import tvm_ffi
from tvm_ffi import register_object
from tvm.ir import Expr, Range, Scriptable, Stmt, Var
@@ -251,8 +250,3 @@ def run_dispatch(call, sctx):
f"target={sctx.target.kind.name} scope={sctx.scope_kind}\n"
f"{error}\ncall: {view.call}"
) from error
-
-
-@tvm_ffi.register_global_func("tirx.f_op_dispatcher")
-def f_op_dispatcher(call, sctx):
- return run_dispatch(call, sctx)
diff --git a/python/tvm/tirx/tile/instruction.py
b/python/tvm/tirx/tile/instruction.py
index 755a409fb1..92082203cf 100644
--- a/python/tvm/tirx/tile/instruction.py
+++ b/python/tvm/tirx/tile/instruction.py
@@ -76,13 +76,11 @@ class Instruction:
def register(self):
if self.name in SPECS:
return SPECS[self.name]._factory
+ from .dispatch import run_dispatch
+
register_op_attr(self.name, "TCallEffectKind", 3)
register_op_attr(self.name, "TFixedReturnType", PrimType("void"))
- register_op_attr(
- self.name,
- "TIRxOpCategory",
- "tile_composite" if ".compose." in self.name else "tile_primitive",
- )
+ register_op_attr(self.name, "FTileOpDispatch", run_dispatch)
register_op_attr(self.name, "TScriptPrinterName", self.name)
op = Op.get(self.name)
op.set_signature(args=[arg.name for arg in self.operands])
@@ -367,11 +365,6 @@ class TensorCall:
return Call(self.op, args, attrs=self.call.attrs, ty=self.call.ty,
loc=self.call.loc)
-@tvm_ffi.register_global_func("tirx.TensorCallScope")
-def _scope(call):
- return ExecScope(SPECS[call.op.name].backend, str(call.attrs.scope))
-
-
for _name in (
"ScopeAttrs",
"MemoryAttrs",
diff --git a/src/backend/cuda/codegen/intrin_rule_cuda.cc
b/src/backend/cuda/codegen/intrin_rule_cuda.cc
index e2c5dcc5a1..4bbbccf694 100644
--- a/src/backend/cuda/codegen/intrin_rule_cuda.cc
+++ b/src/backend/cuda/codegen/intrin_rule_cuda.cc
@@ -290,7 +290,6 @@ void RegisterCudaIntrinRules() {
sig::arg<IntExpr>("lane", "The source thread id."),
sig::arg<IntExpr>("width", "The warp thread width, must be a power
of 2."))
.set_attr<FInferType>(tvm::op_attr::kInferType,
FInferType::FromNative<&InferTypeReturnArgType<1>>())
- .set_attr<tirx::TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("cuda"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.cuda.__shfl_sync"))
@@ -304,7 +303,6 @@ void RegisterCudaIntrinRules() {
sig::arg<IntExpr>("delta", "The source lane id offset to be added."),
sig::arg<IntExpr>("width", "The warp thread width, must be a power
of 2."))
.set_attr<FInferType>(tvm::op_attr::kInferType,
FInferType::FromNative<&InferTypeReturnArgType<1>>())
- .set_attr<tirx::TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("cuda"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.cuda.__shfl_up_sync"))
@@ -318,7 +316,6 @@ void RegisterCudaIntrinRules() {
sig::arg<IntExpr>("delta", "The source lane id offset to be
subtracted."),
sig::arg<IntExpr>("width", "The warp thread width, must be a power
of 2."))
.set_attr<FInferType>(tvm::op_attr::kInferType,
FInferType::FromNative<&InferTypeReturnArgType<1>>())
- .set_attr<tirx::TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("cuda"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.cuda.__shfl_down_sync"))
@@ -332,7 +329,6 @@ void RegisterCudaIntrinRules() {
sig::arg<IntExpr>("lane_mask", "The lane mask."),
sig::arg<IntExpr>("width", "The warp thread width, must be a power
of 2."))
.set_attr<FInferType>(tvm::op_attr::kInferType,
FInferType::FromNative<&InferTypeReturnArgType<1>>())
- .set_attr<tirx::TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("cuda"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.cuda.__shfl_xor_sync"))
@@ -342,7 +338,6 @@ void RegisterCudaIntrinRules() {
OpDef("tirx.cuda.__activemask")
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::UInt(32))
- .set_attr<tirx::TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("cuda"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.cuda.__activemask"))
diff --git a/src/backend/cuda/op/iket.cc b/src/backend/cuda/op/iket.cc
index e4c65a9444..24e12b5bc5 100644
--- a/src/backend/cuda/op/iket.cc
+++ b/src/backend/cuda/op/iket.cc
@@ -41,7 +41,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.iket_mark")
.signature(sig::arg("name", "The name."), sig::var_args("args"))
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Void())
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("cuda"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
@@ -57,7 +56,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.iket_range_start")
.signature(sig::arg("name", "The name."), sig::var_args("args"))
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::UInt(32))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("cuda"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
@@ -73,7 +71,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.iket_range_end")
.signature(sig::arg<IntExpr>("token", "The token."),
sig::var_args("args"))
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Void())
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("cuda"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
@@ -89,7 +86,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.iket_range_push")
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Void())
.signature(sig::arg("name", "The name."), sig::var_args("args"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("cuda"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
@@ -104,7 +100,6 @@ const Op& iket_range_pop_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.iket_range_pop")
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Void())
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("cuda"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
@@ -120,7 +115,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.iket_sentinel_token")
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::UInt(32))
.signature(sig::arg("name", "The name."))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("cuda"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
@@ -132,7 +126,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::UInt(32))
.signature(sig::arg("event_id", "The event identifier."),
sig::arg("source_code", "The source code."),
sig::var_args("args"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("cuda"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
diff --git a/src/backend/cuda/op/target_builtin.cc
b/src/backend/cuda/op/target_builtin.cc
index 0f893fe6be..21c6f292d2 100644
--- a/src/backend/cuda/op/target_builtin.cc
+++ b/src/backend/cuda/op/target_builtin.cc
@@ -106,7 +106,6 @@ void RegisterCudaTargetBuiltins() {
true)
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.cuda.decl_tmem"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
@@ -122,7 +121,6 @@ void RegisterCudaTargetBuiltins() {
sig::arg<IntExpr>("index_c", "The C fragment index."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.cuda.bmma_sync"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -137,7 +135,6 @@ void RegisterCudaTargetBuiltins() {
sig::arg<IntExpr>("cp_size", "The copy size."),
sig::var_args("args"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("s_tir"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
@@ -152,7 +149,6 @@ void RegisterCudaTargetBuiltins() {
sig::arg<IntExpr>("dst_stride", "The destination stride."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.cuda.mma_store"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -162,7 +158,6 @@ void RegisterCudaTargetBuiltins() {
sig::arg<IntExpr>("offset", "The offset."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.cuda.mma_fill"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -174,7 +169,6 @@ void RegisterCudaTargetBuiltins() {
sig::arg<IntExpr>("group_id", "The group identifier."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.cuda.timer_init"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -187,7 +181,6 @@ void RegisterCudaTargetBuiltins() {
sig::arg("leader_cond", "The leader condition."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.cuda.timer_start"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -200,7 +193,6 @@ void RegisterCudaTargetBuiltins() {
sig::arg("leader_cond", "The leader condition."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.cuda.timer_end"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -212,7 +204,6 @@ void RegisterCudaTargetBuiltins() {
sig::arg("leader_cond", "The leader condition."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.cuda.timer_finalize"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -231,9 +222,8 @@ TVM_FFI_NO_INLINE void RegisterDeviceIntrinsicAttrs(OpDef&
def, const char* op_n
printer_name.compare(printer_name.size() - 5, 5, "_warp") == 0))) {
printer_name[printer_name.rfind('_')] = '.';
}
- def.set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
-
.set_attr<TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
- ffi::String(op_namespace))
+
def.set_attr<TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
+ ffi::String(op_namespace))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(effect_kind));
if (std::string(op_namespace) != "cuda") {
def.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
diff --git a/src/backend/cuda/op/tcgen05.cc b/src/backend/cuda/op/tcgen05.cc
index 35380e8d5a..11327e6198 100644
--- a/src/backend/cuda/op/tcgen05.cc
+++ b/src/backend/cuda/op/tcgen05.cc
@@ -73,7 +73,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.tcgen05_encode_instr_descriptor")
.signature(sig::arg("desc"),
sig::call_attrs<backend::cuda::TCGen05InstrDescriptorAttrs>())
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("cuda"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
@@ -82,7 +81,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.tcgen05_encode_instr_descriptor_block_scaled")
.signature(sig::arg("desc"),
sig::call_attrs<backend::cuda::TCGen05InstrDescriptorBlockScaledAttrs>())
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("cuda"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
@@ -91,7 +89,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.tcgen05_encode_matrix_descriptor")
.signature(sig::arg("desc"), sig::arg("addr"), sig::arg<IntExpr>("ldo"),
sig::arg<IntExpr>("sdo"), sig::arg<IntExpr>("swizzle"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("cuda"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
diff --git a/src/backend/cuda/op/tensormap.cc b/src/backend/cuda/op/tensormap.cc
index a2e48ec4b4..b2256fd538 100644
--- a/src/backend/cuda/op/tensormap.cc
+++ b/src/backend/cuda/op/tensormap.cc
@@ -66,7 +66,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::var_args<PrimExpr>("args"),
sig::call_attrs<TensorMapEncodeTiledAttr>())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.cuda.tensormap_encode_tiled"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
diff --git a/src/backend/metal/codegen/intrin_rule_metal.cc
b/src/backend/metal/codegen/intrin_rule_metal.cc
index a3cb9cbaf1..3b2bc27b34 100644
--- a/src/backend/metal/codegen/intrin_rule_metal.cc
+++ b/src/backend/metal/codegen/intrin_rule_metal.cc
@@ -162,7 +162,6 @@ void RegisterMetalIntrinRules() {
.signature(sig::arg("var", "The variable to sync."),
sig::arg<IntExpr>("lane", "The source thread id."))
.set_attr<FInferType>(tvm::op_attr::kInferType,
FInferType::FromNative<&InferTypeReturnArgType<0>>())
- .set_attr<tirx::TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("metal"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.metal.simd_shuffle"))
@@ -173,7 +172,6 @@ void RegisterMetalIntrinRules() {
.signature(sig::arg("var", "The variable to sync."),
sig::arg<IntExpr>("delta", "The source lane id offset to be added."))
.set_attr<FInferType>(tvm::op_attr::kInferType,
FInferType::FromNative<&InferTypeReturnArgType<0>>())
- .set_attr<tirx::TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("metal"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.metal.simd_shuffle_up"))
@@ -184,7 +182,6 @@ void RegisterMetalIntrinRules() {
.signature(sig::arg("var", "The variable to sync."),
sig::arg<IntExpr>("delta", "The source lane id offset to be
subtracted."))
.set_attr<FInferType>(tvm::op_attr::kInferType,
FInferType::FromNative<&InferTypeReturnArgType<0>>())
- .set_attr<tirx::TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("metal"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.metal.simd_shuffle_down"))
diff --git a/src/backend/metal/op/cooperative_tensor.cc
b/src/backend/metal/op/cooperative_tensor.cc
index d003dedb71..d3946a74a9 100644
--- a/src/backend/metal/op/cooperative_tensor.cc
+++ b/src/backend/metal/op/cooperative_tensor.cc
@@ -43,7 +43,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("cols", "The number of columns."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.metal.cooperative_tensor_fill"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -67,7 +66,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("operand_role", "The matrix operand role."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.metal.cooperative_tensor_load"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -91,7 +89,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("operand_role", "The matrix operand role."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.metal.cooperative_tensor_store"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -116,7 +113,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TScriptPrinterName>(
tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.metal.cooperative_tensor_multiply_accumulate"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
diff --git a/src/backend/metal/op/target_builtin.cc
b/src/backend/metal/op/target_builtin.cc
index 6f375b3785..4273304937 100644
--- a/src/backend/metal/op/target_builtin.cc
+++ b/src/backend/metal/op/target_builtin.cc
@@ -47,7 +47,6 @@ void RegisterMetalTargetBuiltins() {
sig::arg<IntExpr>("row", "The row index."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.metal.make_filled_simdgroup_matrix"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -60,7 +59,6 @@ void RegisterMetalTargetBuiltins() {
sig::arg("transpose_matrix", "Whether to transpose the
matrix."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.metal.simdgroup_load"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -73,7 +71,6 @@ void RegisterMetalTargetBuiltins() {
sig::arg("transpose_matrix", "Whether to transpose the
matrix."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.metal.simdgroup_store"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -86,7 +83,6 @@ void RegisterMetalTargetBuiltins() {
sig::arg("c", "The C operand."), sig::arg<IntExpr>("index_c", "The C
fragment index."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.metal.simdgroup_multiply_accumulate"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
diff --git a/src/backend/opencl/op/memory.cc b/src/backend/opencl/op/memory.cc
index 5d88c68a84..bd4ca4cb22 100644
--- a/src/backend/opencl/op/memory.cc
+++ b/src/backend/opencl/op/memory.cc
@@ -39,7 +39,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("z", "The third input value."),
sig::arg<IntExpr>("channel_size", "The number of channels."),
sig::arg<PrimExpr>("value", "The value to use."))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TVectorizable>(tvm::tirx::op_attr::kVectorizable, true)
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -57,7 +56,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("z", "The third input value."),
sig::arg<IntExpr>("channel_size", "The number of channels."),
sig::arg<PrimExpr>("element_index", "The element index within
a texture channel."))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TVectorizable>(tvm::tirx::op_attr::kVectorizable, true)
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -73,7 +71,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("storage_scope", "The storage scope."),
sig::arg<IntExpr>("ndim", "The number of dimensions."),
sig::arg("shape", "The shape."), sig::var_args("args"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
diff --git a/src/backend/trn/op/target_builtin.cc
b/src/backend/trn/op/target_builtin.cc
index 9295bb01d1..83e44d41ee 100644
--- a/src/backend/trn/op/target_builtin.cc
+++ b/src/backend/trn/op/target_builtin.cc
@@ -56,7 +56,6 @@ void RegisterTRNTargetBuiltins() {
.signature()
.set_attr<FRegionGetBodyParams>(tvm::op_attr::kRegionGetBodyParams,
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>())
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("nki"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
@@ -81,7 +80,6 @@ TVM_FFI_NO_INLINE NKIIntrinsicNames
MakeNKIIntrinsicNames(const char* op_name) {
TVM_FFI_NO_INLINE void RegisterNKIIntrinsicAttrs(OpDef& def, const
std::string& printer_name) {
def.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Void())
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("nki"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
diff --git a/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
b/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
index ad2910a96b..70565ce6d2 100644
--- a/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
+++ b/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
@@ -188,7 +188,6 @@ void RegisterWebGPUIntrinRules() {
.signature(sig::arg("var", "The variable to sync."),
sig::arg<IntExpr>("lane", "The source thread id."))
.set_attr<FInferType>(tvm::op_attr::kInferType,
FInferType::FromNative<&InferTypeReturnArgType<0>>())
- .set_attr<tirx::TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("webgpu"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.webgpu.subgroup_shuffle"))
@@ -199,7 +198,6 @@ void RegisterWebGPUIntrinRules() {
.signature(sig::arg("var", "The variable to sync."),
sig::arg<IntExpr>("delta", "The source lane id offset to be added."))
.set_attr<FInferType>(tvm::op_attr::kInferType,
FInferType::FromNative<&InferTypeReturnArgType<0>>())
- .set_attr<tirx::TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("webgpu"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.webgpu.subgroup_shuffle_up"))
@@ -210,7 +208,6 @@ void RegisterWebGPUIntrinRules() {
.signature(sig::arg("var", "The variable to sync."),
sig::arg<IntExpr>("delta", "The source lane id offset to be
subtracted."))
.set_attr<FInferType>(tvm::op_attr::kInferType,
FInferType::FromNative<&InferTypeReturnArgType<0>>())
- .set_attr<tirx::TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>(tvm::tirx::op_attr::kDeviceIntrinsicNamespace,
ffi::String("webgpu"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.webgpu.subgroup_shuffle_down"))
diff --git a/src/tirx/op/abi.cc b/src/tirx/op/abi.cc
index 6eab069826..1dd16b6670 100644
--- a/src/tirx/op/abi.cc
+++ b/src/tirx/op/abi.cc
@@ -64,7 +64,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("func_name", "The function name."),
sig::var_args("args"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.call_extern"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -79,7 +78,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("func_name", "The function name."),
sig::var_args("args"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.call_pure_extern"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
@@ -94,7 +92,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg<IntExpr>("intrin_id", "The intrinsic identifier."),
sig::var_args("args"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.call_llvm_intrin"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -109,7 +106,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg<IntExpr>("intrin_id", "The intrinsic identifier."),
sig::var_args("args"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.call_llvm_pure_intrin"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>(tvm::tirx::op_attr::kVectorizable, true);
@@ -124,7 +120,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.call_spirv_pure_glsl450")
.signature(sig::arg<IntExpr>("intrin_id", "The intrinsic identifier."),
sig::var_args<PrimExpr>("args"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
@@ -140,7 +135,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("field", "The field index."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.abi_field_get"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kReadState));
}
@@ -158,7 +152,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("value", "The value to use."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.abi_field_set"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kUpdateState));
}
@@ -173,7 +166,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Void())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.throw_last_error"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -191,7 +183,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("num", "The number of entries."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.stack_alloca"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -207,7 +198,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::var_args<IntExpr>("args"))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.stack_make_shape"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -227,7 +217,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("elem_offset", "The element offset."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.stack_make_dltensor"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -241,7 +230,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.call_packed")
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Int(32))
.signature(sig::arg("func_name", "The function name."),
sig::var_args("args"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
@@ -275,7 +263,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Int(32))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.call_ffi_kernel"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -289,7 +276,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.call_cpacked")
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Int(32))
.signature(sig::arg("func_name", "The function name."),
sig::var_args("args"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
@@ -308,7 +294,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("args_stack", "The argument stack."),
sig::arg<IntExpr>("begin", "The start index."),
sig::arg<IntExpr>("end", "The end index."))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
@@ -327,7 +312,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("args_stack", "The argument stack."),
sig::arg<IntExpr>("begin", "The start index."),
sig::arg<IntExpr>("end", "The end index."))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
@@ -344,7 +328,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("dtype_bits_hint", "The data type bit-width
hint."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.alloc_workspace"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TGlobalSymbol>(tvm::tirx::op_attr::kGlobalSymbol,
"TVMBackendAllocWorkspace")
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -358,7 +341,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("ptr", "The pointer."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.free_workspace"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TGlobalSymbol>(tvm::tirx::op_attr::kGlobalSymbol,
"TVMBackendFreeWorkspace")
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
diff --git a/src/tirx/op/annotation.cc b/src/tirx/op/annotation.cc
index ea2ececdfc..7cacbe3519 100644
--- a/src/tirx/op/annotation.cc
+++ b/src/tirx/op/annotation.cc
@@ -42,7 +42,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Void())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.kernel_replace_point"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -58,7 +57,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg<PrimExpr>("predicate", "The predicate."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.ignore_loop_partition"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
diff --git a/src/tirx/op/cpu.cc b/src/tirx/op/cpu.cc
index d53af22a3e..3a26a810f0 100644
--- a/src/tirx/op/cpu.cc
+++ b/src/tirx/op/cpu.cc
@@ -43,7 +43,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Void())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.cpu_parallel_barrier"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
diff --git a/src/tirx/op/gpu.cc b/src/tirx/op/gpu.cc
index 65cd6f9839..080147ef51 100644
--- a/src/tirx/op/gpu.cc
+++ b/src/tirx/op/gpu.cc
@@ -54,7 +54,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Void())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.gpu_thread_return"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kControlJump));
}
@@ -71,7 +70,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<PrimExpr>("pred", "The predicate."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.gpu_thread_filter"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
@@ -89,7 +87,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
FInferType::FromNative<&InferTypeReturnArgType<0>>())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.gpu_active_thread_selector"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -106,7 +103,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg<PrimExpr>("cond", "The condition."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.gpu_thread_invariant"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
@@ -122,7 +118,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Void())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.gpu_storage_sync"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -143,7 +138,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
FInferType::FromNative<&InferTypeReturnArgType<1>>())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.gpu_warp_shuffle"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -163,7 +157,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
FInferType::FromNative<&InferTypeReturnArgType<1>>())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.gpu_warp_shuffle_up"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -183,7 +176,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
FInferType::FromNative<&InferTypeReturnArgType<1>>())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.gpu_warp_shuffle_down"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -204,7 +196,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
FInferType::FromNative<&InferTypeReturnArgType<1>>())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.gpu_warp_shuffle_xor"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -219,7 +210,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::UInt(32))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.gpu_warp_activemask"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -240,7 +230,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Void())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.gpu_thread_allreduce"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -258,7 +247,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Int(32))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.gpu_dp4a"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
@@ -273,7 +261,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("stride", "The stride."), sig::arg("layout", "The
layout."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.gpu_load_matrix_sync"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kReadState));
}
@@ -291,7 +278,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("index_c", "The C fragment index."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.gpu_mma_sync"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -305,7 +291,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("index", "The index."), sig::arg("value", "The
value to use."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.gpu_fill_fragment"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -320,7 +305,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("stride", "The stride."), sig::arg("layout", "The
layout."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.gpu_store_matrix_sync"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
diff --git a/src/tirx/op/memory.cc b/src/tirx/op/memory.cc
index ed590a2937..7ab9034dd8 100644
--- a/src/tirx/op/memory.cc
+++ b/src/tirx/op/memory.cc
@@ -175,7 +175,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.reinterpret"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
@@ -191,7 +190,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FInferType>(tvm::op_attr::kInferType,
FInferType::FromNative<&InferTypeAddressOf>())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.address_of"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
@@ -207,7 +205,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.isnullptr"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
@@ -222,7 +219,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("ptr", "The pointer."), sig::arg<IntExpr>("rw", "The
read/write mode."),
sig::arg<IntExpr>("locality", "The locality hint."),
sig::arg<IntExpr>("cache_type", "The cache policy."))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -238,7 +234,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("byte_offset", "Offset in bytes."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.ptr_byte_offset"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
@@ -250,7 +245,6 @@ const Op& static_handle_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.static_handle")
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kSpecialCallArg));
}
@@ -267,7 +261,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ffi::String("tirx.handle_add_byte_offset"))
.set_attr<FInferType>(tvm::op_attr::kInferType,
FInferType::FromNative<&InferTypePointerOffset<0,
-1>>())
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
@@ -280,7 +273,6 @@ const Op& atomic_add_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.atomic_add")
.signature(sig::arg("ptr", "The pointer."), sig::arg("value", "The value
to use."))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -297,7 +289,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Void())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.assume_aligned"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kEmbedInfo));
}
@@ -312,7 +303,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Int(32))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.undef"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kReadState));
}
@@ -330,7 +320,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
FInferType::FromNative<&InferTypeMaskedLoad>())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.masked_load"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kReadState));
}
@@ -348,7 +337,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TFixedReturnType>(tvm::op_attr::kFixedReturnType,
PrimType::Void())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.masked_store"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kUpdateState));
}
@@ -360,7 +348,6 @@ const Op& alloc_tensor_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.alloc_tensor")
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<FInferType>(tvm::op_attr::kInferType,
FInferType::FromNative<&InferTypeTensor<true>>())
.set_validator(ffi::reflection::NativeFunctionView<void(
@@ -382,7 +369,6 @@ const Op& decl_tensor_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.decl_tensor")
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<FInferType>(tvm::op_attr::kInferType,
FInferType::FromNative<&InferTypeTensor<false>>())
.set_validator(ffi::reflection::NativeFunctionView<void(
@@ -406,7 +392,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("tensor", "The tensor variable."))
.set_attr<FInferType>(tvm::op_attr::kInferType,
FInferType::FromNative<&InferTypeTensorDataPtr>())
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
diff --git a/src/tirx/op/prim_metadata.cc b/src/tirx/op/prim_metadata.cc
index 6b0cca62b5..8bee3ac278 100644
--- a/src/tirx/op/prim_metadata.cc
+++ b/src/tirx/op/prim_metadata.cc
@@ -30,223 +30,179 @@ namespace tirx {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("prim.likely")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.likely"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.likely"));
OpDef("prim.if_then_else")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.if_then_else"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.if_then_else"));
OpDef("prim.vscale")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.vscale"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.vscale"));
OpDef("prim.ceil")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.ceil"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.ceil"));
OpDef("prim.log2")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.log2"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.log2"));
OpDef("prim.clz")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.clz"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.clz"));
OpDef("prim.pow")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.pow"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.pow"));
OpDef("prim.fabs")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.fabs"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.fabs"));
OpDef("prim.fmod")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.fmod"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.fmod"));
OpDef("prim.floor")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.floor"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.floor"));
OpDef("prim.round")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.round"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.round"));
OpDef("prim.nearbyint")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.nearbyint"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.nearbyint"));
OpDef("prim.trunc")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.trunc"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.trunc"));
OpDef("prim.exp")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.exp"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.exp"));
OpDef("prim.exp2")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.exp2"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.exp2"));
OpDef("prim.exp10")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.exp10"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.exp10"));
OpDef("prim.erf")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.erf"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.erf"));
OpDef("prim.tanh")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.tanh"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.tanh"));
OpDef("prim.sigmoid")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.sigmoid"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.sigmoid"));
OpDef("prim.sqrt")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.sqrt"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.sqrt"));
OpDef("prim.rsqrt")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.rsqrt"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.rsqrt"));
OpDef("prim.log")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.log"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.log"));
OpDef("prim.log1p")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.log1p"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.log1p"));
OpDef("prim.log10")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.log10"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.log10"));
OpDef("prim.tan")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.tan"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.tan"));
OpDef("prim.cos")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.cos"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.cos"));
OpDef("prim.cosh")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.cosh"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.cosh"));
OpDef("prim.sin")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.sin"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.sin"));
OpDef("prim.sinh")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.sinh"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.sinh"));
OpDef("prim.asin")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.asin"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.asin"));
OpDef("prim.acos")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.acos"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.acos"));
OpDef("prim.atan")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.atan"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.atan"));
OpDef("prim.acosh")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.acosh"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.acosh"));
OpDef("prim.asinh")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.asinh"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.asinh"));
OpDef("prim.atanh")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.atanh"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.atanh"));
OpDef("prim.atan2")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.atan2"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.atan2"));
OpDef("prim.nextafter")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.nextafter"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.nextafter"));
OpDef("prim.hypot")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.hypot"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.hypot"));
OpDef("prim.copysign")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.copysign"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.copysign"));
OpDef("prim.ldexp")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.ldexp"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.ldexp"));
OpDef("prim.isnan")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.isnan"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.isnan"));
OpDef("prim.popcount")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.popcount"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.popcount"));
OpDef("prim.fma")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.fma"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.fma"));
OpDef("prim.assume")
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
- ffi::String("tirx.assume"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+ ffi::String("tirx.assume"));
}
} // namespace tirx
diff --git a/src/tirx/op/region.cc b/src/tirx/op/region.cc
index cbe8be5bc0..42a2a1aec5 100644
--- a/src/tirx/op/region.cc
+++ b/src/tirx/op/region.cc
@@ -62,8 +62,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRegionGetBodyParams>(tvm::op_attr::kRegionGetBodyParams,
FRegionGetBodyParams::FromNative<&LaunchThreadBodyParams>())
.set_attr<FRegionValidate>(tvm::op_attr::kRegionValidate,
-
FRegionValidate::FromNative<&ValidateBuiltinRegion>())
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+
FRegionValidate::FromNative<&ValidateBuiltinRegion>());
}
const Op& device_entry_op() {
@@ -91,8 +90,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRegionGetBodyParams>(tvm::op_attr::kRegionGetBodyParams,
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>())
.set_attr<FRegionValidate>(tvm::op_attr::kRegionValidate,
-
FRegionValidate::FromNative<&ValidateDeviceEntry>())
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+
FRegionValidate::FromNative<&ValidateDeviceEntry>());
}
const Op& device_context_op() {
@@ -108,8 +106,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRegionGetBodyParams>(tvm::op_attr::kRegionGetBodyParams,
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>())
.set_attr<FRegionValidate>(tvm::op_attr::kRegionValidate,
-
FRegionValidate::FromNative<&ValidateBuiltinRegion>())
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+
FRegionValidate::FromNative<&ValidateBuiltinRegion>());
}
const Op& compute_scope_op() {
@@ -125,8 +122,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRegionGetBodyParams>(tvm::op_attr::kRegionGetBodyParams,
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>())
.set_attr<FRegionValidate>(tvm::op_attr::kRegionValidate,
-
FRegionValidate::FromNative<&ValidateBuiltinRegion>())
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+
FRegionValidate::FromNative<&ValidateBuiltinRegion>());
}
const Op& parallel_launch_op() {
@@ -142,8 +138,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRegionGetBodyParams>(tvm::op_attr::kRegionGetBodyParams,
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>())
.set_attr<FRegionValidate>(tvm::op_attr::kRegionValidate,
-
FRegionValidate::FromNative<&ValidateBuiltinRegion>())
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+
FRegionValidate::FromNative<&ValidateBuiltinRegion>());
}
} // namespace tirx
diff --git a/src/tirx/op/vector.cc b/src/tirx/op/vector.cc
index 0070d08ca4..29661611d6 100644
--- a/src/tirx/op/vector.cc
+++ b/src/tirx/op/vector.cc
@@ -63,7 +63,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
FInferType::FromNative<&InferTypeVectorPart<false>>())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.vector_high"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
@@ -80,7 +79,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
FInferType::FromNative<&InferTypeVectorPart<false>>())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.vector_low"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
@@ -98,7 +96,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
FInferType::FromNative<&InferTypeVectorPart<true>>())
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.vector_combine"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
@@ -114,7 +111,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("limit", "The limit value."))
.set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
ffi::String("tirx.get_active_lane_mask"))
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
diff --git a/src/tirx/tile/dispatch.cc b/src/tirx/tile/dispatch.cc
index c0f8b65636..a5cccd7443 100644
--- a/src/tirx/tile/dispatch.cc
+++ b/src/tirx/tile/dispatch.cc
@@ -21,10 +21,13 @@
namespace tvm::tirx::tile {
Function Dispatch(const Call& call, const DispatchContext& context) {
- static auto dispatcher = ffi::Function::GetGlobal("tirx.f_op_dispatcher");
- TVM_FFI_ICHECK(dispatcher.has_value())
- << "Internal Error: tirx.f_op_dispatcher is not registered";
- Function implementation = dispatcher.value()(call, context).cast<Function>();
+ Op op = call->op.as_or_throw<Op>();
+ TVM_FFI_CHECK(Op::HasAttrMap(op_attr::kTileOpDispatch), ValueError)
+ << "No FTileOpDispatch registered for " << op->name;
+ static const auto& dispatchers =
Op::GetAttrMap<FTileOpDispatch>(op_attr::kTileOpDispatch);
+ TVM_FFI_CHECK(dispatchers.count(op), ValueError)
+ << "No FTileOpDispatch registered for " << op->name;
+ Function implementation = dispatchers[op](call, context);
TVM_FFI_ICHECK(implementation.defined()) << "TIRx dispatcher did not return
a Function";
return implementation;
}
diff --git a/src/tirx/tile/lower.cc b/src/tirx/tile/lower.cc
index 3df367cc0b..aaabf822d5 100644
--- a/src/tirx/tile/lower.cc
+++ b/src/tirx/tile/lower.cc
@@ -23,6 +23,7 @@
*/
#include <tvm/backend/cuda/attr.h>
+#include <tvm/ffi/reflection/accessor.h>
#include <tvm/ir/function.h>
#include <tvm/ir/op.h>
#include <tvm/ir/prim/op.h>
@@ -35,7 +36,6 @@
#include <tvm/tirx/op/gpu.h>
#include <tvm/tirx/op/memory.h>
#include <tvm/tirx/op/region.h>
-#include <tvm/tirx/op_attr_types.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/tile/dispatch.h>
@@ -64,11 +64,10 @@ class NoOpCallVerifier : public Verifier<NoOpCallVerifier> {
using Verifier::Visit;
void Dispatch_(const CallNode* call, ffi::reflection::AccessPath path) final
{
- if (auto op = call->op.as<Op>()) {
- static const auto& categories =
- Op::GetAttrMap<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory);
- auto category = categories.get(op.value(), ffi::String(""));
- Verify(category != "tile_primitive" && category != "tile_composite")
+ if (auto op = call->op.as<Op>(); op &&
Op::HasAttrMap(tile::op_attr::kTileOpDispatch)) {
+ static const auto& dispatchers =
+
Op::GetAttrMap<tile::FTileOpDispatch>(tile::op_attr::kTileOpDispatch);
+ Verify(!dispatchers.count(op.value()))
<< "Unlowered tensor instruction " << op.value()->name << " at " <<
path;
}
Verifier::Dispatch_(call, path);
@@ -383,15 +382,21 @@ class TileDispatcher : public StmtExprMutator {
UnchangedOr<Stmt> Mutate_(const EvaluateNode* stmt, InplaceMode
inplace_mode) final {
const auto* op = stmt->value.as<CallNode>();
- if (!op || !op->op.as<Op>()) return StmtExprMutator::Mutate_(stmt,
inplace_mode);
- static const auto& categories =
Op::GetAttrMap<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory);
- auto category = categories.get(op->op.as_or_throw<Op>(), ffi::String(""));
- if (category != "tile_primitive" && category != "tile_composite") {
+ if (!op || !op->op.as<Op>() ||
!Op::HasAttrMap(tile::op_attr::kTileOpDispatch)) {
+ return StmtExprMutator::Mutate_(stmt, inplace_mode);
+ }
+ static const auto& dispatchers =
+ Op::GetAttrMap<tile::FTileOpDispatch>(tile::op_attr::kTileOpDispatch);
+ if (!dispatchers.count(op->op.as_or_throw<Op>())) {
return StmtExprMutator::Mutate_(stmt, inplace_mode);
}
op->op.as_or_throw<Op>().Validate(op);
- static auto get_scope =
ffi::Function::GetGlobalRequired("tirx.TensorCallScope");
- ExecScope scope = get_scope(ffi::GetRef<Call>(op)).cast<ExecScope>();
+ TVM_FFI_CHECK(op->attrs.has_value(), ValueError)
+ << "A tile instruction must have scope attributes";
+ auto scope_name =
+ ffi::reflection::FieldGetter(op->attrs->GetTypeKey(),
"scope")(op->attrs.value())
+ .cast<ffi::String>();
+ ExecScope scope(target_->kind->name, scope_name);
ffi::Map<ffi::String, ffi::Array<PrimExpr>> inter_map, intra_map;
ffi::String scope_kind = scope->name();
if (cuda_analysis_) {
diff --git a/src/tirx/transform/split_host_device.cc
b/src/tirx/transform/split_host_device.cc
index 829a927c9b..79e1780a10 100644
--- a/src/tirx/transform/split_host_device.cc
+++ b/src/tirx/transform/split_host_device.cc
@@ -67,8 +67,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRegionGetBodyParams>(tvm::op_attr::kRegionGetBodyParams,
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>())
.set_attr<FRegionValidate>(tvm::op_attr::kRegionValidate,
-
FRegionValidate::FromNative<&ValidateDeviceScopeRegion>())
- .set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory,
ffi::String("builtin"));
+
FRegionValidate::FromNative<&ValidateDeviceScopeRegion>());
}
// Device-region annotation
diff --git a/tests/python/tirx/iket/test_iket_profiler.py
b/tests/python/tirx/iket/test_iket_profiler.py
index 15e046463b..d6e5e68de2 100644
--- a/tests/python/tirx/iket/test_iket_profiler.py
+++ b/tests/python/tirx/iket/test_iket_profiler.py
@@ -408,7 +408,7 @@ def test_public_interface_is_official_only():
)
def test_annotation_ops_are_cuda_owned(name):
op = tvm.ir.Op.get(f"tirx.cuda.iket_{name}")
- assert op.get_attr("TIRxOpCategory") == "device_intrin"
+ assert op.get_attr("FTileOpDispatch") is None
assert op.get_attr("TDeviceIntrinsicNamespace") == "cuda"
diff --git a/tests/python/tirx/test_op.py b/tests/python/tirx/test_op.py
index 396f8abbd8..7684758327 100644
--- a/tests/python/tirx/test_op.py
+++ b/tests/python/tirx/test_op.py
@@ -34,7 +34,7 @@ def test_instruction_is_a_void_call():
assert isinstance(call, Call)
assert call.ty == tvm.ir.PrimType("void")
assert call.op.get_attr("TCallEffectKind") == 3
- assert call.op.get_attr("TIRxOpCategory") == "tile_primitive"
+ assert callable(call.op.get_attr("FTileOpDispatch"))
assert call.attrs.scope == "warp"
assert call.attrs.vec_bits == 128
assert not hasattr(call, "config")
@@ -54,8 +54,7 @@ def test_roundtrip_and_exactly_one_evaluate(composite):
assert len(nodes) == 1
call = nodes[0].value
assert isinstance(call, Call)
- expected = "tile_composite" if composite else "tile_primitive"
- assert call.op.get_attr("TIRxOpCategory") == expected
+ assert callable(call.op.get_attr("FTileOpDispatch"))
parsed = tvm.script.from_source(func.script(), extra_vars={"T": T})
assert_structural_equal(func, parsed)
assert_structural_equal(func, load_json(save_json(func)))
diff --git a/tests/python/tirx/test_op_namespace_cleanup.py
b/tests/python/tirx/test_op_namespace_cleanup.py
index d6acb75cb7..dd07ea4589 100644
--- a/tests/python/tirx/test_op_namespace_cleanup.py
+++ b/tests/python/tirx/test_op_namespace_cleanup.py
@@ -36,7 +36,7 @@ def _tile_calls(func):
if (
isinstance(stmt, tvm.ir.Call)
and isinstance(stmt.op, Op)
- and stmt.op.get_attr("TIRxOpCategory") in {"tile_primitive",
"tile_composite"}
+ and stmt.op.get_attr("FTileOpDispatch") is not None
):
calls.append(stmt)
@@ -114,7 +114,7 @@ def
test_tensormap_encoding_accepts_attribute_keywords_and_defaults():
def test_kernel_replace_point_is_builtin_marker_not_tile_primitive():
- assert _op_attr("tirx.kernel_replace_point", "TIRxOpCategory") == "builtin"
+ assert _op_attr("tirx.kernel_replace_point", "FTileOpDispatch") is None
assert "tirx.tile.kernel_replace_point" not in Op.list_op_names()
assert hasattr(T, "kernel_replace_point")
assert not hasattr(T.cuda.tile, "kernel_replace_point")
@@ -177,7 +177,7 @@ def
test_device_intrinsic_namespaces_are_canonical_and_classified():
(call.op.name, _op_attr(call.op.name, "TDeviceIntrinsicNamespace"))
for call in calls
] == expected
for op_name, namespace in expected:
- assert _op_attr(op_name, "TIRxOpCategory") == "device_intrin"
+ assert _op_attr(op_name, "FTileOpDispatch") is None
assert _op_attr(op_name, "TDeviceIntrinsicNamespace") == namespace
@@ -296,7 +296,7 @@ def
test_device_intrinsic_printer_roundtrips_canonical_namespaces():
("tirx.metal.simd_shuffle_up", "metal"),
("tirx.metal.simd_shuffle_down", "metal"),
]:
- assert _op_attr(op_name, "TIRxOpCategory") == "device_intrin"
+ assert _op_attr(op_name, "FTileOpDispatch") is None
assert _op_attr(op_name, "TDeviceIntrinsicNamespace") == namespace
assert _op_attr(op_name, "TCallEffectKind") in (1, 3)
@@ -311,11 +311,8 @@ def
test_device_intrinsic_printer_roundtrips_canonical_namespaces():
assert_structural_equal(device_namespaces, reparsed)
-def test_registered_tirx_ops_have_exactly_one_category():
- if _op_attr("prim.sqrt", "TIRxOpCategory") is None:
- pytest.skip("TIRx op categories require a rebuilt C++ runtime")
-
- categories = {"builtin", "tile_primitive", "tile_composite",
"device_intrin"}
+def test_registered_tirx_ops_use_dispatch_attributes():
+ assert not Op.get("prim.sqrt").has_attr("TIRxOpCategory")
device_namespaces = {
"cuda",
"ptx",
@@ -351,38 +348,28 @@ def test_registered_tirx_ops_have_exactly_one_category():
"tirx.zero",
}
- missing = []
- invalid = []
lingering_flat_tile = []
for op_name in sorted(name for name in Op.list_op_names() if
name.startswith("tirx.")):
- category = _op_attr(op_name, "TIRxOpCategory")
+ dispatcher = _op_attr(op_name, "FTileOpDispatch")
device_namespace = _op_attr(op_name, "TDeviceIntrinsicNamespace")
- if category is None:
- missing.append(op_name)
- continue
- if category not in categories:
- invalid.append((op_name, category))
- continue
if op_name in flat_tile_only_names:
lingering_flat_tile.append(op_name)
- if category in {"tile_primitive", "tile_composite"}:
+ if dispatcher is not None:
+ assert callable(dispatcher), op_name
if not op_name.startswith(("tirx.cuda.tile.", "tirx.trn.tile.")):
lingering_flat_tile.append(op_name)
assert device_namespace is None, op_name
- elif category == "device_intrin":
+ elif device_namespace is not None:
assert device_namespace in device_namespaces, op_name
printer_name = _op_attr(op_name, "TScriptPrinterName")
assert printer_name is not None, op_name
assert printer_name.startswith("tirx." + device_namespace + "."),
op_name
assert _has_path(T, printer_name.removeprefix("tirx.")), op_name
else:
- assert category == "builtin"
- assert device_namespace is None, op_name
+ assert not op_name.startswith(("tirx.cuda.tile.",
"tirx.trn.tile.")), op_name
- assert not missing
- assert not invalid
assert not lingering_flat_tile
diff --git a/tests/python/tirx/tile/cuda/gemm/test_gemm_mma_m16n8k_.py
b/tests/python/tirx/tile/cuda/gemm/test_gemm_mma_m16n8k_.py
index 337912519e..1b56c1684e 100644
--- a/tests/python/tirx/tile/cuda/gemm/test_gemm_mma_m16n8k_.py
+++ b/tests/python/tirx/tile/cuda/gemm/test_gemm_mma_m16n8k_.py
@@ -379,7 +379,7 @@ def _lower(func):
def test_cuda_mma_sync_is_a_registered_instruction():
op = tvm.ir.Op.get("tirx.cuda.tile.mma_sync")
- assert op.get_attr("TIRxOpCategory") == "tile_primitive"
+ assert callable(op.get_attr("FTileOpDispatch"))
assert op.get_attr("TCallEffectKind") == 3
diff --git a/tests/python/tirx/tile/test_dispatcher.py
b/tests/python/tirx/tile/test_dispatcher.py
index 45364cd96a..56af11f507 100644
--- a/tests/python/tirx/tile/test_dispatcher.py
+++ b/tests/python/tirx/tile/test_dispatcher.py
@@ -20,8 +20,10 @@ from types import SimpleNamespace
import pytest
import tvm
+from tvm.ir import Call, Evaluate, Op, PrimType, Return, const, make_node,
register_op_attr
from tvm.script import tirx as T
-from tvm.tirx import decl_tensor
+from tvm.tirx import Function, decl_tensor
+from tvm.tirx.tile import DispatchContext
from tvm.tirx.tile.dispatch import run_dispatch
from tvm.tirx.tile.instruction import SPECS
@@ -77,3 +79,84 @@ def test_instruction_scope_contract(backend, scope):
raw = tvm.ir.Call(call.op, call.args, attrs=attrs, ty=call.ty)
with pytest.raises(ValueError, match=f"unsupported scope '{scope}'"):
raw.validate()
+
+
+def _custom_call(name):
+ op_name = f"test.tile_dispatch.{name}"
+ register_op_attr(op_name, "TFixedReturnType", PrimType("void"))
+ register_op_attr(op_name, "TCallEffectKind", 3)
+ op = Op.get(op_name)
+ op.set_signature(args=["token"])
+ op.set_attrs_type_key("tirx.tensor.ScopeAttrs")
+ return Call(op, [const(7)], attrs=make_node("tirx.tensor.ScopeAttrs"))
+
+
+def _lower(body):
+ func = Function([], body).with_attr("target", tvm.target.Target("cuda"))
+ return tvm.tirx.transform.TileDispatch()(tvm.IRModule({"main":
func}))["main"]
+
+
+def test_op_attribute_dispatches_without_python_instruction_spec():
+ # Initialize the pass's cached attribute view before adding a new Op.
+ _lower(Evaluate(0))
+ call = _custom_call("custom")
+ assert call.op.name not in SPECS
+ seen = []
+
+ @register_op_attr(call.op.name, "FTileOpDispatch")
+ def dispatch(actual_call, context):
+ assert actual_call.same_as(call)
+ assert isinstance(context, DispatchContext)
+ assert context.target.kind.name == "cuda"
+ assert context.exec_scope.backend == "cuda"
+ assert context.scope_kind == "thread"
+ seen.append(actual_call)
+ return Function([], Evaluate(23))
+
+ lowered = _lower(Evaluate(call))
+ assert len(seen) == 1
+ assert lowered.body.seq[0].value == 23
+
+ # The live native attribute map observes a replacement on the same Op.
+ register_op_attr(
+ call.op.name,
+ "FTileOpDispatch",
+ lambda actual_call, context: Function([], Evaluate(41)),
+ override=True,
+ )
+ assert _lower(Evaluate(call)).body.seq[0].value == 41
+
+
+def test_dispatch_recursively_expands_tile_implementations():
+ outer = _custom_call("outer")
+ inner = _custom_call("inner")
+ register_op_attr(
+ outer.op.name, "FTileOpDispatch", lambda call, context: Function([],
Evaluate(inner))
+ )
+ register_op_attr(
+ inner.op.name, "FTileOpDispatch", lambda call, context: Function([],
Evaluate(31))
+ )
+ assert _lower(Evaluate(outer)).body.seq[0].value == 31
+
+
+def test_op_without_dispatch_attribute_is_preserved():
+ call = _custom_call("ordinary")
+ assert call.op.get_attr("FTileOpDispatch") is None
+ assert _lower(Evaluate(call)).body.seq[0].value.same_as(call)
+
+
+def test_residual_tile_call_is_rejected():
+ call = _custom_call("residual")
+ register_op_attr(
+ call.op.name, "FTileOpDispatch", lambda call, context: Function([],
Evaluate(0))
+ )
+ # Dispatch expands statement-boundary Calls; other positions must be
rejected.
+ with pytest.raises(tvm.error.InternalError, match="Failed to lower the
TIRx program"):
+ _lower(Return(call))
+
+
+def test_dispatch_requires_function_result():
+ call = _custom_call("invalid_result")
+ register_op_attr(call.op.name, "FTileOpDispatch", lambda call, context:
Evaluate(0))
+ with pytest.raises(TypeError, match="tirx.Function"):
+ _lower(Evaluate(call))
