This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new d55759ec1e [REFACTOR][TIRx] Define region body parameters on
operations (#20575)
d55759ec1e is described below
commit d55759ec1ebed554d9782a52e43e5730b8a96656
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 6 21:52:03 2026 -0400
[REFACTOR][TIRx] Define region body parameters on operations (#20575)
Let region operations define fresh typed body parameters from their
operands and attributes through FRegionGetBodyParams. Every region
operation registers the hook, with an empty array for zero parameters.
Attribute presence identifies region support without invoking the hook;
generic construction rejects missing registration and checks explicit
parameters while retaining their identities.
Use the shared region construction path for thread launches, with a
declared operand signature and operation-owned body parameters, and
remove the dedicated native builder forwarding. Register zero-parameter
hooks for the existing TIRx, S-TIR, and Trainium regions, and document
Region operations and Call operations in the existing builtin header.
---
include/tvm/tirx/builtin.h | 48 ++++++++++++++---
include/tvm/tirx/op_attr_types.h | 16 ++++++
include/tvm/tirx/script/ir_builder/ir.h | 13 ++---
include/tvm/tirx/stmt.h | 15 ++++++
.../tvm/tirx/script/ir_builder/parser_protocol.py | 12 +++--
python/tvm/tirx/stmt.py | 6 ++-
src/backend/trn/op/target_builtin.cc | 2 +
src/s_tir/stmt.cc | 10 +++-
src/tirx/ir/stmt.cc | 63 +++++++++++++---------
src/tirx/op/builtin.cc | 21 ++++++++
src/tirx/script/ir_builder/ir.cc | 18 +++----
src/tirx/transform/split_host_device.cc | 3 ++
12 files changed, 166 insertions(+), 61 deletions(-)
diff --git a/include/tvm/tirx/builtin.h b/include/tvm/tirx/builtin.h
index d7a48f0592..20d24fa89d 100644
--- a/include/tvm/tirx/builtin.h
+++ b/include/tvm/tirx/builtin.h
@@ -19,9 +19,9 @@
/*!
* \file tvm/tirx/builtin.h
- * \brief TIR builtin intrinsics.
+ * \brief TIRx builtin region and call operations.
*
- * TIR builtin intrinsics are stored as tvm:Op.
+ * TIRx builtin operations are stored as tvm::Op.
* They are processed in the same way as we process Ops.
*
* It is not necessary to create a function for every Op,
@@ -39,23 +39,54 @@
namespace tvm {
namespace tirx {
-/*! \brief Collection of builtin intrinsics as ops */
+/*! \brief Collection of builtin region and call operations as Ops. */
namespace builtin {
/*!
- * \brief Thread launch region: args are a StringImm tag and integer extent.
+ * \name Region operations
+ * \brief Operations used by RegionStmt to enclose a lexical body.
+ *
+ * Every region operation registers FRegionGetBodyParams, returning fresh typed
+ * body parameters or an empty array for no parameters. Attribute presence
alone
+ * identifies region support; classification does not invoke the hook. Operands
+ * and attributes are evaluated outside the body-parameter scope.
+ * \{
+ */
+/*!
+ * \brief Thread launch region: operands are a nonempty StringImm tag and a
+ * scalar signed or unsigned integer extent wider than one bit.
* The sole body parameter is a fresh thread-index PrimVar matching the extent
type.
* Tags starting with vthread denote virtual threads. There are no attrs or
results.
*/
TVM_DLL const Op& launch_thread();
-/*! \brief Mark a user-facing device entry region. */
+/*!
+ * \brief Mark a user-facing device entry containing device scope definitions.
+ * Takes no operands, body parameters, attributes or results.
+ */
TVM_DLL const Op& device_entry();
-/*! \brief Supply the device type and ID for lowering the region body. */
+/*!
+ * \brief Supply lexical device context for allocation and packed-call
lowering.
+ * Operands are integer device type and device ID; there are no body
parameters,
+ * attributes or results. The region does not change the active runtime device.
+ */
TVM_DLL const Op& device_context();
-/*! \brief Outline the region body as a named CPU compute helper. */
+/*!
+ * \brief Outline the body as a CPU compute helper named by a StringImm
operand.
+ * There are no body parameters, attributes or results.
+ */
TVM_DLL const Op& compute_scope();
-/*! \brief Launch a CPU worker team around the region body. */
+/*!
+ * \brief Launch a CPU worker team around parallel loops and team barriers.
+ * Takes no operands, body parameters, attributes or results.
+ */
TVM_DLL const Op& parallel_launch();
+/*! \} */
+
+/*!
+ * \name Call operations
+ * \brief Operations invoked through Call with explicit operands and result
types.
+ * \{
+ */
/*!
* \brief Allocate a buffer: alloc_tensor(shape, dtype, scope) -> TensorType.
*
@@ -719,6 +750,7 @@ enum TVMStructFieldKind : int {
* \brief Print the content of a buffer during runtime.
*/
TVM_DLL const Op& print_buffer();
+/*! \} */
} // namespace builtin
} // namespace tirx
} // namespace tvm
diff --git a/include/tvm/tirx/op_attr_types.h b/include/tvm/tirx/op_attr_types.h
index 8ca859ae72..08f977ffdc 100644
--- a/include/tvm/tirx/op_attr_types.h
+++ b/include/tvm/tirx/op_attr_types.h
@@ -29,11 +29,27 @@
#define TVM_TIR_OP_ATTR_TYPES_H_
#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/native_function.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/expr.h>
namespace tvm {
namespace tirx {
+/*!
+ * \brief Construct fresh typed variables for a region's lexical body
parameters.
+ *
+ * The input carries only the operation, operands and attributes, without a
body
+ * or builder state. Parameters are ordered, distinct definitions and may have
+ * name hints. Every region operation must register this hook, returning an
empty
+ * array when it has no body parameters. Presence of the attribute identifies
+ * region support without invoking the hook or allocating variables.
+ */
+using FRegionGetBodyParams =
+ ffi::reflection::NativeFunctionView<ffi::Array<Var>(const CallNode* call)>;
+
+/*! \brief Shared FRegionGetBodyParams implementation for regions without body
parameters. */
+inline ffi::Array<Var> RegionNoBodyParams(const CallNode*) { return {}; }
+
/*!
* \brief Global symbol of the op after lowering.
*/
diff --git a/include/tvm/tirx/script/ir_builder/ir.h
b/include/tvm/tirx/script/ir_builder/ir.h
index 2c35dcfda0..9599b60d3c 100644
--- a/include/tvm/tirx/script/ir_builder/ir.h
+++ b/include/tvm/tirx/script/ir_builder/ir.h
@@ -320,20 +320,13 @@ TensorVar AllocTensor(ffi::Array<PrimExpr> shape,
PrimType dtype = PrimType::Flo
* \brief Construct a result-free region with lexical body parameters.
* \param op The region operation.
* \param args Operands evaluated outside the region body.
- * \param body_params Variables defined at body entry.
+ * \param body_params Existing body parameters, or omitted to generate them
from the operation.
* \param attrs Additional operation attributes.
*/
-RegionFrame Region(Op op, ffi::Array<Expr> args, ffi::Array<Var> body_params =
{},
+RegionFrame Region(Op op, ffi::Array<Expr> args,
+ ffi::Optional<ffi::Array<Var>> body_params = std::nullopt,
DictAttrs attrs = DictAttrs());
-/*!
- * \brief Launch a thread with a fresh lexical variable.
- * \param thread_tag The thread axis tag.
- * \param extent The thread extent, which also determines the variable's type.
- * \return The result RegionFrame.
- */
-RegionFrame LaunchThread(ffi::String thread_tag, PrimExpr extent);
-
/*!
* \brief Store data in a buffer.
* \param buffer The buffer.
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 5b414cef47..42dc15efe9 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -115,6 +115,21 @@ class Bind : public Stmt {
TVM_DEFINE_OBJECT_REF_COW_METHOD(BindNode);
};
+/*!
+ * \brief Check FRegionGetBodyParams registration without invoking it or
creating variables.
+ */
+TVM_DLL bool IsRegionOp(const Op& op);
+
+/*!
+ * \brief Validate region operands and construct fresh typed body parameters.
+ * \param op The region operation.
+ * \param args Operands evaluated in the enclosing scope.
+ * \param attrs Attributes evaluated in the enclosing scope.
+ * \return Parameters from the required FRegionGetBodyParams hook.
+ * \throws ValueError if the operation has no registered region hook.
+ */
+TVM_DLL ffi::Array<Var> GetRegionBodyParams(Op op, ffi::Array<Expr> args,
DictAttrs attrs);
+
/*!
* \brief A single-body statement whose semantics are defined by an operator.
* Operands are evaluated in the enclosing scope. Body parameters are
definitions
diff --git a/python/tvm/tirx/script/ir_builder/parser_protocol.py
b/python/tvm/tirx/script/ir_builder/parser_protocol.py
index a61379f4bc..75e9e93589 100644
--- a/python/tvm/tirx/script/ir_builder/parser_protocol.py
+++ b/python/tvm/tirx/script/ir_builder/parser_protocol.py
@@ -916,7 +916,13 @@ def region(
body_params: Sequence[Var] | None = None,
attrs: _ir.DictAttrs | dict[str, Any] | None = None,
) -> frame.RegionFrame:
- """Construct a result-free region with explicit fresh body parameters.
+ """Construct a result-free region with operation-defined body parameters.
+
+ When ``body_params`` is omitted, the operation's ``FRegionGetBodyParams``
+ hook creates fresh typed variables. Every region operation must register
+ the hook, returning an empty array for no body parameters. Missing hooks
+ reject construction even with explicit parameters. Explicit parameters must
+ match the hook's count and types and retain their identities.
Operands and attributes belong to the enclosing scope. Entering the frame
returns one parameter directly, or a sequence for zero or multiple
parameters.
@@ -927,7 +933,7 @@ def region(
op = _ir.Op.get(op)
if attrs is None or isinstance(attrs, dict):
attrs = _ir.make_node("ir.DictAttrs", **(attrs or {}))
- return _ffi_api.Region(op, args, [] if body_params is None else
body_params, attrs)
+ return _ffi_api.Region(op, args, body_params, attrs)
def device_context(device_type: Expr, device_id: Expr) -> frame.RegionFrame:
@@ -962,7 +968,7 @@ def launch_thread(thread_tag: str, extent: Expr) ->
frame.RegionFrame:
with T.launch_thread("threadIdx.x", 32) as tx:
T.evaluate(tx)
"""
- return _ffi_api.LaunchThread(thread_tag, extent)
+ return region("tirx.launch_thread", [_StringImm(thread_tag), extent])
# --------------------------------------
diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py
index 8c16ca078b..74087da99a 100644
--- a/python/tvm/tirx/stmt.py
+++ b/python/tvm/tirx/stmt.py
@@ -278,7 +278,11 @@ class BufferStore(Stmt):
class RegionStmt(Stmt):
"""An operation with enclosing-scope operands and one lexical body.
- ``body_params`` define variables visible only within ``body``.
``result_vars``
+ ``body_params`` define variables visible only within ``body``. Their count
+ and types must match the operation's required ``FRegionGetBodyParams``
hook.
+ Zero-parameter regions register an empty-return hook; operations without a
+ hook do not support region construction. Explicit parameter
+ identities and their references in ``body`` are preserved. ``result_vars``
define variables after the region in the enclosing sequence. Attributes are
evaluated outside the body-parameter scope. Direct construction and JSON
serialization support result variables; structured script syntax currently
diff --git a/src/backend/trn/op/target_builtin.cc
b/src/backend/trn/op/target_builtin.cc
index 439e47dab6..613b7d3d9c 100644
--- a/src/backend/trn/op/target_builtin.cc
+++ b/src/backend/trn/op/target_builtin.cc
@@ -50,6 +50,8 @@ void RegisterTRNTargetBuiltins() {
OpDef("tirx.nki.tensorized_instruction", "Tensorize the NKI instructions in
the region body.")
.signature()
+ .set_attr<FRegionGetBodyParams>("FRegionGetBodyParams",
+
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>())
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
ffi::String("nki"))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
diff --git a/src/s_tir/stmt.cc b/src/s_tir/stmt.cc
index 7b17a48eaa..f2f2f86041 100644
--- a/src/s_tir/stmt.cc
+++ b/src/s_tir/stmt.cc
@@ -27,6 +27,7 @@
#include <tvm/ir/prim/op.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/sym/analyzer.h>
+#include <tvm/tirx/op_attr_types.h>
namespace tvm {
namespace s_tir {
@@ -54,9 +55,14 @@ const Op& manual_sync() {
}
TVM_FFI_STATIC_INIT_BLOCK() {
- OpDef("s_tir.async_copy_scope", "Mark eligible copies for asynchronous
lowering.").signature();
+ OpDef("s_tir.async_copy_scope", "Mark eligible copies for asynchronous
lowering.")
+ .signature()
+ .set_attr<FRegionGetBodyParams>("FRegionGetBodyParams",
+
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>());
OpDef("s_tir.manual_sync", "Use explicitly authored synchronization within
the body.")
- .signature();
+ .signature()
+ .set_attr<FRegionGetBodyParams>("FRegionGetBodyParams",
+
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>());
OpDef("s_tir.async_commit", "Commit asynchronous copies to a queue.")
.signature(sig::arg<IntImm>("queue_id"))
.set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void())
diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc
index 7694f664f5..4ab20240b1 100644
--- a/src/tirx/ir/stmt.cc
+++ b/src/tirx/ir/stmt.cc
@@ -722,19 +722,47 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
// RegionStmt
+bool IsRegionOp(const Op& op) {
+ if (!Op::HasAttrMap("FRegionGetBodyParams")) return false;
+ static auto get_body_params =
Op::GetAttrMap<FRegionGetBodyParams>("FRegionGetBodyParams");
+ return get_body_params.count(op);
+}
+
+ffi::Array<Var> GetRegionBodyParams(Op op, ffi::Array<Expr> args, DictAttrs
attrs) {
+ TVM_FFI_CHECK(IsRegionOp(op), ValueError)
+ << op->name << " does not support region construction:
FRegionGetBodyParams is required";
+ CallNode signature(op);
+ signature.args = std::move(args);
+ signature.attrs = std::move(attrs);
+ op.Validate(&signature);
+ static auto get_body_params =
Op::GetAttrMap<FRegionGetBodyParams>("FRegionGetBodyParams");
+ ffi::Array<Var> params =
get_body_params[op].CallExpected(&signature).value();
+ std::unordered_set<const VarNode*> definitions;
+ for (const Var& param : params) {
+ TVM_FFI_CHECK(!param->ty.as<MissingType>().has_value(), ValueError)
+ << "FRegionGetBodyParams for " << op->name << " must return typed
variables";
+ TVM_FFI_CHECK(definitions.insert(param.get()).second, ValueError)
+ << "FRegionGetBodyParams for " << op->name << " must return distinct
definitions";
+ }
+ return params;
+}
+
RegionStmt::RegionStmt(Op op, ffi::Array<Expr> args, ffi::Array<Var>
body_params, DictAttrs attrs,
Stmt body, ffi::Array<Var> result_vars, Span span)
: Stmt(ffi::UnsafeInit{}) {
TVM_FFI_CHECK(op.defined() && body.defined(), ValueError)
<< "RegionStmt requires an operator and a body";
- CallNode signature(op);
- signature.args = args;
- signature.attrs = attrs;
- op.Validate(&signature);
+ ffi::Array<Var> expected_params = GetRegionBodyParams(op, args, attrs);
+ TVM_FFI_CHECK_EQ(body_params.size(), expected_params.size(), ValueError)
+ << op->name << " expects " << expected_params.size() << " body
parameters";
+ for (size_t i = 0; i < body_params.size(); ++i) {
+ TVM_FFI_CHECK(ffi::StructuralEqual()(body_params[i]->ty,
expected_params[i]->ty), ValueError)
+ << op->name << " body parameter " << i << " has a type inconsistent
with its contract";
+ }
if (op.same_as(tirx::builtin::device_context()) ||
op.same_as(tirx::builtin::compute_scope()) ||
op.same_as(tirx::builtin::parallel_launch())) {
- TVM_FFI_CHECK(body_params.empty() && result_vars.empty() &&
attrs->dict.empty(), ValueError)
- << op->name << " expects no body parameters, results or attributes";
+ TVM_FFI_CHECK(result_vars.empty() && attrs->dict.empty(), ValueError)
+ << op->name << " expects no results or attributes";
}
std::unordered_set<const VarNode*> definitions;
for (const auto& vars : {body_params, result_vars}) {
@@ -745,31 +773,14 @@ RegionStmt::RegionStmt(Op op, ffi::Array<Expr> args,
ffi::Array<Var> body_params
}
static const Op device_scope = Op::Get("tirx.device_scope");
if (op.same_as(tirx::builtin::device_entry()) || op.same_as(device_scope)) {
- TVM_FFI_CHECK(body_params.empty() && result_vars.empty(), ValueError)
- << op->name << " expects no body parameters or results";
+ TVM_FFI_CHECK(result_vars.empty(), ValueError) << op->name << " expects no
results";
if (op.same_as(tirx::builtin::device_entry())) {
TVM_FFI_CHECK(attrs->dict.empty(), ValueError) << "device_entry expects
no attrs";
}
}
if (op.same_as(tirx::builtin::launch_thread())) {
- TVM_FFI_CHECK(
- args.size() == 2 && body_params.size() == 1 && result_vars.empty() &&
attrs->dict.empty(),
- ValueError)
- << "launch_thread expects tag, extent, one body parameter, and no
results or attrs";
- auto tag = args[0].as<StringImm>();
- auto extent = args[1].as<PrimExpr>();
- auto var = body_params[0].as<PrimVar>();
- TVM_FFI_CHECK(tag && !tag.value()->value.empty(), ValueError)
- << "launch_thread expects a nonempty StringImm thread tag";
- auto integer = [](PrimType ty) {
- return ty.IsScalar() &&
- (ty.MatchesCode(DLDataTypeCode::kDLInt) ||
ty.MatchesCode(DLDataTypeCode::kDLUInt)) &&
- ty.bits() > 1;
- };
- TVM_FFI_CHECK(extent && var && integer(extent.value().ty()) &&
integer(var.value().ty()) &&
- extent.value().ty() == var.value().ty(),
- ValueError)
- << "launch_thread expects matching scalar integer extent and body
parameter types";
+ TVM_FFI_CHECK(result_vars.empty() && attrs->dict.empty(), ValueError)
+ << "launch_thread expects no results or attrs";
}
auto n = ffi::make_object<RegionStmtNode>(std::move(op), std::move(body));
n->args = std::move(args);
diff --git a/src/tirx/op/builtin.cc b/src/tirx/op/builtin.cc
index b4332fcb15..8e3d1cf932 100644
--- a/src/tirx/op/builtin.cc
+++ b/src/tirx/op/builtin.cc
@@ -34,6 +34,16 @@ namespace tvm {
namespace tirx {
namespace builtin {
+ffi::Array<Var> LaunchThreadBodyParams(const CallNode* call) {
+ auto tag = call->args[0].as_or_throw<StringImm>();
+ PrimType dtype = call->args[1].as_or_throw<IntExpr>().ty();
+ TVM_FFI_CHECK(!tag->value.empty(), ValueError)
+ << "launch_thread expects a nonempty StringImm thread tag";
+ TVM_FFI_CHECK_GT(dtype.bits(), 1, ValueError)
+ << "launch_thread extent must have a scalar integer type wider than one
bit";
+ return {PrimVar("", dtype)};
+}
+
template <size_t N>
static Type InferTypeReturnArgType(const CallNode* call) {
TVM_FFI_CHECK_GT(call->args.size(), N, ValueError)
@@ -275,18 +285,29 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.device_entry", "Mark a device entry containing scope
definitions.")
.signature()
+ .set_attr<FRegionGetBodyParams>("FRegionGetBodyParams",
+
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>())
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"));
OpDef("tirx.launch_thread", "Bind a thread index within a body with a launch
extent.")
+ .signature(sig::arg<StringImm>("tag"), sig::arg<IntExpr>("extent"))
+ .set_attr<FRegionGetBodyParams>("FRegionGetBodyParams",
+
FRegionGetBodyParams::FromNative<&LaunchThreadBodyParams>())
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"));
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_attr<FRegionGetBodyParams>("FRegionGetBodyParams",
+
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>())
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"));
OpDef("tirx.compute_scope", "Outline a named CPU compute region.")
.signature(sig::arg<StringImm>("name"))
+ .set_attr<FRegionGetBodyParams>("FRegionGetBodyParams",
+
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>())
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"));
OpDef("tirx.parallel_launch", "Launch a CPU worker team around a region.")
.signature()
+ .set_attr<FRegionGetBodyParams>("FRegionGetBodyParams",
+
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>())
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"));
OpDef("tirx.thread_return")
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.thread_return"))
diff --git a/src/tirx/script/ir_builder/ir.cc b/src/tirx/script/ir_builder/ir.cc
index f4d614e113..8d3c8c73fa 100644
--- a/src/tirx/script/ir_builder/ir.cc
+++ b/src/tirx/script/ir_builder/ir.cc
@@ -414,22 +414,19 @@ Var Bind(Expr value, ffi::Optional<Type> type_annotation,
ffi::Optional<Var> var
return bind_var;
}
-RegionFrame Region(Op op, ffi::Array<Expr> args, ffi::Array<Var> body_params,
DictAttrs attrs) {
+RegionFrame Region(Op op, ffi::Array<Expr> args,
ffi::Optional<ffi::Array<Var>> body_params,
+ DictAttrs attrs) {
+ TVM_FFI_CHECK(tvm::tirx::IsRegionOp(op), ValueError)
+ << op->name << " does not support region construction:
FRegionGetBodyParams is required";
+ auto params = body_params.has_value() ? body_params.value()
+ : tvm::tirx::GetRegionBodyParams(op,
args, attrs);
auto n = ffi::make_object<RegionFrameNode>(std::move(op));
n->args = std::move(args);
- n->body_params = std::move(body_params);
+ n->body_params = std::move(params);
n->attrs = std::move(attrs);
return RegionFrame(n);
}
-RegionFrame LaunchThread(ffi::String thread_tag, PrimExpr extent) {
- PrimType dtype = extent.ty();
- TVM_FFI_CHECK(dtype.IsScalar() && dtype.MatchesCode(kDLInt, kDLUInt),
ValueError)
- << "launch_thread extent must have a scalar integer type";
- PrimVar var("", dtype);
- return Region(tvm::tirx::builtin::launch_thread(), {StringImm(thread_tag),
extent}, {var});
-}
-
WhileFrame While(PrimExpr condition) {
ffi::ObjectPtr<WhileFrameNode> n =
ffi::make_object<WhileFrameNode>(condition);
return WhileFrame(n);
@@ -702,7 +699,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("script.ir_builder.tirx.Else", Else)
.def("script.ir_builder.tirx.DeclTensor", DeclTensor)
.def("script.ir_builder.tirx.Region", Region)
- .def("script.ir_builder.tirx.LaunchThread", LaunchThread)
.def("script.ir_builder.tirx.BufferStore", BufferStore)
.def("script.ir_builder.tirx.Evaluate", Evaluate)
.def("script.ir_builder.tirx.Ptr", Ptr);
diff --git a/src/tirx/transform/split_host_device.cc
b/src/tirx/transform/split_host_device.cc
index cb8ebb7bb1..4526b42577 100644
--- a/src/tirx/transform/split_host_device.cc
+++ b/src/tirx/transform/split_host_device.cc
@@ -35,6 +35,7 @@
#include <tvm/tirx/attrs.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>
+#include <tvm/tirx/op_attr_types.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>
@@ -50,6 +51,8 @@ namespace tirx {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.device_scope", "Internal host/device splitting boundary.")
.signature(sig::call_attrs<DictAttrsNode>())
+ .set_attr<FRegionGetBodyParams>("FRegionGetBodyParams",
+
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>())
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"));
}
