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 1cb00e0651 [REFACTOR][IR] Promote TensorRegionType to shared IR
(#20600)
1cb00e0651 is described below
commit 1cb00e0651afd297b269aaa5a94cf7a43da5e131
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Oct 8 11:16:50 2026 -0400
[REFACTOR][IR] Promote TensorRegionType to shared IR (#20600)
Move the tensor-region type from TIRx into shared IR as
`TensorRegionType`, alongside `TensorRegion`. Update C++ and Python
APIs, reflection keys, constructors and script printing while retaining
TIRx-owned buffer slicing.
Preserve legacy JSON loading through `tvm.ir.load_json` by translating
the former `BufferRegionType` key.
---
docs/reference/api/python/script/ir_builder.rst | 2 +-
docs/tirx/api/tirx.rst | 5 ++--
include/tvm/ir/type.h | 19 ++++++++++++++
include/tvm/tirx/expr.h | 2 +-
include/tvm/tirx/type.h | 19 --------------
python/tvm/ir/__init__.py | 1 +
python/tvm/ir/json_compact.py | 11 ++++----
python/tvm/ir/type.py | 8 ++++++
python/tvm/script/ir_builder/__init__.py | 2 ++
python/tvm/tirx/__init__.py | 2 +-
python/tvm/tirx/stmt.py | 10 +------
src/ir/type.cc | 16 +++++++++++
src/script/printer/ir/prim_type.cc | 10 +++++++
src/script/printer/script_printer.cc | 1 +
src/tirx/ir/buffer.cc | 4 +--
src/tirx/ir/type.cc | 35 -------------------------
src/tirx/script/printer/buffer.cc | 9 -------
17 files changed, 72 insertions(+), 84 deletions(-)
diff --git a/docs/reference/api/python/script/ir_builder.rst
b/docs/reference/api/python/script/ir_builder.rst
index 805f69c695..58d29f9ddd 100644
--- a/docs/reference/api/python/script/ir_builder.rst
+++ b/docs/reference/api/python/script/ir_builder.rst
@@ -23,7 +23,7 @@ tvm.script.ir_builder
.. automodule:: tvm.script.ir_builder
:members:
:imported-members:
- :exclude-members: Call, DataTypeImm, FuncType, GenericConst, MissingType,
Op, PrimType, Range, StringImm, StringType, TensorRegion, Tuple, Type
+ :exclude-members: Call, DataTypeImm, FuncType, GenericConst, MissingType,
Op, PrimType, Range, StringImm, StringType, TensorRegion, TensorRegionType,
Tuple, Type
tvm.relax.script.ir_builder.distributed
***************************************
diff --git a/docs/tirx/api/tirx.rst b/docs/tirx/api/tirx.rst
index 8edffc8496..d427c660c9 100644
--- a/docs/tirx/api/tirx.rst
+++ b/docs/tirx/api/tirx.rst
@@ -25,8 +25,9 @@ here so the same objects are not expanded twice.
For C++ construction, include ``tvm/tirx/expr.h`` for ``TensorVar``, buffer
loads, and buffer-region constructors. Include ``tvm/tirx/type.h`` for
-``TensorType``, ``BufferRegionType``, and ``TensorMapType``. Buffer regions
-use the shared ``TensorRegion`` expression from ``tvm/ir/expr.h``.
+``TensorType`` and ``TensorMapType``. Buffer regions use the shared
+``TensorRegion`` expression from ``tvm/ir/expr.h`` and ``TensorRegionType``
+from ``tvm/ir/type.h``.
.. automodule:: tvm.tirx
:members:
diff --git a/include/tvm/ir/type.h b/include/tvm/ir/type.h
index 21132604c4..e3821058ed 100644
--- a/include/tvm/ir/type.h
+++ b/include/tvm/ir/type.h
@@ -57,6 +57,25 @@ class AnyType : public Type {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AnyType, Type, AnyTypeNode);
};
+/*! \brief The type of a multi-dimensional tensor region expression. */
+class TensorRegionTypeNode : public TypeNode {
+ public:
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<TensorRegionTypeNode>();
+ }
+
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.TensorRegionType",
TensorRegionTypeNode, TypeNode);
+};
+
+/*! \brief Managed reference to TensorRegionTypeNode. */
+class TensorRegionType : public Type {
+ public:
+ TVM_DLL TensorRegionType();
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorRegionType, Type,
TensorRegionTypeNode);
+};
+
/*! \brief Semantic string type; its physical representation is context
dependent. */
class StringTypeNode : public TypeNode {
public:
diff --git a/include/tvm/tirx/expr.h b/include/tvm/tirx/expr.h
index e5c5a3ab22..6d8ce878d0 100644
--- a/include/tvm/tirx/expr.h
+++ b/include/tvm/tirx/expr.h
@@ -242,7 +242,7 @@ TVM_DLL tirx::TensorVar
TensorWithOffsetAlignment(ffi::Array<PrimExpr> shape, Pr
TVM_DLL TensorLoad MakeTensorLoad(TensorVar buffer, ffi::Array<PrimExpr>
indices,
Span span = Span());
-/*! \brief Construct a region with buffer rank validation and
BufferRegionType. */
+/*! \brief Construct a region with buffer rank validation and
TensorRegionType. */
TVM_DLL TensorRegion BufferRegion(TensorVar buffer, ffi::Array<Range> region,
Span span = Span());
/*! \brief Select the entire buffer. */
TVM_DLL TensorRegion FullBufferRegion(TensorVar buffer);
diff --git a/include/tvm/tirx/type.h b/include/tvm/tirx/type.h
index c2208d5f2a..a84dc6c21c 100644
--- a/include/tvm/tirx/type.h
+++ b/include/tvm/tirx/type.h
@@ -169,25 +169,6 @@ class TensorType : public Type {
}
};
-/*! \brief The type of a multi-dimensional buffer region expression. */
-class BufferRegionTypeNode : public TypeNode {
- public:
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<BufferRegionTypeNode>();
- }
-
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferRegionType",
BufferRegionTypeNode, TypeNode);
-};
-
-/*! \brief Managed reference to BufferRegionTypeNode. */
-class BufferRegionType : public Type {
- public:
- TVM_DLL BufferRegionType();
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BufferRegionType, Type,
BufferRegionTypeNode);
-};
-
/*!
* \brief The type of tensor map.
* \sa TensorMapType
diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py
index 7151ef6f1e..345b294468 100644
--- a/python/tvm/ir/__init__.py
+++ b/python/tvm/ir/__init__.py
@@ -41,6 +41,7 @@ from .type import (
PointerType,
PrimType,
StringType,
+ TensorRegionType,
TupleType,
Type,
)
diff --git a/python/tvm/ir/json_compact.py b/python/tvm/ir/json_compact.py
index 085dc52ec6..34107c469e 100644
--- a/python/tvm/ir/json_compact.py
+++ b/python/tvm/ir/json_compact.py
@@ -34,6 +34,7 @@ _PRIM_TYPE_KEY_RENAMES = {
"arith.SplitExpr": "sym.SplitExpr",
"arith.SumExpr": "sym.SumExpr",
"tirx.BufferRegion": "ir.TensorRegion",
+ "tirx.BufferRegionType": "ir.TensorRegionType",
"tirx.IterVar": "s_tir.IterVar",
"tirx.SBlock": "s_tir.SBlock",
"tirx.SBlockRealize": "s_tir.SBlockRealize",
@@ -139,7 +140,7 @@ def upgrade_json(json_str):
# written before the canonical Var field was renamed to `name`. Rewriting
# nodes in place preserves node indices and shared references.
nodes = data.get("nodes", [])
- buffer_region_type = None
+ tensor_region_type = None
for node in nodes:
if node.get("type") == "tirx.BufferRegion":
fields = node.get("data")
@@ -150,10 +151,10 @@ def upgrade_json(json_str):
# an Expr, it had only buffer/region; supply that form's defaults
# by appending a type node so existing graph indices stay intact.
if "ty" not in fields:
- if buffer_region_type is None:
- buffer_region_type = len(nodes)
- nodes.append({"type": "tirx.BufferRegionType", "data":
{"span": 0}})
- fields["ty"] = buffer_region_type
+ if tensor_region_type is None:
+ tensor_region_type = len(nodes)
+ nodes.append({"type": "ir.TensorRegionType", "data":
{"span": 0}})
+ fields["ty"] = tensor_region_type
fields.setdefault("span", 0)
node["type"] = _PRIM_TYPE_KEY_RENAMES.get(node.get("type"),
node.get("type"))
if node.get("type") == "relax.expr.Var":
diff --git a/python/tvm/ir/type.py b/python/tvm/ir/type.py
index 54172b637f..dfaf28623e 100644
--- a/python/tvm/ir/type.py
+++ b/python/tvm/ir/type.py
@@ -70,6 +70,14 @@ class AnyType(Type):
self.__init_handle_by_constructor__(_ffi_api.AnyType, span)
+@tvm_ffi.register_object("ir.TensorRegionType")
+class TensorRegionType(Type):
+ """The type of a multi-dimensional :class:`tvm.ir.TensorRegion`."""
+
+ def __init__(self):
+ self.__init_handle_by_constructor__(_ffi_api.TensorRegionType)
+
+
@tvm_ffi.register_object("ir.OpaqueType")
class OpaqueType(Type):
"""Type marker for opaque values that must be removed from finished IR."""
diff --git a/python/tvm/script/ir_builder/__init__.py
b/python/tvm/script/ir_builder/__init__.py
index 1a6eed82fe..172dcc08b4 100644
--- a/python/tvm/script/ir_builder/__init__.py
+++ b/python/tvm/script/ir_builder/__init__.py
@@ -28,6 +28,7 @@ from tvm.ir import (
StringImm,
StringType,
TensorRegion,
+ TensorRegionType,
Tuple,
Type,
make_node,
@@ -80,6 +81,7 @@ __all__ = [
"StringImm",
"StringType",
"TensorRegion",
+ "TensorRegionType",
"Tuple",
"Type",
"at_",
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index 217fee603d..6e7eda8b05 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -52,7 +52,7 @@ from .stmt import TensorStore, RegionStmt
from .stmt import SeqStmt
from .stmt import IfThenElse, Evaluate, stmt_seq, stmt_list
-from .stmt import BufferRegion, BufferRegionType
+from .stmt import BufferRegion
from .stmt import ScopeIdDefStmt
from .tile_primitive import DispatchContext, TilePrimitiveCall
diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py
index 86792486d9..48c0e03d79 100644
--- a/python/tvm/tirx/stmt.py
+++ b/python/tvm/tirx/stmt.py
@@ -33,7 +33,7 @@ from typing import Any
import tvm_ffi
-from tvm.ir import DictAttrs, Expr, Op, Range, Span, StringImm, TensorRegion,
Type, Var, make_node
+from tvm.ir import DictAttrs, Expr, Op, Range, Span, StringImm, TensorRegion,
Var, make_node
from tvm.runtime import Object, Scriptable
from . import _ffi_api
@@ -409,14 +409,6 @@ class Evaluate(Stmt):
self.__init_handle_by_constructor__(_ffi_api.Evaluate, value, span) #
type: ignore
-@tvm_ffi.register_object("tirx.BufferRegionType")
-class BufferRegionType(Type):
- """The TIRX subscript type of a buffer-backed
:class:`tvm.ir.TensorRegion`."""
-
- def __init__(self) -> None:
- self.__init_handle_by_constructor__(_ffi_api.BufferRegionType) #
type: ignore
-
-
def BufferRegion(buffer: Var, region: list[Range]) -> TensorRegion:
"""Construct a buffer-backed tensor region with TIRX subscript semantics.
diff --git a/src/ir/type.cc b/src/ir/type.cc
index dfa2601f07..13e9743694 100644
--- a/src/ir/type.cc
+++ b/src/ir/type.cc
@@ -292,6 +292,22 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("ir.AnyType", [](Span span) { return AnyType(span); });
}
+TensorRegionType::TensorRegionType() : Type(ffi::UnsafeInit{}) {
+ static const auto singleton = ffi::make_object<TensorRegionTypeNode>();
+ data_ = singleton;
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ TensorRegionTypeNode::RegisterReflection();
+ refl::TypeAttrDef<TensorRegionTypeNode>()
+ .attr(refl::type_attr::kStructuralVisit,
ffi::FStructuralVisit::FromNative<&TypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
ffi::FStructuralMutate::FromNative<&TypeMutate>())
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ ffi::FStructuralMutate::FromNative<&TypeMaybeInplaceMutate>());
+ refl::GlobalDef().def("ir.TensorRegionType", []() { return
TensorRegionType(); });
+}
+
OpaqueType::OpaqueType() : Type(ffi::UnsafeInit{}) { data_ =
ffi::make_object<OpaqueTypeNode>(); }
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/script/printer/ir/prim_type.cc
b/src/script/printer/ir/prim_type.cc
index b3ef7ccec3..dc74aadde6 100644
--- a/src/script/printer/ir/prim_type.cc
+++ b/src/script/printer/ir/prim_type.cc
@@ -117,6 +117,16 @@ TVM_FFI_STATIC_INIT_BLOCK() {
kDocTranslate, FDocTranslate::FromNative<&StringTypeDocTranslate>());
}
+ffi::Optional<ExprDoc> TensorRegionTypeDocTranslate(DocTranslatorObj*,
ffi::AnyView,
+ const ffi::Object*) {
+ return NamespaceDoc("ir")->Attr("TensorRegionType")->Call({});
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::TypeAttrDef<TensorRegionTypeNode>().attr(
+ kDocTranslate,
FDocTranslate::FromNative<&TensorRegionTypeDocTranslate>());
+}
+
ffi::Optional<ExprDoc> TupleTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
const ffi::Object*) {
const auto* ty =
diff --git a/src/script/printer/script_printer.cc
b/src/script/printer/script_printer.cc
index 06dbf1eb05..1d450ceac5 100644
--- a/src/script/printer/script_printer.cc
+++ b/src/script/printer/script_printer.cc
@@ -69,6 +69,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
RegisterScriptRepr<RangeNode>();
RegisterScriptRepr<StringImmNode>();
RegisterScriptRepr<TensorLoadNode>();
+ RegisterScriptRepr<TensorRegionTypeNode>();
RegisterScriptRepr<TupleTypeNode>();
RegisterScriptRepr<VarNode>();
RegisterScriptRepr<prim::AddNode>();
diff --git a/src/tirx/ir/buffer.cc b/src/tirx/ir/buffer.cc
index 0a1fa5b789..475fe0d736 100644
--- a/src/tirx/ir/buffer.cc
+++ b/src/tirx/ir/buffer.cc
@@ -169,7 +169,7 @@ using IndexDiv = prim::FloorDivNode;
TensorRegion BufferRegion(TensorVar buffer, ffi::Array<Range> region, Span
span) {
TVM_FFI_ICHECK_EQ(buffer->shape.size(), region.size())
<< "Buffer rank and region dimension mismatch";
- return TensorRegion(std::move(buffer), std::move(region),
BufferRegionType(), std::move(span));
+ return TensorRegion(std::move(buffer), std::move(region),
TensorRegionType(), std::move(span));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -203,7 +203,7 @@ TensorRegion BufferRegionFromPoint(TensorVar buffer,
ffi::Array<PrimExpr> indice
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::TypeAttrDef<TensorTypeNode>().def("__subscript_expr_realize__",
RealizeBufferSubscript);
- refl::TypeAttrDef<BufferRegionTypeNode>().def("__subscript_expr_realize__",
+ refl::TypeAttrDef<TensorRegionTypeNode>().def("__subscript_expr_realize__",
RealizeBufferRegionSubscript);
}
diff --git a/src/tirx/ir/type.cc b/src/tirx/ir/type.cc
index d95db688e1..f2be37aebb 100644
--- a/src/tirx/ir/type.cc
+++ b/src/tirx/ir/type.cc
@@ -193,21 +193,6 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
TensorTypeMaybeInplaceM
return ffi::Unchanged();
}
-TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>>
BufferRegionTypeVisit(
- ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
- return std::nullopt;
-}
-
-TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
BufferRegionTypeMutate(
- ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged();
-}
-
-TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
BufferRegionTypeMaybeInplaceMutate(
- ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged();
-}
-
} // namespace
TensorType::TensorType(ffi::String storage_scope, PrimType dtype,
ffi::Array<PrimExpr> shape,
@@ -253,26 +238,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-// TensorRegion
-BufferRegionType::BufferRegionType() : Type(ffi::UnsafeInit{}) {
- static ffi::ObjectPtr<BufferRegionTypeNode> singleton =
ffi::make_object<BufferRegionTypeNode>();
- data_ = singleton;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- BufferRegionTypeNode::RegisterReflection();
- refl::TypeAttrDef<BufferRegionTypeNode>()
- .attr(refl::type_attr::kStructuralVisit,
- ffi::FStructuralVisit::FromNative<&BufferRegionTypeVisit>())
- .attr(refl::type_attr::kStructuralMutate,
- ffi::FStructuralMutate::FromNative<&BufferRegionTypeMutate>())
- .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
-
ffi::FStructuralMutate::FromNative<&BufferRegionTypeMaybeInplaceMutate>());
-
- refl::GlobalDef().def("tirx.BufferRegionType", []() { return
BufferRegionType(); });
-}
-
TensorMapType::TensorMapType(Span span) : Type(ffi::UnsafeInit{}) {
ffi::ObjectPtr<TensorMapTypeNode> n = ffi::make_object<TensorMapTypeNode>();
n->span = std::move(span);
diff --git a/src/tirx/script/printer/buffer.cc
b/src/tirx/script/printer/buffer.cc
index 1c21dc68db..51bc27fb06 100644
--- a/src/tirx/script/printer/buffer.cc
+++ b/src/tirx/script/printer/buffer.cc
@@ -233,18 +233,9 @@ ffi::Optional<ExprDoc>
TensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView
{TupleDoc(shape), LiteralDoc::DataType(buffer->dtype->dtype,
std::nullopt)}, keys, values);
}
-ffi::Optional<ExprDoc> BufferRegionTypeDocTranslate(DocTranslatorObj*,
ffi::AnyView,
- const ffi::Object*) {
- return NamespaceDoc("ir")
- ->Attr("make_node")
- ->Call({LiteralDoc::Str("tirx.BufferRegionType", std::nullopt)});
-}
-
TVM_FFI_STATIC_INIT_BLOCK() {
ffi::reflection::TypeAttrDef<tirx::TensorTypeNode>().attr(
kDocTranslate, FDocTranslate::FromNative<&TensorTypeDocTranslate>());
- ffi::reflection::TypeAttrDef<tirx::BufferRegionTypeNode>().attr(
- kDocTranslate,
FDocTranslate::FromNative<&BufferRegionTypeDocTranslate>());
}
} // namespace
