tirx: represent buffer parameters with BufferType

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;

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}
  *

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,19 +113,18 @@ class PrimFunc : public BaseFunc {
    *
    * \param ret_type The return type of the function.
    *
-   * \param buffer_map The buffer map for parameter buffer unpacking.
-   * This contains buffer objects as they appear in the body of the
-   * PrimFunc.  (e.g. a buffer of shape ``[1024]`` originally
-   * generated as a tensor of shape ``[32, 32]``)
-   *
    * \param attrs Additional function attributes.
    *
    * \param span The location of this object in the source code.
    */
   TVM_DLL PrimFunc(ffi::Array<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());
 
+  /*! \brief Compatibility constructor that folds legacy buffer bindings into params. */
+  TVM_DLL PrimFunc(ffi::Array<tirx::Var> params, Stmt body, Type ret_type,
+                   ffi::Map<tirx::Var, BufferVar> buffer_map, DictAttrs attrs = DictAttrs(),
+                   Span span = Span());
+
   TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrimFunc, BaseFunc, PrimFuncNode);
   TVM_DEFINE_OBJECT_REF_COW_METHOD(PrimFuncNode);
 };

python/tvm/backend/cuda/tile_primitive/copy_async/tma.py

@@ -2208,7 +2208,7 @@ def impl():
 
     if selector_bind is not None:
         body = tvm.tirx.SeqStmt([selector_bind, impl.body])
-        impl = PrimFunc([], body, ret_type=None, buffer_map={}).with_attr("global_symbol", "impl")
+        impl = PrimFunc([], body, ret_type=None).with_attr("global_symbol", "impl")
     return impl
 
 

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"
 
 

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):

python/tvm/te/operation.py

@@ -389,9 +389,9 @@ def before_split(a: T.handle, b: T.handle) -> None:
         C = te.extern_primfunc([A, B], func)
     """
 
-    # dt_access_map and primfunc.buffer_map are unordered, so use order from primfunc.params
+    # Preserve the function parameter order while selecting BufferType annotations.
     dt_access_map = tvm.arith._ffi_api.DomainTouchedAccessMap(primfunc)
-    ordered_buffers = [primfunc.buffer_map[param] for param in primfunc.params]
+    ordered_buffers = [param for param in primfunc.params if tvm.tirx.is_buffer_var(param)]
     in_buffers = [buf for buf in ordered_buffers if len(dt_access_map[buf][0])]
     out_buffers = [buf for buf in ordered_buffers if len(dt_access_map[buf][1])]
     assert in_buffers, "PrimFunc has no input buffers"

python/tvm/tirx/function.py

@@ -50,9 +50,6 @@ class PrimFunc(BaseFunc, Scriptable):
     ret_type: tvm.ir.Type
         The return type annotation of the function.
 
-    buffer_map : Map[tvm.tirx.Var, tvm.tirx.Buffer]
-        The buffer binding map.
-
     attrs: Optional[tvm.Attrs]
         Attributes of the function, can be None
 
@@ -73,11 +70,9 @@ def __init__(self, params, body, ret_type=None, buffer_map=None, attrs=None, spa
         for x in params:
             x = tvm.runtime.convert(x) if not isinstance(x, Object) else x
             if is_buffer_var(x):
-                var = Var(x.name, ty="handle")
-                param_list.append(var)
-                buffer_map[var] = x
-            elif isinstance(x, Var):
                 param_list.append(x)
+            elif isinstance(x, Var):
+                param_list.append(buffer_map.get(x, x))
             else:
                 raise TypeError("params can only contain Var or Buffer")
 
@@ -89,7 +84,6 @@ def __init__(self, params, body, ret_type=None, buffer_map=None, attrs=None, spa
             param_list,
             body,
             ret_type,
-            buffer_map,
             attrs,
             span,
         )  # type: ignore
@@ -113,10 +107,9 @@ def with_body(self, new_body, span=None):
         return PrimFunc(
             self.params,
             new_body,
-            self.ret_type,
-            self.buffer_map,
-            self.attrs,
-            span,
+            ret_type=self.ret_type,
+            attrs=self.attrs,
+            span=span,
         )
 
     def specialize(self, param_map: Mapping[Var, Expr | Buffer]):