From 846f405c7e763358991107e602bfd16f2b637300 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Sun, 2 Aug 2026 08:13:04 +0000 Subject: [PATCH 1/7] [TIRx] Mechanically migrate buffer_map to BufferType params --- .../tvm/relax/distributed/axis_group_graph.h | 12 +-- include/tvm/script/ir_builder/base.h | 2 +- include/tvm/tirx/function.h | 80 ++++++------------- .../cuda/tile_primitive/copy_async/tma.py | 2 +- python/tvm/relax/analysis/analysis.py | 9 ++- python/tvm/s_tir/dlight/benchmark/extract.py | 4 +- python/tvm/te/operation.py | 4 +- python/tvm/tirx/function.py | 17 ++-- python/tvm/tirx/script/parser/parser.py | 32 +++++++- python/tvm/tirx/transform/transform.py | 7 +- src/arith/domain_touched.cc | 6 +- src/arith/ir_mutator_with_analyzer.cc | 12 ++- src/arith/ir_mutator_with_analyzer.h | 4 +- src/backend/trn/codegen/codegen_trn.cc | 2 +- .../trn/transform/lower_trainium_layout.cc | 13 ++- src/relax/analysis/layout_transformation.cc | 2 +- src/relax/analysis/tir_op_pattern_kind.cc | 4 +- .../backend/adreno/annotate_custom_storage.cc | 4 +- src/relax/backend/vm/codegen_vm_tir.cc | 2 +- src/relax/distributed/axis_group_graph.cc | 2 +- .../lower_global_view_to_local_view.cc | 26 +++--- src/relax/transform/dataflow_inplace.cc | 6 +- src/relax/transform/fuse_tir.cc | 43 +++++----- src/relax/transform/meta_schedule.cc | 1 - .../transform/rewrite_dataflow_reshape.cc | 2 +- .../transform/split_call_tir_by_pattern.cc | 38 +++------ .../transform/split_layout_rewrite_preproc.cc | 19 ++--- .../sblock_buffer_access_lca_detector.cc | 2 +- .../backend/adreno/inject_texture_alloc.cc | 2 +- src/s_tir/backend/adreno/texture_flatten.cc | 3 +- src/s_tir/meta_schedule/arg_info.cc | 2 +- .../feature_extractor/per_store_feature.cc | 2 +- .../meta_schedule/postproc/rewrite_layout.cc | 2 +- src/s_tir/schedule/analysis/analysis.cc | 2 +- .../schedule/primitive/blockize_tensorize.cc | 6 +- .../primitive/layout_transformation.cc | 4 +- src/s_tir/transform/inject_permuted_layout.cc | 5 +- .../transform/inject_software_pipeline.cc | 2 +- src/s_tir/transform/lower_thread_allreduce.cc | 3 +- .../plan_update_buffer_allocation_location.cc | 4 +- .../remove_weight_layout_rewrite_block.cc | 16 ++-- src/s_tir/transform/renew_defs.cc | 29 ++++--- src/target/llvm/codegen_llvm.cc | 4 +- src/te/operation/create_primfunc.cc | 18 ++--- src/tirx/analysis/verify_memory.cc | 2 +- src/tirx/analysis/verify_ssa.cc | 2 +- src/tirx/analysis/verify_well_formed.cc | 2 +- src/tirx/ir/data_type_rewriter.cc | 9 ++- src/tirx/ir/function.cc | 40 +++++++--- src/tirx/ir/script/script_complete.cc | 2 +- src/tirx/ir/specialize.cc | 67 ++++++++++++---- src/tirx/ir/tir_visitor_with_path.cc | 21 ++--- src/tirx/script/builder/frame.cc | 52 ++++++------ src/tirx/script/builder/ir.cc | 6 +- src/tirx/script/printer/buffer.cc | 10 +-- src/tirx/script/printer/function.cc | 57 ++++--------- src/tirx/script/printer/utils.h | 2 +- src/tirx/transform/flatten_buffer.cc | 13 ++- .../transform/force_narrow_index_to_i32.cc | 7 +- .../transform/inline_private_functions.cc | 7 +- src/tirx/transform/ir_utils.cc | 31 ++++--- src/tirx/transform/lower_tirx_cleanup.cc | 14 +++- src/tirx/transform/make_packed_api.cc | 6 +- src/tirx/transform/split_host_device.cc | 2 +- src/tirx/transform/stmt_simplify.cc | 2 +- src/tirx/transform/storage_rewrite.cc | 68 +++++++--------- src/tirx/transform/tvm_ffi_binder.cc | 11 ++- src/tirx/transform/tvm_ffi_binder.h | 2 + .../transform/unsupported_dtype_legalize.cc | 3 +- .../python/arith/test_arith_domain_touched.py | 4 +- .../python/relax/test_analysis_well_formed.py | 8 +- tests/python/relax/test_blockbuilder_core.py | 7 +- ...m_specialize_primfunc_based_on_callsite.py | 10 ++- .../test_s_tir_analysis_identify_memcpy.py | 22 ++--- .../analysis/test_sblock_access_region.py | 8 +- .../analysis/test_sblock_buffer_access_lca.py | 11 ++- .../s_tir/base/test_tir_te_extern_primfunc.py | 8 +- .../test_tir_schedule_decompose_padding.py | 4 +- .../test_tir_schedule_rolling_buffer.py | 4 +- .../test_tir_schedule_transform_layout.py | 8 +- tests/python/s_tir/test_s_tir_renew_defs.py | 6 +- .../test_tir_analysis_verify_well_formed.py | 8 +- tests/python/tirx-base/test_tir_nodes.py | 6 +- tests/python/tirx-base/test_tir_specialize.py | 13 ++- .../test_tir_transform_convert_ssa.py | 27 ++----- ...tir_transform_force_narrow_index_to_i32.py | 6 +- ...ir_transform_pointer_value_type_rewrite.py | 2 +- tests/python/tirx/test_jit.py | 34 ++++---- .../tirx/transform/test_stmt_functor.py | 2 +- .../tvmscript/test_tvmscript_complete.py | 6 +- .../test_tvmscript_ir_builder_tir.py | 23 ++---- ...test_tvmscript_printer_structural_equal.py | 10 +-- .../tvmscript/test_tvmscript_printer_tir.py | 58 +++++--------- 93 files changed, 594 insertions(+), 600 deletions(-) diff --git a/include/tvm/relax/distributed/axis_group_graph.h b/include/tvm/relax/distributed/axis_group_graph.h index e7cd828b1f16..5b47fe4f6ce4 100644 --- a/include/tvm/relax/distributed/axis_group_graph.h +++ b/include/tvm/relax/distributed/axis_group_graph.h @@ -71,14 +71,14 @@ class BufferAxisGraphExtractor : public StmtExprVisitor { BufferAxisGraphExtractor extractor; extractor(prim_func->body); ffi::Map inverse_buffer_map; - for (const auto& pr : prim_func->buffer_map) { - inverse_buffer_map.Set(pr.second, pr.first); - } std::vector> tir_var_axis_group_list; std::unordered_set visited; - for (const auto& pr : prim_func->buffer_map) { - Var param = pr.first; - BufferVar buffer = pr.second; + for (const Var& param : prim_func->params) { + if (!param->ty.as()) { + continue; + } + BufferVar buffer(param); + inverse_buffer_map.Set(buffer, param); for (int i = 0; i < static_cast(buffer->shape.size()); i++) { if (extractor.buffer_axis_graph_.count({buffer, i})) { std::vector buffer_axis_group; diff --git a/include/tvm/script/ir_builder/base.h b/include/tvm/script/ir_builder/base.h index 010a091ab2de..7e0d34a42f80 100644 --- a/include/tvm/script/ir_builder/base.h +++ b/include/tvm/script/ir_builder/base.h @@ -39,7 +39,7 @@ namespace ir_builder { * * \example * - * The `T::MatchBuffer` below adds an element in `PrimFuncNode::buffer_map`: + * The `T::MatchBuffer` below annotates a PrimFunc parameter with BufferType: * * \code {.cpp} * diff --git a/include/tvm/tirx/function.h b/include/tvm/tirx/function.h index 414e6c40b929..74ece4538d23 100644 --- a/include/tvm/tirx/function.h +++ b/include/tvm/tirx/function.h @@ -38,6 +38,25 @@ namespace tvm { namespace tirx { +/*! \brief Return a checked buffer view when a parameter carries BufferType. */ +inline ffi::Optional AsBufferVar(const Var& var) { + if (var->ty.as()) { + return BufferVar(var); + } + return std::nullopt; +} + +/*! \brief Derive a keyed view of buffer parameters from their BufferType annotations. */ +inline ffi::Map BufferParamMap(const ffi::Array& params) { + ffi::Map result; + for (const Var& param : params) { + if (auto buffer = AsBufferVar(param)) { + result.Set(param, buffer.value()); + } + } + return result; +} + /*! * \brief Primitive functions that contains TIR statements. * @@ -52,52 +71,6 @@ class PrimFuncNode : public BaseFuncNode { ffi::Array params; /*! \brief The return type of the function. */ Type ret_type = Type::Missing(); - /*! - * \brief Maps some parameters to specific buffer data structures. - * - * buffer_map provides a way to express data structure's field and shape - * constraints. The provided information is used in the program analysis - * and the code generation. - * - * - It defines the vars in the buffer (m, n) in the cases below when - * they appears in the buffer_map for the first time. - * - When a var appears multiple times, they translate into runtime - * assertion to check the field constraint. - * - * \code - * - * # The corresponding fields of f are as follows - * # - * # - f.params = [a, b] - * # - f.buffer_map = {a: A, b: B} - * # - A = decl_buffer(shape=[m, n]) - * # - B = decl_buffer(shape=[m, n]) - * - * def f(a, b): - * m, n = var(), var() - * A = bind_buffer(a, shape=[m, n]) - * B = bind_buffer(b, shape=[m, n]) - * # body - * - * \endcode - * - * buffer_map is a sugar to express: - * - Parameter unpacking: e.g. I can load a.shape[0] to get value of m - * - Constraint checking: a.shape[0] must equal b.shape[0] because they - * both corresponds to m. - - * While we could have express parameter unpacking and constraint using - * normal statements, making buffer_map as first class citizen of PrimFunc - * will make program analysis much easier. - * - * Prior to buffer flattening, which is performed FlattenBuffer for - * TIR-based schedules, these buffer objects are used directly in - * the body of the function. After buffer flattening, these buffer - * objects remain unflattened for use in argument validation, but - * all usage in the body of the function is done through a - * flattened alias of the buffer. - */ - ffi::Map buffer_map; /*! \brief The body of the function */ tirx::Stmt body; @@ -106,8 +79,6 @@ class PrimFuncNode : public BaseFuncNode { refl::ObjectDef() .def_ro("params", &PrimFuncNode::params, refl::AttachFieldFlag::SEqHashDefRecursive()) .def_ro("ret_type", &PrimFuncNode::ret_type) - .def_ro("buffer_map", &PrimFuncNode::buffer_map, - refl::AttachFieldFlag::SEqHashDefRecursive()) .def_ro("body", &PrimFuncNode::body); refl::TypeAttrDef() .def("__s_equal__", &PrimFuncNode::SEqual) @@ -122,7 +93,6 @@ class PrimFuncNode : public BaseFuncNode { return equal(attrs, other->attrs, false, "attrs") && equal(params, other->params, true, "params") && equal(ret_type, other->ret_type, false, "ret_type") && - equal(buffer_map, other->buffer_map, true, "buffer_map") && equal(body, other->body, false, "body"); } @@ -131,7 +101,6 @@ class PrimFuncNode : public BaseFuncNode { hash_value = hash(attrs, hash_value, false); hash_value = hash(params, hash_value, true); hash_value = hash(ret_type, hash_value, false); - hash_value = hash(buffer_map, hash_value, true); hash_value = hash(body, hash_value, false); return hash_value; } @@ -163,19 +132,18 @@ class PrimFunc : public BaseFunc { * * \param ret_type The return type of the function. * - * \param buffer_map The buffer map for parameter buffer unpacking. - * This contains buffer objects as they appear in the body of the - * PrimFunc. (e.g. a buffer of shape ``[1024]`` originally - * generated as a tensor of shape ``[32, 32]``) - * * \param attrs Additional function attributes. * * \param span The location of this object in the source code. */ TVM_DLL PrimFunc(ffi::Array params, Stmt body, Type ret_type = VoidType(), - ffi::Map buffer_map = ffi::Map(), DictAttrs attrs = DictAttrs(), Span span = Span()); + /*! \brief Compatibility constructor that folds legacy buffer bindings into params. */ + TVM_DLL PrimFunc(ffi::Array params, Stmt body, Type ret_type, + ffi::Map buffer_map, DictAttrs attrs = DictAttrs(), + Span span = Span()); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrimFunc, BaseFunc, PrimFuncNode); TVM_DEFINE_OBJECT_REF_COW_METHOD(PrimFuncNode); }; diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py b/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py index f5d678567c4d..0abba7675bbc 100644 --- a/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py +++ b/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py @@ -2208,7 +2208,7 @@ def impl(): if selector_bind is not None: body = tvm.tirx.SeqStmt([selector_bind, impl.body]) - impl = PrimFunc([], body, ret_type=None, buffer_map={}).with_attr("global_symbol", "impl") + impl = PrimFunc([], body, ret_type=None).with_attr("global_symbol", "impl") return impl diff --git a/python/tvm/relax/analysis/analysis.py b/python/tvm/relax/analysis/analysis.py index 987b7eedb344..7f0a6c2c3950 100644 --- a/python/tvm/relax/analysis/analysis.py +++ b/python/tvm/relax/analysis/analysis.py @@ -527,10 +527,11 @@ def check_well_formed(obj: IRModule | Function, check_ty: bool = True) -> bool: def _get_prim_func_default_dtype(func: PrimFunc): - """Detect default index dtype from function buffer map""" - for _, v in func.buffer_map.items(): - for value in v.shape: - return value.ty + """Detect default index dtype from BufferType-annotated parameters.""" + for param in func.params: + if tirx.is_buffer_var(param): + for value in param.shape: + return value.ty return "int64" diff --git a/python/tvm/s_tir/dlight/benchmark/extract.py b/python/tvm/s_tir/dlight/benchmark/extract.py index a1eb56361e7b..486252aedb6b 100644 --- a/python/tvm/s_tir/dlight/benchmark/extract.py +++ b/python/tvm/s_tir/dlight/benchmark/extract.py @@ -180,7 +180,9 @@ def extract_func_info_from_prim_func( func_args = [] dym_var = {} for param in func.params: - buffer = func.buffer_map[param] + if not tvm.tirx.is_buffer_var(param): + continue + buffer = param shape = [] for dim in buffer.shape: if isinstance(dim, tvm.tirx.IntImm): diff --git a/python/tvm/te/operation.py b/python/tvm/te/operation.py index 73c8d7f26301..f7d88f813105 100644 --- a/python/tvm/te/operation.py +++ b/python/tvm/te/operation.py @@ -389,9 +389,9 @@ def before_split(a: T.handle, b: T.handle) -> None: C = te.extern_primfunc([A, B], func) """ - # dt_access_map and primfunc.buffer_map are unordered, so use order from primfunc.params + # Preserve the function parameter order while selecting BufferType annotations. dt_access_map = tvm.arith._ffi_api.DomainTouchedAccessMap(primfunc) - ordered_buffers = [primfunc.buffer_map[param] for param in primfunc.params] + ordered_buffers = [param for param in primfunc.params if tvm.tirx.is_buffer_var(param)] in_buffers = [buf for buf in ordered_buffers if len(dt_access_map[buf][0])] out_buffers = [buf for buf in ordered_buffers if len(dt_access_map[buf][1])] assert in_buffers, "PrimFunc has no input buffers" diff --git a/python/tvm/tirx/function.py b/python/tvm/tirx/function.py index bbacc0472fa1..cae64bdda61b 100644 --- a/python/tvm/tirx/function.py +++ b/python/tvm/tirx/function.py @@ -50,9 +50,6 @@ class PrimFunc(BaseFunc, Scriptable): ret_type: tvm.ir.Type The return type annotation of the function. - buffer_map : Map[tvm.tirx.Var, tvm.tirx.Buffer] - The buffer binding map. - attrs: Optional[tvm.Attrs] Attributes of the function, can be None @@ -73,11 +70,9 @@ def __init__(self, params, body, ret_type=None, buffer_map=None, attrs=None, spa for x in params: x = tvm.runtime.convert(x) if not isinstance(x, Object) else x if is_buffer_var(x): - var = Var(x.name, ty="handle") - param_list.append(var) - buffer_map[var] = x - elif isinstance(x, Var): param_list.append(x) + elif isinstance(x, Var): + param_list.append(buffer_map.get(x, x)) else: raise TypeError("params can only contain Var or Buffer") @@ -89,7 +84,6 @@ def __init__(self, params, body, ret_type=None, buffer_map=None, attrs=None, spa param_list, body, ret_type, - buffer_map, attrs, span, ) # type: ignore @@ -113,10 +107,9 @@ def with_body(self, new_body, span=None): return PrimFunc( self.params, new_body, - self.ret_type, - self.buffer_map, - self.attrs, - span, + ret_type=self.ret_type, + attrs=self.attrs, + span=span, ) def specialize(self, param_map: Mapping[Var, Expr | Buffer]): diff --git a/python/tvm/tirx/script/parser/parser.py b/python/tvm/tirx/script/parser/parser.py index 3e62e8b0559e..d8a7445a9e1a 100644 --- a/python/tvm/tirx/script/parser/parser.py +++ b/python/tvm/tirx/script/parser/parser.py @@ -651,15 +651,39 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None: # - kwarg: arg | None # - defaults: list[expr] # - posonlyargs: list[arg] + # Buffer annotations may refer to scalar parameters that occur + # later in the signature. Predeclare any independently + # evaluable scalar/pointer parameters so those forward + # references resolve without changing ABI parameter order. + evaluated_annotations = {} for arg in node.args.args: if arg.annotation is None: - self.report_error(arg, "Type annotation required for function parameters.") + continue try: ann = self.eval_expr(arg.annotation) + if isinstance(ann, _OptionalAnnotation) or ann is _constexpr_sentinel: + evaluated_annotations[arg.arg] = ann + continue + if callable(ann) and not isinstance(ann, Expr): + ann = ann() except Exception: # pylint: disable=broad-except - ann = func_annotation.get(arg.arg, None) - if ann is None: - raise + continue + evaluated_annotations[arg.arg] = ann + if isinstance(ann, tvm.tirx.Var) and not is_buffer_var(ann): + self.var_table.add(arg.arg, ann) + + for arg in node.args.args: + if arg.annotation is None: + self.report_error(arg, "Type annotation required for function parameters.") + if arg.arg in evaluated_annotations: + ann = evaluated_annotations[arg.arg] + else: + try: + ann = self.eval_expr(arg.annotation) + except Exception: # pylint: disable=broad-except + ann = func_annotation.get(arg.arg, None) + if ann is None: + raise if isinstance(ann, _OptionalAnnotation): if not _is_jit_function(node): self.report_error( diff --git a/python/tvm/tirx/transform/transform.py b/python/tvm/tirx/transform/transform.py index 907324454916..66dc7f3b5da7 100644 --- a/python/tvm/tirx/transform/transform.py +++ b/python/tvm/tirx/transform/transform.py @@ -248,10 +248,9 @@ def ConvertSSA(): def MakePackedAPI(): """Transform the PrimFuncs in the module to a packed func API. - Prior to this pass, the PrimFunc may have Buffer arguments defined - in the `PrimFuncNode::buffer_map`. This pass consumes the - `buffer_map`, using it to generate arguments that implement - the packed based TVM FFI API. + Prior to this pass, the PrimFunc may have parameters annotated with + `BufferType`. This pass consumes those annotations to generate + arguments that implement the packed based TVM FFI API. For static shapes, the `BufferType::shape`, `BufferType::strides`, and `BufferType::elem_offset` fields are used to diff --git a/src/arith/domain_touched.cc b/src/arith/domain_touched.cc index 94a90be1098e..e26fb448da7c 100644 --- a/src/arith/domain_touched.cc +++ b/src/arith/domain_touched.cc @@ -141,9 +141,11 @@ Region DomainTouched(const Stmt& stmt, const BufferVar& buffer, bool consider_lo ffi::Map> DomainTouchedAccessMap(const PrimFunc& func) { auto buffer_access_map = BufferTouchedDomain(func->body).GetAccessedBufferRegions(); ffi::Map> ret; - auto& buffer_map = func->buffer_map; for (auto& var : func->params) { - auto& buffer = buffer_map[var]; + if (!var->ty.as()) { + continue; + } + BufferVar buffer(var); auto& access = buffer_access_map[buffer.get()]; ffi::Array> loads, stores, combined; for (std::vector& touch : std::get(access).set) { diff --git a/src/arith/ir_mutator_with_analyzer.cc b/src/arith/ir_mutator_with_analyzer.cc index 65f20e042943..260c053f755f 100644 --- a/src/arith/ir_mutator_with_analyzer.cc +++ b/src/arith/ir_mutator_with_analyzer.cc @@ -153,10 +153,14 @@ void EnterConstraintFacts(WithGroup* constraints, AnalyzerObj } // namespace -void IRMutatorWithAnalyzer::MarkBufferMapShapes(const tirx::PrimFunc& func) { - // Mark the all the symbolic buffer shape values in the buffer map as positive value. - for (auto kv : func->buffer_map) { - for (PrimExpr shape : kv.second->shape) { +void IRMutatorWithAnalyzer::MarkBufferParamShapes(const tirx::PrimFunc& func) { + // Mark all symbolic buffer-parameter shape values as positive. + for (const tirx::Var& param : func->params) { + if (!param->ty.as()) { + continue; + } + tirx::BufferVar buffer(param); + for (PrimExpr shape : buffer->shape) { analyzer_->MarkGlobalNonNegValue(shape); } } diff --git a/src/arith/ir_mutator_with_analyzer.h b/src/arith/ir_mutator_with_analyzer.h index 47b52da78c14..1d2f8ef1d633 100644 --- a/src/arith/ir_mutator_with_analyzer.h +++ b/src/arith/ir_mutator_with_analyzer.h @@ -69,12 +69,12 @@ class IRMutatorWithAnalyzer : public tirx::StmtExprMutator { protected: /*! - * \brief Mark the all the buffer shape values in the buffer map as positive value. + * \brief Mark all buffer-parameter shape values as positive values. * * \note call this function before Visit function's body to maximize * simplification efficiency */ - void MarkBufferMapShapes(const tirx::PrimFunc& func); + void MarkBufferParamShapes(const tirx::PrimFunc& func); /*! * \brief Use internal bound information to perform inter map simplification of indices. diff --git a/src/backend/trn/codegen/codegen_trn.cc b/src/backend/trn/codegen/codegen_trn.cc index fbad6be2a31a..ed96deb5aea3 100644 --- a/src/backend/trn/codegen/codegen_trn.cc +++ b/src/backend/trn/codegen/codegen_trn.cc @@ -115,7 +115,7 @@ void CodeGenTrainium::AddFunction(const GlobalVar& gvar, const PrimFunc& func) { LOG(FATAL) << "Trainium codegen currently only support buffer arguments"; }; std::string vid = AllocVarID(v.get()); - if (auto buffer = func->buffer_map.Get(v)) { + if (auto buffer = tirx::BufferParamMap(func->params).Get(v)) { var_idmap_[buffer.value().get()] = vid; } if (i >= static_cast(num_inputs.value())) { diff --git a/src/backend/trn/transform/lower_trainium_layout.cc b/src/backend/trn/transform/lower_trainium_layout.cc index 7c6f069d57ab..feef2333ca1a 100644 --- a/src/backend/trn/transform/lower_trainium_layout.cc +++ b/src/backend/trn/transform/lower_trainium_layout.cc @@ -346,7 +346,18 @@ namespace transform { Pass LowerTrainiumLayout() { auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) { auto* n = f.CopyOnWrite(); - std::tie(n->body, n->buffer_map) = TrainiumLayoutApplier::Lower(n->body, n->buffer_map); + auto [body, buffer_map] = + TrainiumLayoutApplier::Lower(n->body, tirx::BufferParamMap(n->params)); + ffi::Array params; + for (const Var& param : n->params) { + if (auto buffer = buffer_map.Get(param)) { + params.push_back(buffer.value().var()); + } else { + params.push_back(param); + } + } + n->body = std::move(body); + n->params = std::move(params); n->body = TrainiumBufferOffsetRemover::Remove(n->body); return f; }; diff --git a/src/relax/analysis/layout_transformation.cc b/src/relax/analysis/layout_transformation.cc index d3aef0ee1829..f5b4cf0085b1 100644 --- a/src/relax/analysis/layout_transformation.cc +++ b/src/relax/analysis/layout_transformation.cc @@ -562,7 +562,7 @@ class PrimFuncAnalyzer : public StmtExprVisitor { size_t first_write_index = func->params.size() - write_transformations.size(); for (size_t i = 0; i < write_transformations.size(); ++i) { auto param = func->params[first_write_index + i]; - ffi::Optional param_buf = func->buffer_map.Get(param); + ffi::Optional param_buf = tirx::AsBufferVar(param); TVM_FFI_ICHECK(param_buf.has_value()); TVM_FFI_ICHECK_EQ(param_buf.value()->shape.size(), write_transformations[i]->initial_indices.size()) diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc index 9290f0056490..43e703d6435b 100644 --- a/src/relax/analysis/tir_op_pattern_kind.cc +++ b/src/relax/analysis/tir_op_pattern_kind.cc @@ -39,7 +39,7 @@ class PatternKindAnalyzer : public StmtExprVisitor { public: explicit PatternKindAnalyzer(const tirx::PrimFunc& func) { for (const tirx::Var& param : func->params) { - ffi::Optional param_buf = func->buffer_map.Get(param); + ffi::Optional param_buf = tirx::AsBufferVar(param); if (param_buf.has_value()) { param_buffers_.insert(param_buf.value()); } @@ -535,7 +535,7 @@ bool HasReshapePattern(const PrimFunc& func) { ffi::Array buffer_args; for (const auto& param : func->params) { - if (auto buffer = func->buffer_map.Get(param)) { + if (auto buffer = tirx::AsBufferVar(param)) { buffer_args.push_back(buffer.value()); } } diff --git a/src/relax/backend/adreno/annotate_custom_storage.cc b/src/relax/backend/adreno/annotate_custom_storage.cc index 237fe8ca6016..644f46ed0075 100644 --- a/src/relax/backend/adreno/annotate_custom_storage.cc +++ b/src/relax/backend/adreno/annotate_custom_storage.cc @@ -231,8 +231,8 @@ * - Fusion * - FoldVDeviceScopeChange: There existed some ToVDevice copies from texture to buffer * This pass removes the copes and updates producer scope to global. - * - SpecializePrimFuncBasedOnCallSite: Finally we update the buffer maps according to - * VDevice scopes. + * - SpecializePrimFuncBasedOnCallSite: Finally we update the buffer parameter annotations + * according to VDevice scopes. * */ diff --git a/src/relax/backend/vm/codegen_vm_tir.cc b/src/relax/backend/vm/codegen_vm_tir.cc index 48191f32a3a0..75fd5da0bee7 100644 --- a/src/relax/backend/vm/codegen_vm_tir.cc +++ b/src/relax/backend/vm/codegen_vm_tir.cc @@ -199,7 +199,7 @@ class CodeGenVMTIR : public ExprFunctor(const Expr&)> { ffi::Array tir_params = {ctx_ptr_, reg_anylist_handle_, const_anylist_handle_, func_anylist_handle_}; ffi::String tir_func_name = system_lib_prefix_.value_or("") + "__vmtir__" + gsymbol.value(); - tirx::PrimFunc tir_func(tir_params, body, ret_type, {}); + tirx::PrimFunc tir_func(tir_params, body, ret_type); tir_func = WithAttr(tir_func, "global_symbol", tir_func_name); tir_func = WithAttr(tir_func, tvm::attr::kSTir, true); registers_num_ = 0; diff --git a/src/relax/distributed/axis_group_graph.cc b/src/relax/distributed/axis_group_graph.cc index 042d4593866d..40f285498c67 100644 --- a/src/relax/distributed/axis_group_graph.cc +++ b/src/relax/distributed/axis_group_graph.cc @@ -357,7 +357,7 @@ void BuildAxisGraphCallTIR(const Var& output_var, const Call& call, const tirx:: ffi::Array input_list = call->args[1].as_or_throw()->fields; input_list.push_back(output_var); for (int i = 0; i < static_cast(input_list.size()); i++) { - if (func->buffer_map.count(func->params[i])) { + if (func->params[i]->ty.as()) { input_var_to_relax_expr.Set(func->params[i], input_list[i]); } } diff --git a/src/relax/distributed/transform/lower_global_view_to_local_view.cc b/src/relax/distributed/transform/lower_global_view_to_local_view.cc index 7f9628398b43..a2da47cf9d43 100644 --- a/src/relax/distributed/transform/lower_global_view_to_local_view.cc +++ b/src/relax/distributed/transform/lower_global_view_to_local_view.cc @@ -156,19 +156,23 @@ class DistributedBufferCompactor : StmtExprMutator { const std::vector& sharding_specs, PrimFunc prim_func) { prim_func = s_tir::RenewDefs(prim_func); DistributedBufferCompactor compactor(sharding_specs, prim_func); - ffi::Map new_func_buffer_map; + ffi::Array new_params; ffi::Map replace_buffer_map; - for (const auto& pr : prim_func->buffer_map) { - BufferVar shard_buffer = compactor.ShardBuffer(pr.second); - new_func_buffer_map.Set(pr.first, shard_buffer); - if (!shard_buffer.same_as(pr.second)) { - replace_buffer_map.Set(pr.second, shard_buffer); + for (const Var& param : prim_func->params) { + if (!param->ty.as()) { + new_params.push_back(param); + continue; + } + BufferVar buffer(param); + BufferVar shard_buffer = compactor.ShardBuffer(buffer); + new_params.push_back(shard_buffer.var()); + if (!shard_buffer.same_as(buffer)) { + replace_buffer_map.Set(buffer, shard_buffer); } } Stmt new_body = compactor(prim_func->body); new_body = DistBufferReplacer::BufferReplace(new_body, replace_buffer_map); - PrimFunc new_func(prim_func->params, new_body, prim_func->ret_type, new_func_buffer_map, - prim_func->attrs, prim_func->span); + PrimFunc new_func(new_params, new_body, prim_func->ret_type, prim_func->attrs, prim_func->span); return std::make_tuple(new_func, compactor.add_allreduce_kind_); } @@ -184,10 +188,10 @@ class DistributedBufferCompactor : StmtExprMutator { std::unordered_set visited; for (int i = 0, j = 0; i < static_cast(prim_func->params.size()); i++) { Var param_var = prim_func->params[i]; - if (!prim_func->buffer_map.count(param_var)) { + if (!param_var->ty.as()) { continue; } - BufferVar param_buffer = prim_func->buffer_map[param_var]; + BufferVar param_buffer(param_var); ShardingSpec spec = sharding_specs_[j++]; for (int mesh_dim = 0; mesh_dim < static_cast(spec.first->shape.size()); mesh_dim++) { @@ -409,7 +413,7 @@ class LowerTIRToLocalView : public ExprMutator { for (size_t i = 0; i < args.size(); ++i) { const Expr& arg = args[i]; const tirx::Var& param = prim_func->params[i]; - if (prim_func->buffer_map.count(param)) { + if (param->ty.as()) { const auto* ty = GetTypeAs(arg); TVM_FFI_CHECK(ty, TypeError) << "Expected buffer parameter " << param << " to receive a distributed tensor, but " diff --git a/src/relax/transform/dataflow_inplace.cc b/src/relax/transform/dataflow_inplace.cc index d6a129722697..ca26a22a05f8 100644 --- a/src/relax/transform/dataflow_inplace.cc +++ b/src/relax/transform/dataflow_inplace.cc @@ -966,8 +966,8 @@ class ModuleInplaceTransformer : public ExprMutator { var_subst_map.Set(output_var, inplace_var); // also do the same with the buffer vars - auto output_buffer = old_primfunc->buffer_map.at(output_var); - auto inplace_buffer = old_primfunc->buffer_map.at(inplace_var); + auto output_buffer = tirx::BufferParamMap(old_primfunc->params).at(output_var); + auto inplace_buffer = tirx::BufferParamMap(old_primfunc->params).at(inplace_var); var_subst_map.Set(output_buffer.var(), inplace_buffer.var()); buffer_subst_map.Set(output_buffer, inplace_buffer); } @@ -983,7 +983,7 @@ class ModuleInplaceTransformer : public ExprMutator { }); // remove the now-unused outputs from the buffer map - auto new_buffer_map = old_primfunc->buffer_map; + auto new_buffer_map = tirx::BufferParamMap(old_primfunc->params); for (size_t i = 0; i < num_outs; i++) { new_buffer_map.erase(old_primfunc->params[num_params - num_outs + i]); } diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc index a7145144b839..69729afe2a89 100644 --- a/src/relax/transform/fuse_tir.cc +++ b/src/relax/transform/fuse_tir.cc @@ -465,7 +465,7 @@ class RelaxToTIRVarMapCollector : public ExprVisitor { void CollectVarMapping(const CallNode* call, const Expr& lhs_var, bool in_place) { GlobalVar gv = call->args[0].as_or_throw(); tirx::PrimFunc prim_func_ = mod_->Lookup(gv).as_or_throw(); - const auto& buffer_map = prim_func_->buffer_map; + const auto& buffer_map = tirx::BufferParamMap(prim_func_->params); const auto& tir_args = prim_func_->params; const auto& relax_args = call->args[1].as_or_throw()->fields; @@ -626,8 +626,8 @@ class FusedTIRConstructor : public ExprVisitor { int out_idx = 0; for (size_t i = 0; i < buffers.size(); ++i) { // Do not add output vars for in-place inputs - // (i.e., already listed in the buffer map. This would result - // in duplicates in the buffer map otherwise) + // (i.e., already listed among the buffer parameters, which would + // otherwise result in duplicate parameters) if (auto it = buffer_to_idx.find(buffers[i]); it != buffer_to_idx.end()) { auto idx = (*it).second; TVM_FFI_ICHECK(!inplace_indices_.count(idx)) @@ -810,9 +810,9 @@ class FusedTIRConstructor : public ExprVisitor { for (size_t i = 0; i < call_args.size(); ++i) { const Expr& arg = call_args[i]; const tirx::Var& param = func->params[i]; - if (func->buffer_map.count(param)) { + if (auto buffer = tirx::AsBufferVar(param)) { arg_list.push_back(arg); - buffer_list.push_back(func->buffer_map.at(param)); + buffer_list.push_back(buffer.value()); } else { auto prim_arg = arg.as(); TVM_FFI_CHECK(prim_arg.has_value(), TypeError) @@ -826,20 +826,21 @@ class FusedTIRConstructor : public ExprVisitor { MapArgsToBuffer(arg_list, buffer_list); } - static ffi::Array GetPrimFuncOutputParams(const tirx::PrimFunc& func, - const ffi::Array& output_indices) { + static ffi::Array GetPrimFuncOutputParams( + const tirx::PrimFunc& func, const ffi::Array& output_indices) { size_t n = func->params.size(); size_t output_size = output_indices.size(); TVM_FFI_ICHECK_GE(n, output_size); - ffi::Array ret; + ffi::Array ret; for (int64_t idx : output_indices) { int i = static_cast(idx); const tirx::Var& param = func->params[static_cast(i)]; - TVM_FFI_ICHECK(param->ty.as()) - << "The output params of a PrimFunc must be buffer handles, but parameter " << i - << " has type " << param->ty; - ret.push_back(param); + auto buffer = tirx::AsBufferVar(param); + TVM_FFI_ICHECK(buffer.has_value()) + << "The output params of a PrimFunc must be buffers, but parameter " << i << " has type " + << param->ty; + ret.push_back(buffer.value()); } return ret; } @@ -871,10 +872,9 @@ class FusedTIRConstructor : public ExprVisitor { } } - ffi::Array output_params = GetPrimFuncOutputParams(func, output_idxs); + ffi::Array output_params = GetPrimFuncOutputParams(func, output_idxs); for (size_t i = 0; i < output_size; ++i) { - const tirx::Var& param = output_params[i]; - const tirx::BufferVar& buffer = func->buffer_map.at(param); + const tirx::BufferVar& buffer = output_params[i]; // if this is an inplace output, do not do an intermediate allocation if (output_idxs[i] < num_inputs) { @@ -988,8 +988,13 @@ class FusedTIRConstructor : public ExprVisitor { body = subst.Substitute(body); body = tirx::SBlock({}, {}, {}, "root", std::move(body), std::nullopt, alloc_buffers); body = tirx::SBlockRealize({}, IntImm::Bool(true), body.as_or_throw()); - tirx::PrimFunc func(func_info_.params, body, VoidType(), func_info_.buffer_map, - DictAttrs(attr_map)); + ffi::Array params = func_info_.params.Map([&](const tirx::Var& param) { + if (auto buffer = func_info_.buffer_map.Get(param)) { + return buffer.value().var(); + } + return param; + }); + tirx::PrimFunc func(params, body, VoidType(), DictAttrs(attr_map)); // Renew function defs to prevent using the same symbolic vars in different functions return s_tir::RenewDefs(func); } @@ -1032,9 +1037,9 @@ class FusedTIRConstructor : public ExprVisitor { * function */ ffi::Map buffer_subst_map; - /*! \brief The `buffer_map` in the fused function*/ + /*! \brief Buffer annotations keyed by their placeholder parameters. */ ffi::Map buffer_map; - /*! \brief The output buffers in the function buffer_map*/ + /*! \brief The output buffers among the function parameters. */ std::unordered_set output_buffers; /*! \brief The name of the fused function */ std::string global_name = "fused"; diff --git a/src/relax/transform/meta_schedule.cc b/src/relax/transform/meta_schedule.cc index 161f459038f2..bd035d970f39 100644 --- a/src/relax/transform/meta_schedule.cc +++ b/src/relax/transform/meta_schedule.cc @@ -134,7 +134,6 @@ Pass MetaScheduleApplyDatabase(ffi::Optional work_dir, bool enable_ tirx::PrimFunc new_prim_func = tirx::PrimFunc(/*params=*/tuned_prim_func->params, /*body=*/tuned_prim_func->body, /*ret_type=*/tuned_prim_func->ret_type, - /*buffer_map=*/tuned_prim_func->buffer_map, /*attrs=*/prim_func->attrs); new_prim_func = WithAttr(std::move(new_prim_func), tirx::attr::kIsScheduled, true); result.Set(gv, new_prim_func); diff --git a/src/relax/transform/rewrite_dataflow_reshape.cc b/src/relax/transform/rewrite_dataflow_reshape.cc index 97a3658f5347..c73da19ae6a4 100644 --- a/src/relax/transform/rewrite_dataflow_reshape.cc +++ b/src/relax/transform/rewrite_dataflow_reshape.cc @@ -39,7 +39,7 @@ namespace relax { std::vector GetUsedTensorArgIndices(const tirx::PrimFunc& fn, size_t num_args) { std::vector indices; for (size_t i = 0; i < num_args; ++i) { - if (auto buffer = fn->buffer_map.Get(fn->params[i])) { + if (auto buffer = tirx::AsBufferVar(fn->params[i])) { auto buffer_var = buffer.value().var(); if (tirx::UsesVar(fn->body, [=](const tirx::VarNode* var) { return var == buffer_var.get(); })) { diff --git a/src/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc index 18618840d98c..adacbfc14862 100644 --- a/src/relax/transform/split_call_tir_by_pattern.cc +++ b/src/relax/transform/split_call_tir_by_pattern.cc @@ -70,9 +70,8 @@ class ForMatcher : public TensorizeComparator { } // Get evaluated symbols, buffers from the pattern. for (const auto& arg : pattern_->params) { - auto it = pattern_->buffer_map.find(arg); - if (it != pattern_->buffer_map.end()) { - auto itt = rhs_buffer_map_.find((*it).second); + if (auto buffer = tirx::AsBufferVar(arg)) { + auto itt = rhs_buffer_map_.find(buffer.value()); TVM_FFI_ICHECK(itt != rhs_buffer_map_.end()); evaluated_buffers.push_back(itt->second); } @@ -392,7 +391,7 @@ class TIRPatternMatcher { for (const TIRPattern& pattern : patterns_) { tirx::PrimFunc pattern_func = pattern; ffi::Array pattern_symbolic_vars; - int buffer_count = pattern_func->buffer_map.size(); + int buffer_count = tirx::BufferParamMap(pattern_func->params).size(); for (int i = buffer_count; i < static_cast(pattern_func->params.size()); i++) { pattern_symbolic_vars.push_back(pattern_func->params[i]); } @@ -614,7 +613,8 @@ std::pair> SplitFunctions( for (const auto& buffer : func1_args) { TVM_FFI_ICHECK(partitioner.input1.find(buffer) != partitioner.input1.end()); for (size_t i = 0; i < func->params.size(); i++) { - if (func->buffer_map[func->params[i]].same_as(buffer)) { + auto param_buffer = tirx::AsBufferVar(func->params[i]); + if (param_buffer.has_value() && param_buffer.value().same_as(buffer)) { new_params1.push_back(func->params[i]); arg_partition1.push_back(i); break; @@ -622,23 +622,17 @@ std::pair> SplitFunctions( } } arg_partition->push_back(arg_partition1); - new_params1.push_back(Var("output", PointerType::VoidPointerTy())); - ffi::Map new_buffer_map1; - for (const auto& kv : func->buffer_map) { - if (partitioner.input1.count(kv.second)) { - new_buffer_map1.Set(kv.first, kv.second); - } - } - new_buffer_map1.Set(new_params1.back(), partitioner.intermediate_buffer); - PrimFunc func1 = PrimFunc(new_params1, body1, func->ret_type, new_buffer_map1, func->attrs); + new_params1.push_back(partitioner.intermediate_buffer.var()); + PrimFunc func1 = PrimFunc(new_params1, body1, func->ret_type, func->attrs); func1 = WithAttr(func1, kLibraryKernel, library_code); // Step 4. Craft the second function. ffi::Array new_params2; std::vector arg_partition2; - new_params2.push_back(Var("input", PointerType::VoidPointerTy())); + new_params2.push_back(partitioner.intermediate_buffer.var()); for (int i = 0; i < static_cast(func->params.size()); i++) { Var param = func->params[i]; - if (partitioner.input2.count(func->buffer_map[param])) { + auto param_buffer = tirx::AsBufferVar(param); + if (param_buffer.has_value() && partitioner.input2.count(param_buffer.value())) { new_params2.push_back(param); if (i != static_cast(func->params.size()) - 1) { arg_partition2.push_back(i); @@ -646,14 +640,7 @@ std::pair> SplitFunctions( } } arg_partition->push_back(arg_partition2); - ffi::Map new_buffer_map2; - new_buffer_map2.Set(new_params2[0], partitioner.intermediate_buffer); - for (const auto& kv : func->buffer_map) { - if (partitioner.input2.count(kv.second)) { - new_buffer_map2.Set(kv.first, kv.second); - } - } - PrimFunc func2 = PrimFunc(new_params2, body2, func->ret_type, new_buffer_map2, func->attrs); + PrimFunc func2 = PrimFunc(new_params2, body2, func->ret_type, func->attrs); return {func1, func2}; } } // namespace tirx @@ -751,7 +738,8 @@ class SplitMutator : public ExprMutator { if (lib_func->IsInstance()) return ffi::GetRef(op); TVM_FFI_ICHECK(lib_func->IsInstance()); builder_->UpdateFunction(gv, lib_func); - tirx::BufferVar intermediate_buffer = func1->buffer_map.at(func1->params.back()); + tirx::BufferVar intermediate_buffer = + tirx::BufferParamMap(func1->params).at(func1->params.back()); PrimType dtype = intermediate_buffer->dtype; Call call1(Type::Missing(), call_dps_packed_, {lib_func, Tuple(args1)}, call->attrs, {TensorType(ShapeExpr(intermediate_buffer->shape), dtype)}); diff --git a/src/relax/transform/split_layout_rewrite_preproc.cc b/src/relax/transform/split_layout_rewrite_preproc.cc index ff340d153fec..1c920155d4cd 100644 --- a/src/relax/transform/split_layout_rewrite_preproc.cc +++ b/src/relax/transform/split_layout_rewrite_preproc.cc @@ -62,15 +62,12 @@ class SplitPrimFuncLayoutRewrite : public StmtMutator { // Step 2: Create the params for the new PrimFunc ffi::Array params; - ffi::Map buffer_map; for (const auto& info : rewrite_infos_) { - params.push_back(Var(info.pre_rewrite_buffer.name(), PointerType::VoidPointerTy())); - buffer_map.Set(params.back(), info.pre_rewrite_buffer); + params.push_back(info.pre_rewrite_buffer.var()); } for (const auto& info : rewrite_infos_) { - params.push_back(Var(info.post_rewrite_buffer.name(), PointerType::VoidPointerTy())); - buffer_map.Set(params.back(), info.post_rewrite_buffer); + params.push_back(info.post_rewrite_buffer.var()); } // Step 3: Create the body for the new PrimFunc @@ -94,7 +91,7 @@ class SplitPrimFuncLayoutRewrite : public StmtMutator { } } DictAttrs attrs(dict); - PrimFunc func = PrimFunc(params, body, VoidType(), buffer_map, attrs); + PrimFunc func = PrimFunc(params, body, VoidType(), attrs); return s_tir::RenewDefs(func); } @@ -102,11 +99,10 @@ class SplitPrimFuncLayoutRewrite : public StmtMutator { PrimFunc create_compute_func() const { // Step 1: Create the params for the new PrimFunc ffi::Array params = original_func_->params; - ffi::Map buffer_map = original_func_->buffer_map; for (const auto& info : rewrite_infos_) { const Var& param = params[info.buffer_index]; - TVM_FFI_ICHECK(buffer_map[param] == info.pre_rewrite_buffer); - buffer_map.Set(param, info.post_rewrite_buffer); + TVM_FFI_ICHECK(tirx::AsBufferVar(param).value() == info.pre_rewrite_buffer); + params.Set(info.buffer_index, info.post_rewrite_buffer.var()); } // Step 2: Create the body for the new PrimFunc @@ -140,7 +136,7 @@ class SplitPrimFuncLayoutRewrite : public StmtMutator { } } DictAttrs attrs(dict); - PrimFunc func = PrimFunc(original_func_->params, body, VoidType(), buffer_map, attrs); + PrimFunc func = PrimFunc(params, body, VoidType(), attrs); return s_tir::RenewDefs(func); } @@ -193,7 +189,8 @@ class SplitPrimFuncLayoutRewrite : public StmtMutator { const BufferVar& preproc_buffer = op->reads[0]->buffer; int buffer_index = -1; for (size_t i = 0; i < original_func_->params.size(); ++i) { - const BufferVar& buffer = original_func_->buffer_map[original_func_->params[i]]; + const BufferVar& buffer = + tirx::BufferParamMap(original_func_->params)[original_func_->params[i]]; if (buffer == preproc_buffer) { buffer_index = i; break; diff --git a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc index b0de1d8f2115..bae15155981e 100644 --- a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc +++ b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc @@ -45,7 +45,7 @@ class LCADetector : public StmtExprVisitor { public: static ffi::Map> Detect(const PrimFunc& func) { LCADetector detector; - for (const auto& kv : func->buffer_map) { + for (const auto& kv : tirx::BufferParamMap(func->params)) { const BufferVar& buffer = kv.second; detector.buffer_var_map_.emplace(buffer.get(), buffer.get()); } diff --git a/src/s_tir/backend/adreno/inject_texture_alloc.cc b/src/s_tir/backend/adreno/inject_texture_alloc.cc index d60589e9ef30..27544a2d80f1 100644 --- a/src/s_tir/backend/adreno/inject_texture_alloc.cc +++ b/src/s_tir/backend/adreno/inject_texture_alloc.cc @@ -48,7 +48,7 @@ class TextureAllocInjector : public arith::IRMutatorWithAnalyzer { arith::Analyzer ana; auto pass = TextureAllocInjector(ana); auto writer = func.CopyOnWrite(); - pass.MarkBufferMapShapes(func); + pass.MarkBufferParamShapes(func); writer->body = pass.VisitStmt(func->body); return func; } diff --git a/src/s_tir/backend/adreno/texture_flatten.cc b/src/s_tir/backend/adreno/texture_flatten.cc index 5d14e314bdd2..7c042a2518ee 100644 --- a/src/s_tir/backend/adreno/texture_flatten.cc +++ b/src/s_tir/backend/adreno/texture_flatten.cc @@ -159,7 +159,8 @@ PrimFunc TextureFlattenHandler(PrimFunc func) { auto fptr = func.CopyOnWrite(); IRVisitorWithAnalyzer bound_analyzer; bound_analyzer(fptr->body); - fptr->body = TextureFlattener(fptr->buffer_map, &bound_analyzer)(std::move(fptr->body)); + fptr->body = + TextureFlattener(tirx::BufferParamMap(fptr->params), &bound_analyzer)(std::move(fptr->body)); return func; } diff --git a/src/s_tir/meta_schedule/arg_info.cc b/src/s_tir/meta_schedule/arg_info.cc index 94fec3bdf988..375b24a69448 100644 --- a/src/s_tir/meta_schedule/arg_info.cc +++ b/src/s_tir/meta_schedule/arg_info.cc @@ -96,7 +96,7 @@ ffi::Array ArgInfo::FromPrimFunc(const tirx::PrimFunc& func) { ffi::Array result; result.reserve(func->params.size()); for (const tirx::Var& arg : func->params) { - if (ffi::Optional _buffer = func->buffer_map.Get(arg)) { + if (ffi::Optional _buffer = tirx::BufferParamMap(func->params).Get(arg)) { tirx::BufferVar buffer = _buffer.value(); result.push_back(TensorInfo(/*dtype=*/buffer->dtype->dtype, /*shape=*/AsVector(buffer->shape))); diff --git a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc index a8578f6567b7..b6050584fe85 100644 --- a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc +++ b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc @@ -1288,7 +1288,7 @@ class PerStoreFeatureCollector : private StmtVisitor { for (const auto& kv : mod->functions) { if (const PrimFuncNode* func = kv.second.as()) { collector(func->body); - for (const auto& it : func->buffer_map) { + for (const auto& it : tirx::BufferParamMap(func->params)) { collector.HandleBufferAlloc(it.second); } } diff --git a/src/s_tir/meta_schedule/postproc/rewrite_layout.cc b/src/s_tir/meta_schedule/postproc/rewrite_layout.cc index f6a439b650d1..49dd31a460f0 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_layout.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_layout.cc @@ -129,7 +129,7 @@ ffi::Array CollectLayoutFreeBuffers(const PrimFuncNode* func) { for (int64_t index : layout_free_buffer_index) { TVM_FFI_ICHECK(static_cast(index) < func->params.size()); const Var& param = func->params[index]; - layout_free_buffers.push_back(func->buffer_map.at(param)); + layout_free_buffers.push_back(tirx::BufferParamMap(func->params).at(param)); } LayoutFreeBufferCollector collector; diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc index 31ef367cb786..f86597264ae3 100644 --- a/src/s_tir/schedule/analysis/analysis.cc +++ b/src/s_tir/schedule/analysis/analysis.cc @@ -1314,7 +1314,7 @@ void AddShapeVarBounds(const ScheduleState& state, const StmtSRefNode* sref, sref = sref->parent; } const PrimFuncNode* f = GetRootPrimFunc(state->mod, sref->stmt, nullptr); - for (const auto& kv : f->buffer_map) { + for (const auto& kv : tirx::BufferParamMap(f->params)) { const BufferVar& buffer = kv.second; for (const PrimExpr& e : buffer->shape) { analyzer->MarkGlobalNonNegValue(e); diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc b/src/s_tir/schedule/primitive/blockize_tensorize.cc index 4f2f15a230c0..f3cd1fb57750 100644 --- a/src/s_tir/schedule/primitive/blockize_tensorize.cc +++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc @@ -796,8 +796,8 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int std::unordered_map impl2desc; TVM_FFI_ICHECK_EQ(intrin_desc->params.size(), intrin_impl->params.size()); for (int i = 0, n = intrin_desc->params.size(); i < n; ++i) { - const BufferVar& desc = intrin_desc->buffer_map[intrin_desc->params[i]]; - const BufferVar& impl = intrin_impl->buffer_map[intrin_impl->params[i]]; + const BufferVar& desc = tirx::BufferParamMap(intrin_desc->params)[intrin_desc->params[i]]; + const BufferVar& impl = tirx::BufferParamMap(intrin_impl->params)[intrin_impl->params[i]]; impl2desc[impl] = desc; } std::unordered_map impl2cur; @@ -821,7 +821,7 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int ffi::Array match_buffer_regions; match_buffer_regions.reserve(intrin_impl->params.size()); for (int i = 0, n = intrin_impl->params.size(); i < n; ++i) { - const BufferVar& impl = intrin_impl->buffer_map.at(intrin_impl->params[i]); + const BufferVar& impl = tirx::BufferParamMap(intrin_impl->params).at(intrin_impl->params[i]); const BufferVar& cur = impl2cur.at(impl); const ffi::Array& old_region = impl2region.at(impl); const std::vector& indices_base = comparator.buffer_indices_.at(cur); diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc b/src/s_tir/schedule/primitive/layout_transformation.cc index 9c8f835a0a55..872f99ab22cf 100644 --- a/src/s_tir/schedule/primitive/layout_transformation.cc +++ b/src/s_tir/schedule/primitive/layout_transformation.cc @@ -1229,7 +1229,7 @@ void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_ index_map, opt_inverse, padding_predicate, pad_value); SBlock new_scope_block = new_stmt.as_or_throw(); - // Step 4: Rewrite buffer_map of the PrimFunc if necessary. + // Step 4: Rewrite the PrimFunc buffer parameter if necessary. if (!defining_site_sref.has_value()) { GlobalVar g_var; const auto* old_func = GetRootPrimFunc(self->mod, scope_block, &g_var); @@ -1237,7 +1237,7 @@ void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_ ffi::MapObj* new_map = new_mod->functions.CopyOnWrite(); ffi::Map new_buffer_map; - for (auto [var, buffer] : old_func->buffer_map) { + for (auto [var, buffer] : tirx::BufferParamMap(old_func->params)) { if (buffer.same_as(old_buffer)) { buffer = new_buffer; } diff --git a/src/s_tir/transform/inject_permuted_layout.cc b/src/s_tir/transform/inject_permuted_layout.cc index 6c2ad82710a7..f83d9b45f116 100644 --- a/src/s_tir/transform/inject_permuted_layout.cc +++ b/src/s_tir/transform/inject_permuted_layout.cc @@ -70,8 +70,9 @@ class PermutedLayoutInjector : private IRMutatorWithAnalyzer { private: explicit PermutedLayoutInjector(PrimFunc func, const Analyzer& analyzer) : IRMutatorWithAnalyzer(analyzer) { - buffer_map_.insert(func->buffer_map.begin(), func->buffer_map.end()); - for (const auto& [_, buffer] : func->buffer_map) { + auto param_buffers = tirx::BufferParamMap(func->params); + buffer_map_.insert(param_buffers.begin(), param_buffers.end()); + for (const auto& [_, buffer] : param_buffers) { buffer_map_.insert({buffer.var(), buffer}); } } diff --git a/src/s_tir/transform/inject_software_pipeline.cc b/src/s_tir/transform/inject_software_pipeline.cc index 2f20f10303ec..1cb9a48de07f 100644 --- a/src/s_tir/transform/inject_software_pipeline.cc +++ b/src/s_tir/transform/inject_software_pipeline.cc @@ -1065,7 +1065,7 @@ class PipelineInjector : private StmtExprMutator { static Stmt Inject(const PrimFunc& func) { auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); PipelineInjector injector(global_symbol); - for (const auto& kv : func->buffer_map) { + for (const auto& kv : tirx::BufferParamMap(func->params)) { const BufferVar& buffer = kv.second; injector.buffer_data_to_buffer_.Set(buffer.var(), buffer); } diff --git a/src/s_tir/transform/lower_thread_allreduce.cc b/src/s_tir/transform/lower_thread_allreduce.cc index 027c124d433f..68b4678ec96b 100644 --- a/src/s_tir/transform/lower_thread_allreduce.cc +++ b/src/s_tir/transform/lower_thread_allreduce.cc @@ -966,7 +966,8 @@ Pass LowerThreadAllreduce() { auto target = f->GetAttr(tvm::attr::kTarget); TVM_FFI_ICHECK(target.has_value()) << "LowerThreadAllreduce: Require the target attribute"; const TargetNode* target_node = target.as(); - ThreadAllreduceBuilder thread_all_reduce(target_node, f->params, f->buffer_map); + ThreadAllreduceBuilder thread_all_reduce(target_node, f->params, + tirx::BufferParamMap(f->params)); n->body = thread_all_reduce(n->body); // Post-process: apply deferred remappings for flat IR DeferredRemapper remapper(thread_all_reduce.alloc_remap_, thread_all_reduce.var_remap_, diff --git a/src/s_tir/transform/plan_update_buffer_allocation_location.cc b/src/s_tir/transform/plan_update_buffer_allocation_location.cc index ccf5400cae19..ee6dc4a7eeba 100644 --- a/src/s_tir/transform/plan_update_buffer_allocation_location.cc +++ b/src/s_tir/transform/plan_update_buffer_allocation_location.cc @@ -56,7 +56,7 @@ class BufferAllocateOrderCollector : public StmtExprVisitor { public: static ffi::Array Collect(const PrimFunc& func) { BufferAllocateOrderCollector collector; - for (const auto& kv : func->buffer_map) { + for (const auto& kv : tirx::BufferParamMap(func->params)) { collector.buffer_alloc_recorder_.push_back(kv.second); } collector(func->body); @@ -114,7 +114,7 @@ class BufferAllocationLocator : public StmtExprMutator { collector(func->body); managed_allocations_ = collector.managed_allocations; - for (const auto& kv : func->buffer_map) { + for (const auto& kv : tirx::BufferParamMap(func->params)) { const BufferVar& buffer = kv.second; arg_buffer_vars.emplace(buffer.get()); buffer_data_to_buffer_.Set(buffer.var(), buffer); diff --git a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc index 8cb514954040..076218e6bab0 100644 --- a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc +++ b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc @@ -126,16 +126,22 @@ class WeightLayoutRewriteBlockRemover : public StmtMutator { PrimFuncNode* n = f_.CopyOnWrite(); - ffi::Map buffer_map; - for (const auto& [param, buffer] : f_->buffer_map) { + ffi::Array params; + for (const tirx::Var& param : f_->params) { + auto opt_buffer = tirx::AsBufferVar(param); + if (!opt_buffer.has_value()) { + params.push_back(param); + continue; + } + BufferVar buffer = opt_buffer.value(); auto it = buf_map.find(buffer); if (it != buf_map.end()) { - buffer_map.Set(param, (*it).second); + params.push_back((*it).second.var()); } else { - buffer_map.Set(param, buffer); + params.push_back(param); } } - n->buffer_map = std::move(buffer_map); + n->params = std::move(params); return f_; } }; diff --git a/src/s_tir/transform/renew_defs.cc b/src/s_tir/transform/renew_defs.cc index 519a82d8be35..273ee423031d 100644 --- a/src/s_tir/transform/renew_defs.cc +++ b/src/s_tir/transform/renew_defs.cc @@ -48,15 +48,16 @@ class RenewDefMutator : public StmtExprMutator { public: static PrimFunc Transform(const PrimFunc& func) { RenewDefMutator generator; - // Redefine params - ffi::Array params; + // Redefine scalar parameters first, because they may occur in a buffer + // parameter's type annotation. for (const auto& param : func->params) { - params.push_back(generator.ReDefineVar(param)); + if (!tirx::AsBufferVar(param)) { + generator.ReDefineVar(param); + } } for (const auto& param : func->params) { - auto it = func->buffer_map.find(param); - if (it != func->buffer_map.end()) { - const BufferVar& buffer = (*it).second; + if (auto opt_buffer = tirx::AsBufferVar(param)) { + const BufferVar& buffer = opt_buffer.value(); for (const PrimExpr& e : buffer->shape) { if (auto var = e.as()) { if (generator.remap_.count(var.value()) == 0) { @@ -66,22 +67,20 @@ class RenewDefMutator : public StmtExprMutator { } } } - // Redefine buffers in order + // Redefine buffer parameters in order, preserving the original signature. // TODO(Siyuan Feng): checking var is used after define - ffi::Map buffer_map; + ffi::Array params; for (const auto& param : func->params) { - auto it = func->buffer_map.find(param); - if (it != func->buffer_map.end()) { - const BufferVar& buffer = (*it).second; - Var new_param = generator.VisitExpr(param).as_or_throw(); - BufferVar new_buffer = generator.DefineBuffer(buffer); - buffer_map.Set(new_param, new_buffer); + if (auto opt_buffer = tirx::AsBufferVar(param)) { + params.push_back(generator.DefineBuffer(opt_buffer.value())); + } else { + params.push_back(generator.VisitExpr(param).as_or_throw()); } } // Visit body Stmt body = generator(func->body); // Recreate function - return PrimFunc(params, body, func->ret_type, buffer_map, func->attrs, func->span); + return PrimFunc(params, body, func->ret_type, func->attrs, func->span); } private: diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc index f8fd8d1b8a54..07dd1a054c05 100644 --- a/src/target/llvm/codegen_llvm.cc +++ b/src/target/llvm/codegen_llvm.cc @@ -269,8 +269,8 @@ llvm::Function* CodeGenLLVM::DeclareFunctionInternal(const GlobalVar& gvar, cons return it->second; } - TVM_FFI_ICHECK_EQ(func->buffer_map.size(), 0U) - << "Cannot codegen function with buffer_map, please lower them first"; + TVM_FFI_ICHECK_EQ(tirx::BufferParamMap(func->params).size(), 0U) + << "Cannot codegen BufferType-annotated parameters; please lower them first"; std::vector param_types; is_restricted_ = func->HasNonzeroAttr(tirx::attr::kNoAlias); diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc index 90c8922c3c97..9d23f16d0473 100644 --- a/src/te/operation/create_primfunc.cc +++ b/src/te/operation/create_primfunc.cc @@ -133,7 +133,7 @@ class LayoutFreePlaceholdersNormalizer : public StmtMutator { PrimFunc Process(PrimFunc func) { for (int i = 0, n = func->params.size(); i < n; ++i) { if (auto v = func->params[i].as()) { - if (ffi::Optional buffer = func->buffer_map.Get(v.value())) { + if (ffi::Optional buffer = tirx::BufferParamMap(func->params).Get(v.value())) { buffer2index_[buffer.value()] = i; } } @@ -753,19 +753,15 @@ void RewriteStageToBlock(const te::Operation& op, CreateFuncInfo* info, PrimFunc GenerateAndCompletePrimFunc(const ffi::Array& arg_list, const ffi::Array& root_stmts, CreateFuncInfo* info) { ffi::Array parameters; - ffi::Map buffer_map; for (const te::Tensor& tensor : arg_list) { auto it = info->tensor2buffers.find(tensor); TVM_FFI_ICHECK(it != info->tensor2buffers.end()); - Var arg("var_" + tensor->GetNameHint(), PointerType::VoidPointerTy()); - parameters.push_back(arg); - buffer_map.Set(arg, it->second); + parameters.push_back(it->second.var()); } PrimFunc func = WithAttrs( PrimFunc(/*params=*/std::move(parameters), /*body=*/SeqStmt::Flatten(root_stmts), - /*ret_type=*/VoidType(), - /*buffer_map=*/std::move(buffer_map)), + /*ret_type=*/VoidType()), {{"global_symbol", ffi::String("main")}, {"tirx.noalias", true}, {tvm::attr::kSTir, true}}); const auto fcomplete = tvm::ffi::Function::GetGlobal("script.Complete"); TVM_FFI_ICHECK(fcomplete.has_value()); @@ -819,15 +815,12 @@ TVM_FFI_STATIC_INIT_BLOCK() { PrimFunc GenerateAndCompletePrimFunc(const ffi::Array& arg_tir_var_list, const ffi::Array& root_stmts, CreateFuncInfo* info) { ffi::Array parameters; - ffi::Map buffer_map; for (const ffi::ObjectRef& arg : arg_tir_var_list) { if (auto opt_tensor = arg.as()) { te::Tensor tensor = opt_tensor.value(); auto it = info->tensor2buffers.find(tensor); TVM_FFI_ICHECK(it != info->tensor2buffers.end()); - Var param("var_" + tensor->GetNameHint(), PointerType::VoidPointerTy()); - parameters.push_back(param); - buffer_map.Set(param, it->second); + parameters.push_back(it->second.var()); } else if (auto var = arg.as()) { parameters.push_back(var.value()); } @@ -835,8 +828,7 @@ PrimFunc GenerateAndCompletePrimFunc(const ffi::Array& arg_tir_v PrimFunc func = WithAttrs( PrimFunc(/*params=*/std::move(parameters), /*body=*/SeqStmt::Flatten(root_stmts), - /*ret_type=*/VoidType(), - /*buffer_map=*/std::move(buffer_map)), + /*ret_type=*/VoidType()), {{"global_symbol", ffi::String("main")}, {"tirx.noalias", true}, {tvm::attr::kSTir, true}}); const auto fcomplete = tvm::ffi::Function::GetGlobal("script.Complete"); TVM_FFI_ICHECK(fcomplete.has_value()); diff --git a/src/tirx/analysis/verify_memory.cc b/src/tirx/analysis/verify_memory.cc index 2656bc521bd1..ae4d0211e281 100644 --- a/src/tirx/analysis/verify_memory.cc +++ b/src/tirx/analysis/verify_memory.cc @@ -103,7 +103,7 @@ class MemoryAccessVerifier final : protected StmtExprVisitor { /// Check if the value of a Variable comes from function argument. bool IsFromFunctionArgs(const VarNode* var) const { const VarNode* V = var; - for (auto kv : func_->buffer_map) { + for (auto kv : tirx::BufferParamMap(func_->params)) { if (V == kv.second.get()) return true; } diff --git a/src/tirx/analysis/verify_ssa.cc b/src/tirx/analysis/verify_ssa.cc index b14dd6062e3d..5ba2013be2bf 100644 --- a/src/tirx/analysis/verify_ssa.cc +++ b/src/tirx/analysis/verify_ssa.cc @@ -93,7 +93,7 @@ class SSAVerifier final : public StmtExprVisitor { MarkDef(param, param); } - for (auto kv : func->buffer_map) { + for (auto kv : tirx::BufferParamMap(func->params)) { this->DefineBuffer(kv.second); } this->VisitStmt(func->body); diff --git a/src/tirx/analysis/verify_well_formed.cc b/src/tirx/analysis/verify_well_formed.cc index a2bb35aae6d5..2fc5c3a4b796 100644 --- a/src/tirx/analysis/verify_well_formed.cc +++ b/src/tirx/analysis/verify_well_formed.cc @@ -239,7 +239,7 @@ class UndefinedVarVerifier : public Verifier { /*! \brief Verify that buffers with a declaration are not used outside their declared scope. * * When a buffer is declared via one of the following sites: - * - PrimFunc buffer_map (function parameter buffers) + * - BufferType-annotated PrimFunc parameters * - DeclBuffer statement * - SBlock::alloc_buffers * - SBlock::match_buffers diff --git a/src/tirx/ir/data_type_rewriter.cc b/src/tirx/ir/data_type_rewriter.cc index a8cb4f62e1ff..ca264bda684f 100644 --- a/src/tirx/ir/data_type_rewriter.cc +++ b/src/tirx/ir/data_type_rewriter.cc @@ -635,14 +635,16 @@ PrimFunc IndexDataTypeNormalizer::Rewrite(PrimFunc func) { buffer_remap_.clear(); ivmap_.clear(); // start rewrite - ffi::Map new_buffer_map = func->buffer_map; - for (const auto& [var, buffer] : func->buffer_map) { - new_buffer_map.Set(var, VisitBufferDef(buffer, /*alloc_data=*/true)); + for (const auto& [_, buffer] : tirx::BufferParamMap(func->params)) { + VisitBufferDef(buffer, /*alloc_data=*/true); } // remap params bool is_enabled = true; std::swap(is_enabled_, is_enabled); ffi::Array params = func->params.Map([this](Var param) { + if (auto buffer = AsBufferVar(param)) { + return buffer_remap_.Get(buffer.value()).value_or(buffer.value()).var(); + } if (auto param_ty = param->ty.as(); param_ty && param_ty.value().MatchesCode(DLDataTypeCode::kDLInt)) { return this->VisitPrimExpr(param.as_or_throw()).as_or_throw(); @@ -654,7 +656,6 @@ PrimFunc IndexDataTypeNormalizer::Rewrite(PrimFunc func) { PrimFuncNode* new_func = func.CopyOnWrite(); new_func->params = std::move(params); - new_func->buffer_map = std::move(new_buffer_map); new_func->body = VisitStmt(std::move(new_func->body)); return func; } diff --git a/src/tirx/ir/function.cc b/src/tirx/ir/function.cc index c991e572dc28..4b4ab4878022 100644 --- a/src/tirx/ir/function.cc +++ b/src/tirx/ir/function.cc @@ -38,12 +38,26 @@ TVM_FFI_STATIC_INIT_BLOCK() { } namespace { +ffi::Array NormalizeBufferParams(ffi::Array params, + const ffi::Map& buffer_map) { + ffi::Array normalized; + normalized.reserve(params.size()); + for (const Var& param : params) { + if (auto buffer = buffer_map.Get(param)) { + normalized.push_back(buffer.value().var()); + } else { + normalized.push_back(param); + } + } + return normalized; +} + tvm::Type InferType(const PrimFunc& prim_func) { ffi::Array params; for (const auto& param : prim_func->params) { tvm::Type param_ty = [&]() -> tvm::Type { - if (auto opt_buf = prim_func->buffer_map.Get(param)) { - auto buf = opt_buf.value(); + if (param->ty.as()) { + BufferVar buf(param); relax::ShapeExpr shape( buf->shape.Map([](PrimExpr dim) { return cast(PrimType::Int(64), dim); })); return relax::TensorType(shape, buf->dtype); @@ -79,8 +93,8 @@ tvm::Type InferType(const PrimFunc& prim_func) { } // namespace // Get the function type of a PrimFunc -PrimFunc::PrimFunc(ffi::Array params, Stmt body, Type ret_type, - ffi::Map buffer_map, DictAttrs attrs, Span span) { +PrimFunc::PrimFunc(ffi::Array params, Stmt body, Type ret_type, DictAttrs attrs, + Span span) { if (ret_type.IsMissing()) { ret_type = VoidType(); } @@ -89,7 +103,6 @@ PrimFunc::PrimFunc(ffi::Array params, Stmt body, Type ret_type, n->params = std::move(params); n->body = std::move(body); n->ret_type = std::move(ret_type); - n->buffer_map = std::move(buffer_map); n->attrs = std::move(attrs); n->ty = relax::FuncType::OpaqueFunc(); n->span = std::move(span); @@ -98,6 +111,11 @@ PrimFunc::PrimFunc(ffi::Array params, Stmt body, Type ret_type, (*this)->ty = InferType(*this); } +PrimFunc::PrimFunc(ffi::Array params, Stmt body, Type ret_type, + ffi::Map buffer_map, DictAttrs attrs, Span span) + : PrimFunc(NormalizeBufferParams(std::move(params), buffer_map), std::move(body), + std::move(ret_type), std::move(attrs), std::move(span)) {} + FuncType PrimFuncNode::func_type_annotation() const { ffi::Array param_types; for (auto param : this->params) { @@ -121,7 +139,9 @@ TensorIntrin::TensorIntrin(PrimFunc desc, PrimFunc impl) { TVM_FFI_CHECK_EQ(desc->params.size(), impl->params.size(), ValueError) << "The number of parameters of the description and the implementation of the " "tensor intrinsic doesn't match."; - auto is_handle = [](const Var& param) { return param->ty.as() != nullptr; }; + auto is_handle = [](const Var& param) { + return param->ty.as() != nullptr || param->ty.as() != nullptr; + }; for (size_t i = 0; i < desc->params.size(); i++) { TVM_FFI_CHECK(is_handle(desc->params[i]), ValueError) << "Parameters of the description of the " @@ -130,8 +150,6 @@ TensorIntrin::TensorIntrin(PrimFunc desc, PrimFunc impl) { << "Parameters of the implementation of " "the tensor intrinsic should be handle only."; } - TVM_FFI_ICHECK_EQ(desc->buffer_map.size(), impl->buffer_map.size()); - ffi::ObjectPtr n = ffi::make_object(); n->desc = std::move(desc); n->impl = std::move(impl); @@ -164,9 +182,9 @@ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() .def("tirx.PrimFunc", - [](ffi::Array params, Stmt body, Type ret_type, - ffi::Map buffer_map, DictAttrs attrs, - Span span) { return PrimFunc(params, body, ret_type, buffer_map, attrs, span); }) + [](ffi::Array params, Stmt body, Type ret_type, DictAttrs attrs, Span span) { + return PrimFunc(params, body, ret_type, attrs, span); + }) .def("tirx.TensorIntrin", [](PrimFunc desc_func, PrimFunc intrin_func) { return TensorIntrin(desc_func, intrin_func); diff --git a/src/tirx/ir/script/script_complete.cc b/src/tirx/ir/script/script_complete.cc index ae128f20ad08..1d4ffdb7ef9e 100644 --- a/src/tirx/ir/script/script_complete.cc +++ b/src/tirx/ir/script/script_complete.cc @@ -126,7 +126,7 @@ class ScriptCompleter : public StmtMutator { PrimFunc ScriptComplete(PrimFunc func, const ffi::Array& root_allocates, bool s_tir) { ffi::Map buffer_var_map; - for (const auto& pair : func->buffer_map) { + for (const auto& pair : tirx::BufferParamMap(func->params)) { const BufferVar& buffer = pair.second; buffer_var_map.Set(buffer.var(), buffer); } diff --git a/src/tirx/ir/specialize.cc b/src/tirx/ir/specialize.cc index 687b37dd3a78..9aec2a325341 100644 --- a/src/tirx/ir/specialize.cc +++ b/src/tirx/ir/specialize.cc @@ -31,6 +31,7 @@ #include #include +#include #include "../transform/ir_utils.h" #include "functor_common.h" @@ -78,10 +79,27 @@ class PrimFuncSpecializer : public StmtExprMutator { static PrimFunc Specialize(PrimFunc f, const VarMap& var_map) { PrimFuncSpecializer specializer(var_map); + for (const Var& param : f->params) { + auto buffer = tirx::AsBufferVar(param); + auto replacement = var_map.find(param); + if (!buffer || replacement == var_map.end()) { + continue; + } + if (auto replacement_var = replacement->second.as()) { + if (auto replacement_buffer = tirx::AsBufferVar(replacement_var.value())) { + if (IsParam(f, replacement_var.value())) { + specializer.buffer_aliases_[buffer.value()] = replacement_buffer.value(); + } else { + specializer.constrained_buffer_params_.insert(param.get()); + } + } + } + } + // Updating BufferVar map ffi::Map buffer_map; bool buffer_map_updated = false; - for (const auto& it : f->buffer_map) { + for (const auto& it : tirx::BufferParamMap(f->params)) { const Var& var = it.first; const BufferVar& buffer = it.second; BufferVar new_buffer = specializer.MutateBuffer(buffer); @@ -92,13 +110,18 @@ class PrimFuncSpecializer : public StmtExprMutator { } } - // Updating parmeters + // Updating parameters ffi::Array params; bool param_updated = false; for (const auto& var : f->params) { // Remove parmeters which has been specialized. - if (var_map.find(var) == var_map.end()) { - params.push_back(var); + if (var_map.find(var) == var_map.end() || + specializer.constrained_buffer_params_.count(var.get())) { + if (auto buffer = buffer_map.Get(var)) { + params.push_back(buffer.value().var()); + } else { + params.push_back(var); + } } else { param_updated = true; } @@ -108,7 +131,7 @@ class PrimFuncSpecializer : public StmtExprMutator { Stmt body = specializer(f->body); if (param_updated || buffer_map_updated || !f->body.same_as(body)) { - return PrimFunc(params, body, f->ret_type, buffer_map, f->attrs, f->span); + return PrimFunc(params, body, f->ret_type, f->attrs, f->span); } else { return f; } @@ -160,6 +183,9 @@ class PrimFuncSpecializer : public StmtExprMutator { BufferVar VisitBufferUse(const BufferVar& buffer) final { return GetNewBuffer(buffer); } Expr VisitExpr_(const VarNode* op) final { + if (constrained_buffer_params_.count(op)) { + return ffi::GetRef(op); + } auto it = var_map_.find(ffi::GetRef(op)); if (it == var_map_.end()) { return ffi::GetRef(op); @@ -189,6 +215,10 @@ class PrimFuncSpecializer : public StmtExprMutator { private: BufferVar MutateBuffer(const BufferVar& buffer) { + if (auto it = buffer_aliases_.find(buffer); it != buffer_aliases_.end()) { + return it->second; + } + ffi::Optional specialized_storage_scope; if (auto it = var_map_.find(buffer.var()); it != var_map_.end()) { if (const auto* new_var = it->second.as()) { @@ -283,7 +313,7 @@ class PrimFuncSpecializer : public StmtExprMutator { << "mutation must occur at the buffer's point of definition " << "(see discussion on https://github.com/apache/tvm/pull/14565 for more details). " << "Please add a definition for this buffer, " - << "either in the PrimFunc's buffer_map, " + << "either as a BufferType-annotated PrimFunc parameter, " << "in a tirx::SBlock's alloc_buffer, " << "or in a DeclBuffer statement."; @@ -307,6 +337,10 @@ class PrimFuncSpecializer : public StmtExprMutator { const VarMap& var_map_; /*! \brief map from old buffer to mutated buffer */ std::unordered_map buffer_map_; + /*! \brief Direct aliases between buffer parameters. */ + std::unordered_map buffer_aliases_; + /*! \brief Buffer parameters constrained by a concrete, non-parameter buffer. */ + std::unordered_set constrained_buffer_params_; }; /*! @@ -317,7 +351,7 @@ class PrimFuncSpecializer : public StmtExprMutator { * \param var_map The var mapping to be updated. * \note This function will match target buffer's shape, strides and element_offset * For example, we define a buffer in PrimFunc: - * A = T.match_buffer(a, [m, n]) + * A: T.Buffer([m, n]) * * Then we match it with a buffer B = tirx.decl_buffer((8, 16)) * @@ -331,10 +365,10 @@ void UpdateSpecializeVarMap(const PrimFunc& func, const Var& param, const Buffer // preliminaries tirx::ExprDeepEqual equal; - auto it = func->buffer_map.find(param); - TVM_FFI_CHECK(it != func->buffer_map.end(), ValueError) - << "specialize expects param to be in PrimFunc's buffer_map"; - const BufferVar& buf_to_specialize = (*it).second; + auto opt_buffer = tirx::AsBufferVar(param); + TVM_FFI_CHECK(opt_buffer, ValueError) + << "specialize expects param to have a BufferType annotation"; + const BufferVar& buf_to_specialize = opt_buffer.value(); // build var mapping using specific_buf's parameters auto build_var_mapping = [&](const Expr& new_expr, const Expr& old_expr) { @@ -379,9 +413,8 @@ void UpdateSpecializeVarMap(const PrimFunc& func, const Var& param, const Buffer build_var_mapping(specific_buf->strides[i], buf_to_specialize->strides[i]); } build_var_mapping(specific_buf->elem_offset, buf_to_specialize->elem_offset); - // The buffer identity owns the pointer projection instead of storing a - // separate data Var. Remap the typed Var itself so buffer_data uses retain - // the alias relationship established by specialization. + // The specializer distinguishes a concrete buffer constraint (which keeps + // the parameter) from another function parameter (which aliases it). build_var_mapping(specific_buf.var(), buf_to_specialize.var()); // Check data_alignment and offset_factor. @@ -407,9 +440,9 @@ void UpdateSpecializeVarMap(const PrimFunc& func, const Var& param, const Expr& // check param is in PrimFunc's parameters TVM_FFI_CHECK(IsParam(func, param), ValueError) << "Specialize expects param to be in PrimFunc's params"; - // specialize a param not in buffer_map - TVM_FFI_CHECK_EQ(func->buffer_map.count(param), 0, ValueError) - << "Specialize expects param to not be in PrimFunc's buffer_map"; + // Specialize a scalar parameter rather than a buffer parameter. + TVM_FFI_CHECK(!tirx::AsBufferVar(param), ValueError) + << "Specialize expects param to not have a BufferType annotation"; // build var mapping using specific_expr (*var_map)[param] = specific_expr; } diff --git a/src/tirx/ir/tir_visitor_with_path.cc b/src/tirx/ir/tir_visitor_with_path.cc index dd9bcc88d15f..4223f51d9ac5 100644 --- a/src/tirx/ir/tir_visitor_with_path.cc +++ b/src/tirx/ir/tir_visitor_with_path.cc @@ -77,28 +77,22 @@ void TIRVisitorWithPath::Visit(const IRModule& mod, AccessPath path) { } void TIRVisitorWithPath::Visit(const PrimFunc& func, AccessPath path) { - // The implicit definitions from a PrimFunc::buffer_map are pretty - // weird. They only apply if no previous definition of that - // variable has occurred. Therefore, to ensure that we only avoid - // duplicate calls to VisitVarDef, these semantics need to be - // checked. + // BufferType metadata may introduce symbolic dimensions. Define those + // symbols before entering the buffer parameter itself. std::vector, DefContext>> context; auto ppath = path->Attr("params"); for (size_t i = 0; i < func->params.size(); i++) { const Var& param = func->params[i]; - if (param->ty.as()) { - context.push_back(WithDef(BufferVar(param), ppath->ArrayItem(i))); - } else { + if (!param->ty.as()) { context.push_back(WithDef(param, ppath->ArrayItem(i))); } } - auto buffer_map_path = path->Attr("buffer_map"); for (size_t i = 0; i < func->params.size(); i++) { - if (auto opt = func->buffer_map.Get(func->params[i])) { + if (auto opt = AsBufferVar(func->params[i])) { auto buf = opt.value(); - auto buf_path = buffer_map_path->MapItem(ppath->ArrayItem(i)); + auto buf_path = ppath->ArrayItem(i)->Attr("ty"); for (auto& def : WithMatchBufferDefs(buf, buf_path)) { context.push_back(std::move(def)); @@ -109,9 +103,8 @@ void TIRVisitorWithPath::Visit(const PrimFunc& func, AccessPath path) { // Only after all the implicit definitions have been visited can we // visit the buffer definition itself. for (size_t i = 0; i < func->params.size(); i++) { - if (auto opt = func->buffer_map.Get(func->params[i])) { - auto buf_path = buffer_map_path->MapItem(ppath->ArrayItem(i)); - context.push_back(WithDef(opt.value(), buf_path)); + if (auto opt = AsBufferVar(func->params[i])) { + context.push_back(WithDef(opt.value(), ppath->ArrayItem(i))); } } diff --git a/src/tirx/script/builder/frame.cc b/src/tirx/script/builder/frame.cc index a253ee8cc3b6..640842f7cb4a 100644 --- a/src/tirx/script/builder/frame.cc +++ b/src/tirx/script/builder/frame.cc @@ -113,39 +113,41 @@ void PrimFuncFrameNode::ExitWithScope() { } // s_tir-mode normalization: drop stale default layouts (see comment on // STirBufferLayoutNormalizer above) and rewrite body references coherently. - ffi::Map effective_buffer_map = buffer_map; ffi::Array effective_root_alloc_buffers = root_alloc_buffers; tvm::tirx::Stmt body = AsStmt(stmts); - if (s_tir) { - STirBufferLayoutNormalizer normalizer; - ffi::Map new_buffer_map; - for (const auto& kv : buffer_map) { - tvm::tirx::BufferVar buf = kv.second; - if (buf->layout.has_value()) { - ffi::ObjectPtr type = tvm::tirx::CopyBufferType(buf); - type->layout = std::nullopt; - tvm::tirx::BufferVar new_buf = tvm::tirx::RebuildBufferVar(buf, std::move(type)); - normalizer.Register(buf, new_buf); - new_buffer_map.Set(kv.first, new_buf); - } else { - new_buffer_map.Set(kv.first, buf); - } + STirBufferLayoutNormalizer normalizer; + ffi::Array effective_args; + for (const tvm::tirx::Var& arg : args) { + ffi::Optional opt_buffer = buffer_map.Get(arg); + if (!opt_buffer.has_value() && arg->ty.as()) { + opt_buffer = tvm::tirx::BufferVar(arg); } - if (!normalizer.Empty()) { - body = normalizer(std::move(body)); - ffi::Array new_root_alloc_buffers; - for (const tvm::tirx::BufferVar& buf : root_alloc_buffers) { - new_root_alloc_buffers.push_back(normalizer.Lookup(buf)); - } - effective_buffer_map = std::move(new_buffer_map); - effective_root_alloc_buffers = std::move(new_root_alloc_buffers); + if (!opt_buffer.has_value()) { + effective_args.push_back(arg); + continue; + } + tvm::tirx::BufferVar buffer = opt_buffer.value(); + if (s_tir && buffer->layout.has_value()) { + ffi::ObjectPtr type = tvm::tirx::CopyBufferType(buffer); + type->layout = std::nullopt; + tvm::tirx::BufferVar new_buffer = tvm::tirx::RebuildBufferVar(buffer, std::move(type)); + normalizer.Register(buffer, new_buffer); + buffer = new_buffer; + } + effective_args.push_back(buffer.var()); + } + if (!normalizer.Empty()) { + body = normalizer(std::move(body)); + ffi::Array new_root_alloc_buffers; + for (const tvm::tirx::BufferVar& buffer : root_alloc_buffers) { + new_root_alloc_buffers.push_back(normalizer.Lookup(buffer)); } + effective_root_alloc_buffers = std::move(new_root_alloc_buffers); } tvm::tirx::PrimFunc func( - /*params=*/args, + /*params=*/effective_args, /*body=*/body, /*ret_type=*/ret_type.value_or(TupleType::Empty()), - /*buffer_map=*/effective_buffer_map, /*attrs=*/attrs.defined() ? DictAttrs(attrs) : DictAttrs(), /*span=*/tvm::Span()); func = tvm::tirx::ScriptComplete(func, effective_root_alloc_buffers, s_tir); diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc index 97e24e0f6211..1538ec6836ed 100644 --- a/src/tirx/script/builder/ir.cc +++ b/src/tirx/script/builder/ir.cc @@ -90,11 +90,7 @@ Var Arg(ffi::String name, Var var) { BufferVar Arg(ffi::String name, BufferVar buffer) { PrimFuncFrame frame = FindPrimFuncFrame("T.Arg"); details::Namer::Name(buffer, name); - // A buffer parameter is an opaque ABI handle. The buffer's data pointer - // carries the exact pointee type used within the function body. - Var handle(buffer.name() + "_handle", PointerType::VoidPointerTy()); - frame->args.push_back(handle); - frame->buffer_map.Set(handle, buffer); + frame->args.push_back(buffer.var()); return buffer; } diff --git a/src/tirx/script/printer/buffer.cc b/src/tirx/script/printer/buffer.cc index 943b21100447..37f72b2d3377 100644 --- a/src/tirx/script/printer/buffer.cc +++ b/src/tirx/script/printer/buffer.cc @@ -318,11 +318,11 @@ ExprDoc BufferDecl(const tirx::BufferVar& buffer, const ffi::String& method, ExprDoc BufferAttn(const tirx::BufferVar& buffer, const AccessPath& p, const Frame& frame, const IRDocsifier& d) { ffi::Map attrs = - BufferAttrs(buffer, p, frame, d, BufferVarDefinition::DataPointer); - ExprDoc shape = attrs.Get("shape").value(); - ExprDoc dtype = - attrs.Get("dtype").value_or(LiteralDoc::DataType(buffer->dtype->dtype, p->Attr("dtype"))); - return TIR(d, "Buffer")->Call({shape, dtype}, {}, {}); + BufferAttrs(buffer, p, frame, d, BufferVarDefinition::MatchBuffer); + if (!attrs.count("dtype")) { + attrs.Set("dtype", LiteralDoc::DataType(buffer->dtype->dtype, p->Attr("dtype"))); + } + return BufferCall(TIR(d, "Buffer"), attrs, {}); } ffi::Array BufferIndices(const ffi::Array& indices, const AccessPath& p, diff --git a/src/tirx/script/printer/function.cc b/src/tirx/script/printer/function.cc index a1eddde6f2d7..76c059aa122f 100644 --- a/src/tirx/script/printer/function.cc +++ b/src/tirx/script/printer/function.cc @@ -71,10 +71,6 @@ int CountVarOccurrence(const tirx::PrimFunc& f, const tirx::Var& v) { for (const tirx::Var& v : f->params) { counter.VisitVar(v); } - for (const auto& pair : f->buffer_map) { - counter.VisitVar(pair.first); - counter.VisitBuffer(pair.second); - } return counter.count; } @@ -87,36 +83,27 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) return obj->IsInstance() || obj->IsInstance(); }); int n_args = func->params.size(); - std::unordered_map buffer_data_counter; - for (const auto& pair : func->buffer_map) { - const tirx::VarNode* buffer_var = pair.second.get(); - if (!buffer_data_counter.count(buffer_var)) { - buffer_data_counter.insert({buffer_var, 0}); - } - ++buffer_data_counter.at(buffer_var); - } // Step 1. Handle `func->params` ffi::Array args; args.reserve(n_args); - std::unordered_set buffer_inlined; + std::unordered_map scalar_param_docs; + for (const tirx::Var& param : func->params) { + if (!param->ty.as()) { + scalar_param_docs.emplace(param.get(), DefineVar(param, *f, d)); + } + } for (int i = 0; i < n_args; ++i) { tirx::Var var = func->params[i]; AccessPath var_p = p->Attr("params")->ArrayItem(i); - if (d->cfg->syntax_sugar && CountVarOccurrence(func, var) == 2 && - func->buffer_map.count(var)) { - tirx::BufferVar buffer = func->buffer_map[var]; - bool s_tir = func->attrs->dict.count(tvm::attr::kSTir); - if (IsSimpleBuffer(buffer, s_tir) && buffer_data_counter.at(buffer.get()) == 1) { - AccessPath buffer_p = p->Attr("buffer_map")->MapItem(var); - IdDoc lhs = DefineBuffer(buffer, *f, d); - ExprDoc annotation = BufferAttn(buffer, buffer_p, *f, d); - args.push_back(AssignDoc(lhs, std::nullopt, annotation)); - buffer_inlined.insert(buffer.get()); - continue; - } + if (var->ty.as()) { + tirx::BufferVar buffer(var); + IdDoc lhs = DefineBuffer(buffer, *f, d); + ExprDoc annotation = BufferAttn(buffer, var_p->Attr("ty"), *f, d); + args.push_back(AssignDoc(lhs, std::nullopt, annotation)); + continue; } ExprDoc a = d->AsDoc(var->ty, var_p->Attr("ty")); - args.push_back(AssignDoc(DefineVar(var, *f, d), std::nullopt, a)); + args.push_back(AssignDoc(scalar_param_docs.at(var.get()), std::nullopt, a)); } // Step 2. Handle `func->attrs` if (!func->attrs->dict.empty()) { @@ -147,23 +134,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) ->Call({d->AsDoc(DictAttrs(new_attrs), p->Attr("attrs"))}))); } } - // Step 3. Handle `func->buffer_map` - for (int i = 0; i < n_args; ++i) { - tirx::Var param = func->params[i]; - if (func->buffer_map.count(param)) { - tirx::BufferVar buffer = func->buffer_map[param]; - if (buffer_inlined.count(buffer.get())) { - continue; - } - ExprDoc param_doc = args[i]->lhs; - AccessPath buffer_p = p->Attr("buffer_map")->MapItem(param); - ExprDoc lhs = DefineBuffer(buffer, *f, d); - ExprDoc rhs = BufferDecl(buffer, "match_buffer", {param_doc}, buffer_p, *f, d, - BufferVarDefinition::MatchBuffer); - (*f)->stmts.push_back(AssignDoc(lhs, rhs, std::nullopt)); - } - } - // Step 4. Handle `func->body` + // Step 3. Handle `func->body` ffi::Optional implicit_root_block = [&]() -> ffi::Optional { const tirx::SBlockRealizeNode* root_block_realize = func->body.as(); diff --git a/src/tirx/script/printer/utils.h b/src/tirx/script/printer/utils.h index 296c24cd7ae9..0c15655ad591 100644 --- a/src/tirx/script/printer/utils.h +++ b/src/tirx/script/printer/utils.h @@ -286,7 +286,7 @@ enum class BufferVarDefinition { // The data pointer is defined along with the buffer, along with any // buffer parameters (shape/stride/elem_offset) that have not // previously been defined. For example, - // `BlockNode::match_buffers`, or the `PrimFuncNode::buffer_map`. + // `BlockNode::match_buffers`, or a BufferType-annotated PrimFunc parameter. MatchBuffer, }; diff --git a/src/tirx/transform/flatten_buffer.cc b/src/tirx/transform/flatten_buffer.cc index 3dbb703f994c..b57e05f3e7fe 100644 --- a/src/tirx/transform/flatten_buffer.cc +++ b/src/tirx/transform/flatten_buffer.cc @@ -62,20 +62,19 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { static PrimFunc Flatten(PrimFunc func) { arith::Analyzer ana; auto pass = BufferFlattener(ana); - pass.MarkBufferMapShapes(func); - for (const auto& [param, buffer] : func->buffer_map) { + pass.MarkBufferParamShapes(func); + for (const auto& [param, buffer] : tirx::BufferParamMap(func->params)) { pass.extern_buffers_.insert(buffer); pass.Define(buffer); } auto body = pass.VisitStmt(func->body); - // The buffers in func->buffer_map are deliberately left - // unflattened, as they are used for validation of user-provided - // arguments. The flattened buffers used in the updated - // function body alias the argument buffers. + // Buffer parameters are deliberately left unflattened, as they are used + // for validation of user-provided arguments. The flattened buffers used + // in the updated function body alias the argument buffers. for (size_t i = func->params.size(); i > 0; i--) { auto handle = func->params[i - 1]; - if (auto opt = func->buffer_map.Get(handle)) { + if (auto opt = tirx::BufferParamMap(func->params).Get(handle)) { auto old_buf = opt.value(); if (pass.buffers_used_.count(old_buf)) { auto new_buf = pass.Lookup(old_buf).flattened; diff --git a/src/tirx/transform/force_narrow_index_to_i32.cc b/src/tirx/transform/force_narrow_index_to_i32.cc index 5a4eb54d9e17..1e7e8ff810c1 100644 --- a/src/tirx/transform/force_narrow_index_to_i32.cc +++ b/src/tirx/transform/force_narrow_index_to_i32.cc @@ -37,11 +37,10 @@ class Int32DTypeNarrower : public IndexDataTypeNormalizer { public: static PrimFunc RewriteDataType(PrimFunc func) { // Check if the integer parameter buffers have dtype other than int32. - for (auto it : func->buffer_map) { + for (auto it : tirx::BufferParamMap(func->params)) { if (it.second->dtype.MatchesCode(DLDataTypeCode::kDLInt) && it.second->dtype.bits() > 32) { - TVM_FFI_THROW(InternalError) - << "The buffer " << it.second << " in the function buffer map has dtype " - << it.second->dtype << ". The function is " << func; + TVM_FFI_THROW(InternalError) << "The buffer parameter " << it.second << " has dtype " + << it.second->dtype << ". The function is " << func; } } diff --git a/src/tirx/transform/inline_private_functions.cc b/src/tirx/transform/inline_private_functions.cc index abe5d2850a54..c6ac2f3081c5 100644 --- a/src/tirx/transform/inline_private_functions.cc +++ b/src/tirx/transform/inline_private_functions.cc @@ -115,7 +115,7 @@ bool IsInlinablePrimFunc(const GlobalVar& gvar, const PrimFunc& prim_func, // We do not currently support inlining of functions that accept // buffer arguments. - bool has_buffer_arguments = prim_func->buffer_map.size(); + bool has_buffer_arguments = tirx::BufferParamMap(prim_func->params).size(); if (has_buffer_arguments) return false; // We do not currently support inlining of schedulable TIR @@ -228,9 +228,10 @@ class PrimFuncInliner : StmtExprMutator { << callee->params << "), but is called with " << args.size() << " arguments (" << args << ")"; - TVM_FFI_ICHECK(callee->buffer_map.empty()) + TVM_FFI_ICHECK(tirx::BufferParamMap(callee->params).empty()) << "Inlining of PrimFuncs with buffer arguments is not yet supported, " - << "but callee " << gvar << " has non-empty buffer map " << callee->buffer_map; + << "but callee " << gvar << " has BufferType-annotated parameters " + << tirx::BufferParamMap(callee->params); ffi::Map> param_map; for (size_t i = 0; i < callee->params.size(); i++) { diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc index e0e4dab78e39..36e08c515048 100644 --- a/src/tirx/transform/ir_utils.cc +++ b/src/tirx/transform/ir_utils.cc @@ -110,7 +110,7 @@ class IRConvertSSA final : public StmtExprMutator { for (const auto& var : func->params) { defined_params.insert(var.get()); } - for (const auto& [var, buffer] : func->buffer_map) { + for (const auto& [var, buffer] : tirx::BufferParamMap(func->params)) { static_cast(var); // gcc 7.x bug, https://gcc.gnu.org/bugzilla/show_bug.cgi?id=81767 auto check_expr = [&](const PrimExpr& expr) { auto* var_ptr = expr.as(); @@ -134,27 +134,23 @@ class IRConvertSSA final : public StmtExprMutator { } } - // Update the buffer map, based on the redefined parameters + // Update the buffer parameters, based on the redefined parameters + bool buffer_params_changed = false; auto buffer_map = [&]() { ffi::Map buffer_map; bool made_change = false; - for (const auto& [var, buffer] : func->buffer_map) { + for (const auto& [var, buffer] : tirx::BufferParamMap(func->params)) { auto new_var = GetRemappedVar(var); - if (defined_.count(buffer.get())) { - Var new_buffer_var = MakeNewVar(buffer.var()); - PushVarRemap(buffer.var(), new_buffer_var); - } else { - defined_.insert(buffer.get()); - } auto new_buf = GetRemappedBuffer(buffer); made_change = made_change || !var.same_as(new_var) || !buffer.same_as(new_buf); buffer_map.Set(new_var, new_buf); } if (made_change) { + buffer_params_changed = true; return buffer_map; } else { - return func->buffer_map; + return tirx::BufferParamMap(func->params); } }(); @@ -183,10 +179,19 @@ class IRConvertSSA final : public StmtExprMutator { auto body = VisitStmt(func->body); + if (buffer_params_changed) { + params = params.Map([&](const Var& param) -> Var { + if (auto buffer = buffer_map.Get(param)) { + return buffer.value().var(); + } + return param; + }); + } + // If anything changed, update the returned function - if (!params.same_as(func->params) || !buffer_map.same_as(func->buffer_map) || - !attrs.same_as(func->attrs) || !body.same_as(func->body)) { - func = PrimFunc(params, body, func->ret_type, buffer_map, attrs); + if (!params.same_as(func->params) || buffer_params_changed || !attrs.same_as(func->attrs) || + !body.same_as(func->body)) { + func = PrimFunc(params, body, func->ret_type, attrs); } // Pop function-scope remaps in reverse order diff --git a/src/tirx/transform/lower_tirx_cleanup.cc b/src/tirx/transform/lower_tirx_cleanup.cc index 4bf6176495bc..0bf0eb140d03 100644 --- a/src/tirx/transform/lower_tirx_cleanup.cc +++ b/src/tirx/transform/lower_tirx_cleanup.cc @@ -404,8 +404,18 @@ Pass LowerTIRxCleanup() { auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) { Target target = ResolveTarget(f); auto* n = f.CopyOnWrite(); - std::tie(n->body, n->buffer_map) = - LayoutApplier::Flatten(n->body, n->params, n->buffer_map, target); + auto [body, buffer_map] = + LayoutApplier::Flatten(n->body, n->params, BufferParamMap(n->params), target); + ffi::Array params; + for (const Var& param : n->params) { + if (auto buffer = buffer_map.Get(param)) { + params.push_back(buffer.value().var()); + } else { + params.push_back(param); + } + } + n->body = std::move(body); + n->params = std::move(params); n->body = BufferOffsetRemover::Remove(n->body); return f; }; diff --git a/src/tirx/transform/make_packed_api.cc b/src/tirx/transform/make_packed_api.cc index 8ad6dbaa03dc..c69cbc5c9829 100644 --- a/src/tirx/transform/make_packed_api.cc +++ b/src/tirx/transform/make_packed_api.cc @@ -230,8 +230,9 @@ PrimFunc MakePackedAPI(PrimFunc func) { IntImm device_type(PrimType::Int(32), target_device_type); // Create TVMFFIABIBuilder and decode all packed args - TVMFFIABIBuilder binder(name_hint, func_ptr->params, func_ptr->buffer_map, v_packed_args, - v_num_packed_args, device_type, device_id.as_or_throw()); + TVMFFIABIBuilder binder(name_hint, func_ptr->params, BufferParamMap(func_ptr->params), + v_packed_args, v_num_packed_args, device_type, + device_id.as_or_throw()); binder.DecodeAllParams(); auto result = binder.Finalize(); @@ -279,7 +280,6 @@ PrimFunc MakePackedAPI(PrimFunc func) { << "In PrimFunc " << name_hint << " variables " << undefined << " are used, but are not passed in as API arguments"; - func_ptr->buffer_map = ffi::Map(); func_ptr->ret_type = PrimType::Int(32); // return the function. diff --git a/src/tirx/transform/split_host_device.cc b/src/tirx/transform/split_host_device.cc index cd87fc608700..334c459c8728 100644 --- a/src/tirx/transform/split_host_device.cc +++ b/src/tirx/transform/split_host_device.cc @@ -172,7 +172,7 @@ class HostDeviceSplitter : public StmtMutator { } else { std::unordered_map param_order; for (size_t i = 0; i < cur_func_->params.size(); ++i) { - param_order[cur_func_->buffer_map[cur_func_->params[i]].var()] = i; + param_order[tirx::BufferParamMap(cur_func_->params)[cur_func_->params[i]].var()] = i; } // sort by original order std::sort(params.begin(), params.end(), diff --git a/src/tirx/transform/stmt_simplify.cc b/src/tirx/transform/stmt_simplify.cc index 1be9222ec293..7875adcb95c8 100644 --- a/src/tirx/transform/stmt_simplify.cc +++ b/src/tirx/transform/stmt_simplify.cc @@ -104,7 +104,7 @@ class StmtSimplifier : public IRMutatorWithAnalyzer { analyzer->rewrite_simplify.SetEnabledExtensions(config->GetEnabledExtensions()); StmtSimplifier simplifier(analyzer, config); - simplifier.MarkBufferMapShapes(func); + simplifier.MarkBufferParamShapes(func); func.CopyOnWrite()->body = simplifier(func->body); return func; } diff --git a/src/tirx/transform/storage_rewrite.cc b/src/tirx/transform/storage_rewrite.cc index ba5e893c9093..b40732c9fac0 100644 --- a/src/tirx/transform/storage_rewrite.cc +++ b/src/tirx/transform/storage_rewrite.cc @@ -1185,8 +1185,8 @@ class StoragePlanRewriter : public StmtExprMutator { */ struct BufferVarInfo { enum DeclarationLocation { - kPrimFuncParam = (1 << 0), - kPrimFuncBufferMap = (1 << 1), + kPrimFuncBufferParam = (1 << 0), + kPrimFuncPointerParam = (1 << 1), kAllocBufferNode = (1 << 2), kLetNode = (1 << 3), kDeclBufferNode = (1 << 4), @@ -1201,8 +1201,8 @@ struct BufferVarInfo { /* The extent of the buffer. * * If multidimensional, the extent of the last dimension of the buffer. If the - * size is unknown (e.g. pointer arguments to PrimFunc with no corresponding - * entry in buffer_map), then extent is zero. + * size is unknown (e.g. pointer arguments to PrimFunc without a BufferType + * annotation), then extent is zero. */ PrimExpr extent; @@ -1274,7 +1274,7 @@ class VectorTypeAccessChecker : public StmtExprVisitor { * * @param params The parameters passed to a PrimFunc * - * @param buffer_map The buffer_map associated with a PrimFunc + * @param buffer_map A derived view of the PrimFunc's buffer parameters * * @param allow_untyped_handles If a buffer or pointer variable is * missing a type annotation, assume that it has the same underlying @@ -1294,18 +1294,21 @@ class VectorTypeAccessChecker : public StmtExprVisitor { buffer_aliases_.Set(buffer_var, buffer_var); PrimType dtype = buffer->dtype; PrimExpr extent = buffer->shape.size() ? buffer->shape[buffer->shape.size() - 1] : 0; - OnArrayDeclaration(buffer_var, dtype, extent, BufferVarInfo::kPrimFuncParam); + OnArrayDeclaration(buffer_var, dtype, extent, BufferVarInfo::kPrimFuncBufferParam); } // If a pointer parameter isn't in the buffer map, then we want to // track the parameter itself. for (Var buffer_var : params) { + if (buffer_map.count(buffer_var)) { + continue; + } if (auto buffer_type = buffer_var->ty.as()) { BufferVar buffer(buffer_var); buffer_aliases_.Set(buffer_var, buffer_var); PrimExpr extent = buffer->shape.size() ? buffer->shape[buffer->shape.size() - 1] : PrimExpr(0); - OnArrayDeclaration(buffer_var, buffer->dtype, extent, BufferVarInfo::kPrimFuncParam); + OnArrayDeclaration(buffer_var, buffer->dtype, extent, BufferVarInfo::kPrimFuncBufferParam); continue; } auto pointer_type = GetPointerType(buffer_var->ty); @@ -1313,10 +1316,10 @@ class VectorTypeAccessChecker : public StmtExprVisitor { (buffer_map.count(buffer_var) == 0)) { PrimType dtype = pointer_type.value(); PrimExpr extent = 0; - OnArrayDeclaration(buffer_var, dtype, extent, BufferVarInfo::kPrimFuncBufferMap); + OnArrayDeclaration(buffer_var, dtype, extent, BufferVarInfo::kPrimFuncPointerParam); } else if (pointer_type.has_value() && allow_untyped_pointers_ && (buffer_map.count(buffer_var) == 0)) { - OnArrayDeclaration(buffer_var, PrimType::Void(), 0, BufferVarInfo::kPrimFuncBufferMap); + OnArrayDeclaration(buffer_var, PrimType::Void(), 0, BufferVarInfo::kPrimFuncPointerParam); } } } @@ -1578,11 +1581,11 @@ class VectorTypeRewriter : public StmtExprMutator { * @param checker The VectorTypeAccessChecker that has previously read out * information from the PrimFunc * - * @param rewrite_params Whether pointer-type parameters passed into the - * function should be rewritten from scalar types to vectorized types. + * @param rewrite_buffer_params Whether BufferType-annotated parameters should + * be rewritten from scalar element types to vectorized element types. * - * @param rewrite_buffer_map Whether buffers present in the buffer_map should - * have their data variable be rewritten from scalar types to vectorized types. + * @param rewrite_pointer_params Whether pointer-typed parameters should be + * rewritten from scalar types to vectorized types. * * @param rewrite_alloc_buffer_node Whether the buffer variable associated with * AllocBufferNodes should be rewritten from scalar types to vectorized types. @@ -1594,17 +1597,17 @@ class VectorTypeRewriter : public StmtExprMutator { * should be re-written. */ VectorTypeRewriter(const std::unordered_map& info_map, - const ffi::Map& buffer_aliases, bool rewrite_params = true, - bool rewrite_buffer_map = true, bool rewrite_alloc_buffer_node = true, + const ffi::Map& buffer_aliases, bool rewrite_buffer_params = true, + bool rewrite_pointer_params = true, bool rewrite_alloc_buffer_node = true, bool rewrite_indices = true, bool rewrite_let_node = true, bool rewrite_scalar_read_to_vector_shuffle = true) : rewrite_indices_(rewrite_indices), buffer_aliases_(buffer_aliases) { int rewrite_mask = 0; - if (rewrite_params) { - rewrite_mask |= BufferVarInfo::kPrimFuncParam; + if (rewrite_buffer_params) { + rewrite_mask |= BufferVarInfo::kPrimFuncBufferParam; } - if (rewrite_buffer_map) { - rewrite_mask |= BufferVarInfo::kPrimFuncBufferMap; + if (rewrite_pointer_params) { + rewrite_mask |= BufferVarInfo::kPrimFuncPointerParam; } if (rewrite_alloc_buffer_node) { rewrite_mask |= BufferVarInfo::kAllocBufferNode; @@ -1903,17 +1906,6 @@ class VectorTypeRewriter : public StmtExprMutator { } } n->params = new_params; - - // Remap the BufferVar objects in PrimFunc::buffer_map so that the - // buffers use the new buffer variables - ffi::Map new_buffer_map; - for (const auto& pair : n->buffer_map) { - Var key = pair.first; - BufferVar old_buffer = pair.second; - BufferVar new_buffer = RemapBuffer(old_buffer); - new_buffer_map.Set(key, new_buffer); - } - n->buffer_map = new_buffer_map; } private: @@ -1941,16 +1933,17 @@ class VectorTypeRewriter : public StmtExprMutator { // Rewrite allocates, pointer parameters, and buffer map into vectorized versions // if each access into a buffer is the same vector type. PrimFunc PointerValueTypeRewrite(PrimFunc f, bool allow_untyped_pointers = false, - bool rewrite_params = true, bool rewrite_buffer_map = true, + bool rewrite_buffer_params = true, + bool rewrite_pointer_params = true, bool rewrite_alloc_buffer_node = true, bool rewrite_indices = true, bool rewrite_let_node = true, bool rewrite_scalar_read_to_vector_shuffle = true) { - VectorTypeAccessChecker checker(f->params, f->buffer_map, allow_untyped_pointers, - rewrite_scalar_read_to_vector_shuffle); + VectorTypeAccessChecker checker(f->params, tirx::BufferParamMap(f->params), + allow_untyped_pointers, rewrite_scalar_read_to_vector_shuffle); checker(f->body); - VectorTypeRewriter rewriter(checker.info_map_, checker.buffer_aliases_, rewrite_params, - rewrite_buffer_map, rewrite_alloc_buffer_node, rewrite_indices, + VectorTypeRewriter rewriter(checker.info_map_, checker.buffer_aliases_, rewrite_buffer_params, + rewrite_pointer_params, rewrite_alloc_buffer_node, rewrite_indices, rewrite_let_node, rewrite_scalar_read_to_vector_shuffle); PrimFuncNode* n = f.CopyOnWrite(); n->body = rewriter(std::move(n->body)); @@ -1980,8 +1973,9 @@ Pass StorageRewrite() { reuse_require_exact_matched_dtype = true; } auto* n = f.CopyOnWrite(); - n->body = StoragePlanRewriter().Rewrite(std::move(n->body), n->params, n->buffer_map, true, - enable_reuse, reuse_require_exact_matched_dtype); + n->body = + StoragePlanRewriter().Rewrite(std::move(n->body), n->params, BufferParamMap(n->params), + true, enable_reuse, reuse_require_exact_matched_dtype); // Parameters may not be rewritten, but internal allocations may. return PointerValueTypeRewrite(std::move(f), true, false, false, true, true, true, false); }; diff --git a/src/tirx/transform/tvm_ffi_binder.cc b/src/tirx/transform/tvm_ffi_binder.cc index 2404cd3ae20c..030075e178cc 100644 --- a/src/tirx/transform/tvm_ffi_binder.cc +++ b/src/tirx/transform/tvm_ffi_binder.cc @@ -557,6 +557,14 @@ void TVMFFIABIBuilder::DecodeParam(int param_index) { ffi::reflection::AccessPath param_path = ffi::reflection::AccessPath::Root()->Extend(AccessStep::ArrayItem(param_index)); + if (buffer_map_.count(param)) { + Var handle(param->name + ".handle", PointerType::VoidPointerTy()); + Expr handle_value = DecodeParamOpaqueHandle(param_index, type_index.as_or_throw()); + BindPointer(handle, handle_value, param_path, true); + buffer_handles_.emplace(param.get(), handle); + return; + } + if (param->ty.as()) { Expr handle_value = DecodeParamOpaqueHandle(param_index, type_index.as_or_throw()); Expr pointer_value = Call(param->ty, builtin::reinterpret(), {handle_value}); @@ -599,10 +607,11 @@ void TVMFFIABIBuilder::DecodeAllParams() { Var param = params_[i]; if (buffer_map_.count(param)) { BufferVar buffer = buffer_map_[param]; + Var handle = buffer_handles_.at(param.get()); ffi::reflection::AccessPath param_path = ffi::reflection::AccessPath::Root() ->Extend(AccessStep::ArrayItem(i)) ->Attr(ffi::String(buffer.name())); - Expr data = DecodeParamDLTensor(buffer, device_type_, device_id_, param, + Expr data = DecodeParamDLTensor(buffer, device_type_, device_id_, handle, func_name_ + "." + param->name, param_path); decl_buffers_.push_back(DeclBuffer(buffer, data)); } diff --git a/src/tirx/transform/tvm_ffi_binder.h b/src/tirx/transform/tvm_ffi_binder.h index c6b95291c49d..7c3c1e7bd3db 100644 --- a/src/tirx/transform/tvm_ffi_binder.h +++ b/src/tirx/transform/tvm_ffi_binder.h @@ -407,6 +407,8 @@ class TVMFFIABIBuilder { ffi::Array params_; /*! \brief The buffer map from parameters to buffers. */ ffi::Map buffer_map_; + /*! \brief Raw packed-ABI handles decoded for buffer-typed parameters. */ + std::unordered_map buffer_handles_; /*! \brief The packed args variable. */ Var v_packed_args_; /*! \brief The expected device type expression. */ diff --git a/src/tirx/transform/unsupported_dtype_legalize.cc b/src/tirx/transform/unsupported_dtype_legalize.cc index 1bca19fbc7a7..058c74802f06 100644 --- a/src/tirx/transform/unsupported_dtype_legalize.cc +++ b/src/tirx/transform/unsupported_dtype_legalize.cc @@ -557,7 +557,8 @@ class FP8ComputeLegalizer : public ComputeLegalizer { class StorageLegalizer : public StmtExprMutator { public: PrimFunc Legalize(PrimFunc func) { - TVM_FFI_ICHECK_EQ(func->buffer_map.size(), 0) << "This pass must be called after MakePackedAPI"; + TVM_FFI_ICHECK_EQ(tirx::BufferParamMap(func->params).size(), 0) + << "This pass must be called after MakePackedAPI"; auto* n = func.CopyOnWrite(); n->params = n->params.Map([this](Var var) { return this->RemapVarDef(var); }); n->body = this->VisitStmt(std::move(n->body)); diff --git a/tests/python/arith/test_arith_domain_touched.py b/tests/python/arith/test_arith_domain_touched.py index ed7d4a990136..a5deb7058783 100644 --- a/tests/python/arith/test_arith_domain_touched.py +++ b/tests/python/arith/test_arith_domain_touched.py @@ -34,7 +34,7 @@ def scalar_func(a: T.handle, b: T.handle): def test_domain_touched(): func = scalar_func - a, b = [func.buffer_map[var] for var in func.params] + a, b = [var for var in func.params if tvm.tirx.is_buffer_var(var)] ir = func.body a_domain_r = tvm.arith._ffi_api.DomainTouched(ir, a, True, False) @@ -82,7 +82,7 @@ def func(a: T.handle, b: T.handle, n: T.int32): for i in T.serial(n): A[i * m : (i + 1) * m : 1] = A[i * m : (i + 1) * m : 1] + B[i * m : (i + 1) * m : 1] - a, b = [func.buffer_map[var] for var in func.params[:2]] + a, b = [var for var in func.params[:2] if tvm.tirx.is_buffer_var(var)] assert tvm.arith._ffi_api.DomainTouched(func.body, a, True, False)[0].extent.value == 128 assert tvm.arith._ffi_api.DomainTouched(func.body, a, True, False)[0].extent.value == 128 diff --git a/tests/python/relax/test_analysis_well_formed.py b/tests/python/relax/test_analysis_well_formed.py index ed37b367e221..23752550c9f1 100644 --- a/tests/python/relax/test_analysis_well_formed.py +++ b/tests/python/relax/test_analysis_well_formed.py @@ -688,10 +688,10 @@ def test_pass_dltensor_arg_to_tir(): In TIR, a `DLTensor*` argument with unknown shape and dtype is represented as a `tirx.Var` with - `tvm::PointerType::VoidPointerTy()`, and with no entry in the - `PrimFuncNode::buffer_map`. In Relax, this is represented as - `R.Tensor`. Calls from Relax to TIR that pass a tensor of unknown - rank/shape are well-formed. + `tvm::PointerType::VoidPointerTy()`, rather than a buffer-typed + parameter. In Relax, this is represented as `R.Tensor`. Calls + from Relax to TIR that pass a tensor of unknown rank/shape are + well-formed. In the test case below, a TIR function accepts an arbitrary `R.Tensor`, and returns a boolean value based on inspection of the diff --git a/tests/python/relax/test_blockbuilder_core.py b/tests/python/relax/test_blockbuilder_core.py index 8df1ca822da8..b76acc1ee917 100644 --- a/tests/python/relax/test_blockbuilder_core.py +++ b/tests/python/relax/test_blockbuilder_core.py @@ -385,11 +385,10 @@ def test_call_te_unique_tensor_name(): f_matmul = bb.finalize()["matmul"] param_A = f_matmul.params[0] param_B = f_matmul.params[1] - buffer_A = f_matmul.buffer_map[param_A] - buffer_B = f_matmul.buffer_map[param_B] assert param_A.name != param_B.name - assert buffer_A.name != buffer_B.name - assert not buffer_A.same_as(buffer_B) + assert tvm.tirx.is_buffer_var(param_A) + assert tvm.tirx.is_buffer_var(param_B) + assert not param_A.same_as(param_B) def test_call_te_with_unsupported_shape_arg(): diff --git a/tests/python/relax/test_transform_specialize_primfunc_based_on_callsite.py b/tests/python/relax/test_transform_specialize_primfunc_based_on_callsite.py index 72bcf8a53d19..b9c9b16e1ea9 100644 --- a/tests/python/relax/test_transform_specialize_primfunc_based_on_callsite.py +++ b/tests/python/relax/test_transform_specialize_primfunc_based_on_callsite.py @@ -44,7 +44,9 @@ def visit_call_(self, call: relax.Call) -> None: # pylint: disable=arguments-re pfunc = self.mod[call.args[0]] if not self.is_matched: # All scopes should be global in before pass - for _, buf in pfunc.buffer_map.items(): + for buf in pfunc.params: + if not tvm.tirx.is_buffer_var(buf): + continue assert "global" == buf.data.ty.storage_scope, ( f"expected to be global scoped, but got {val.data.ty.storage_scope}" ) @@ -54,12 +56,12 @@ def visit_call_(self, call: relax.Call) -> None: # pylint: disable=arguments-re assert isinstance(arg_ty, relax.TensorType), ( f"Expected TensorType but git {type(arg_ty)}" ) - buf = pfunc.buffer_map[pfunc.params[idx]] + buf = pfunc.params[idx] assert arg_ty.vdevice.memory_scope == buf.data.ty.storage_scope, ( f"scope mismatched after specialization {arg_ty.vdevice.memory_scope} vs {buf.data.ty.storage_scope}" ) if isinstance(call.ty_args[0], relax.TensorType): - buf = pfunc.buffer_map[pfunc.params[-1]] + buf = pfunc.params[-1] assert call.ty_args[0].vdevice.memory_scope == buf.data.ty.storage_scope, ( f"scope mismatched after specialization {call.ty_args[0].vdevice.memory_scope} vs {buf.data.ty.storage_scope}" ) @@ -68,7 +70,7 @@ def visit_call_(self, call: relax.Call) -> None: # pylint: disable=arguments-re f"Expected TupleType but git {type(call.ty_args[0])}" ) for idx, ty in enumerate(call.ty_args[0].fields): - buf = pfunc.buffer_map[pfunc.params[len(call.args[1]) + idx]] + buf = pfunc.params[len(call.args[1]) + idx] assert ty.vdevice.memory_scope == buf.data.ty.storage_scope, ( f"scope mismatched after specialization {ty.vdevice.memory_scope} vs {buf.data.ty.storage_scope}" ) diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py b/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py index 9e27c2208053..acd33f9814ca 100644 --- a/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py +++ b/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py @@ -55,7 +55,7 @@ def func(A: T.Buffer(1024, "float32"), B: T.Buffer(1024, "float32")): for i in T.serial(1024): B[i] = A[i] - A, B = func.buffer_map.values() + A, B = [param for param in func.params if tvm.tirx.is_buffer_var(param)] expected = (A[0:1024], B[0:1024]) _check_memcpy_results(func, expected) @@ -117,7 +117,7 @@ def func(A: T.Buffer(1024, "float32"), B: T.Buffer((32, 32), "float32")): for i in T.serial(1024): B[i // 32, i % 32] = A[i] - A, B = func.buffer_map.values() + A, B = [param for param in func.params if tvm.tirx.is_buffer_var(param)] expected = (A[0:1024], B[0:32, 0:32]) _check_memcpy_results(func, expected) @@ -130,7 +130,7 @@ def func(A: T.Buffer((32, 32), "float32"), B: T.Buffer(1024, "float32")): for i in T.serial(1024): B[i] = A[i // 32, i % 32] - A, B = func.buffer_map.values() + A, B = [param for param in func.params if tvm.tirx.is_buffer_var(param)] expected = (A[0:32, 0:32], B[0:1024]) _check_memcpy_results(func, expected) @@ -149,7 +149,7 @@ def func(A: T.Buffer(1024, "float32"), B: T.Buffer(1024, "float32")): for i, j in T.grid(32, 32): B[i * 32 + j] = A[i * 32 + j] - A, B = func.buffer_map.values() + A, B = [param for param in func.params if tvm.tirx.is_buffer_var(param)] i = func.body.loop_var expected = [ (A[0:1024], B[0:1024]), @@ -171,7 +171,7 @@ def func(A: T.Buffer(1024, "float32"), B: T.Buffer(1024, "float32")): for i, j in T.grid(32, 32): B[i * 32 + j] = A[j + i * 32] - A, B = func.buffer_map.values() + A, B = [param for param in func.params if tvm.tirx.is_buffer_var(param)] i = func.body.loop_var expected = [ (A[0:1024], B[0:1024]), @@ -188,7 +188,7 @@ def func(A: T.Buffer(1024, "float32"), B: T.Buffer((32, 32), "float32")): for i, j in T.grid(32, 32): B[i, j] = A[i * 32 + j] - A, B = func.buffer_map.values() + A, B = [param for param in func.params if tvm.tirx.is_buffer_var(param)] i = func.body.loop_var expected = [ (A[0:1024], B[0:32, 0:32]), @@ -205,7 +205,7 @@ def func(A: T.Buffer((32, 32), "float32"), B: T.Buffer(1024, "float32")): for i, j in T.grid(32, 32): B[i * 32 + j] = A[i, j] - A, B = func.buffer_map.values() + A, B = [param for param in func.params if tvm.tirx.is_buffer_var(param)] i = func.body.loop_var expected = [ (A[0:32, 0:32], B[0:1024]), @@ -222,7 +222,7 @@ def func(A: T.Buffer((32, 32), "float32"), B: T.Buffer((32, 32), "float32")): for i, j in T.grid(32, 32): B[i, j] = A[i, j] - A, B = func.buffer_map.values() + A, B = [param for param in func.params if tvm.tirx.is_buffer_var(param)] i = func.body.loop_var expected = [ (A[0:32, 0:32], B[0:32, 0:32]), @@ -281,7 +281,7 @@ def func(A: T.Buffer((32, 32), "float32"), B: T.Buffer((32, 32), "float32")): for i, j in T.grid(32, 32): B[j, i] = A[j, i] - A, B = func.buffer_map.values() + A, B = [param for param in func.params if tvm.tirx.is_buffer_var(param)] expected = [ (A[0:32, 0:32], B[0:32, 0:32]), "Mismatch .* number of src indices touched", @@ -301,7 +301,7 @@ def func(A: T.Buffer((32, 32), "float32"), B: T.Buffer(32, "float32")): for i, j in T.grid(32, 32): B[j] = A[i, j] - A, B = func.buffer_map.values() + A, B = [param for param in func.params if tvm.tirx.is_buffer_var(param)] i = func.body.loop_var expected = [ "does not form a bijective transform", @@ -322,7 +322,7 @@ def func(A: T.Buffer(32, "float32"), B: T.Buffer((32, 32), "float32")): for i, j in T.grid(32, 32): B[i, j] = A[j] - A, B = func.buffer_map.values() + A, B = [param for param in func.params if tvm.tirx.is_buffer_var(param)] i = func.body.loop_var expected = [ "does not form a bijective transform", diff --git a/tests/python/s_tir/analysis/test_sblock_access_region.py b/tests/python/s_tir/analysis/test_sblock_access_region.py index daf94301986c..15a84fae301c 100644 --- a/tests/python/s_tir/analysis/test_sblock_access_region.py +++ b/tests/python/s_tir/analysis/test_sblock_access_region.py @@ -277,7 +277,9 @@ def test_opaque_access_with_tvm_access_ptr(): def test_decl_buffer_alias_is_not_an_opaque_access(): block = decl_buffer_alias_func.body.block - buffer_var_map = {buf: buf for buf in decl_buffer_alias_func.buffer_map.values()} + buffer_var_map = { + buf: buf for buf in decl_buffer_alias_func.params if tvm.tirx.is_buffer_var(buf) + } reads, writes, opaque = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(block.reads, reads) @@ -401,7 +403,7 @@ def func( output[vi, vs] = storage[seq_id, history_id, vs] block = func.body.block.body.body.body.block - buffer_var_map = {buf: buf for buf in func.buffer_map.values()} + buffer_var_map = {buf: buf for buf in func.params if tvm.tirx.is_buffer_var(buf)} ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(block.reads, ret[0]) tvm.ir.assert_structural_equal(block.writes, ret[1]) @@ -427,7 +429,7 @@ def func( C[vi, vs1] = A[vi1, vs2] + B[vi2, vs3] block = func.body.block.body.body.body.block - buffer_var_map = {buf: buf for buf in func.buffer_map.values()} + buffer_var_map = {buf: buf for buf in func.params if tvm.tirx.is_buffer_var(buf)} ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(block.reads, ret[0]) tvm.ir.assert_structural_equal(block.writes, ret[1]) diff --git a/tests/python/s_tir/analysis/test_sblock_buffer_access_lca.py b/tests/python/s_tir/analysis/test_sblock_buffer_access_lca.py index 87b578d8764b..08f72af25d62 100644 --- a/tests/python/s_tir/analysis/test_sblock_buffer_access_lca.py +++ b/tests/python/s_tir/analysis/test_sblock_buffer_access_lca.py @@ -14,7 +14,6 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# ruff: noqa: F401 import tvm from tvm import s_tir from tvm.script import tirx as T @@ -109,7 +108,7 @@ def global_buffer_with_blockidx( def test_buffer_load_store(): func = buffer_load_store_func - A, B = [func.buffer_map[x] for x in func.params] + A, B = [x for x in func.params if tvm.tirx.is_buffer_var(x)] C, D = func.body.block.alloc_buffers lca = s_tir.analysis.detect_buffer_access_lca(func) @@ -133,7 +132,7 @@ def test_buffer_load_store(): def test_opaque_access(): func = buffer_opaque_access - B, C = [func.buffer_map[x] for x in func.params] + B, C = [x for x in func.params if tvm.tirx.is_buffer_var(x)] lca = s_tir.analysis.detect_buffer_access_lca(func) # Cannot detect buffer A since it is define by low-level Allocate @@ -148,14 +147,14 @@ def test_opaque_access(): def test_lca_func_root(): func = lca_is_func_root - (A,) = [func.buffer_map[x] for x in func.params] + (A,) = [x for x in func.params if tvm.tirx.is_buffer_var(x)] lca = s_tir.analysis.detect_buffer_access_lca(func) assert lca[A] is None def test_match_buffer(): func = match_buffer_func - A, B = [func.buffer_map[x] for x in func.params] + A, B = [x for x in func.params if tvm.tirx.is_buffer_var(x)] lca = s_tir.analysis.detect_buffer_access_lca(func) root_block = func.body.block @@ -171,7 +170,7 @@ def test_match_buffer(): def test_global_buffer_with_blockidx(): func = global_buffer_with_blockidx - A, B = [func.buffer_map[x] for x in func.params] + A, B = [x for x in func.params if tvm.tirx.is_buffer_var(x)] lca = s_tir.analysis.detect_buffer_access_lca(func) root_block = func.body.block diff --git a/tests/python/s_tir/base/test_tir_te_extern_primfunc.py b/tests/python/s_tir/base/test_tir_te_extern_primfunc.py index d1cd710ed5bc..be10c72262e6 100644 --- a/tests/python/s_tir/base/test_tir_te_extern_primfunc.py +++ b/tests/python/s_tir/base/test_tir_te_extern_primfunc.py @@ -192,7 +192,9 @@ def test_te_extern_call(self, func, params, verify): ir_mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main")) prim_func = ir_mod["main"] - buf_name_map = {buf.name: buf for buf in func.buffer_map.values()} + buf_name_map = { + param.name: param for param in func.params if tvm.tirx.is_buffer_var(param) + } input_tensors = [te.placeholder(buf_name_map[name].shape) for name in params] output = te.extern_primfunc(input_tensors, prim_func) rt_prim_func = te.create_prim_func(tensors_from_extern_op(output, prim_func)) @@ -219,8 +221,8 @@ def tensors_from_extern_op(extern, func): buffer_to_tensor = {**input_binds, **output_binds} ordered_tensors = [] for var in func.params: - buf = func.buffer_map[var] - ordered_tensors.append(buffer_to_tensor[buf]) + if tvm.tirx.is_buffer_var(var): + ordered_tensors.append(buffer_to_tensor[var]) return ordered_tensors diff --git a/tests/python/s_tir/schedule/test_tir_schedule_decompose_padding.py b/tests/python/s_tir/schedule/test_tir_schedule_decompose_padding.py index 29d879dce266..e2a94759b579 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_decompose_padding.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_decompose_padding.py @@ -31,8 +31,8 @@ def check_decompose_padding(origin, scheduled, expected, check_run=False): assert_structural_equal_ignore_global_symbol(scheduled, expected) if check_run: - in_buffer = origin.buffer_map[origin.params[0]] - out_buffer = origin.buffer_map[origin.params[1]] + in_buffer = origin.params[0] + out_buffer = origin.params[1] in_shape = [int(_) for _ in in_buffer.shape] out_shape = [int(_) for _ in out_buffer.shape] x = tvm.runtime.tensor(np.random.uniform(0, 64, in_shape).astype(in_buffer.dtype)) diff --git a/tests/python/s_tir/schedule/test_tir_schedule_rolling_buffer.py b/tests/python/s_tir/schedule/test_tir_schedule_rolling_buffer.py index e04576e48d73..dcb807fa4cee 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_rolling_buffer.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_rolling_buffer.py @@ -35,8 +35,8 @@ def check_rolling_buffer( assert_structural_equal_ignore_global_symbol(scheduled, expected) verify_trace_roundtrip(sch, origin) if check_run: - in_buffer = origin.buffer_map[origin.params[0]] - out_buffer = origin.buffer_map[origin.params[1]] + in_buffer = origin.params[0] + out_buffer = origin.params[1] in_shape = [int(_) for _ in in_buffer.shape] out_shape = [int(_) for _ in out_buffer.shape] x = tvm.runtime.tensor(np.random.uniform(0, 64, in_shape).astype(in_buffer.dtype)) diff --git a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py index 5a6a97f618c3..4252f88f72b4 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py @@ -304,7 +304,9 @@ def ref(B: T.Buffer((8, 8, 16, 16), "float32"), C: T.Buffer((128, 128), "float32 # T.reads(B[vi // 16 + vi_o, vj // 16 + vj_o, vi % 16, vj % 16]) # C[...] = B[vi // 16 + vi_o, vj // 16 + vj_o, vi % 16, vj % 16] + T.float32(1) - expected = tvm.tirx.PrimFunc(list(ref.buffer_map.values()), ref.body.block.body) + expected = tvm.tirx.PrimFunc( + [param for param in ref.params if tvm.tirx.is_buffer_var(param)], ref.body.block.body + ) actual_block = sch.get(block_outer) actual = tvm.tirx.PrimFunc( [actual_block.reads[0].buffer, actual_block.writes[0].buffer], @@ -1176,7 +1178,7 @@ def after(a: T.handle, b: T.handle, c: T.handle): # pylint: enable=invalid-name,line-too-long,too-many-locals # fmt: on # pylint: disable=invalid-name - _, _, n, _ = before.buffer_map[before.params[1]].ty.shape + _, _, n, _ = before.params[1].ty.shape sch = tvm.s_tir.Schedule(before) block = sch.get_sblock("NT_matmul") sch.transform_layout( @@ -1226,7 +1228,7 @@ def after(a: T.handle, b: T.handle, c: T.handle): # pylint: enable=invalid-name,line-too-long,too-many-locals # fmt: on # pylint: disable=invalid-name - _, _, n, _ = before.buffer_map[before.params[1]].ty.shape + _, _, n, _ = before.params[1].ty.shape sch = tvm.s_tir.Schedule(before) block = sch.get_sblock("NT_matmul") sch.transform_block_layout( diff --git a/tests/python/s_tir/test_s_tir_renew_defs.py b/tests/python/s_tir/test_s_tir_renew_defs.py index 82f0109150f1..ff3b418e3313 100644 --- a/tests/python/s_tir/test_s_tir_renew_defs.py +++ b/tests/python/s_tir/test_s_tir_renew_defs.py @@ -27,7 +27,7 @@ def _check_func_signature_remap(lhs: PrimFunc, rhs: PrimFunc): assert lhs != rhs for x, y in zip(lhs.params, rhs.params): assert x != y - assert lhs.buffer_map[x] != rhs.buffer_map[y] + assert tvm.tirx.is_buffer_var(x) == tvm.tirx.is_buffer_var(y) def _check_buffer_decl(lhs: Buffer, rhs: Buffer): @@ -169,7 +169,7 @@ def symbolic_func(a: T.handle, b: T.handle, n: T.int32): tvm.ir.assert_structural_equal(f1, f2) -def test_buffer_map(): +def test_buffer_params(): @T.prim_func(s_tir=True) def main(a: T.handle, b: T.handle): m = T.int64() @@ -183,7 +183,7 @@ def main(a: T.handle, b: T.handle): f1 = main f2 = tvm.s_tir.renew_defs(main) tvm.ir.assert_structural_equal(f1, f2) - assert f1.buffer_map[f1.params[1]].shape[0] != f2.buffer_map[f2.params[1]].shape[0] + assert f1.params[1].shape[0] != f2.params[1].shape[0] def test_gather(): diff --git a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py index 29a7d9cace20..45f33a65fd0b 100644 --- a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py +++ b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py @@ -376,8 +376,8 @@ def func(): assert "was re-defined at" in error_msg -def test_buffer_in_buffer_map_is_well_formed(): - """Buffers defined via function parameter buffer_map are in scope for the body.""" +def test_buffer_param_is_well_formed(): + """BufferType-annotated parameters are in scope for the body.""" @T.prim_func(s_tir=True) def func(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")): @@ -444,7 +444,7 @@ def test_error_undeclared_buffer_in_schedulable_tir(): i = tvm.tirx.Var("i", "int32") # Create an undeclared buffer using an explicit data pointer that is NOT - # in the buffer_map and NOT wrapped with DeclBuffer. + # a function parameter and NOT wrapped with DeclBuffer. B_data = tvm.tirx.Var("B_data", tvm.ir.PointerType(tvm.ir.PrimType("float32"))) B = tvm.tirx.decl_buffer([n], "float32", name="B", data=B_data) @@ -466,7 +466,7 @@ def test_error_undeclared_buffer_in_schedulable_tir(): prim_func = tvm.tirx.PrimFunc( params=[A, B_data], body=tvm.tirx.For(i, 0, n, tvm.tirx.ForKind.SERIAL, block_realize), - # Note: B is NOT in buffer_map, so its declaration scope is only + # Note: B is NOT a function parameter, so its declaration scope is only # within a DeclBuffer node (which we intentionally omit here). ) diff --git a/tests/python/tirx-base/test_tir_nodes.py b/tests/python/tirx-base/test_tir_nodes.py index 48b92eb4850b..4af595a78375 100644 --- a/tests/python/tirx-base/test_tir_nodes.py +++ b/tests/python/tirx-base/test_tir_nodes.py @@ -309,9 +309,11 @@ def test_prim_func(): func = tvm.tirx.PrimFunc([x, y, b], stmt) # make sure we can print - assert func.buffer_map[func.params[2]].same_as(b) + assert func.params[2].same_as(b) + assert not hasattr(func, "buffer_map") - assert len(func.buffer_map) == 1 + assert sum(tvm.tirx.is_buffer_var(param) for param in func.params) == 1 + assert func.with_body(tvm.tirx.Evaluate(0)).params[2].same_as(b) f2 = func.with_attr({"calling_conv": 1, "tirx.noalias": True}) assert f2.attrs["calling_conv"] == 1 assert not func.attrs diff --git a/tests/python/tirx-base/test_tir_specialize.py b/tests/python/tirx-base/test_tir_specialize.py index efbe9babd6b6..f47a6dc591ae 100644 --- a/tests/python/tirx-base/test_tir_specialize.py +++ b/tests/python/tirx-base/test_tir_specialize.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=missing-function-docstring, missing-module-docstring -# ruff: noqa: F401, F841 +# ruff: noqa: F401 import pytest @@ -194,7 +194,7 @@ def test_specialize_matmul(): def test_specialize_elemwise(): a, c = element_wise.params - C = element_wise.buffer_map[c] + C = c # fully specialized func = element_wise.specialize({a: tvm.tirx.decl_buffer((128, 64))}) assert_structural_equal_ignore_global_symbol(func, element_wise_128_64) @@ -280,17 +280,14 @@ def before(A: T.Buffer([16, 16], "float32"), B: T.Buffer([16, 16], "float32")): for i in range(256): B_flat[i] = A_flat[i] * 2.0 - # well-formed checker complains about multiple nested definitions of B_flat - # since it appears in the buffer map twice - @T.prim_func(private=True, check_well_formed=False, s_tir=True) - def expected(A: T.Buffer([16, 16], "float32"), B_handle: T.handle): - B = T.match_buffer(B_handle, [16, 16], "float32", data=A.data) + @T.prim_func(private=True, s_tir=True) + def expected(A: T.Buffer([16, 16], "float32")): A_flat = T.decl_buffer([256], "float32", data=A.data) B_flat = T.decl_buffer([256], "float32", data=A.data) for i in range(256): B_flat[i] = A_flat[i] * 2.0 - A = before.buffer_map[before.params[0]] + A = before.params[0] B_handle = before.params[1] param_map = {B_handle: A} after = before.specialize(param_map) diff --git a/tests/python/tirx-transform/test_tir_transform_convert_ssa.py b/tests/python/tirx-transform/test_tir_transform_convert_ssa.py index eb7166a7fb14..639f7bf156f7 100644 --- a/tests/python/tirx-transform/test_tir_transform_convert_ssa.py +++ b/tests/python/tirx-transform/test_tir_transform_convert_ssa.py @@ -198,7 +198,7 @@ def func_b(a: T.handle("float32")): def test_reused_buffer_parameter(): - """De-duplicate buffer_map across entire module""" + """De-duplicate buffer parameters across the entire module.""" @T.prim_func(private=True, s_tir=True) def func(A: T.Buffer(1, "float32")): @@ -490,32 +490,23 @@ def test_track_forward_declarations_in_attr_stmt(): stmt, ) - A_handle = tirx.Var("A_handle", "handle") - B_handle = tirx.Var("B_handle", "handle") - - before = tirx.PrimFunc( - [A_handle, B_handle], - stmt, - buffer_map={A_handle: A, B_handle: B}, - ) + before = tirx.PrimFunc([A, B], stmt) mod = tvm.IRModule.from_expr(before) after = tvm.tirx.transform.ConvertSSA()(mod) tvm.ir.assert_structural_equal(after["main"], before) -def test_shared_shape_var_in_buffer_map_and_alloc_buffer(): - """Shape var shared across buffer_map entries and AllocBuffer should not be renamed. +def test_shared_shape_var_in_buffer_params_and_alloc_buffer(): + """Shape var shared across buffer params and AllocBuffer should not be renamed. - When the same Var (e.g., `n`) appears in multiple buffer_map - entries (A and B both have shape [n]), ConvertSSA should not treat + When the same Var (e.g., `n`) appears in multiple buffer parameter + annotations (A and B both have shape [n]), ConvertSSA should not treat the second occurrence as a redefinition. All uses of `n` in the function body (including AllocBuffer shapes) must remain the same Var object so that MakePackedAPI can bind it from the DLTensor shape. """ n = tirx.Var("n", "int32") - A_handle = tirx.Var("A_handle", "handle") - B_handle = tirx.Var("B_handle", "handle") A = tirx.decl_buffer((n,), "float32", "A") B = tirx.decl_buffer((n,), "float32", "B") @@ -523,11 +514,7 @@ def test_shared_shape_var_in_buffer_map_and_alloc_buffer(): C = tirx.decl_buffer((n,), "float32", "C") body = tirx.SeqStmt([tirx.AllocBuffer(C), tirx.Evaluate(1)]) - before = tirx.PrimFunc( - [A_handle, B_handle], - body, - buffer_map={A_handle: A, B_handle: B}, - ) + before = tirx.PrimFunc([A, B], body) mod = tvm.IRModule.from_expr(before) after = tvm.tirx.transform.ConvertSSA()(mod) diff --git a/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py b/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py index 8bf1965a6434..af5b71941713 100644 --- a/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py +++ b/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# ruff: noqa: F811, F841 +# ruff: noqa: F841 import pytest import tvm @@ -205,7 +205,7 @@ def expected(A: T.Buffer((128,), "int16"), B: T.Buffer((128,), "int16")): tvm.ir.assert_structural_equal(after, expected) -def test_fail_on_buffer_map(): +def test_fail_on_buffer_param(): @T.prim_func(private=True, s_tir=True) def func(A: T.Buffer((128,), "int64"), B: T.Buffer((128,), "int64")): for i in T.serial(0, 16): @@ -219,7 +219,7 @@ def func(A: T.Buffer((128,), "int64"), B: T.Buffer((128,), "int64")): tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"] -def test_fail_on_buffer_map(): +def test_fail_on_internal_buffer(): @T.prim_func(private=True, s_tir=True) def func(A: T.Buffer((128,), "int32"), B: T.Buffer((128,), "int32")): C = T.sblock_alloc_buffer((128,), "int64") diff --git a/tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py b/tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py index c11fdf402444..2b15b48471cf 100644 --- a/tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py +++ b/tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py @@ -157,7 +157,7 @@ def main(A: T.Buffer((16,), "float32")): After = transform(Before) assert tvm.tirx.analysis.verify_well_formed(After) func = After["main"] - assert next(iter(func.buffer_map.values())).ty.dtype == tvm.ir.PrimType("float32x4") + assert func.params[0].ty.dtype == tvm.ir.PrimType("float32x4") decl_buffers = [] buffer_stores = [] diff --git a/tests/python/tirx/test_jit.py b/tests/python/tirx/test_jit.py index d3cd32d13480..ca9a91846f32 100644 --- a/tests/python/tirx/test_jit.py +++ b/tests/python/tirx/test_jit.py @@ -247,7 +247,7 @@ def k( condition, candidate = op_call.config["src_selector"][0] assert isinstance(condition, tvm.tirx.LT) assert int(condition.b) == 4 - assert candidate.same_as(specialized.buffer_map[specialized.params[1]]) + assert candidate.same_as(specialized.params[1]) def test_optional_param_present_and_absent_ir(): @@ -275,10 +275,10 @@ def expected_absent(out_h: T.handle): absent = kernel.specialize(a=None) assert_structural_equal(present, expected_present, map_free_vars=True) assert_structural_equal(absent, expected_absent, map_free_vars=True) - assert [param.name for param in present.params] == ["a", "out_h"] - assert [param.name for param in absent.params] == ["out_h"] - assert {param.name for param in present.buffer_map} == {"a", "out_h"} - assert {param.name for param in absent.buffer_map} == {"out_h"} + assert [param.name for param in present.params] == ["A", "out"] + assert [param.name for param in absent.params] == ["out"] + assert all(tvm.tirx.is_buffer_var(param) for param in present.params) + assert all(tvm.tirx.is_buffer_var(param) for param in absent.params) def test_optional_specialization_cache_includes_presence(): @@ -318,22 +318,22 @@ def kernel( out[0] = out[0] + B[0] assert [param.name for param in kernel.specialize().params] == [ - "first_h", - "a", + "first", + "A", "scale", - "b", - "out_h", + "B", + "out", ] assert [param.name for param in kernel.specialize(a=None).params] == [ - "first_h", + "first", "scale", - "b", - "out_h", + "B", + "out", ] assert [param.name for param in kernel.specialize(a=None, b=None).params] == [ - "first_h", + "first", "scale", - "out_h", + "out", ] @@ -384,8 +384,8 @@ def kernel(a: T.Optional(T.handle), out_h: T.handle): absent = kernel.specialize(a=None) assert len(present.params) == 2 assert len(absent.params) == 1 - assert len(present.buffer_map) == 1 - assert len(absent.buffer_map) == 1 + assert sum(tvm.tirx.is_buffer_var(param) for param in present.params) == 1 + assert sum(tvm.tirx.is_buffer_var(param) for param in absent.params) == 1 def test_compile_time_bool_ops_and_if_expression_short_circuit(): @@ -402,7 +402,7 @@ def kernel(a: T.Optional(T.handle), out_h: T.handle): out[0] = 3 if a is None else fail_if_evaluated() absent = kernel.specialize(a=None) - assert [param.name for param in absent.params] == ["out_h"] + assert [param.name for param in absent.params] == ["out"] def test_runtime_tir_if_cannot_guard_absent_optional_param(): diff --git a/tests/python/tirx/transform/test_stmt_functor.py b/tests/python/tirx/transform/test_stmt_functor.py index 8d602cce3560..ae7e1aa8fa00 100644 --- a/tests/python/tirx/transform/test_stmt_functor.py +++ b/tests/python/tirx/transform/test_stmt_functor.py @@ -1204,7 +1204,7 @@ def selector( op_call = selector.body assert isinstance(op_call, tir.TilePrimitiveCall) - original_b = selector.buffer_map[selector.params[1]] + original_b = selector.params[1] seen = [] post_order_visit(selector.body, seen.append) assert any(isinstance(node, Var) and node.same_as(selector.params[-1]) for node in seen) diff --git a/tests/python/tvmscript/test_tvmscript_complete.py b/tests/python/tvmscript/test_tvmscript_complete.py index 9d56f4bdfd60..b4b9736c9bb5 100644 --- a/tests/python/tvmscript/test_tvmscript_complete.py +++ b/tests/python/tvmscript/test_tvmscript_complete.py @@ -110,7 +110,7 @@ def func_with_part_access_region(a: T.handle, b: T.handle, c: T.handle) -> None: def test_complete_matmul(): func = matmul - A, B, C = [func.buffer_map[x] for x in func.params] + A, B, C = [x for x in func.params if tvm.tirx.is_buffer_var(x)] block = func.body.block.body.body.body.body.block assert isinstance(block, tvm.tirx.SBlock) @@ -130,7 +130,7 @@ def test_complete_matmul(): def test_complete_matmul_original(): func = matmul_original - A, B, C = [func.buffer_map[x] for x in func.params] + A, B, C = [x for x in func.params if tvm.tirx.is_buffer_var(x)] block1 = func.body.block.body.body.body[0].block assert isinstance(block1, tvm.tirx.SBlock) @@ -158,7 +158,7 @@ def test_complete_matmul_original(): def _check_elementwise(func): - A, B, C = [func.buffer_map[x] for x in func.params] + A, B, C = [x for x in func.params if tvm.tirx.is_buffer_var(x)] root_block = func.body.block assert len(root_block.reads) == 0 diff --git a/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py b/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py index f5ace5999ac4..dd02640f30d3 100644 --- a/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py +++ b/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py @@ -43,7 +43,6 @@ def test_ir_builder_tir_primfunc_base(): params=[], body=tirx.Evaluate(0), ret_type=None, - buffer_map=None, attrs=tvm.ir.make_node("ir.DictAttrs", s_tir=True), ) @@ -80,29 +79,19 @@ def test_ir_builder_tir_primfunc_complete(): prim_func_actual = ib.get() # the expected prim_func - c_handle, c_buffer = ( - tirx.Var("c_handle", tvm.ir.PointerType(tvm.ir.PrimType("void"))), - tirx.decl_buffer((128, 128), "float32", name="c", layout=None), - ) - d_handle, d_buffer = ( - tirx.Var("d", tvm.ir.PointerType(tvm.ir.PrimType("void"))), - tirx.decl_buffer((64, 64), "int64", name="d", layout=None), - ) - e_handle, e_buffer = ( - tirx.Var("e_handle", tvm.ir.PointerType(tvm.ir.PrimType("void"))), - tirx.decl_buffer((1024,), "int8", name="e", layout=None), - ) + c_buffer = tirx.decl_buffer((128, 128), "float32", name="c", layout=None) + d_buffer = tirx.decl_buffer((64, 64), "int64", name="d", layout=None) + e_buffer = tirx.decl_buffer((1024,), "int8", name="e", layout=None) prim_func_expected = tirx.PrimFunc( params=[ tirx.Var("a", tvm.ir.PointerType(tvm.ir.PrimType("void"))), tirx.Var("b", "int64"), - c_handle, - d_handle, - e_handle, + c_buffer, + d_buffer, + e_buffer, ], body=tirx.Evaluate(0), ret_type=tvm.ir.PrimType("int64"), - buffer_map={c_handle: c_buffer, d_handle: d_buffer, e_handle: e_buffer}, attrs=tvm.ir.make_node("ir.DictAttrs", key="value", s_tir=True), ) diff --git a/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py b/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py index f25f53c67ca1..b82e860c8aad 100644 --- a/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py +++ b/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py @@ -54,7 +54,7 @@ def test_prim_type_hidden_path_exact_message(): ) -def test_prim_func_buffer_map(): +def test_prim_func_buffer_param(): @T.prim_func(s_tir=True) def func1(a: T.handle, b: T.handle): A = T.match_buffer(a, (128, 128)) @@ -74,15 +74,15 @@ def func2(a: T.handle, b: T.handle): func1, func2, AccessPath.root() - .attr("buffer_map") - .map_item(func1.params[1]) + .attr("params") + .array_item(1) .attr("ty") .attr("shape") .array_item(1) .attr("value"), AccessPath.root() - .attr("buffer_map") - .map_item(func2.params[1]) + .attr("params") + .array_item(1) .attr("ty") .attr("shape") .array_item(1) diff --git a/tests/python/tvmscript/test_tvmscript_printer_tir.py b/tests/python/tvmscript/test_tvmscript_printer_tir.py index ac0d76166737..9908a51120da 100644 --- a/tests/python/tvmscript/test_tvmscript_printer_tir.py +++ b/tests/python/tvmscript/test_tvmscript_printer_tir.py @@ -33,16 +33,12 @@ def _assert_print(obj, expected): def test_prim_func(): - a = tirx.Var("a", "handle") - b = tirx.Var("b", "handle") + A = tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A") + B = tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B") func = ( tirx.PrimFunc( - params=[a, b], + params=[A, B], ret_type=None, - buffer_map={ - a: tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A"), - b: tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B"), - }, body=tirx.Evaluate(0), ) .with_attr("global_symbol", "main") @@ -60,18 +56,14 @@ def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")) ) -def test_prim_func_no_sugar_inlined_buffer(): - a = tirx.Var("a", "handle") - b = tirx.Var("b", "handle") +def test_prim_func_buffer_data_use(): + A = tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A") + B = tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B") func = ( tirx.PrimFunc( - params=[a, b], + params=[A, B], ret_type=None, - buffer_map={ - a: tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A"), - b: tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B"), - }, - body=tirx.Evaluate(a), + body=tirx.Evaluate(A.data), ) .with_attr("global_symbol", "main") .with_attr("s_tir", True) @@ -83,25 +75,20 @@ def test_prim_func_no_sugar_inlined_buffer(): # from tvm.tirx.layout import Axis @T.prim_func(s_tir=True) -def main(a: T.handle, B: T.Buffer((256, 256), "float32")): - A = T.match_buffer(a, (128, 128)) - T.evaluate(a) +def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")): + T.evaluate(A.data) """, ) def test_prim_func_buffer_data_argument_is_scope_hint(): - a = tirx.Var("a", "handle") - b = tirx.Var("b", "handle") buffer_data = tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A").data + A = tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A", data=buffer_data) + B = tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B", data=buffer_data) func = ( tirx.PrimFunc( - params=[a, b], + params=[A, B], ret_type=None, - buffer_map={ - a: tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A", data=buffer_data), - b: tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B", data=buffer_data), - }, body=tirx.Evaluate(0), ) .with_attr("global_symbol", "main") @@ -823,15 +810,11 @@ def main(): def test_private_primfunc(): from tvm.script import tirx as T - a = tirx.Var("a", "handle") - b = tirx.Var("b", "handle") + A = tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A") + B = tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B") func = tirx.PrimFunc( - params=[a, b], + params=[A, B], ret_type=None, - buffer_map={ - a: tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A"), - b: tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B"), - }, body=tirx.Evaluate(0), ).with_attr("s_tir", True) _assert_print( @@ -984,25 +967,24 @@ def func(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")) def test_predicated_buffer_load_store(): a = tirx.Var("a", "handle") b = tirx.Var("b", "handle") - buffer_map = { + buffers = { a: tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A"), b: tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B"), } buffer_load = tirx.BufferLoad( - buffer=buffer_map[b], + buffer=buffers[b], indices=[0, tirx.Ramp(0, 4, 4)], predicate=tirx.Broadcast(tirx.IntImm("bool", 0), 4), ) body = tirx.BufferStore( - buffer=buffer_map[a], + buffer=buffers[a], value=buffer_load, indices=[0, tirx.Ramp(0, 2, 4)], predicate=tirx.Broadcast(tirx.IntImm("bool", 0), 4), ) func = tirx.PrimFunc( - params=[a, b], + params=[buffers[a], buffers[b]], ret_type=None, - buffer_map=buffer_map, body=body, ).with_attr("s_tir", True) From 14b9545313c6f1b9a394beb0ee6ad007c5e7632c Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Sun, 2 Aug 2026 13:40:00 +0000 Subject: [PATCH 2/7] [TIRx] Apply mechanical buffer parameter cleanup --- include/tvm/tirx/function.h | 10 +--------- src/relax/analysis/layout_transformation.cc | 2 +- src/relax/analysis/tir_op_pattern_kind.cc | 4 ++-- src/relax/transform/fuse_tir.cc | 4 ++-- src/relax/transform/rewrite_dataflow_reshape.cc | 2 +- src/relax/transform/split_call_tir_by_pattern.cc | 6 +++--- src/relax/transform/split_layout_rewrite_preproc.cc | 2 +- .../transform/remove_weight_layout_rewrite_block.cc | 2 +- src/s_tir/transform/renew_defs.cc | 6 +++--- src/tirx/ir/data_type_rewriter.cc | 2 +- src/tirx/ir/specialize.cc | 8 ++++---- src/tirx/ir/tir_visitor_with_path.cc | 4 ++-- tests/python/s_tir/base/test_tir_te_extern_primfunc.py | 4 +--- 13 files changed, 23 insertions(+), 33 deletions(-) diff --git a/include/tvm/tirx/function.h b/include/tvm/tirx/function.h index 74ece4538d23..0ed90d57e740 100644 --- a/include/tvm/tirx/function.h +++ b/include/tvm/tirx/function.h @@ -38,19 +38,11 @@ namespace tvm { namespace tirx { -/*! \brief Return a checked buffer view when a parameter carries BufferType. */ -inline ffi::Optional AsBufferVar(const Var& var) { - if (var->ty.as()) { - return BufferVar(var); - } - return std::nullopt; -} - /*! \brief Derive a keyed view of buffer parameters from their BufferType annotations. */ inline ffi::Map BufferParamMap(const ffi::Array& params) { ffi::Map result; for (const Var& param : params) { - if (auto buffer = AsBufferVar(param)) { + if (auto buffer = param.as()) { result.Set(param, buffer.value()); } } diff --git a/src/relax/analysis/layout_transformation.cc b/src/relax/analysis/layout_transformation.cc index f5b4cf0085b1..207e3117023b 100644 --- a/src/relax/analysis/layout_transformation.cc +++ b/src/relax/analysis/layout_transformation.cc @@ -562,7 +562,7 @@ class PrimFuncAnalyzer : public StmtExprVisitor { size_t first_write_index = func->params.size() - write_transformations.size(); for (size_t i = 0; i < write_transformations.size(); ++i) { auto param = func->params[first_write_index + i]; - ffi::Optional param_buf = tirx::AsBufferVar(param); + ffi::Optional param_buf = param.as(); TVM_FFI_ICHECK(param_buf.has_value()); TVM_FFI_ICHECK_EQ(param_buf.value()->shape.size(), write_transformations[i]->initial_indices.size()) diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc index 43e703d6435b..c7b27ac2d2fb 100644 --- a/src/relax/analysis/tir_op_pattern_kind.cc +++ b/src/relax/analysis/tir_op_pattern_kind.cc @@ -39,7 +39,7 @@ class PatternKindAnalyzer : public StmtExprVisitor { public: explicit PatternKindAnalyzer(const tirx::PrimFunc& func) { for (const tirx::Var& param : func->params) { - ffi::Optional param_buf = tirx::AsBufferVar(param); + ffi::Optional param_buf = param.as(); if (param_buf.has_value()) { param_buffers_.insert(param_buf.value()); } @@ -535,7 +535,7 @@ bool HasReshapePattern(const PrimFunc& func) { ffi::Array buffer_args; for (const auto& param : func->params) { - if (auto buffer = tirx::AsBufferVar(param)) { + if (auto buffer = param.as()) { buffer_args.push_back(buffer.value()); } } diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc index 69729afe2a89..d801133c75b7 100644 --- a/src/relax/transform/fuse_tir.cc +++ b/src/relax/transform/fuse_tir.cc @@ -810,7 +810,7 @@ class FusedTIRConstructor : public ExprVisitor { for (size_t i = 0; i < call_args.size(); ++i) { const Expr& arg = call_args[i]; const tirx::Var& param = func->params[i]; - if (auto buffer = tirx::AsBufferVar(param)) { + if (auto buffer = param.as()) { arg_list.push_back(arg); buffer_list.push_back(buffer.value()); } else { @@ -836,7 +836,7 @@ class FusedTIRConstructor : public ExprVisitor { for (int64_t idx : output_indices) { int i = static_cast(idx); const tirx::Var& param = func->params[static_cast(i)]; - auto buffer = tirx::AsBufferVar(param); + auto buffer = param.as(); TVM_FFI_ICHECK(buffer.has_value()) << "The output params of a PrimFunc must be buffers, but parameter " << i << " has type " << param->ty; diff --git a/src/relax/transform/rewrite_dataflow_reshape.cc b/src/relax/transform/rewrite_dataflow_reshape.cc index c73da19ae6a4..464be7d774cf 100644 --- a/src/relax/transform/rewrite_dataflow_reshape.cc +++ b/src/relax/transform/rewrite_dataflow_reshape.cc @@ -39,7 +39,7 @@ namespace relax { std::vector GetUsedTensorArgIndices(const tirx::PrimFunc& fn, size_t num_args) { std::vector indices; for (size_t i = 0; i < num_args; ++i) { - if (auto buffer = tirx::AsBufferVar(fn->params[i])) { + if (auto buffer = fn->params[i].as()) { auto buffer_var = buffer.value().var(); if (tirx::UsesVar(fn->body, [=](const tirx::VarNode* var) { return var == buffer_var.get(); })) { diff --git a/src/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc index adacbfc14862..290be4c22cef 100644 --- a/src/relax/transform/split_call_tir_by_pattern.cc +++ b/src/relax/transform/split_call_tir_by_pattern.cc @@ -70,7 +70,7 @@ class ForMatcher : public TensorizeComparator { } // Get evaluated symbols, buffers from the pattern. for (const auto& arg : pattern_->params) { - if (auto buffer = tirx::AsBufferVar(arg)) { + if (auto buffer = arg.as()) { auto itt = rhs_buffer_map_.find(buffer.value()); TVM_FFI_ICHECK(itt != rhs_buffer_map_.end()); evaluated_buffers.push_back(itt->second); @@ -613,7 +613,7 @@ std::pair> SplitFunctions( for (const auto& buffer : func1_args) { TVM_FFI_ICHECK(partitioner.input1.find(buffer) != partitioner.input1.end()); for (size_t i = 0; i < func->params.size(); i++) { - auto param_buffer = tirx::AsBufferVar(func->params[i]); + auto param_buffer = func->params[i].as(); if (param_buffer.has_value() && param_buffer.value().same_as(buffer)) { new_params1.push_back(func->params[i]); arg_partition1.push_back(i); @@ -631,7 +631,7 @@ std::pair> SplitFunctions( new_params2.push_back(partitioner.intermediate_buffer.var()); for (int i = 0; i < static_cast(func->params.size()); i++) { Var param = func->params[i]; - auto param_buffer = tirx::AsBufferVar(param); + auto param_buffer = param.as(); if (param_buffer.has_value() && partitioner.input2.count(param_buffer.value())) { new_params2.push_back(param); if (i != static_cast(func->params.size()) - 1) { diff --git a/src/relax/transform/split_layout_rewrite_preproc.cc b/src/relax/transform/split_layout_rewrite_preproc.cc index 1c920155d4cd..94ab0dcb3a99 100644 --- a/src/relax/transform/split_layout_rewrite_preproc.cc +++ b/src/relax/transform/split_layout_rewrite_preproc.cc @@ -101,7 +101,7 @@ class SplitPrimFuncLayoutRewrite : public StmtMutator { ffi::Array params = original_func_->params; for (const auto& info : rewrite_infos_) { const Var& param = params[info.buffer_index]; - TVM_FFI_ICHECK(tirx::AsBufferVar(param).value() == info.pre_rewrite_buffer); + TVM_FFI_ICHECK(param.as().value() == info.pre_rewrite_buffer); params.Set(info.buffer_index, info.post_rewrite_buffer.var()); } diff --git a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc index 076218e6bab0..a987c1abeb2c 100644 --- a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc +++ b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc @@ -128,7 +128,7 @@ class WeightLayoutRewriteBlockRemover : public StmtMutator { ffi::Array params; for (const tirx::Var& param : f_->params) { - auto opt_buffer = tirx::AsBufferVar(param); + auto opt_buffer = param.as(); if (!opt_buffer.has_value()) { params.push_back(param); continue; diff --git a/src/s_tir/transform/renew_defs.cc b/src/s_tir/transform/renew_defs.cc index 273ee423031d..4c7a19ba28e5 100644 --- a/src/s_tir/transform/renew_defs.cc +++ b/src/s_tir/transform/renew_defs.cc @@ -51,12 +51,12 @@ class RenewDefMutator : public StmtExprMutator { // Redefine scalar parameters first, because they may occur in a buffer // parameter's type annotation. for (const auto& param : func->params) { - if (!tirx::AsBufferVar(param)) { + if (!param.as()) { generator.ReDefineVar(param); } } for (const auto& param : func->params) { - if (auto opt_buffer = tirx::AsBufferVar(param)) { + if (auto opt_buffer = param.as()) { const BufferVar& buffer = opt_buffer.value(); for (const PrimExpr& e : buffer->shape) { if (auto var = e.as()) { @@ -71,7 +71,7 @@ class RenewDefMutator : public StmtExprMutator { // TODO(Siyuan Feng): checking var is used after define ffi::Array params; for (const auto& param : func->params) { - if (auto opt_buffer = tirx::AsBufferVar(param)) { + if (auto opt_buffer = param.as()) { params.push_back(generator.DefineBuffer(opt_buffer.value())); } else { params.push_back(generator.VisitExpr(param).as_or_throw()); diff --git a/src/tirx/ir/data_type_rewriter.cc b/src/tirx/ir/data_type_rewriter.cc index ca264bda684f..45b84e4edd8c 100644 --- a/src/tirx/ir/data_type_rewriter.cc +++ b/src/tirx/ir/data_type_rewriter.cc @@ -642,7 +642,7 @@ PrimFunc IndexDataTypeNormalizer::Rewrite(PrimFunc func) { bool is_enabled = true; std::swap(is_enabled_, is_enabled); ffi::Array params = func->params.Map([this](Var param) { - if (auto buffer = AsBufferVar(param)) { + if (auto buffer = param.as()) { return buffer_remap_.Get(buffer.value()).value_or(buffer.value()).var(); } if (auto param_ty = param->ty.as(); diff --git a/src/tirx/ir/specialize.cc b/src/tirx/ir/specialize.cc index 9aec2a325341..e0e70024813e 100644 --- a/src/tirx/ir/specialize.cc +++ b/src/tirx/ir/specialize.cc @@ -80,13 +80,13 @@ class PrimFuncSpecializer : public StmtExprMutator { static PrimFunc Specialize(PrimFunc f, const VarMap& var_map) { PrimFuncSpecializer specializer(var_map); for (const Var& param : f->params) { - auto buffer = tirx::AsBufferVar(param); + auto buffer = param.as(); auto replacement = var_map.find(param); if (!buffer || replacement == var_map.end()) { continue; } if (auto replacement_var = replacement->second.as()) { - if (auto replacement_buffer = tirx::AsBufferVar(replacement_var.value())) { + if (auto replacement_buffer = replacement_var.value().as()) { if (IsParam(f, replacement_var.value())) { specializer.buffer_aliases_[buffer.value()] = replacement_buffer.value(); } else { @@ -365,7 +365,7 @@ void UpdateSpecializeVarMap(const PrimFunc& func, const Var& param, const Buffer // preliminaries tirx::ExprDeepEqual equal; - auto opt_buffer = tirx::AsBufferVar(param); + auto opt_buffer = param.as(); TVM_FFI_CHECK(opt_buffer, ValueError) << "specialize expects param to have a BufferType annotation"; const BufferVar& buf_to_specialize = opt_buffer.value(); @@ -441,7 +441,7 @@ void UpdateSpecializeVarMap(const PrimFunc& func, const Var& param, const Expr& TVM_FFI_CHECK(IsParam(func, param), ValueError) << "Specialize expects param to be in PrimFunc's params"; // Specialize a scalar parameter rather than a buffer parameter. - TVM_FFI_CHECK(!tirx::AsBufferVar(param), ValueError) + TVM_FFI_CHECK(!param.as(), ValueError) << "Specialize expects param to not have a BufferType annotation"; // build var mapping using specific_expr (*var_map)[param] = specific_expr; diff --git a/src/tirx/ir/tir_visitor_with_path.cc b/src/tirx/ir/tir_visitor_with_path.cc index 4223f51d9ac5..4b39167651fb 100644 --- a/src/tirx/ir/tir_visitor_with_path.cc +++ b/src/tirx/ir/tir_visitor_with_path.cc @@ -90,7 +90,7 @@ void TIRVisitorWithPath::Visit(const PrimFunc& func, AccessPath path) { } for (size_t i = 0; i < func->params.size(); i++) { - if (auto opt = AsBufferVar(func->params[i])) { + if (auto opt = func->params[i].as()) { auto buf = opt.value(); auto buf_path = ppath->ArrayItem(i)->Attr("ty"); @@ -103,7 +103,7 @@ void TIRVisitorWithPath::Visit(const PrimFunc& func, AccessPath path) { // Only after all the implicit definitions have been visited can we // visit the buffer definition itself. for (size_t i = 0; i < func->params.size(); i++) { - if (auto opt = AsBufferVar(func->params[i])) { + if (auto opt = func->params[i].as()) { context.push_back(WithDef(opt.value(), ppath->ArrayItem(i))); } } diff --git a/tests/python/s_tir/base/test_tir_te_extern_primfunc.py b/tests/python/s_tir/base/test_tir_te_extern_primfunc.py index be10c72262e6..65cb259c988f 100644 --- a/tests/python/s_tir/base/test_tir_te_extern_primfunc.py +++ b/tests/python/s_tir/base/test_tir_te_extern_primfunc.py @@ -192,9 +192,7 @@ def test_te_extern_call(self, func, params, verify): ir_mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main")) prim_func = ir_mod["main"] - buf_name_map = { - param.name: param for param in func.params if tvm.tirx.is_buffer_var(param) - } + buf_name_map = {param.name: param for param in func.params if tvm.tirx.is_buffer_var(param)} input_tensors = [te.placeholder(buf_name_map[name].shape) for name in params] output = te.extern_primfunc(input_tensors, prim_func) rt_prim_func = te.create_prim_func(tensors_from_extern_op(output, prim_func)) From b8b66af7e7a2ca136b935ce3351397f09f1f54d6 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Sun, 2 Aug 2026 13:41:04 +0000 Subject: [PATCH 3/7] [FIX][TIRx] Preserve buffer parameter binding semantics --- .../tvm/relax/distributed/axis_group_graph.h | 6 +- python/tvm/tirx/script/parser/parser.py | 91 +++++++++++++++++-- src/tirx/ir/function.cc | 22 ++++- src/tirx/script/builder/frame.cc | 12 +++ 4 files changed, 122 insertions(+), 9 deletions(-) diff --git a/include/tvm/relax/distributed/axis_group_graph.h b/include/tvm/relax/distributed/axis_group_graph.h index 5b47fe4f6ce4..6d446e3f7f35 100644 --- a/include/tvm/relax/distributed/axis_group_graph.h +++ b/include/tvm/relax/distributed/axis_group_graph.h @@ -71,6 +71,11 @@ class BufferAxisGraphExtractor : public StmtExprVisitor { BufferAxisGraphExtractor extractor; extractor(prim_func->body); ffi::Map inverse_buffer_map; + for (const Var& param : prim_func->params) { + if (param->ty.as()) { + inverse_buffer_map.Set(BufferVar(param), param); + } + } std::vector> tir_var_axis_group_list; std::unordered_set visited; for (const Var& param : prim_func->params) { @@ -78,7 +83,6 @@ class BufferAxisGraphExtractor : public StmtExprVisitor { continue; } BufferVar buffer(param); - inverse_buffer_map.Set(buffer, param); for (int i = 0; i < static_cast(buffer->shape.size()); i++) { if (extractor.buffer_axis_graph_.count({buffer, i})) { std::vector buffer_axis_group; diff --git a/python/tvm/tirx/script/parser/parser.py b/python/tvm/tirx/script/parser/parser.py index d8a7445a9e1a..17fff86d200c 100644 --- a/python/tvm/tirx/script/parser/parser.py +++ b/python/tvm/tirx/script/parser/parser.py @@ -272,6 +272,66 @@ def _is_jit_function(node: doc.FunctionDef) -> bool: return False +def _loaded_names(node: doc.expr) -> set[str]: + return { + child.id + for child in ast.walk(from_doc(node)) + if isinstance(child, ast.Name) and isinstance(child.ctx, ast.Load) + } + + +def _stored_names(node: doc.expr) -> set[str]: + return { + child.id + for child in ast.walk(from_doc(node)) + if isinstance(child, ast.Name) and isinstance(child.ctx, ast.Store) + } + + +def _bind_forward_signature_vars(self: Parser, node: doc.FunctionDef) -> list[doc.stmt]: + """Bind body-defined PrimVars referenced by the function signature. + + Free symbolic dimensions are printed as ordinary top-level PrimVar + declarations in the function body, after their first textual use in a + buffer annotation. Resolve exactly those declaration dependencies before + evaluating the signature, and leave all unrelated body statements in their + original order. + """ + + signature_names = set() + for arg in node.args.args: + if arg.annotation is not None: + signature_names.update(_loaded_names(arg.annotation)) + if node.returns is not None: + signature_names.update(_loaded_names(node.returns)) + + remaining_body = [] + for statement in node.body: + if ( + not isinstance(statement, doc.Assign) + or len(statement.targets) != 1 + or not (_stored_names(statement.targets[0]) & signature_names) + ): + remaining_body.append(statement) + continue + try: + value = self.eval_expr(statement.value) + except Exception: # pylint: disable=broad-except + remaining_body.append(statement) + continue + values = value if isinstance(value, list | tuple) else [value] + if not all( + isinstance(item, tvm.ir.Var) + and isinstance(item.ty, PrimType) + and not is_buffer_var(item) + for item in values + ): + remaining_body.append(statement) + continue + self.eval_assign(target=statement.targets[0], source=value, bind_value=bind_assign_value) + return remaining_body + + @dispatch.register(token="tirx", type_name="For") def visit_for(self: Parser, node: doc.For) -> None: """The for visiting method for tirx. @@ -636,6 +696,7 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None: prim_func_ctx = T.prim_func(is_private=privacy, s_tir=s_tir, persistent=persistent) with prim_func_ctx: T.func_name(node.name) + remaining_body = _bind_forward_signature_vars(self, node) if node.returns is not None: ret_type = self.eval_expr(node.returns) if callable(ret_type) and not isinstance(ret_type, Expr): @@ -713,7 +774,7 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None: continue param = T.arg(arg.arg, ann) self.var_table.add(arg.arg, param) - self.visit_body(node.body) + self.visit_body(remaining_body) self.function_annotations = supplied_annotation @@ -954,6 +1015,7 @@ def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar ret_type = None with self.var_table.with_frame(): + _bind_forward_signature_vars(self, node) if node.returns is not None: ret_type = self.eval_expr(node.returns) if callable(ret_type) and not isinstance(ret_type, Expr): @@ -961,18 +1023,35 @@ def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar if isinstance(ret_type, Expr): ret_type = ret_type.ty - arg_annotations = [] + evaluated_annotations = {} for arg in node.args.args: if arg.annotation is None: - self.report_error(arg, "Type annotation required for function parameters.") + continue try: ann = self.eval_expr(arg.annotation) if callable(ann) and not isinstance(ann, Expr): ann = ann() except Exception: # pylint: disable=broad-except - ann = func_annotation.get(arg.arg, None) - if ann is None: - raise + continue + evaluated_annotations[arg.arg] = ann + if isinstance(ann, tvm.tirx.Var) and not is_buffer_var(ann): + self.var_table.add(arg.arg, ann) + + arg_annotations = [] + for arg in node.args.args: + if arg.annotation is None: + self.report_error(arg, "Type annotation required for function parameters.") + if arg.arg in evaluated_annotations: + ann = evaluated_annotations[arg.arg] + else: + try: + ann = self.eval_expr(arg.annotation) + if callable(ann) and not isinstance(ann, Expr): + ann = ann() + except Exception: # pylint: disable=broad-except + ann = func_annotation.get(arg.arg, None) + if ann is None: + raise IRBuilder.name(arg.arg, ann) arg_annotations.append(ann) diff --git a/src/tirx/ir/function.cc b/src/tirx/ir/function.cc index 4b4ab4878022..b1c4acedc733 100644 --- a/src/tirx/ir/function.cc +++ b/src/tirx/ir/function.cc @@ -28,6 +28,7 @@ #include #include #include +#include namespace tvm { namespace tirx { @@ -52,6 +53,22 @@ ffi::Array NormalizeBufferParams(ffi::Array params, return normalized; } +Stmt NormalizeBufferParamUses(Stmt body, const ffi::Map& buffer_map) { + if (!body.defined()) { + return body; + } + ffi::Map replacements; + for (const auto& [param, buffer] : buffer_map) { + if (!param.same_as(buffer.var()) && !param->ty.as()) { + Expr data = buffer.data(); + replacements.Set(param, ffi::StructuralEqual()(param->ty, data->ty) + ? data + : tvm::reinterpret(param->ty, std::move(data))); + } + } + return replacements.empty() ? body : Substitute(std::move(body), replacements); +} + tvm::Type InferType(const PrimFunc& prim_func) { ffi::Array params; for (const auto& param : prim_func->params) { @@ -113,8 +130,9 @@ PrimFunc::PrimFunc(ffi::Array params, Stmt body, Type ret_type, DictA PrimFunc::PrimFunc(ffi::Array params, Stmt body, Type ret_type, ffi::Map buffer_map, DictAttrs attrs, Span span) - : PrimFunc(NormalizeBufferParams(std::move(params), buffer_map), std::move(body), - std::move(ret_type), std::move(attrs), std::move(span)) {} + : PrimFunc(NormalizeBufferParams(std::move(params), buffer_map), + NormalizeBufferParamUses(std::move(body), buffer_map), std::move(ret_type), + std::move(attrs), std::move(span)) {} FuncType PrimFuncNode::func_type_annotation() const { ffi::Array param_types; diff --git a/src/tirx/script/builder/frame.cc b/src/tirx/script/builder/frame.cc index 640842f7cb4a..7a41ad7a152a 100644 --- a/src/tirx/script/builder/frame.cc +++ b/src/tirx/script/builder/frame.cc @@ -117,8 +117,10 @@ void PrimFuncFrameNode::ExitWithScope() { tvm::tirx::Stmt body = AsStmt(stmts); STirBufferLayoutNormalizer normalizer; ffi::Array effective_args; + ffi::Map param_replacements; for (const tvm::tirx::Var& arg : args) { ffi::Optional opt_buffer = buffer_map.Get(arg); + bool replaces_legacy_param = opt_buffer.has_value(); if (!opt_buffer.has_value() && arg->ty.as()) { opt_buffer = tvm::tirx::BufferVar(arg); } @@ -135,6 +137,13 @@ void PrimFuncFrameNode::ExitWithScope() { buffer = new_buffer; } effective_args.push_back(buffer.var()); + if (replaces_legacy_param && !arg.same_as(buffer.var()) && + !arg->ty.as()) { + tvm::Expr data = buffer.data(); + param_replacements.Set(arg, ffi::StructuralEqual()(arg->ty, data->ty) + ? data + : tvm::reinterpret(arg->ty, std::move(data))); + } } if (!normalizer.Empty()) { body = normalizer(std::move(body)); @@ -144,6 +153,9 @@ void PrimFuncFrameNode::ExitWithScope() { } effective_root_alloc_buffers = std::move(new_root_alloc_buffers); } + if (!param_replacements.empty()) { + body = tvm::tirx::Substitute(std::move(body), param_replacements); + } tvm::tirx::PrimFunc func( /*params=*/effective_args, /*body=*/body, From 87fda62f39add1de02f461e545d798df4b240754 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Sun, 2 Aug 2026 23:32:18 +0000 Subject: [PATCH 4/7] [REFACTOR][TIRx] Use typed buffer parameters directly --- include/tvm/tirx/function.h | 11 - python/tvm/tirx/script/parser/parser.py | 219 ++++++++++++------ src/backend/trn/codegen/codegen_trn.cc | 2 +- .../trn/transform/lower_trainium_layout.cc | 40 ++-- src/relax/transform/dataflow_inplace.cc | 16 +- src/relax/transform/fuse_tir.cc | 3 +- .../transform/split_call_tir_by_pattern.cc | 9 +- .../transform/split_layout_rewrite_preproc.cc | 3 +- .../sblock_buffer_access_lca_detector.cc | 7 +- src/s_tir/backend/adreno/texture_flatten.cc | 17 +- src/s_tir/meta_schedule/arg_info.cc | 7 +- .../feature_extractor/per_store_feature.cc | 6 +- .../meta_schedule/postproc/rewrite_layout.cc | 2 +- src/s_tir/schedule/analysis/analysis.cc | 9 +- .../schedule/primitive/blockize_tensorize.cc | 6 +- .../primitive/layout_transformation.cc | 15 +- src/s_tir/transform/inject_permuted_layout.cc | 8 +- .../transform/inject_software_pipeline.cc | 7 +- src/s_tir/transform/lower_thread_allreduce.cc | 9 +- .../plan_update_buffer_allocation_location.cc | 15 +- src/target/llvm/codegen_llvm.cc | 6 +- src/te/operation/create_primfunc.cc | 6 +- src/tirx/analysis/verify_memory.cc | 4 +- src/tirx/analysis/verify_ssa.cc | 6 +- src/tirx/ir/data_type_rewriter.cc | 6 +- src/tirx/ir/script/script_complete.cc | 7 +- src/tirx/ir/specialize.cc | 31 +-- src/tirx/script/printer/buffer.cc | 9 +- src/tirx/script/printer/function.cc | 14 +- src/tirx/transform/flatten_buffer.cc | 20 +- .../transform/force_narrow_index_to_i32.cc | 10 +- .../transform/inline_private_functions.cc | 14 +- src/tirx/transform/ir_utils.cc | 46 ++-- src/tirx/transform/lower_tirx_cleanup.cc | 47 ++-- src/tirx/transform/make_packed_api.cc | 5 +- src/tirx/transform/split_host_device.cc | 2 +- src/tirx/transform/storage_rewrite.cc | 66 ++---- src/tirx/transform/tvm_ffi_binder.cc | 28 +-- src/tirx/transform/tvm_ffi_binder.h | 16 +- .../transform/unsupported_dtype_legalize.cc | 6 +- .../tvmscript/test_tvmscript_printer_tir.py | 15 ++ 41 files changed, 397 insertions(+), 378 deletions(-) diff --git a/include/tvm/tirx/function.h b/include/tvm/tirx/function.h index 0ed90d57e740..35594f4f7dfd 100644 --- a/include/tvm/tirx/function.h +++ b/include/tvm/tirx/function.h @@ -38,17 +38,6 @@ namespace tvm { namespace tirx { -/*! \brief Derive a keyed view of buffer parameters from their BufferType annotations. */ -inline ffi::Map BufferParamMap(const ffi::Array& params) { - ffi::Map result; - for (const Var& param : params) { - if (auto buffer = param.as()) { - result.Set(param, buffer.value()); - } - } - return result; -} - /*! * \brief Primitive functions that contains TIR statements. * diff --git a/python/tvm/tirx/script/parser/parser.py b/python/tvm/tirx/script/parser/parser.py index 17fff86d200c..603949380444 100644 --- a/python/tvm/tirx/script/parser/parser.py +++ b/python/tvm/tirx/script/parser/parser.py @@ -163,7 +163,13 @@ def bind_for_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> A raise NotImplementedError -def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any: +def bind_assign_value( + self: Parser, + node: doc.expr, + var_name: str, + value: Any, + prim_var_declarations: dict[str, str] | None = None, +) -> Any: """Value binding methods when parsing assign statement. e.g. binding vi, vj, vk with T.axis.remap("SSR", [i, j, k]), when parsing vi, vj, vk = T.axis.remap("SSR", [i, j, k]). @@ -187,6 +193,27 @@ def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) - res : Any The bound value. """ + if var_name in (prim_var_declarations or {}): + previous = self.var_table.get().get(var_name) + signature_prim_vars = getattr(self, "_signature_prim_vars", []) + if isinstance(previous, tvm.ir.Var) and any( + previous.same_as(var) for var in signature_prim_vars + ): + if not ( + isinstance(previous, tvm.ir.Var) + and isinstance(previous.ty, PrimType) + and isinstance(value, tvm.ir.Var) + and isinstance(value.ty, PrimType) + and not is_buffer_var(previous) + and not is_buffer_var(value) + ): + self.report_error(node, f"{var_name} is already bound to a non-PrimVar value") + if previous.ty != value.ty: + self.report_error( + node, + f"Expected the same dtype for PrimVars but got {value.ty} vs {previous.ty}", + ) + return previous if isinstance(value, T.scalar_wrapper): # pylint: disable=protected-access # special case for scalar, name the buffer, but the var is used as BufferLoad assert isinstance(value.scalar, T.BufferLoad) @@ -272,64 +299,103 @@ def _is_jit_function(node: doc.FunctionDef) -> bool: return False -def _loaded_names(node: doc.expr) -> set[str]: - return { - child.id - for child in ast.walk(from_doc(node)) - if isinstance(child, ast.Name) and isinstance(child.ctx, ast.Load) - } - - -def _stored_names(node: doc.expr) -> set[str]: - return { - child.id - for child in ast.walk(from_doc(node)) - if isinstance(child, ast.Name) and isinstance(child.ctx, ast.Store) - } +def _prim_var_declaration_dtype(node: doc.expr) -> str | None: + """Return the dtype of literal ``T.dtype()`` declaration syntax.""" + if not ( + isinstance(node, doc.Call) + and not node.args + and not node.keywords + and isinstance(node.func, doc.Attribute) + and isinstance(node.func.value, doc.Name) + and node.func.value.id == "T" + ): + return None + constructor = getattr(T, node.func.attr, None) + if not isinstance(constructor, T.DtypeConstructor) and constructor is not T.bool: + return None + return str(constructor().ty.dtype) + + +def _collect_prim_var_declarations(target: doc.expr, value: doc.expr) -> dict[str, str]: + """Pair assignment targets with literal PrimVar declaration dtypes.""" + if isinstance(target, doc.Name): + dtype = _prim_var_declaration_dtype(value) + return {target.id: dtype} if dtype is not None else {} + if isinstance(target, doc.Tuple | doc.List) and isinstance(value, doc.Tuple | doc.List): + if len(target.elts) != len(value.elts): + return {} + declarations = {} + for lhs, rhs in zip(target.elts, value.elts): + declarations.update(_collect_prim_var_declarations(lhs, rhs)) + return declarations + return {} + + +def _function_prim_var_declarations(node: doc.FunctionDef) -> dict[str, str]: + """Collect body declaration dtypes without evaluating or binding them.""" + declarations = {} + for statement in node.body: + if isinstance(statement, doc.Assign) and len(statement.targets) == 1: + declarations.update( + _collect_prim_var_declarations(statement.targets[0], statement.value) + ) + return declarations -def _bind_forward_signature_vars(self: Parser, node: doc.FunctionDef) -> list[doc.stmt]: - """Bind body-defined PrimVars referenced by the function signature. +def _eval_signature_annotation( + self: Parser, node: doc.expr, declaration_dtypes: dict[str, str] +) -> Any: + """Evaluate string expressions in function-parameter buffer shapes. - Free symbolic dimensions are printed as ordinary top-level PrimVar - declarations in the function body, after their first textual use in a - buffer annotation. Resolve exactly those declaration dependencies before - evaluating the signature, and leave all unrelated body statements in their - original order. + String expressions are the sole exception to ordinary definition-before-use + ordering. Their first occurrence creates any missing PrimVars, while a + later ``n = T.dtype()`` declaration reuses the same object. """ - signature_names = set() - for arg in node.args.args: - if arg.annotation is not None: - signature_names.update(_loaded_names(arg.annotation)) - if node.returns is not None: - signature_names.update(_loaded_names(node.returns)) + class ShapeStringRewriter(ast.NodeTransformer): + def visit_Constant(self, constant): # pylint: disable=invalid-name + if not isinstance(constant.value, str): + return constant + expression = ast.parse(constant.value, mode="eval").body + for child in ast.walk(expression): + if not isinstance(child, ast.Name) or not isinstance(child.ctx, ast.Load): + continue + if child.id not in self_parser.var_table.get(): + var = tvm.tirx.Var(child.id, declaration_dtypes.get(child.id, "int64")) + self_parser.var_table.add(child.id, var, allow_shadowing=False) + signature_prim_vars = getattr(self_parser, "_signature_prim_vars", []) + signature_prim_vars.append(var) + self_parser._signature_prim_vars = signature_prim_vars + for child in ast.walk(expression): + ast.copy_location(child, constant) + return expression + + class BufferShapeRewriter(ast.NodeTransformer): + def visit_Call(self, call): # pylint: disable=invalid-name + is_buffer = ( + isinstance(call.func, ast.Attribute) + and call.func.attr == "Buffer" + and isinstance(call.func.value, ast.Name) + and call.func.value.id == "T" + ) + if not is_buffer: + return self.generic_visit(call) + call.func = self.visit(call.func) + if call.args: + call.args[0] = ShapeStringRewriter().visit(call.args[0]) + call.args[1:] = [self.visit(arg) for arg in call.args[1:]] + for keyword in call.keywords: + if keyword.arg == "shape": + keyword.value = ShapeStringRewriter().visit(keyword.value) + else: + keyword.value = self.visit(keyword.value) + return call - remaining_body = [] - for statement in node.body: - if ( - not isinstance(statement, doc.Assign) - or len(statement.targets) != 1 - or not (_stored_names(statement.targets[0]) & signature_names) - ): - remaining_body.append(statement) - continue - try: - value = self.eval_expr(statement.value) - except Exception: # pylint: disable=broad-except - remaining_body.append(statement) - continue - values = value if isinstance(value, list | tuple) else [value] - if not all( - isinstance(item, tvm.ir.Var) - and isinstance(item.ty, PrimType) - and not is_buffer_var(item) - for item in values - ): - remaining_body.append(statement) - continue - self.eval_assign(target=statement.targets[0], source=value, bind_value=bind_assign_value) - return remaining_body + self_parser = self + python_node = from_doc(deepcopy(node)) + python_node = BufferShapeRewriter().visit(python_node) + ast.fix_missing_locations(python_node) + return self.eval_expr(doc.to_doc(python_node)) @dispatch.register(token="tirx", type_name="For") @@ -506,7 +572,12 @@ def visit_assign(self: Parser, node: doc.Assign) -> None: except TypeError: pass # rhs not compatible with buffer_store, fall through # otherwise - self.eval_assign(target=lhs, source=rhs, bind_value=bind_assign_value) + declarations = _collect_prim_var_declarations(lhs, node.value) + self.eval_assign( + target=lhs, + source=rhs, + bind_value=partial(bind_assign_value, prim_var_declarations=declarations), + ) @dispatch.register(token="tirx", type_name="AugAssign") @@ -696,14 +767,7 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None: prim_func_ctx = T.prim_func(is_private=privacy, s_tir=s_tir, persistent=persistent) with prim_func_ctx: T.func_name(node.name) - remaining_body = _bind_forward_signature_vars(self, node) - if node.returns is not None: - ret_type = self.eval_expr(node.returns) - if callable(ret_type) and not isinstance(ret_type, Expr): - ret_type = ret_type() - if isinstance(ret_type, Expr): - ret_type = ret_type.ty - T.func_ret(ret_type) + declaration_dtypes = _function_prim_var_declarations(node) with self.with_dispatch_token("tirx"): # TODO: handle different types of arguments: # - vararg: arg | None @@ -733,6 +797,14 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None: if isinstance(ann, tvm.tirx.Var) and not is_buffer_var(ann): self.var_table.add(arg.arg, ann) + if node.returns is not None: + ret_type = _eval_signature_annotation(self, node.returns, declaration_dtypes) + if callable(ret_type) and not isinstance(ret_type, Expr): + ret_type = ret_type() + if isinstance(ret_type, Expr): + ret_type = ret_type.ty + T.func_ret(ret_type) + for arg in node.args.args: if arg.annotation is None: self.report_error(arg, "Type annotation required for function parameters.") @@ -740,7 +812,9 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None: ann = evaluated_annotations[arg.arg] else: try: - ann = self.eval_expr(arg.annotation) + ann = _eval_signature_annotation( + self, arg.annotation, declaration_dtypes + ) except Exception: # pylint: disable=broad-except ann = func_annotation.get(arg.arg, None) if ann is None: @@ -774,7 +848,7 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None: continue param = T.arg(arg.arg, ann) self.var_table.add(arg.arg, param) - self.visit_body(remaining_body) + self.visit_body(node.body) self.function_annotations = supplied_annotation @@ -1015,13 +1089,7 @@ def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar ret_type = None with self.var_table.with_frame(): - _bind_forward_signature_vars(self, node) - if node.returns is not None: - ret_type = self.eval_expr(node.returns) - if callable(ret_type) and not isinstance(ret_type, Expr): - ret_type = ret_type() - if isinstance(ret_type, Expr): - ret_type = ret_type.ty + declaration_dtypes = _function_prim_var_declarations(node) evaluated_annotations = {} for arg in node.args.args: @@ -1037,6 +1105,13 @@ def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar if isinstance(ann, tvm.tirx.Var) and not is_buffer_var(ann): self.var_table.add(arg.arg, ann) + if node.returns is not None: + ret_type = _eval_signature_annotation(self, node.returns, declaration_dtypes) + if callable(ret_type) and not isinstance(ret_type, Expr): + ret_type = ret_type() + if isinstance(ret_type, Expr): + ret_type = ret_type.ty + arg_annotations = [] for arg in node.args.args: if arg.annotation is None: @@ -1045,7 +1120,7 @@ def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar ann = evaluated_annotations[arg.arg] else: try: - ann = self.eval_expr(arg.annotation) + ann = _eval_signature_annotation(self, arg.annotation, declaration_dtypes) if callable(ann) and not isinstance(ann, Expr): ann = ann() except Exception: # pylint: disable=broad-except diff --git a/src/backend/trn/codegen/codegen_trn.cc b/src/backend/trn/codegen/codegen_trn.cc index ed96deb5aea3..d973e10e8cae 100644 --- a/src/backend/trn/codegen/codegen_trn.cc +++ b/src/backend/trn/codegen/codegen_trn.cc @@ -115,7 +115,7 @@ void CodeGenTrainium::AddFunction(const GlobalVar& gvar, const PrimFunc& func) { LOG(FATAL) << "Trainium codegen currently only support buffer arguments"; }; std::string vid = AllocVarID(v.get()); - if (auto buffer = tirx::BufferParamMap(func->params).Get(v)) { + if (auto buffer = v.as()) { var_idmap_[buffer.value().get()] = vid; } if (i >= static_cast(num_inputs.value())) { diff --git a/src/backend/trn/transform/lower_trainium_layout.cc b/src/backend/trn/transform/lower_trainium_layout.cc index feef2333ca1a..2c78845eed76 100644 --- a/src/backend/trn/transform/lower_trainium_layout.cc +++ b/src/backend/trn/transform/lower_trainium_layout.cc @@ -55,29 +55,34 @@ static bool IsTrainiumLayout(const TileLayoutNode* layout) { class TrainiumLayoutApplier : public arith::IRMutatorWithAnalyzer { public: - static std::pair> Lower( - const Stmt& stmt, const ffi::Map buffer_map) { + static std::pair> Lower(const Stmt& stmt, const ffi::Array& params) { arith::Analyzer ana; TrainiumLayoutApplier storage_lower(ana); - std::unordered_map new_buffer_map; + ffi::Array new_params; + new_params.reserve(params.size()); std::vector> param_flattened_buffers; - for (const auto& kv : buffer_map) { - if (kv.second->layout.has_value()) { - BufferVar flattened = storage_lower.GetFlattenedBuffer(kv.second); - auto type = CopyBufferType(kv.second); + for (const Var& param : params) { + auto buffer = param.as(); + if (!buffer) { + new_params.push_back(param); + continue; + } + if (buffer.value()->layout.has_value()) { + BufferVar flattened = storage_lower.GetFlattenedBuffer(buffer.value()); + auto type = CopyBufferType(buffer.value()); type->layout = std::nullopt; - BufferVar buffer = RebuildBufferVar(kv.second, std::move(type)); - param_flattened_buffers.emplace_back(flattened, buffer); - new_buffer_map[kv.first] = buffer; + BufferVar source = RebuildBufferVar(buffer.value(), std::move(type)); + param_flattened_buffers.emplace_back(flattened, source); + new_params.push_back(source.var()); } else { - new_buffer_map[kv.first] = kv.second; + new_params.push_back(buffer.value().var()); } } auto new_stmt = storage_lower(stmt); for (const auto& [buf, source] : param_flattened_buffers) { new_stmt = SeqStmt::Flatten(DeclBuffer(buf, source.data()), std::move(new_stmt)); } - return std::make_pair(new_stmt, ffi::Map(new_buffer_map)); + return std::make_pair(new_stmt, new_params); } protected: @@ -346,16 +351,7 @@ namespace transform { Pass LowerTrainiumLayout() { auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) { auto* n = f.CopyOnWrite(); - auto [body, buffer_map] = - TrainiumLayoutApplier::Lower(n->body, tirx::BufferParamMap(n->params)); - ffi::Array params; - for (const Var& param : n->params) { - if (auto buffer = buffer_map.Get(param)) { - params.push_back(buffer.value().var()); - } else { - params.push_back(param); - } - } + auto [body, params] = TrainiumLayoutApplier::Lower(n->body, n->params); n->body = std::move(body); n->params = std::move(params); n->body = TrainiumBufferOffsetRemover::Remove(n->body); diff --git a/src/relax/transform/dataflow_inplace.cc b/src/relax/transform/dataflow_inplace.cc index ca26a22a05f8..b7cbcaa9379f 100644 --- a/src/relax/transform/dataflow_inplace.cc +++ b/src/relax/transform/dataflow_inplace.cc @@ -956,7 +956,7 @@ class ModuleInplaceTransformer : public ExprMutator { // var's buffers // 2. For each output var, replace its instances with the corresponding inplace index var // 3. Do the same for the *buffer vars* corresponding to the output vars - // 4. Remove the output vars from the param list and buffer map + // 4. Remove the output vars from the param list ffi::Map buffer_subst_map; ffi::Map var_subst_map; for (size_t i = 0; i < num_outs; i++) { @@ -966,8 +966,8 @@ class ModuleInplaceTransformer : public ExprMutator { var_subst_map.Set(output_var, inplace_var); // also do the same with the buffer vars - auto output_buffer = tirx::BufferParamMap(old_primfunc->params).at(output_var); - auto inplace_buffer = tirx::BufferParamMap(old_primfunc->params).at(inplace_var); + auto output_buffer = output_var.as_or_throw(); + auto inplace_buffer = inplace_var.as_or_throw(); var_subst_map.Set(output_buffer.var(), inplace_buffer.var()); buffer_subst_map.Set(output_buffer, inplace_buffer); } @@ -982,19 +982,13 @@ class ModuleInplaceTransformer : public ExprMutator { return std::nullopt; }); - // remove the now-unused outputs from the buffer map - auto new_buffer_map = tirx::BufferParamMap(old_primfunc->params); - for (size_t i = 0; i < num_outs; i++) { - new_buffer_map.erase(old_primfunc->params[num_params - num_outs + i]); - } - // now get rid of the last num_outputs arguments // (couldn't do earlier or else it would have thrown off the indexing) ffi::Array new_params(old_primfunc->params.begin(), old_primfunc->params.begin() + (num_params - num_outs)); - tirx::PrimFunc new_primfunc(new_params, new_body, old_primfunc->ret_type, new_buffer_map, - old_primfunc->attrs, old_primfunc->span); + tirx::PrimFunc new_primfunc(new_params, new_body, old_primfunc->ret_type, old_primfunc->attrs, + old_primfunc->span); // note: this might be a good time to get rid of the old legalized function, but we don't do it // now because later ops might need the same one. Instead, we will clean up at the end diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc index d801133c75b7..6fb6e9829a39 100644 --- a/src/relax/transform/fuse_tir.cc +++ b/src/relax/transform/fuse_tir.cc @@ -465,7 +465,6 @@ class RelaxToTIRVarMapCollector : public ExprVisitor { void CollectVarMapping(const CallNode* call, const Expr& lhs_var, bool in_place) { GlobalVar gv = call->args[0].as_or_throw(); tirx::PrimFunc prim_func_ = mod_->Lookup(gv).as_or_throw(); - const auto& buffer_map = tirx::BufferParamMap(prim_func_->params); const auto& tir_args = prim_func_->params; const auto& relax_args = call->args[1].as_or_throw()->fields; @@ -504,7 +503,7 @@ class RelaxToTIRVarMapCollector : public ExprVisitor { }; for (size_t i = 0; i < tir_args.size(); ++i) { const auto& tir_var = tir_args[i]; - if (auto tir_buffer = buffer_map.Get(tir_var)) { + if (auto tir_buffer = tir_var.as()) { if (i < num_inputs) { const auto& relax_var = relax_args[i]; ValidateBufferCompatibility(tir_buffer.value(), relax_var); diff --git a/src/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc index 290be4c22cef..01953213ffc0 100644 --- a/src/relax/transform/split_call_tir_by_pattern.cc +++ b/src/relax/transform/split_call_tir_by_pattern.cc @@ -391,7 +391,11 @@ class TIRPatternMatcher { for (const TIRPattern& pattern : patterns_) { tirx::PrimFunc pattern_func = pattern; ffi::Array pattern_symbolic_vars; - int buffer_count = tirx::BufferParamMap(pattern_func->params).size(); + int buffer_count = 0; + while (buffer_count < static_cast(pattern_func->params.size()) && + pattern_func->params[buffer_count]->ty.as()) { + ++buffer_count; + } for (int i = buffer_count; i < static_cast(pattern_func->params.size()); i++) { pattern_symbolic_vars.push_back(pattern_func->params[i]); } @@ -738,8 +742,7 @@ class SplitMutator : public ExprMutator { if (lib_func->IsInstance()) return ffi::GetRef(op); TVM_FFI_ICHECK(lib_func->IsInstance()); builder_->UpdateFunction(gv, lib_func); - tirx::BufferVar intermediate_buffer = - tirx::BufferParamMap(func1->params).at(func1->params.back()); + tirx::BufferVar intermediate_buffer = func1->params.back().as_or_throw(); PrimType dtype = intermediate_buffer->dtype; Call call1(Type::Missing(), call_dps_packed_, {lib_func, Tuple(args1)}, call->attrs, {TensorType(ShapeExpr(intermediate_buffer->shape), dtype)}); diff --git a/src/relax/transform/split_layout_rewrite_preproc.cc b/src/relax/transform/split_layout_rewrite_preproc.cc index 94ab0dcb3a99..b2a47810a9a6 100644 --- a/src/relax/transform/split_layout_rewrite_preproc.cc +++ b/src/relax/transform/split_layout_rewrite_preproc.cc @@ -189,8 +189,7 @@ class SplitPrimFuncLayoutRewrite : public StmtMutator { const BufferVar& preproc_buffer = op->reads[0]->buffer; int buffer_index = -1; for (size_t i = 0; i < original_func_->params.size(); ++i) { - const BufferVar& buffer = - tirx::BufferParamMap(original_func_->params)[original_func_->params[i]]; + BufferVar buffer = original_func_->params[i].as_or_throw(); if (buffer == preproc_buffer) { buffer_index = i; break; diff --git a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc index bae15155981e..4fb7535c64b4 100644 --- a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc +++ b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc @@ -45,9 +45,10 @@ class LCADetector : public StmtExprVisitor { public: static ffi::Map> Detect(const PrimFunc& func) { LCADetector detector; - for (const auto& kv : tirx::BufferParamMap(func->params)) { - const BufferVar& buffer = kv.second; - detector.buffer_var_map_.emplace(buffer.get(), buffer.get()); + for (const Var& param : func->params) { + if (auto buffer = param.as()) { + detector.buffer_var_map_.emplace(buffer.value().get(), buffer.value().get()); + } } // The root node must be explicitly present in the list of diff --git a/src/s_tir/backend/adreno/texture_flatten.cc b/src/s_tir/backend/adreno/texture_flatten.cc index 7c042a2518ee..3cb8c28136fe 100644 --- a/src/s_tir/backend/adreno/texture_flatten.cc +++ b/src/s_tir/backend/adreno/texture_flatten.cc @@ -48,11 +48,12 @@ using runtime::IsTextureStorage; class TextureLoweringBase : public StmtExprMutator { public: - explicit TextureLoweringBase(const ffi::Map& extern_buffer_map, - IRVisitorWithAnalyzer* bound_analyzer) + explicit TextureLoweringBase(const ffi::Array& params, IRVisitorWithAnalyzer* bound_analyzer) : bound_analyzer_{bound_analyzer} { - for (auto kv : extern_buffer_map) { - extern_buf_.insert(kv.second); + for (const Var& param : params) { + if (auto buffer = param.as()) { + extern_buf_.insert(buffer.value()); + } } } @@ -84,9 +85,8 @@ class TextureLoweringBase : public StmtExprMutator { class TextureFlattener : public TextureLoweringBase { public: using StmtExprMutator::VisitStmt_; - explicit TextureFlattener(const ffi::Map& extern_buffer_map, - IRVisitorWithAnalyzer* bound_analyzer) - : TextureLoweringBase(extern_buffer_map, bound_analyzer) {} + explicit TextureFlattener(const ffi::Array& params, IRVisitorWithAnalyzer* bound_analyzer) + : TextureLoweringBase(params, bound_analyzer) {} Stmt VisitStmt_(const BufferStoreNode* op) final { Stmt stmt = StmtExprMutator::VisitStmt_(op); @@ -159,8 +159,7 @@ PrimFunc TextureFlattenHandler(PrimFunc func) { auto fptr = func.CopyOnWrite(); IRVisitorWithAnalyzer bound_analyzer; bound_analyzer(fptr->body); - fptr->body = - TextureFlattener(tirx::BufferParamMap(fptr->params), &bound_analyzer)(std::move(fptr->body)); + fptr->body = TextureFlattener(fptr->params, &bound_analyzer)(std::move(fptr->body)); return func; } diff --git a/src/s_tir/meta_schedule/arg_info.cc b/src/s_tir/meta_schedule/arg_info.cc index 375b24a69448..c36a952961e4 100644 --- a/src/s_tir/meta_schedule/arg_info.cc +++ b/src/s_tir/meta_schedule/arg_info.cc @@ -96,10 +96,9 @@ ffi::Array ArgInfo::FromPrimFunc(const tirx::PrimFunc& func) { ffi::Array result; result.reserve(func->params.size()); for (const tirx::Var& arg : func->params) { - if (ffi::Optional _buffer = tirx::BufferParamMap(func->params).Get(arg)) { - tirx::BufferVar buffer = _buffer.value(); - result.push_back(TensorInfo(/*dtype=*/buffer->dtype->dtype, - /*shape=*/AsVector(buffer->shape))); + if (auto buffer = arg.as()) { + result.push_back(TensorInfo(/*dtype=*/buffer.value()->dtype->dtype, + /*shape=*/AsVector(buffer.value()->shape))); } else { TVM_FFI_THROW(ValueError) << "Unsupported argument type: " << arg; } diff --git a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc index b6050584fe85..acf609f61cd0 100644 --- a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc +++ b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc @@ -1288,8 +1288,10 @@ class PerStoreFeatureCollector : private StmtVisitor { for (const auto& kv : mod->functions) { if (const PrimFuncNode* func = kv.second.as()) { collector(func->body); - for (const auto& it : tirx::BufferParamMap(func->params)) { - collector.HandleBufferAlloc(it.second); + for (const Var& param : func->params) { + if (auto buffer = param.as()) { + collector.HandleBufferAlloc(buffer.value()); + } } } } diff --git a/src/s_tir/meta_schedule/postproc/rewrite_layout.cc b/src/s_tir/meta_schedule/postproc/rewrite_layout.cc index 49dd31a460f0..2dce55cbe864 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_layout.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_layout.cc @@ -129,7 +129,7 @@ ffi::Array CollectLayoutFreeBuffers(const PrimFuncNode* func) { for (int64_t index : layout_free_buffer_index) { TVM_FFI_ICHECK(static_cast(index) < func->params.size()); const Var& param = func->params[index]; - layout_free_buffers.push_back(tirx::BufferParamMap(func->params).at(param)); + layout_free_buffers.push_back(param.as_or_throw()); } LayoutFreeBufferCollector collector; diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc index f86597264ae3..7d1c87055326 100644 --- a/src/s_tir/schedule/analysis/analysis.cc +++ b/src/s_tir/schedule/analysis/analysis.cc @@ -1314,10 +1314,11 @@ void AddShapeVarBounds(const ScheduleState& state, const StmtSRefNode* sref, sref = sref->parent; } const PrimFuncNode* f = GetRootPrimFunc(state->mod, sref->stmt, nullptr); - for (const auto& kv : tirx::BufferParamMap(f->params)) { - const BufferVar& buffer = kv.second; - for (const PrimExpr& e : buffer->shape) { - analyzer->MarkGlobalNonNegValue(e); + for (const Var& param : f->params) { + if (auto buffer = param.as()) { + for (const PrimExpr& e : buffer.value()->shape) { + analyzer->MarkGlobalNonNegValue(e); + } } } } diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc b/src/s_tir/schedule/primitive/blockize_tensorize.cc index f3cd1fb57750..16a551df74a2 100644 --- a/src/s_tir/schedule/primitive/blockize_tensorize.cc +++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc @@ -796,8 +796,8 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int std::unordered_map impl2desc; TVM_FFI_ICHECK_EQ(intrin_desc->params.size(), intrin_impl->params.size()); for (int i = 0, n = intrin_desc->params.size(); i < n; ++i) { - const BufferVar& desc = tirx::BufferParamMap(intrin_desc->params)[intrin_desc->params[i]]; - const BufferVar& impl = tirx::BufferParamMap(intrin_impl->params)[intrin_impl->params[i]]; + BufferVar desc = intrin_desc->params[i].as_or_throw(); + BufferVar impl = intrin_impl->params[i].as_or_throw(); impl2desc[impl] = desc; } std::unordered_map impl2cur; @@ -821,7 +821,7 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int ffi::Array match_buffer_regions; match_buffer_regions.reserve(intrin_impl->params.size()); for (int i = 0, n = intrin_impl->params.size(); i < n; ++i) { - const BufferVar& impl = tirx::BufferParamMap(intrin_impl->params).at(intrin_impl->params[i]); + BufferVar impl = intrin_impl->params[i].as_or_throw(); const BufferVar& cur = impl2cur.at(impl); const ffi::Array& old_region = impl2region.at(impl); const std::vector& indices_base = comparator.buffer_indices_.at(cur); diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc b/src/s_tir/schedule/primitive/layout_transformation.cc index 872f99ab22cf..39f6f57b7a2b 100644 --- a/src/s_tir/schedule/primitive/layout_transformation.cc +++ b/src/s_tir/schedule/primitive/layout_transformation.cc @@ -1236,16 +1236,15 @@ void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_ IRModuleNode* new_mod = self->mod.CopyOnWrite(); ffi::MapObj* new_map = new_mod->functions.CopyOnWrite(); - ffi::Map new_buffer_map; - for (auto [var, buffer] : tirx::BufferParamMap(old_func->params)) { - if (buffer.same_as(old_buffer)) { - buffer = new_buffer; + ffi::Array new_params = old_func->params.Map([&](Var param) -> Var { + if (auto buffer = param.as(); buffer && buffer.value().same_as(old_buffer)) { + return new_buffer.var(); } - new_buffer_map.Set(var, buffer); - } + return param; + }); - PrimFunc ref_new_func(old_func->params, old_func->body, old_func->ret_type, new_buffer_map, - old_func->attrs, old_func->span); + PrimFunc ref_new_func(new_params, old_func->body, old_func->ret_type, old_func->attrs, + old_func->span); new_map->at(g_var) = std::move(ref_new_func); } diff --git a/src/s_tir/transform/inject_permuted_layout.cc b/src/s_tir/transform/inject_permuted_layout.cc index f83d9b45f116..0ed602945ad6 100644 --- a/src/s_tir/transform/inject_permuted_layout.cc +++ b/src/s_tir/transform/inject_permuted_layout.cc @@ -70,10 +70,10 @@ class PermutedLayoutInjector : private IRMutatorWithAnalyzer { private: explicit PermutedLayoutInjector(PrimFunc func, const Analyzer& analyzer) : IRMutatorWithAnalyzer(analyzer) { - auto param_buffers = tirx::BufferParamMap(func->params); - buffer_map_.insert(param_buffers.begin(), param_buffers.end()); - for (const auto& [_, buffer] : param_buffers) { - buffer_map_.insert({buffer.var(), buffer}); + for (const Var& param : func->params) { + if (auto buffer = param.as()) { + buffer_map_.insert({buffer.value().var(), buffer.value()}); + } } } diff --git a/src/s_tir/transform/inject_software_pipeline.cc b/src/s_tir/transform/inject_software_pipeline.cc index 1cb9a48de07f..35140437656c 100644 --- a/src/s_tir/transform/inject_software_pipeline.cc +++ b/src/s_tir/transform/inject_software_pipeline.cc @@ -1065,9 +1065,10 @@ class PipelineInjector : private StmtExprMutator { static Stmt Inject(const PrimFunc& func) { auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); PipelineInjector injector(global_symbol); - for (const auto& kv : tirx::BufferParamMap(func->params)) { - const BufferVar& buffer = kv.second; - injector.buffer_data_to_buffer_.Set(buffer.var(), buffer); + for (const Var& param : func->params) { + if (auto buffer = param.as()) { + injector.buffer_data_to_buffer_.Set(buffer.value().var(), buffer.value()); + } } injector.fragment_info_ = GetTensorCoreFragmentInfo(func->body); return injector(func->body); diff --git a/src/s_tir/transform/lower_thread_allreduce.cc b/src/s_tir/transform/lower_thread_allreduce.cc index 68b4678ec96b..a4f1d0dcda70 100644 --- a/src/s_tir/transform/lower_thread_allreduce.cc +++ b/src/s_tir/transform/lower_thread_allreduce.cc @@ -58,14 +58,10 @@ ffi::Optional GetBufferDataVar(const ffi::Any& data) { class ThreadAllreduceBuilder final : public StmtExprMutator { public: - explicit ThreadAllreduceBuilder(const TargetNode* target, const ffi::Array& params, - const ffi::Map& buffer_map) + explicit ThreadAllreduceBuilder(const TargetNode* target, const ffi::Array& params) : target_(target), warp_size_(target->GetAttr("thread_warp_size", 1).value()), max_num_threads_(target->GetAttr("max_num_threads", -1).value()) { - for (const auto& [_, buffer] : buffer_map) { - buffer_aliases_.Set(buffer.var(), buffer.var()); - } for (const Var& param : params) { if (param->ty.as()) { buffer_aliases_.Set(param, param); @@ -966,8 +962,7 @@ Pass LowerThreadAllreduce() { auto target = f->GetAttr(tvm::attr::kTarget); TVM_FFI_ICHECK(target.has_value()) << "LowerThreadAllreduce: Require the target attribute"; const TargetNode* target_node = target.as(); - ThreadAllreduceBuilder thread_all_reduce(target_node, f->params, - tirx::BufferParamMap(f->params)); + ThreadAllreduceBuilder thread_all_reduce(target_node, f->params); n->body = thread_all_reduce(n->body); // Post-process: apply deferred remappings for flat IR DeferredRemapper remapper(thread_all_reduce.alloc_remap_, thread_all_reduce.var_remap_, diff --git a/src/s_tir/transform/plan_update_buffer_allocation_location.cc b/src/s_tir/transform/plan_update_buffer_allocation_location.cc index ee6dc4a7eeba..af33de6f20fc 100644 --- a/src/s_tir/transform/plan_update_buffer_allocation_location.cc +++ b/src/s_tir/transform/plan_update_buffer_allocation_location.cc @@ -56,8 +56,10 @@ class BufferAllocateOrderCollector : public StmtExprVisitor { public: static ffi::Array Collect(const PrimFunc& func) { BufferAllocateOrderCollector collector; - for (const auto& kv : tirx::BufferParamMap(func->params)) { - collector.buffer_alloc_recorder_.push_back(kv.second); + for (const Var& param : func->params) { + if (auto buffer = param.as()) { + collector.buffer_alloc_recorder_.push_back(buffer.value()); + } } collector(func->body); return std::move(collector.buffer_alloc_recorder_); @@ -114,10 +116,11 @@ class BufferAllocationLocator : public StmtExprMutator { collector(func->body); managed_allocations_ = collector.managed_allocations; - for (const auto& kv : tirx::BufferParamMap(func->params)) { - const BufferVar& buffer = kv.second; - arg_buffer_vars.emplace(buffer.get()); - buffer_data_to_buffer_.Set(buffer.var(), buffer); + for (const Var& param : func->params) { + if (auto buffer = param.as()) { + arg_buffer_vars.emplace(buffer.value().get()); + buffer_data_to_buffer_.Set(buffer.value().var(), buffer.value()); + } } // create buffers to be allocated at each stmts for (const auto& buffer : buffer_alloc_recorder) { diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc index 07dd1a054c05..dc454de5ad3c 100644 --- a/src/target/llvm/codegen_llvm.cc +++ b/src/target/llvm/codegen_llvm.cc @@ -269,8 +269,10 @@ llvm::Function* CodeGenLLVM::DeclareFunctionInternal(const GlobalVar& gvar, cons return it->second; } - TVM_FFI_ICHECK_EQ(tirx::BufferParamMap(func->params).size(), 0U) - << "Cannot codegen BufferType-annotated parameters; please lower them first"; + for (const Var& param : func->params) { + TVM_FFI_ICHECK(!param->ty.as()) + << "Cannot codegen BufferType-annotated parameter " << param << "; please lower it first"; + } std::vector param_types; is_restricted_ = func->HasNonzeroAttr(tirx::attr::kNoAlias); diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc index 9d23f16d0473..7a8c4ae80b18 100644 --- a/src/te/operation/create_primfunc.cc +++ b/src/te/operation/create_primfunc.cc @@ -132,10 +132,8 @@ class LayoutFreePlaceholdersNormalizer : public StmtMutator { public: PrimFunc Process(PrimFunc func) { for (int i = 0, n = func->params.size(); i < n; ++i) { - if (auto v = func->params[i].as()) { - if (ffi::Optional buffer = tirx::BufferParamMap(func->params).Get(v.value())) { - buffer2index_[buffer.value()] = i; - } + if (auto buffer = func->params[i].as()) { + buffer2index_[buffer.value()] = i; } } PrimFuncNode* f = func.CopyOnWrite(); diff --git a/src/tirx/analysis/verify_memory.cc b/src/tirx/analysis/verify_memory.cc index ae4d0211e281..569428461439 100644 --- a/src/tirx/analysis/verify_memory.cc +++ b/src/tirx/analysis/verify_memory.cc @@ -103,8 +103,8 @@ class MemoryAccessVerifier final : protected StmtExprVisitor { /// Check if the value of a Variable comes from function argument. bool IsFromFunctionArgs(const VarNode* var) const { const VarNode* V = var; - for (auto kv : tirx::BufferParamMap(func_->params)) { - if (V == kv.second.get()) return true; + for (const Var& param : func_->params) { + if (param->ty.as() && V == param.get()) return true; } while (true) { diff --git a/src/tirx/analysis/verify_ssa.cc b/src/tirx/analysis/verify_ssa.cc index 5ba2013be2bf..f9b56dd9e7b5 100644 --- a/src/tirx/analysis/verify_ssa.cc +++ b/src/tirx/analysis/verify_ssa.cc @@ -93,8 +93,10 @@ class SSAVerifier final : public StmtExprVisitor { MarkDef(param, param); } - for (auto kv : tirx::BufferParamMap(func->params)) { - this->DefineBuffer(kv.second); + for (const Var& param : func->params) { + if (auto buffer = param.as()) { + this->DefineBuffer(buffer.value()); + } } this->VisitStmt(func->body); } diff --git a/src/tirx/ir/data_type_rewriter.cc b/src/tirx/ir/data_type_rewriter.cc index 45b84e4edd8c..d8a6fff6cc64 100644 --- a/src/tirx/ir/data_type_rewriter.cc +++ b/src/tirx/ir/data_type_rewriter.cc @@ -635,8 +635,10 @@ PrimFunc IndexDataTypeNormalizer::Rewrite(PrimFunc func) { buffer_remap_.clear(); ivmap_.clear(); // start rewrite - for (const auto& [_, buffer] : tirx::BufferParamMap(func->params)) { - VisitBufferDef(buffer, /*alloc_data=*/true); + for (const Var& param : func->params) { + if (auto buffer = param.as()) { + VisitBufferDef(buffer.value(), /*alloc_data=*/true); + } } // remap params bool is_enabled = true; diff --git a/src/tirx/ir/script/script_complete.cc b/src/tirx/ir/script/script_complete.cc index 1d4ffdb7ef9e..3e7bd254e7c5 100644 --- a/src/tirx/ir/script/script_complete.cc +++ b/src/tirx/ir/script/script_complete.cc @@ -126,9 +126,10 @@ class ScriptCompleter : public StmtMutator { PrimFunc ScriptComplete(PrimFunc func, const ffi::Array& root_allocates, bool s_tir) { ffi::Map buffer_var_map; - for (const auto& pair : tirx::BufferParamMap(func->params)) { - const BufferVar& buffer = pair.second; - buffer_var_map.Set(buffer.var(), buffer); + for (const Var& param : func->params) { + if (auto buffer = param.as()) { + buffer_var_map.Set(buffer.value().var(), buffer.value()); + } } for (const auto& alloc : root_allocates) { buffer_var_map.Set(alloc.var(), alloc); diff --git a/src/tirx/ir/specialize.cc b/src/tirx/ir/specialize.cc index e0e70024813e..3890829dc14a 100644 --- a/src/tirx/ir/specialize.cc +++ b/src/tirx/ir/specialize.cc @@ -96,32 +96,23 @@ class PrimFuncSpecializer : public StmtExprMutator { } } - // Updating BufferVar map - ffi::Map buffer_map; - bool buffer_map_updated = false; - for (const auto& it : tirx::BufferParamMap(f->params)) { - const Var& var = it.first; - const BufferVar& buffer = it.second; - BufferVar new_buffer = specializer.MutateBuffer(buffer); - buffer_map.Set(var, new_buffer); - if (!new_buffer.same_as(buffer)) { - buffer_map_updated = true; - specializer.buffer_map_[buffer] = new_buffer; - } - } - // Updating parameters ffi::Array params; bool param_updated = false; for (const auto& var : f->params) { + Var new_var = var; + if (auto buffer = var.as()) { + BufferVar new_buffer = specializer.MutateBuffer(buffer.value()); + new_var = new_buffer.var(); + if (!new_buffer.same_as(buffer.value())) { + param_updated = true; + specializer.buffer_map_[buffer.value()] = new_buffer; + } + } // Remove parmeters which has been specialized. if (var_map.find(var) == var_map.end() || specializer.constrained_buffer_params_.count(var.get())) { - if (auto buffer = buffer_map.Get(var)) { - params.push_back(buffer.value().var()); - } else { - params.push_back(var); - } + params.push_back(new_var); } else { param_updated = true; } @@ -130,7 +121,7 @@ class PrimFuncSpecializer : public StmtExprMutator { // Updating function body Stmt body = specializer(f->body); - if (param_updated || buffer_map_updated || !f->body.same_as(body)) { + if (param_updated || !f->body.same_as(body)) { return PrimFunc(params, body, f->ret_type, f->attrs, f->span); } else { return f; diff --git a/src/tirx/script/printer/buffer.cc b/src/tirx/script/printer/buffer.cc index 37f72b2d3377..ae2001d28794 100644 --- a/src/tirx/script/printer/buffer.cc +++ b/src/tirx/script/printer/buffer.cc @@ -85,10 +85,17 @@ ffi::Map BufferAttrs(tirx::BufferVar buffer, const AccessP for (int i = 0; i < n; ++i) { PrimExpr e = shape[i]; AccessPath e_p = shape_p->ArrayItem(i); + bool contains_new_var = false; + tirx::PostOrderVisit(e, [&](const ffi::ObjectRef& obj) { + if (const auto* var = obj.as()) { + contains_new_var = contains_new_var || !d->IsVarDefined(ffi::GetRef(var)); + } + }); if (is_new_var(e)) { add_out_of_line_var_def(e.as_or_throw(), e_p); } - results.push_back(d->AsDoc(e, e_p)); + ExprDoc result = d->AsDoc(e, e_p); + results.push_back(contains_new_var ? ExprStringDoc(result, e_p) : result); } kwargs.Set("shape", TupleDoc(results)); } diff --git a/src/tirx/script/printer/function.cc b/src/tirx/script/printer/function.cc index 76c059aa122f..3e30f004f9dc 100644 --- a/src/tirx/script/printer/function.cc +++ b/src/tirx/script/printer/function.cc @@ -105,6 +105,13 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) ExprDoc a = d->AsDoc(var->ty, var_p->Attr("ty")); args.push_back(AssignDoc(scalar_param_docs.at(var.get()), std::nullopt, a)); } + ffi::Optional ret_type = std::nullopt; + if (!func->ret_type.IsMissing()) { + const auto* as_tuple = func->ret_type.as(); + if (!as_tuple || as_tuple->fields.size()) { + ret_type = d->AsDoc(func->ret_type, p->Attr("ret_type")); + } + } // Step 2. Handle `func->attrs` if (!func->attrs->dict.empty()) { // for global symbol, don't display it if it matches the func name @@ -172,13 +179,6 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) } else { AsDocBody(func->body, p->Attr("body"), f->get(), d); } - ffi::Optional ret_type = std::nullopt; - if (!func->ret_type.IsMissing()) { - const auto* as_tuple = func->ret_type.as(); - if (!as_tuple || as_tuple->fields.size()) { - ret_type = d->AsDoc(func->ret_type, p->Attr("ret_type")); - } - } // Step 5. Determine if we need to display the private annotation in the decorator ExprDoc decorator = TIR(d, "prim_func"); ffi::Array kwargs_keys; diff --git a/src/tirx/transform/flatten_buffer.cc b/src/tirx/transform/flatten_buffer.cc index b57e05f3e7fe..d901d914342e 100644 --- a/src/tirx/transform/flatten_buffer.cc +++ b/src/tirx/transform/flatten_buffer.cc @@ -63,9 +63,11 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { arith::Analyzer ana; auto pass = BufferFlattener(ana); pass.MarkBufferParamShapes(func); - for (const auto& [param, buffer] : tirx::BufferParamMap(func->params)) { - pass.extern_buffers_.insert(buffer); - pass.Define(buffer); + for (const Var& param : func->params) { + if (auto buffer = param.as()) { + pass.extern_buffers_.insert(buffer.value()); + pass.Define(buffer.value()); + } } auto body = pass.VisitStmt(func->body); @@ -73,13 +75,11 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { // for validation of user-provided arguments. The flattened buffers used // in the updated function body alias the argument buffers. for (size_t i = func->params.size(); i > 0; i--) { - auto handle = func->params[i - 1]; - if (auto opt = tirx::BufferParamMap(func->params).Get(handle)) { - auto old_buf = opt.value(); - if (pass.buffers_used_.count(old_buf)) { - auto new_buf = pass.Lookup(old_buf).flattened; - if (!old_buf.same_as(new_buf)) { - body = SeqStmt::Flatten(DeclBuffer(new_buf, old_buf.data()), std::move(body)); + if (auto old_buf = func->params[i - 1].as()) { + if (pass.buffers_used_.count(old_buf.value())) { + auto new_buf = pass.Lookup(old_buf.value()).flattened; + if (!old_buf.value().same_as(new_buf)) { + body = SeqStmt::Flatten(DeclBuffer(new_buf, old_buf.value().data()), std::move(body)); } } } diff --git a/src/tirx/transform/force_narrow_index_to_i32.cc b/src/tirx/transform/force_narrow_index_to_i32.cc index 1e7e8ff810c1..2e444c273ba6 100644 --- a/src/tirx/transform/force_narrow_index_to_i32.cc +++ b/src/tirx/transform/force_narrow_index_to_i32.cc @@ -37,10 +37,12 @@ class Int32DTypeNarrower : public IndexDataTypeNormalizer { public: static PrimFunc RewriteDataType(PrimFunc func) { // Check if the integer parameter buffers have dtype other than int32. - for (auto it : tirx::BufferParamMap(func->params)) { - if (it.second->dtype.MatchesCode(DLDataTypeCode::kDLInt) && it.second->dtype.bits() > 32) { - TVM_FFI_THROW(InternalError) << "The buffer parameter " << it.second << " has dtype " - << it.second->dtype << ". The function is " << func; + for (const Var& param : func->params) { + if (auto buffer = param.as(); + buffer && buffer.value()->dtype.MatchesCode(DLDataTypeCode::kDLInt) && + buffer.value()->dtype.bits() > 32) { + TVM_FFI_THROW(InternalError) << "The buffer parameter " << buffer.value() << " has dtype " + << buffer.value()->dtype << ". The function is " << func; } } diff --git a/src/tirx/transform/inline_private_functions.cc b/src/tirx/transform/inline_private_functions.cc index c6ac2f3081c5..4b3cafb42d41 100644 --- a/src/tirx/transform/inline_private_functions.cc +++ b/src/tirx/transform/inline_private_functions.cc @@ -115,8 +115,9 @@ bool IsInlinablePrimFunc(const GlobalVar& gvar, const PrimFunc& prim_func, // We do not currently support inlining of functions that accept // buffer arguments. - bool has_buffer_arguments = tirx::BufferParamMap(prim_func->params).size(); - if (has_buffer_arguments) return false; + for (const Var& param : prim_func->params) { + if (param->ty.as()) return false; + } // We do not currently support inlining of schedulable TIR // functions. To support this use case, repeated names in @@ -228,10 +229,11 @@ class PrimFuncInliner : StmtExprMutator { << callee->params << "), but is called with " << args.size() << " arguments (" << args << ")"; - TVM_FFI_ICHECK(tirx::BufferParamMap(callee->params).empty()) - << "Inlining of PrimFuncs with buffer arguments is not yet supported, " - << "but callee " << gvar << " has BufferType-annotated parameters " - << tirx::BufferParamMap(callee->params); + for (const Var& param : callee->params) { + TVM_FFI_ICHECK(!param->ty.as()) + << "Inlining of PrimFuncs with buffer arguments is not yet supported, " + << "but callee " << gvar << " has BufferType-annotated parameter " << param; + } ffi::Map> param_map; for (size_t i = 0; i < callee->params.size(); i++) { diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc index 36e08c515048..cf129e91a145 100644 --- a/src/tirx/transform/ir_utils.cc +++ b/src/tirx/transform/ir_utils.cc @@ -110,49 +110,42 @@ class IRConvertSSA final : public StmtExprMutator { for (const auto& var : func->params) { defined_params.insert(var.get()); } - for (const auto& [var, buffer] : tirx::BufferParamMap(func->params)) { - static_cast(var); // gcc 7.x bug, https://gcc.gnu.org/bugzilla/show_bug.cgi?id=81767 + for (const Var& param : func->params) { + auto buffer = param.as(); + if (!buffer) continue; auto check_expr = [&](const PrimExpr& expr) { auto* var_ptr = expr.as(); if (!var_ptr) return; if (defined_params.count(var_ptr)) return; - // Buffer_map shape vars use "match" semantics: first occurrence + // Buffer-parameter shape vars use "match" semantics: first occurrence // defines the var, subsequent occurrences (in other buffers) are // just consistent uses of the same var -- not redefinitions. if (!defined_.count(var_ptr)) { defined_.insert(var_ptr); } }; - for (const auto& dim : buffer->shape) { + for (const auto& dim : buffer.value()->shape) { check_expr(dim); } - for (const auto& stride : buffer->strides) { + for (const auto& stride : buffer.value()->strides) { check_expr(stride); } - check_expr(buffer->elem_offset); + check_expr(buffer.value()->elem_offset); } } // Update the buffer parameters, based on the redefined parameters bool buffer_params_changed = false; - auto buffer_map = [&]() { - ffi::Map buffer_map; - bool made_change = false; - for (const auto& [var, buffer] : tirx::BufferParamMap(func->params)) { - auto new_var = GetRemappedVar(var); - auto new_buf = GetRemappedBuffer(buffer); - - made_change = made_change || !var.same_as(new_var) || !buffer.same_as(new_buf); - buffer_map.Set(new_var, new_buf); - } - if (made_change) { - buffer_params_changed = true; - return buffer_map; - } else { - return tirx::BufferParamMap(func->params); + for (size_t i = 0; i < func->params.size(); ++i) { + if (auto buffer = func->params[i].as()) { + BufferVar new_buffer = GetRemappedBuffer(buffer.value()); + if (!new_buffer.same_as(buffer.value()) || !params[i].same_as(new_buffer)) { + buffer_params_changed = true; + params.Set(i, new_buffer.var()); + } } - }(); + } auto attrs = [&]() -> DictAttrs { ffi::Map dict; @@ -179,15 +172,6 @@ class IRConvertSSA final : public StmtExprMutator { auto body = VisitStmt(func->body); - if (buffer_params_changed) { - params = params.Map([&](const Var& param) -> Var { - if (auto buffer = buffer_map.Get(param)) { - return buffer.value().var(); - } - return param; - }); - } - // If anything changed, update the returned function if (!params.same_as(func->params) || buffer_params_changed || !attrs.same_as(func->attrs) || !body.same_as(func->body)) { diff --git a/src/tirx/transform/lower_tirx_cleanup.cc b/src/tirx/transform/lower_tirx_cleanup.cc index 0bf0eb140d03..d20cf7a0ea04 100644 --- a/src/tirx/transform/lower_tirx_cleanup.cc +++ b/src/tirx/transform/lower_tirx_cleanup.cc @@ -46,36 +46,36 @@ namespace tirx { class LayoutApplier : public arith::IRMutatorWithAnalyzer { public: - static std::pair> Flatten( - const Stmt& stmt, const ffi::Array& params, - const ffi::Map buffer_map, const Target& target) { + static std::pair> Flatten(const Stmt& stmt, const ffi::Array& params, + const Target& target) { arith::Analyzer ana; LayoutApplier storage_lower(ana, target); + ffi::Array new_params; + new_params.reserve(params.size()); + std::vector> param_flattened_buffers; for (const Var& param : params) { - if (param->ty.as()) { - storage_lower.buffer_aliases_.Set(param, param); + auto buffer = param.as(); + if (!buffer) { + new_params.push_back(param); + continue; } - } - std::unordered_map new_buffer_map; - std::vector> param_flattened_buffers; - for (const auto& kv : buffer_map) { - storage_lower.buffer_aliases_.Set(kv.second.var(), kv.second.var()); - if (kv.second->layout.has_value()) { - BufferVar flattened = storage_lower.GetFlattenedBuffer(kv.second); - auto type = CopyBufferType(kv.second); + storage_lower.buffer_aliases_.Set(buffer.value().var(), buffer.value().var()); + if (buffer.value()->layout.has_value()) { + BufferVar flattened = storage_lower.GetFlattenedBuffer(buffer.value()); + auto type = CopyBufferType(buffer.value()); type->layout = std::nullopt; - BufferVar buffer = RebuildBufferVar(kv.second, std::move(type)); - param_flattened_buffers.emplace_back(flattened, buffer); - new_buffer_map[kv.first] = buffer; + BufferVar source = RebuildBufferVar(buffer.value(), std::move(type)); + param_flattened_buffers.emplace_back(flattened, source); + new_params.push_back(source.var()); } else { - new_buffer_map[kv.first] = kv.second; + new_params.push_back(buffer.value().var()); } } auto new_stmt = storage_lower(stmt); for (const auto& [buf, source] : param_flattened_buffers) { new_stmt = SeqStmt::Flatten(DeclBuffer(buf, source.data()), std::move(new_stmt)); } - return std::make_pair(new_stmt, ffi::Map(new_buffer_map)); + return std::make_pair(new_stmt, new_params); } protected: @@ -404,16 +404,7 @@ Pass LowerTIRxCleanup() { auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) { Target target = ResolveTarget(f); auto* n = f.CopyOnWrite(); - auto [body, buffer_map] = - LayoutApplier::Flatten(n->body, n->params, BufferParamMap(n->params), target); - ffi::Array params; - for (const Var& param : n->params) { - if (auto buffer = buffer_map.Get(param)) { - params.push_back(buffer.value().var()); - } else { - params.push_back(param); - } - } + auto [body, params] = LayoutApplier::Flatten(n->body, n->params, target); n->body = std::move(body); n->params = std::move(params); n->body = BufferOffsetRemover::Remove(n->body); diff --git a/src/tirx/transform/make_packed_api.cc b/src/tirx/transform/make_packed_api.cc index c69cbc5c9829..0673692f97cf 100644 --- a/src/tirx/transform/make_packed_api.cc +++ b/src/tirx/transform/make_packed_api.cc @@ -230,9 +230,8 @@ PrimFunc MakePackedAPI(PrimFunc func) { IntImm device_type(PrimType::Int(32), target_device_type); // Create TVMFFIABIBuilder and decode all packed args - TVMFFIABIBuilder binder(name_hint, func_ptr->params, BufferParamMap(func_ptr->params), - v_packed_args, v_num_packed_args, device_type, - device_id.as_or_throw()); + TVMFFIABIBuilder binder(name_hint, func_ptr->params, v_packed_args, v_num_packed_args, + device_type, device_id.as_or_throw()); binder.DecodeAllParams(); auto result = binder.Finalize(); diff --git a/src/tirx/transform/split_host_device.cc b/src/tirx/transform/split_host_device.cc index 334c459c8728..fc6890666a49 100644 --- a/src/tirx/transform/split_host_device.cc +++ b/src/tirx/transform/split_host_device.cc @@ -172,7 +172,7 @@ class HostDeviceSplitter : public StmtMutator { } else { std::unordered_map param_order; for (size_t i = 0; i < cur_func_->params.size(); ++i) { - param_order[tirx::BufferParamMap(cur_func_->params)[cur_func_->params[i]].var()] = i; + param_order[cur_func_->params[i].as_or_throw().var()] = i; } // sort by original order std::sort(params.begin(), params.end(), diff --git a/src/tirx/transform/storage_rewrite.cc b/src/tirx/transform/storage_rewrite.cc index b40732c9fac0..00116b968f1c 100644 --- a/src/tirx/transform/storage_rewrite.cc +++ b/src/tirx/transform/storage_rewrite.cc @@ -95,14 +95,10 @@ struct PrimTypeEqual { // class LinearAccessPatternFinder final : public StmtExprVisitor { public: - LinearAccessPatternFinder(const ffi::Array& params, - const ffi::Map& buffer_map) { - for (const auto& [_, buffer] : buffer_map) { - buffer_aliases_.Set(buffer.var(), buffer.var()); - } + explicit LinearAccessPatternFinder(const ffi::Array& params) { for (const Var& param : params) { - if (param->ty.as()) { - buffer_aliases_.Set(param, param); + if (auto buffer = param.as()) { + buffer_aliases_.Set(buffer.value().var(), buffer.value().var()); } } } @@ -468,11 +464,11 @@ class StoragePlanRewriter : public StmtExprMutator { using StmtEntry = LinearAccessPatternFinder::StmtEntry; using AllocEntry = LinearAccessPatternFinder::AllocEntry; - Stmt Rewrite(Stmt stmt, const ffi::Array& params, const ffi::Map& buffer_map, - bool detect_inplace, bool enable_reuse, bool reuse_require_exact_matched_dtype) { + Stmt Rewrite(Stmt stmt, const ffi::Array& params, bool detect_inplace, bool enable_reuse, + bool reuse_require_exact_matched_dtype) { detect_inplace_ = detect_inplace; // plan the rewrite - LinearAccessPatternFinder finder(params, buffer_map); + LinearAccessPatternFinder finder(params); finder(stmt); this->LivenessAnalysis(finder.linear_seq_); this->PlanMemory(finder.linear_seq_, finder.alloc_info_, enable_reuse, @@ -1274,51 +1270,30 @@ class VectorTypeAccessChecker : public StmtExprVisitor { * * @param params The parameters passed to a PrimFunc * - * @param buffer_map A derived view of the PrimFunc's buffer parameters - * * @param allow_untyped_handles If a buffer or pointer variable is * missing a type annotation, assume that it has the same underlying * type as it is later accessed, with scalar element types. */ - VectorTypeAccessChecker(const ffi::Array& params, - const ffi::Map& buffer_map, - bool allow_untyped_pointers = false, + VectorTypeAccessChecker(const ffi::Array& params, bool allow_untyped_pointers = false, bool detect_scalar_read_patterns = true) : allow_untyped_pointers_(allow_untyped_pointers), detect_scalar_read_patterns_(detect_scalar_read_patterns) { - // If a parameter is in the buffer map, we want to track the - // version in the map. - for (auto it : buffer_map) { - BufferVar& buffer = it.second; - Var buffer_var = buffer.var(); - buffer_aliases_.Set(buffer_var, buffer_var); - PrimType dtype = buffer->dtype; - PrimExpr extent = buffer->shape.size() ? buffer->shape[buffer->shape.size() - 1] : 0; - OnArrayDeclaration(buffer_var, dtype, extent, BufferVarInfo::kPrimFuncBufferParam); - } - - // If a pointer parameter isn't in the buffer map, then we want to - // track the parameter itself. for (Var buffer_var : params) { - if (buffer_map.count(buffer_var)) { - continue; - } - if (auto buffer_type = buffer_var->ty.as()) { - BufferVar buffer(buffer_var); + if (auto buffer = buffer_var.as()) { buffer_aliases_.Set(buffer_var, buffer_var); - PrimExpr extent = - buffer->shape.size() ? buffer->shape[buffer->shape.size() - 1] : PrimExpr(0); - OnArrayDeclaration(buffer_var, buffer->dtype, extent, BufferVarInfo::kPrimFuncBufferParam); + PrimExpr extent = buffer.value()->shape.size() + ? buffer.value()->shape[buffer.value()->shape.size() - 1] + : PrimExpr(0); + OnArrayDeclaration(buffer_var, buffer.value()->dtype, extent, + BufferVarInfo::kPrimFuncBufferParam); continue; } auto pointer_type = GetPointerType(buffer_var->ty); - if (pointer_type.has_value() && !pointer_type.value().IsVoid() && - (buffer_map.count(buffer_var) == 0)) { + if (pointer_type.has_value() && !pointer_type.value().IsVoid()) { PrimType dtype = pointer_type.value(); PrimExpr extent = 0; OnArrayDeclaration(buffer_var, dtype, extent, BufferVarInfo::kPrimFuncPointerParam); - } else if (pointer_type.has_value() && allow_untyped_pointers_ && - (buffer_map.count(buffer_var) == 0)) { + } else if (pointer_type.has_value() && allow_untyped_pointers_) { OnArrayDeclaration(buffer_var, PrimType::Void(), 0, BufferVarInfo::kPrimFuncPointerParam); } } @@ -1930,7 +1905,7 @@ class VectorTypeRewriter : public StmtExprMutator { arith::Analyzer analyzer_; }; -// Rewrite allocates, pointer parameters, and buffer map into vectorized versions +// Rewrite allocates, pointer parameters, and buffer parameters into vectorized versions // if each access into a buffer is the same vector type. PrimFunc PointerValueTypeRewrite(PrimFunc f, bool allow_untyped_pointers = false, bool rewrite_buffer_params = true, @@ -1938,8 +1913,8 @@ PrimFunc PointerValueTypeRewrite(PrimFunc f, bool allow_untyped_pointers = false bool rewrite_alloc_buffer_node = true, bool rewrite_indices = true, bool rewrite_let_node = true, bool rewrite_scalar_read_to_vector_shuffle = true) { - VectorTypeAccessChecker checker(f->params, tirx::BufferParamMap(f->params), - allow_untyped_pointers, rewrite_scalar_read_to_vector_shuffle); + VectorTypeAccessChecker checker(f->params, allow_untyped_pointers, + rewrite_scalar_read_to_vector_shuffle); checker(f->body); VectorTypeRewriter rewriter(checker.info_map_, checker.buffer_aliases_, rewrite_buffer_params, @@ -1973,9 +1948,8 @@ Pass StorageRewrite() { reuse_require_exact_matched_dtype = true; } auto* n = f.CopyOnWrite(); - n->body = - StoragePlanRewriter().Rewrite(std::move(n->body), n->params, BufferParamMap(n->params), - true, enable_reuse, reuse_require_exact_matched_dtype); + n->body = StoragePlanRewriter().Rewrite(std::move(n->body), n->params, true, enable_reuse, + reuse_require_exact_matched_dtype); // Parameters may not be rewritten, but internal allocations may. return PointerValueTypeRewrite(std::move(f), true, false, false, true, true, true, false); }; diff --git a/src/tirx/transform/tvm_ffi_binder.cc b/src/tirx/transform/tvm_ffi_binder.cc index 030075e178cc..a4da6880e0c3 100644 --- a/src/tirx/transform/tvm_ffi_binder.cc +++ b/src/tirx/transform/tvm_ffi_binder.cc @@ -44,12 +44,10 @@ using ffi::reflection::AccessStep; // ============================================================ TVMFFIABIBuilder::TVMFFIABIBuilder(const ffi::String& func_name, const ffi::Array& params, - const ffi::Map& buffer_map, const Var& v_packed_args, const Var& v_num_packed_args, const PrimExpr& device_type, const PrimExpr& device_id) : func_name_(func_name), params_(params), - buffer_map_(buffer_map), v_packed_args_(v_packed_args), device_type_(device_type), device_id_(device_id) { @@ -59,21 +57,20 @@ TVMFFIABIBuilder::TVMFFIABIBuilder(const ffi::String& func_name, const ffi::Arra for (size_t i = 0; i < params.size(); ++i) { if (i > 0) os << ", "; Var param = params[i]; - if (buffer_map.count(param)) { - BufferVar buf = buffer_map[param]; - std::string buf_name = buf.name(); + if (auto buf = param.as()) { + std::string buf_name = buf.value().name(); os << buf_name << ": Tensor(["; - for (size_t j = 0; j < buf->shape.size(); ++j) { + for (size_t j = 0; j < buf.value()->shape.size(); ++j) { if (j > 0) os << ", "; std::ostringstream shape_os; - if (auto var = buf->shape[j].as()) { + if (auto var = buf.value()->shape[j].as()) { shape_os << ((*var)->name.empty() ? "v" : (*var)->name.c_str()); } else { - shape_os << buf->shape[j]; + shape_os << buf.value()->shape[j]; } os << shape_os.str(); } - os << "], " << buf->dtype->dtype << ")"; + os << "], " << buf.value()->dtype->dtype << ")"; param_names_[static_cast(i)] = buf_name; } else { os << param->name << ": "; @@ -486,7 +483,7 @@ Expr TVMFFIABIBuilder::LoadTVMFFIAnyUnionValue(const Var& v_packed_args, int par Expr TVMFFIABIBuilder::DecodeParamOpaqueHandle(int param_index, const PrimExpr& type_index) { // ── Type check: accept handle-like types ─────────────────── - std::string expected_type = buffer_map_.count(params_[param_index]) ? "Tensor" : "pointer"; + std::string expected_type = params_[param_index]->ty.as() ? "Tensor" : "pointer"; EmitTypeIndexCheck(param_index, type_index == ffi::TypeIndex::kTVMFFINone || type_index == ffi::TypeIndex::kTVMFFIOpaquePtr || @@ -557,7 +554,7 @@ void TVMFFIABIBuilder::DecodeParam(int param_index) { ffi::reflection::AccessPath param_path = ffi::reflection::AccessPath::Root()->Extend(AccessStep::ArrayItem(param_index)); - if (buffer_map_.count(param)) { + if (param->ty.as()) { Var handle(param->name + ".handle", PointerType::VoidPointerTy()); Expr handle_value = DecodeParamOpaqueHandle(param_index, type_index.as_or_throw()); BindPointer(handle, handle_value, param_path, true); @@ -605,15 +602,14 @@ void TVMFFIABIBuilder::DecodeAllParams() { // Phase 2: Bind DLTensor buffers (shape, strides, dtype, device checks) for (int i = 0; i < num_args; ++i) { Var param = params_[i]; - if (buffer_map_.count(param)) { - BufferVar buffer = buffer_map_[param]; + if (auto buffer = param.as()) { Var handle = buffer_handles_.at(param.get()); ffi::reflection::AccessPath param_path = ffi::reflection::AccessPath::Root() ->Extend(AccessStep::ArrayItem(i)) - ->Attr(ffi::String(buffer.name())); - Expr data = DecodeParamDLTensor(buffer, device_type_, device_id_, handle, + ->Attr(ffi::String(buffer.value().name())); + Expr data = DecodeParamDLTensor(buffer.value(), device_type_, device_id_, handle, func_name_ + "." + param->name, param_path); - decl_buffers_.push_back(DeclBuffer(buffer, data)); + decl_buffers_.push_back(DeclBuffer(buffer.value(), data)); } } } diff --git a/src/tirx/transform/tvm_ffi_binder.h b/src/tirx/transform/tvm_ffi_binder.h index 7c3c1e7bd3db..7f29a365f5f2 100644 --- a/src/tirx/transform/tvm_ffi_binder.h +++ b/src/tirx/transform/tvm_ffi_binder.h @@ -62,11 +62,11 @@ namespace tirx { * - init_nest: Binds, DeclBuffers for shape/strides arrays, AttrStmts — * all value-loading code that defines variables. * - asserts: AssertStmts — all validation checks. - * - decl_buffers: DeclBuffer for buffer_map entries — buffer declarations. + * - decl_buffers: DeclBuffer for buffer-typed parameters — buffer declarations. * * ## Calling Protocol * - * 1. Construct with function metadata (func_name, params, buffer_map, v_packed_args, + * 1. Construct with function metadata (func_name, params, v_packed_args, * v_num_packed_args). The constructor emits arg count and null-pointer checks. * 2. Call DecodeAllParams(device_type, device_id) * - Decodes, type-checks, and binds all packed arguments @@ -100,7 +100,7 @@ class TVMFFIABIBuilder { std::vector init_nest; /*! \brief Validation checks (all AssertStmts). */ std::vector asserts; - /*! \brief BufferVar declarations for buffer_map entries. */ + /*! \brief BufferVar declarations for buffer-typed parameters. */ std::vector decl_buffers; }; @@ -112,16 +112,14 @@ class TVMFFIABIBuilder { * * \param func_name The function name. * \param params The function parameters. - * \param buffer_map The buffer map from parameters to buffers. * \param v_packed_args The packed args variable (used for struct_get calls). * \param v_num_packed_args The variable holding the actual number of packed args. * \param device_type The expected device type expression. * \param device_id The device id variable (may be defined during buffer binding). */ TVMFFIABIBuilder(const ffi::String& func_name, const ffi::Array& params, - const ffi::Map& buffer_map, const Var& v_packed_args, - const Var& v_num_packed_args, const PrimExpr& device_type, - const PrimExpr& device_id); + const Var& v_packed_args, const Var& v_num_packed_args, + const PrimExpr& device_type, const PrimExpr& device_id); /*! * \brief Decode all packed arguments: type-check, load values, bind buffers. @@ -391,7 +389,7 @@ class TVMFFIABIBuilder { std::vector init_nest_; /*! \brief Validation checks: all AssertStmts. */ std::vector asserts_; - /*! \brief BufferVar declarations for buffer_map entries. */ + /*! \brief BufferVar declarations for buffer-typed parameters. */ std::vector decl_buffers_; /*! \brief Deferred constant-expression assertions for display-var substitution. */ std::vector pending_const_asserts_; @@ -405,8 +403,6 @@ class TVMFFIABIBuilder { std::string func_signature_; /*! \brief The function parameters. */ ffi::Array params_; - /*! \brief The buffer map from parameters to buffers. */ - ffi::Map buffer_map_; /*! \brief Raw packed-ABI handles decoded for buffer-typed parameters. */ std::unordered_map buffer_handles_; /*! \brief The packed args variable. */ diff --git a/src/tirx/transform/unsupported_dtype_legalize.cc b/src/tirx/transform/unsupported_dtype_legalize.cc index 058c74802f06..c9ba1736b010 100644 --- a/src/tirx/transform/unsupported_dtype_legalize.cc +++ b/src/tirx/transform/unsupported_dtype_legalize.cc @@ -557,8 +557,10 @@ class FP8ComputeLegalizer : public ComputeLegalizer { class StorageLegalizer : public StmtExprMutator { public: PrimFunc Legalize(PrimFunc func) { - TVM_FFI_ICHECK_EQ(tirx::BufferParamMap(func->params).size(), 0) - << "This pass must be called after MakePackedAPI"; + for (const Var& param : func->params) { + TVM_FFI_ICHECK(!param->ty.as()) + << "This pass must be called after MakePackedAPI"; + } auto* n = func.CopyOnWrite(); n->params = n->params.Map([this](Var var) { return this->RemapVarDef(var); }); n->body = this->VisitStmt(std::move(n->body)); diff --git a/tests/python/tvmscript/test_tvmscript_printer_tir.py b/tests/python/tvmscript/test_tvmscript_printer_tir.py index 9908a51120da..32047b0e3bb4 100644 --- a/tests/python/tvmscript/test_tvmscript_printer_tir.py +++ b/tests/python/tvmscript/test_tvmscript_printer_tir.py @@ -56,6 +56,21 @@ def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")) ) +def test_prim_func_symbolic_buffer_param_roundtrip(): + n = tirx.Var("n", "int32") + A = tirx.decl_buffer(shape=[n + 1, n], dtype="float32", name="A", layout=None) + func = ( + tirx.PrimFunc(params=[A], body=tirx.Evaluate(n)) + .with_attr("global_symbol", "main") + .with_attr("s_tir", True) + ) + + source = func.script() + assert 'T.Buffer(("n + 1", n)' in source + assert source.index("n = T.int32()") < source.index("T.evaluate(n)") + tvm.ir.assert_structural_equal(tvm.script.from_source(source), func) + + def test_prim_func_buffer_data_use(): A = tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A") B = tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B") From 4d6349dd8c32269fe9777915fccd0dce3b64f086 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Mon, 3 Aug 2026 01:14:05 +0000 Subject: [PATCH 5/7] [FIX][TIRx] Preserve local buffer shape expressions --- src/tirx/script/printer/buffer.cc | 8 +++++--- .../python/tvmscript/test_tvmscript_printer_tir.py | 13 +++++++++++++ 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/src/tirx/script/printer/buffer.cc b/src/tirx/script/printer/buffer.cc index ae2001d28794..0538287ec931 100644 --- a/src/tirx/script/printer/buffer.cc +++ b/src/tirx/script/printer/buffer.cc @@ -29,7 +29,8 @@ namespace printer { ffi::Map BufferAttrs(tirx::BufferVar buffer, const AccessPath& buffer_p, const Frame& frame, const IRDocsifier& d, BufferVarDefinition var_definitions, - ffi::Optional data = std::nullopt) { + ffi::Optional data = std::nullopt, + bool stringify_undefined_shape = false) { using tvm::tirx::Var; using tvm::tirx::VarNode; ffi::Map kwargs; @@ -95,7 +96,8 @@ ffi::Map BufferAttrs(tirx::BufferVar buffer, const AccessP add_out_of_line_var_def(e.as_or_throw(), e_p); } ExprDoc result = d->AsDoc(e, e_p); - results.push_back(contains_new_var ? ExprStringDoc(result, e_p) : result); + results.push_back(stringify_undefined_shape && contains_new_var ? ExprStringDoc(result, e_p) + : result); } kwargs.Set("shape", TupleDoc(results)); } @@ -325,7 +327,7 @@ ExprDoc BufferDecl(const tirx::BufferVar& buffer, const ffi::String& method, ExprDoc BufferAttn(const tirx::BufferVar& buffer, const AccessPath& p, const Frame& frame, const IRDocsifier& d) { ffi::Map attrs = - BufferAttrs(buffer, p, frame, d, BufferVarDefinition::MatchBuffer); + BufferAttrs(buffer, p, frame, d, BufferVarDefinition::MatchBuffer, std::nullopt, true); if (!attrs.count("dtype")) { attrs.Set("dtype", LiteralDoc::DataType(buffer->dtype->dtype, p->Attr("dtype"))); } diff --git a/tests/python/tvmscript/test_tvmscript_printer_tir.py b/tests/python/tvmscript/test_tvmscript_printer_tir.py index 32047b0e3bb4..e9078972175c 100644 --- a/tests/python/tvmscript/test_tvmscript_printer_tir.py +++ b/tests/python/tvmscript/test_tvmscript_printer_tir.py @@ -71,6 +71,19 @@ def test_prim_func_symbolic_buffer_param_roundtrip(): tvm.ir.assert_structural_equal(tvm.script.from_source(source), func) +def test_prim_func_symbolic_alloc_buffer_roundtrip(): + size = tirx.Var("size", "int32") + buf = tirx.decl_buffer(shape=[size], dtype="float32", name="buf", layout=None) + func = tirx.PrimFunc( + params=[], + body=tirx.SeqStmt([tirx.AllocBuffer(buf), tirx.Evaluate(tirx.BufferLoad(buf, [0]))]), + ).with_attr("s_tir", True) + + source = func.script() + assert "T.alloc_buffer((size,))" in source + tvm.ir.assert_structural_equal(tvm.script.from_source(source, check_well_formed=False), func) + + def test_prim_func_buffer_data_use(): A = tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A") B = tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B") From 84724a405fc22a17d91c1c02a6516c316e55bc21 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Mon, 3 Aug 2026 01:57:57 +0000 Subject: [PATCH 6/7] [FIX][TIRx] Parse dependent parameters left to right --- python/tvm/tirx/script/parser/parser.py | 199 ++++++++++-------- src/tirx/script/printer/buffer.cc | 19 +- src/tirx/script/printer/function.cc | 17 +- src/tirx/script/printer/utils.h | 5 +- .../tvmscript/test_tvmscript_parser_tir.py | 99 +++++++++ 5 files changed, 237 insertions(+), 102 deletions(-) diff --git a/python/tvm/tirx/script/parser/parser.py b/python/tvm/tirx/script/parser/parser.py index 603949380444..fe6319d29a56 100644 --- a/python/tvm/tirx/script/parser/parser.py +++ b/python/tvm/tirx/script/parser/parser.py @@ -194,26 +194,8 @@ def bind_assign_value( The bound value. """ if var_name in (prim_var_declarations or {}): - previous = self.var_table.get().get(var_name) - signature_prim_vars = getattr(self, "_signature_prim_vars", []) - if isinstance(previous, tvm.ir.Var) and any( - previous.same_as(var) for var in signature_prim_vars - ): - if not ( - isinstance(previous, tvm.ir.Var) - and isinstance(previous.ty, PrimType) - and isinstance(value, tvm.ir.Var) - and isinstance(value.ty, PrimType) - and not is_buffer_var(previous) - and not is_buffer_var(value) - ): - self.report_error(node, f"{var_name} is already bound to a non-PrimVar value") - if previous.ty != value.ty: - self.report_error( - node, - f"Expected the same dtype for PrimVars but got {value.ty} vs {previous.ty}", - ) - return previous + if _get_signature_prim_var(self, var_name) is not None: + return _reuse_signature_prim_var(self, node, var_name, value) if isinstance(value, T.scalar_wrapper): # pylint: disable=protected-access # special case for scalar, name the buffer, but the var is used as BufferLoad assert isinstance(value.scalar, T.BufferLoad) @@ -316,6 +298,20 @@ def _prim_var_declaration_dtype(node: doc.expr) -> str | None: return str(constructor().ty.dtype) +def _prim_var_annotation_dtype(node: doc.expr) -> str | None: + """Return the dtype of a scalar ``T.dtype`` parameter annotation.""" + if not ( + isinstance(node, doc.Attribute) + and isinstance(node.value, doc.Name) + and node.value.id == "T" + ): + return None + constructor = getattr(T, node.attr, None) + if not isinstance(constructor, T.DtypeConstructor) and constructor is not T.bool: + return None + return str(constructor().ty.dtype) + + def _collect_prim_var_declarations(target: doc.expr, value: doc.expr) -> dict[str, str]: """Pair assignment targets with literal PrimVar declaration dtypes.""" if isinstance(target, doc.Name): @@ -332,8 +328,13 @@ def _collect_prim_var_declarations(target: doc.expr, value: doc.expr) -> dict[st def _function_prim_var_declarations(node: doc.FunctionDef) -> dict[str, str]: - """Collect body declaration dtypes without evaluating or binding them.""" - declarations = {} + """Collect signature-related PrimVar dtypes without binding any values.""" + declarations = { + arg.arg: dtype + for arg in node.args.args + if arg.annotation is not None + if (dtype := _prim_var_annotation_dtype(arg.annotation)) is not None + } for statement in node.body: if isinstance(statement, doc.Assign) and len(statement.targets) == 1: declarations.update( @@ -342,8 +343,58 @@ def _function_prim_var_declarations(node: doc.FunctionDef) -> dict[str, str]: return declarations +@contextlib.contextmanager +def _signature_prim_var_scope(self: Parser): + """Track string-defined signature PrimVars within one function parse.""" + previous = getattr(self, "_signature_prim_vars", None) + self._signature_prim_vars = [] + try: + yield + finally: + if previous is None: + del self._signature_prim_vars + else: + self._signature_prim_vars = previous + + +def _get_signature_prim_var(self: Parser, var_name: str) -> tvm.ir.Var | None: + """Return a string-defined PrimVar from the current function signature.""" + previous = self.var_table.get().get(var_name) + signature_prim_vars = getattr(self, "_signature_prim_vars", []) + if isinstance(previous, tvm.ir.Var) and any( + previous.same_as(var) for var in signature_prim_vars + ): + return previous + return None + + +def _reuse_signature_prim_var(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any: + """Reuse a PrimVar first defined by a string expression in this signature.""" + previous = _get_signature_prim_var(self, var_name) + if previous is None: + return value + if not ( + isinstance(previous.ty, PrimType) + and isinstance(value, tvm.ir.Var) + and isinstance(value.ty, PrimType) + and not is_buffer_var(previous) + and not is_buffer_var(value) + ): + self.report_error(node, f"{var_name} is already bound to a non-PrimVar value") + if previous.ty != value.ty: + self.report_error( + node, + f"Expected the same dtype for PrimVars but got {value.ty} vs {previous.ty}", + ) + return previous + + def _eval_signature_annotation( - self: Parser, node: doc.expr, declaration_dtypes: dict[str, str] + self: Parser, + node: doc.expr, + declaration_dtypes: dict[str, str], + *, + define_missing: bool = True, ) -> Any: """Evaluate string expressions in function-parameter buffer shapes. @@ -360,7 +411,7 @@ def visit_Constant(self, constant): # pylint: disable=invalid-name for child in ast.walk(expression): if not isinstance(child, ast.Name) or not isinstance(child.ctx, ast.Load): continue - if child.id not in self_parser.var_table.get(): + if define_missing and child.id not in self_parser.var_table.get(): var = tvm.tirx.Var(child.id, declaration_dtypes.get(child.id, "int64")) self_parser.var_table.add(child.id, var, allow_shadowing=False) signature_prim_vars = getattr(self_parser, "_signature_prim_vars", []) @@ -763,7 +814,7 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None: s_tir = find_decorator_annotation(node, "s_tir", default=False) persistent = find_decorator_annotation(node, "persistent", default=False) self.function_annotations = None - with self.var_table.with_frame(): + with self.var_table.with_frame(), _signature_prim_var_scope(self): prim_func_ctx = T.prim_func(is_private=privacy, s_tir=s_tir, persistent=persistent) with prim_func_ctx: T.func_name(node.name) @@ -776,49 +827,15 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None: # - kwarg: arg | None # - defaults: list[expr] # - posonlyargs: list[arg] - # Buffer annotations may refer to scalar parameters that occur - # later in the signature. Predeclare any independently - # evaluable scalar/pointer parameters so those forward - # references resolve without changing ABI parameter order. - evaluated_annotations = {} for arg in node.args.args: if arg.annotation is None: - continue + self.report_error(arg, "Type annotation required for function parameters.") try: - ann = self.eval_expr(arg.annotation) - if isinstance(ann, _OptionalAnnotation) or ann is _constexpr_sentinel: - evaluated_annotations[arg.arg] = ann - continue - if callable(ann) and not isinstance(ann, Expr): - ann = ann() + ann = _eval_signature_annotation(self, arg.annotation, declaration_dtypes) except Exception: # pylint: disable=broad-except - continue - evaluated_annotations[arg.arg] = ann - if isinstance(ann, tvm.tirx.Var) and not is_buffer_var(ann): - self.var_table.add(arg.arg, ann) - - if node.returns is not None: - ret_type = _eval_signature_annotation(self, node.returns, declaration_dtypes) - if callable(ret_type) and not isinstance(ret_type, Expr): - ret_type = ret_type() - if isinstance(ret_type, Expr): - ret_type = ret_type.ty - T.func_ret(ret_type) - - for arg in node.args.args: - if arg.annotation is None: - self.report_error(arg, "Type annotation required for function parameters.") - if arg.arg in evaluated_annotations: - ann = evaluated_annotations[arg.arg] - else: - try: - ann = _eval_signature_annotation( - self, arg.annotation, declaration_dtypes - ) - except Exception: # pylint: disable=broad-except - ann = func_annotation.get(arg.arg, None) - if ann is None: - raise + ann = func_annotation.get(arg.arg, None) + if ann is None: + raise if isinstance(ann, _OptionalAnnotation): if not _is_jit_function(node): self.report_error( @@ -846,8 +863,19 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None: # TIRJit.specialize() and lives in an outer var_table # frame; do not register a runtime PrimFunc param. continue + ann = _reuse_signature_prim_var(self, arg.annotation, arg.arg, ann) param = T.arg(arg.arg, ann) self.var_table.add(arg.arg, param) + + if node.returns is not None: + ret_type = _eval_signature_annotation( + self, node.returns, declaration_dtypes, define_missing=False + ) + if callable(ret_type) and not isinstance(ret_type, Expr): + ret_type = ret_type() + if isinstance(ret_type, Expr): + ret_type = ret_type.ty + T.func_ret(ret_type) self.visit_body(node.body) self.function_annotations = supplied_annotation @@ -1088,48 +1116,35 @@ def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar func_annotation = supplied_annotation.get(node.name, {}) ret_type = None - with self.var_table.with_frame(): + with self.var_table.with_frame(), _signature_prim_var_scope(self): declaration_dtypes = _function_prim_var_declarations(node) - evaluated_annotations = {} + arg_annotations = [] for arg in node.args.args: if arg.annotation is None: - continue + self.report_error(arg, "Type annotation required for function parameters.") try: - ann = self.eval_expr(arg.annotation) + ann = _eval_signature_annotation(self, arg.annotation, declaration_dtypes) if callable(ann) and not isinstance(ann, Expr): ann = ann() except Exception: # pylint: disable=broad-except - continue - evaluated_annotations[arg.arg] = ann - if isinstance(ann, tvm.tirx.Var) and not is_buffer_var(ann): - self.var_table.add(arg.arg, ann) + ann = func_annotation.get(arg.arg, None) + if ann is None: + raise + + ann = _reuse_signature_prim_var(self, arg.annotation, arg.arg, ann) + IRBuilder.name(arg.arg, ann) + self.var_table.add(arg.arg, ann) + arg_annotations.append(ann) if node.returns is not None: - ret_type = _eval_signature_annotation(self, node.returns, declaration_dtypes) + ret_type = _eval_signature_annotation( + self, node.returns, declaration_dtypes, define_missing=False + ) if callable(ret_type) and not isinstance(ret_type, Expr): ret_type = ret_type() if isinstance(ret_type, Expr): ret_type = ret_type.ty - arg_annotations = [] - for arg in node.args.args: - if arg.annotation is None: - self.report_error(arg, "Type annotation required for function parameters.") - if arg.arg in evaluated_annotations: - ann = evaluated_annotations[arg.arg] - else: - try: - ann = _eval_signature_annotation(self, arg.annotation, declaration_dtypes) - if callable(ann) and not isinstance(ann, Expr): - ann = ann() - except Exception: # pylint: disable=broad-except - ann = func_annotation.get(arg.arg, None) - if ann is None: - raise - - IRBuilder.name(arg.arg, ann) - arg_annotations.append(ann) - func_signature = tvm.tirx.PrimFunc(arg_annotations, None, ret_type=ret_type) return I.decl_function(node.name, func_signature) diff --git a/src/tirx/script/printer/buffer.cc b/src/tirx/script/printer/buffer.cc index 0538287ec931..2cdcf6c8ef1f 100644 --- a/src/tirx/script/printer/buffer.cc +++ b/src/tirx/script/printer/buffer.cc @@ -26,11 +26,11 @@ namespace tvm { namespace script { namespace printer { -ffi::Map BufferAttrs(tirx::BufferVar buffer, const AccessPath& buffer_p, - const Frame& frame, const IRDocsifier& d, - BufferVarDefinition var_definitions, - ffi::Optional data = std::nullopt, - bool stringify_undefined_shape = false) { +ffi::Map BufferAttrs( + tirx::BufferVar buffer, const AccessPath& buffer_p, const Frame& frame, const IRDocsifier& d, + BufferVarDefinition var_definitions, ffi::Optional data = std::nullopt, + bool stringify_undefined_shape = false, + const std::unordered_set& stringify_shape_vars = {}) { using tvm::tirx::Var; using tvm::tirx::VarNode; ffi::Map kwargs; @@ -89,7 +89,8 @@ ffi::Map BufferAttrs(tirx::BufferVar buffer, const AccessP bool contains_new_var = false; tirx::PostOrderVisit(e, [&](const ffi::ObjectRef& obj) { if (const auto* var = obj.as()) { - contains_new_var = contains_new_var || !d->IsVarDefined(ffi::GetRef(var)); + contains_new_var = contains_new_var || !d->IsVarDefined(ffi::GetRef(var)) || + stringify_shape_vars.count(var); } }); if (is_new_var(e)) { @@ -325,9 +326,11 @@ ExprDoc BufferDecl(const tirx::BufferVar& buffer, const ffi::String& method, } ExprDoc BufferAttn(const tirx::BufferVar& buffer, const AccessPath& p, const Frame& frame, - const IRDocsifier& d) { + const IRDocsifier& d, + const std::unordered_set& stringify_shape_vars) { ffi::Map attrs = - BufferAttrs(buffer, p, frame, d, BufferVarDefinition::MatchBuffer, std::nullopt, true); + BufferAttrs(buffer, p, frame, d, BufferVarDefinition::MatchBuffer, std::nullopt, true, + stringify_shape_vars); if (!attrs.count("dtype")) { attrs.Set("dtype", LiteralDoc::DataType(buffer->dtype->dtype, p->Attr("dtype"))); } diff --git a/src/tirx/script/printer/function.cc b/src/tirx/script/printer/function.cc index 3e30f004f9dc..a3095af61005 100644 --- a/src/tirx/script/printer/function.cc +++ b/src/tirx/script/printer/function.cc @@ -87,9 +87,11 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) ffi::Array args; args.reserve(n_args); std::unordered_map scalar_param_docs; + std::unordered_set pending_scalar_params; for (const tirx::Var& param : func->params) { if (!param->ty.as()) { scalar_param_docs.emplace(param.get(), DefineVar(param, *f, d)); + pending_scalar_params.insert(param.get()); } } for (int i = 0; i < n_args; ++i) { @@ -97,13 +99,26 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) AccessPath var_p = p->Attr("params")->ArrayItem(i); if (var->ty.as()) { tirx::BufferVar buffer(var); + std::unordered_set stringify_shape_vars; + for (const PrimExpr& shape : buffer->shape) { + tirx::PostOrderVisit(shape, [&](const ffi::ObjectRef& obj) { + if (const auto* shape_var = obj.as(); + shape_var && pending_scalar_params.count(shape_var)) { + stringify_shape_vars.insert(shape_var); + } + }); + } IdDoc lhs = DefineBuffer(buffer, *f, d); - ExprDoc annotation = BufferAttn(buffer, var_p->Attr("ty"), *f, d); + ExprDoc annotation = BufferAttn(buffer, var_p->Attr("ty"), *f, d, stringify_shape_vars); args.push_back(AssignDoc(lhs, std::nullopt, annotation)); + for (const tirx::VarNode* shape_var : stringify_shape_vars) { + pending_scalar_params.erase(shape_var); + } continue; } ExprDoc a = d->AsDoc(var->ty, var_p->Attr("ty")); args.push_back(AssignDoc(scalar_param_docs.at(var.get()), std::nullopt, a)); + pending_scalar_params.erase(var.get()); } ffi::Optional ret_type = std::nullopt; if (!func->ret_type.IsMissing()) { diff --git a/src/tirx/script/printer/utils.h b/src/tirx/script/printer/utils.h index 0c15655ad591..d298d67c994d 100644 --- a/src/tirx/script/printer/utils.h +++ b/src/tirx/script/printer/utils.h @@ -35,6 +35,7 @@ #include #include +#include #include #include @@ -313,10 +314,12 @@ ExprDoc BufferDecl(const tirx::BufferVar& buffer, const ffi::String& method, * \param p The object path * \param f The frame * \param d The IRDocsifier + * \param stringify_shape_vars Variables whose first shape use must be stringified * \return The ExprDoc corresponding to the buffer declaration */ ExprDoc BufferAttn(const tirx::BufferVar& buffer, const AccessPath& p, const Frame& frame, - const IRDocsifier& d); + const IRDocsifier& d, + const std::unordered_set& stringify_shape_vars = {}); /*! * \brief Print the creation of a Var diff --git a/tests/python/tvmscript/test_tvmscript_parser_tir.py b/tests/python/tvmscript/test_tvmscript_parser_tir.py index 04d111d25054..113535ed1dda 100644 --- a/tests/python/tvmscript/test_tvmscript_parser_tir.py +++ b/tests/python/tvmscript/test_tvmscript_parser_tir.py @@ -40,6 +40,105 @@ def test_tir_buffer_proxy(): ) +def test_tir_bound_prim_param_reused_in_dependent_annotations(): + func = tvm.script.from_source( + """ +@T.prim_func +def main( + n: T.int32, + direct: T.Buffer((n,), "float32"), + string_direct: T.Buffer(("n",), "float32"), + compound: T.Buffer(("n + 1",), "float32"), +) -> T.Buffer(("n",), "float32"): + return string_direct +""" + ) + + n, direct, string_direct, compound = func.params + assert direct.ty.shape[0].same_as(n) + assert string_direct.ty.shape[0].same_as(n) + assert compound.ty.shape[0].a.same_as(n) + assert func.ret_type.shape[0].same_as(n) + + +def test_tir_bound_prim_param_reused_in_declared_function_signature(): + mod = tvm.script.from_source( + """ +@I.ir_module +class Module: + @T.prim_func + def main(n: T.int32, A: T.Buffer(("n + 1",), "float32")): + T.evaluate(n) +""" + ) + + n, A = mod["main"].params + assert A.ty.shape[0].a.same_as(n) + + +def test_tir_string_defined_symbol_adopted_by_later_prim_param(): + func = tvm.script.from_source( + """ +@T.prim_func +def main(A: T.Buffer(("n",), "float32"), n: T.int32): + T.evaluate(n) +""" + ) + + A, n = func.params + assert A.ty.shape[0].same_as(n) + assert str(n.ty.dtype) == "int32" + + mod = tvm.script.from_source( + """ +@I.ir_module +class Module: + @T.prim_func + def main(A: T.Buffer(("n",), "float32"), n: T.int32): + T.evaluate(n) +""" + ) + + A, n = mod["main"].params + assert A.ty.shape[0].same_as(n) + assert str(n.ty.dtype) == "int32" + + +@pytest.mark.parametrize( + "source", + [ + """ +@T.prim_func +def main(A: T.Buffer((n,), "float32"), n: T.int32): + T.evaluate(n) +""", + """ +@I.ir_module +class Module: + @T.prim_func + def main(A: T.Buffer((n,), "float32"), n: T.int32): + T.evaluate(n) +""", + ], +) +def test_tir_direct_later_prim_param_is_undefined(source): + with pytest.raises(tvm.error.DiagnosticError): + tvm.script.from_source(source) + + +def test_tir_return_annotation_does_not_define_symbolic_var(): + with pytest.raises(tvm.error.DiagnosticError): + tvm.script.from_source( + """ +@T.prim_func +def main() -> T.Buffer(("n",), "float32"): + n = T.int32() + A = T.alloc_buffer((n,), "float32") + return A +""" + ) + + def test_tir_ptr_proxy(): ptr_0 = T.handle("int32", "global") assert ( From 2ed1560c7551ae019e55adb5d649e98973bbd060 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Mon, 3 Aug 2026 07:11:12 +0000 Subject: [PATCH 7/7] [FIX][TIRx] Bind signature match variables locally --- python/tvm/tirx/script/parser/parser.py | 66 +++++++++---------- .../tvmscript/test_tvmscript_parser_tir.py | 37 +++++++++++ 2 files changed, 69 insertions(+), 34 deletions(-) diff --git a/python/tvm/tirx/script/parser/parser.py b/python/tvm/tirx/script/parser/parser.py index fe6319d29a56..832ee14415e3 100644 --- a/python/tvm/tirx/script/parser/parser.py +++ b/python/tvm/tirx/script/parser/parser.py @@ -168,7 +168,7 @@ def bind_assign_value( node: doc.expr, var_name: str, value: Any, - prim_var_declarations: dict[str, str] | None = None, + prim_var_declarations: set[str] | None = None, ) -> Any: """Value binding methods when parsing assign statement. e.g. binding vi, vj, vk with T.axis.remap("SSR", [i, j, k]), when parsing @@ -193,7 +193,7 @@ def bind_assign_value( res : Any The bound value. """ - if var_name in (prim_var_declarations or {}): + if var_name in (prim_var_declarations or set()): if _get_signature_prim_var(self, var_name) is not None: return _reuse_signature_prim_var(self, node, var_name, value) if isinstance(value, T.scalar_wrapper): # pylint: disable=protected-access @@ -281,8 +281,8 @@ def _is_jit_function(node: doc.FunctionDef) -> bool: return False -def _prim_var_declaration_dtype(node: doc.expr) -> str | None: - """Return the dtype of literal ``T.dtype()`` declaration syntax.""" +def _is_prim_var_declaration(node: doc.expr) -> bool: + """Return whether an expression is literal ``T.dtype()`` declaration syntax.""" if not ( isinstance(node, doc.Call) and not node.args @@ -291,11 +291,9 @@ def _prim_var_declaration_dtype(node: doc.expr) -> str | None: and isinstance(node.func.value, doc.Name) and node.func.value.id == "T" ): - return None + return False constructor = getattr(T, node.func.attr, None) - if not isinstance(constructor, T.DtypeConstructor) and constructor is not T.bool: - return None - return str(constructor().ty.dtype) + return isinstance(constructor, T.DtypeConstructor) or constructor is T.bool def _prim_var_annotation_dtype(node: doc.expr) -> str | None: @@ -312,35 +310,28 @@ def _prim_var_annotation_dtype(node: doc.expr) -> str | None: return str(constructor().ty.dtype) -def _collect_prim_var_declarations(target: doc.expr, value: doc.expr) -> dict[str, str]: - """Pair assignment targets with literal PrimVar declaration dtypes.""" +def _prim_var_declaration_names(target: doc.expr, value: doc.expr) -> set[str]: + """Return targets paired with literal ``T.dtype()`` declarations.""" if isinstance(target, doc.Name): - dtype = _prim_var_declaration_dtype(value) - return {target.id: dtype} if dtype is not None else {} + return {target.id} if _is_prim_var_declaration(value) else set() if isinstance(target, doc.Tuple | doc.List) and isinstance(value, doc.Tuple | doc.List): if len(target.elts) != len(value.elts): - return {} - declarations = {} + return set() + declarations = set() for lhs, rhs in zip(target.elts, value.elts): - declarations.update(_collect_prim_var_declarations(lhs, rhs)) + declarations.update(_prim_var_declaration_names(lhs, rhs)) return declarations - return {} + return set() -def _function_prim_var_declarations(node: doc.FunctionDef) -> dict[str, str]: - """Collect signature-related PrimVar dtypes without binding any values.""" - declarations = { +def _signature_prim_var_dtypes(node: doc.FunctionDef) -> dict[str, str]: + """Collect scalar parameter dtypes without inspecting the function body.""" + return { arg.arg: dtype for arg in node.args.args if arg.annotation is not None if (dtype := _prim_var_annotation_dtype(arg.annotation)) is not None } - for statement in node.body: - if isinstance(statement, doc.Assign) and len(statement.targets) == 1: - declarations.update( - _collect_prim_var_declarations(statement.targets[0], statement.value) - ) - return declarations @contextlib.contextmanager @@ -392,7 +383,7 @@ def _reuse_signature_prim_var(self: Parser, node: doc.expr, var_name: str, value def _eval_signature_annotation( self: Parser, node: doc.expr, - declaration_dtypes: dict[str, str], + signature_dtypes: dict[str, str], *, define_missing: bool = True, ) -> Any: @@ -404,6 +395,11 @@ def _eval_signature_annotation( """ class ShapeStringRewriter(ast.NodeTransformer): + def visit_Name(self, name): # pylint: disable=invalid-name + if isinstance(name.ctx, ast.Load) and name.id not in self_parser.var_table.get(): + raise NameError(f"name '{name.id}' is not defined") + return name + def visit_Constant(self, constant): # pylint: disable=invalid-name if not isinstance(constant.value, str): return constant @@ -412,7 +408,9 @@ def visit_Constant(self, constant): # pylint: disable=invalid-name if not isinstance(child, ast.Name) or not isinstance(child.ctx, ast.Load): continue if define_missing and child.id not in self_parser.var_table.get(): - var = tvm.tirx.Var(child.id, declaration_dtypes.get(child.id, "int64")) + # TIR match-scope indices default to int32. A later scalar + # parameter keeps its explicitly declared dtype. + var = tvm.tirx.Var(child.id, signature_dtypes.get(child.id, "int32")) self_parser.var_table.add(child.id, var, allow_shadowing=False) signature_prim_vars = getattr(self_parser, "_signature_prim_vars", []) signature_prim_vars.append(var) @@ -623,7 +621,7 @@ def visit_assign(self: Parser, node: doc.Assign) -> None: except TypeError: pass # rhs not compatible with buffer_store, fall through # otherwise - declarations = _collect_prim_var_declarations(lhs, node.value) + declarations = _prim_var_declaration_names(lhs, node.value) self.eval_assign( target=lhs, source=rhs, @@ -818,7 +816,7 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None: prim_func_ctx = T.prim_func(is_private=privacy, s_tir=s_tir, persistent=persistent) with prim_func_ctx: T.func_name(node.name) - declaration_dtypes = _function_prim_var_declarations(node) + signature_dtypes = _signature_prim_var_dtypes(node) with self.with_dispatch_token("tirx"): # TODO: handle different types of arguments: # - vararg: arg | None @@ -831,7 +829,7 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None: if arg.annotation is None: self.report_error(arg, "Type annotation required for function parameters.") try: - ann = _eval_signature_annotation(self, arg.annotation, declaration_dtypes) + ann = _eval_signature_annotation(self, arg.annotation, signature_dtypes) except Exception: # pylint: disable=broad-except ann = func_annotation.get(arg.arg, None) if ann is None: @@ -869,7 +867,7 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None: if node.returns is not None: ret_type = _eval_signature_annotation( - self, node.returns, declaration_dtypes, define_missing=False + self, node.returns, signature_dtypes, define_missing=False ) if callable(ret_type) and not isinstance(ret_type, Expr): ret_type = ret_type() @@ -1117,14 +1115,14 @@ def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar ret_type = None with self.var_table.with_frame(), _signature_prim_var_scope(self): - declaration_dtypes = _function_prim_var_declarations(node) + signature_dtypes = _signature_prim_var_dtypes(node) arg_annotations = [] for arg in node.args.args: if arg.annotation is None: self.report_error(arg, "Type annotation required for function parameters.") try: - ann = _eval_signature_annotation(self, arg.annotation, declaration_dtypes) + ann = _eval_signature_annotation(self, arg.annotation, signature_dtypes) if callable(ann) and not isinstance(ann, Expr): ann = ann() except Exception: # pylint: disable=broad-except @@ -1139,7 +1137,7 @@ def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar if node.returns is not None: ret_type = _eval_signature_annotation( - self, node.returns, declaration_dtypes, define_missing=False + self, node.returns, signature_dtypes, define_missing=False ) if callable(ret_type) and not isinstance(ret_type, Expr): ret_type = ret_type() diff --git a/tests/python/tvmscript/test_tvmscript_parser_tir.py b/tests/python/tvmscript/test_tvmscript_parser_tir.py index 113535ed1dda..a4baa9b4568a 100644 --- a/tests/python/tvmscript/test_tvmscript_parser_tir.py +++ b/tests/python/tvmscript/test_tvmscript_parser_tir.py @@ -104,6 +104,43 @@ def main(A: T.Buffer(("n",), "float32"), n: T.int32): assert str(n.ty.dtype) == "int32" +def test_tir_string_defined_symbol_preserves_later_prim_param_dtype(): + func = tvm.script.from_source( + """ +@T.prim_func +def main(A: T.Buffer(("n",), "float32"), n: T.int64): + T.evaluate(n) +""" + ) + + A, n = func.params + assert A.ty.shape[0].same_as(n) + assert str(n.ty.dtype) == "int64" + + +def test_tir_string_defined_symbol_does_not_take_dtype_from_body(): + with pytest.raises(tvm.error.DiagnosticError): + tvm.script.from_source( + """ +@T.prim_func +def main(A: T.Buffer(("n",), "float32")): + n = T.int64() + T.evaluate(n) +""" + ) + + +def test_tir_direct_use_before_string_definition_is_undefined(): + with pytest.raises(tvm.error.DiagnosticError): + tvm.script.from_source( + """ +@T.prim_func +def main(A: T.Buffer((n, "n"), "float32")): + T.evaluate(0) +""" + ) + + @pytest.mark.parametrize( "source", [