This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new eab49d73f [FEAT] Expose registered ops and types at definition sites
(#96)
eab49d73f is described below
commit eab49d73fc70fe825c7e84a5570843285dd5019d
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 27 15:39:34 2026 -0400
[FEAT] Expose registered ops and types at definition sites (#96)
Add definition-site script exposure for registered operators and IR
types through `Op::ExposeScript`, `Type::ExposeScript`, and canonical
Rust and Python completion hooks.
Stable Python receivers defer early native requests until the script
builder initializer connects the exposure callbacks. Completed
definitions publish operator exports and type aliases without global
registry scans, preserving explicit builders and raw backend operator
construction. Native helpers cache an optional receiver on first use so
standalone compiler loading also works without Python callbacks.
---
include/tvm/ir/op.h | 17 +-
include/tvm/ir/type.h | 25 +++
python/tvm_tirx/backend/cuda/op.py | 23 +-
python/tvm_tirx/backend/cuda/ptx/engine.py | 2 +
python/tvm_tirx/backend/cuda/script.py | 9 +-
python/tvm_tirx/backend/opencl/__init__.py | 14 +-
python/tvm_tirx/backend/opencl/script.py | 8 +-
python/tvm_tirx/base.py | 27 +++
python/tvm_tirx/ir/op.py | 73 +------
python/tvm_tirx/ir/type.py | 57 -----
python/tvm_tirx/script/ir_builder/__init__.py | 15 +-
python/tvm_tirx/script/ir_builder/ir.py | 13 +-
python/tvm_tirx/tirx/script/ir_builder/__init__.py | 38 +++-
python/tvm_tirx/tirx/script/ir_builder/ir.py | 5 +-
python/tvm_tirx/tirx/script/ir_builder/op.py | 83 +++++++-
python/tvm_tirx/tirx/tile/instruction.py | 1 +
rust/src/ir/op_def.rs | 7 +-
rust/src/ir/type_def.rs | 5 +-
rust/src/lib.rs | 5 +-
src/backend/cuda/op/iket.cc | 21 +-
src/backend/cuda/op/legacy.cc | 12 +-
src/backend/cuda/op/target_builtin.cc | 234 ++++++++++++++-------
src/backend/cuda/op/tcgen05.cc | 9 +-
src/backend/cuda/op/tensormap.cc | 3 +-
src/backend/opencl/op/memory.cc | 9 +-
src/ir/op.cc | 8 +
src/ir/prim/builtin.cc | 18 +-
src/ir/prim/math.cc | 114 ++++++----
src/ir/type.cc | 17 ++
src/tirx/ir/type.cc | 2 +
src/tirx/op/abi.cc | 54 +++--
src/tirx/op/annotation.cc | 9 +-
src/tirx/op/cpu.cc | 3 +-
src/tirx/op/gpu.cc | 48 +++--
src/tirx/op/memory.cc | 45 ++--
src/tirx/op/region.cc | 18 +-
src/tirx/op/vector.cc | 12 +-
37 files changed, 704 insertions(+), 359 deletions(-)
diff --git a/include/tvm/ir/op.h b/include/tvm/ir/op.h
index edf5992a7..120d9a366 100644
--- a/include/tvm/ir/op.h
+++ b/include/tvm/ir/op.h
@@ -129,7 +129,8 @@ class OpObj : public ExprObj {
* TVM_FFI_STATIC_INIT_BLOCK() {
* OpDef("example.identity", "Return the input expression.")
* .signature(sig::arg("value", "The input expression."))
- * .set_attr(kPure, true);
+ * .set_attr(kPure, true)
+ * .expose_script();
* }
* // Copies the handle; the canonical node is shared.
* Op op = Op::Get("example.identity");
@@ -168,6 +169,15 @@ class Op : public Expr {
return HasAttrMap(key.name);
}
+ /*!
+ * \brief Request a registered operator's builder in its script namespace.
+ * \param op_name The canonical registered operator name.
+ * \note Call after its signature and attributes are complete. The first
call caches the optional
+ * script receiver, including absence for standalone native use.
Install the receiver before
+ * the first call to enable exposure; errors from an installed
receiver propagate.
+ */
+ TVM_TIRX_DLL static void ExposeScript(const ffi::String& op_name);
+
/*!
* \brief Look up a registered operator by name.
* \param op_name The canonical name.
@@ -472,6 +482,11 @@ class OpDef {
OpDef& operator=(OpDef&&) = delete;
/*! \brief Get the canonical operator. \return A copied handle, safe after
this builder dies. */
Op op() const { return op_; }
+ /*!
+ * \brief Request script exposure after this operator's signature and
attributes are complete.
+ * \note Does nothing if the optional receiver was absent on
Op::ExposeScript's first call.
+ */
+ void expose_script() const { Op::ExposeScript(op_->name); }
/*!
* \brief Append a required value-argument descriptor.
* \param name Argument name.
diff --git a/include/tvm/ir/type.h b/include/tvm/ir/type.h
index cadf78d0c..5f72f9cfc 100644
--- a/include/tvm/ir/type.h
+++ b/include/tvm/ir/type.h
@@ -74,6 +74,31 @@ class TypeObj : public ffi::Object {
*/
class Type : public ffi::ObjectRef {
public:
+ /*!
+ * \brief Request a registered IR type's constructor in its script namespace.
+ * \param type_key The registered type key.
+ * \note Call after ObjectDef finalization and all reflection and
constructor registration.
+ * The first call caches the optional script receiver, including
absence for standalone
+ * native use. Install the receiver before the first call to enable
exposure; errors from
+ * an installed receiver propagate.
+ */
+ TVM_TIRX_DLL static void ExposeScript(const ffi::String& type_key);
+
+ /*!
+ * \brief Request script exposure of a registered IR object type.
+ * \tparam TypeObj The object class, not its reference wrapper.
+ * \note Finish the ObjectDef statement and constructor registration before
calling.
+ * Does nothing if the optional receiver was absent on the string
helper's first call.
+ * \code
+ * ffi::reflection::ObjectDef<FooTypeObj>().def_ro("value",
&FooTypeObj::value);
+ * Type::ExposeScript<FooTypeObj>();
+ * \endcode
+ */
+ template <typename TypeObj>
+ static inline void ExposeScript() {
+ ExposeScript(TypeObj::_type_key);
+ }
+
/*! \brief Construct a MissingType for type information not yet populated. */
TVM_TIRX_DLL static Type Missing();
diff --git a/python/tvm_tirx/backend/cuda/op.py
b/python/tvm_tirx/backend/cuda/op.py
index dbebcd77e..e7b4e6b69 100644
--- a/python/tvm_tirx/backend/cuda/op.py
+++ b/python/tvm_tirx/backend/cuda/op.py
@@ -27,7 +27,7 @@ import tvm_ffi
from tvm_tirx import tirx
from tvm_tirx.ir import Attrs, Call, Op, StringImm, const
from tvm_tirx.ir.location import UNKNOWN_LOC, Location
-from tvm_tirx.ir.op import _init_op_api, _make_op_api
+from tvm_tirx.ir.op import _make_op_api
from tvm_tirx.runtime import MemoryScope
from tvm_tirx.tirx.op import bitwise_and, call_intrin, ptr_byte_offset
@@ -1605,4 +1605,23 @@ def tensormap_encode_tiled(
)
-_init_op_api("tirx.cuda", __name__)
+def __getattr__(name):
+ try:
+ op = Op.get(f"tirx.cuda.{name}")
+ except AttributeError as err:
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
from err
+ call = _make_op_api(op, __name__)
+ globals()[name] = call
+ return call
+
+
+# Copy the completed script exposures into the raw API without scanning the Op
registry.
+from tvm_tirx.tirx.script.ir_builder import op as _builder_op # noqa: E402
+
+for _name in vars(_builder_op._get_script_namespace("cuda")):
+ if _name not in globals():
+ try:
+ __getattr__(_name)
+ except AttributeError:
+ # Nested script namespaces are not raw operator callables.
+ pass
diff --git a/python/tvm_tirx/backend/cuda/ptx/engine.py
b/python/tvm_tirx/backend/cuda/ptx/engine.py
index c23a9a3c5..69054fa14 100644
--- a/python/tvm_tirx/backend/cuda/ptx/engine.py
+++ b/python/tvm_tirx/backend/cuda/ptx/engine.py
@@ -107,6 +107,7 @@ def register_table(table: dict[str, InstructionEntry]) ->
None:
get_global_func("script.printer.PTXCallDocTranslate")(),
)
register_codegen(f"ptx.{entry.name}")(_make_codegen(entry))
+ Op.get(entry.op_name).expose_script()
def register_addr() -> None:
@@ -124,6 +125,7 @@ def register_addr() -> None:
get_global_func("script.printer.PTXCallDocTranslate")(),
)
register_codegen("ptx.addr")(_unconsumed_addr_codegen)
+ op.expose_script()
def _unconsumed_addr_codegen(*_args):
diff --git a/python/tvm_tirx/backend/cuda/script.py
b/python/tvm_tirx/backend/cuda/script.py
index 1098ebe65..a2ea39638 100644
--- a/python/tvm_tirx/backend/cuda/script.py
+++ b/python/tvm_tirx/backend/cuda/script.py
@@ -161,7 +161,6 @@ __activemask = _auto_emit_void_call(_cuda_op.__activemask)
tile = _make_tensor_namespace()
tile.apply = _auto_emit_void_call(tile.apply)
tile.ewise = _auto_emit_void_call(tile.ewise)
-_ir.op._init_op_api("tirx.cuda", __name__, wrapper=_auto_emit_void_call)
class NVSHMEMNamespace:
@@ -247,3 +246,11 @@ class PTXNamespace(_PTXNamespace):
def _family(self, token):
chain = super()._family(token)
return None if chain is None else _ScriptInstrChain(chain._cands)
+
+
+# Preserve completed definition-site exports alongside the explicit builders.
+import sys as _sys # noqa: E402
+
+from tvm_tirx.tirx.script.ir_builder.op import register_script_namespace as
_register_namespace # noqa: E402
+
+_register_namespace("cuda", _sys.modules[__name__])
diff --git a/python/tvm_tirx/backend/opencl/__init__.py
b/python/tvm_tirx/backend/opencl/__init__.py
index 633d54059..397818c06 100644
--- a/python/tvm_tirx/backend/opencl/__init__.py
+++ b/python/tvm_tirx/backend/opencl/__init__.py
@@ -38,7 +38,6 @@ def _detect_target_from_device(dev):
def register_backend():
"""Register OpenCL-owned Python semantics."""
- from tvm_tirx.ir.op import _init_op_api
from tvm_tirx.target.detect_target import register_device_target_detector
from tvm_tirx.tirx.script.ir_builder import op as builder_op
@@ -56,9 +55,20 @@ def register_backend():
register_device_target_detector("opencl", _detect_target_from_device)
from . import script
- _init_op_api("tirx.opencl", __name__)
builder_op.register_script_namespace("opencl", script)
return None
+def __getattr__(name):
+ from tvm_tirx.ir.op import Op, _make_op_api
+
+ try:
+ op = Op.get(f"tirx.opencl.{name}")
+ except AttributeError as err:
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
from err
+ call = _make_op_api(op, __name__)
+ globals()[name] = call
+ return call
+
+
__all__ = ["register_backend"]
diff --git a/python/tvm_tirx/backend/opencl/script.py
b/python/tvm_tirx/backend/opencl/script.py
index 52712824f..16000c047 100644
--- a/python/tvm_tirx/backend/opencl/script.py
+++ b/python/tvm_tirx/backend/opencl/script.py
@@ -16,7 +16,9 @@
# under the License.
"""OpenCL TVMScript operator exports."""
-from tvm_tirx.ir.op import _init_op_api
-from tvm_tirx.script.ir_builder.base import auto_emit_void_call as
_auto_emit_void_call
+# Preserve completed definition-site exports alongside the explicit builders.
+import sys as _sys
-_init_op_api("tirx.opencl", __name__, wrapper=_auto_emit_void_call)
+from tvm_tirx.tirx.script.ir_builder.op import register_script_namespace as
_register_namespace
+
+_register_namespace("opencl", _sys.modules[__name__])
diff --git a/python/tvm_tirx/base.py b/python/tvm_tirx/base.py
index 0110d676b..9e9fb728d 100644
--- a/python/tvm_tirx/base.py
+++ b/python/tvm_tirx/base.py
@@ -21,13 +21,40 @@
import importlib
import importlib.util
import os
+from collections import deque
from pathlib import Path
from tvm_ffi import libinfo as tvm_ffi_libinfo
+from tvm_ffi import register_global_func
from tvm_ffi.libinfo import load_lib_ctypes
from . import libinfo
+# These receivers must exist before native static initialization. Keep them
stable:
+# native callers may cache them before the script builder has finished
importing.
+_EXPOSE_SCRIPT_OP_CALLBACK = None
+_EXPOSE_SCRIPT_TYPE_CALLBACK = None
+_PENDING_SCRIPT_EXPOSURES = deque()
+
+
+@register_global_func("script.expose_script_op")
+def _expose_script_op(name):
+ """Request script exposure after the operator's registration is
complete."""
+ if _EXPOSE_SCRIPT_OP_CALLBACK is None:
+ _PENDING_SCRIPT_EXPOSURES.append(("op", name))
+ else:
+ _EXPOSE_SCRIPT_OP_CALLBACK(name)
+
+
+@register_global_func("script.expose_script_type")
+def _expose_script_type(type_key):
+ """Request script exposure after the type's reflection and constructor are
ready."""
+ if _EXPOSE_SCRIPT_TYPE_CALLBACK is None:
+ _PENDING_SCRIPT_EXPOSURES.append(("type", type_key))
+ else:
+ _EXPOSE_SCRIPT_TYPE_CALLBACK(type_key)
+
+
# ----------------------------
# library loading
# ----------------------------
diff --git a/python/tvm_tirx/ir/op.py b/python/tvm_tirx/ir/op.py
index aeb2405c2..ab16de9ca 100644
--- a/python/tvm_tirx/ir/op.py
+++ b/python/tvm_tirx/ir/op.py
@@ -17,10 +17,7 @@
# pylint: disable=invalid-name
"""Primitive operators in the TVM IR."""
-import keyword
-import sys
from collections.abc import Sequence
-from types import ModuleType, SimpleNamespace
import tvm_ffi
@@ -170,6 +167,14 @@ class Op(Expr):
type_tail = [] if var_ty_args is None else
list(parse_entry(var_ty_args, "var_ty_args"))
self._set_signature(arg_names, arg_docs, ty_arg_names, ty_arg_docs,
value_tail, type_tail)
+ def expose_script(self) -> None:
+ """Expose the completed operator definition in TVMScript.
+
+ Call after setting the signature and attributes. Early requests wait
+ until the script builder is initialized; explicit builders take
precedence.
+ """
+ tvm_ffi.get_global_func("script.expose_script_op")(self.name)
+
class OpAttrMap:
"""The values of one operator attribute, keyed by operator.
@@ -311,65 +316,3 @@ def _make_op_api(op, module_name):
call.__module__ = module_name
call.__doc__ = op.doc or f"Construct a call to {op.name}."
return call
-
-
-def _init_op_api(namespace, target_module_name=None, *, wrapper=None):
- """Initialize registered Op callables in an already loaded Python module.
-
- Like :func:`tvm_ffi.init_ffi_api`, the registry prefix comes first and the
- target defaults to that module name. For example, a backend module uses
- ``_init_op_api("tirx.cuda", __name__)``. Dotted suffixes require existing
- module or SimpleNamespace containers; underscores are ordinary name parts.
-
- Generated functions accept registered positional/named operands plus
- ``attrs``, ``ty_args``, ``loc`` and ``ty``.
- Region operations construct a RegionFrame and accept ``body_params``
instead
- of Call result/type arguments.
- Omitting the result invokes an available Op inference hook; without one,
- Call retains a missing type.
- Explicit results and inference errors are preserved; Call.validate checks
- the operator contract separately. Attribute
- keywords construct the registered attrs schema, when one is declared.
-
- Existing Python callables retain ownership of their names; only missing
- exposures are generated. Existing functions must accept printed calls or
- have an appropriate exceptional printer hook. Non-callable collisions fail
- before any functions are installed. An optional ``wrapper`` adapts newly
- generated Call constructors at export time; existing exposures and region
- builders are untouched.
- Returns None.
- """
- if not namespace or any(
- not part.isidentifier() or keyword.iskeyword(part) for part in
namespace.split(".")
- ):
- raise ValueError(f"Invalid Op namespace {namespace!r}")
- target = sys.modules[target_module_name or namespace]
- pending = []
- destinations = set()
- for name in sorted(Op.list_op_names()):
- if not name.startswith(namespace + "."):
- continue
- parts = name[len(namespace) + 1 :].split(".")
- if any(not part.isidentifier() or keyword.iskeyword(part) for part in
parts):
- raise ValueError(f"Op {name!r} has no Python attribute spelling")
- container = target
- for part in parts[:-1]:
- container = vars(container).get(part)
- if not isinstance(container, ModuleType | SimpleNamespace):
- raise ValueError(f"Op {name!r} requires an existing namespace
at {part!r}")
- destination = (id(container), parts[-1])
- if destination in destinations:
- raise ValueError(f"Op {name!r} aliases another exposure
destination")
- destinations.add(destination)
- op = Op.get(name)
- if parts[-1] in vars(container):
- current = vars(container)[parts[-1]]
- if not callable(current):
- raise ValueError(f"Op {name!r} conflicts with an existing
Python attribute")
- else:
- call = _make_op_api(op, target.__name__)
- if wrapper is not None and op.get_attr("FRegionGetBodyParams") is
None:
- call = wrapper(call)
- pending.append((container, parts[-1], call))
- for container, name, call in pending:
- setattr(container, name, call)
diff --git a/python/tvm_tirx/ir/type.py b/python/tvm_tirx/ir/type.py
index e742d9073..9601eefbd 100644
--- a/python/tvm_tirx/ir/type.py
+++ b/python/tvm_tirx/ir/type.py
@@ -16,10 +16,7 @@
# under the License.
"""Unified type system in the project."""
-import keyword
-import sys
from collections.abc import Callable
-from types import ModuleType
from typing import ClassVar
import numpy as np
@@ -265,57 +262,3 @@ def _normalize_type(value) -> Type:
if isinstance(value, str | tvm_ffi.dtype):
return PrimType(value)
raise TypeError(f"Expected an IR type or primitive dtype, got
{type(value).__name__}")
-
-
-def _init_type_api(namespace: str, target_module: str | ModuleType) -> None:
- """Expose loaded Type bindings as short aliases in a builder namespace.
-
- Dialect bindings precede root ``ir`` bindings. Explicit exports always win.
- This only aliases registered classes: it never imports bindings, constructs
- instances, or generates constructors. Call it again after loading
additional
- bindings to fill missing names or replace generated root aliases with
dialect
- aliases. Module-local provenance distinguishes generated aliases from
exports.
-
- Parameters
- ----------
- namespace : str
- Registered type-key namespace to expose, such as ``"tirx"``. Its type
- aliases take precedence over root ``"ir"`` aliases with the same name.
- target_module : str or types.ModuleType
- Module receiving the aliases, or its name in ``sys.modules``. The
- module must already be loaded when passed by name.
- """
- module = sys.modules[target_module] if isinstance(target_module, str) else
target_module
- exports = vars(module)
- generated = exports.setdefault("_generated_type_aliases", {})
- root_index = Type.__tvm_ffi_type_info__.type_index
- type_keys = [str(key) for key in
tvm_ffi.registry.get_registered_type_keys()]
- dialect_aliases = set()
- for prefix in dict.fromkeys((namespace, "ir")):
- for type_key in type_keys:
- owner, _, name = type_key.rpartition(".")
- if owner != prefix or type_key == "ir.Type":
- continue
- info =
tvm_ffi.core._lookup_or_register_type_info_from_type_key(type_key)
- if root_index not in info.type_ancestors:
- continue
- alias = name[:-4] if name.endswith("Type") else name
- if keyword.iskeyword(alias):
- alias = name
- if prefix == namespace:
- # Reserve unbound dialect names so a later repeat can fill
them.
- dialect_aliases.add(alias)
- elif alias in dialect_aliases:
- continue
- if info.type_cls is None or not alias:
- continue
- if alias in exports:
- previous = generated.get(alias)
- if previous is None or exports[alias] is not previous[1]:
- continue
- if prefix != namespace or previous[0] != "ir" or namespace ==
"ir":
- continue
- exports[alias] = info.type_cls
- generated[alias] = (prefix, info.type_cls)
- if "__all__" in exports and alias not in exports["__all__"]:
- exports["__all__"] = [*exports["__all__"], alias]
diff --git a/python/tvm_tirx/script/ir_builder/__init__.py
b/python/tvm_tirx/script/ir_builder/__init__.py
index 68b9c2e1f..3adbfa624 100644
--- a/python/tvm_tirx/script/ir_builder/__init__.py
+++ b/python/tvm_tirx/script/ir_builder/__init__.py
@@ -35,7 +35,6 @@ from tvm_tirx.ir import (
Var,
make_node,
)
-from tvm_tirx.ir.type import _init_type_api
from .base import (
MISSING,
@@ -48,6 +47,7 @@ from .base import (
from .frame import IRModuleFrame
from .ir import (
Lambda,
+ _get_dialect_builder,
constexpr,
dtype,
dynamic,
@@ -57,7 +57,6 @@ from .ir import (
module_global_infos,
module_set_attr,
)
-from .ir import _get_dialect_builder as __getattr__
from .parser_protocol import (
decl_function,
ir_module,
@@ -66,7 +65,7 @@ from .parser_protocol import (
)
# Keep source namespaces independent of imported helper modules and lazy
builders.
-__all__ = [
+_exports = [
"MISSING",
"AlreadyEmitted",
"Bind",
@@ -107,4 +106,12 @@ __all__ = [
]
-_init_type_api("ir", __name__)
+def __getattr__(name):
+ if name == "__all__":
+ # A wildcard import needs the same completed type exposures as named
access.
+ import importlib
+
+ importlib.import_module("tvm_tirx.tirx.script.ir_builder")
+
+ return _exports
+ return _get_dialect_builder(name)
diff --git a/python/tvm_tirx/script/ir_builder/ir.py
b/python/tvm_tirx/script/ir_builder/ir.py
index 771a75dea..2a73fc33a 100644
--- a/python/tvm_tirx/script/ir_builder/ir.py
+++ b/python/tvm_tirx/script/ir_builder/ir.py
@@ -22,7 +22,6 @@ import tvm_ffi
from tvm_tirx.ir import DataTypeImm, GlobalInfo, LambdaExpr, Location, Var
from tvm_tirx.ir.location import UNKNOWN_LOC
-from tvm_tirx.ir.type import _init_type_api
from . import _ffi_api
from .base import IRBuilder, LocationEntry
@@ -190,6 +189,10 @@ def _get_dialect_builder(name: str):
module = importlib.import_module(f"tvm_tirx.script.ir_builder.{name}")
setattr(sys.modules["tvm_tirx.script.ir_builder"], name, module)
return module
+ if not name.startswith("_"):
+ importlib.import_module("tvm_tirx.tirx.script.ir_builder")
+ if name in globals():
+ return globals()[name]
raise AttributeError(f"module 'tvm_tirx.script.ir_builder' has no
attribute {name!r}")
@@ -202,4 +205,10 @@ def Lambda(parameter_types, function, *, ret_type=None):
# pylint: disable=inva
return LambdaExpr(parameter_types, function, ret_type=ret_type)
-_init_type_api("ir", __name__)
+def __getattr__(name):
+ if name == "__all__":
+ import importlib
+
+ importlib.import_module("tvm_tirx.tirx.script.ir_builder")
+ return [key for key in globals() if not key.startswith("_")]
+ return _get_dialect_builder(name)
diff --git a/python/tvm_tirx/tirx/script/ir_builder/__init__.py
b/python/tvm_tirx/tirx/script/ir_builder/__init__.py
index 1f41384dc..37683e278 100644
--- a/python/tvm_tirx/tirx/script/ir_builder/__init__.py
+++ b/python/tvm_tirx/tirx/script/ir_builder/__init__.py
@@ -16,7 +16,8 @@
# under the License.
"""Concrete TIRx construction operations over the shared native IRBuilder
stack."""
-from tvm_tirx.ir.type import _init_type_api
+from tvm_tirx import base as _base
+from tvm_tirx.script import ir_builder as _shared
from tvm_tirx.script.ir_builder import * # noqa: F403 - Re-export the public
namespace.
from tvm_tirx.script.ir_builder import __all__ as _shared_exports
from tvm_tirx.script.ir_builder import constexpr, dtype, dynamic
@@ -28,6 +29,7 @@ from .ir import boolean as bool
from .op import * # noqa: F403 - Re-export the public namespace.
from .op import __all__ as _op_exports
from .op import _get_script_namespace as __getattr__
+from .op import expose_script_op, expose_script_type
from .parser_protocol import * # noqa: F403 - Re-export the public namespace.
from .parser_protocol import __all__ as _protocol_exports
from .utils import (
@@ -54,4 +56,36 @@ __all__ = [
]
-_init_type_api("tirx", __name__)
+def _init_expose_script():
+ """Connect completed definitions to the initialized script builder
namespace."""
+ # A retry can inherit successful type aliases from modules that survived
+ # the first import. Preserve their provenance in the same export order.
+ generated = globals().setdefault("_generated_type_aliases", {})
+ for module, exports in ((_shared, _shared_exports), (_native,
_ir_exports)):
+ aliases = vars(module).get("_generated_type_aliases", {})
+ for name in exports:
+ previous = aliases.get(name)
+ if (
+ previous is not None
+ and globals().get(name) is previous[1]
+ and name not in _op_exports
+ and name not in _protocol_exports
+ ):
+ generated[name] = previous
+ else:
+ generated.pop(name, None)
+ # Explicit builders take precedence over generated bindings. Install both
+ # callbacks before replay so reentrant requests can run immediately.
+ _base._EXPOSE_SCRIPT_OP_CALLBACK = expose_script_op
+ _base._EXPOSE_SCRIPT_TYPE_CALLBACK = expose_script_type
+ while _base._PENDING_SCRIPT_EXPOSURES:
+ kind, name = _base._PENDING_SCRIPT_EXPOSURES[0]
+ if kind == "op":
+ expose_script_op(name)
+ else:
+ expose_script_type(name)
+ # Retain a failed request so initialization can be retried.
+ _base._PENDING_SCRIPT_EXPOSURES.popleft()
+
+
+_init_expose_script()
diff --git a/python/tvm_tirx/tirx/script/ir_builder/ir.py
b/python/tvm_tirx/tirx/script/ir_builder/ir.py
index 46ed3cc65..f20db9a5b 100644
--- a/python/tvm_tirx/tirx/script/ir_builder/ir.py
+++ b/python/tvm_tirx/tirx/script/ir_builder/ir.py
@@ -31,7 +31,7 @@ from tvm_tirx import ir as _ir
from tvm_tirx import tirx as tir
from tvm_tirx.ir import FloatImm, IntImm, Range, TensorLoad, TensorRegion, Type
from tvm_tirx.ir.location import UNKNOWN_LOC, Location
-from tvm_tirx.ir.type import _init_type_api, _normalize_type
+from tvm_tirx.ir.type import _normalize_type
from tvm_tirx.runtime import MemoryScope
from tvm_tirx.script.ir_builder.base import MISSING, AlreadyEmitted, IRBuilder
from tvm_tirx.script.ir_builder.ir import Lambda, meta_var
@@ -1312,6 +1312,3 @@ __all__ = [
"uint64x64",
"void",
]
-
-
-_init_type_api("tirx", __name__)
diff --git a/python/tvm_tirx/tirx/script/ir_builder/op.py
b/python/tvm_tirx/tirx/script/ir_builder/op.py
index f1233f0be..a6334cc29 100644
--- a/python/tvm_tirx/tirx/script/ir_builder/op.py
+++ b/python/tvm_tirx/tirx/script/ir_builder/op.py
@@ -18,6 +18,7 @@
from __future__ import annotations
+import keyword
import sys
from types import SimpleNamespace
from typing import Any
@@ -223,12 +224,22 @@ _SCRIPT_NAMESPACES = {}
def _get_script_namespace(name: str) -> object:
- """Return an explicitly registered backend construction namespace."""
+ """Return a generated construction binding or registered backend
namespace."""
if name in _SCRIPT_NAMESPACES:
return _SCRIPT_NAMESPACES[name]
raise AttributeError(f"No script namespace {name!r}")
+def _merge_script_namespace(generated, explicit):
+ """Keep definition-site exports when a backend installs its explicit
namespace."""
+ for name, value in vars(generated).items():
+ current = getattr(explicit, name, None)
+ if current is None:
+ setattr(explicit, name, value)
+ elif isinstance(value, SimpleNamespace):
+ _merge_script_namespace(value, current)
+
+
def register_script_namespace(name: str, namespace: object) -> object:
"""Register a construction namespace and return it.
@@ -239,6 +250,9 @@ def register_script_namespace(name: str, namespace: object)
-> object:
namespace : object
Construction namespace object.
"""
+ previous = _SCRIPT_NAMESPACES.get(name)
+ if isinstance(previous, SimpleNamespace) and previous is not namespace:
+ _merge_script_namespace(previous, namespace)
_SCRIPT_NAMESPACES[name] = namespace
globals()[name] = namespace
if "__all__" in globals() and name not in __all__:
@@ -320,6 +334,13 @@ def _script_slot(name: str):
parser.register_namespace(part, value)
else:
setattr(container, part, value)
+ exports = vars(container).get("__all__")
+ if isinstance(exports, list) and part not in exports:
+ exports.append(part)
+ if container is sys.modules["tvm_tirx.tirx.script.ir_builder"]:
+ # Keep successful bindings if a later exposure aborts the
builder's
+ # first import. Its next import can recover them through
__getattr__.
+ _SCRIPT_NAMESPACES[part] = value
*namespaces, leaf = parts
for part in namespaces:
@@ -329,6 +350,32 @@ def _script_slot(name: str):
return leaf, lookup, bind
+def _bind_type_alias(module, name, cls, type_key):
+ """Publish one completed type, preserving explicit and dialect-owned
aliases."""
+ exports = vars(module)
+ generated = exports.setdefault("_generated_type_aliases", {})
+ existing = exports.get(name)
+ existing_info = getattr(existing, "__tvm_ffi_type_info__", None)
+ previous = generated.get(name)
+ replaces_root = (
+ previous is not None
+ and existing is previous[1]
+ and previous[0].startswith("ir.")
+ and type_key.startswith("tirx.")
+ )
+ if existing is None or getattr(existing_info, "type_key", None) ==
type_key or replaces_root:
+ exports[name] = cls
+ if existing is None or (previous is not None and existing is
previous[1]):
+ generated[name] = (type_key, cls)
+ else:
+ # Re-exposing an explicit class must not make its alias
replaceable.
+ generated.pop(name, None)
+ module_all = getattr(module, "__all__", None)
+ if isinstance(module_all, list) and name not in module_all:
+ module_all.append(name)
+ return exports.get(name)
+
+
def expose_script_type(type_key: str) -> Any:
"""Give a registered Type the class its printed constructor names, and
return it.
@@ -366,7 +413,34 @@ def expose_script_type(type_key: str) -> Any:
cls = tvm_ffi.register_object(type_key)(type(name, (base,), attrs))
# The printer drops the "Type" suffix (DefaultTypeTranslate).
if len(name) > 4 and name.endswith("Type"):
- name = name[:-4]
+ short_name = name[:-4]
+ if keyword.iskeyword(short_name):
+ # Retain attribute lookup under the printed short name, while
+ # publishing a valid Python alias such as NoneType as well.
+ leaf, lookup, bind = _script_slot(f"{namespace}.{short_name}")
+ if lookup(leaf) is None:
+ bind(leaf, cls)
+ else:
+ name = short_name
+ # Core types also belong to the shared builders. Publish each completed
+ # definition directly, without scanning the object registry at module init.
+ modules = []
+ if namespace == "ir":
+ modules.extend(("tvm_tirx.script.ir_builder",
"tvm_tirx.script.ir_builder.ir"))
+ if namespace in ("ir", "tirx"):
+ modules.append("tvm_tirx.tirx.script.ir_builder.ir")
+ for module_name in modules:
+ module = sys.modules.get(module_name)
+ if module is not None:
+ _bind_type_alias(module, name, cls, type_key)
+ if namespace in ("ir", "tirx"):
+ module = sys.modules["tvm_tirx.tirx.script.ir_builder"]
+ result = _bind_type_alias(module, name, cls, type_key)
+ if not callable(result):
+ raise ValueError(f"Type {type_key!r} conflicts with a non-callable
Python attribute")
+ if result is cls:
+ _SCRIPT_NAMESPACES[name] = cls
+ return result
leaf, lookup, bind = _script_slot(f"{namespace}.{name}")
existing = lookup(leaf)
existing_info = getattr(existing, "__tvm_ffi_type_info__", None)
@@ -378,11 +452,6 @@ def expose_script_type(type_key: str) -> Any:
return cls
-# Rust operators and types call these when they are registered.
-tvm_ffi.register_global_func("tirx.script.expose_script_op", expose_script_op)
-tvm_ffi.register_global_func("tirx.script.expose_script_type",
expose_script_type)
-
-
abs = _auto_emit_void_call(_tir_op.abs) # pylint: disable=redefined-builtin
diff --git a/python/tvm_tirx/tirx/tile/instruction.py
b/python/tvm_tirx/tirx/tile/instruction.py
index aa3e2fbbe..8ad02bb1d 100644
--- a/python/tvm_tirx/tirx/tile/instruction.py
+++ b/python/tvm_tirx/tirx/tile/instruction.py
@@ -113,6 +113,7 @@ class Instruction:
]
construct.__signature__ = Signature(parameters, return_annotation=Call)
self._factory = construct
+ op.expose_script()
return construct
def make(self, *values, **kw):
diff --git a/rust/src/ir/op_def.rs b/rust/src/ir/op_def.rs
index 05b055f6d..e732ff9a2 100644
--- a/rust/src/ir/op_def.rs
+++ b/rust/src/ir/op_def.rs
@@ -92,11 +92,12 @@ impl OpDef {
Self(registered.and_then(|_| Op::get(name)))
}
- /// The operator, or the error of the step that failed. Once Python has
imported `tvm_tirx`,
- /// this also gives the operator its TVMScript builder
(`tirx.script.expose_script_op`).
+ /// The completed operator, or the error of the step that failed. Request
its TVMScript
+ /// builder through `script.expose_script_op` after its signature and
attributes are ready.
+ /// An absent receiver is a no-op; an installed receiver's error is
returned.
pub fn op(self) -> Result<Op> {
let op = self.0?;
- if let Ok(expose) =
tvm_ffi::Function::get_global("tirx.script.expose_script_op") {
+ if let Ok(expose) =
tvm_ffi::Function::get_global("script.expose_script_op") {
expose.call_tuple((op.name.clone(),))?;
}
Ok(op)
diff --git a/rust/src/ir/type_def.rs b/rust/src/ir/type_def.rs
index 5f8165d36..d32c59d26 100644
--- a/rust/src/ir/type_def.rs
+++ b/rust/src/ir/type_def.rs
@@ -108,7 +108,8 @@ impl Placeholder for Expr {
/// Register the type of the reference `R` of a
[`register_types!`](crate::register_types) block:
/// `__ffi_new__`, then the `__ffi_init__` tvm-ffi generates from the
reflected fields, then, once
-/// Python has imported `tvm_tirx`, its TVMScript class
(`tirx.script.expose_script_type`).
+/// the Python receiver is available, request its TVMScript class
(`script.expose_script_type`).
+/// An absent receiver is a no-op; an installed receiver's error is returned.
#[doc(hidden)]
pub fn __register_type<R>() -> Result<()>
where
@@ -118,7 +119,7 @@ where
let new = Function::from_typed(|| -> Result<R> { Ok(R::placeholder()) });
register_type_attr(type_index, "__ffi_new__", Any::from(new))?;
tvm_ffi::cached_global_func!("ffi._RegisterFFIInit").call_tuple((type_index,))?;
- if let Ok(expose) = Function::get_global("tirx.script.expose_script_type")
{
+ if let Ok(expose) = Function::get_global("script.expose_script_type") {
expose.call_tuple((String::from(support::type_key(type_index).as_str()),))?;
}
Ok(())
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index 01c5db7ef..9ae6cc118 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -146,7 +146,10 @@
//! for `my_ext.<name>`) is generated by
`tvm_tirx.tirx.script.expose_script_op`, unless Python has
//! one written for it, as for a C++ operator, and a type of a
//! [`register_types!`](crate::register_types) block gets the class its
TVMScript constructor names
-//! (`T.<Name>` for `tirx.<Name>Type`) from `expose_script_type`.
+//! (`T.<Name>` for `tirx.<Name>Type`) from `expose_script_type`. Both
completion hooks use the
+//! canonical `script.expose_script_op` and `script.expose_script_type` FFI
receivers. Python
+//! installs these before loading native libraries and defers exposure until
its script builders
+//! are ready. Without receivers, native registration proceeds without script
exposure.
//!
//! A library of operators and types of its own, a `cdylib` built on the crate
as
//! `examples/define_ir.rs` is, registers them when Python loads it, by
`tvm_ffi.load_module` or
diff --git a/src/backend/cuda/op/iket.cc b/src/backend/cuda/op/iket.cc
index 9c6e329b6..b0f55410f 100644
--- a/src/backend/cuda/op/iket.cc
+++ b/src/backend/cuda/op/iket.cc
@@ -41,7 +41,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.iket_mark")
.signature(sig::arg("name", "The name."), sig::var_args("args"))
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& iket_range_start_op() {
@@ -53,7 +54,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.iket_range_start")
.signature(sig::arg("name", "The name."), sig::var_args("args"))
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& iket_range_end_op() {
@@ -65,7 +67,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.iket_range_end")
.signature(sig::arg<IntExpr>("token", "The token."),
sig::var_args("args"))
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& iket_range_push_op() {
@@ -77,7 +80,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.iket_range_push")
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
.signature(sig::arg("name", "The name."), sig::var_args("args"))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& iket_range_pop_op() {
@@ -88,7 +92,8 @@ const Op& iket_range_pop_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.iket_range_pop")
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& iket_sentinel_token_op() {
@@ -100,7 +105,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.iket_sentinel_token")
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32))
.signature(sig::arg("name", "The name."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -108,7 +114,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr(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(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
} // namespace cuda
diff --git a/src/backend/cuda/op/legacy.cc b/src/backend/cuda/op/legacy.cc
index 57f54fa22..22782d6d5 100644
--- a/src/backend/cuda/op/legacy.cc
+++ b/src/backend/cuda/op/legacy.cc
@@ -35,26 +35,30 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("src_ptr", "The source pointer."),
sig::arg<IntExpr>("src_offset", "The source offset."),
sig::arg<IntExpr>("dst_stride", "The destination stride."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
OpDef("tirx.cuda.mma_fill")
.signature(sig::arg<IntExpr>("local_size", "The local allocation size."),
sig::arg("local_ptr", "The local pointer."),
sig::arg<IntExpr>("offset", "The offset."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
OpDef("tirx.ptx_legacy.ldmatrix")
.signature(sig::arg("trans"), sig::arg<IntExpr>("num"),
sig::arg("dtype"),
sig::arg("local_ptr"), sig::arg<IntExpr>("local_offset"),
sig::arg("smem_ptr"),
sig::arg<IntExpr>("smem_offset"))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
OpDef("tirx.ptx_legacy.mma")
.signature(sig::arg("shape"), sig::arg("a_layout"),
sig::arg("b_layout"), sig::arg("a_dtype"),
sig::arg("b_dtype"), sig::arg("c_dtype"), sig::arg("a_ptr"),
sig::arg<IntExpr>("a_offset"), sig::arg("b_ptr"),
sig::arg<IntExpr>("b_offset"),
sig::arg("acc_ptr"), sig::arg<IntExpr>("c_offset"),
sig::arg("saturate"),
sig::var_args("args"))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
} // namespace tirx
diff --git a/src/backend/cuda/op/target_builtin.cc
b/src/backend/cuda/op/target_builtin.cc
index d23d27df9..2857a7f51 100644
--- a/src/backend/cuda/op/target_builtin.cc
+++ b/src/backend/cuda/op/target_builtin.cc
@@ -102,7 +102,8 @@ void RegisterCudaTargetBuiltins() {
.set_validator(ffi::reflection::NativeFunctionView<void(
const CallObj*)>::FromNative<&ValidateDeclTmem>(),
true)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("tirx.cuda.bmma_sync")
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
@@ -114,7 +115,8 @@ void RegisterCudaTargetBuiltins() {
sig::arg<IntExpr>("index_b", "The B fragment index."),
sig::arg("fragment_c", "The C fragment."),
sig::arg<IntExpr>("index_c", "The C fragment index."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
// Raw legacy cp.async form emitted by InjectPTXAsyncCopy (and round-tripped
by
// the T.s_tir.cp_async_raw.legacy 6-arg surface). It carries the element
dtype in Call.dtype
@@ -125,7 +127,8 @@ void RegisterCudaTargetBuiltins() {
sig::arg("src_ptr", "The source pointer."),
sig::arg<IntExpr>("src_offset", "The source offset."),
sig::arg<IntExpr>("cp_size", "The copy size."),
sig::var_args("args"))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
OpDef("tirx.cuda.timer_init")
.signature(sig::arg("profiler_buffer", "The profiler buffer."),
@@ -133,7 +136,8 @@ void RegisterCudaTargetBuiltins() {
sig::arg<IntExpr>("profiler_write_offset", "The profiler
write offset."),
sig::arg<IntExpr>("num_groups", "The number of groups."),
sig::arg<IntExpr>("group_id", "The group identifier."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
OpDef("tirx.cuda.timer_start")
.signature(sig::arg("event_type", "The event type."),
@@ -142,7 +146,8 @@ void RegisterCudaTargetBuiltins() {
sig::arg<IntExpr>("profiler_write_offset", "The profiler
write offset."),
sig::arg<IntExpr>("profiler_write_stride", "The profiler
write stride."),
sig::arg("leader_cond", "The leader condition."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
OpDef("tirx.cuda.timer_end")
.signature(sig::arg("event_type", "The event type."),
@@ -151,7 +156,8 @@ void RegisterCudaTargetBuiltins() {
sig::arg<IntExpr>("profiler_write_offset", "The profiler
write offset."),
sig::arg<IntExpr>("profiler_write_stride", "The profiler
write stride."),
sig::arg("leader_cond", "The leader condition."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
OpDef("tirx.cuda.timer_finalize")
.signature(sig::arg("profiler_buffer", "The profiler buffer."),
@@ -159,7 +165,8 @@ void RegisterCudaTargetBuiltins() {
sig::arg<IntExpr>("profiler_write_offset", "The profiler
write offset."),
sig::arg<IntExpr>("profiler_write_stride", "The profiler
write stride."),
sig::arg("leader_cond", "The leader condition."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
RegisterDeviceIntrinsics();
}
@@ -182,219 +189,290 @@ void RegisterDeviceIntrinsics() {
"block_dim", "cluster_dim"}) {
RegisterDeviceIntrinsic(OpDef(std::string("tirx.cuda.") + name),
CallEffectKind::kPure,
sig::arg<StringImm>("axis"))
- .set_native_fn_attr<&InferTypeCudaAxis>(tvm::op_attr::kInferType);
+ .set_native_fn_attr<&InferTypeCudaAxis>(tvm::op_attr::kInferType)
+ .expose_script();
}
for (const char* name : {"linear_thread_id", "warp_id", "lane_id",
"warpgroup_id",
"warp_in_warpgroup", "thread_in_warpgroup",
"cta_pair_id"}) {
RegisterDeviceIntrinsic(OpDef(std::string("tirx.cuda.") + name),
CallEffectKind::kPure)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32))
+ .expose_script();
}
RegisterDeviceIntrinsic(OpDef("tirx.cuda.any_sync"), CallEffectKind::kPure,
sig::arg<IntExpr>("mask"), sig::arg("pred"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.atomic_add"),
CallEffectKind::kOpaque,
sig::arg("res_addr"), sig::arg("value"))
-
.set_native_fn_attr<&InferTypeReturnArgType<1>>(tvm::op_attr::kInferType);
+ .set_native_fn_attr<&InferTypeReturnArgType<1>>(tvm::op_attr::kInferType)
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.atomic_cas"),
CallEffectKind::kOpaque, sig::arg("ptr"),
sig::arg("old_val"), sig::arg("new_val"))
-
.set_native_fn_attr<&InferTypeReturnArgType<1>>(tvm::op_attr::kInferType);
+ .set_native_fn_attr<&InferTypeReturnArgType<1>>(tvm::op_attr::kInferType)
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.wait_until"),
CallEffectKind::kOpaque, sig::arg("dst"),
sig::arg("ptr"), sig::arg("condition"),
sig::arg("scope"),
- sig::arg("space"), sig::arg("ptx_type"),
sig::arg<IntExpr>("backoff_ns"));
+ sig::arg("space"), sig::arg("ptx_type"),
sig::arg<IntExpr>("backoff_ns"))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.ballot_sync"),
CallEffectKind::kOpaque,
sig::arg<IntExpr>("mask"), sig::arg("pred"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.bfloat1622float2"),
CallEffectKind::kOpaque,
sig::arg("packed"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(64));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(64))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.bfloat162float"),
CallEffectKind::kOpaque,
sig::arg("src"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Float(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Float(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.clock64"), CallEffectKind::kOpaque)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(64));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(64))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.cluster_sync"),
CallEffectKind::kOpaque)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.cta_reduce"),
CallEffectKind::kOpaque, sig::arg("value"),
sig::arg("op"), sig::arg<IntExpr>("num_warps"),
sig::arg("scratch"))
-
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType);
+ .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.cta_sync"), CallEffectKind::kOpaque)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.cvta_generic_to_shared"),
CallEffectKind::kOpaque,
sig::arg("ptr"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.elect_sync"),
CallEffectKind::kOpaque)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.fadd2_rn"),
CallEffectKind::kOpaque, sig::arg("a"),
sig::arg("b"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(64));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(64))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.fdividef"), CallEffectKind::kPure,
sig::arg("x"),
sig::arg("y"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Float(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Float(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.ffs_u32"), CallEffectKind::kOpaque,
sig::arg<IntExpr>("value"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.float22bfloat162_rn"),
CallEffectKind::kOpaque,
sig::arg("v0"), sig::arg("v1"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.float22bfloat162_rn_from_float2"),
CallEffectKind::kOpaque, sig::arg("packed"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.float22half2"),
CallEffectKind::kOpaque, sig::arg("dst"),
sig::arg("src"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.float2_x"),
CallEffectKind::kOpaque, sig::arg("packed"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Float(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Float(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.float2_y"),
CallEffectKind::kOpaque, sig::arg("packed"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Float(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Float(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.float8tohalf8"),
CallEffectKind::kOpaque,
sig::arg("src_addr"), sig::arg("dst_addr"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.float_as_uint"),
CallEffectKind::kOpaque, sig::arg("x"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.fmul2_rn"),
CallEffectKind::kOpaque, sig::arg("a"),
sig::arg("b"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(64));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(64))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.fp8x4_e4m3_from_float4"),
CallEffectKind::kOpaque,
sig::arg("x"), sig::arg("y"), sig::arg("z"),
sig::arg("w"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.func_call"),
CallEffectKind::kOpaque,
sig::arg("func_name"), sig::var_args("args"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.get_tmem_addr"),
CallEffectKind::kOpaque,
sig::arg("addr"), sig::arg<IntExpr>("row_offset"),
sig::arg<IntExpr>("col_offset"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.grid_sync"),
CallEffectKind::kOpaque)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.half2float"),
CallEffectKind::kOpaque, sig::arg("src"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Float(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Float(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.half8tofloat8"),
CallEffectKind::kOpaque,
sig::arg("src_addr"), sig::arg("dst_addr"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.hmax2"), CallEffectKind::kOpaque,
sig::arg("a"),
sig::arg("b"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.hmin2"), CallEffectKind::kOpaque,
sig::arg("a"),
sig::arg("b"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.make_float2"),
CallEffectKind::kOpaque, sig::arg("x"),
sig::arg("y"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(64));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(64))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.mbarrier_wait"),
CallEffectKind::kOpaque,
sig::arg("bar"), sig::arg<IntExpr>("phase"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.mbarrier_wait_acquire_cluster"),
CallEffectKind::kOpaque,
sig::arg("bar"), sig::arg<IntExpr>("phase"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.mov_sreg"), CallEffectKind::kPure,
sig::arg<IntExpr>("bits"), sig::arg("reg_name"))
- .set_native_fn_attr<&InferTypeMovSreg>(tvm::op_attr::kInferType);
+ .set_native_fn_attr<&InferTypeMovSreg>(tvm::op_attr::kInferType)
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.nano_sleep"),
CallEffectKind::kOpaque,
sig::arg<IntExpr>("time"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.printf"), CallEffectKind::kOpaque,
sig::arg("fmt"),
sig::var_args("args"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.reduce_add_sync_u32"),
CallEffectKind::kOpaque,
sig::arg<IntExpr>("mask"),
sig::arg<IntExpr>("value"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.reduce_min_sync_u32"),
CallEffectKind::kOpaque,
sig::arg<IntExpr>("mask"),
sig::arg<IntExpr>("value"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.runtime_instr_desc"),
CallEffectKind::kOpaque,
sig::arg("desc"), sig::arg<IntExpr>("sf_id"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.smem_addr_from_uint64"),
CallEffectKind::kOpaque,
sig::arg<IntExpr>("cluster_addr"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.syncthreads_and"),
CallEffectKind::kOpaque,
sig::arg("cond"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(64));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(64))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.syncthreads_or"),
CallEffectKind::kOpaque,
sig::arg("cond"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(64));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(64))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.thread_fence"),
CallEffectKind::kOpaque)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.thread_rank"),
CallEffectKind::kPure)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.trap_when_assert_failed"),
CallEffectKind::kOpaque,
sig::arg("cond"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.uint_as_float"),
CallEffectKind::kOpaque,
sig::arg<IntExpr>("bits"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Float(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Float(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.warp_reduce"),
CallEffectKind::kOpaque,
sig::arg("value"), sig::arg("op"),
sig::arg<IntExpr>("width"))
-
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType);
+ .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.warp_sync"),
CallEffectKind::kOpaque)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.warpgroup_sync"),
CallEffectKind::kOpaque,
sig::arg<IntExpr>("bar_no"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.wgmma_encode_matrix_descriptor"),
CallEffectKind::kOpaque, sig::arg("desc"),
sig::arg("addr"),
sig::arg<IntExpr>("ldo"), sig::arg<IntExpr>("sdo"),
sig::arg<IntExpr>("swizzle"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.wgmma_noop_barrier"),
CallEffectKind::kOpaque,
sig::arg("reg"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.barrier_all"),
CallEffectKind::kOpaque)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.fence"), CallEffectKind::kOpaque)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.getmem_nbi"),
CallEffectKind::kOpaque,
sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
sig::arg<IntExpr>("pe"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.getmem_nbi_block"),
CallEffectKind::kOpaque,
sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
- sig::arg<IntExpr>("pe"));
+ sig::arg<IntExpr>("pe"))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.getmem_nbi_warp"),
CallEffectKind::kOpaque,
sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
sig::arg<IntExpr>("pe"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.my_pe"), CallEffectKind::kOpaque)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.n_pes"), CallEffectKind::kOpaque)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32));
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32))
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.putmem_nbi"),
CallEffectKind::kOpaque,
sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
sig::arg<IntExpr>("pe"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.putmem_nbi_block"),
CallEffectKind::kOpaque,
sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
sig::arg<IntExpr>("pe"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.putmem_nbi_warp"),
CallEffectKind::kOpaque,
sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
sig::arg<IntExpr>("pe"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.putmem_signal_nbi"),
CallEffectKind::kOpaque,
sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
sig::arg("sig_addr"), sig::arg<IntExpr>("signal"),
sig::arg("sig_op"),
sig::arg<IntExpr>("pe"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.putmem_signal_nbi_block"),
CallEffectKind::kOpaque,
sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
sig::arg("sig_addr"), sig::arg<IntExpr>("signal"),
sig::arg("sig_op"),
sig::arg<IntExpr>("pe"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.putmem_signal_nbi_warp"),
CallEffectKind::kOpaque,
sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
sig::arg("sig_addr"), sig::arg<IntExpr>("signal"),
sig::arg("sig_op"),
sig::arg<IntExpr>("pe"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.quiet"), CallEffectKind::kOpaque)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.signal_op"),
CallEffectKind::kOpaque,
sig::arg("sig_addr"), sig::arg<IntExpr>("signal"),
sig::arg("sig_op"),
sig::arg<IntExpr>("pe"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.wait_until"),
CallEffectKind::kOpaque,
sig::arg("ivar"), sig::arg("cmp"),
sig::arg<IntExpr>("cmp_value"),
sig::arg("type"))
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
}
} // namespace
diff --git a/src/backend/cuda/op/tcgen05.cc b/src/backend/cuda/op/tcgen05.cc
index c50f6eddc..3079d82a5 100644
--- a/src/backend/cuda/op/tcgen05.cc
+++ b/src/backend/cuda/op/tcgen05.cc
@@ -74,17 +74,20 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.tcgen05_encode_instr_descriptor")
.signature(sig::arg("desc"),
sig::call_attrs<backend::cuda::TCGen05InstrDescriptorAttrs>())
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
OpDef("tirx.cuda.tcgen05_encode_instr_descriptor_block_scaled")
.signature(sig::arg("desc"),
sig::call_attrs<backend::cuda::TCGen05InstrDescriptorBlockScaledAttrs>())
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
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(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
- .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
+ .expose_script();
}
} // namespace cuda
diff --git a/src/backend/cuda/op/tensormap.cc b/src/backend/cuda/op/tensormap.cc
index f9ad36f95..6ec9fddef 100644
--- a/src/backend/cuda/op/tensormap.cc
+++ b/src/backend/cuda/op/tensormap.cc
@@ -63,7 +63,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32))
.signature(sig::arg("descriptor", "The descriptor."), sig::arg("data",
"The input data."),
sig::var_args<PrimExpr>("args"),
sig::call_attrs<TensorMapEncodeTiledAttr>())
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
} // namespace cuda
diff --git a/src/backend/opencl/op/memory.cc b/src/backend/opencl/op/memory.cc
index 6e8406984..fef318a0b 100644
--- a/src/backend/opencl/op/memory.cc
+++ b/src/backend/opencl/op/memory.cc
@@ -40,7 +40,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("channel_size", "The number of channels."),
sig::arg<PrimExpr>("value", "The value to use."))
.set_attr(tvm::tirx::op_attr::kVectorizable, true)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& texture2d_load_op() {
@@ -56,7 +57,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("channel_size", "The number of channels."),
sig::arg<PrimExpr>("element_index", "The element index within
a texture channel."))
.set_attr(tvm::tirx::op_attr::kVectorizable, true)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& nd_mem_alloc_with_scope_op() {
@@ -69,7 +71,8 @@ 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(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
} // namespace opencl
diff --git a/src/ir/op.cc b/src/ir/op.cc
index 503f12590..61a0b6011 100644
--- a/src/ir/op.cc
+++ b/src/ir/op.cc
@@ -25,11 +25,19 @@
#include <tvm/ffi/extra/serialization.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/op.h>
+#include <optional>
+
namespace tvm {
+void Op::ExposeScript(const ffi::String& op_name) {
+ static std::optional<ffi::Function> cached =
ffi::Function::GetGlobal("script.expose_script_op");
+ if (cached) (*cached)(op_name);
+}
+
// Owns canonical Ops and mutable attribute columns for the lifetime of the
process.
class OpRegistry {
public:
diff --git a/src/ir/prim/builtin.cc b/src/ir/prim/builtin.cc
index f91ecf056..56b38a7a0 100644
--- a/src/ir/prim/builtin.cc
+++ b/src/ir/prim/builtin.cc
@@ -205,7 +205,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kExprAnnotation)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.if_then_else")
.signature(sig::arg("condition", "The condition."),
@@ -216,28 +217,33 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_native_fn_attr<&IfThenElseMutate>(tirx::op_attr::kStructuralMutateWithConstraintContext)
.set_native_fn_attr<&IfThenElseMaybeInplaceMutate>(
tirx::op_attr::kStructuralMaybeInplaceMutateWithConstraintContext)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.vscale")
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.ceil")
.signature(sig::arg("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.log2")
.signature(sig::arg("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.clz")
.signature(sig::arg("x", "The input value."))
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
} // namespace prim
diff --git a/src/ir/prim/math.cc b/src/ir/prim/math.cc
index 253c77587..80ab990e4 100644
--- a/src/ir/prim/math.cc
+++ b/src/ir/prim/math.cc
@@ -388,206 +388,242 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<PrimExpr>("y", "The second input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.fabs")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.fmod")
.signature(sig::arg<PrimExpr>("x", "The input value."),
sig::arg<PrimExpr>("y", "The second input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.floor")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.round")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.nearbyint")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.trunc")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.exp")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.exp2")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.exp10")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.erf")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.tanh")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.sigmoid")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.sqrt")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.rsqrt")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.log")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.log1p")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.log10")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.tan")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.cos")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.cosh")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.sin")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.sinh")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.asin")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.acos")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.atan")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.acosh")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.asinh")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.atanh")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.atan2")
.signature(sig::arg<PrimExpr>("x1", "The first input value."),
sig::arg<PrimExpr>("x2", "The second input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.nextafter")
.signature(sig::arg<PrimExpr>("x1", "The first input value."),
sig::arg<PrimExpr>("x2", "The second input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.hypot")
.signature(sig::arg<PrimExpr>("x1", "The first input value."),
sig::arg<PrimExpr>("x2", "The second input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.copysign")
.signature(sig::arg<PrimExpr>("x1", "The first input value."),
sig::arg<PrimExpr>("x2", "The second input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.ldexp")
.signature(sig::arg<PrimExpr>("x1", "The first input value."),
sig::arg<PrimExpr>("x2", "The second input value."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.isnan")
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_native_fn_attr<&InferTypeIsNaN>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
OpDef("prim.popcount")
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.signature(sig::arg<PrimExpr>("x", "The input value."))
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.fma")
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
@@ -595,12 +631,14 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<PrimExpr>("y", "The second input value."),
sig::arg<PrimExpr>("z", "The third input value."))
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
OpDef("prim.assume")
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Bool())
.signature(sig::arg<PrimExpr>("cond", "The condition."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kEmbedInfo);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kEmbedInfo)
+ .expose_script();
ffi::reflection::GlobalDef()
.def("prim.infinity", prim::infinity)
diff --git a/src/ir/type.cc b/src/ir/type.cc
index 069c38350..fe9666804 100644
--- a/src/ir/type.cc
+++ b/src/ir/type.cc
@@ -29,10 +29,17 @@
#include <tvm/ir/type.h>
#include <cstdint>
+#include <optional>
#include <unordered_map>
namespace tvm {
+void Type::ExposeScript(const ffi::String& type_key) {
+ static std::optional<ffi::Function> cached =
+ ffi::Function::GetGlobal("script.expose_script_type");
+ if (cached) (*cached)(type_key);
+}
+
namespace {
DLDataType ScalableVectorDType(DLDataTypeCode code, int bits, int lanes) {
@@ -273,6 +280,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.attr(ffi::type_attr::kStructuralMaybeInplaceMutate,
ffi::FStructuralMutate::FromNative<&TypeMaybeInplaceMutate>());
refl::GlobalDef().def("ir.MissingType", []() { return MissingType(); });
+ Type::ExposeScript<MissingTypeObj>();
}
NoneType::NoneType() : Type(ffi::UnsafeInit{}) {
@@ -289,6 +297,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.attr(ffi::type_attr::kStructuralMaybeInplaceMutate,
ffi::FStructuralMutate::FromNative<&TypeMaybeInplaceMutate>());
refl::GlobalDef().def("ir.NoneType", []() { return NoneType(); });
+ Type::ExposeScript<NoneTypeObj>();
}
AnyType::AnyType(Location loc) : Type(ffi::UnsafeInit{}) {
@@ -306,6 +315,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.attr(ffi::type_attr::kStructuralMaybeInplaceMutate,
ffi::FStructuralMutate::FromNative<&TypeMaybeInplaceMutate>());
refl::GlobalDef().def("ir.AnyType", [](Location loc) { return AnyType(loc);
});
+ Type::ExposeScript<AnyTypeObj>();
}
TensorRegionType::TensorRegionType() : Type(ffi::UnsafeInit{}) {
@@ -322,6 +332,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.attr(ffi::type_attr::kStructuralMaybeInplaceMutate,
ffi::FStructuralMutate::FromNative<&TypeMaybeInplaceMutate>());
refl::GlobalDef().def("ir.TensorRegionType", []() { return
TensorRegionType(); });
+ Type::ExposeScript<TensorRegionTypeObj>();
}
OpaqueType::OpaqueType() : Type(ffi::UnsafeInit{}) { data_ =
ffi::make_object<OpaqueTypeObj>(); }
@@ -337,6 +348,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ffi::FStructuralMutate::FromNative<&OpaqueTypeMaybeInplaceMutate>());
refl::GlobalDef().def("ir.OpaqueType", []() { return OpaqueType(); });
+ Type::ExposeScript<OpaqueTypeObj>();
}
// PrimType
@@ -365,6 +377,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ffi::FStructuralMutate::FromNative<&PrimTypeMaybeInplaceMutate>());
refl::GlobalDef().def("ir.PrimType", [](DLDataType dtype) { return
PrimType(dtype); });
+ Type::ExposeScript<PrimTypeObj>();
}
PrimType PrimType::Int(int bits, int lanes) {
@@ -422,6 +435,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.attr(ffi::type_attr::kStructuralMaybeInplaceMutate,
ffi::FStructuralMutate::FromNative<&TypeMaybeInplaceMutate>());
refl::GlobalDef().def("ir.StringType", []() { return StringType(); });
+ Type::ExposeScript<StringTypeObj>();
}
// PtrType
@@ -446,6 +460,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("ir.PtrType", [](Type element_type,
runtime::MemoryScope scope = {}) {
return PtrType(element_type, scope);
});
+ Type::ExposeScript<PtrTypeObj>();
}
PtrType PtrType::VoidPtrType(runtime::MemoryScope scope) {
@@ -474,6 +489,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("ir.FuncType", [](tvm::ffi::Array<Type> arg_types,
Type ret_type) {
return FuncType(arg_types, ret_type);
});
+ Type::ExposeScript<FuncTypeObj>();
}
TupleType::TupleType(ffi::Array<Type> fields, Location loc) :
Type(ffi::UnsafeInit{}) {
@@ -495,6 +511,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def(
"ir.TupleType", [](ffi::Array<Type> fields, Location loc) { return
TupleType(fields, loc); });
+ Type::ExposeScript<TupleTypeObj>();
}
TupleType TupleType::Empty() { return TupleType(ffi::Array<Type>()); }
diff --git a/src/tirx/ir/type.cc b/src/tirx/ir/type.cc
index bae5fbc14..aa0c86f78 100644
--- a/src/tirx/ir/type.cc
+++ b/src/tirx/ir/type.cc
@@ -188,6 +188,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
return TensorType(std::move(shape),
std::move(dtype), std::move(scope),
data_alignment, std::move(layout),
std::move(loc));
});
+ Type::ExposeScript<TensorTypeObj>();
}
TensorMapType::TensorMapType(Location loc) : Type(ffi::UnsafeInit{}) {
@@ -208,6 +209,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ffi::FStructuralMutate::FromNative<&TensorMapTypeMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.TensorMapType", [](Location loc) { return
TensorMapType(loc); });
+ Type::ExposeScript<TensorMapTypeObj>();
}
} // namespace tvm::tirx
diff --git a/src/tirx/op/abi.cc b/src/tirx/op/abi.cc
index 470f21fc5..ff3ed68dc 100644
--- a/src/tirx/op/abi.cc
+++ b/src/tirx/op/abi.cc
@@ -61,7 +61,8 @@ const Op& call_extern_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.call_extern")
.signature(sig::arg("func_name", "The function name."),
sig::var_args("args"))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& call_pure_extern_op() {
@@ -72,7 +73,8 @@ const Op& call_pure_extern_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.call_pure_extern")
.signature(sig::arg("func_name", "The function name."),
sig::var_args("args"))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
const Op& call_llvm_intrin_op() {
@@ -83,7 +85,8 @@ const Op& call_llvm_intrin_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.call_llvm_intrin")
.signature(sig::arg<IntExpr>("intrin_id", "The intrinsic identifier."),
sig::var_args("args"))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& call_llvm_pure_intrin_op() {
@@ -95,7 +98,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.call_llvm_pure_intrin")
.signature(sig::arg<IntExpr>("intrin_id", "The intrinsic identifier."),
sig::var_args("args"))
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
- .set_attr(tvm::tirx::op_attr::kVectorizable, true);
+ .set_attr(tvm::tirx::op_attr::kVectorizable, true)
+ .expose_script();
}
const Op& call_spirv_pure_glsl450_op() {
@@ -107,7 +111,8 @@ 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(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
const Op& abi_field_get_op() {
@@ -119,7 +124,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.abi_field_get")
.signature(sig::arg("arr", "The array."), sig::arg<IntExpr>("index",
"The index."),
sig::arg<IntExpr>("field", "The field index."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kReadState);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kReadState)
+ .expose_script();
}
const Op& abi_field_set_op() {
@@ -133,7 +139,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("arr", "The array."), sig::arg<IntExpr>("index",
"The index."),
sig::arg<IntExpr>("field", "The field index."),
sig::arg("value", "The value to use."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kUpdateState);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kUpdateState)
+ .expose_script();
}
const Op& throw_last_error_op() {
@@ -144,7 +151,8 @@ const Op& throw_last_error_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.throw_last_error")
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& stack_alloca_op() {
@@ -157,7 +165,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_native_fn_attr<&InferTypeStackAlloca>(tvm::op_attr::kInferType)
.signature(sig::arg("dtype_str", "The data type name."),
sig::arg<IntExpr>("num", "The number of entries."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& stack_make_shape_op() {
@@ -169,7 +178,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.stack_make_shape")
.set_attr(tvm::op_attr::kFixedReturnType, PtrType(PrimType::Int(64)))
.signature(sig::var_args<IntExpr>("args"))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& stack_make_dltensor_op() {
@@ -185,7 +195,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("ndim", "The number of dimensions."),
sig::arg("arr_dtype", "The array data type."),
sig::arg<IntExpr>("elem_offset", "The element offset."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& call_packed_op() {
@@ -197,7 +208,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.call_packed")
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32))
.signature(sig::arg("func_name", "The function name."),
sig::var_args("args"))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -225,7 +237,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("kernel", "The kernel."), sig::var_args("args"),
sig::call_attrs<CallFFIKernelAttr>())
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& call_cpacked_op() {
@@ -237,7 +250,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.call_cpacked")
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32))
.signature(sig::arg("func_name", "The function name."),
sig::var_args("args"))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& call_packed_lowered_op() {
@@ -252,7 +266,8 @@ 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(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& call_cpacked_lowered_op() {
@@ -267,7 +282,8 @@ 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(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -279,7 +295,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("dtype_code_hint", "The data type code
hint."),
sig::arg<IntExpr>("dtype_bits_hint", "The data type bit-width
hint."))
.set_attr(tvm::tirx::op_attr::kGlobalSymbol, "TVMBackendAllocWorkspace")
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -289,7 +306,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("device_id", "The device index."),
sig::arg("ptr", "The pointer."))
.set_attr(tvm::tirx::op_attr::kGlobalSymbol, "TVMBackendFreeWorkspace")
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
} // namespace tirx
diff --git a/src/tirx/op/annotation.cc b/src/tirx/op/annotation.cc
index 49a169fde..987783f15 100644
--- a/src/tirx/op/annotation.cc
+++ b/src/tirx/op/annotation.cc
@@ -83,7 +83,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_native_fn_attr<&AssumeAlwaysMutate>(op_attr::kStructuralMutateWithConstraintContext)
.set_native_fn_attr<&AssumeAlwaysMaybeInplaceMutate>(
op_attr::kStructuralMaybeInplaceMutateWithConstraintContext)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kEmbedInfo);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kEmbedInfo)
+ .expose_script();
}
const Op& kernel_replace_point_op() {
@@ -94,7 +95,8 @@ const Op& kernel_replace_point_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.kernel_replace_point")
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& ignore_loop_partition_op() {
@@ -106,7 +108,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.ignore_loop_partition")
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Bool())
.signature(sig::arg<PrimExpr>("predicate", "The predicate."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
} // namespace tirx
diff --git a/src/tirx/op/cpu.cc b/src/tirx/op/cpu.cc
index cc652174e..169c0b1cd 100644
--- a/src/tirx/op/cpu.cc
+++ b/src/tirx/op/cpu.cc
@@ -40,7 +40,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cpu_parallel_barrier")
.signature()
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
} // namespace tirx
diff --git a/src/tirx/op/gpu.cc b/src/tirx/op/gpu.cc
index 824e269ad..ba4a378c4 100644
--- a/src/tirx/op/gpu.cc
+++ b/src/tirx/op/gpu.cc
@@ -51,7 +51,8 @@ PrimExpr gpu_thread_return(Location loc) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.gpu_thread_return")
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kControlJump);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kControlJump)
+ .expose_script();
}
const Op& gpu_thread_filter_op() {
@@ -64,7 +65,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Bool())
.signature(sig::arg<PrimVar>("var", "The thread-axis variable."),
sig::arg<PrimExpr>("pred", "The predicate."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
const Op& gpu_active_thread_selector_op() {
@@ -77,7 +79,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg<PrimVar>("var", "The thread-axis variable."),
sig::arg<PrimExpr>("pred", "The predicate."))
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& gpu_thread_invariant_op() {
@@ -89,7 +92,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.gpu_thread_invariant")
.set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType)
.signature(sig::arg<PrimExpr>("cond", "The condition."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
const Op& gpu_storage_sync_op() {
@@ -101,7 +105,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.gpu_storage_sync")
.signature(sig::arg<StringImm>("storage_scope", "The storage scope."))
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& gpu_warp_shuffle_op() {
@@ -117,7 +122,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("width", "The width."),
sig::arg<IntExpr>("warp_size", "The number of threads per
warp."))
.set_native_fn_attr<&InferTypeReturnArgType<1>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& gpu_warp_shuffle_up_op() {
@@ -132,7 +138,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("offset", "The offset."),
sig::arg<IntExpr>("width", "The width."),
sig::arg<IntExpr>("warp_size", "The number of threads per warp."))
.set_native_fn_attr<&InferTypeReturnArgType<1>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& gpu_warp_shuffle_down_op() {
@@ -147,7 +154,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("offset", "The offset."),
sig::arg<IntExpr>("width", "The width."),
sig::arg<IntExpr>("warp_size", "The number of threads per warp."))
.set_native_fn_attr<&InferTypeReturnArgType<1>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& gpu_warp_shuffle_xor_op() {
@@ -163,7 +171,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("width", "The width."),
sig::arg<IntExpr>("warp_size", "The number of threads per
warp."))
.set_native_fn_attr<&InferTypeReturnArgType<1>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& gpu_warp_activemask_op() {
@@ -174,7 +183,8 @@ const Op& gpu_warp_activemask_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.gpu_warp_activemask")
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& gpu_thread_allreduce_op() {
@@ -191,7 +201,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<Expr>("destinations", "The destination tensor
loads."),
sig::arg<Expr>("thread_axes", "The reduction thread axes."))
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& gpu_dp4a_op() {
@@ -205,7 +216,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<PrimExpr>("vec2", "The second input vector."),
sig::arg<PrimExpr>("acc", "The accumulator."))
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -216,7 +228,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("n", "The N dimension."), sig::arg<IntExpr>("k",
"The K dimension."),
sig::arg<IntExpr>("index", "The index."), sig::arg("buffer_ptr",
"The buffer pointer."),
sig::arg<IntExpr>("stride", "The stride."), sig::arg("layout", "The
layout."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kReadState);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kReadState)
+ .expose_script();
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -230,7 +243,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("index_b", "The B fragment index."),
sig::arg("fragment_c", "The C fragment."),
sig::arg<IntExpr>("index_c", "The C fragment index."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -240,7 +254,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("fragment", "The matrix fragment."), sig::arg<IntExpr>("m",
"The M dimension."),
sig::arg<IntExpr>("n", "The N dimension."), sig::arg<IntExpr>("k",
"The K dimension."),
sig::arg<IntExpr>("index", "The index."), sig::arg("value", "The
value to use."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -251,7 +266,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("n", "The N dimension."), sig::arg<IntExpr>("k",
"The K dimension."),
sig::arg<IntExpr>("index", "The index."), sig::arg("buffer_ptr",
"The buffer pointer."),
sig::arg<IntExpr>("stride", "The stride."), sig::arg("layout", "The
layout."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/tirx/op/memory.cc b/src/tirx/op/memory.cc
index ba3e553c6..b2a425649 100644
--- a/src/tirx/op/memory.cc
+++ b/src/tirx/op/memory.cc
@@ -172,7 +172,8 @@ Expr reinterpret(Type target_ty, Expr value, Location loc) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.reinterpret")
.signature(sig::arg("x", "The input value."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
const Op& address_of_op() {
@@ -184,7 +185,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.address_of")
.signature(sig::arg("obj", "The referenced object."))
.set_native_fn_attr<&InferTypeAddressOf>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
const Op& isnullptr_op() {
@@ -196,7 +198,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.isnullptr")
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Bool())
.signature(sig::arg("x", "The input value."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
const Op& prefetch_op() {
@@ -210,7 +213,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<IntExpr>("locality", "The locality hint."),
sig::arg<IntExpr>("cache_type", "The cache policy."))
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& ptr_byte_offset_op() {
@@ -222,7 +226,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.ptr_byte_offset")
.signature(sig::arg("data", "Base pointer."),
sig::arg<IntExpr>("byte_offset", "Offset in bytes."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
const Op& static_handle_op() {
@@ -232,7 +237,8 @@ const Op& static_handle_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.static_handle")
- .set_attr(tvm::op_attr::kCallEffectKind,
CallEffectKind::kSpecialCallArg);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kSpecialCallArg)
+ .expose_script();
}
const Op& handle_add_byte_offset_op() {
@@ -244,7 +250,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.handle_add_byte_offset")
.signature(sig::arg("handle", "The handle."),
sig::arg<IntExpr>("offset", "The offset."))
.set_native_fn_attr<&InferTypePointerOffset<0,
-1>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
const Op& atomic_add_op() {
@@ -255,7 +262,8 @@ 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(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& assume_aligned_op() {
@@ -268,7 +276,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg<TensorVar>("tensor", "The tensor whose base address
is aligned."),
sig::arg<IntImm>("alignment_bytes", "The constant byte
alignment."))
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kEmbedInfo);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kEmbedInfo)
+ .expose_script();
}
const Op& undef_op() {
@@ -279,7 +288,8 @@ const Op& undef_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.undef")
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kReadState);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kReadState)
+ .expose_script();
}
const Op& masked_load_op() {
@@ -292,7 +302,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg<TensorVar>("tensor", "The tensor."),
sig::arg<PrimExpr>("index", "The index."),
sig::var_args<PrimExpr>("args"))
.set_native_fn_attr<&InferTypeMaskedLoad>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kReadState);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kReadState)
+ .expose_script();
}
const Op& masked_store_op() {
@@ -306,7 +317,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg<PrimExpr>("value", "The value to use."),
sig::arg<PrimExpr>("index", "The index."),
sig::var_args<PrimExpr>("args"))
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kUpdateState);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kUpdateState)
+ .expose_script();
}
const Op& alloc_tensor_op() {
@@ -323,7 +335,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::var_args<tvm::Tuple>("placement", "Optional
producer-owned placement."),
sig::ty_arg<TensorType>("tensor_type"))
.set_attrs_type_key(DictAttrsObj::_type_key)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .expose_script();
}
const Op& decl_tensor_op() {
@@ -339,7 +352,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
true)
.signature(sig::arg("data", "The existing data pointer."),
sig::ty_arg<TensorType>("tensor_type"))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
const Op& tensor_data_ptr_op() {
@@ -351,7 +365,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.tensor_data_ptr")
.signature(sig::arg("tensor", "The tensor variable."))
.set_native_fn_attr<&InferTypeTensorDataPtr>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/tirx/op/region.cc b/src/tirx/op/region.cc
index 1afd492f2..d6d0ef29a 100644
--- a/src/tirx/op/region.cc
+++ b/src/tirx/op/region.cc
@@ -270,7 +270,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_native_fn_attr<&LaunchThreadMutate>(op_attr::kStructuralMutateWithConstraintContext)
.set_native_fn_attr<&LaunchThreadMaybeInplaceMutate>(
op_attr::kStructuralMaybeInplaceMutateWithConstraintContext)
-
.set_native_fn_attr<&ValidateBuiltinRegion>(tvm::op_attr::kRegionValidate);
+
.set_native_fn_attr<&ValidateBuiltinRegion>(tvm::op_attr::kRegionValidate)
+ .expose_script();
}
void ValidateDeviceScopeRegion(const RegionStmtObj* region) {
@@ -281,7 +282,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.device_scope", "Internal host/device splitting boundary.")
.signature(sig::var_args<Expr>("launch_values"),
sig::call_attrs<DictAttrsObj>())
.set_native_fn_attr<&RegionNoBodyParams>(tvm::op_attr::kRegionGetBodyParams)
-
.set_native_fn_attr<&ValidateDeviceScopeRegion>(tvm::op_attr::kRegionValidate);
+
.set_native_fn_attr<&ValidateDeviceScopeRegion>(tvm::op_attr::kRegionValidate)
+ .expose_script();
}
const Op& device_entry_op() {
@@ -306,7 +308,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.device_entry", "Enter a device kernel with independent launch
configuration.")
.signature(sig::var_args("launch_values"),
sig::call_attrs<DictAttrsObj>())
.set_native_fn_attr<&RegionNoBodyParams>(tvm::op_attr::kRegionGetBodyParams)
- .set_native_fn_attr<&ValidateDeviceEntry>(tvm::op_attr::kRegionValidate);
+ .set_native_fn_attr<&ValidateDeviceEntry>(tvm::op_attr::kRegionValidate)
+ .expose_script();
}
const Op& device_context_op() {
@@ -318,7 +321,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.device_context", "Supply the device type and ID within a
region.")
.signature(sig::arg<IntExpr>("device_type"),
sig::arg<IntExpr>("device_id"))
.set_native_fn_attr<&RegionNoBodyParams>(tvm::op_attr::kRegionGetBodyParams)
-
.set_native_fn_attr<&ValidateBuiltinRegion>(tvm::op_attr::kRegionValidate);
+
.set_native_fn_attr<&ValidateBuiltinRegion>(tvm::op_attr::kRegionValidate)
+ .expose_script();
}
const Op& compute_scope_op() {
@@ -330,7 +334,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.compute_scope", "Outline a named CPU compute region.")
.signature(sig::arg<StringImm>("name"))
.set_native_fn_attr<&RegionNoBodyParams>(tvm::op_attr::kRegionGetBodyParams)
-
.set_native_fn_attr<&ValidateBuiltinRegion>(tvm::op_attr::kRegionValidate);
+
.set_native_fn_attr<&ValidateBuiltinRegion>(tvm::op_attr::kRegionValidate)
+ .expose_script();
}
const Op& parallel_launch_op() {
@@ -342,7 +347,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.parallel_launch", "Launch a CPU worker team around a region.")
.signature()
.set_native_fn_attr<&RegionNoBodyParams>(tvm::op_attr::kRegionGetBodyParams)
-
.set_native_fn_attr<&ValidateBuiltinRegion>(tvm::op_attr::kRegionValidate);
+
.set_native_fn_attr<&ValidateBuiltinRegion>(tvm::op_attr::kRegionValidate)
+ .expose_script();
}
} // namespace tirx
diff --git a/src/tirx/op/vector.cc b/src/tirx/op/vector.cc
index 621d08ea9..4edd42879 100644
--- a/src/tirx/op/vector.cc
+++ b/src/tirx/op/vector.cc
@@ -59,7 +59,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.vector_high")
.signature(sig::arg<PrimExpr>("vec", "The input vector."))
.set_native_fn_attr<&InferTypeVectorPart<false>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
const Op& vector_low_op() {
@@ -71,7 +72,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.vector_low")
.signature(sig::arg<PrimExpr>("vec", "The input vector."))
.set_native_fn_attr<&InferTypeVectorPart<false>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
const Op& vector_combine_op() {
@@ -84,7 +86,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg<PrimExpr>("vec1", "The first input vector."),
sig::arg<PrimExpr>("vec2", "The second input vector."))
.set_native_fn_attr<&InferTypeVectorPart<true>>(tvm::op_attr::kInferType)
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
const Op& get_active_lane_mask_op() {
@@ -96,7 +99,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.get_active_lane_mask")
.signature(sig::arg<IntExpr>("base", "The base value."),
sig::arg<IntExpr>("limit", "The limit value."))
- .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure);
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure)
+ .expose_script();
}
} // namespace tirx
