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 b365abb18b [REFACTOR][IR] Localize attribute ownership and fix
documentation (#20801)
b365abb18b is described below
commit b365abb18b20ed27bf47879b49aa253459f3af85
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 19 22:29:38 2026 -0400
[REFACTOR][IR] Localize attribute ownership and fix documentation (#20801)
Use `script.namespace` for S-TIR construction, printing and lowering,
removing the redundant `s_tir` flag. Give Trainium its own
`trn.num_input_args` contract and annotate kernel arguments in a backend
pass after the ordinary `SplitHostDevice` pass and before
`MakePackedAPI`, deriving each kernel's input count from the arguments
it actually captures.
Move external-codegen constant metadata into Relax and remove unused
C-host device-context plumbing.
Repair duplicate Python API registrations and tutorial section
boundaries so canonical IR objects and tutorial headings render
correctly.
---
docs/arch/tvmscript.rst | 4 +-
docs/how_to/tutorials/optimize_llm.py | 1 +
docs/reference/api/python/script/ir_builder.rst | 5 +-
docs/reference/api/python/script/script.rst | 7 +-
include/tvm/relax/expr.h | 10 ++
python/tvm/s_tir/pipeline.py | 2 +-
src/relax/backend/vm/vm_shape_lower.cc | 1 -
src/relax/transform/compute_prim_value.cc | 1 -
src/relax/transform/fuse_tir.cc | 1 -
src/relax/transform/run_codegen.cc | 3 +-
src/s_tir/script/ir_builder/frame.cc | 15 ---
src/te/operation/create_function.cc | 4 +-
.../relax/script/test_relax_script_printer.py | 6 +-
.../python/s_tir/base/test_compilation_pipeline.py | 2 +
.../s_tir/script/test_s_tir_script_ir_builder.py | 4 +-
.../s_tir/script/test_s_tir_script_namespace.py | 32 +----
.../s_tir/script/test_s_tir_script_printer.py | 4 +-
.../script/test_script_constructor_contracts.py | 2 +-
tirx/include/tvm/backend/trn/attr.h | 41 ++++++
tirx/include/tvm/ir/function.h | 22 +---
tirx/include/tvm/ir/module.h | 13 --
tirx/python/tvm_tirx/backend/trn/pipeline.py | 1 +
.../tvm_tirx/backend/trn/transform/__init__.py | 9 +-
.../trn/transform/annotate_kernel_arguments.py | 137 +++++++++++++++++++++
tirx/src/backend/cuda/codegen/codegen_cuda_host.cc | 2 +-
tirx/src/backend/trn/codegen/codegen_trn.cc | 7 +-
tirx/src/codegen/c/codegen_c_host.cc | 14 +--
tirx/src/codegen/c/codegen_c_host.h | 4 +-
tirx/src/tirx/transform/split_host_device.cc | 38 ++----
tirx/tests/python/tirx/codegen/test_nki.py | 6 +-
30 files changed, 253 insertions(+), 145 deletions(-)
diff --git a/docs/arch/tvmscript.rst b/docs/arch/tvmscript.rst
index e4f49c8909..43cf0ccc2e 100644
--- a/docs/arch/tvmscript.rst
+++ b/docs/arch/tvmscript.rst
@@ -140,7 +140,9 @@ lookup in the active module.
Public script namespaces expose decorators and construction operations. The
same
underlying builders can also be used directly from Python. The transpiler
therefore
needs no separate mutable IR representation: it generates calls to this
construction
-protocol.
+protocol. ``script.namespace`` records the construction namespace for printing
+and dialect lowering. S-TIR functions use ``"s_tir"``; their builders own block
+completion and validation.
Ordinary IR constructors can be used directly in parsed source. Shared exports
such as ``I.Call`` and ``T.Range`` use the same constructor contracts as
diff --git a/docs/how_to/tutorials/optimize_llm.py
b/docs/how_to/tutorials/optimize_llm.py
index c8f4cc1868..263a4c26ce 100644
--- a/docs/how_to/tutorials/optimize_llm.py
+++ b/docs/how_to/tutorials/optimize_llm.py
@@ -45,6 +45,7 @@ TinyLlama model from Hugging Face and deploy it on various
devices.
# - **Build and Universal Deployment**: Build the optimized model to a
deployable module to the
# universal runtime, and execute it on different devices, such as CPU, GPU,
or other accelerators.
#
+
######################################################################
# Construct the model architecture
# --------------------------------
diff --git a/docs/reference/api/python/script/ir_builder.rst
b/docs/reference/api/python/script/ir_builder.rst
index 564e9c1db6..3b3d84ce4d 100644
--- a/docs/reference/api/python/script/ir_builder.rst
+++ b/docs/reference/api/python/script/ir_builder.rst
@@ -20,10 +20,13 @@ tvm_tirx.script.ir_builder
tvm_tirx.script.ir_builder
**************************
+
+Shared IR classes are documented in :mod:`tvm_tirx.ir`.
+
.. automodule:: tvm_tirx.script.ir_builder
:members:
:imported-members:
- :exclude-members: Call, DataTypeImm, FuncType, GenericConst, MissingType,
Op, PrimType, Range, StringImm, StringType, TensorRegion, TensorRegionType,
Tuple, Type, Var
+ :exclude-members: Bind, Call, DataTypeImm, FuncType, GenericConst,
MissingType, NoneType, Op, PrimType, Range, StringImm, StringType,
TensorRegion, TensorRegionType, Tuple, Type, Var
tvm.relax.script.ir_builder.distributed
***************************************
diff --git a/docs/reference/api/python/script/script.rst
b/docs/reference/api/python/script/script.rst
index 8fd2ae0545..43cc6676cf 100644
--- a/docs/reference/api/python/script/script.rst
+++ b/docs/reference/api/python/script/script.rst
@@ -29,10 +29,11 @@ tvm.relax.script
****************
Operator attributes re-exported by this namespace are documented in
:mod:`tvm.relax.op`.
+The shared ``NoneType`` class is documented in :mod:`tvm_tirx.ir`.
.. automodule:: tvm.relax.script
:members:
- :exclude-members: ExternFunc, ShapeExpr, TupleGetItem, Range, function,
macro,
+ :exclude-members: ExternFunc, NoneType, ShapeExpr, TupleGetItem, Range,
function, macro,
ArgmaxArgminAttrs, ArgsortAttrs, AstypeAttrs, CallInplacePackedAttrs,
CallTIRInplaceAttrs,
CallTIRWithGradAttrs, ConcatAttrs, EinsumAttrs, ExpandDimsAttrs,
FlipAttrs,
GatherElementsAttrs, GatherNDAttrs, HintOnDeviceAttrs, IndexPutAttrs,
InitAttrs,
@@ -56,8 +57,8 @@ The canonical ``MemoryScope`` class is documented in
:mod:`tvm_tirx.runtime`.
.. automodule:: tvm_tirx.script.tirx
:members:
- :exclude-members: AlreadyEmitted, DataTypeImm, FuncType, GenericConst,
IRBuilder,
- IRModuleFrame, MemoryScope, MissingType, Op, PrimType, StringImm,
StringType, TensorLoad, TensorRegion,
+ :exclude-members: AlreadyEmitted, Bind, DataTypeImm, FuncType,
GenericConst, IRBuilder,
+ IRModuleFrame, MemoryScope, MissingType, NoneType, Op, PrimType,
StringImm, StringType, TensorLoad, TensorRegion,
TensorRegionType, Type, Range, meta_var, Var, Call, CommReducer, Reduce,
FloatImm, IntImm, Cast, Add, Sub, Mul, Div, Mod, FloorDiv, FloorMod, LShift,
RShift, BitwiseAnd, BitwiseOr, BitwiseXor, BitwiseNot, Min, Max, EQ, NE, LT,
LE, GT, GE, And, Or, Not, Select, Ramp, Broadcast, Shuffle, CallEffectKind,
IterVar, ComposeLayout, DtypeConstructor, ThreadScope, Iter, Layout,
LetAnnotation, LocalVectorAnnotation, TileLayout, CallFFIKernelAttr,
TensorMapEncodeTiledAttr, function, jit, inli [...]
.. autofunction:: tvm_tirx.script.tirx.function
diff --git a/include/tvm/relax/expr.h b/include/tvm/relax/expr.h
index ec27709cb3..7374ab124d 100644
--- a/include/tvm/relax/expr.h
+++ b/include/tvm/relax/expr.h
@@ -495,6 +495,16 @@ TVM_DLL Expr GetShapeOf(const Expr& expr);
// TODO(@sunggg): Investigate the exact usage of kComposite,
kPartitionedFromPattern, and
// kPrimitive.
namespace attr {
+/*!
+ * \brief All the named Tensors accumulated during compilation by external
codegen.
+ * Generally the associated runtime::Module will indicate it requires bindings
for these names,
+ * and during module initialization these bindings will be recovered from a
ConstLoaderModule.
+ *
+ * Type: ffi::Map<ffi::String, Tensor>
+ */
+inline constexpr ffi::reflection::AttrKey<ffi::Map<ffi::String, ffi::Tensor>>
kConstNameToConstant{
+ "const_name_to_constant"};
+
/*! \brief Mark the function as a primitive function. */
inline constexpr ffi::reflection::AttrKey<bool> kPrimitive{"Primitive"};
/*!
diff --git a/python/tvm/s_tir/pipeline.py b/python/tvm/s_tir/pipeline.py
index 56390ae988..dea38d3196 100644
--- a/python/tvm/s_tir/pipeline.py
+++ b/python/tvm/s_tir/pipeline.py
@@ -127,7 +127,7 @@ def _partition_dialects(mod):
for gv, func in mod.functions.items():
group = (
scheduled
- if isinstance(func, tirx.Function) and func.attrs.get("s_tir",
False)
+ if isinstance(func, tirx.Function) and
func.attrs.get("script.namespace") == "s_tir"
else native
)
group[gv] = func
diff --git a/src/relax/backend/vm/vm_shape_lower.cc
b/src/relax/backend/vm/vm_shape_lower.cc
index 0be9aa8f72..977c3739d1 100644
--- a/src/relax/backend/vm/vm_shape_lower.cc
+++ b/src/relax/backend/vm/vm_shape_lower.cc
@@ -750,7 +750,6 @@ class VMShapeLowerMutator
// the shape_func to indicate that this is a host function
// This could require us to attach target to the relax function here.
tirx::Function shape_func(params, body, ret_type);
- shape_func = WithAttr(std::move(shape_func), tvm::attr::kSTir.name, true);
shape_func =
WithAttr(std::move(shape_func), tvm::attr::kScriptNamespace.name,
ffi::String("s_tir"));
if
(!shape_func->attrs.GetAttr<tvm::Target>(tvm::attr::kTarget.name).has_value()) {
diff --git a/src/relax/transform/compute_prim_value.cc
b/src/relax/transform/compute_prim_value.cc
index 18b4d5733c..1bb1edb000 100644
--- a/src/relax/transform/compute_prim_value.cc
+++ b/src/relax/transform/compute_prim_value.cc
@@ -118,7 +118,6 @@ class PrimExprComputeInjector : public ExprMutator {
tirx::Function func(param_vars, tvm::SeqStmt(body), ret_ty,
DictAttrs({{tvm::tirx::attr::kIsHostFunc.name, true},
- {tvm::attr::kSTir.name, true},
{tvm::attr::kScriptNamespace.name,
ffi::String("s_tir")}}));
func = tirx::RenewDef(func);
diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc
index bf4daa52ec..8d75beece6 100644
--- a/src/relax/transform/fuse_tir.cc
+++ b/src/relax/transform/fuse_tir.cc
@@ -866,7 +866,6 @@ class FusedTIRConstructor : public ExprVisitor {
tirx::Function ConstructFunc() {
ffi::Map<ffi::String, Any> attr_map;
attr_map.Set(tvm::tirx::attr::kNoAlias.name, true);
- attr_map.Set(tvm::attr::kSTir.name, true);
attr_map.Set(tvm::attr::kScriptNamespace.name, ffi::String("s_tir"));
auto subst =
ffi::make_object<tirx::FuseTIRTensorSubstitutor>(func_info_.tensor_subst_map,
func_info_.symbolic_var_remap);
diff --git a/src/relax/transform/run_codegen.cc
b/src/relax/transform/run_codegen.cc
index b6b5b7b9e6..f425768157 100644
--- a/src/relax/transform/run_codegen.cc
+++ b/src/relax/transform/run_codegen.cc
@@ -100,7 +100,8 @@ class CodeGenRunner : ExprMutator {
TVM_FFI_ICHECK(!constants.count(name)) << "More than one constant with
the name " << name;
constants.Set(name, constant->value.cast<ffi::Tensor>());
}
- out_mod = WithAttr(out_mod, tvm::attr::kConstNameToConstant.name,
std::move(constants));
+ out_mod =
+ WithAttr(out_mod, tvm::relax::attr::kConstNameToConstant.name,
std::move(constants));
}
// TODO(@tvm-team): Implicit pass dependency. Revisit when we have a
better way to handle this.
diff --git a/src/s_tir/script/ir_builder/frame.cc
b/src/s_tir/script/ir_builder/frame.cc
index c0da42aaf5..4ad1cf37a0 100644
--- a/src/s_tir/script/ir_builder/frame.cc
+++ b/src/s_tir/script/ir_builder/frame.cc
@@ -33,20 +33,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
FunctionFrameObj::RegisterReflection();
SBlockFrameObj::RegisterReflection();
BlockInitFrameObj::RegisterReflection();
- tirx::FunctionFrameObj::RegisterAttrValidator(
- tvm::attr::kSTir.name, [](const tirx::FunctionFrameObj* frame, const
ffi::Any& value) {
- TVM_FFI_CHECK(frame->IsInstance<FunctionFrameObj>(), ValueError)
- << "The s_tir attribute requires Ts.function";
- bool enabled = false;
- if (auto flag = value.as<bool>()) {
- enabled = *flag;
- } else if (auto flag = value.as<int64_t>()) {
- enabled = *flag != 0;
- } else if (auto flag = value.as<tvm::IntImm>()) {
- enabled = !tvm::prim::IsZero(*flag);
- }
- TVM_FFI_CHECK(enabled, ValueError) << "Ts.function cannot disable the
s_tir attribute";
- });
}
tvm::tirx::Function FunctionFrameObj::FinalizeFunction(tvm::tirx::Function
func) {
@@ -54,7 +40,6 @@ tvm::tirx::Function
FunctionFrameObj::FinalizeFunction(tvm::tirx::Function func)
TVM_FFI_CHECK(!is_declaration || root_alloc_tensors.empty(), ValueError)
<< "A function declaration cannot allocate tensors";
if (!is_declaration) {
- func = WithAttr(std::move(func), tvm::attr::kSTir.name, true);
func =
tvm::s_tir::ScriptComplete(std::move(func), root_alloc_tensors,
root_allocated_addresses);
}
diff --git a/src/te/operation/create_function.cc
b/src/te/operation/create_function.cc
index db722e9b1f..732a331bfb 100644
--- a/src/te/operation/create_function.cc
+++ b/src/te/operation/create_function.cc
@@ -916,7 +916,7 @@ Function GenerateAndCompleteFunction(const
ffi::Array<te::Tensor>& arg_list,
/*ret_type=*/VoidType()),
{{tvm::attr::kGlobalSymbol.name,
ffi::String("main")},
{tvm::tirx::attr::kNoAlias.name, true},
- {tvm::attr::kSTir.name, true}});
+ {tvm::attr::kScriptNamespace.name,
ffi::String("s_tir")}});
const auto fcomplete =
tvm::ffi::Function::GetGlobal("s_tir.script.Complete");
TVM_FFI_ICHECK(fcomplete.has_value());
func = (*fcomplete)(std::move(func), info->root_alloc).cast<Function>();
@@ -988,7 +988,7 @@ Function GenerateAndCompleteFunction(const
ffi::Array<ffi::ObjectRef>& arg_tir_v
/*ret_type=*/VoidType()),
{{tvm::attr::kGlobalSymbol.name,
ffi::String("main")},
{tvm::tirx::attr::kNoAlias.name, true},
- {tvm::attr::kSTir.name, true}});
+ {tvm::attr::kScriptNamespace.name,
ffi::String("s_tir")}});
const auto fcomplete =
tvm::ffi::Function::GetGlobal("s_tir.script.Complete");
TVM_FFI_ICHECK(fcomplete.has_value());
func = (*fcomplete)(std::move(func), info->root_alloc).cast<Function>();
diff --git a/tests/python/relax/script/test_relax_script_printer.py
b/tests/python/relax/script/test_relax_script_printer.py
index 2862600443..ded8bb4be2 100644
--- a/tests/python/relax/script/test_relax_script_printer.py
+++ b/tests/python/relax/script/test_relax_script_printer.py
@@ -195,7 +195,7 @@ class Module:
T.module_global_infos({"mesh": [R.device_mesh((2, 2), T.Range(T.int32(0),
T.int32(4))), R.device_mesh((1,), T.Range(T.int32(4), T.int32(5)))]})
@Ts.function
def tir_func(x: T.Tensor((T.int64(128), T.int64(128)), T.float32), y:
T.Tensor((T.int64(128), T.int64(128)), T.float32)):
- T.func_attr({"s_tir": True, "script.namespace": "s_tir",
"tirx.noalias": True})
+ T.func_attr({"script.namespace": "s_tir", "tirx.noalias": True})
with Ts.sblock("root"):
Ts.reads()
Ts.writes()
@@ -970,7 +970,7 @@ from __future__ import annotations
class Module:
@Ts.function
def tir_func(x: T.Tensor((T.int64(128),), T.float32), y:
T.Tensor((T.int64(128),), T.float32)):
- T.func_attr({"s_tir": True, "script.namespace": "s_tir"})
+ T.func_attr({"script.namespace": "s_tir"})
@R.function
def foo(x: R.Tensor((128,), dtype="float32")) -> R.Tensor((128,),
dtype="float32"):
@@ -994,7 +994,7 @@ from __future__ import annotations
class Module:
@Ts.function
def tir_func(x: T.Tensor((T.int64(128),), T.float32), y:
T.Tensor((T.int64(128),), T.float32)):
- T.func_attr({"s_tir": True, "script.namespace": "s_tir"})
+ T.func_attr({"script.namespace": "s_tir"})
@R.function
def foo(x: R.Tensor((128,), dtype="float32")) -> R.Tensor((128,),
dtype="float32"):
diff --git a/tests/python/s_tir/base/test_compilation_pipeline.py
b/tests/python/s_tir/base/test_compilation_pipeline.py
index d076fe110b..a23c7ee503 100644
--- a/tests/python/s_tir/base/test_compilation_pipeline.py
+++ b/tests/python/s_tir/base/test_compilation_pipeline.py
@@ -41,6 +41,8 @@ def test_default_pipeline_selects_dialect(mixed):
for i in range(16):
B[i] = A[i] + T.float32(2)
+ assert "s_tir" not in scheduled.attrs
+ assert scheduled.attrs["script.namespace"] == "s_tir"
functions = {"scheduled": scheduled}
if mixed:
functions["native"] = native
diff --git a/tests/python/s_tir/script/test_s_tir_script_ir_builder.py
b/tests/python/s_tir/script/test_s_tir_script_ir_builder.py
index c069eab951..82e179e1d3 100644
--- a/tests/python/s_tir/script/test_s_tir_script_ir_builder.py
+++ b/tests/python/s_tir/script/test_s_tir_script_ir_builder.py
@@ -42,7 +42,7 @@ def test_ir_builder_tir_function_base():
params=[],
body=tvm_tirx.ir.Evaluate(0),
ret_type=None,
- attrs=tvm_tirx.ir.make_node("ir.DictAttrs", **{"s_tir": True,
"script.namespace": "s_tir"}),
+ attrs=tvm_tirx.ir.make_node("ir.DictAttrs", **{"script.namespace":
"s_tir"}),
)
# Check if the generated ir is expected
@@ -80,7 +80,7 @@ def test_ir_builder_tir_function_complete():
body=tvm_tirx.ir.Evaluate(0),
ret_type=tvm_tirx.ir.PrimType("int64"),
attrs=tvm_tirx.ir.make_node(
- "ir.DictAttrs", **{"key": "value", "s_tir": True,
"script.namespace": "s_tir"}
+ "ir.DictAttrs", **{"key": "value", "script.namespace": "s_tir"}
),
)
diff --git a/tests/python/s_tir/script/test_s_tir_script_namespace.py
b/tests/python/s_tir/script/test_s_tir_script_namespace.py
index 482383c835..764357be01 100644
--- a/tests/python/s_tir/script/test_s_tir_script_namespace.py
+++ b/tests/python/s_tir/script/test_s_tir_script_namespace.py
@@ -44,11 +44,11 @@ def test_shared_operations_and_aliases():
v = Ts.axis.spatial(A.shape[0], i)
A[v] = S.float32(1)
- assert shared.attrs["s_tir"]
+ assert shared.attrs["script.namespace"] == "s_tir"
assert shared.params[0].ty.layout is None
assert Ts.Tensor is T.Tensor
assert Ts.serial is T.serial
- assert Ts.bind is T.bind
+ assert Ts.Bind is T.Bind
assert Ts.tile is T.tile
tvm_tirx.ir.assert_structural_equal(
shared,
@@ -74,8 +74,8 @@ def test_mixed_module_roundtrip():
for i in T.serial(4):
A[i] = 2.0
- assert Mixed["scheduled"].attrs["s_tir"]
- assert not Mixed["direct"].attrs.get("s_tir", False)
+ assert Mixed["scheduled"].attrs["script.namespace"] == "s_tir"
+ assert Mixed["direct"].attrs.get("script.namespace") != "s_tir"
script = Mixed.script()
assert "@Ts.function" in script
assert "@T.function" in script
@@ -172,7 +172,7 @@ def test_s_tir_options_and_helpers():
def scheduled(A: Ts.Tensor((4,), "float32")):
fill(A)
- assert scheduled.attrs["s_tir"]
+ assert scheduled.attrs["script.namespace"] == "s_tir"
assert "global_symbol" not in scheduled.attrs
tvm_tirx.ir.assert_structural_equal(
scheduled,
@@ -223,27 +223,5 @@ def test_import_order(first):
)
[email protected]("value", [True, False])
-def test_tirx_cannot_change_dialect_with_attribute(value):
- with pytest.raises(ValueError, match="Ts.function"):
- tvm_tirx.script.from_source(
- "@T.function(check_well_formed=False)\n"
- f"def main():\n T.func_attr({{'s_tir': {value}}})\n
T.evaluate(0)\n",
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx,
"Ts": tvm.s_tir.script},
- )
-
-
[email protected]("value", [True, False])
-def test_tirx_cannot_change_dialect_from_exit_callback(value):
- from tvm_tirx.script.ir_builder import IRBuilder
- from tvm_tirx.tirx.script import ir_builder as builder
-
- with pytest.raises(ValueError, match="Ts.function"):
- with IRBuilder():
- with builder.function() as frame:
- frame.add_callback(lambda: builder.func_attr({"s_tir": value}))
- builder.evaluate(0)
-
-
if __name__ == "__main__":
tvm_tirx.testing.main()
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 ab1a0d8cf2..9f2886a1c6 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
@@ -142,7 +142,7 @@ def test_remap():
@Ts.function
def main():
- T.func_attr({"s_tir": True, "script.namespace": "s_tir"})
+ T.func_attr({"script.namespace": "s_tir"})
with Ts.sblock("root"):
Ts.reads()
Ts.writes()
@@ -190,7 +190,7 @@ def test_root_block():
@Ts.function
def main():
- T.func_attr({"s_tir": True, "script.namespace": "s_tir"})
+ T.func_attr({"script.namespace": "s_tir"})
with Ts.sblock("root"):
Ts.reads()
Ts.writes()
diff --git a/tests/python/s_tir/script/test_script_constructor_contracts.py
b/tests/python/s_tir/script/test_script_constructor_contracts.py
index 94be4cd150..daf986a753 100644
--- a/tests/python/s_tir/script/test_script_constructor_contracts.py
+++ b/tests/python/s_tir/script/test_script_constructor_contracts.py
@@ -42,7 +42,7 @@ def test_cast_script_positional_and_keyword_type():
S.evaluate(S.cast(x, ty="float32"))
assert S.cast is T.cast
- assert positional.attrs["s_tir"]
+ assert "s_tir" not in positional.attrs
assert positional.attrs["script.namespace"] == "s_tir"
ir.assert_structural_equal(positional, keyword)
cast_expr = positional.body[0].value
diff --git a/tirx/include/tvm/backend/trn/attr.h
b/tirx/include/tvm/backend/trn/attr.h
new file mode 100644
index 0000000000..791e2669b4
--- /dev/null
+++ b/tirx/include/tvm/backend/trn/attr.h
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file tvm/backend/trn/attr.h
+ * \brief Trainium backend attribute keys.
+ */
+#ifndef TVM_BACKEND_TRN_ATTR_H_
+#define TVM_BACKEND_TRN_ATTR_H_
+
+#include <tvm/ffi/reflection/accessor.h>
+
+#include <cstdint>
+
+namespace tvm {
+namespace backend {
+namespace trn {
+namespace attr {
+/*! \brief Number of input arguments preceding the mutable output arguments. */
+inline constexpr ffi::reflection::AttrKey<int64_t>
kNumInputArgs{"trn.num_input_args"};
+} // namespace attr
+} // namespace trn
+} // namespace backend
+} // namespace tvm
+
+#endif // TVM_BACKEND_TRN_ATTR_H_
diff --git a/tirx/include/tvm/ir/function.h b/tirx/include/tvm/ir/function.h
index 0511324526..95992db20b 100644
--- a/tirx/include/tvm/ir/function.h
+++ b/tirx/include/tvm/ir/function.h
@@ -252,31 +252,13 @@ inline constexpr ffi::reflection::AttrKey<Target>
kTarget{"target"};
inline constexpr ffi::reflection::AttrKey<ffi::String>
kGlobalSymbol{"global_symbol"};
/*!
- * \brief The function uses s_tir (apache-derived TIR) semantics:
- * parser fills layout=None, ScriptComplete wraps body in a root SBlock,
- * independent of the script construction namespace.
- * Default (attr absent or False) is tirx semantics.
- *
- * Type: IntImm (bool dtype)
- */
-inline constexpr ffi::reflection::AttrKey<bool> kSTir{"s_tir"};
-
-/*!
- * \brief Script construction namespace used to select the function decorator.
- * Absent metadata selects the function kind's default namespace. This does not
- * change the function's semantic dialect.
+ * \brief Script namespace used to select construction, printing, and dialect
lowering.
+ * Absent metadata selects the function kind's default namespace.
*
* Type: String
*/
inline constexpr ffi::reflection::AttrKey<ffi::String>
kScriptNamespace{"script.namespace"};
-/*!
- * \brief Number of inputs of the Function
- *
- * Type: Int
- */
-inline constexpr ffi::reflection::AttrKey<int64_t> kNumInputs{"num_inputs"};
-
} // namespace attr
inline LinkageType BaseFuncNode::GetLinkageType() const {
diff --git a/tirx/include/tvm/ir/module.h b/tirx/include/tvm/ir/module.h
index c0a157fcd7..9a22f91a46 100644
--- a/tirx/include/tvm/ir/module.h
+++ b/tirx/include/tvm/ir/module.h
@@ -353,19 +353,6 @@ inline constexpr
ffi::reflection::AttrKey<ffi::Array<ffi::Module>> kExternalMods
*/
inline constexpr ffi::reflection::AttrKey<ffi::String>
kSystemLibPrefix{"system_lib_prefix"};
-/*!
- * \brief All the named Tensors accumulated during compilation by external
codegen.
- * Generally the associated runtime::Module will indicate it requires bindings
for these names,
- * and during module initialization these bindings will be recovered from a
ConstLoaderModule.
- *
- * Type: ffi::Map<ffi::String, Tensor>
- */
-inline constexpr ffi::reflection::AttrKey<ffi::Map<ffi::String, ffi::Tensor>>
kConstNameToConstant{
- "const_name_to_constant"};
-
-inline constexpr ffi::reflection::AttrKey<ffi::Map<GlobalVar, ffi::String>>
kDeviceContexts{
- "device_contexts"};
-
} // namespace attr
} // namespace tvm
#endif // TVM_IR_MODULE_H_
diff --git a/tirx/python/tvm_tirx/backend/trn/pipeline.py
b/tirx/python/tvm_tirx/backend/trn/pipeline.py
index afccc76832..146b451ff8 100644
--- a/tirx/python/tvm_tirx/backend/trn/pipeline.py
+++ b/tirx/python/tvm_tirx/backend/trn/pipeline.py
@@ -56,6 +56,7 @@ def trn_pipeline(*, target=None):
tirx.transform.RemoveNoOp(),
tirx.transform.AnnotateEntryFunc(),
tirx.transform.SplitHostDevice(),
+ trn_transform.AnnotateKernelArguments(),
tirx.transform.MakePackedAPI(),
]
return tvm_tirx.ir.transform.Sequential(passes)(mod)
diff --git a/tirx/python/tvm_tirx/backend/trn/transform/__init__.py
b/tirx/python/tvm_tirx/backend/trn/transform/__init__.py
index 3f311cc17f..d0df43a06c 100644
--- a/tirx/python/tvm_tirx/backend/trn/transform/__init__.py
+++ b/tirx/python/tvm_tirx/backend/trn/transform/__init__.py
@@ -23,6 +23,7 @@ import tvm_tirx
from tvm_tirx import tirx
_LAZY_TRANSFORMS = {
+ "AnnotateKernelArguments": ".annotate_kernel_arguments",
"TrnNaiveAllocator": ".naive_allocator",
"TrnPrivateTensorAlloc": ".private_tensor_alloc",
}
@@ -90,4 +91,10 @@ def __getattr__(name):
return getattr(import_module(target, __name__), name)
-__all__ = ["LowerTIRx", "LowerTrainiumLayout", "TrnNaiveAllocator",
"TrnPrivateTensorAlloc"]
+__all__ = [
+ "AnnotateKernelArguments",
+ "LowerTIRx",
+ "LowerTrainiumLayout",
+ "TrnNaiveAllocator",
+ "TrnPrivateTensorAlloc",
+]
diff --git
a/tirx/python/tvm_tirx/backend/trn/transform/annotate_kernel_arguments.py
b/tirx/python/tvm_tirx/backend/trn/transform/annotate_kernel_arguments.py
new file mode 100644
index 0000000000..374aeb87e1
--- /dev/null
+++ b/tirx/python/tvm_tirx/backend/trn/transform/annotate_kernel_arguments.py
@@ -0,0 +1,137 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Apply the Trainium input/output contract after host/device splitting."""
+
+import tvm_ffi
+
+import tvm_tirx
+from tvm_tirx import tirx
+from tvm_tirx.ir import Bind, Call, GlobalVar, Op, Var
+from tvm_tirx.ir.transform import module_pass
+
+
+@module_pass(opt_level=0, name="AnnotateKernelArguments")
+class AnnotateKernelArguments:
+ """Order kernel inputs before outputs and annotate their per-kernel input
count.
+
+ Run after ``SplitHostDevice`` and before ``MakePackedAPI``. The source
+ ``trn.num_input_args`` contract remains on each host function until its
+ kernel calls have been matched to the surviving host parameters and
captures.
+ Functions without extracted kernels retain their own contract.
+ """
+
+ def transform_module(self, mod, _ctx):
+ attr = "trn.num_input_args"
+ kernels = {
+ gv.name_hint: gv
+ for gv, func in mod.functions.items()
+ if isinstance(func, tirx.Function)
+ and func.attrs.get("target") is not None
+ and func.attrs["target"].kind.name == "trn"
+ and func.attrs.get("tirx.is_global_func", False)
+ and attr not in func.attrs
+ }
+ updates = {}
+ for gv, host in mod.functions.items():
+ if not isinstance(host, tirx.Function) or attr not in host.attrs:
+ continue
+ num_inputs = host.attrs[attr]
+ if not isinstance(num_inputs, int) or not 0 <= num_inputs <=
len(host.params):
+ raise ValueError("trn.num_input_args must be between zero and
the parameter count")
+ param_order = {param: i for i, param in enumerate(host.params)}
+ has_kernel = False
+ bindings = {}
+ tvm_ffi.structural_walk(
+ host.body, (Bind, lambda bind: bindings.update({bind.var:
bind.value}))
+ )
+
+ def reorder(call):
+ nonlocal has_kernel
+ offset = 0
+ if isinstance(call.op, GlobalVar):
+ name = call.op.name_hint
+ elif isinstance(call.op, Op) and call.op.name ==
"tirx.call_ffi_kernel":
+ name = call.args[0].value
+ offset = 1
+ else:
+ return call
+ if name not in kernels:
+ return call
+ has_kernel = True
+ kernel_gv = kernels[name]
+ kernel = mod[kernel_gv]
+ args = list(call.args[offset : offset + len(kernel.params)])
+
+ def source_position(arg):
+ # A device region may capture a tensor alias declared on
the host.
+ # Follow its pointer projection back to the source
parameter.
+ while True:
+ if isinstance(arg, Var):
+ if arg in param_order:
+ return param_order[arg]
+ if arg not in bindings:
+ return None
+ arg = bindings[arg]
+ elif (
+ isinstance(arg, Call)
+ and arg.args
+ and isinstance(arg.op, Op)
+ and arg.op.name in ("tirx.tensor_data_ptr",
"tirx.decl_tensor")
+ ):
+ arg = arg.args[0]
+ else:
+ return None
+
+ positions = [source_position(arg) for arg in args]
+
+ # Captures are additional inputs. Only surviving host outputs
are mutable
+ # outputs; unused host inputs must not inflate the kernel's
input count.
+ def sort_key(i):
+ position = positions[i]
+ if position is None:
+ return (1, i)
+ return (2 if position >= num_inputs else 0, position)
+
+ order = sorted(range(len(args)), key=sort_key)
+ input_count = sum(pos is None or pos < num_inputs for pos in
positions)
+ updates[kernel_gv] = tirx.Function(
+ [kernel.params[i] for i in order],
+ kernel.body,
+ kernel.ret_type,
+ kernel.attrs,
+ kernel.loc,
+ ).with_attr(attr, input_count)
+ return Call(
+ call.op,
+ [
+ *call.args[:offset],
+ *(args[i] for i in order),
+ *call.args[offset + len(args) :],
+ ],
+ attrs=call.attrs,
+ ty_args=call.ty_args,
+ loc=call.loc,
+ ty=call.ty,
+ )
+
+ updated = host.with_body(
+ tvm_ffi.structural_map(host.body, (Call, reorder),
order="post")
+ )
+ if has_kernel:
+ updates[gv] = updated.without_attr(attr)
+ mod.update(tvm_tirx.IRModule(updates))
+ return mod
diff --git a/tirx/src/backend/cuda/codegen/codegen_cuda_host.cc
b/tirx/src/backend/cuda/codegen/codegen_cuda_host.cc
index 867ec37123..f3a5ce0b1b 100644
--- a/tirx/src/backend/cuda/codegen/codegen_cuda_host.cc
+++ b/tirx/src/backend/cuda/codegen/codegen_cuda_host.cc
@@ -50,7 +50,7 @@ class CodeGenCUDAHost : public CodeGenCHost {
using CodeGenCHost::Dispatch_;
void Init(Target target) {
- CodeGenCHost::Init(false, true, true, target->str(), {});
+ CodeGenCHost::Init(false, true, true, target->str());
// Retain C-host wrapper initialization, with an FFI/CUDA-only prelude.
decl_stream.str("");
decl_stream.clear();
diff --git a/tirx/src/backend/trn/codegen/codegen_trn.cc
b/tirx/src/backend/trn/codegen/codegen_trn.cc
index 260d9f3c82..48a49c2d34 100644
--- a/tirx/src/backend/trn/codegen/codegen_trn.cc
+++ b/tirx/src/backend/trn/codegen/codegen_trn.cc
@@ -22,6 +22,7 @@
*/
#include "codegen_trn.h"
+#include <tvm/backend/trn/attr.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ir/function.h>
#include <tvm/ir/prim/expr.h>
@@ -114,8 +115,10 @@ void CodeGenTrainium::AddFunction(const GlobalVar& gvar,
const Function& func) {
this->stream << "def " << static_cast<std::string>(global_symbol.value()) <<
"(";
// Buffer arguments
- auto num_inputs = func->GetAttr(tvm::attr::kNumInputs);
- TVM_FFI_ICHECK(num_inputs.has_value());
+ auto num_inputs = func->GetAttr(tvm::backend::trn::attr::kNumInputArgs);
+ TVM_FFI_CHECK(num_inputs.has_value(), ValueError) << "Trainium requires
trn.num_input_args";
+ TVM_FFI_CHECK_GE(num_inputs.value(), 0, ValueError);
+ TVM_FFI_CHECK_LE(num_inputs.value(),
static_cast<int64_t>(func->params.size()), ValueError);
std::vector<std::string> output_vids;
size_t num_buffer = 0;
for (size_t i = 0; i < func->params.size(); ++i, ++num_buffer) {
diff --git a/tirx/src/codegen/c/codegen_c_host.cc
b/tirx/src/codegen/c/codegen_c_host.cc
index cc37cb66c8..152000c639 100644
--- a/tirx/src/codegen/c/codegen_c_host.cc
+++ b/tirx/src/codegen/c/codegen_c_host.cc
@@ -32,7 +32,6 @@
#include <algorithm>
#include <string>
-#include <unordered_set>
#include <utility>
#include <vector>
@@ -44,7 +43,7 @@ CodeGenCHost::CodeGenCHost() {
}
void CodeGenCHost::Init(bool output_ssa, bool emit_asserts, bool
emit_fwd_func_decl,
- std::string target_str, const
std::unordered_set<std::string>& devices) {
+ std::string target_str) {
emit_asserts_ = emit_asserts;
emit_fwd_func_decl_ = emit_fwd_func_decl;
declared_globals_.clear();
@@ -378,17 +377,8 @@ ffi::Module BuildCHost(IRModule mod, Target target) {
bool emit_asserts = true;
bool emit_fwd_func_decl = true;
- std::unordered_set<std::string> devices;
- if (mod->GetAttr(tvm::attr::kDeviceContexts) != nullptr) {
- ffi::Map<GlobalVar, ffi::String> device_contexts =
- mod->GetAttr(tvm::attr::kDeviceContexts).value();
- for (auto const& context : device_contexts) {
- devices.insert(context.second.data());
- }
- }
-
CodeGenCHost cg;
- cg.Init(output_ssa, emit_asserts, emit_fwd_func_decl, target->str(),
devices);
+ cg.Init(output_ssa, emit_asserts, emit_fwd_func_decl, target->str());
cg.SetConstantsByteAlignment(target->GetAttr<int64_t>("constants-byte-alignment").value_or(16));
std::vector<std::pair<GlobalVar, Function>> funcs;
diff --git a/tirx/src/codegen/c/codegen_c_host.h
b/tirx/src/codegen/c/codegen_c_host.h
index e8fd3bfde3..5421ff3def 100644
--- a/tirx/src/codegen/c/codegen_c_host.h
+++ b/tirx/src/codegen/c/codegen_c_host.h
@@ -28,7 +28,6 @@
#include <string>
#include <unordered_map>
-#include <unordered_set>
#include <utility>
#include <vector>
@@ -41,8 +40,7 @@ namespace codegen {
class CodeGenCHost : public CodeGenC {
public:
CodeGenCHost();
- void Init(bool output_ssa, bool emit_asserts, bool emit_fwd_func_decl,
std::string target_str,
- const std::unordered_set<std::string>& devices);
+ void Init(bool output_ssa, bool emit_asserts, bool emit_fwd_func_decl,
std::string target_str);
void InitGlobalContext();
diff --git a/tirx/src/tirx/transform/split_host_device.cc
b/tirx/src/tirx/transform/split_host_device.cc
index bb6875a12d..b75042b17a 100644
--- a/tirx/src/tirx/transform/split_host_device.cc
+++ b/tirx/src/tirx/transform/split_host_device.cc
@@ -157,27 +157,17 @@ class HostDeviceSplitter : public StmtExprMutator {
// Sort first by variable type, then by variable name
std::vector<Var> params{undefined.begin(), undefined.end()};
- if (device_target->kind->name != "trn") {
- std::sort(params.begin(), params.end(), [](const Var& a, const Var& b)
{
- auto sort_key = [](const Var& var) {
- bool is_handle =
- var->ty.as<PtrTypeNode>() != nullptr ||
var->ty.as<TensorTypeNode>() != nullptr;
- return std::tuple{
- !is_handle,
- var->name,
- };
+ std::sort(params.begin(), params.end(), [](const Var& a, const Var& b) {
+ auto sort_key = [](const Var& var) {
+ bool is_handle =
+ var->ty.as<PtrTypeNode>() != nullptr ||
var->ty.as<TensorTypeNode>() != nullptr;
+ return std::tuple{
+ !is_handle,
+ var->name,
};
- return sort_key(a) < sort_key(b);
- });
- } else {
- std::unordered_map<Var, int, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
param_order;
- for (size_t i = 0; i < cur_func_->params.size(); ++i) {
-
param_order[cur_func_->params[i].as_or_throw<tvm::tirx::TensorVar>()] = i;
- }
- // sort by original order
- std::sort(params.begin(), params.end(),
- [&](const Var& a, const Var& b) { return param_order[a] <
param_order[b]; });
- }
+ };
+ return sort_key(a) < sort_key(b);
+ });
return {params, tensors};
}();
@@ -246,19 +236,11 @@ class HostDeviceSplitter : public StmtExprMutator {
device_func = WithAttr(std::move(device_func),
tvm::attr::kScriptNamespace.name,
script_namespace.value());
}
- bool is_stir = cur_func_->attrs->dict.count(tvm::attr::kSTir.name);
- if (is_stir) {
- device_func = WithAttr(std::move(device_func), tvm::attr::kSTir.name,
true);
- }
if (auto launch_params = cur_func_->GetAttr<ffi::Array<ffi::String>>(
tvm::tirx::attr::kKernelLaunchParams.name)) {
device_func = WithAttr(std::move(device_func),
tvm::tirx::attr::kKernelLaunchParams.name,
launch_params.value());
}
- auto num_inputs = cur_func_->GetAttr(tvm::attr::kNumInputs);
- if (num_inputs.has_value()) {
- device_func = WithAttr(std::move(device_func),
tvm::attr::kNumInputs.name, num_inputs);
- }
GlobalVar kernel_symbol_global = var_supply_();
if
(region->attrs->dict.count(tvm::backend::cuda::attr::kLaunchFields.name)) {
Stmt launch = MakeCudaKernelLaunch(kernel_symbol_global, &device_func,
call_args, region);
diff --git a/tirx/tests/python/tirx/codegen/test_nki.py
b/tirx/tests/python/tirx/codegen/test_nki.py
index e6493d7da1..bebc8121dd 100644
--- a/tirx/tests/python/tirx/codegen/test_nki.py
+++ b/tirx/tests/python/tirx/codegen/test_nki.py
@@ -40,7 +40,7 @@ def test_nki_add_1():
# fmt: off
@T.function
def func(A: T.Tensor((128, 512), "float32"), B: T.Tensor((128, 512),
"float32")):
- T.func_attr({"num_inputs": 1})
+ T.func_attr({"trn.num_input_args": 1})
T.device_entry()
A_sbuf = T.alloc_tensor(ty_args=[T.Tensor((128, 512), "float32",
scope="trn.sbuf")])
B_sbuf = T.alloc_tensor(ty_args=[T.Tensor((128, 512), "float32",
scope="trn.sbuf")])
@@ -94,7 +94,7 @@ def test_nki_add_2():
# fmt: off
@T.function
def func(A: T.Tensor((128, 2048), "float32"), B: T.Tensor((128, 2048),
"float32")):
- T.func_attr({"num_inputs": 1})
+ T.func_attr({"trn.num_input_args": 1})
T.device_entry()
A_sbuf = T.alloc_tensor(ty_args=[T.Tensor((128, 512), "float32",
scope="trn.sbuf")])
B_sbuf = T.alloc_tensor(ty_args=[T.Tensor((128, 512), "float32",
scope="trn.sbuf")])
@@ -174,7 +174,7 @@ def test_nki_matmul_1():
rhs: T.Tensor((K, N), "float16"),
result: T.Tensor((M, N), "float16"),
):
- T.func_attr({"num_inputs": 2})
+ T.func_attr({"trn.num_input_args": 2})
result_tiles = T.alloc_tensor(
ty_args=[
T.Tensor(
