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 cf29951491 [REFACTOR][TIRx] Emit tensor declarations without a
dedicated frame (#20587)
cf29951491 is described below
commit cf29951491280c1c19917ae35f35a5998b8be382
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 7 11:42:31 2026 -0400
[REFACTOR][TIRx] Emit tensor declarations without a dedicated frame (#20587)
---
include/tvm/tirx/script/ir_builder/frame.h | 35 -----------------------
include/tvm/tirx/script/ir_builder/ir.h | 15 +++++-----
python/tvm/tirx/buffer.py | 6 ++--
python/tvm/tirx/script/ir_builder/frame.py | 7 -----
python/tvm/tirx/script/ir_builder/ir.py | 8 +-----
src/tirx/script/ir_builder/frame.cc | 30 -------------------
src/tirx/script/ir_builder/ir.cc | 46 ++++++++++++++++++------------
7 files changed, 39 insertions(+), 108 deletions(-)
diff --git a/include/tvm/tirx/script/ir_builder/frame.h
b/include/tvm/tirx/script/ir_builder/frame.h
index a1621d1cb7..dd4fba7374 100644
--- a/include/tvm/tirx/script/ir_builder/frame.h
+++ b/include/tvm/tirx/script/ir_builder/frame.h
@@ -461,41 +461,6 @@ class ElseFrame : public TIRFrame {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ElseFrame, TIRFrame,
ElseFrameNode);
};
-class DeclTensorFrameNode : public TIRFrameNode {
- public:
- explicit DeclTensorFrameNode(ffi::UnsafeInit tag) : buffer(tag) {}
-
- explicit DeclTensorFrameNode(tvm::tirx::TensorVar buffer) :
buffer(std::move(buffer)) {}
-
- /*! \brief The declared buffer. */
- tvm::tirx::TensorVar buffer;
- /*! \brief Physical pointer expression backing the declaration. */
- ffi::Optional<Expr> data;
- /*! \brief The buffer allocated or not. */
- bool allocated;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<DeclTensorFrameNode>()
- .def_ro("buffer", &DeclTensorFrameNode::buffer)
- .def_ro("data", &DeclTensorFrameNode::data)
- .def_ro("allocated", &DeclTensorFrameNode::allocated);
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.DeclTensorFrame",
DeclTensorFrameNode,
- TIRFrameNode);
-
- public:
- void ExitWithScope() final;
-};
-
-class DeclTensorFrame : public TIRFrame {
- public:
- explicit DeclTensorFrame(ffi::ObjectPtr<DeclTensorFrameNode> data) :
TIRFrame(data) {
- TVM_FFI_ICHECK(data != nullptr);
- }
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DeclTensorFrame, TIRFrame,
DeclTensorFrameNode);
-};
-
} // namespace tirx
} // namespace ir_builder
} // namespace script
diff --git a/include/tvm/tirx/script/ir_builder/ir.h
b/include/tvm/tirx/script/ir_builder/ir.h
index 9599b60d3c..0263aa709e 100644
--- a/include/tvm/tirx/script/ir_builder/ir.h
+++ b/include/tvm/tirx/script/ir_builder/ir.h
@@ -284,7 +284,7 @@ ThenFrame Then();
ElseFrame Else();
/*!
- * \brief The buffer declaration frame.
+ * \brief Declare a tensor binding or allocate its storage.
* \param shape The type of the buffer prior to flattening.
* \param dtype The data type in the content of the buffer.
* \param buffer_name The name of the buffer.
@@ -295,14 +295,13 @@ ElseFrame Else();
* \param align The alignment requirement of data pointer in bytes.
* \param offset_factor The factor of elem_offset field.
* \param layout The layout of the buffer.
- * \return The declaration frame.
+ * \return The declared tensor.
*/
-DeclTensorFrame DeclTensor(ffi::Array<PrimExpr> shape, PrimType dtype,
ffi::String buffer_name,
- ffi::Optional<Expr> data,
ffi::Optional<ffi::Array<PrimExpr>> strides,
- ffi::Optional<PrimExpr> elem_offset, ffi::String
storage_scope,
- int align, int offset_factor,
- ffi::Optional<Layout> layout = std::nullopt,
- ffi::Optional<PrimExpr> allocated_addr =
std::nullopt);
+TensorVar DeclTensor(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::String
buffer_name,
+ ffi::Optional<Expr> data,
ffi::Optional<ffi::Array<PrimExpr>> strides,
+ ffi::Optional<PrimExpr> elem_offset, ffi::String
storage_scope, int align,
+ int offset_factor, ffi::Optional<Layout> layout =
std::nullopt,
+ ffi::Optional<PrimExpr> allocated_addr = std::nullopt);
/*!
* \brief Statement-level buffer allocation (binds a buffer-returning
allocation Call).
diff --git a/python/tvm/tirx/buffer.py b/python/tvm/tirx/buffer.py
index 1984196595..c738f3ee73 100644
--- a/python/tvm/tirx/buffer.py
+++ b/python/tvm/tirx/buffer.py
@@ -310,7 +310,7 @@ class _TensorMethods:
Returns
-------
- view : DeclTensorFrame
+ view : Var
The corresponding view buffer.
"""
@@ -347,7 +347,7 @@ class _TensorMethods:
Returns
-------
- local : DeclTensorFrame
+ local : Var
The corresponding local buffer.
"""
return _buffer_view.local(self, *shape, layout=layout)
@@ -362,7 +362,7 @@ class _TensorMethods:
Returns
-------
- permuted : DeclTensorFrame
+ permuted : Var
The buffer with permuted dimensions.
"""
return _buffer_view.permute(self, *dims)
diff --git a/python/tvm/tirx/script/ir_builder/frame.py
b/python/tvm/tirx/script/ir_builder/frame.py
index 9ddb9a4d46..cf521ddb60 100644
--- a/python/tvm/tirx/script/ir_builder/frame.py
+++ b/python/tvm/tirx/script/ir_builder/frame.py
@@ -84,13 +84,6 @@ class ThenFrame(TIRFrame): ...
class ElseFrame(TIRFrame): ...
-@_register_object("script.ir_builder.tirx.DeclTensorFrame")
-class DeclTensorFrame(TIRFrame):
- def __enter__(self) -> Var:
- super().__enter__()
- return self.buffer
-
-
@_register_object("script.ir_builder.tirx.RegionFrame")
class RegionFrame(TIRFrame):
"""A result-free region whose body parameters are lexical bindings."""
diff --git a/python/tvm/tirx/script/ir_builder/ir.py
b/python/tvm/tirx/script/ir_builder/ir.py
index 16ef9c9c9d..b359e7354a 100644
--- a/python/tvm/tirx/script/ir_builder/ir.py
+++ b/python/tvm/tirx/script/ir_builder/ir.py
@@ -21,7 +21,6 @@ import functools
import inspect
import threading
from collections.abc import Callable
-from functools import partial
from numbers import Integral
from typing import TYPE_CHECKING, Any, TypeVar
@@ -731,7 +730,7 @@ def decl_tensor(
if strides is None:
strides = []
dtype = _normalize_prim_type(dtype)
- decl_frame = _ffi_api.DeclTensor( # type: ignore[attr-defined] # pylint:
disable=no-member
+ buf = _ffi_api.DeclTensor( # type: ignore[attr-defined] # pylint:
disable=no-member
shape,
dtype,
"",
@@ -744,11 +743,6 @@ def decl_tensor(
_get_layout(layout, shape, scope),
allocated_addr,
)
- if isinstance(decl_frame, frame.DeclTensorFrame):
- decl_frame.add_callback(partial(decl_frame.__exit__, None, None, None))
- buf = decl_frame.__enter__()
- else:
- buf = decl_frame
_record_meta_resource(buf, skip_frames=2)
return buf
diff --git a/src/tirx/script/ir_builder/frame.cc
b/src/tirx/script/ir_builder/frame.cc
index 5939eaff2d..c532d10e6a 100644
--- a/src/tirx/script/ir_builder/frame.cc
+++ b/src/tirx/script/ir_builder/frame.cc
@@ -47,7 +47,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
IfFrameNode::RegisterReflection();
ThenFrameNode::RegisterReflection();
ElseFrameNode::RegisterReflection();
- DeclTensorFrameNode::RegisterReflection();
}
namespace {
@@ -272,35 +271,6 @@ void ElseFrameNode::ExitWithScope() {
FindIfFrame("T.else_")->else_stmts = stmts;
}
-void DeclTensorFrameNode::ExitWithScope() {
- TIRFrameNode::ExitWithScope();
- if (allocated) {
- TVM_FFI_ICHECK(data.has_value());
- AddToParent(tvm::tirx::SeqStmt::Flatten(
- tvm::tirx::Bind(buffer,
- tvm::Call(buffer.type(),
tvm::tirx::builtin::decl_tensor(),
- {data.value(),
tvm::Tuple(buffer->shape),
-
tvm::DataTypeImm(buffer->dtype->dtype),
- tvm::StringImm(buffer.scope())},
- {}, {}, source_span),
- source_span),
- AsStmt(stmts)),
- source_span);
- } else {
- // data is undefined in `decl_tensor(...)`, lower to `alloc_tensor(...)`.
- AddToParent(
- tvm::tirx::SeqStmt::Flatten(
- tvm::tirx::Bind(buffer.var(),
- Call(buffer.type(),
tvm::tirx::builtin::alloc_tensor(),
- {tvm::Tuple(buffer->shape),
DataTypeImm(buffer->dtype->dtype),
- StringImm(buffer.scope())},
- DictAttrs(), {}, source_span),
- source_span),
- AsStmt(stmts)),
- source_span);
- }
-}
-
} // namespace tirx
} // namespace ir_builder
} // namespace script
diff --git a/src/tirx/script/ir_builder/ir.cc b/src/tirx/script/ir_builder/ir.cc
index d3becf2a3c..e386b70926 100644
--- a/src/tirx/script/ir_builder/ir.cc
+++ b/src/tirx/script/ir_builder/ir.cc
@@ -534,11 +534,11 @@ tvm::tirx::Stmt BufferStore(TensorVar buffer, PrimExpr
value, ffi::Array<PrimExp
return store;
}
-DeclTensorFrame DeclTensor(ffi::Array<PrimExpr> shape, PrimType dtype,
ffi::String buffer_name,
- ffi::Optional<Expr> data,
ffi::Optional<ffi::Array<PrimExpr>> strides,
- ffi::Optional<PrimExpr> elem_offset, ffi::String
storage_scope,
- int align, int offset_factor, ffi::Optional<Layout>
layout,
- ffi::Optional<PrimExpr> allocated_addr) {
+TensorVar DeclTensor(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::String
buffer_name,
+ ffi::Optional<Expr> data,
ffi::Optional<ffi::Array<PrimExpr>> strides,
+ ffi::Optional<PrimExpr> elem_offset, ffi::String
storage_scope, int align,
+ int offset_factor, ffi::Optional<Layout> layout,
+ ffi::Optional<PrimExpr> allocated_addr) {
std::string scope = static_cast<std::string>(storage_scope);
if (scope.empty()) {
scope = "global";
@@ -566,21 +566,31 @@ DeclTensorFrame DeclTensor(ffi::Array<PrimExpr> shape,
PrimType dtype, ffi::Stri
}
}
- ffi::ObjectPtr<DeclTensorFrameNode> n =
ffi::make_object<DeclTensorFrameNode>(
- TensorDecl(shape, dtype, buffer_name, data, strides, elem_offset,
storage_scope, align,
- offset_factor, layout, allocated_addr_arr));
+ TensorVar buffer = TensorDecl(shape, dtype, buffer_name, data, strides,
elem_offset,
+ storage_scope, align, offset_factor, layout,
allocated_addr_arr);
+ if (scope == "tmem") {
+ // Tensor memory is externally allocated; make its address-to-pointer
binding explicit.
+ data =
+ Call(buffer.DataPointerType(), tvm::tirx::builtin::reinterpret(),
{allocated_addr.value()});
+ }
+ Span span = IRBuilder::Current()->GetCurrentSourceSpan();
if (data.has_value()) {
- n->data = data.value();
- } else if (scope == "tmem") {
- // Tensor memory is an externally allocated address space. Make that
- // address-to-pointer relationship explicit so every DeclTensor has a
- // physical data binding.
- n->data = Call(n->buffer.DataPointerType(),
tvm::tirx::builtin::reinterpret(),
- {allocated_addr.value()});
+ AddToParent(tvm::tirx::Bind(buffer.var(),
+ Call(buffer.type(),
tvm::tirx::builtin::decl_tensor(),
+ {data.value(), tvm::Tuple(buffer->shape),
+ DataTypeImm(buffer->dtype->dtype),
StringImm(buffer.scope())},
+ {}, {}, span),
+ span));
+ } else {
+ // Without a backing pointer, declare and allocate the tensor together.
+ AddToParent(tvm::tirx::Bind(buffer.var(),
+ Call(buffer.type(),
tvm::tirx::builtin::alloc_tensor(),
+ {tvm::Tuple(buffer->shape),
DataTypeImm(buffer->dtype->dtype),
+ StringImm(buffer.scope())},
+ DictAttrs(), {}, span),
+ span));
}
- // For tmem, even without `data`, we should not emit an Allocate node.
- n->allocated = (scope == "tmem") || data.has_value();
- return DeclTensorFrame(n);
+ return buffer;
}
TensorVar AllocTensor(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::String
storage_scope,
