This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx-standalone-root in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 229aec2065c7a79636c249577904eac98d10e984 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. --- 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 +- 12 files changed, 211 insertions(+), 83 deletions(-) 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(
