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 a55f5d5d62 [REFACTOR][TIRx] Use regions for CPU and runtime context
scopes (#20567)
a55f5d5d62 is described below
commit a55f5d5d623a3bf2444b89a7b6d289a89a1844a1
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 6 13:25:05 2026 -0400
[REFACTOR][TIRx] Use regions for CPU and runtime context scopes (#20567)
Represent device context, CPU compute outlining, and explicit parallel
launches with dedicated RegionStmt ops. Pair device type and ID as
ordinary operands, retain compute-helper naming and existing worker-team
behavior, and migrate packed-API production, lowering, codegen, and
script syntax together.
Reuse typed operator signatures at region construction and preserve
lexical workspace cleanup. Remove the superseded attribute encodings.
---
include/tvm/tirx/builtin.h | 6 ++
include/tvm/tirx/stmt.h | 9 ---
python/tvm/tirx/op.py | 2 +-
.../tvm/tirx/script/ir_builder/parser_protocol.py | 21 ++++++
src/target/llvm/codegen_cpu.cc | 22 ++++---
src/target/llvm/codegen_cpu.h | 3 +-
src/target/llvm/codegen_llvm.cc | 4 ++
src/target/source/codegen_c.cc | 4 ++
src/tirx/ir/stmt.cc | 5 ++
src/tirx/op/builtin.cc | 12 ++++
src/tirx/script/printer/stmt.cc | 11 ++++
src/tirx/transform/ir_utils.cc | 8 +--
src/tirx/transform/lower_tvm_builtin.cc | 74 ++++++++++++----------
src/tirx/transform/make_packed_api.cc | 16 ++---
.../s_tir/script/test_s_tir_script_printer.py | 67 ++++++++++----------
.../test_tir_transform_lower_tvm_builtin.py | 32 ++++------
.../test_tir_transform_make_packed_api.py | 12 ++--
.../test_tir_transform_storage_rewrite.py | 5 +-
18 files changed, 188 insertions(+), 125 deletions(-)
diff --git a/include/tvm/tirx/builtin.h b/include/tvm/tirx/builtin.h
index 139f8a3972..688198c60c 100644
--- a/include/tvm/tirx/builtin.h
+++ b/include/tvm/tirx/builtin.h
@@ -50,6 +50,12 @@ TVM_DLL const Op& launch_thread();
/*! \brief Mark a user-facing device entry region. */
TVM_DLL const Op& device_entry();
+/*! \brief Supply the device type and ID for lowering the region body. */
+TVM_DLL const Op& device_context();
+/*! \brief Outline the region body as a named CPU compute helper. */
+TVM_DLL const Op& compute_scope();
+/*! \brief Launch a CPU worker team around the region body. */
+TVM_DLL const Op& parallel_launch();
/*!
* \brief Allocate a buffer: alloc_tensor(shape, dtype, scope) -> TensorType.
*
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 6331344f52..4894b8b55f 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -852,15 +852,6 @@ class ScopeIdDefStmt : public Stmt {
/*! \brief Statement attribute and loop annotation keys. */
namespace attr {
-/*!
- * \brief Mark the scope as when computation start to happen.
- * This can hint some code generator to create a new function for compute.
- */
-constexpr const char* compute_scope = "compute_scope";
-/*! \brief The allocation device for global malloc in host. */
-constexpr const char* device_id = "device_id";
-/*! \brief The device type. */
-constexpr const char* device_type = "device_type";
/*!
* \brief For annotation: maximum work for automatic unrolling.
*
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index 3ab04608b7..9014497141 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -918,7 +918,7 @@ def cpu_parallel_barrier():
"""Synchronize all workers in the current CPU parallel launch.
Every worker must reach this operation at the same program point. Place it
- inside a ``pragma_parallel_launch_point`` scope, outside parallel loops,
+ inside a ``parallel_launch`` region, outside parallel loops,
whose iteration counts can differ between workers. Writes before the
barrier
are visible to all workers after it.
diff --git a/python/tvm/tirx/script/ir_builder/parser_protocol.py
b/python/tvm/tirx/script/ir_builder/parser_protocol.py
index 4f88013707..34eb8b90e3 100644
--- a/python/tvm/tirx/script/ir_builder/parser_protocol.py
+++ b/python/tvm/tirx/script/ir_builder/parser_protocol.py
@@ -980,6 +980,24 @@ def region(
return _ffi_api.Region(op, args, [] if body_params is None else
body_params, attrs)
+def device_context(device_type: Expr, device_id: Expr) -> frame.RegionFrame:
+ """Supply lexical device context for allocation and packed-call lowering.
+
+ This region does not change the active runtime device.
+ """
+ return region("tirx.device_context", [device_type, device_id])
+
+
+def compute_scope(name: str) -> frame.RegionFrame:
+ """Outline the body as a named CPU compute helper."""
+ return region("tirx.compute_scope", [_StringImm(name)])
+
+
+def parallel_launch() -> frame.RegionFrame:
+ """Launch a CPU worker team around parallel loops and team barriers."""
+ return region("tirx.parallel_launch", [])
+
+
def launch_thread(thread_tag: str, extent: Expr) -> frame.RegionFrame:
"""Launch a hardware or virtual thread with a fresh lexical variable.
@@ -1017,8 +1035,10 @@ __all__ = [
"buffer_store",
"call_global_var_",
"check_well_formed_",
+ "compute_scope",
"continue_",
"decl_mutable_cell_",
+ "device_context",
"device_entry",
"else_",
"emit",
@@ -1043,6 +1063,7 @@ __all__ = [
"not_",
"or_",
"parallel",
+ "parallel_launch",
"prim_func",
"range_",
"region",
diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index bb42c87b5b..62d0113ac9 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -508,7 +508,7 @@ llvm::BasicBlock* CodeGenCPU::CheckCallSuccess(llvm::Value*
retcode) {
return end_block;
}
-void CodeGenCPU::CreateComputeScope(const AttrStmtNode* op) {
+void CodeGenCPU::CreateComputeScope(const RegionStmtNode* op) {
EmitDebugLocation(op);
/*! \brief maintain states that should be guarded when step into compute
scope */
struct ComputeScopeStates {
@@ -555,8 +555,7 @@ void CodeGenCPU::CreateComputeScope(const AttrStmtNode* op)
{
// $xxx_compute_ functions are not global. They should be marked as static
(via InternalLinkage)
// to call them correctly on MIPS platform (CALL16 issue)
// Linkage ld Error: CALL16 reloc at 0x290 not against global symbol
- const StringImmNode* value = op->value.as<StringImmNode>();
- TVM_FFI_ICHECK(value != nullptr);
+ const auto* value = op->args[0].as<StringImmNode>();
llvm::Function* fcompute = llvm::Function::Create(ftype,
llvm::Function::InternalLinkage,
MakeStringRef(value->value), module_.get());
SetComputeScopeAttributes(fcompute);
@@ -1131,18 +1130,25 @@ void CodeGenCPU::Dispatch_(const AssertStmtNode* op) {
CodeGenLLVM::Dispatch_(op);
}
+void CodeGenCPU::Dispatch_(const RegionStmtNode* op) {
+ EmitDebugLocation(op);
+ if (op->op.same_as(tirx::builtin::compute_scope())) {
+ CreateComputeScope(op);
+ } else if (op->op.same_as(tirx::builtin::parallel_launch())) {
+ CreateParallelLaunch(op->body, 0, "pragma_parallel");
+ } else {
+ CodeGenLLVM::Dispatch_(op);
+ }
+}
+
void CodeGenCPU::Dispatch_(const AttrStmtNode* op) {
EmitDebugLocation(op);
- if (op->attr_key == tirx::attr::compute_scope) {
- this->CreateComputeScope(op);
- } else if (tirx::attr::IsPragmaKey(op->attr_key)) {
+ if (tirx::attr::IsPragmaKey(op->attr_key)) {
if (op->attr_key == "pragma_parallel_stride_pattern") {
TVM_FFI_THROW(ValueError)
<< "pragma_parallel_stride_pattern is retired; annotate every
affected parallel For "
<< "with parallel_stride_pattern=True, including later loops in the
same launch "
<< "outside the former attribute body";
- } else if (op->attr_key == "pragma_parallel_launch_point") {
- CreateParallelLaunch(op->body, 0, "pragma_parallel");
} else if (op->attr_key == "pragma_parallel_barrier_when_finish") {
TVM_FFI_THROW(ValueError)
<< "pragma_parallel_barrier_when_finish is retired; place
cpu_parallel_barrier() "
diff --git a/src/target/llvm/codegen_cpu.h b/src/target/llvm/codegen_cpu.h
index a73210d151..dba07332d8 100644
--- a/src/target/llvm/codegen_cpu.h
+++ b/src/target/llvm/codegen_cpu.h
@@ -72,6 +72,7 @@ class CodeGenCPU : public CodeGenLLVM {
void AddMainFunction(const std::string& entry_func_name) override;
std::unique_ptr<llvm::Module> Finish() override;
void Dispatch_(const AssertStmtNode* op) override;
+ void Dispatch_(const RegionStmtNode* op) override;
void Dispatch_(const AttrStmtNode* op) override;
void Dispatch_(const ForNode* op) override;
llvm::Value* CreateIntrinsic(const CallNode* op) override;
@@ -147,7 +148,7 @@ class CodeGenCPU : public CodeGenLLVM {
// Create parallel launch
void CreateParallelLaunch(const Stmt& body, int num_task, std::string name =
"");
// Create a new compute scope.
- void CreateComputeScope(const AttrStmtNode* op);
+ void CreateComputeScope(const RegionStmtNode* op);
// Check if the call to packed function is successful
// if not directly finalize function and pass on return code.
// return the end block after the check
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 33a00a7d10..9d8f6da7b2 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -2256,6 +2256,10 @@ void CodeGenLLVM::Dispatch_(const RegionStmtNode* op) {
With<sym::ConstraintContext> thread_scope(analyzer_, var >= 0 && var <
extent);
this->Dispatch(op->body);
var_map_.erase(var.get());
+ } else if (op->op.same_as(tirx::builtin::device_context()) ||
+ op->op.same_as(tirx::builtin::compute_scope()) ||
+ op->op.same_as(tirx::builtin::parallel_launch())) {
+ this->Dispatch(op->body);
} else {
TVM_FFI_THROW(ValueError) << "Unsupported region op " << op->op;
}
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index aa6313b76b..7e693f5c24 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -1303,6 +1303,10 @@ void CodeGenC::Dispatch_(const RegionStmtNode* op) {
BindThreadIndex(var, tag);
this->PrintStmt(op->body);
var_idmap_.erase(var.get());
+ } else if (op->op.same_as(tirx::builtin::device_context()) ||
+ op->op.same_as(tirx::builtin::compute_scope()) ||
+ op->op.same_as(tirx::builtin::parallel_launch())) {
+ this->PrintStmt(op->body);
} else {
TVM_FFI_THROW(ValueError) << "Unsupported region op " << op->op;
}
diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc
index d11518e973..874861b318 100644
--- a/src/tirx/ir/stmt.cc
+++ b/src/tirx/ir/stmt.cc
@@ -813,6 +813,11 @@ RegionStmt::RegionStmt(Op op, ffi::Array<Expr> args,
ffi::Array<Var> body_params
signature.args = args;
signature.attrs = attrs;
op.Validate(&signature);
+ 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";
+ }
std::unordered_set<const VarNode*> definitions;
for (const auto& vars : {body_params, result_vars}) {
for (const Var& var : vars) {
diff --git a/src/tirx/op/builtin.cc b/src/tirx/op/builtin.cc
index d791c0d0b3..b4332fcb15 100644
--- a/src/tirx/op/builtin.cc
+++ b/src/tirx/op/builtin.cc
@@ -191,6 +191,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TVM_DEFINE_CACHED_OP_GETTER(reinterpret, "tirx.reinterpret")
TVM_DEFINE_CACHED_OP_GETTER(launch_thread, "tirx.launch_thread")
TVM_DEFINE_CACHED_OP_GETTER(device_entry, "tirx.device_entry")
+TVM_DEFINE_CACHED_OP_GETTER(device_context, "tirx.device_context")
+TVM_DEFINE_CACHED_OP_GETTER(compute_scope, "tirx.compute_scope")
+TVM_DEFINE_CACHED_OP_GETTER(parallel_launch, "tirx.parallel_launch")
TVM_DEFINE_CACHED_OP_GETTER(thread_return, "tirx.thread_return")
TVM_DEFINE_CACHED_OP_GETTER(filter, "tirx.filter")
TVM_DEFINE_CACHED_OP_GETTER(selector, "tirx.selector")
@@ -276,6 +279,15 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.launch_thread", "Bind a thread index within a body with a launch
extent.")
.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<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"));
+ OpDef("tirx.compute_scope", "Outline a named CPU compute region.")
+ .signature(sig::arg<StringImm>("name"))
+ .set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"));
+ OpDef("tirx.parallel_launch", "Launch a CPU worker team around a region.")
+ .signature()
+ .set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"));
OpDef("tirx.thread_return")
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.thread_return"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc
index 5088fe5289..0ffc17c34d 100644
--- a/src/tirx/script/printer/stmt.cc
+++ b/src/tirx/script/printer/stmt.cc
@@ -367,6 +367,17 @@ ffi::Optional<ExprDoc>
RegionStmtDocTranslate(DocTranslatorObj* d, ffi::AnyView
->Attr("launch_thread")
->Call({LiteralDoc::Str(stmt->args[0].as_or_throw<StringImm>()->value,
std::nullopt),
d->Translate(stmt->args[1].as_or_throw<PrimExpr>()).value()});
+ } else if (stmt->op.same_as(tirx::builtin::device_context())) {
+ rhs = NamespaceDoc("tirx")
+ ->Attr("device_context")
+ ->Call({d->Translate(stmt->args[0]).value(),
d->Translate(stmt->args[1]).value()});
+ } else if (stmt->op.same_as(tirx::builtin::compute_scope())) {
+ rhs =
+ NamespaceDoc("tirx")
+ ->Attr("compute_scope")
+
->Call({LiteralDoc::Str(stmt->args[0].as_or_throw<StringImm>()->value,
std::nullopt)});
+ } else if (stmt->op.same_as(tirx::builtin::parallel_launch())) {
+ rhs = NamespaceDoc("tirx")->Attr("parallel_launch")->Call({});
} else {
ffi::Array<ExprDoc> args;
for (const Expr& arg : stmt->args)
args.push_back(d->Translate(arg).value());
diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc
index 0362497aa7..a13c7130c6 100644
--- a/src/tirx/transform/ir_utils.cc
+++ b/src/tirx/transform/ir_utils.cc
@@ -304,12 +304,8 @@ UnchangedOr<Stmt> IRConvertSSA::Mutate_(const
AttrStmtNode* op, InplaceMode inpl
function_scope_var_remap_.insert({var.get(), new_var});
var = new_var;
} else {
- // The AttrStmt refers to an undefined variable. This is
- // allowed for some attributes, such as
- // "pragma_parallel_launch_point", which annotates a variable
- // that is about to occur in a ForNode. In these cases, the
- // ForNode and the AttrStmt must continue using the same
- // variable defintion.
+ // An attribute can refer to a variable defined by a later ForNode.
+ // Preserve the same variable definition in the attribute and the loop.
//
// Preserve the annotated variable's identity for later definitions
// and independent functions, without introducing a lexical body binding.
diff --git a/src/tirx/transform/lower_tvm_builtin.cc
b/src/tirx/transform/lower_tvm_builtin.cc
index a2c1414546..2046074656 100644
--- a/src/tirx/transform/lower_tvm_builtin.cc
+++ b/src/tirx/transform/lower_tvm_builtin.cc
@@ -337,39 +337,49 @@ class BuiltinLower : public StmtExprMutator {
return SeqStmt({alloc_bind, alloc_nullptr_check});
}
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- if (op->attr_key == attr::device_id) {
- auto cache = device_id_;
- device_id_ = op->value.as_or_throw<PrimExpr>();
- Stmt out = scope_.WithNewScope([&]() -> Stmt {
- Stmt body = this->Mutate(op->body,
inplace_mode).ValueOrUnchanged(op->body);
- return AppendPendingFrees(body);
- });
- device_id_ = cache;
- return out;
- } else if (op->attr_key == attr::device_type) {
- auto cache = device_type_;
- device_type_ = op->value.as_or_throw<PrimExpr>();
- Stmt out = scope_.WithNewScope([&]() -> Stmt {
- Stmt body = this->Mutate(op->body,
inplace_mode).ValueOrUnchanged(op->body);
- return AppendPendingFrees(body);
+ UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) final {
+ if (op->op.same_as(builtin::device_context())) {
+ PrimExpr device_type = this->Mutate(op->args[0], inplace_mode)
+ .ValueOrUnchanged(op->args[0])
+ .as_or_throw<PrimExpr>();
+ PrimExpr device_id = this->Mutate(op->args[1], inplace_mode)
+ .ValueOrUnchanged(op->args[1])
+ .as_or_throw<PrimExpr>();
+ auto saved_type = device_type_;
+ auto saved_id = device_id_;
+ device_type_ = device_type;
+ device_id_ = device_id;
+ Stmt body = scope_.WithNewScope([&]() -> Stmt {
+ return AppendPendingFrees(this->Mutate(op->body,
inplace_mode).ValueOrUnchanged(op->body));
});
- device_type_ = cache;
- return out;
- } else {
- return scope_.WithNewScope([&]() -> Stmt {
- Stmt visited =
- StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- if (!scope_.Current().pending_frees.empty()) {
- const auto* attr = visited.as<AttrStmtNode>();
- if (attr) {
- return AttrStmt(attr->node, attr->attr_key, attr->value,
AppendPendingFrees(attr->body),
- attr->span);
- }
+ device_type_ = saved_type;
+ device_id_ = saved_id;
+ return body;
+ }
+ return scope_.WithNewScope([&]() -> Stmt {
+ RegionStmt region = StmtExprMutator::Mutate_(op, inplace_mode)
+ .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
+ .as_or_throw<RegionStmt>();
+ if (!scope_.Current().pending_frees.empty()) {
+ region.CopyOnWrite()->body = AppendPendingFrees(region->body);
+ }
+ return region;
+ });
+ }
+
+ UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
+ return scope_.WithNewScope([&]() -> Stmt {
+ Stmt visited =
+ StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
+ if (!scope_.Current().pending_frees.empty()) {
+ const auto* attr = visited.as<AttrStmtNode>();
+ if (attr) {
+ return AttrStmt(attr->node, attr->attr_key, attr->value,
AppendPendingFrees(attr->body),
+ attr->span);
}
- return visited;
- });
- }
+ }
+ return visited;
+ });
}
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
auto min_result = this->Mutate(op->min, inplace_mode);
@@ -467,7 +477,7 @@ class BuiltinLower : public StmtExprMutator {
StringImm GetDeviceMethodName(const char* method_name) const {
TVM_FFI_ICHECK(device_type_) << "Method " << method_name << " requires the
device type, "
- << "but occurred outside of a \"device_type\"
annotation";
+ << "but occurred outside of a
\"device_context\" region";
auto as_int = device_type_.as<IntImmNode>();
TVM_FFI_ICHECK(as_int) << "Method " << method_name
diff --git a/src/tirx/transform/make_packed_api.cc
b/src/tirx/transform/make_packed_api.cc
index 7ccca70c4c..b3b554f067 100644
--- a/src/tirx/transform/make_packed_api.cc
+++ b/src/tirx/transform/make_packed_api.cc
@@ -231,7 +231,6 @@ PrimFunc MakePackedAPI(PrimFunc func) {
}
auto* func_ptr = func.CopyOnWrite();
- const Stmt nop = Evaluate(0);
// Data field definitions
Var v_self_handle("self_handle", PointerType::VoidPointerTy());
@@ -251,8 +250,6 @@ PrimFunc MakePackedAPI(PrimFunc func) {
auto result = binder.Finalize();
bool need_set_device = result.var_defs.count(device_id.get());
- std::vector<Stmt> seq_check;
-
// signature: (void* handle, TVMFFIAny* packed_args, int num_args,
TVMFFIAny* v_result)
ffi::Array<Var> args{v_self_handle, v_packed_args, v_num_packed_args,
v_result};
@@ -265,13 +262,10 @@ PrimFunc MakePackedAPI(PrimFunc func) {
Stmt body = ffi::make_object<ReturnRewriter>(v_result)
->Mutate(func_ptr->body.value(), InplaceMode::kAllow)
.ValueOrUnchanged(func_ptr->body.value());
- body = AttrStmt(0, attr::compute_scope, StringImm(name_hint + "_compute_"),
body);
+ body = RegionStmt(builtin::compute_scope(), {StringImm(name_hint +
"_compute_")}, {}, DictAttrs(),
+ body);
// Set device context
if (need_set_device) {
- ffi::Any node = ffi::String("default");
- seq_check.push_back(AttrStmt(node, attr::device_id,
device_id.as_or_throw<PrimExpr>(), nop));
- seq_check.push_back(AttrStmt(node, attr::device_type, device_type, nop));
-
if (runtime::DeviceAPI::NeedSetDevice(target_device_type)) {
Stmt set_device = Evaluate(Call(PrimType::Int(32),
builtin::tvm_call_packed(),
{StringImm(runtime::symbol::tvm_set_device), device_type,
@@ -287,7 +281,11 @@ PrimFunc MakePackedAPI(PrimFunc func) {
// Tensor declarations and alignment assumptions are ordinary statements,
// not scopes with a body hole for MergeNest to fill.
body = SeqStmt::Flatten(result.decl_buffers, body);
- body = MergeNest({std::move(result.init_nest), seq_check,
std::move(result.asserts)}, body);
+ body = MergeNest(std::move(result.asserts), body);
+ if (need_set_device) {
+ body = RegionStmt(builtin::device_context(), {device_type, device_id}, {},
DictAttrs(), body);
+ }
+ body = MergeNest(std::move(result.init_nest), body);
func_ptr->body = body;
func_ptr->params = args;
diff --git a/tests/python/s_tir/script/test_s_tir_script_printer.py
b/tests/python/s_tir/script/test_s_tir_script_printer.py
index 7d9e1ec37b..102232b06c 100644
--- a/tests/python/s_tir/script/test_s_tir_script_printer.py
+++ b/tests/python/s_tir/script/test_s_tir_script_printer.py
@@ -1149,26 +1149,26 @@ def opt_conv_tensorcore_mod_host():
T.evaluate(T.tvm_struct_set(stack_value, 1, 12, T.cast(dev_id,
"int64")))
stack_tcode[1] = 0
T.evaluate(T.tvm_call_packed_lowered("__tvm_set_device", stack_value,
0, 2))
- T.attr(0, "compute_scope", "default_function_compute_")
- T.evaluate(T.tvm_struct_set(stack_value, 0, 12, A))
- stack_tcode[0] = 3
- T.evaluate(T.tvm_struct_set(stack_value, 1, 12, W))
- stack_tcode[1] = 3
- T.evaluate(T.tvm_struct_set(stack_value, 2, 12, Conv))
- stack_tcode[2] = 3
- T.evaluate(T.tvm_struct_set(stack_value, 3, 12, T.cast(196, "int64")))
- stack_tcode[3] = 0
- T.evaluate(T.tvm_struct_set(stack_value, 4, 12, T.cast(2, "int64")))
- stack_tcode[4] = 0
- T.evaluate(T.tvm_struct_set(stack_value, 5, 12, T.cast(4, "int64")))
- stack_tcode[5] = 0
- T.evaluate(T.tvm_struct_set(stack_value, 6, 12, T.cast(4, "int64")))
- stack_tcode[6] = 0
- T.evaluate(T.tvm_struct_set(stack_value, 7, 12, T.cast(2, "int64")))
- stack_tcode[7] = 0
- T.evaluate(T.tvm_struct_set(stack_value, 8, 12, T.cast(32, "int64")))
- stack_tcode[8] = 0
- T.evaluate(T.tvm_call_packed_lowered("default_function_kernel0",
stack_value, 0, 9))
+ with T.compute_scope("default_function_compute_"):
+ T.evaluate(T.tvm_struct_set(stack_value, 0, 12, A))
+ stack_tcode[0] = 3
+ T.evaluate(T.tvm_struct_set(stack_value, 1, 12, W))
+ stack_tcode[1] = 3
+ T.evaluate(T.tvm_struct_set(stack_value, 2, 12, Conv))
+ stack_tcode[2] = 3
+ T.evaluate(T.tvm_struct_set(stack_value, 3, 12, T.cast(196,
"int64")))
+ stack_tcode[3] = 0
+ T.evaluate(T.tvm_struct_set(stack_value, 4, 12, T.cast(2,
"int64")))
+ stack_tcode[4] = 0
+ T.evaluate(T.tvm_struct_set(stack_value, 5, 12, T.cast(4,
"int64")))
+ stack_tcode[5] = 0
+ T.evaluate(T.tvm_struct_set(stack_value, 6, 12, T.cast(4,
"int64")))
+ stack_tcode[6] = 0
+ T.evaluate(T.tvm_struct_set(stack_value, 7, 12, T.cast(2,
"int64")))
+ stack_tcode[7] = 0
+ T.evaluate(T.tvm_struct_set(stack_value, 8, 12, T.cast(32,
"int64")))
+ stack_tcode[8] = 0
+ T.evaluate(T.tvm_call_packed_lowered("default_function_kernel0",
stack_value, 0, 9))
return opt_conv_tensorcore_mod_host
@@ -1927,21 +1927,20 @@ def tvm_struct_set_generated_in_cpp():
class Module:
@Ts.prim_func
def tir_packed_call(A: T.Tensor(16)):
- T.attr(0, "device_id", 0)
- T.attr(0, "device_type", 0)
- T.evaluate(
- T.tvm_call_cpacked(
- "tvm_test_cpacked",
- T.tvm_stack_make_array(
- A.data,
- T.tvm_stack_make_shape(16),
- T.reinterpret(T.uint64(0), dtype="handle"),
- T.uint32(1),
- T.Cast("float32", 0),
- 0,
- ),
+ with T.device_context(0, 0):
+ T.evaluate(
+ T.tvm_call_cpacked(
+ "tvm_test_cpacked",
+ T.tvm_stack_make_array(
+ A.data,
+ T.tvm_stack_make_shape(16),
+ T.reinterpret(T.uint64(0), dtype="handle"),
+ T.uint32(1),
+ T.Cast("float32", 0),
+ 0,
+ ),
+ )
)
- )
return tvm.tirx.transform.LowerTVMBuiltin()(Module)
diff --git
a/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
b/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
index 7ed2f330be..82b7938d91 100644
--- a/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
+++ b/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
@@ -40,8 +40,8 @@ def test_lower_call_packed():
C: T.Tensor((64, 64), "float32"),
):
T.func_attr({"target": tvm.target.Target("llvm")})
- T.attr("", "device_id", T.int32(0))
- T.call_packed("tvm.test_matmul", A, B, C)
+ with T.device_context(1, 0):
+ T.call_packed("tvm.test_matmul", A, B, C)
@I.ir_module(check_well_formed=False)
class Expected:
@@ -218,11 +218,10 @@ def test_lower_device_allocate():
@T.prim_func
def main():
T.func_attr({"target": T.target("llvm")})
- T.attr("dummy", "device_type", 2) # kDLCuda
- T.attr("dummy", "device_id", 0)
- ptr = T.alloc_tensor((16,), "float32")
- buf = T.decl_tensor(16, "float32", data=ptr.data)
- buf[0] = 0.0
+ with T.device_context(2, 0): # kDLCuda
+ ptr = T.alloc_tensor((16,), "float32")
+ buf = T.decl_tensor(16, "float32", data=ptr.data)
+ buf[0] = 0.0
After = tvm.tirx.transform.LowerTVMBuiltin()(Before)
# Verify the lowered module can be printed (no crash)
@@ -242,11 +241,10 @@ def test_lower_cpu_allocation():
@T.prim_func
def main():
T.func_attr({"target": T.target("llvm")})
- T.attr("dummy", "device_type", 1) # kDLCPU
- T.attr("dummy", "device_id", 0)
- ptr = T.alloc_tensor((16,), "float32")
- buf = T.decl_tensor(16, "float32", data=ptr.data)
- buf[0] = 0.0
+ with T.device_context(1, 0): # kDLCPU
+ ptr = T.alloc_tensor((16,), "float32")
+ buf = T.decl_tensor(16, "float32", data=ptr.data)
+ buf[0] = 0.0
@I.ir_module
class Expected:
@@ -268,8 +266,7 @@ def test_lower_allocate_requires_device_id():
class Before:
@T.prim_func
def main():
- T.func_attr({"target": T.target("llvm")})
- T.attr("dummy", "device_type", 2) # kDLCuda
+ T.func_attr({"target": T.target("cuda"), "tirx.is_host_func":
True})
ptr = T.alloc_tensor((16,), "float32")
buf = T.decl_tensor(16, "float32", data=ptr.data)
buf[0] = 0.0
@@ -281,8 +278,8 @@ def test_lower_allocate_requires_device_id():
def test_lower_allocate_requires_device_type():
"""If device type is missing, error.
- The device type can be inferred either from the `"device_type"`
- statement attribute, or from the `"target"` function attribute.
+ The device type can be inferred either from a device_context
+ region or from the `"target"` function attribute.
Here, we provide neither. The `"tirx.is_host_func"` attribute is
provided as otherwise the function would be skipped altogether by
LowerTVMBuiltin.
@@ -293,7 +290,6 @@ def test_lower_allocate_requires_device_type():
@T.prim_func
def main():
T.func_attr({"tirx.is_host_func": True})
- T.attr("dummy", "device_id", 0)
ptr = T.alloc_tensor((1024 * 1024,), "float32")
buf = T.decl_tensor(1024 * 1024, "float32", data=ptr.data)
buf[0] = 0.0
@@ -306,7 +302,7 @@ def test_lower_cpu_alloc_with_function_attr():
"""CPU allocations can be handled at codegen time
Like `test_lower_cpu_allocation`, but the device type is taken from
- the function attribute. The `AttrStmt` can override the device
+ the function attribute. A device_context region can override the device
type for allocations within its scope, but it defaults to the
function's target.
"""
diff --git a/tests/python/tirx-transform/test_tir_transform_make_packed_api.py
b/tests/python/tirx-transform/test_tir_transform_make_packed_api.py
index bcb5ce14b9..00efe234ab 100644
--- a/tests/python/tirx-transform/test_tir_transform_make_packed_api.py
+++ b/tests/python/tirx-transform/test_tir_transform_make_packed_api.py
@@ -43,7 +43,9 @@ def _find_compute_scope(func):
result = None
def _visitor(stmt):
- if isinstance(stmt, tirx.AttrStmt) and stmt.attr_key ==
"compute_scope":
+ if isinstance(stmt, tirx.RegionStmt) and stmt.op.same_as(
+ tvm.ir.Op.get("tirx.compute_scope")
+ ):
nonlocal result
result = stmt
@@ -194,7 +196,7 @@ def test_zero_arg_function():
"TypeError",
["Expected ", "0", " arguments", " when calling:\n `",
"func_without_arg()", "`"],
)
- with T.attr(0, "compute_scope", "func_without_arg_compute_"):
+ with T.compute_scope("func_without_arg_compute_"):
T.tvm_struct_set(result, 0, 13, 1)
T.tvm_struct_set(result, 0, 14, 0)
T.tvm_struct_set(result, 0, 15, T.Cast("int64", T.int64(42)))
@@ -292,7 +294,7 @@ def test_int_parameter():
],
)
arg: T.let[T.int32] = T.Cast("int32", T.tvm_struct_get(args, 0,
15, "int64"))
- with T.attr(0, "compute_scope", "main_compute_"):
+ with T.compute_scope("main_compute_"):
if arg > 0:
T.tvm_struct_set(result, 0, 13, 1)
T.tvm_struct_set(result, 0, 14, 0)
@@ -359,7 +361,7 @@ def test_bool_parameter():
],
)
arg: T.let[T.bool] = T.Cast("bool", T.tvm_struct_get(args, 0, 15,
"int64"))
- with T.attr(0, "compute_scope", "main_compute_"):
+ with T.compute_scope("main_compute_"):
if arg:
T.tvm_struct_set(result, 0, 13, 1)
T.tvm_struct_set(result, 0, 14, 0)
@@ -430,7 +432,7 @@ def test_float_parameter():
T.Cast("float32", T.tvm_struct_get(args, 0, 15, "float64")),
T.Cast("float32", T.tvm_struct_get(args, 0, 15, "int64")),
)
- with T.attr(0, "compute_scope", "main_compute_"):
+ with T.compute_scope("main_compute_"):
if arg > T.float32(0.0):
T.tvm_struct_set(result, 0, 13, 1)
T.tvm_struct_set(result, 0, 14, 0)
diff --git a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py
b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py
index 63c7365a36..c9da13728a 100644
--- a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py
+++ b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py
@@ -197,7 +197,7 @@ def test_parallel_alloc():
@T.prim_func
def func2(n: T.int32):
for t in T.serial(n):
- with T.attr(T.int32(1), "pragma_scope", "parallel_launch_point"):
+ with T.parallel_launch():
for i in T.parallel(n):
for j in range(10):
A = T.alloc_tensor((n,))
@@ -206,7 +206,8 @@ def test_parallel_alloc():
mod = tvm.IRModule.from_expr(func2)
body = tvm.tirx.transform.StorageRewrite()(mod)["func2"]
- assert _is_buffer_binding(body.body.body.body.body[0], "tirx.alloc_tensor")
+ # The launch body owns the allocation executed by each worker.
+ assert _is_buffer_binding(body.body.body.body[0], "tirx.alloc_tensor")
def test_while_alloc():
