You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/08/05 01:17:10 UTC
(tvm) branch main updated: tirx: represent buffer parameters with BufferType (#20086)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 68ba2b31c2 tirx: represent buffer parameters with BufferType (#20086)
68ba2b31c2 is described below
commit 68ba2b31c2a6d202fcd44e5780ef10ce08721dd5
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Wed Aug 5 09:17:05 2026 +0800
tirx: represent buffer parameters with BufferType (#20086)
This phases out PrimFunc.buffer_map entirely. Buffer parameters now
carry their BufferType annotation directly in PrimFunc.params, with no
compatibility constructor or derived PrimFunc buffer map.
The migration updates construction and visitor paths across TIRx, TE,
S-TIR, Relax, printers/parsers, packed-ABI lowering, specialization,
storage rewriting, and related tests and documentation. BufferType
signature shapes use left-to-right match-scope binding: the first quoted
shape expression defines missing symbolic variables, and later
dimensions, parameters, and matching scalar declarations reuse those
exact variables.
---
include/tvm/relax/distributed/axis_group_graph.h | 14 +-
include/tvm/script/ir_builder/base.h | 2 +-
include/tvm/tirx/function.h | 56 -----
.../backend/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 | 25 +-
python/tvm/tirx/script/parser/parser.py | 252 +++++++++++++++++++--
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 +-
src/backend/trn/transform/lower_trainium_layout.cc | 33 +--
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/backend/vm/vm_shape_lower.cc | 7 +-
src/relax/distributed/axis_group_graph.cc | 2 +-
.../transform/lower_global_view_to_local_view.cc | 26 ++-
src/relax/op/tensor/inspect.cc | 4 +-
src/relax/transform/compute_prim_value.cc | 2 +-
src/relax/transform/dataflow_inplace.cc | 16 +-
src/relax/transform/fuse_tir.cc | 44 ++--
src/relax/transform/meta_schedule.cc | 1 -
src/relax/transform/rewrite_dataflow_reshape.cc | 2 +-
src/relax/transform/split_call_tir_by_pattern.cc | 41 ++--
.../transform/split_layout_rewrite_preproc.cc | 18 +-
.../analysis/sblock_buffer_access_lca_detector.cc | 7 +-
src/s_tir/backend/adreno/inject_texture_alloc.cc | 2 +-
src/s_tir/backend/adreno/texture_flatten.cc | 16 +-
src/s_tir/meta_schedule/arg_info.cc | 7 +-
.../feature_extractor/per_store_feature.cc | 6 +-
src/s_tir/meta_schedule/postproc/rewrite_layout.cc | 2 +-
src/s_tir/schedule/analysis/analysis.cc | 9 +-
src/s_tir/schedule/primitive/blockize_tensorize.cc | 6 +-
.../schedule/primitive/layout_transformation.cc | 17 +-
src/s_tir/transform/inject_permuted_layout.cc | 7 +-
src/s_tir/transform/inject_software_pipeline.cc | 7 +-
src/s_tir/transform/lower_thread_allreduce.cc | 8 +-
.../plan_update_buffer_allocation_location.cc | 15 +-
.../remove_weight_layout_rewrite_block.cc | 16 +-
src/s_tir/transform/renew_defs.cc | 41 ++--
src/target/llvm/codegen_llvm.cc | 6 +-
src/te/operation/create_primfunc.cc | 22 +-
src/tirx/analysis/verify_memory.cc | 4 +-
src/tirx/analysis/verify_ssa.cc | 6 +-
src/tirx/analysis/verify_well_formed.cc | 2 +-
src/tirx/ir/data_type_rewriter.cc | 11 +-
src/tirx/ir/function.cc | 21 +-
src/tirx/ir/script/script_complete.cc | 7 +-
src/tirx/ir/specialize.cc | 80 ++++---
src/tirx/ir/tir_visitor_with_path.cc | 21 +-
src/tirx/ir/tir_visitor_with_path.h | 29 ++-
src/tirx/script/builder/frame.cc | 64 ++++--
src/tirx/script/builder/ir.cc | 6 +-
src/tirx/script/printer/buffer.cc | 42 +++-
src/tirx/script/printer/function.cc | 144 ++++--------
src/tirx/script/printer/utils.h | 58 +----
src/tirx/transform/flatten_buffer.cc | 29 ++-
src/tirx/transform/force_narrow_index_to_i32.cc | 11 +-
src/tirx/transform/inline_private_functions.cc | 13 +-
src/tirx/transform/ir_utils.cc | 67 +++---
src/tirx/transform/lower_tirx_cleanup.cc | 41 ++--
src/tirx/transform/make_packed_api.cc | 5 +-
src/tirx/transform/split_host_device.cc | 2 +-
src/tirx/transform/stmt_simplify.cc | 2 +-
src/tirx/transform/storage_rewrite.cc | 110 ++++-----
src/tirx/transform/tvm_ffi_binder.cc | 35 +--
src/tirx/transform/tvm_ffi_binder.h | 18 +-
src/tirx/transform/unsupported_dtype_legalize.cc | 5 +-
tests/python/arith/test_arith_domain_touched.py | 4 +-
tests/python/relax/test_analysis_well_formed.py | 8 +-
tests/python/relax/test_blockbuilder_core.py | 7 +-
...nsform_specialize_primfunc_based_on_callsite.py | 10 +-
.../test_s_tir_analysis_identify_memcpy.py | 22 +-
.../s_tir/analysis/test_sblock_access_region.py | 8 +-
.../analysis/test_sblock_buffer_access_lca.py | 11 +-
.../s_tir/base/test_tir_te_extern_primfunc.py | 6 +-
.../test_tir_schedule_decompose_padding.py | 4 +-
.../schedule/test_tir_schedule_rolling_buffer.py | 4 +-
.../schedule/test_tir_schedule_transform_layout.py | 8 +-
tests/python/s_tir/test_s_tir_renew_defs.py | 17 +-
.../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 | 49 ++--
...test_tir_transform_force_narrow_index_to_i32.py | 6 +-
...est_tir_transform_pointer_value_type_rewrite.py | 2 +-
tests/python/tirx/test_jit.py | 34 +--
tests/python/tirx/transform/test_stmt_functor.py | 2 +-
tests/python/tvmscript/test_tvmscript_complete.py | 6 +-
.../tvmscript/test_tvmscript_ir_builder_tir.py | 23 +-
.../python/tvmscript/test_tvmscript_parser_tir.py | 136 +++++++++++
.../test_tvmscript_printer_structural_equal.py | 10 +-
.../python/tvmscript/test_tvmscript_printer_tir.py | 101 +++++----
98 files changed, 1199 insertions(+), 940 deletions(-)
diff --git a/include/tvm/relax/distributed/axis_group_graph.h b/include/tvm/relax/distributed/axis_group_graph.h
index e7cd828b1f..6d446e3f7f 100644
--- a/include/tvm/relax/distributed/axis_group_graph.h
+++ b/include/tvm/relax/distributed/axis_group_graph.h
@@ -71,14 +71,18 @@ class BufferAxisGraphExtractor : public StmtExprVisitor {
BufferAxisGraphExtractor extractor;
extractor(prim_func->body);
ffi::Map<BufferVar, Var> inverse_buffer_map;
- for (const auto& pr : prim_func->buffer_map) {
- inverse_buffer_map.Set(pr.second, pr.first);
+ for (const Var& param : prim_func->params) {
+ if (param->ty.as<BufferTypeNode>()) {
+ inverse_buffer_map.Set(BufferVar(param), param);
+ }
}
std::vector<std::vector<TIRVarAxis>> tir_var_axis_group_list;
std::unordered_set<BufferAxis, BufferAxisHash> 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<BufferTypeNode>()) {
+ continue;
+ }
+ BufferVar buffer(param);
for (int i = 0; i < static_cast<int>(buffer->shape.size()); i++) {
if (extractor.buffer_axis_graph_.count({buffer, i})) {
std::vector<BufferAxis> buffer_axis_group;
diff --git a/include/tvm/script/ir_builder/base.h b/include/tvm/script/ir_builder/base.h
index 010a091ab2..7e0d34a42f 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 414e6c40b9..d335ee9515 100644
--- a/include/tvm/tirx/function.h
+++ b/include/tvm/tirx/function.h
@@ -52,52 +52,6 @@ class PrimFuncNode : public BaseFuncNode {
ffi::Array<tirx::Var> 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<tirx::Var, BufferVar> buffer_map;
/*! \brief The body of the function */
tirx::Stmt body;
@@ -106,8 +60,6 @@ class PrimFuncNode : public BaseFuncNode {
refl::ObjectDef<PrimFuncNode>()
.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<PrimFuncNode>()
.def("__s_equal__", &PrimFuncNode::SEqual)
@@ -122,7 +74,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 +82,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,17 +113,11 @@ 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<tirx::Var> params, Stmt body, Type ret_type = VoidType(),
- ffi::Map<tirx::Var, BufferVar> buffer_map = ffi::Map<tirx::Var, BufferVar>(),
DictAttrs attrs = DictAttrs(), Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrimFunc, BaseFunc, 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 f5d678567c..0abba7675b 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 _emit_plan(
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 987b7eedb3..7f0a6c2c39 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 a1eb56361e..486252aedb 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 73c8d7f263..f7d88f8131 100644
--- a/python/tvm/te/operation.py
+++ b/python/tvm/te/operation.py
@@ -389,9 +389,9 @@ def extern_primfunc(input_tensors: list[_tensor.Tensor], primfunc: tvm.tirx.Prim
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 bbacc0472f..38c701075f 100644
--- a/python/tvm/tirx/function.py
+++ b/python/tvm/tirx/function.py
@@ -31,7 +31,7 @@ from tvm.runtime import Object, Scriptable
from ..runtime._tensor import Tensor
from . import _ffi_api
-from .buffer import Buffer, is_buffer_var
+from .buffer import Buffer
from .expr import Expr, Var
@@ -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
@@ -60,7 +57,7 @@ class PrimFunc(BaseFunc, Scriptable):
The location of this itervar in the source code.
"""
- def __init__(self, params, body, ret_type=None, buffer_map=None, attrs=None, span=None):
+ def __init__(self, params, body, ret_type=None, attrs=None, span=None):
# Legacy compatibility: expand body-carrying leaf stmt wrappers
# (e.g. DeclBuffer/AllocBuffer forms) into SeqStmt form.
from .stmt import _normalize_legacy_stmt
@@ -69,17 +66,11 @@ class PrimFunc(BaseFunc, Scriptable):
if ret_type is None:
ret_type = tvm.ir.Type.missing()
param_list = []
- buffer_map = {} if buffer_map is None else buffer_map
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)
- else:
+ if not isinstance(x, Var):
raise TypeError("params can only contain Var or Buffer")
+ param_list.append(x)
if attrs is None:
attrs = tvm.ir.make_node("ir.DictAttrs")
@@ -89,7 +80,6 @@ class PrimFunc(BaseFunc, Scriptable):
param_list,
body,
ret_type,
- buffer_map,
attrs,
span,
) # type: ignore
@@ -113,10 +103,9 @@ class PrimFunc(BaseFunc, Scriptable):
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 3e62e8b055..324f59cdb5 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: 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
vi, vj, vk = T.axis.remap("SSR", [i, j, k]).
@@ -187,6 +193,12 @@ 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 set()):
+ # A quoted Buffer shape may have already created this PrimVar. In that
+ # case ``n = T.int32()`` is match-like syntax: bind the Python name to
+ # the signature Var instead of emitting a second definition.
+ if _get_signature_match_var(self, var_name) is not None:
+ return _reuse_signature_match_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)
@@ -272,6 +284,191 @@ def _is_jit_function(node: doc.FunctionDef) -> bool:
return False
+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
+ and not node.keywords
+ and isinstance(node.func, doc.Attribute)
+ and isinstance(node.func.value, doc.Name)
+ and node.func.value.id == "T"
+ ):
+ return False
+ constructor = getattr(T, node.func.attr, None)
+ return isinstance(constructor, T.DtypeConstructor) or constructor is T.bool
+
+
+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 _prim_var_declaration_names(target: doc.expr, value: doc.expr) -> set[str]:
+ """Return targets paired with literal ``T.dtype()`` declarations.
+
+ This classifies only the assignment currently being visited. In
+ particular, it is not a predeclaration pass over the function body.
+ """
+ if isinstance(target, doc.Name):
+ 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 set()
+ declarations = set()
+ for lhs, rhs in zip(target.elts, value.elts):
+ declarations.update(_prim_var_declaration_names(lhs, rhs))
+ return declarations
+ return set()
+
+
+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
+ }
+
+
+@contextlib.contextmanager
+def _signature_match_var_scope(self: Parser):
+ """Track PrimVars introduced by Buffer-shape match expressions."""
+ previous = getattr(self, "_signature_match_vars", None)
+ self._signature_match_vars = {}
+ try:
+ yield
+ finally:
+ if previous is None:
+ del self._signature_match_vars
+ else:
+ self._signature_match_vars = previous
+
+
+def _get_signature_match_var(self: Parser, var_name: str) -> tvm.ir.Var | None:
+ """Return the active PrimVar created by a signature match expression."""
+ current = self.var_table.get().get(var_name)
+ match_var = getattr(self, "_signature_match_vars", {}).get(var_name)
+ if (
+ isinstance(current, tvm.ir.Var)
+ and isinstance(match_var, tvm.ir.Var)
+ and current.same_as(match_var)
+ ):
+ return current
+ return None
+
+
+def _reuse_signature_match_var(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
+ """Match a later scalar declaration to a signature-defined PrimVar."""
+ previous = _get_signature_match_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,
+ signature_dtypes: dict[str, str],
+ *,
+ define_missing: bool = True,
+) -> Any:
+ """Evaluate a function annotation using Buffer-shape match semantics.
+
+ Function parameters are parsed from left to right. An ordinary name in an
+ annotation is therefore a lookup and must already be bound. A quoted
+ expression in ``T.Buffer``'s shape is the one exception: it is a match
+ scope, so its first occurrence creates every missing PrimVar before the
+ expression is evaluated. Later dimensions and parameters reuse those
+ exact Var objects.
+
+ ``signature_dtypes`` comes only from scalar parameter annotations, never
+ from the function body. It lets a match expression that precedes a scalar
+ parameter create the Var with that parameter's declared dtype. Otherwise
+ TIR match variables default to int32. Return annotations pass
+ ``define_missing=False`` and consequently cannot introduce variables.
+ """
+
+ class ShapeStringRewriter(ast.NodeTransformer):
+ def visit_Name(self, name): # pylint: disable=invalid-name
+ # Direct AST names follow normal definition-before-use ordering.
+ 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
+ # Replace the quoted expression with its AST. Register its names
+ # first so evaluating the rewritten annotation resolves all uses to
+ # the canonical match-scope Vars.
+ 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 define_missing and child.id not in self_parser.var_table.get():
+ # 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)
+ self_parser._signature_match_vars[child.id] = var
+ 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)
+ # Only Buffer.shape has match-scope semantics. Dtype and other
+ # Buffer options remain ordinary expressions.
+ 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
+
+ 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")
def visit_for(self: Parser, node: doc.For) -> None:
"""The for visiting method for tirx.
@@ -446,7 +643,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 = _prim_var_declaration_names(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")
@@ -632,17 +834,12 @@ 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_match_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)
- 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)
+ # This is a signature-only dtype index, not lookahead into the body.
+ signature_dtypes = _signature_prim_var_dtypes(node)
with self.with_dispatch_token("tirx"):
# TODO: handle different types of arguments:
# - vararg: arg | None
@@ -655,7 +852,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 = self.eval_expr(arg.annotation)
+ 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:
@@ -687,8 +884,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_match_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, signature_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
@@ -929,20 +1137,15 @@ 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():
- 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
+ with self.var_table.with_frame(), _signature_match_var_scope(self):
+ 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 = self.eval_expr(arg.annotation)
+ 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
@@ -950,8 +1153,19 @@ def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar
if ann is None:
raise
+ ann = _reuse_signature_match_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, signature_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
+
func_signature = tvm.tirx.PrimFunc(arg_annotations, None, ret_type=ret_type)
return I.decl_function(node.name, func_signature)
diff --git a/python/tvm/tirx/transform/transform.py b/python/tvm/tirx/transform/transform.py
index 9073244549..66dc7f3b5d 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 94a90be109..e26fb448da 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<BufferVar, ffi::Array<ffi::ObjectRef>> DomainTouchedAccessMap(const PrimFunc& func) {
auto buffer_access_map = BufferTouchedDomain(func->body).GetAccessedBufferRegions();
ffi::Map<BufferVar, ffi::Array<ffi::ObjectRef>> ret;
- auto& buffer_map = func->buffer_map;
for (auto& var : func->params) {
- auto& buffer = buffer_map[var];
+ if (!var->ty.as<BufferTypeNode>()) {
+ continue;
+ }
+ BufferVar buffer(var);
auto& access = buffer_access_map[buffer.get()];
ffi::Array<ffi::Array<IntSet>> loads, stores, combined;
for (std::vector<IntSet>& touch : std::get<LoadAccess>(access).set) {
diff --git a/src/arith/ir_mutator_with_analyzer.cc b/src/arith/ir_mutator_with_analyzer.cc
index 65f20e0429..260c053f75 100644
--- a/src/arith/ir_mutator_with_analyzer.cc
+++ b/src/arith/ir_mutator_with_analyzer.cc
@@ -153,10 +153,14 @@ void EnterConstraintFacts(WithGroup<ConstraintContext>* 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<tirx::BufferTypeNode>()) {
+ 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 47b52da78c..1d2f8ef1d6 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 fbad6be2a3..d973e10e8c 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 = v.as<tirx::BufferVar>()) {
var_idmap_[buffer.value().get()] = vid;
}
if (i >= static_cast<size_t>(num_inputs.value())) {
diff --git a/src/backend/trn/transform/lower_trainium_layout.cc b/src/backend/trn/transform/lower_trainium_layout.cc
index 7c6f069d57..2c78845eed 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<Stmt, ffi::Map<Var, BufferVar>> Lower(
- const Stmt& stmt, const ffi::Map<tirx::Var, BufferVar> buffer_map) {
+ static std::pair<Stmt, ffi::Array<Var>> Lower(const Stmt& stmt, const ffi::Array<Var>& params) {
arith::Analyzer ana;
TrainiumLayoutApplier storage_lower(ana);
- std::unordered_map<Var, BufferVar> new_buffer_map;
+ ffi::Array<Var> new_params;
+ new_params.reserve(params.size());
std::vector<std::pair<BufferVar, BufferVar>> 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<BufferVar>();
+ 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<Var, BufferVar>(new_buffer_map));
+ return std::make_pair(new_stmt, new_params);
}
protected:
@@ -346,7 +351,9 @@ 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, params] = TrainiumLayoutApplier::Lower(n->body, n->params);
+ 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 d3aef0ee18..207e311702 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<BufferVar> param_buf = func->buffer_map.Get(param);
+ ffi::Optional<BufferVar> param_buf = param.as<BufferVar>();
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 9290f00564..c7b27ac2d2 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<BufferVar> param_buf = func->buffer_map.Get(param);
+ ffi::Optional<BufferVar> param_buf = param.as<BufferVar>();
if (param_buf.has_value()) {
param_buffers_.insert(param_buf.value());
}
@@ -535,7 +535,7 @@ bool HasReshapePattern(const PrimFunc& func) {
ffi::Array<BufferVar> buffer_args;
for (const auto& param : func->params) {
- if (auto buffer = func->buffer_map.Get(param)) {
+ if (auto buffer = param.as<BufferVar>()) {
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 237fe8ca60..644f46ed00 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 48191f32a3..75fd5da0be 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<ffi::Optional<Expr>(const Expr&)> {
ffi::Array<tirx::Var> 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/backend/vm/vm_shape_lower.cc b/src/relax/backend/vm/vm_shape_lower.cc
index cd4b1e952c..e7d5287493 100644
--- a/src/relax/backend/vm/vm_shape_lower.cc
+++ b/src/relax/backend/vm/vm_shape_lower.cc
@@ -709,9 +709,6 @@ class VMShapeLowerMutator
// construct a PrimFunc that compute the shape.
ffi::Array<PrimExpr> buffer_shape{heap_size_};
tirx::BufferVar buffer = tirx::decl_buffer(buffer_shape, PrimType(ShapeDType()), "H", "global");
- tirx::Var heap("heap", PointerType::VoidPointerTy());
- ffi::Map<tirx::Var, tirx::BufferVar> buffer_map;
- buffer_map.Set(heap, buffer);
ffi::Map<tirx::Var, PrimExpr> var_map;
for (const auto& [expr, slot] : slot_map_) {
@@ -731,13 +728,13 @@ class VMShapeLowerMutator
}
tirx::Stmt body = tirx::SeqStmt::Flatten(seq);
- ffi::Array<tirx::Var> params{heap};
+ ffi::Array<tirx::Var> params{buffer.var()};
Type ret_type = VoidType();
// TODO(relax-team): Consider attach the target attribute to
// the shape_func to indicate that this is a host function
// This could require us to attach target to the relax function here.
- tirx::PrimFunc shape_func(params, body, ret_type, buffer_map);
+ tirx::PrimFunc shape_func(params, body, ret_type);
shape_func = WithAttr(std::move(shape_func), tvm::attr::kSTir, true);
if (!shape_func->attrs.GetAttr<tvm::Target>(tvm::attr::kTarget).has_value()) {
// kTarget and kIsHostFunc are mutually exclusive
diff --git a/src/relax/distributed/axis_group_graph.cc b/src/relax/distributed/axis_group_graph.cc
index 042d459386..40f285498c 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<Expr> input_list = call->args[1].as_or_throw<Tuple>()->fields;
input_list.push_back(output_var);
for (int i = 0; i < static_cast<int>(input_list.size()); i++) {
- if (func->buffer_map.count(func->params[i])) {
+ if (func->params[i]->ty.as<tirx::BufferTypeNode>()) {
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 7f9628398b..a2da47cf9d 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<ShardingSpec>& sharding_specs, PrimFunc prim_func) {
prim_func = s_tir::RenewDefs(prim_func);
DistributedBufferCompactor compactor(sharding_specs, prim_func);
- ffi::Map<Var, BufferVar> new_func_buffer_map;
+ ffi::Array<Var> new_params;
ffi::Map<BufferVar, BufferVar> 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<BufferTypeNode>()) {
+ 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<BufferAxis, BufferAxisHash> visited;
for (int i = 0, j = 0; i < static_cast<int>(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<BufferTypeNode>()) {
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<int>(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<tirx::BufferTypeNode>()) {
const auto* ty = GetTypeAs<DTensorTypeNode>(arg);
TVM_FFI_CHECK(ty, TypeError)
<< "Expected buffer parameter " << param << " to receive a distributed tensor, but "
diff --git a/src/relax/op/tensor/inspect.cc b/src/relax/op/tensor/inspect.cc
index d4aec119bb..6a6de40cd1 100644
--- a/src/relax/op/tensor/inspect.cc
+++ b/src/relax/op/tensor/inspect.cc
@@ -101,7 +101,7 @@ tirx::PrimFunc GetDLTensorField(tirx::builtin::TVMStructFieldKind field, PrimTyp
DictAttrs attrs({{"tirx.is_scheduled", true}, {"tirx.is_host_func", true}});
- tirx::PrimFunc func(ffi::Array<tirx::Var>{dlpack_handle}, body, field_ty, {}, attrs);
+ tirx::PrimFunc func(ffi::Array<tirx::Var>{dlpack_handle}, body, field_ty, attrs);
FuncType ty({TensorType(std::nullopt, kUnknownNDim)}, field_ty);
func->ty = ty;
@@ -286,7 +286,7 @@ Expr LegalizeTensorShape(const BlockBuilder& bb, const Call& call) {
DictAttrs attrs({{"tirx.is_scheduled", true}, {"tirx.is_host_func", true}});
- tirx::PrimFunc func({dlpack_handle, axis}, body, field_ty, {}, attrs);
+ tirx::PrimFunc func({dlpack_handle, axis}, body, field_ty, attrs);
FuncType ty({TensorType(std::nullopt, kUnknownNDim), axis->ty.as_or_throw<PrimType>()},
field_ty);
diff --git a/src/relax/transform/compute_prim_value.cc b/src/relax/transform/compute_prim_value.cc
index 1e582f13a8..d2e748bb43 100644
--- a/src/relax/transform/compute_prim_value.cc
+++ b/src/relax/transform/compute_prim_value.cc
@@ -104,7 +104,7 @@ class PrimExprComputeInjector : public ExprMutator {
auto param_vars = tirx::UndefinedVars(node);
tirx::Stmt body = tirx::Return(node);
- tirx::PrimFunc func(param_vars, body, ret_ty, {},
+ tirx::PrimFunc func(param_vars, body, ret_ty,
DictAttrs({{tirx::attr::kIsHostFunc, true}, {tvm::attr::kSTir, true}}));
func = s_tir::RenewDefs(func);
diff --git a/src/relax/transform/dataflow_inplace.cc b/src/relax/transform/dataflow_inplace.cc
index d6a1297226..b7cbcaa937 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<tirx::BufferVar, tirx::BufferVar> buffer_subst_map;
ffi::Map<tirx::Var, tirx::Var> 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 = old_primfunc->buffer_map.at(output_var);
- auto inplace_buffer = old_primfunc->buffer_map.at(inplace_var);
+ auto output_buffer = output_var.as_or_throw<tirx::BufferVar>();
+ auto inplace_buffer = inplace_var.as_or_throw<tirx::BufferVar>();
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 = old_primfunc->buffer_map;
- 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<tirx::Var> 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 a7145144b8..6fb6e9829a 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<GlobalVar>();
tirx::PrimFunc prim_func_ = mod_->Lookup(gv).as_or_throw<tirx::PrimFunc>();
- const auto& buffer_map = prim_func_->buffer_map;
const auto& tir_args = prim_func_->params;
const auto& relax_args = call->args[1].as_or_throw<Tuple>()->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<tirx::BufferVar>()) {
if (i < num_inputs) {
const auto& relax_var = relax_args[i];
ValidateBufferCompatibility(tir_buffer.value(), relax_var);
@@ -626,8 +625,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 +809,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 = param.as<tirx::BufferVar>()) {
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<PrimExpr>();
TVM_FFI_CHECK(prim_arg.has_value(), TypeError)
@@ -826,20 +825,21 @@ class FusedTIRConstructor : public ExprVisitor {
MapArgsToBuffer(arg_list, buffer_list);
}
- static ffi::Array<tirx::Var> GetPrimFuncOutputParams(const tirx::PrimFunc& func,
- const ffi::Array<int64_t>& output_indices) {
+ static ffi::Array<tirx::BufferVar> GetPrimFuncOutputParams(
+ const tirx::PrimFunc& func, const ffi::Array<int64_t>& output_indices) {
size_t n = func->params.size();
size_t output_size = output_indices.size();
TVM_FFI_ICHECK_GE(n, output_size);
- ffi::Array<tirx::Var> ret;
+ ffi::Array<tirx::BufferVar> ret;
for (int64_t idx : output_indices) {
int i = static_cast<int>(idx);
const tirx::Var& param = func->params[static_cast<size_t>(i)];
- TVM_FFI_ICHECK(param->ty.as<PointerTypeNode>())
- << "The output params of a PrimFunc must be buffer handles, but parameter " << i
- << " has type " << param->ty;
- ret.push_back(param);
+ auto buffer = param.as<tirx::BufferVar>();
+ 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 +871,9 @@ class FusedTIRConstructor : public ExprVisitor {
}
}
- ffi::Array<tirx::Var> output_params = GetPrimFuncOutputParams(func, output_idxs);
+ ffi::Array<tirx::BufferVar> 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 +987,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::SBlock>());
- tirx::PrimFunc func(func_info_.params, body, VoidType(), func_info_.buffer_map,
- DictAttrs(attr_map));
+ ffi::Array<tirx::Var> 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 +1036,9 @@ class FusedTIRConstructor : public ExprVisitor {
* function
*/
ffi::Map<tirx::BufferVar, tirx::BufferVar> buffer_subst_map;
- /*! \brief The `buffer_map` in the fused function*/
+ /*! \brief Buffer annotations keyed by their placeholder parameters. */
ffi::Map<tirx::Var, tirx::BufferVar> buffer_map;
- /*! \brief The output buffers in the function buffer_map*/
+ /*! \brief The output buffers among the function parameters. */
std::unordered_set<const tirx::VarNode*> 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 161f459038..bd035d970f 100644
--- a/src/relax/transform/meta_schedule.cc
+++ b/src/relax/transform/meta_schedule.cc
@@ -134,7 +134,6 @@ Pass MetaScheduleApplyDatabase(ffi::Optional<ffi::String> 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 97a3658f53..464be7d774 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<size_t> GetUsedTensorArgIndices(const tirx::PrimFunc& fn, size_t num_args) {
std::vector<size_t> indices;
for (size_t i = 0; i < num_args; ++i) {
- if (auto buffer = fn->buffer_map.Get(fn->params[i])) {
+ if (auto buffer = fn->params[i].as<tirx::BufferVar>()) {
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 18618840d9..01953213ff 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 = arg.as<tirx::BufferVar>()) {
+ 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,11 @@ class TIRPatternMatcher {
for (const TIRPattern& pattern : patterns_) {
tirx::PrimFunc pattern_func = pattern;
ffi::Array<Var> pattern_symbolic_vars;
- int buffer_count = pattern_func->buffer_map.size();
+ int buffer_count = 0;
+ while (buffer_count < static_cast<int>(pattern_func->params.size()) &&
+ pattern_func->params[buffer_count]->ty.as<tirx::BufferTypeNode>()) {
+ ++buffer_count;
+ }
for (int i = buffer_count; i < static_cast<int>(pattern_func->params.size()); i++) {
pattern_symbolic_vars.push_back(pattern_func->params[i]);
}
@@ -614,7 +617,8 @@ std::pair<PrimFunc, ffi::Optional<PrimFunc>> 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 = func->params[i].as<tirx::BufferVar>();
+ 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 +626,17 @@ std::pair<PrimFunc, ffi::Optional<PrimFunc>> SplitFunctions(
}
}
arg_partition->push_back(arg_partition1);
- new_params1.push_back(Var("output", PointerType::VoidPointerTy()));
- ffi::Map<Var, BufferVar> 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<Var> new_params2;
std::vector<int> 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<int>(func->params.size()); i++) {
Var param = func->params[i];
- if (partitioner.input2.count(func->buffer_map[param])) {
+ auto param_buffer = param.as<tirx::BufferVar>();
+ if (param_buffer.has_value() && partitioner.input2.count(param_buffer.value())) {
new_params2.push_back(param);
if (i != static_cast<int>(func->params.size()) - 1) {
arg_partition2.push_back(i);
@@ -646,14 +644,7 @@ std::pair<PrimFunc, ffi::Optional<PrimFunc>> SplitFunctions(
}
}
arg_partition->push_back(arg_partition2);
- ffi::Map<Var, BufferVar> 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 +742,7 @@ class SplitMutator : public ExprMutator {
if (lib_func->IsInstance<tirx::PrimFuncNode>()) return ffi::GetRef<Call>(op);
TVM_FFI_ICHECK(lib_func->IsInstance<ExternFuncNode>());
builder_->UpdateFunction(gv, lib_func);
- tirx::BufferVar intermediate_buffer = func1->buffer_map.at(func1->params.back());
+ tirx::BufferVar intermediate_buffer = func1->params.back().as_or_throw<tirx::BufferVar>();
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 ff340d153f..b2a47810a9 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<Var> params;
- ffi::Map<Var, BufferVar> 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<Var> params = original_func_->params;
- ffi::Map<Var, BufferVar> 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(param.as<tirx::BufferVar>().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,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 = original_func_->buffer_map[original_func_->params[i]];
+ BufferVar buffer = original_func_->params[i].as_or_throw<BufferVar>();
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 b0de1d8f21..4fb7535c64 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<BufferVar, ffi::Optional<Stmt>> Detect(const PrimFunc& func) {
LCADetector detector;
- for (const auto& kv : func->buffer_map) {
- 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<BufferVar>()) {
+ 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/inject_texture_alloc.cc b/src/s_tir/backend/adreno/inject_texture_alloc.cc
index d60589e9ef..27544a2d80 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 5d14e314bd..3cb8c28136 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<Var, BufferVar>& extern_buffer_map,
- IRVisitorWithAnalyzer* bound_analyzer)
+ explicit TextureLoweringBase(const ffi::Array<Var>& 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<BufferVar>()) {
+ 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<Var, BufferVar>& extern_buffer_map,
- IRVisitorWithAnalyzer* bound_analyzer)
- : TextureLoweringBase(extern_buffer_map, bound_analyzer) {}
+ explicit TextureFlattener(const ffi::Array<Var>& params, IRVisitorWithAnalyzer* bound_analyzer)
+ : TextureLoweringBase(params, bound_analyzer) {}
Stmt VisitStmt_(const BufferStoreNode* op) final {
Stmt stmt = StmtExprMutator::VisitStmt_(op);
@@ -159,7 +159,7 @@ 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(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 94fec3bdf9..c36a952961 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> ArgInfo::FromPrimFunc(const tirx::PrimFunc& func) {
ffi::Array<ArgInfo> result;
result.reserve(func->params.size());
for (const tirx::Var& arg : func->params) {
- if (ffi::Optional<tirx::BufferVar> _buffer = func->buffer_map.Get(arg)) {
- tirx::BufferVar buffer = _buffer.value();
- result.push_back(TensorInfo(/*dtype=*/buffer->dtype->dtype,
- /*shape=*/AsVector<PrimExpr, int64_t>(buffer->shape)));
+ if (auto buffer = arg.as<tirx::BufferVar>()) {
+ result.push_back(TensorInfo(/*dtype=*/buffer.value()->dtype->dtype,
+ /*shape=*/AsVector<PrimExpr, int64_t>(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 a8578f6567..acf609f61c 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<PrimFuncNode>()) {
collector(func->body);
- for (const auto& it : func->buffer_map) {
- collector.HandleBufferAlloc(it.second);
+ for (const Var& param : func->params) {
+ if (auto buffer = param.as<BufferVar>()) {
+ 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 f6a439b650..2dce55cbe8 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<BufferVar> CollectLayoutFreeBuffers(const PrimFuncNode* func) {
for (int64_t index : layout_free_buffer_index) {
TVM_FFI_ICHECK(static_cast<size_t>(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(param.as_or_throw<BufferVar>());
}
LayoutFreeBufferCollector collector;
diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc
index 31ef367cb7..7d1c870553 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 : f->buffer_map) {
- 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<BufferVar>()) {
+ 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 4f2f15a230..16a551df74 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<BufferVar, BufferVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> 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]];
+ BufferVar desc = intrin_desc->params[i].as_or_throw<BufferVar>();
+ BufferVar impl = intrin_impl->params[i].as_or_throw<BufferVar>();
impl2desc[impl] = desc;
}
std::unordered_map<BufferVar, BufferVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> impl2cur;
@@ -821,7 +821,7 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int
ffi::Array<MatchBufferRegion> 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]);
+ BufferVar impl = intrin_impl->params[i].as_or_throw<BufferVar>();
const BufferVar& cur = impl2cur.at(impl);
const ffi::Array<Range>& old_region = impl2region.at(impl);
const std::vector<PrimExpr>& 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 9c8f835a0a..39f6f57b7a 100644
--- a/src/s_tir/schedule/primitive/layout_transformation.cc
+++ b/src/s_tir/schedule/primitive/layout_transformation.cc
@@ -1229,23 +1229,22 @@ 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<SBlock>();
- // 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);
IRModuleNode* new_mod = self->mod.CopyOnWrite();
ffi::MapObj* new_map = new_mod->functions.CopyOnWrite();
- ffi::Map<Var, BufferVar> new_buffer_map;
- for (auto [var, buffer] : old_func->buffer_map) {
- if (buffer.same_as(old_buffer)) {
- buffer = new_buffer;
+ ffi::Array<Var> new_params = old_func->params.Map([&](Var param) -> Var {
+ if (auto buffer = param.as<BufferVar>(); 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 6c2ad82710..0ed602945a 100644
--- a/src/s_tir/transform/inject_permuted_layout.cc
+++ b/src/s_tir/transform/inject_permuted_layout.cc
@@ -70,9 +70,10 @@ 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) {
- buffer_map_.insert({buffer.var(), buffer});
+ for (const Var& param : func->params) {
+ if (auto buffer = param.as<BufferVar>()) {
+ 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 2f20f10303..3514043765 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<ffi::String>(tvm::attr::kGlobalSymbol);
PipelineInjector injector(global_symbol);
- for (const auto& kv : func->buffer_map) {
- 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<BufferVar>()) {
+ 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 027c124d43..a4f1d0dcda 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<Var> GetBufferDataVar(const ffi::Any& data) {
class ThreadAllreduceBuilder final : public StmtExprMutator {
public:
- explicit ThreadAllreduceBuilder(const TargetNode* target, const ffi::Array<Var>& params,
- const ffi::Map<Var, BufferVar>& buffer_map)
+ explicit ThreadAllreduceBuilder(const TargetNode* target, const ffi::Array<Var>& params)
: target_(target),
warp_size_(target->GetAttr<int64_t>("thread_warp_size", 1).value()),
max_num_threads_(target->GetAttr<int64_t>("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<BufferTypeNode>()) {
buffer_aliases_.Set(param, param);
@@ -966,7 +962,7 @@ Pass LowerThreadAllreduce() {
auto target = f->GetAttr<Target>(tvm::attr::kTarget);
TVM_FFI_ICHECK(target.has_value()) << "LowerThreadAllreduce: Require the target attribute";
const TargetNode* target_node = target.as<TargetNode>();
- ThreadAllreduceBuilder thread_all_reduce(target_node, f->params, f->buffer_map);
+ 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 ccf5400cae..af33de6f20 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<BufferVar> Collect(const PrimFunc& func) {
BufferAllocateOrderCollector collector;
- for (const auto& kv : func->buffer_map) {
- collector.buffer_alloc_recorder_.push_back(kv.second);
+ for (const Var& param : func->params) {
+ if (auto buffer = param.as<BufferVar>()) {
+ 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 : func->buffer_map) {
- 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<BufferVar>()) {
+ 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/s_tir/transform/remove_weight_layout_rewrite_block.cc b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc
index 8cb5149540..a987c1abeb 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<tirx::Var, BufferVar> buffer_map;
- for (const auto& [param, buffer] : f_->buffer_map) {
+ ffi::Array<tirx::Var> params;
+ for (const tirx::Var& param : f_->params) {
+ auto opt_buffer = param.as<BufferVar>();
+ 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 519a82d8be..1e6d312904 100644
--- a/src/s_tir/transform/renew_defs.cc
+++ b/src/s_tir/transform/renew_defs.cc
@@ -48,40 +48,41 @@ class RenewDefMutator : public StmtExprMutator {
public:
static PrimFunc Transform(const PrimFunc& func) {
RenewDefMutator generator;
- // Redefine params
- ffi::Array<Var> 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 (!param.as<BufferVar>()) {
+ 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 = param.as<BufferVar>()) {
+ const BufferVar& buffer = opt_buffer.value();
for (const PrimExpr& e : buffer->shape) {
- if (auto var = e.as<PrimVar>()) {
- if (generator.remap_.count(var.value()) == 0) {
- generator.ReDefineVar(var.value());
+ PostOrderVisit(e, [&generator](const ffi::ObjectRef& obj) {
+ if (auto var = obj.as<Var>()) {
+ if (generator.remap_.count(var.value()) == 0) {
+ generator.ReDefineVar(var.value());
+ }
}
- }
+ });
}
}
}
- // Redefine buffers in order
+ // Redefine buffer parameters in order, preserving the original signature.
// TODO(Siyuan Feng): checking var is used after define
- ffi::Map<tirx::Var, BufferVar> buffer_map;
+ ffi::Array<Var> 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<Var>();
- BufferVar new_buffer = generator.DefineBuffer(buffer);
- buffer_map.Set(new_param, new_buffer);
+ if (auto opt_buffer = param.as<BufferVar>()) {
+ params.push_back(generator.DefineBuffer(opt_buffer.value()));
+ } else {
+ params.push_back(generator.VisitExpr(param).as_or_throw<Var>());
}
}
// 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:
@@ -184,7 +185,7 @@ class RenewDefMutator : public StmtExprMutator {
}
};
- // shape is USED (references existing definitions like buffer_map shape vars),
+ // shape is USED (references existing definitions like buffer-parameter shape vars),
// remap via VisitExpr to avoid creating spurious new var definitions
auto visit_expr = [this](const PrimExpr& e) -> PrimExpr { return this->VisitPrimExpr(e); };
ffi::Array<PrimExpr> shape = buffer->shape.Map(visit_expr);
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index f8fd8d1b8a..dc454de5ad 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(func->buffer_map.size(), 0U)
- << "Cannot codegen function with buffer_map, please lower them first";
+ for (const Var& param : func->params) {
+ TVM_FFI_ICHECK(!param->ty.as<tirx::BufferTypeNode>())
+ << "Cannot codegen BufferType-annotated parameter " << param << "; please lower it first";
+ }
std::vector<llvm::Type*> 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 90c8922c3c..7a8c4ae80b 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<Var>()) {
- if (ffi::Optional<BufferVar> buffer = func->buffer_map.Get(v.value())) {
- buffer2index_[buffer.value()] = i;
- }
+ if (auto buffer = func->params[i].as<BufferVar>()) {
+ buffer2index_[buffer.value()] = i;
}
}
PrimFuncNode* f = func.CopyOnWrite();
@@ -753,19 +751,15 @@ void RewriteStageToBlock(const te::Operation& op, CreateFuncInfo* info,
PrimFunc GenerateAndCompletePrimFunc(const ffi::Array<te::Tensor>& arg_list,
const ffi::Array<Stmt>& root_stmts, CreateFuncInfo* info) {
ffi::Array<Var> parameters;
- ffi::Map<Var, BufferVar> 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 +813,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
PrimFunc GenerateAndCompletePrimFunc(const ffi::Array<ffi::ObjectRef>& arg_tir_var_list,
const ffi::Array<Stmt>& root_stmts, CreateFuncInfo* info) {
ffi::Array<Var> parameters;
- ffi::Map<Var, BufferVar> buffer_map;
for (const ffi::ObjectRef& arg : arg_tir_var_list) {
if (auto opt_tensor = arg.as<te::Tensor>()) {
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<tirx::PrimVar>()) {
parameters.push_back(var.value());
}
@@ -835,8 +826,7 @@ PrimFunc GenerateAndCompletePrimFunc(const ffi::Array<ffi::ObjectRef>& 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 2656bc521b..5694284614 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 : func_->buffer_map) {
- if (V == kv.second.get()) return true;
+ for (const Var& param : func_->params) {
+ if (param->ty.as<BufferTypeNode>() && V == param.get()) return true;
}
while (true) {
diff --git a/src/tirx/analysis/verify_ssa.cc b/src/tirx/analysis/verify_ssa.cc
index b14dd6062e..f9b56dd9e7 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 : func->buffer_map) {
- this->DefineBuffer(kv.second);
+ for (const Var& param : func->params) {
+ if (auto buffer = param.as<BufferVar>()) {
+ this->DefineBuffer(buffer.value());
+ }
}
this->VisitStmt(func->body);
}
diff --git a/src/tirx/analysis/verify_well_formed.cc b/src/tirx/analysis/verify_well_formed.cc
index a2bb35aae6..2fc5c3a4b7 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<UndefinedVarVerifier> {
/*! \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 a8cb4f62e1..d8a6fff6cc 100644
--- a/src/tirx/ir/data_type_rewriter.cc
+++ b/src/tirx/ir/data_type_rewriter.cc
@@ -635,14 +635,18 @@ PrimFunc IndexDataTypeNormalizer::Rewrite(PrimFunc func) {
buffer_remap_.clear();
ivmap_.clear();
// start rewrite
- ffi::Map<Var, BufferVar> 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 Var& param : func->params) {
+ if (auto buffer = param.as<BufferVar>()) {
+ VisitBufferDef(buffer.value(), /*alloc_data=*/true);
+ }
}
// remap params
bool is_enabled = true;
std::swap(is_enabled_, is_enabled);
ffi::Array<Var> params = func->params.Map([this](Var param) {
+ if (auto buffer = param.as<BufferVar>()) {
+ return buffer_remap_.Get(buffer.value()).value_or(buffer.value()).var();
+ }
if (auto param_ty = param->ty.as<PrimType>();
param_ty && param_ty.value().MatchesCode(DLDataTypeCode::kDLInt)) {
return this->VisitPrimExpr(param.as_or_throw<PrimExpr>()).as_or_throw<Var>();
@@ -654,7 +658,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 c991e572dc..463a15084f 100644
--- a/src/tirx/ir/function.cc
+++ b/src/tirx/ir/function.cc
@@ -42,8 +42,8 @@ tvm::Type InferType(const PrimFunc& prim_func) {
ffi::Array<tvm::Type> 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<BufferTypeNode>()) {
+ 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 +79,8 @@ tvm::Type InferType(const PrimFunc& prim_func) {
} // namespace
// Get the function type of a PrimFunc
-PrimFunc::PrimFunc(ffi::Array<tirx::Var> params, Stmt body, Type ret_type,
- ffi::Map<tirx::Var, BufferVar> buffer_map, DictAttrs attrs, Span span) {
+PrimFunc::PrimFunc(ffi::Array<tirx::Var> params, Stmt body, Type ret_type, DictAttrs attrs,
+ Span span) {
if (ret_type.IsMissing()) {
ret_type = VoidType();
}
@@ -89,7 +89,6 @@ PrimFunc::PrimFunc(ffi::Array<tirx::Var> 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);
@@ -121,7 +120,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<PointerTypeNode>() != nullptr; };
+ auto is_handle = [](const Var& param) {
+ return param->ty.as<PointerTypeNode>() != nullptr || param->ty.as<BufferTypeNode>() != 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 +131,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<TensorIntrinNode> n = ffi::make_object<TensorIntrinNode>();
n->desc = std::move(desc);
n->impl = std::move(impl);
@@ -164,9 +163,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
.def("tirx.PrimFunc",
- [](ffi::Array<tirx::Var> params, Stmt body, Type ret_type,
- ffi::Map<tirx::Var, BufferVar> buffer_map, DictAttrs attrs,
- Span span) { return PrimFunc(params, body, ret_type, buffer_map, attrs, span); })
+ [](ffi::Array<tirx::Var> 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 ae128f20ad..3e7bd254e7 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<BufferVar>& root_allocates, bool s_tir) {
ffi::Map<Var, BufferVar> buffer_var_map;
- for (const auto& pair : func->buffer_map) {
- const BufferVar& buffer = pair.second;
- buffer_var_map.Set(buffer.var(), buffer);
+ for (const Var& param : func->params) {
+ if (auto buffer = param.as<BufferVar>()) {
+ 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 687b37dd3a..3890829dc1 100644
--- a/src/tirx/ir/specialize.cc
+++ b/src/tirx/ir/specialize.cc
@@ -31,6 +31,7 @@
#include <tvm/tirx/stmt_functor.h>
#include <functional>
+#include <unordered_set>
#include "../transform/ir_utils.h"
#include "functor_common.h"
@@ -78,27 +79,40 @@ class PrimFuncSpecializer : public StmtExprMutator {
static PrimFunc Specialize(PrimFunc f, const VarMap& var_map) {
PrimFuncSpecializer specializer(var_map);
- // Updating BufferVar map
- ffi::Map<Var, BufferVar> buffer_map;
- bool buffer_map_updated = false;
- for (const auto& it : f->buffer_map) {
- 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;
+ for (const Var& param : f->params) {
+ auto buffer = param.as<BufferVar>();
+ auto replacement = var_map.find(param);
+ if (!buffer || replacement == var_map.end()) {
+ continue;
+ }
+ if (auto replacement_var = replacement->second.as<Var>()) {
+ if (auto replacement_buffer = replacement_var.value().as<BufferVar>()) {
+ if (IsParam(f, replacement_var.value())) {
+ specializer.buffer_aliases_[buffer.value()] = replacement_buffer.value();
+ } else {
+ specializer.constrained_buffer_params_.insert(param.get());
+ }
+ }
}
}
- // Updating parmeters
+ // Updating parameters
ffi::Array<Var> params;
bool param_updated = false;
for (const auto& var : f->params) {
+ Var new_var = var;
+ if (auto buffer = var.as<BufferVar>()) {
+ 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()) {
- params.push_back(var);
+ if (var_map.find(var) == var_map.end() ||
+ specializer.constrained_buffer_params_.count(var.get())) {
+ params.push_back(new_var);
} else {
param_updated = true;
}
@@ -107,8 +121,8 @@ class PrimFuncSpecializer : public StmtExprMutator {
// Updating function body
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);
+ if (param_updated || !f->body.same_as(body)) {
+ return PrimFunc(params, body, f->ret_type, f->attrs, f->span);
} else {
return f;
}
@@ -160,6 +174,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<Var>(op);
+ }
auto it = var_map_.find(ffi::GetRef<Var>(op));
if (it == var_map_.end()) {
return ffi::GetRef<Var>(op);
@@ -189,6 +206,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<ffi::String> specialized_storage_scope;
if (auto it = var_map_.find(buffer.var()); it != var_map_.end()) {
if (const auto* new_var = it->second.as<VarNode>()) {
@@ -283,7 +304,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 +328,10 @@ class PrimFuncSpecializer : public StmtExprMutator {
const VarMap& var_map_;
/*! \brief map from old buffer to mutated buffer */
std::unordered_map<BufferVar, BufferVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> buffer_map_;
+ /*! \brief Direct aliases between buffer parameters. */
+ std::unordered_map<BufferVar, BufferVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> buffer_aliases_;
+ /*! \brief Buffer parameters constrained by a concrete, non-parameter buffer. */
+ std::unordered_set<const VarNode*> constrained_buffer_params_;
};
/*!
@@ -317,7 +342,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 +356,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 = param.as<BufferVar>();
+ 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 +404,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 +431,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(!param.as<BufferVar>(), 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 dd9bcc88d1..4b39167651 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<std::variant<DefContext<Var>, DefContext<BufferVar>>> 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<BufferTypeNode>()) {
- context.push_back(WithDef(BufferVar(param), ppath->ArrayItem(i)));
- } else {
+ if (!param->ty.as<BufferTypeNode>()) {
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 = func->params[i].as<BufferVar>()) {
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 = func->params[i].as<BufferVar>()) {
+ context.push_back(WithDef(opt.value(), ppath->ArrayItem(i)));
}
}
diff --git a/src/tirx/ir/tir_visitor_with_path.h b/src/tirx/ir/tir_visitor_with_path.h
index 85c42570f6..e9735a89b8 100644
--- a/src/tirx/ir/tir_visitor_with_path.h
+++ b/src/tirx/ir/tir_visitor_with_path.h
@@ -255,16 +255,27 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const Expr&, ffi::reflecti
}
}
};
- auto try_visit_implicit_var_def_array = [&try_visit_implicit_var_def](
- const ffi::Array<PrimExpr>& arr,
- ffi::reflection::AccessPath path) {
- for (size_t i = 0; i < arr.size(); i++) {
- try_visit_implicit_var_def(arr[i], path->ArrayItem(i));
- }
- };
- try_visit_implicit_var_def_array(buf->shape, path->Attr("shape"));
- try_visit_implicit_var_def_array(buf->strides, path->Attr("strides"));
+ // A Buffer shape is a match scope. The first shape expression that
+ // contains an undefined Var defines it, even when the expression is
+ // compound (for example, `n + 1`). Later expressions then see the same
+ // Var in `in_scope_definitions_` and reuse it.
+ auto shape_path = path->Attr("shape");
+ for (size_t i = 0; i < buf->shape.size(); i++) {
+ auto dim_path = shape_path->ArrayItem(i);
+ PostOrderVisit(buf->shape[i], [this, &context, &dim_path](const ffi::ObjectRef& obj) {
+ if (auto opt = obj.as<Var>()) {
+ if (auto var_def = WithDefIfUndefined(opt.value(), dim_path)) {
+ context.push_back(std::move(var_def).value());
+ }
+ }
+ });
+ }
+
+ auto strides_path = path->Attr("strides");
+ for (size_t i = 0; i < buf->strides.size(); i++) {
+ try_visit_implicit_var_def(buf->strides[i], strides_path->ArrayItem(i));
+ }
try_visit_implicit_var_def(buf->elem_offset, path->Attr("elem_offset"));
return context;
diff --git a/src/tirx/script/builder/frame.cc b/src/tirx/script/builder/frame.cc
index a253ee8cc3..7a41ad7a15 100644
--- a/src/tirx/script/builder/frame.cc
+++ b/src/tirx/script/builder/frame.cc
@@ -113,39 +113,53 @@ void PrimFuncFrameNode::ExitWithScope() {
}
// s_tir-mode normalization: drop stale default layouts (see comment on
// STirBufferLayoutNormalizer above) and rewrite body references coherently.
- ffi::Map<tvm::tirx::Var, tvm::tirx::BufferVar> effective_buffer_map = buffer_map;
ffi::Array<tvm::tirx::BufferVar> effective_root_alloc_buffers = root_alloc_buffers;
tvm::tirx::Stmt body = AsStmt(stmts);
- if (s_tir) {
- STirBufferLayoutNormalizer normalizer;
- ffi::Map<tvm::tirx::Var, tvm::tirx::BufferVar> new_buffer_map;
- for (const auto& kv : buffer_map) {
- tvm::tirx::BufferVar buf = kv.second;
- if (buf->layout.has_value()) {
- ffi::ObjectPtr<tvm::tirx::BufferTypeNode> 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<tvm::tirx::Var> effective_args;
+ ffi::Map<tvm::tirx::Var, tvm::Expr> param_replacements;
+ for (const tvm::tirx::Var& arg : args) {
+ ffi::Optional<tvm::tirx::BufferVar> opt_buffer = buffer_map.Get(arg);
+ bool replaces_legacy_param = opt_buffer.has_value();
+ if (!opt_buffer.has_value() && arg->ty.as<tvm::tirx::BufferTypeNode>()) {
+ opt_buffer = tvm::tirx::BufferVar(arg);
}
- if (!normalizer.Empty()) {
- body = normalizer(std::move(body));
- ffi::Array<tvm::tirx::BufferVar> 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<tvm::tirx::BufferTypeNode> 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 (replaces_legacy_param && !arg.same_as(buffer.var()) &&
+ !arg->ty.as<tvm::tirx::BufferTypeNode>()) {
+ 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));
+ ffi::Array<tvm::tirx::BufferVar> 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);
+ }
+ if (!param_replacements.empty()) {
+ body = tvm::tirx::Substitute(std::move(body), param_replacements);
}
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 97e24e0f62..1538ec6836 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 943b211004..e161a03868 100644
--- a/src/tirx/script/printer/buffer.cc
+++ b/src/tirx/script/printer/buffer.cc
@@ -19,6 +19,7 @@
#include <tvm/runtime/device_api.h> // For `kAllocAlignment`
#include <algorithm>
+#include <utility>
#include "./utils.h"
@@ -26,10 +27,11 @@ namespace tvm {
namespace script {
namespace printer {
-ffi::Map<ffi::String, ExprDoc> BufferAttrs(tirx::BufferVar buffer, const AccessPath& buffer_p,
- const Frame& frame, const IRDocsifier& d,
- BufferVarDefinition var_definitions,
- ffi::Optional<Expr> data = std::nullopt) {
+ffi::Map<ffi::String, ExprDoc> BufferAttrs(
+ tirx::BufferVar buffer, const AccessPath& buffer_p, const Frame& frame, const IRDocsifier& d,
+ BufferVarDefinition var_definitions, ffi::Optional<Expr> data = std::nullopt,
+ bool stringify_undefined_shape = false,
+ std::unordered_set<tirx::Var> stringify_shape_vars = {}) {
using tvm::tirx::Var;
using tvm::tirx::VarNode;
ffi::Map<ffi::String, ExprDoc> kwargs;
@@ -85,10 +87,27 @@ ffi::Map<ffi::String, ExprDoc> 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;
+ std::unordered_set<Var> vars_in_shape;
+ tirx::PostOrderVisit(e, [&](const ffi::ObjectRef& obj) {
+ if (const auto* var_node = obj.as<VarNode>()) {
+ Var var = ffi::GetRef<Var>(var_node);
+ vars_in_shape.insert(var);
+ contains_new_var =
+ contains_new_var || !d->IsVarDefined(var) || stringify_shape_vars.count(var);
+ }
+ });
if (is_new_var(e)) {
add_out_of_line_var_def(e.as_or_throw<Var>(), e_p);
}
- results.push_back(d->AsDoc<ExprDoc>(e, e_p));
+ ExprDoc result = d->AsDoc<ExprDoc>(e, e_p);
+ results.push_back(stringify_undefined_shape && contains_new_var ? ExprStringDoc(result, e_p)
+ : result);
+ // A quoted shape expression defines every Var it contains. Do not quote
+ // later dimensions merely because they reuse a Var introduced here.
+ for (const Var& var : vars_in_shape) {
+ stringify_shape_vars.erase(var);
+ }
}
kwargs.Set("shape", TupleDoc(results));
}
@@ -316,13 +335,14 @@ 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, std::unordered_set<tirx::Var> stringify_shape_vars) {
ffi::Map<ffi::String, ExprDoc> 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, std::nullopt, true,
+ std::move(stringify_shape_vars));
+ if (!attrs.count("dtype")) {
+ attrs.Set("dtype", LiteralDoc::DataType(buffer->dtype->dtype, p->Attr("dtype")));
+ }
+ return BufferCall(TIR(d, "Buffer"), attrs, {});
}
ffi::Array<Doc> BufferIndices(const ffi::Array<PrimExpr>& indices, const AccessPath& p,
diff --git a/src/tirx/script/printer/function.cc b/src/tirx/script/printer/function.cc
index a1eddde6f2..f6ba953b86 100644
--- a/src/tirx/script/printer/function.cc
+++ b/src/tirx/script/printer/function.cc
@@ -16,68 +16,16 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/runtime/device_api.h>
#include <tvm/runtime/logging.h>
+#include <utility>
+
#include "./utils.h"
namespace tvm {
namespace script {
namespace printer {
-bool IsSimpleBuffer(const tirx::BufferVar& buf, bool s_tir) {
- if (!buf->strides.empty()) {
- return false;
- }
- for (const PrimExpr& shp_i : buf->shape) {
- if (!tirx::UndefinedVars(shp_i).empty()) {
- return false;
- }
- }
- for (const PrimExpr& stride_i : buf->strides) {
- if (!tirx::UndefinedVars(stride_i).empty()) {
- return false;
- }
- }
- if (!tirx::UndefinedVars(buf->elem_offset).empty()) {
- return false;
- } else if (buf->elem_offset->IsInstance<IntImmNode>()) {
- IntImm elem_offset = buf->elem_offset.as_or_throw<IntImm>();
- if (elem_offset->value != 0) {
- return false;
- }
- }
- if (s_tir) {
- if (buf->layout.has_value() &&
- !ffi::StructuralEqual()(buf->layout, tirx::TileLayoutNode::DefaultLayout(buf->shape))) {
- return false;
- }
- } else {
- if (!buf->layout.has_value() ||
- !ffi::StructuralEqual()(buf->layout, tirx::TileLayoutNode::DefaultLayout(buf->shape))) {
- return false;
- }
- }
- if (!buf->allocated_addr.empty()) {
- return false;
- }
- return buf.scope() == "global" && buf->data_alignment == runtime::kAllocAlignment &&
- buf->offset_factor == 1;
-}
-
-int CountVarOccurrence(const tirx::PrimFunc& f, const tirx::Var& v) {
- OccurrenceCounter counter(v.get());
- counter(f->body);
- 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;
-}
-
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
.set_dispatch<tirx::PrimFunc>("", [](tirx::PrimFunc func, AccessPath p, IRDocsifier d) -> Doc {
With<TIRFrame> f(d, func);
@@ -87,36 +35,57 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
return obj->IsInstance<tirx::VarNode>() || obj->IsInstance<tirx::BufferTypeNode>();
});
int n_args = func->params.size();
- std::unordered_map<const tirx::VarNode*, int> 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<AssignDoc> args;
args.reserve(n_args);
- std::unordered_set<const tirx::VarNode*> buffer_inlined;
+ std::unordered_map<const tirx::VarNode*, ExprDoc> scalar_param_docs;
+ // Define scalar docs up front so a preceding Buffer parameter can render
+ // a reference to a later scalar parameter. `bound_signature_vars`
+ // separately tracks source order: the first shape expression that sees
+ // an unbound Var must be quoted because Buffer shapes are match scopes.
+ std::unordered_set<tirx::Var> bound_signature_vars;
+ for (const tirx::Var& param : func->params) {
+ if (!param->ty.as<tirx::BufferTypeNode>()) {
+ 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::BufferTypeNode>()) {
+ tirx::BufferVar buffer(var);
+ std::unordered_set<tirx::Var> stringify_shape_vars;
+ std::unordered_set<tirx::Var> shape_vars;
+ for (const PrimExpr& shape : buffer->shape) {
+ tirx::PostOrderVisit(shape, [&](const ffi::ObjectRef& obj) {
+ if (const auto* shape_var_node = obj.as<tirx::VarNode>()) {
+ tirx::Var shape_var = ffi::GetRef<tirx::Var>(shape_var_node);
+ shape_vars.insert(shape_var);
+ if (!bound_signature_vars.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, std::move(stringify_shape_vars));
+ args.push_back(AssignDoc(lhs, std::nullopt, annotation));
+ for (const tirx::Var& shape_var : shape_vars) {
+ bound_signature_vars.insert(shape_var);
+ }
+ continue;
}
ExprDoc a = d->AsDoc<ExprDoc>(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));
+ bound_signature_vars.insert(var);
+ }
+ ffi::Optional<ExprDoc> ret_type = std::nullopt;
+ if (!func->ret_type.IsMissing()) {
+ const auto* as_tuple = func->ret_type.as<TupleTypeNode>();
+ if (!as_tuple || as_tuple->fields.size()) {
+ ret_type = d->AsDoc<ExprDoc>(func->ret_type, p->Attr("ret_type"));
+ }
}
// Step 2. Handle `func->attrs`
if (!func->attrs->dict.empty()) {
@@ -147,23 +116,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
->Call({d->AsDoc<ExprDoc>(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<tirx::SBlock> implicit_root_block = [&]() -> ffi::Optional<tirx::SBlock> {
const tirx::SBlockRealizeNode* root_block_realize =
func->body.as<tirx::SBlockRealizeNode>();
@@ -201,13 +154,6 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
} else {
AsDocBody(func->body, p->Attr("body"), f->get(), d);
}
- ffi::Optional<ExprDoc> ret_type = std::nullopt;
- if (!func->ret_type.IsMissing()) {
- const auto* as_tuple = func->ret_type.as<TupleTypeNode>();
- if (!as_tuple || as_tuple->fields.size()) {
- ret_type = d->AsDoc<ExprDoc>(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<ffi::String, void> kwargs_keys;
diff --git a/src/tirx/script/printer/utils.h b/src/tirx/script/printer/utils.h
index 296c24cd7a..2c38ec59e2 100644
--- a/src/tirx/script/printer/utils.h
+++ b/src/tirx/script/printer/utils.h
@@ -35,6 +35,7 @@
#include <string>
#include <unordered_map>
+#include <unordered_set>
#include <utility>
#include <vector>
@@ -286,7 +287,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,
};
@@ -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. The set is
+ * passed by value so entries can be consumed as dimensions are emitted.
* \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, std::unordered_set<tirx::Var> stringify_shape_vars = {});
/*!
* \brief Print the creation of a Var
@@ -335,57 +338,6 @@ LambdaDoc PrintLambda(const ffi::ObjectRef& pred, const ffi::Array<tirx::Var>& v
const AccessPath& vs_p, const PrimExpr& p, const AccessPath& p_p,
const IRDocsifier& d);
-/*! \brief A Var occurrence counter visitor */
-class OccurrenceCounter : public tirx::StmtExprVisitor {
- public:
- /*! \brief The occurrence counter */
- int count = 0;
- /*! \brief The Var to count occurrence */
- const tirx::VarNode* v = nullptr;
-
- void VisitVar(const tirx::Var& var) { VisitExpr(static_cast<const Expr&>(var)); }
-
- void VisitExpr_(const tirx::VarNode* op) final {
- if (op == v) {
- ++count;
- }
- tirx::StmtExprVisitor::VisitExpr_(op);
- }
-
- void VisitStmt_(const tirx::BufferStoreNode* op) final {
- VisitBuffer(op->buffer);
- tirx::StmtExprVisitor::VisitStmt_(op);
- }
-
- void VisitExpr_(const tirx::BufferLoadNode* op) final {
- VisitBuffer(op->buffer);
- tirx::StmtExprVisitor::VisitExpr_(op);
- }
-
- void VisitStmt_(const tirx::AllocBufferNode* op) final {
- VisitBuffer(op->buffer);
- tirx::StmtExprVisitor::VisitStmt_(op);
- }
-
- void VisitStmt_(const tirx::DeclBufferNode* op) final {
- VisitBuffer(op->buffer);
- tirx::StmtExprVisitor::VisitStmt_(op);
- }
-
- void VisitBuffer(const tirx::BufferVar& buffer) {
- VisitExpr(buffer.var());
- for (const PrimExpr& shape_i : buffer->shape) {
- VisitExpr(shape_i);
- }
- for (const PrimExpr& stride_i : buffer->strides) {
- VisitExpr(stride_i);
- }
- VisitExpr(buffer->elem_offset);
- }
-
- explicit OccurrenceCounter(const tirx::VarNode* var) { v = var; }
-};
-
#ifndef TVM_SCRIPT_REPR
#define TVM_SCRIPT_REPR(ObjectType, Method) TVM_REGISTER_SCRIPT_AS_REPR(ObjectType, Method)
#endif
diff --git a/src/tirx/transform/flatten_buffer.cc b/src/tirx/transform/flatten_buffer.cc
index 3dbb703f99..d901d91434 100644
--- a/src/tirx/transform/flatten_buffer.cc
+++ b/src/tirx/transform/flatten_buffer.cc
@@ -62,25 +62,24 @@ 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.extern_buffers_.insert(buffer);
- pass.Define(buffer);
+ pass.MarkBufferParamShapes(func);
+ for (const Var& param : func->params) {
+ if (auto buffer = param.as<BufferVar>()) {
+ pass.extern_buffers_.insert(buffer.value());
+ pass.Define(buffer.value());
+ }
}
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)) {
- 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<BufferVar>()) {
+ 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 5a4eb54d9e..2e444c273b 100644
--- a/src/tirx/transform/force_narrow_index_to_i32.cc
+++ b/src/tirx/transform/force_narrow_index_to_i32.cc
@@ -37,11 +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 : func->buffer_map) {
- 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;
+ for (const Var& param : func->params) {
+ if (auto buffer = param.as<BufferVar>();
+ 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 abe5d2850a..4b3cafb42d 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 = prim_func->buffer_map.size();
- if (has_buffer_arguments) return false;
+ for (const Var& param : prim_func->params) {
+ if (param->ty.as<BufferTypeNode>()) return false;
+ }
// We do not currently support inlining of schedulable TIR
// functions. To support this use case, repeated names in
@@ -228,9 +229,11 @@ class PrimFuncInliner : StmtExprMutator {
<< callee->params << "), but is called with " << args.size() << " arguments (" << args
<< ")";
- TVM_FFI_ICHECK(callee->buffer_map.empty())
- << "Inlining of PrimFuncs with buffer arguments is not yet supported, "
- << "but callee " << gvar << " has non-empty buffer map " << callee->buffer_map;
+ for (const Var& param : callee->params) {
+ TVM_FFI_ICHECK(!param->ty.as<BufferTypeNode>())
+ << "Inlining of PrimFuncs with buffer arguments is not yet supported, "
+ << "but callee " << gvar << " has BufferType-annotated parameter " << param;
+ }
ffi::Map<Var, ffi::Variant<tirx::BufferVar, tvm::Expr>> 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 e0e4dab78e..dabad44f55 100644
--- a/src/tirx/transform/ir_utils.cc
+++ b/src/tirx/transform/ir_utils.cc
@@ -110,53 +110,48 @@ class IRConvertSSA final : public StmtExprMutator {
for (const auto& var : func->params) {
defined_params.insert(var.get());
}
- for (const auto& [var, buffer] : func->buffer_map) {
- static_cast<void>(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<VarNode>();
- if (!var_ptr) return;
+ std::unordered_set<const VarNode*> defined_match_vars;
+ for (const Var& param : func->params) {
+ auto buffer = param.as<BufferVar>();
+ if (!buffer) continue;
+ auto check_var = [&](const Var& var) {
+ const VarNode* var_ptr = var.get();
if (defined_params.count(var_ptr)) return;
+ if (!defined_match_vars.insert(var_ptr).second) 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)) {
+ if (defined_.count(var_ptr)) {
+ Var new_var = MakeNewVar(var);
+ PushVarRemap(var, new_var);
+ } else {
defined_.insert(var_ptr);
}
};
- for (const auto& dim : buffer->shape) {
- check_expr(dim);
+ for (const auto& dim : buffer.value()->shape) {
+ PostOrderVisit(dim, [&](const ffi::ObjectRef& obj) {
+ if (auto var = obj.as<Var>()) check_var(var.value());
+ });
}
- for (const auto& stride : buffer->strides) {
- check_expr(stride);
+ for (const auto& stride : buffer.value()->strides) {
+ if (auto var = stride.as<Var>()) check_var(var.value());
}
- check_expr(buffer->elem_offset);
+ if (auto var = buffer.value()->elem_offset.as<Var>()) check_var(var.value());
}
}
- // Update the buffer map, based on the redefined parameters
- auto buffer_map = [&]() {
- ffi::Map<Var, BufferVar> buffer_map;
- bool made_change = false;
- for (const auto& [var, buffer] : func->buffer_map) {
- 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());
+ // Update the buffer parameters, based on the redefined parameters
+ bool buffer_params_changed = false;
+ for (size_t i = 0; i < func->params.size(); ++i) {
+ if (auto buffer = func->params[i].as<BufferVar>()) {
+ 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 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) {
- return buffer_map;
- } else {
- return func->buffer_map;
}
- }();
+ }
auto attrs = [&]() -> DictAttrs {
ffi::Map<ffi::String, ffi::Any> dict;
@@ -184,9 +179,9 @@ class IRConvertSSA final : public StmtExprMutator {
auto body = VisitStmt(func->body);
// 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 4bf6176495..d20cf7a0ea 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<Stmt, ffi::Map<Var, BufferVar>> Flatten(
- const Stmt& stmt, const ffi::Array<Var>& params,
- const ffi::Map<tirx::Var, BufferVar> buffer_map, const Target& target) {
+ static std::pair<Stmt, ffi::Array<Var>> Flatten(const Stmt& stmt, const ffi::Array<Var>& params,
+ const Target& target) {
arith::Analyzer ana;
LayoutApplier storage_lower(ana, target);
+ ffi::Array<Var> new_params;
+ new_params.reserve(params.size());
+ std::vector<std::pair<BufferVar, BufferVar>> param_flattened_buffers;
for (const Var& param : params) {
- if (param->ty.as<BufferTypeNode>()) {
- storage_lower.buffer_aliases_.Set(param, param);
+ auto buffer = param.as<BufferVar>();
+ if (!buffer) {
+ new_params.push_back(param);
+ continue;
}
- }
- std::unordered_map<Var, BufferVar> new_buffer_map;
- std::vector<std::pair<BufferVar, BufferVar>> 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<Var, BufferVar>(new_buffer_map));
+ return std::make_pair(new_stmt, new_params);
}
protected:
@@ -404,8 +404,9 @@ 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, params] = LayoutApplier::Flatten(n->body, n->params, target);
+ 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 8ad6dbaa03..0673692f97 100644
--- a/src/tirx/transform/make_packed_api.cc
+++ b/src/tirx/transform/make_packed_api.cc
@@ -230,8 +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, func_ptr->buffer_map, v_packed_args,
- v_num_packed_args, device_type, device_id.as_or_throw<PrimExpr>());
+ TVMFFIABIBuilder binder(name_hint, func_ptr->params, v_packed_args, v_num_packed_args,
+ device_type, device_id.as_or_throw<PrimExpr>());
binder.DecodeAllParams();
auto result = binder.Finalize();
@@ -279,7 +279,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<Var, BufferVar>();
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 cd87fc6087..fc6890666a 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<Var, int, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> 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[cur_func_->params[i].as_or_throw<BufferVar>().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 1be9222ec2..7875adcb95 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 ba5e893c90..00116b968f 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<Var>& params,
- const ffi::Map<Var, BufferVar>& buffer_map) {
- for (const auto& [_, buffer] : buffer_map) {
- buffer_aliases_.Set(buffer.var(), buffer.var());
- }
+ explicit LinearAccessPatternFinder(const ffi::Array<Var>& params) {
for (const Var& param : params) {
- if (param->ty.as<BufferTypeNode>()) {
- buffer_aliases_.Set(param, param);
+ if (auto buffer = param.as<BufferVar>()) {
+ 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<Var>& params, const ffi::Map<Var, BufferVar>& buffer_map,
- bool detect_inplace, bool enable_reuse, bool reuse_require_exact_matched_dtype) {
+ Stmt Rewrite(Stmt stmt, const ffi::Array<Var>& 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,
@@ -1185,8 +1181,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 +1197,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,49 +1270,31 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
*
* @param params The parameters passed to a PrimFunc
*
- * @param buffer_map The buffer_map associated with a PrimFunc
- *
* @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<tirx::Var>& params,
- const ffi::Map<Var, BufferVar>& buffer_map,
- bool allow_untyped_pointers = false,
+ VectorTypeAccessChecker(const ffi::Array<tirx::Var>& 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::kPrimFuncParam);
- }
-
- // If a pointer parameter isn't in the buffer map, then we want to
- // track the parameter itself.
for (Var buffer_var : params) {
- if (auto buffer_type = buffer_var->ty.as<BufferType>()) {
- BufferVar buffer(buffer_var);
+ if (auto buffer = buffer_var.as<BufferVar>()) {
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);
+ 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::kPrimFuncBufferMap);
- } 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, dtype, extent, BufferVarInfo::kPrimFuncPointerParam);
+ } else if (pointer_type.has_value() && allow_untyped_pointers_) {
+ OnArrayDeclaration(buffer_var, PrimType::Void(), 0, BufferVarInfo::kPrimFuncPointerParam);
}
}
}
@@ -1578,11 +1556,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 +1572,17 @@ class VectorTypeRewriter : public StmtExprMutator {
* should be re-written.
*/
VectorTypeRewriter(const std::unordered_map<const VarNode*, BufferVarInfo>& info_map,
- const ffi::Map<Var, Var>& buffer_aliases, bool rewrite_params = true,
- bool rewrite_buffer_map = true, bool rewrite_alloc_buffer_node = true,
+ const ffi::Map<Var, Var>& 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 +1881,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<Var, BufferVar> 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:
@@ -1938,19 +1905,20 @@ 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_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,
+ 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_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 +1948,8 @@ 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, 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 2404cd3ae2..a4da6880e0 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<Var>& params,
- const ffi::Map<Var, BufferVar>& 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<BufferVar>()) {
+ 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<PrimVar>()) {
+ if (auto var = buf.value()->shape[j].as<PrimVar>()) {
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<int>(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<BufferTypeNode>() ? "Tensor" : "pointer";
EmitTypeIndexCheck(param_index,
type_index == ffi::TypeIndex::kTVMFFINone ||
type_index == ffi::TypeIndex::kTVMFFIOpaquePtr ||
@@ -557,6 +554,14 @@ void TVMFFIABIBuilder::DecodeParam(int param_index) {
ffi::reflection::AccessPath param_path =
ffi::reflection::AccessPath::Root()->Extend(AccessStep::ArrayItem(param_index));
+ if (param->ty.as<BufferTypeNode>()) {
+ Var handle(param->name + ".handle", PointerType::VoidPointerTy());
+ Expr handle_value = DecodeParamOpaqueHandle(param_index, type_index.as_or_throw<PrimExpr>());
+ BindPointer(handle, handle_value, param_path, true);
+ buffer_handles_.emplace(param.get(), handle);
+ return;
+ }
+
if (param->ty.as<PointerTypeNode>()) {
Expr handle_value = DecodeParamOpaqueHandle(param_index, type_index.as_or_throw<PrimExpr>());
Expr pointer_value = Call(param->ty, builtin::reinterpret(), {handle_value});
@@ -597,14 +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<BufferVar>()) {
+ 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,
+ ->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 c6b95291c4..7f29a365f5 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<Stmt> init_nest;
/*! \brief Validation checks (all AssertStmts). */
std::vector<Stmt> asserts;
- /*! \brief BufferVar declarations for buffer_map entries. */
+ /*! \brief BufferVar declarations for buffer-typed parameters. */
std::vector<Stmt> 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<Var>& params,
- const ffi::Map<Var, BufferVar>& 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<Stmt> init_nest_;
/*! \brief Validation checks: all AssertStmts. */
std::vector<Stmt> asserts_;
- /*! \brief BufferVar declarations for buffer_map entries. */
+ /*! \brief BufferVar declarations for buffer-typed parameters. */
std::vector<Stmt> decl_buffers_;
/*! \brief Deferred constant-expression assertions for display-var substitution. */
std::vector<PendingConstAssert> pending_const_asserts_;
@@ -405,8 +403,8 @@ class TVMFFIABIBuilder {
std::string func_signature_;
/*! \brief The function parameters. */
ffi::Array<Var> params_;
- /*! \brief The buffer map from parameters to buffers. */
- ffi::Map<Var, BufferVar> buffer_map_;
+ /*! \brief Raw packed-ABI handles decoded for buffer-typed parameters. */
+ std::unordered_map<const VarNode*, Var> 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 1bca19fbc7..c9ba1736b0 100644
--- a/src/tirx/transform/unsupported_dtype_legalize.cc
+++ b/src/tirx/transform/unsupported_dtype_legalize.cc
@@ -557,7 +557,10 @@ 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";
+ for (const Var& param : func->params) {
+ TVM_FFI_ICHECK(!param->ty.as<BufferTypeNode>())
+ << "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 ed7d4a9901..a5deb70587 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 test_domain_touched_vector():
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 ed37b367e2..23752550c9 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 8df1ca822d..b76acc1ee9 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 72bcf8a53d..b9c9b16e1e 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 @@ class ValidateBufferScopes(PyExprVisitor): # pylint: disable=abstract-method
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 @@ class ValidateBufferScopes(PyExprVisitor): # pylint: disable=abstract-method
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 @@ class ValidateBufferScopes(PyExprVisitor): # pylint: disable=abstract-method
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 9e27c22080..acd33f9814 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 test_1d():
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 test_1d_input_2d_output_fused_loop():
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 test_2d_input_1d_output_fused_loop():
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 test_1d_input_1d_output_nested_loop():
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 test_1d_input_1d_output_nested_loop_equivalent_expressions():
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 test_1d_input_2d_output_nested_loop():
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 test_2d_input_1d_output_nested_loop():
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 test_2d_input_2d_output_nested_loop():
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 test_2d_input_2d_output_transpose_both():
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 test_cache_read():
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 test_cache_write():
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 daf9430198..15a84fae30 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 test_buffer_access_with_let_binding():
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 test_buffer_access_with_nested_let_binding():
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 87b578d876..08f72af25d 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 d1cd710ed5..65cb259c98 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,7 @@ class TestPrimFuncs:
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 +219,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 29d879dce2..e2a94759b5 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 @@ from tvm.script import tirx as T
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 e04576e48d..dcb807fa4c 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 5a6a97f618..4252f88f72 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 test_simplify():
# 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 test_transform_layout_with_symbolic_bound():
# 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 test_transform_block_layout_with_symbolic_bound():
# 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 82f0109150..23495d7f42 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 test_symbolic_func():
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,18 @@ def test_buffer_map():
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_compound_buffer_param_shape_var():
+ n = tvm.tirx.Var("n", "int32")
+ A = tvm.tirx.decl_buffer((tvm.tirx.max(n, 1),), layout=None)
+ f1 = tvm.tirx.PrimFunc([A], tvm.tirx.Evaluate(n))
+ f2 = tvm.s_tir.renew_defs(f1)
+
+ tvm.ir.assert_structural_equal(f1, f2)
+ assert not f1.body.value.same_as(f2.body.value)
+ assert f2.params[0].shape[0].a.same_as(f2.body.value)
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 29a7d9cace..45f33a65fd 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 test_sequential_redefinition_with_location():
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 48b92eb485..4af595a783 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 efbe9babd6..f47a6dc591 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 test_specialize_buffer_var_to_var():
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 eb7166a7fb..77217232e6 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 test_reused_buffer_obj():
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")):
@@ -225,6 +225,28 @@ def test_reused_buffer_parameter():
tvm.ir.assert_structural_equal(after, expected)
+def test_reused_compound_buffer_shape_var():
+ """De-duplicate implicit Vars nested in buffer parameter shapes."""
+ n = tirx.Var("n", "int32")
+ A = tirx.decl_buffer((tirx.max(n, 1),), layout=None)
+ func = tirx.PrimFunc([A], tirx.Evaluate(n))
+ before = tvm.IRModule(
+ {
+ "func_a": func.with_attr("global_symbol", "func_a"),
+ "func_b": func.with_attr("global_symbol", "func_b"),
+ }
+ )
+
+ after = tvm.tirx.transform.ConvertSSA()(before)
+ func_a = after["func_a"]
+ func_b = after["func_b"]
+ n_a = func_a.params[0].shape[0].a
+ n_b = func_b.params[0].shape[0].a
+ assert not n_a.same_as(n_b)
+ assert n_a.same_as(func_a.body.value)
+ assert n_b.same_as(func_b.body.value)
+
+
def test_no_change_if_already_ssa():
"""A module that is already SSA should be unchanged"""
@@ -490,32 +512,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 +536,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 8bf1965a64..af5b719417 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 test_i16_buffer():
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 test_fail_on_buffer_map():
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 c11fdf4024..2b15b48471 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 test_decl_buffer_alias_chain_uses_flat_root_map():
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 d3cd32d134..ca9a91846f 100644
--- a/tests/python/tirx/test_jit.py
+++ b/tests/python/tirx/test_jit.py
@@ -247,7 +247,7 @@ def test_constexpr_specializes_nested_selector_condition():
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 test_optional_param_present_and_absent_ir():
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 test_multiple_optional_params_preserve_runtime_order():
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 test_compile_time_if_binding_uses_python_scope():
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 test_compile_time_bool_ops_and_if_expression_short_circuit():
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 8d602cce35..ae7e1aa8fa 100644
--- a/tests/python/tirx/transform/test_stmt_functor.py
+++ b/tests/python/tirx/transform/test_stmt_functor.py
@@ -1204,7 +1204,7 @@ def test_op_call_nested_config_visited_and_substituted():
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 9d56f4bdfd..b4b9736c9b 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 f5ace5999a..dd02640f30 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_parser_tir.py b/tests/python/tvmscript/test_tvmscript_parser_tir.py
index 04d111d250..a4baa9b456 100644
--- a/tests/python/tvmscript/test_tvmscript_parser_tir.py
+++ b/tests/python/tvmscript/test_tvmscript_parser_tir.py
@@ -40,6 +40,142 @@ 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"
+
+
+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",
+ [
+ """
+@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 (
diff --git a/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py b/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py
index f25f53c67c..b82e860c8a 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 test_prim_func_buffer_map():
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 ac0d761667..0cd538fb65 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,57 @@ 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_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_compound_buffer_shape_first_use_roundtrip():
+ n = tirx.Var("n", "int32")
+ A = tirx.decl_buffer(shape=[tirx.max(n, 1)], 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(("T.max(n, 1)",)' 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_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")
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 +118,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 +853,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 +1010,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)