[REFACTOR][TIRx] Represent buffer definitions as operation bindings (#20482)

Represent buffer allocation and declaration as ordinary bindings of
buffer-returning operation calls. Their shape, data type, and storage
scope are explicit operands; declaration also carries its existing
pointer operand. Allocation annotations live in call attributes. Update
traversal, lowering, code generation, and script construction and
printing to consume the shared representation.

Author: tqchen

include/tvm/relax/distributed/axis_group_graph.h

@@ -74,7 +74,7 @@ class BufferAxisGraphExtractor : public s_tir::StmtExprVisitor {
     ffi::Map<BufferVar, Var> inverse_buffer_map;
     for (const Var& param : prim_func->params) {
       if (param->ty.as<BufferTypeNode>()) {
-        inverse_buffer_map.Set(BufferVar(param), param);
+        inverse_buffer_map.Set(param.as_or_throw<BufferVar>(), param);
       }
     }
     std::vector<std::vector<TIRVarAxis>> tir_var_axis_group_list;
@@ -83,7 +83,7 @@ class BufferAxisGraphExtractor : public s_tir::StmtExprVisitor {
       if (!param->ty.as<BufferTypeNode>()) {
         continue;
       }
-      BufferVar buffer(param);
+      BufferVar buffer = param.as_or_throw<BufferVar>();
       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/tirx/builtin.h

@@ -41,6 +41,53 @@ namespace tirx {
 
 /*! \brief Collection of builtin intrinsics as ops */
 namespace builtin {
+/*!
+ * \brief Allocate a buffer: alloc_buffer(shape, dtype, scope) -> BufferType.
+ *
+ * Arguments, in order:
+ * - args[0]: shape, Tuple of integer extents (IntImm or symbolic integer expressions).
+ * - args[1]: dtype, DataTypeImm with a DLDataType payload for the element type.
+ * - args[2]: scope, StringImm naming the storage scope.
+ *
+ * DictAttrs directly holds the allocation annotations, defaulting to an empty dictionary.
+ * The BufferType result agrees with the operands and retains buffer access/storage metadata.
+ *
+ * \code
+ * // Example pattern match code for a given Binding:
+ * if (const auto* call = binding->value.as<CallNode>();
+ *     call && call->op.same_as(builtin::alloc_buffer())) {
+ *   tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
+ *   DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
+ *   ffi::String scope = call->args[2].as_or_throw<StringImm>()->value;
+ *   DictAttrs annotations = call->attrs.as_or_throw<DictAttrs>();
+ * }
+ * \endcode
+ */
+TVM_DLL const Op& alloc_buffer();
+/*!
+ * \brief Declare a buffer view: decl_buffer(data, shape, dtype, scope) -> BufferType.
+ *
+ * Arguments, in order:
+ * - args[0]: data, Expr for the existing physical pointer backing the buffer view.
+ * - args[1]: shape, Tuple of integer extents (IntImm or symbolic integer expressions).
+ * - args[2]: dtype, DataTypeImm with a DLDataType payload for the element type.
+ * - args[3]: scope, StringImm naming the storage scope.
+ *
+ * There are no attributes. The BufferType result agrees with the operands and retains
+ * buffer access/storage metadata. The operation binds a view without allocating memory.
+ *
+ * \code
+ * // Example pattern match code for a given Binding:
+ * if (const auto* call = binding->value.as<CallNode>();
+ *     call && call->op.same_as(builtin::decl_buffer())) {
+ *   Expr data = call->args[0];
+ *   tvm::Tuple shape = call->args[1].as_or_throw<tvm::Tuple>();
+ *   DLDataType dtype = call->args[2].as_or_throw<DataTypeImm>()->value;
+ *   ffi::String scope = call->args[3].as_or_throw<StringImm>()->value;
+ * }
+ * \endcode
+ */
+TVM_DLL const Op& decl_buffer();
 /*!
  * \brief Return from a GPU thread without returning a function value.
  */

include/tvm/tirx/expr.h

@@ -117,7 +117,7 @@ class BufferVar : public Var {
    *
    * If flattening changes the type, the result is a fresh BufferVar.  Callers
    * that use it as a view over this buffer must bind the returned variable with
-   * `DeclBuffer(flattened, this->data())`.
+   * a `Bind` of `flattened` to a `decl_buffer` Call over `this->data()`.
    */
   BufferVar GetFlattenedBuffer() const;
 

include/tvm/tirx/script/ir_builder/ir.h

@@ -325,7 +325,7 @@ DeclBufferFrame DeclBuffer(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::Stri
                            ffi::Optional<PrimExpr> allocated_addr = std::nullopt);
 
 /*!
- * \brief Statement-level buffer allocation (creates an AllocBuffer IR node).
+ * \brief Statement-level buffer allocation (binds a buffer-returning allocation Call).
  * \param shape The shape of the buffer to allocate.
  * \param dtype The data type of buffer elements.
  * \param storage_scope The storage scope (e.g., "global", "shared").

include/tvm/tirx/stmt.h

@@ -231,82 +231,6 @@ class BufferStore : public Stmt {
   TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferStoreNode);
 };
 
-/*! \brief Declare a buffer that can be used in the body */
-class DeclBufferNode : public StmtNode {
- public:
-  /*! \brief The buffer being declared */
-  BufferVar buffer;
-  /*! \brief Physical pointer expression backing the declaration. */
-  Expr data;
-
-  static void RegisterReflection() {
-    namespace refl = tvm::ffi::reflection;
-    refl::ObjectDef<DeclBufferNode>()
-        .def_ro("buffer", &DeclBufferNode::buffer, refl::AttachFieldFlag::SEqHashDefSimple())
-        .def_ro("data", &DeclBufferNode::data);
-  }
-  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.DeclBuffer", DeclBufferNode, StmtNode);
-};
-
-/*! \brief Managed reference to DeclBufferNode */
-class DeclBuffer : public Stmt {
- public:
-  TVM_DLL DeclBuffer(BufferVar buffer, Expr data, Span span = Span());
-  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DeclBuffer, Stmt, DeclBufferNode);
-  TVM_DEFINE_OBJECT_REF_COW_METHOD(DeclBufferNode);
-};
-
-/*! \brief Allocate a buffer and declare it in scope */
-class AllocBufferNode : public StmtNode {
- public:
-  /*! \brief The buffer being allocated and declared */
-  BufferVar buffer;
-  /*!
-   * \brief Additional annotations about the allocation.
-   *
-   *  These annotations can be used as auxiliary hint
-   *  to future transformations.
-   */
-  ffi::Map<ffi::String, ffi::Any> annotations;
-
-  static void RegisterReflection() {
-    namespace refl = tvm::ffi::reflection;
-    refl::ObjectDef<AllocBufferNode>()
-        .def_ro("buffer", &AllocBufferNode::buffer, refl::AttachFieldFlag::SEqHashDefSimple())
-        .def_ro("annotations", &AllocBufferNode::annotations);
-  }
-  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.AllocBuffer", AllocBufferNode, StmtNode);
-};
-
-/*! \brief Managed reference to AllocBufferNode */
-class AllocBuffer : public Stmt {
- public:
-  TVM_DLL AllocBuffer(
-      BufferVar buffer,
-      ffi::Map<ffi::String, ffi::Any> annotations = ffi::Map<ffi::String, ffi::Any>(),
-      Span span = Span());
-  /*!
-   * \brief If the buffer's shape is constant, return the total number of elements.
-   * \return The product of all shape extents if all are constant, std::nullopt otherwise.
-   */
-  std::optional<int64_t> ConstantAllocationSize() const {
-    int64_t result = 1;
-    for (const PrimExpr& extent : (*this)->buffer->shape) {
-      if (const auto* int_size = extent.as<IntImmNode>()) {
-        auto product = (result * int_size->value).as<int64_t>();
-        if (!product.has_value()) return std::nullopt;
-        result = *product;
-      } else {
-        return std::nullopt;
-      }
-    }
-    return result;
-  }
-
-  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(AllocBuffer, Stmt, AllocBufferNode);
-  TVM_DEFINE_OBJECT_REF_COW_METHOD(AllocBufferNode);
-};
-
 /*!
  * \brief The container of seq statement.
  *        Represent a sequence of statements.

include/tvm/tirx/stmt_functor.h

@@ -98,12 +98,6 @@ class StmtFunctor<R(const Stmt&, Args...)> {
   virtual R Dispatch_(const ContinueNode* node, Args... args) {
     return DispatchDefault_(node, std::forward<Args>(args)...);
   }
-  virtual R Dispatch_(const AllocBufferNode* node, Args... args) {
-    return DispatchDefault_(node, std::forward<Args>(args)...);
-  }
-  virtual R Dispatch_(const DeclBufferNode* node, Args... args) {
-    return DispatchDefault_(node, std::forward<Args>(args)...);
-  }
   virtual R Dispatch_(const BufferStoreNode* node, Args... args) {
     return DispatchDefault_(node, std::forward<Args>(args)...);
   }
@@ -147,8 +141,6 @@ class StmtFunctor<R(const Stmt&, Args...)> {
     SetDispatch<TSelf, ReturnNode>(vtable);
     SetDispatch<TSelf, BreakNode>(vtable);
     SetDispatch<TSelf, ContinueNode>(vtable);
-    SetDispatch<TSelf, AllocBufferNode>(vtable);
-    SetDispatch<TSelf, DeclBufferNode>(vtable);
     SetDispatch<TSelf, BufferStoreNode>(vtable);
     SetDispatch<TSelf, AssertStmtNode>(vtable);
     SetDispatch<TSelf, SeqStmtNode>(vtable);
@@ -206,8 +198,6 @@ class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor {
   virtual ffi::Optional<VisitInterrupt> Visit_(const ReturnNode* op);
   virtual ffi::Optional<VisitInterrupt> Visit_(const BreakNode* op);
   virtual ffi::Optional<VisitInterrupt> Visit_(const ContinueNode* op);
-  virtual ffi::Optional<VisitInterrupt> Visit_(const AllocBufferNode* op);
-  virtual ffi::Optional<VisitInterrupt> Visit_(const DeclBufferNode* op);
   virtual ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op);
   virtual ffi::Optional<VisitInterrupt> Visit_(const AssertStmtNode* op);
   virtual ffi::Optional<VisitInterrupt> Visit_(const SeqStmtNode* op);
@@ -227,9 +217,6 @@ class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor {
   ffi::Optional<VisitInterrupt> Visit_(const prim::BroadcastNode* op) override;
   ffi::Optional<VisitInterrupt> Visit_(const prim::ShuffleNode* op) override;
 
-  /*! \brief Visit definition metadata as uses, separately from the buffer Var definition. */
-  ffi::Optional<VisitInterrupt> VisitBufferMetadata(const BufferVar& buffer);
-
  protected:
   explicit StmtExprVisitor(const VTable* vtable) : tvm::ExprVisitor(vtable) {}
   static void InitVTable(VTable* vtable);
@@ -270,8 +257,6 @@ class TVM_DLL StmtExprMutator : public tvm::ExprMutator {
   virtual UnchangedOr<Stmt> Mutate_(const ReturnNode* op, InplaceMode inplace_mode);
   virtual UnchangedOr<Stmt> Mutate_(const BreakNode* op, InplaceMode inplace_mode);
   virtual UnchangedOr<Stmt> Mutate_(const ContinueNode* op, InplaceMode inplace_mode);
-  virtual UnchangedOr<Stmt> Mutate_(const AllocBufferNode* op, InplaceMode inplace_mode);
-  virtual UnchangedOr<Stmt> Mutate_(const DeclBufferNode* op, InplaceMode inplace_mode);
   virtual UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode);
   virtual UnchangedOr<Stmt> Mutate_(const AssertStmtNode* op, InplaceMode inplace_mode);
   virtual UnchangedOr<Stmt> Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode);

include/tvm/tirx/type.h

@@ -28,6 +28,8 @@
 #include <tvm/ir/type.h>
 #include <tvm/tirx/layout.h>
 
+#include <optional>
+
 namespace tvm::tirx {
 
 #ifndef TVM_INDEX_DEFAULT_I64
@@ -127,6 +129,12 @@ class BufferTypeNode : public TypeNode {
   /*! \return type of the physical pointer projected by buffer_data. */
   PointerType DataPointerType() const { return PointerType(dtype, storage_scope); }
 
+  /*! \brief Whether this type supports scalar buffer syntax. */
+  TVM_DLL bool IsScalar(bool alloc_or_decl = true) const;
+
+  /*! \brief Return the constant element count, or nullopt for symbolic extents or overflow. */
+  TVM_DLL std::optional<int64_t> ConstantAllocationSize() const;
+
   /*! \brief Determine the offset in the buffer of the given index.
    *
    * Returns the buffer offset, in number of elements of type dtype,

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

@@ -109,14 +109,15 @@
 import operator
 
 import tvm
+from tvm.ir import Call, DataTypeImm, DictAttrs, StringImm, Tuple
 from tvm.runtime import DataType
 from tvm.script import tirx as T
 from tvm.sym import Analyzer
 from tvm.tirx import Buffer, PrimFunc
 from tvm.tirx.layout import ComposeLayout, TCol, TileLayout, TLane
 from tvm.tirx.layout import m as m_axis
 from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
-from tvm.tirx.stmt import AllocBuffer, Evaluate, SeqStmt
+from tvm.tirx.stmt import Bind, Evaluate, SeqStmt
 from tvm.tirx.tile_primitive import TilePrimitiveCall
 
 from ..copy import _single_thread_exec
@@ -679,7 +680,24 @@ def _get_or_create_desc(sctx, s_buf, ldo, sdo, swizzle):
     encode_call = T.cuda.tcgen05.encode_matrix_descriptor(
         desc_buf.data, T.reinterpret("handle", T.uint64(0)), ldo, sdo, swizzle
     )
-    wrap = SeqStmt([AllocBuffer(desc_buf), Evaluate(encode_call)])
+    wrap = SeqStmt(
+        [
+            Bind(
+                desc_buf,
+                Call(
+                    "tirx.alloc_buffer",
+                    [
+                        Tuple(desc_buf.ty.shape),
+                        DataTypeImm(desc_buf.ty.dtype.dtype),
+                        StringImm(desc_buf.scope()),
+                    ],
+                    attrs=DictAttrs({}),
+                    ret_ty=desc_buf.ty,
+                ),
+            ),
+            Evaluate(encode_call),
+        ]
+    )
     sctx.add_post_buffer_def_stmt(s_buf, wrap)
     sctx.cache_set(cache_key, desc_buf)
     return desc_buf

python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py

@@ -26,7 +26,7 @@
 import operator
 
 import tvm
-from tvm.ir import TensorRegion
+from tvm.ir import Call, DataTypeImm, DictAttrs, StringImm, TensorRegion, Tuple
 from tvm.runtime import DataType
 from tvm.script import tirx as T
 from tvm.sym.analyzer import Analyzer
@@ -44,7 +44,7 @@
     tmem_mma_operand_layout,
 )
 from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
-from tvm.tirx.stmt import AllocBuffer, Evaluate, SeqStmt
+from tvm.tirx.stmt import Bind, Evaluate, SeqStmt
 from tvm.tirx.tile_primitive import TilePrimitiveCall
 
 from ...cpp.descriptors import (
@@ -1089,8 +1089,32 @@ def _make_lo_uniform(desc_buf):
         pack = T.ptx.mov.b64(desc_buf[0], desc_lo[0], desc_hi[0])
         return SeqStmt(
             [
-                AllocBuffer(desc_lo),
-                AllocBuffer(desc_hi),
+                Bind(
+                    desc_lo,
+                    Call(
+                        "tirx.alloc_buffer",
+                        [
+                            Tuple(desc_lo.ty.shape),
+                            DataTypeImm(desc_lo.ty.dtype.dtype),
+                            StringImm(desc_lo.scope()),
+                        ],
+                        attrs=DictAttrs({}),
+                        ret_ty=desc_lo.ty,
+                    ),
+                ),
+                Bind(
+                    desc_hi,
+                    Call(
+                        "tirx.alloc_buffer",
+                        [
+                            Tuple(desc_hi.ty.shape),
+                            DataTypeImm(desc_hi.ty.dtype.dtype),
+                            StringImm(desc_hi.scope()),
+                        ],
+                        attrs=DictAttrs({}),
+                        ret_ty=desc_hi.ty,
+                    ),
+                ),
                 Evaluate(unpack),
                 Evaluate(shuffle),
                 Evaluate(pack),
@@ -1112,7 +1136,22 @@ def _make_desc(smem_buf, ldo, sdo, swizzle_val, name):
             sdo,
             swizzle_val,
         )
-        wrap_stmts = [AllocBuffer(desc_buf), Evaluate(encode_call)]
+        wrap_stmts = [
+            Bind(
+                desc_buf,
+                Call(
+                    "tirx.alloc_buffer",
+                    [
+                        Tuple(desc_buf.ty.shape),
+                        DataTypeImm(desc_buf.ty.dtype.dtype),
+                        StringImm(desc_buf.scope()),
+                    ],
+                    attrs=DictAttrs({}),
+                    ret_ty=desc_buf.ty,
+                ),
+            ),
+            Evaluate(encode_call),
+        ]
         if warp_scope:
             wrap_stmts.append(_make_lo_uniform(desc_buf))
         wrap_stmts.append(_krp)