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 c836e8c942 [REFACTOR][TIRx] Make BufferRegion a typed expression
(#20256)
c836e8c942 is described below
commit c836e8c94262b92a6875228c44e5e96ba8f27fb6
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Sep 1 15:41:58 2026 -0400
[REFACTOR][TIRx] Make BufferRegion a typed expression (#20256)
BufferRegion describes a retained buffer footprint rather than a
primitive value. Model it as a typed expression while keeping
value-producing loads explicit.
- add BufferRegionType and expression/statement traversal
- remove the implicit/generic BufferRegion-to-primitive conversion path
and construct point/ramp loads explicitly at primitive consumers
- integrate BufferRegion with eager type-directed subscription so
further slicing yields a narrowed BufferRegion, while full point
indexing yields a TensorLoad
---
include/tvm/tirx/buffer_region.h | 80 +++++++++++++++++++
include/tvm/tirx/expr_functor.h | 5 ++
include/tvm/tirx/stmt.h | 51 +-----------
include/tvm/tirx/stmt_functor.h | 2 +
python/tvm/s_tir/tensor_intrin/rocm.py | 24 ++++--
python/tvm/tirx/__init__.py | 2 +-
python/tvm/tirx/expr.py | 10 +--
python/tvm/tirx/expr_functor.py | 26 +++++++
python/tvm/tirx/op.py | 40 ++++++++--
python/tvm/tirx/script/builder/ir.py | 5 ++
python/tvm/tirx/stmt.py | 12 ++-
python/tvm/tirx/stmt_functor.py | 12 +--
src/tirx/ir/buffer.cc | 2 +-
src/tirx/ir/expr_functor.cc | 19 +++++
src/tirx/ir/stmt.cc | 91 +++++++++++++++++-----
src/tirx/ir/stmt_functor.cc | 20 +++++
src/tirx/ir/tir_visitor_with_path.cc | 4 +
src/tirx/ir/tir_visitor_with_path.h | 3 +
.../test_meta_schedule_trace_apply.py | 8 +-
.../test_s_tir_transform_inject_virtual_thread.py | 12 +--
...s_tir_transform_lower_cross_thread_reduction.py | 16 ++--
.../test_tir_analysis_verify_well_formed.py | 12 +++
tests/python/tirx-base/test_tir_constructor.py | 37 +++++++++
...est_tir_transform_pointer_value_type_rewrite.py | 12 +--
.../test_tir_transform_storage_rewrite.py | 2 +-
.../tirx-transform/test_tir_transform_vectorize.py | 6 +-
tests/python/tirx/test_op_namespace_cleanup.py | 2 +-
tests/python/tirx/test_parser_printer.py | 25 +++++-
28 files changed, 412 insertions(+), 128 deletions(-)
diff --git a/include/tvm/tirx/buffer_region.h b/include/tvm/tirx/buffer_region.h
new file mode 100644
index 0000000000..72978daec8
--- /dev/null
+++ b/include/tvm/tirx/buffer_region.h
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#ifndef TVM_TIRX_BUFFER_REGION_H_
+#define TVM_TIRX_BUFFER_REGION_H_
+
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/expr.h>
+#include <tvm/tirx/buffer.h>
+
+namespace tvm {
+namespace tirx {
+
+/*! \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 Representing a region of multi-dimensional buffer access. */
+class BufferRegionNode : public ExprNode {
+ public:
+ BufferVar buffer;
+ ffi::Array<Range> region;
+
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<BufferRegionNode>()
+ .def_ro("buffer", &BufferRegionNode::buffer,
refl::AttachFieldFlag::SEqHashDefRecursive())
+ .def_ro("region", &BufferRegionNode::region);
+ }
+
+ static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferRegion", BufferRegionNode,
ExprNode);
+};
+
+/*! \brief Managed reference to BufferRegionNode. */
+class BufferRegion : public Expr {
+ public:
+ TVM_DLL explicit BufferRegion(BufferVar buffer, ffi::Array<Range> region,
Span span = Span());
+
+ TVM_DLL static BufferRegion FullRegion(BufferVar buffer);
+ TVM_DLL static BufferRegion FromPoint(BufferVar buffer, ffi::Array<PrimExpr>
indices);
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BufferRegion, Expr,
BufferRegionNode);
+ TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferRegionNode);
+};
+
+} // namespace tirx
+} // namespace tvm
+
+#endif // TVM_TIRX_BUFFER_REGION_H_
diff --git a/include/tvm/tirx/expr_functor.h b/include/tvm/tirx/expr_functor.h
index 72590da21a..d44194a7e7 100644
--- a/include/tvm/tirx/expr_functor.h
+++ b/include/tvm/tirx/expr_functor.h
@@ -27,6 +27,7 @@
#include <tvm/ir/node_functor.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/tirx/buffer_region.h>
#include <utility>
@@ -117,6 +118,7 @@ class ExprFunctor<R(const Expr& n, Args...)> {
virtual R VisitExpr_(const VarNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const TensorLoadNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const OpaqueExprNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const BufferRegionNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const TupleNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const TupleGetItemNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const prim::LetNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
@@ -160,6 +162,7 @@ class ExprFunctor<R(const Expr& n, Args...)> {
IR_EXPR_FUNCTOR_DISPATCH(VarNode);
IR_EXPR_FUNCTOR_DISPATCH(TensorLoadNode);
IR_EXPR_FUNCTOR_DISPATCH(OpaqueExprNode);
+ IR_EXPR_FUNCTOR_DISPATCH(BufferRegionNode);
IR_EXPR_FUNCTOR_DISPATCH(TupleNode);
IR_EXPR_FUNCTOR_DISPATCH(TupleGetItemNode);
IR_EXPR_FUNCTOR_DISPATCH(prim::LetNode);
@@ -211,6 +214,7 @@ class TVM_DLL ExprVisitor : public ExprFunctor<void(const
Expr&)> {
void VisitExpr_(const VarNode* op) override;
void VisitExpr_(const TensorLoadNode* op) override;
void VisitExpr_(const OpaqueExprNode* op) override;
+ void VisitExpr_(const BufferRegionNode* op) override;
void VisitExpr_(const TupleNode* op) override;
void VisitExpr_(const TupleGetItemNode* op) override;
void VisitExpr_(const prim::LetNode* op) override;
@@ -258,6 +262,7 @@ class TVM_DLL ExprMutator : protected
ExprFunctor<Expr(const Expr&)> {
Expr VisitExpr_(const VarNode* op) override;
Expr VisitExpr_(const TensorLoadNode* op) override;
Expr VisitExpr_(const OpaqueExprNode* op) override;
+ Expr VisitExpr_(const BufferRegionNode* op) override;
Expr VisitExpr_(const TupleNode* op) override;
Expr VisitExpr_(const TupleGetItemNode* op) override;
Expr VisitExpr_(const prim::LetNode* op) override;
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 68e1b69e65..cbe20d5c02 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -27,6 +27,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/tirx/buffer.h>
+#include <tvm/tirx/buffer_region.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/layout.h>
@@ -767,56 +768,6 @@ class Continue : public Stmt {
TVM_DEFINE_OBJECT_REF_COW_METHOD(ContinueNode);
};
-/*!
- * \brief Representing the region of multi-dimensional buffer access.
- */
-class BufferRegionNode : public PrimExprConvertibleNode {
- public:
- /*! \brief The buffer of the buffer region. */
- BufferVar buffer;
- /*! \brief The region array of the buffer region. */
- ffi::Array<Range> region;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<BufferRegionNode>()
- .def_ro("buffer", &BufferRegionNode::buffer,
refl::AttachFieldFlag::SEqHashDefRecursive())
- .def_ro("region", &BufferRegionNode::region);
- }
-
- TVM_DLL PrimExpr ToPrimExpr() const final;
-
- static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferRegion", BufferRegionNode,
PrimExprConvertibleNode);
-};
-
-/*!
- * \brief Managed reference to BufferRegionNode.
- * \sa BufferRegionNode
- */
-class BufferRegion : public PrimExprConvertible {
- public:
- TVM_DLL explicit BufferRegion(BufferVar buffer, ffi::Array<Range> region);
-
- /*!
- * \brief Create a BufferRegion which is full region of the given buffer.
- * \param buffer The buffer to generate full BufferRegion.
- * \return The BufferRegion which covers all region of the given buffer
- */
- TVM_DLL static BufferRegion FullRegion(BufferVar buffer);
-
- /*!
- * \brief Create a BufferRegion which is a single point of the given buffer.
- * \param buffer The buffer to generate single point BufferRegion.
- * \param indices The access point indices of the buffer
- * \return The BufferRegion which is the single point of the given buffer.
- */
- TVM_DLL static BufferRegion FromPoint(BufferVar buffer, ffi::Array<PrimExpr>
indices);
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BufferRegion,
PrimExprConvertible, BufferRegionNode);
- TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferRegionNode);
-};
-
/*!
* \brief Match introduces a constraint that the source buffer region can be
remapped to the data
* layout specified by the buffer field. The constraint can be checked in
later part of lowering (or
diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h
index d442f2e842..78ebec8434 100644
--- a/include/tvm/tirx/stmt_functor.h
+++ b/include/tvm/tirx/stmt_functor.h
@@ -343,6 +343,7 @@ class TVM_DLL StmtExprVisitor : public ExprVisitor, public
StmtVisitor {
void VisitExpr(const Expr& e) override { return ExprVisitor::VisitExpr(e); }
void VisitExpr_(const TensorLoadNode* op) override;
+ void VisitExpr_(const BufferRegionNode* op) override;
};
/*!
@@ -362,6 +363,7 @@ class TVM_DLL StmtExprMutator : public ExprMutator, public
StmtMutator {
Expr VisitExpr(const Expr& e) override { return ExprMutator::VisitExpr(e); }
Expr VisitExpr_(const VarNode* op) override;
Expr VisitExpr_(const TensorLoadNode* op) override;
+ Expr VisitExpr_(const BufferRegionNode* op) override;
};
/*!
diff --git a/python/tvm/s_tir/tensor_intrin/rocm.py
b/python/tvm/s_tir/tensor_intrin/rocm.py
index 8573c45304..29749dd443 100644
--- a/python/tvm/s_tir/tensor_intrin/rocm.py
+++ b/python/tvm/s_tir/tensor_intrin/rocm.py
@@ -334,11 +334,11 @@ def get_mfma_intrin(k_dim, in_dtype="float32",
out_dtype="float32", b_transposed
T.writes(C[0:WARP_SIZE, 0:local_size_out])
tx = T.env_thread("threadIdx.x")
T.launch_thread(tx, WARP_SIZE)
- C[tx, 0:local_size_out] = T.call_llvm_pure_intrin(
+ C[tx, T.ramp(0, 1, local_size_out)] = T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id(mfma_intrin),
- A[tx, 0:local_size],
- B[tx, 0:local_size],
- C[tx, 0:local_size_out],
+ A[tx, T.ramp(0, 1, local_size) if local_size > 1 else 0],
+ B[tx, T.ramp(0, 1, local_size) if local_size > 1 else 0],
+ C[tx, T.ramp(0, 1, local_size_out)],
T.int32(0),
T.int32(0),
T.int32(0),
@@ -361,11 +361,19 @@ def get_mfma_intrin(k_dim, in_dtype="float32",
out_dtype="float32", b_transposed
tx = T.env_thread("threadIdx.x")
T.launch_thread(tx, WARP_SIZE)
- C[tx, 0:local_size_out] = T.call_llvm_pure_intrin(
+ C[tx, T.ramp(0, 1, local_size_out)] = T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id(mfma_intrin),
- T.call_intrin("int32", "tirx.reinterpret", A[tx,
0:local_size]),
- T.call_intrin("int32", "tirx.reinterpret", A[tx,
0:local_size]),
- C[tx, 0:local_size_out],
+ T.call_intrin(
+ "int32",
+ "tirx.reinterpret",
+ A[tx, T.ramp(0, 1, local_size) if local_size > 1 else 0],
+ ),
+ T.call_intrin(
+ "int32",
+ "tirx.reinterpret",
+ B[tx, T.ramp(0, 1, local_size) if local_size > 1 else 0],
+ ),
+ C[tx, T.ramp(0, 1, local_size_out)],
T.int32(0),
T.int32(0),
T.int32(0),
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index ae228360ae..48712407b8 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -51,7 +51,7 @@ from .stmt import BufferStore, AllocBuffer, AttrStmt,
DeclBuffer
from .stmt import SeqStmt
from .stmt import IfThenElse, Evaluate, stmt_seq, stmt_list
-from .stmt import BufferRegion, MatchBufferRegion, SBlock, SBlockRealize
+from .stmt import BufferRegion, BufferRegionType, MatchBufferRegion, SBlock,
SBlockRealize
from .stmt import ScopeIdDefStmt
from .tile_primitive import DispatchContext, LambdaExpr, TilePrimitiveCall
diff --git a/python/tvm/tirx/expr.py b/python/tvm/tirx/expr.py
index 7536dcba8b..6c1b6cde94 100644
--- a/python/tvm/tirx/expr.py
+++ b/python/tvm/tirx/expr.py
@@ -83,13 +83,13 @@ def _dtype_is_float(value):
def _is_scalar_operand(value):
- if isinstance(value, ExprOp | int | float) or ir.is_prim_expr(value):
- return True
-
- # BufferRegion is a C++ PrimExprConvertible, but its Python wrapper is not
an ExprOp.
from .stmt import BufferRegion # pylint: disable=import-outside-toplevel
- return isinstance(value, BufferRegion)
+ if isinstance(value, BufferRegion):
+ raise TypeError(
+ "BufferRegion is not a primitive operand; construct a BufferLoad
explicitly"
+ )
+ return isinstance(value, ExprOp | int | float) or ir.is_prim_expr(value)
class ExprOp:
diff --git a/python/tvm/tirx/expr_functor.py b/python/tvm/tirx/expr_functor.py
index ffa84fec61..f6dc958674 100644
--- a/python/tvm/tirx/expr_functor.py
+++ b/python/tvm/tirx/expr_functor.py
@@ -49,6 +49,7 @@ class ExprFunctor:
def __init__(self):
self._dispatch_map = {
+ "tirx.BufferRegion": self.visit_buffer_region_,
"Var": self.visit_var_,
"TensorLoad": self.visit_buffer_load_,
"Tuple": self.visit_tuple_,
@@ -124,6 +125,10 @@ class ExprFunctor:
"""Default visitor for an opaque construction-time expression."""
return self.visit_expr_default_(op)
+ def visit_buffer_region_(self, op):
+ """Default visitor for BufferRegion node."""
+ return self.visit_expr_default_(op)
+
def visit_tuple_(self, op):
"""Default visitor for Tuple node."""
return self.visit_expr_default_(op)
@@ -291,6 +296,12 @@ class ExprVisitor(ExprFunctor):
"""Visitor implementation for an opaque construction-time
expression."""
pass
+ def visit_buffer_region_(self, op):
+ """Visitor implementation for BufferRegion."""
+ for region in op.region:
+ self.visit_expr(region.min)
+ self.visit_expr(region.extent)
+
def visit_tuple_(self, op):
"""Visitor implementation for Tuple."""
_visit_array(op.fields, self.visit_expr)
@@ -477,6 +488,21 @@ class ExprMutator(ExprFunctor):
"""Mutator implementation for an opaque construction-time
expression."""
return op
+ def visit_buffer_region_(self, op):
+ """Mutator implementation for BufferRegion."""
+
+ def mutate_range(old):
+ new_min = self.visit_expr(old.min)
+ new_extent = self.visit_expr(old.extent)
+ if new_min is old.min and new_extent is old.extent:
+ return old
+ return Range.from_min_extent(new_min, new_extent)
+
+ region = [mutate_range(r) for r in op.region]
+ if all(old is new for old, new in zip(op.region, region)):
+ return op
+ return tvm.tirx.BufferRegion(op.buffer, region)
+
def visit_tuple_(self, op):
"""Mutator implementation for Tuple."""
fields = [self.visit_expr(field) for field in op.fields]
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index 4d9627292a..73422f82db 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -61,6 +61,16 @@ def _canonical_device_intrin_name(func_name: str) -> str:
return func_name
+def _reject_buffer_region(value, api_name):
+ """Reject region metadata where a call argument must denote a runtime
value."""
+ if isinstance(value, tirx.BufferRegion):
+ raise TypeError(
+ f"tirx.{api_name} does not accept BufferRegion arguments; "
+ "construct a BufferLoad with explicit indices"
+ )
+ return value
+
+
def _primexpr_ty(expr):
"""Return the runtime primitive type of an expression."""
if isinstance(expr, tvm.ir.PrimType):
@@ -135,7 +145,10 @@ def call_packed_lowered(*args, span=None):
--------
te.extern : Create tensor with extern function call.
"""
- call_args = [_pack_buffer(x) if is_buffer_var(x) else x for x in args]
+ call_args = [
+ _pack_buffer(x) if is_buffer_var(x) else _reject_buffer_region(x,
"call_packed_lowered")
+ for x in args
+ ]
return Call(Op.get("tirx.tvm_call_packed_lowered"), call_args, span=span,
ret_ty="int32")
@@ -161,7 +174,10 @@ def call_cpacked_lowered(*args, span=None):
--------
te.extern : Create tensor with extern function call.
"""
- call_args = [_pack_buffer(x) if is_buffer_var(x) else x for x in args]
+ call_args = [
+ _pack_buffer(x) if is_buffer_var(x) else _reject_buffer_region(x,
"call_cpacked_lowered")
+ for x in args
+ ]
return Call(Op.get("tirx.tvm_call_cpacked_lowered"), call_args, span=span,
ret_ty="int32")
@@ -192,7 +208,10 @@ def call_packed(*args, span=None):
--------
te.extern : Create tensor with extern function call.
"""
- call_args = [_pack_buffer(x) if is_buffer_var(x) else x for x in args]
+ call_args = [
+ _pack_buffer(x) if is_buffer_var(x) else _reject_buffer_region(x,
"call_packed")
+ for x in args
+ ]
return Call(Op.get("tirx.tvm_call_packed"), call_args, span=span,
ret_ty="int32")
@@ -219,7 +238,10 @@ def call_cpacked(*args, span=None):
--------
te.extern : Create tensor with extern function call.
"""
- call_args = [_pack_buffer(x) if is_buffer_var(x) else x for x in args]
+ call_args = [
+ _pack_buffer(x) if is_buffer_var(x) else _reject_buffer_region(x,
"call_cpacked")
+ for x in args
+ ]
return Call(Op.get("tirx.tvm_call_cpacked"), call_args, span=span,
ret_ty="int32")
@@ -253,6 +275,7 @@ def call_intrin(dtype: str | tvm.ir.Type, func_name, *args,
attrs=None, span=Non
"""
if isinstance(func_name, str):
func_name = _canonical_device_intrin_name(func_name)
+ args = tuple(_reject_buffer_region(arg, "call_intrin") for arg in args)
return Call(func_name, args, attrs=attrs, span=span, ret_ty=dtype)
@@ -280,7 +303,7 @@ def call_pure_extern(dtype, func_name, *args, span=None):
"""
return Call(
Op.get("tirx.call_pure_extern"),
- [func_name, *args],
+ [func_name, *(_reject_buffer_region(arg, "call_pure_extern") for arg
in args)],
span=span,
ret_ty=dtype,
)
@@ -310,7 +333,7 @@ def call_extern(dtype, func_name, *args, span=None):
"""
return Call(
Op.get("tirx.call_extern"),
- [func_name, *args],
+ [func_name, *(_reject_buffer_region(arg, "call_extern") for arg in
args)],
span=span,
ret_ty=dtype,
)
@@ -530,6 +553,7 @@ def call_tir(global_var: tvm.ir.GlobalVar, *args):
The call expression.
"""
assert isinstance(global_var, tvm.ir.GlobalVar)
+ args = tuple(_reject_buffer_region(arg, "call_tir") for arg in args)
dtype = "void"
if global_var.ty is not None:
@@ -1288,7 +1312,9 @@ def trace(args, trace_action="tvm.default_trace_action"):
"""
if not isinstance(args, list):
raise Exception("tvm.tirx.trace consumes the args as list type")
- call_args = [_pack_buffer(x) if is_buffer_var(x) else x for x in args]
+ call_args = [
+ _pack_buffer(x) if is_buffer_var(x) else _reject_buffer_region(x,
"trace") for x in args
+ ]
call_args.insert(0, tvm.tirx.StringImm(trace_action))
tracing_value = args[-1]
ret_ty = tracing_value.ty if isinstance(tracing_value, Expr) else
tracing_value.dtype
diff --git a/python/tvm/tirx/script/builder/ir.py
b/python/tvm/tirx/script/builder/ir.py
index c544d35902..1782c5ff6a 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/python/tvm/tirx/script/builder/ir.py
@@ -2390,6 +2390,11 @@ def evaluate(value: Expr) -> None:
value = StringImm(value)
if isinstance(value, bool):
value = IntImm("bool", value)
+ if isinstance(value, tir.BufferRegion):
+ raise TypeError(
+ "T.evaluate does not accept BufferRegion values; "
+ "construct a BufferLoad with explicit indices"
+ )
return _ffi_api.Evaluate(value) # type: ignore[attr-defined] # pylint:
disable=no-member
diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py
index ff3c27a990..46c7931c58 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 Expr, Range, Span
+from tvm.ir import Expr, Range, Span, Type
from tvm.runtime import Object, Scriptable, const
from . import _ffi_api
@@ -606,8 +606,16 @@ 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 structural type of a :class:`BufferRegion` expression."""
+
+ def __init__(self) -> None:
+ self.__init_handle_by_constructor__(_ffi_api.BufferRegionType) #
type: ignore
+
+
@tvm_ffi.register_object("tirx.BufferRegion")
-class BufferRegion(Object, Scriptable):
+class BufferRegion(Expr, Scriptable):
"""BufferRegion node.
Parameters
diff --git a/python/tvm/tirx/stmt_functor.py b/python/tvm/tirx/stmt_functor.py
index a0a7c4a553..83405d7eea 100644
--- a/python/tvm/tirx/stmt_functor.py
+++ b/python/tvm/tirx/stmt_functor.py
@@ -361,12 +361,12 @@ class StmtVisitor(StmtFunctor):
def visit_op_call_(self, op):
"""Visitor implementation for TilePrimitiveCall."""
for arg in op.args:
- if isinstance(arg, tvm.ir.Expr):
+ if isinstance(arg, tvm.tirx.BufferRegion):
+ self.visit_buffer_region_(arg)
+ elif isinstance(arg, tvm.ir.Expr):
self.visit_expr(arg)
elif isinstance(arg, tvm.tirx.Stmt):
self.visit_stmt(arg)
- elif isinstance(arg, tvm.tirx.BufferRegion):
- self.visit_buffer_region_(arg)
for value in op.config.values():
if isinstance(value, tvm.ir.Expr):
self.visit_expr(value)
@@ -838,12 +838,12 @@ class StmtMutator(StmtFunctor):
args_changed = False
for arg in op.args:
- if isinstance(arg, tvm.ir.Expr):
+ if isinstance(arg, tvm.tirx.BufferRegion):
+ new_arg = self.visit_buffer_region_(arg)
+ elif isinstance(arg, tvm.ir.Expr):
new_arg = self.visit_expr(arg)
elif isinstance(arg, tvm.tirx.Stmt):
new_arg = self.visit_stmt(arg)
- elif isinstance(arg, tvm.tirx.BufferRegion):
- new_arg = self.visit_buffer_region_(arg)
else:
new_arg = arg
diff --git a/src/tirx/ir/buffer.cc b/src/tirx/ir/buffer.cc
index a262ce7957..577c163474 100644
--- a/src/tirx/ir/buffer.cc
+++ b/src/tirx/ir/buffer.cc
@@ -99,7 +99,7 @@ ffi::ObjectRef RealizeBufferSubscript(
region.push_back(
Range::FromMinExtent(IntImm(buffer_ty->shape[i].ty(), 0),
buffer_ty->shape[i]));
}
- return BufferRegion(buffer, region);
+ return BufferRegion(buffer, region, span);
}
} // namespace
diff --git a/src/tirx/ir/expr_functor.cc b/src/tirx/ir/expr_functor.cc
index 3399114864..e35354fdd4 100644
--- a/src/tirx/ir/expr_functor.cc
+++ b/src/tirx/ir/expr_functor.cc
@@ -38,6 +38,13 @@ void ExprVisitor::VisitExpr_(const TensorLoadNode* op) {
void ExprVisitor::VisitExpr_(const OpaqueExprNode* op) {}
+void ExprVisitor::VisitExpr_(const BufferRegionNode* op) {
+ VisitArray(op->region, [this](const Range& range) {
+ this->VisitExpr(range->min);
+ this->VisitExpr(range->extent);
+ });
+}
+
void ExprVisitor::VisitExpr_(const TupleNode* op) {
VisitArray(op->fields, [this](const Expr& e) { this->VisitExpr(e); });
}
@@ -120,6 +127,18 @@ Expr ExprMutator::VisitExpr_(const TensorLoadNode* op) {
Expr ExprMutator::VisitExpr_(const OpaqueExprNode* op) { return
ffi::GetRef<OpaqueExpr>(op); }
+Expr ExprMutator::VisitExpr_(const BufferRegionNode* op) {
+ ffi::Array<Range> region = op->region.Map([this](const Range& range) {
+ PrimExpr min = this->VisitPrimExpr(range->min);
+ PrimExpr extent = this->VisitPrimExpr(range->extent);
+ return min.same_as(range->min) && extent.same_as(range->extent)
+ ? range
+ : Range::FromMinExtent(std::move(min), std::move(extent));
+ });
+ return region.same_as(op->region) ? ffi::GetRef<BufferRegion>(op)
+ : BufferRegion(op->buffer,
std::move(region), op->span);
+}
+
Expr ExprMutator::VisitExpr_(const TupleNode* op) {
ffi::Array<Expr> fields =
op->fields.Map([this](const Expr& field) { return
this->VisitExpr(field); });
diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc
index c3ee7b8e97..c72e9d1d3b 100644
--- a/src/tirx/ir/stmt.cc
+++ b/src/tirx/ir/stmt.cc
@@ -34,7 +34,66 @@
namespace tvm {
namespace tirx {
+namespace {
+
+using SubscriptSlice = ffi::Array<ffi::Variant<
+ ffi::Tuple<ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>>,
+ PrimExpr>>;
+
+ffi::ObjectRef RealizeBufferRegionSubscript(Expr value, SubscriptSlice slice,
Span span) {
+ BufferRegion source = value.as_or_throw<BufferRegion>();
+ TVM_FFI_CHECK_LE(slice.size(), source->region.size(), IndexError)
+ << "Too many indices for a " << source->region.size() << "-dimensional
buffer region";
+
+ bool all_points = slice.size() == source->region.size();
+ for (const auto& item : slice) {
+ if (auto descriptor = item.as<ffi::Tuple<ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>,
+ ffi::Optional<PrimExpr>>>()) {
+ all_points = false;
+ ffi::Optional<PrimExpr> step = descriptor.value().get<2>();
+ TVM_FFI_CHECK(!step.has_value() || is_one(step.value()), ValueError)
+ << "BufferRegion slices with a non-unit step are not supported";
+ }
+ }
+
+ if (all_points) {
+ ffi::Array<PrimExpr> indices;
+ indices.reserve(slice.size());
+ for (size_t i = 0; i < slice.size(); ++i) {
+ indices.push_back(source->region[i]->min +
slice[i].as<PrimExpr>().value());
+ }
+ return BufferLoad(source->buffer, indices, span);
+ }
+
+ arith::Analyzer analyzer;
+ ffi::Array<Range> region;
+ region.reserve(source->region.size());
+ for (size_t i = 0; i < slice.size(); ++i) {
+ const Range& old_range = source->region[i];
+ if (auto point = slice[i].as<PrimExpr>()) {
+ PrimExpr new_min = old_range->min + point.value();
+ region.push_back(Range::FromMinExtent(new_min,
IntImm(point.value().ty(), 1)));
+ } else {
+ auto descriptor = slice[i]
+ .as<ffi::Tuple<ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>,
+ ffi::Optional<PrimExpr>>>()
+ .value();
+ PrimExpr start =
descriptor.get<0>().value_or(IntImm(old_range->extent.ty(), 0));
+ PrimExpr stop = descriptor.get<1>().value_or(old_range->extent);
+ region.push_back(
+ Range::FromMinExtent(old_range->min + start, analyzer->Simplify(stop
- start)));
+ }
+ }
+ for (size_t i = slice.size(); i < source->region.size(); ++i) {
+ region.push_back(source->region[i]);
+ }
+ return BufferRegion(source->buffer, region, span);
+}
+
+} // namespace
+
TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
StmtNode::RegisterReflection();
BindNode::RegisterReflection();
@@ -51,6 +110,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ReturnNode::RegisterReflection();
BreakNode::RegisterReflection();
ContinueNode::RegisterReflection();
+ BufferRegionTypeNode::RegisterReflection();
+ refl::TypeAttrDef<BufferRegionTypeNode>().def("__subscript_expr_realize__",
+ RealizeBufferRegionSubscript);
BufferRegionNode::RegisterReflection();
MatchBufferRegionNode::RegisterReflection();
SBlockNode::RegisterReflection();
@@ -486,28 +548,18 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
// BufferRegion
-PrimExpr BufferRegionNode::ToPrimExpr() const {
- // Auto convert to PrimExpr if it is a single point load
- ffi::Array<PrimExpr> indices;
- indices.reserve(this->region.size());
- for (const Range& r : this->region) {
- if (tvm::tirx::is_one(r->extent)) {
- indices.push_back(r->min);
- } else if (r->extent.as<IntImmNode>()) {
- indices.push_back(prim::Ramp(r->min, IntImm(r->min.ty(), 1), r->extent));
- } else {
- TVM_FFI_THROW(ValueError) << "Cannot convert to BufferLoad: "
- << ffi::GetRef<BufferRegion>(this);
- }
- }
- return tirx::BufferLoad(this->buffer, indices);
+BufferRegionType::BufferRegionType() : Type(ffi::UnsafeInit{}) {
+ static ffi::ObjectPtr<BufferRegionTypeNode> singleton =
ffi::make_object<BufferRegionTypeNode>();
+ data_ = singleton;
}
-BufferRegion::BufferRegion(BufferVar buffer, ffi::Array<Range> region) {
+BufferRegion::BufferRegion(BufferVar buffer, ffi::Array<Range> region, Span
span) {
TVM_FFI_ICHECK_EQ(buffer->shape.size(), region.size())
<< "The dimension between " << buffer << " and region " << region
<< " mismatched, the buffer is " << buffer;
ffi::ObjectPtr<BufferRegionNode> node = ffi::make_object<BufferRegionNode>();
+ node->ty = BufferRegionType();
+ node->span = std::move(span);
node->buffer = std::move(buffer);
node->region = std::move(region);
data_ = std::move(node);
@@ -536,9 +588,10 @@ BufferRegion BufferRegion::FromPoint(BufferVar buffer,
ffi::Array<PrimExpr> indi
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.BufferRegion", [](BufferVar buffer,
ffi::Array<Range> region) {
- return BufferRegion(buffer, region);
- });
+ refl::GlobalDef()
+ .def("tirx.BufferRegionType", []() { return BufferRegionType(); })
+ .def("tirx.BufferRegion",
+ [](BufferVar buffer, ffi::Array<Range> region) { return
BufferRegion(buffer, region); });
}
// MatchBufferRegion
diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc
index e52911e51b..ee0161d42d 100644
--- a/src/tirx/ir/stmt_functor.cc
+++ b/src/tirx/ir/stmt_functor.cc
@@ -95,6 +95,11 @@ void StmtExprVisitor::VisitExpr_(const TensorLoadNode* op) {
ExprVisitor::VisitExpr_(op);
}
+void StmtExprVisitor::VisitExpr_(const BufferRegionNode* op) {
+ this->VisitBufferUse(op->buffer);
+ ExprVisitor::VisitExpr_(op);
+}
+
void StmtVisitor::VisitStmt_(const AllocBufferNode* op) {
this->VisitBufferDef(op->buffer, /*alloc_data=*/true);
}
@@ -463,6 +468,21 @@ Expr StmtExprMutator::VisitExpr_(const TensorLoadNode* op)
{
return expr;
}
+Expr StmtExprMutator::VisitExpr_(const BufferRegionNode* op) {
+ BufferVar new_buf = this->VisitBufferUse(op->buffer);
+ ffi::Array<Range> new_region = op->region.Map([this](const Range& range) {
+ PrimExpr min = this->VisitPrimExpr(range->min);
+ PrimExpr extent = this->VisitPrimExpr(range->extent);
+ return min.same_as(range->min) && extent.same_as(range->extent)
+ ? range
+ : Range::FromMinExtent(std::move(min), std::move(extent));
+ });
+ if (new_buf.same_as(op->buffer) && new_region.same_as(op->region)) {
+ return ffi::GetRef<BufferRegion>(op);
+ }
+ return BufferRegion(std::move(new_buf), std::move(new_region), op->span);
+}
+
Stmt StmtMutator::VisitStmt_(const AllocBufferNode* op) {
BufferVar new_buf = this->VisitBufferDef(op->buffer, /*alloc_data=*/true);
diff --git a/src/tirx/ir/tir_visitor_with_path.cc
b/src/tirx/ir/tir_visitor_with_path.cc
index 14bab9aad6..4403fd66a6 100644
--- a/src/tirx/ir/tir_visitor_with_path.cc
+++ b/src/tirx/ir/tir_visitor_with_path.cc
@@ -357,6 +357,10 @@ void TIRVisitorWithPath::VisitExpr_(const TensorLoadNode*
op, AccessPath path) {
Visit(op->indices, path->Attr("indices"));
}
+void TIRVisitorWithPath::VisitExpr_(const BufferRegionNode* op, AccessPath
path) {
+ Visit(ffi::GetRef<BufferRegion>(op), path);
+}
+
void TIRVisitorWithPath::VisitExpr_(const OpaqueExprNode* op, AccessPath path)
{}
void TIRVisitorWithPath::VisitExpr_(const TupleNode* op, AccessPath path) {
diff --git a/src/tirx/ir/tir_visitor_with_path.h
b/src/tirx/ir/tir_visitor_with_path.h
index 698e1ce1cf..b14ad713f3 100644
--- a/src/tirx/ir/tir_visitor_with_path.h
+++ b/src/tirx/ir/tir_visitor_with_path.h
@@ -66,6 +66,8 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const
Expr&, ffi::reflecti
VisitExpr_(tuple, path);
} else if (auto* tuple_get_item = obj.as<TupleGetItemNode>()) {
VisitExpr_(tuple_get_item, path);
+ } else if (auto* buffer_region = obj.as<BufferRegionNode>()) {
+ VisitExpr_(buffer_region, path);
} else if (obj.as<OpaqueExprNode>()) {
VisitExpr(obj, path);
} else {
@@ -152,6 +154,7 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const
Expr&, ffi::reflecti
using ExprFunctor::VisitExpr;
void VisitExpr_(const VarNode* op, ffi::reflection::AccessPath path)
override;
void VisitExpr_(const TensorLoadNode* op, ffi::reflection::AccessPath path)
override;
+ void VisitExpr_(const BufferRegionNode* op, ffi::reflection::AccessPath
path) override;
void VisitExpr_(const OpaqueExprNode* op, ffi::reflection::AccessPath path)
override;
void VisitExpr_(const TupleNode* op, ffi::reflection::AccessPath path)
override;
void VisitExpr_(const TupleGetItemNode* op, ffi::reflection::AccessPath
path) override;
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
b/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
index befa940157..42ec6b8384 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
@@ -1157,12 +1157,12 @@ def get_conv2d_vnni_mod(intrin_id):
A = T.match_buffer(p0[n, ic_outer, oh + kh, ow +
kw, ic_f_inner * 4 : ic_f_inner * 4 + 4], [4], dtype="uint8", offset_factor=1)
B = T.match_buffer(p1[oc_chunk, ic_outer, kh, kw,
ic_f_inner, 0 : 16, 0 : 4], [16, 4], dtype="int8", offset_factor=1)
C = T.match_buffer(conv2d_NCHWc_int8[n, oc_chunk,
oh, ow, 0 : 16], [16], dtype="int32", offset_factor=1)
- A_u8x4: T.uint8x4 = A[0:4]
+ A_u8x4: T.uint8x4 = A[T.ramp(0, 1, 4)]
A_i32: T.int32 = T.reinterpret(A_u8x4,
dtype="int32")
- B_i8x64: T.int8x64 = B[0, 0:64]
+ B_i8x64: T.int8x64 = B[0, T.ramp(0, 1, 64)]
B_i32x16: T.int32x16 = T.reinterpret(B_i8x64,
dtype="int32x16")
- C_i32x16: T.int32x16 = C[0:16]
- C[0:16] =
T.call_llvm_pure_intrin(T.uint32(intrin_id), C_i32x16, T.broadcast(A_i32, 16),
B_i32x16, dtype="int32x16")
+ C_i32x16: T.int32x16 = C[T.ramp(0, 1, 16)]
+ C[T.ramp(0, 1, 16)] =
T.call_llvm_pure_intrin(T.uint32(intrin_id), C_i32x16, T.broadcast(A_i32, 16),
B_i32x16, dtype="int32x16")
for ax0, ax1, ax2, ax3 in T.grid(1, 1, 1, 7):
for ax4_fused in T.vectorized(16):
with T.sblock("T_cast_8"):
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
index 6ba06251c9..471974f16c 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
@@ -165,17 +165,17 @@ def test_vthread_simplified():
vthread = T.env_thread("vthread")
T.launch_thread(vthread, 4)
B = T.alloc_buffer((4,), "int32", scope="shared")
- B[0:4] = T.broadcast(vthread, 4)
+ B[T.ramp(0, 1, 4)] = T.broadcast(vthread, 4)
@T.prim_func(s_tir=True)
def expected_func():
B = T.alloc_buffer((16,), "int32", scope="shared")
# The indices for B should each be a single Ramp node, and
# should not be the sum of a Ramp and Broadcast node.
- B[T.Mul(0, 4) : T.Mul(0, 4) + 4] = T.broadcast(0, 4)
- B[T.Mul(1, 4) : T.Mul(1, 4) + 4] = T.broadcast(1, 4)
- B[T.Mul(2, 4) : T.Mul(2, 4) + 4] = T.broadcast(2, 4)
- B[T.Mul(3, 4) : T.Mul(3, 4) + 4] = T.broadcast(3, 4)
+ B[T.ramp(T.Mul(0, 4), 1, 4)] = T.broadcast(0, 4)
+ B[T.ramp(T.Mul(1, 4), 1, 4)] = T.broadcast(1, 4)
+ B[T.ramp(T.Mul(2, 4), 1, 4)] = T.broadcast(2, 4)
+ B[T.ramp(T.Mul(3, 4), 1, 4)] = T.broadcast(3, 4)
before_mod = tvm.IRModule.from_expr(before_func.with_attr("global_symbol",
"main"))
after_mod = tvm.s_tir.transform.InjectVirtualThread()(before_mod)
@@ -192,7 +192,7 @@ def test_vthread_vectorized():
vthread = T.env_thread("vthread")
T.launch_thread(vthread, 4)
B = T.alloc_buffer((4,), "int32", scope="shared")
- B[0:4] = T.broadcast(vthread, 4)
+ B[T.ramp(0, 1, 4)] = T.broadcast(vthread, 4)
before_mod = tvm.IRModule.from_expr(before_func.with_attr("global_symbol",
"main"))
intermediate_mod = tvm.s_tir.transform.InjectVirtualThread()(before_mod)
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py
b/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py
index 88e4c7ee13..e327242972 100644
---
a/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py
+++
b/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py
@@ -829,12 +829,16 @@ def single_reduction_loop_with_tensorize(
C = T.match_buffer(
output[n, oc_chunk, oh, ow, 0:32], [32], dtype="int32",
offset_factor=1
)
- A_u8x4: T.uint8x4 = A[0:4]
+ A_u8x4: T.uint8x4 = A[T.ramp(0, 1, 4)]
A_i32: T.int32 = T.reinterpret(A_u8x4, dtype="int32")
- B_i8x128 = B[0, 0:128]
+ B_i8x128 = B[0, T.ramp(0, 1, 128)]
B_i32x32: T.int32x32 = T.reinterpret(B_i8x128,
dtype="int32x32")
- C[0:32] = T.call_llvm_pure_intrin(
- 4217, C[0:32], T.broadcast(A_i32, 32), B_i32x32,
dtype="int32x32"
+ C[T.ramp(0, 1, 32)] = T.call_llvm_pure_intrin(
+ 4217,
+ C[T.ramp(0, 1, 32)],
+ T.broadcast(A_i32, 32),
+ B_i32x32,
+ dtype="int32x32",
)
@@ -881,9 +885,9 @@ def nested_reduction_loop_with_inner_match_buffers(
offset_factor=1,
)
C = T.match_buffer(out[yi, xr], [1], dtype="int32",
offset_factor=1)
- A_i8x4: T.int8x4 = A[0:4]
+ A_i8x4: T.int8x4 = A[T.ramp(0, 1, 4)]
A_i32: T.int32 = T.reinterpret(A_i8x4, dtype="int32")
- B_i8x4: T.int8x4 = B[0:4]
+ B_i8x4: T.int8x4 = B[T.ramp(0, 1, 4)]
B_i32: T.int32 = T.reinterpret(B_i8x4, dtype="int32")
C[0] = A_i32 + B_i32 + C[0]
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 9bc3908948..5f1102ea64 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
@@ -47,6 +47,18 @@ def test_pass_simple():
assert
tvm.tirx.analysis.verify_well_formed(tvm.IRModule.from_expr(element_wise))
+def test_buffer_region_bounds_are_visited():
+ data = tvm.tirx.Var(
+ "data", tvm.ir.PointerType(tvm.ir.PrimType("int32"),
storage_scope="global")
+ )
+ buffer = tvm.tirx.decl_buffer([4], "int32", data=data)
+ undefined = tvm.tirx.Var("undefined", "int32")
+ region = tvm.tirx.BufferRegion(buffer,
[tvm.ir.Range.from_min_extent(undefined, 4)])
+ block = tvm.tirx.SBlock([], [region], [], "region", tvm.tirx.Evaluate(0))
+ func = tvm.tirx.PrimFunc([buffer], block)
+ assert not tvm.tirx.analysis.verify_well_formed(func, assert_mode=False)
+
+
def test_fail_use_out_loop_var():
@T.prim_func(check_well_formed=False, s_tir=True)
def element_wise(
diff --git a/tests/python/tirx-base/test_tir_constructor.py
b/tests/python/tirx-base/test_tir_constructor.py
index 37ac416ad8..2c3da5eeb9 100644
--- a/tests/python/tirx-base/test_tir_constructor.py
+++ b/tests/python/tirx-base/test_tir_constructor.py
@@ -20,6 +20,7 @@ import tvm_ffi
import tvm
from tvm import te, topi
+from tvm.script import tirx as T
from tvm.tirx.analysis import expr_deep_equal
from tvm.tirx.expr_functor import ExprMutator
@@ -214,6 +215,42 @@ def test_expr_constructor():
assert x.body == v
+def test_buffer_region_call_wrappers_reject():
+ buffer = tvm.tirx.decl_buffer([4], "int32")
+ region = buffer[0:4]
+ calls = [
+ lambda: tvm.tirx.call_intrin("int32", "tirx.reinterpret", region),
+ lambda: tvm.tirx.call_extern("int32", "consume", region),
+ lambda: tvm.tirx.call_pure_extern("int32", "consume", region),
+ lambda: tvm.tirx.call_packed("consume", region),
+ lambda: tvm.tirx.call_cpacked("consume", region, 0),
+ lambda: tvm.tirx.call_packed_lowered("consume", region),
+ lambda: tvm.tirx.call_cpacked_lowered("consume", region, 0),
+ lambda: tvm.tirx.call_tir(tvm.ir.GlobalVar("callee"), region),
+ lambda: tvm.tirx.trace([region]),
+ lambda: T.evaluate(region),
+ ]
+ for call in calls:
+ with pytest.raises(TypeError, match="construct a BufferLoad with
explicit indices"):
+ call()
+
+ assert not hasattr(region, "to_buffer_load")
+
+
+def test_buffer_region_type_is_singleton():
+ lhs = tvm.tirx.decl_buffer([1], "int32")[0:1]
+ rhs = tvm.tirx.decl_buffer([2], "float32")[0:2]
+ assert isinstance(lhs, tvm.tirx.BufferRegion)
+ assert isinstance(rhs, tvm.tirx.BufferRegion)
+ assert lhs.ty.same_as(rhs.ty)
+
+
+def test_buffer_region_is_not_arithmetic_operand():
+ int_region = tvm.tirx.decl_buffer([4], "int32")[0:4]
+ with pytest.raises(TypeError, match="construct a BufferLoad explicitly"):
+ tvm.tirx.IterVar((0, 4), "i", tvm.tirx.IterVar.DataPar) + int_region
+
+
def test_operator_base_categories_have_primitive_type():
var = tvm.tirx.Var("x", "int32")
buffer = tvm.tirx.decl_buffer([4], "float32")
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 2b15b48471..89fd4011a6 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
@@ -31,7 +31,7 @@ def test_rewrite_to_shuffle_0():
def main(A: T.Buffer((16,), "float32"), B: T.Buffer((4,), "float32")):
A_local = T.alloc_buffer((16,), scope="local")
for i in range(4):
- A_local[i * 4 : i * 4 + 4] = A[i * 4 : i * 4 + 4]
+ A_local[T.ramp(i * 4, 1, 4)] = A[T.ramp(i * 4, 1, 4)]
for i in range(4):
B[i] = A_local[i * 4] + A_local[i * 4 + 1] + A_local[i * 4 +
2] + A_local[i * 4 + 3]
@@ -62,8 +62,8 @@ def test_rewrite_to_shuffle_1():
@T.prim_func(s_tir=True)
def main(A: T.Buffer((8,), "float32"), B: T.Buffer((1,), "float32")):
A_local = T.alloc_buffer((8,), scope="local")
- A_local[0:4] = A[0:4]
- A_local[4:8] = A[4:8]
+ A_local[T.ramp(0, 1, 4)] = A[T.ramp(0, 1, 4)]
+ A_local[T.ramp(4, 1, 4)] = A[T.ramp(4, 1, 4)]
B[0] = (
A_local[0]
+ A_local[1]
@@ -106,7 +106,7 @@ def test_address_of():
def main(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")):
for i in range(4):
T.evaluate(T.address_of(A[i * 4]))
- B[i * 4 : i * 4 + 4] = A[i * 4 : i * 4 + 4]
+ B[T.ramp(i * 4, 1, 4)] = A[T.ramp(i * 4, 1, 4)]
@I.ir_module
class Expected:
@@ -114,7 +114,7 @@ def test_address_of():
def main(A: T.Buffer((16,), "float32"), B: T.Buffer((4,),
"float32x4")):
for i in range(4):
T.evaluate(T.address_of(A[i * 4]))
- B[T.Div(i * 4, 4)] = A[i * 4 : i * 4 + 4]
+ B[T.Div(i * 4, 4)] = A[T.ramp(i * 4, 1, 4)]
After = transform(Before)
tvm.ir.assert_structural_equal(After, Expected)
@@ -152,7 +152,7 @@ def test_decl_buffer_alias_chain_uses_flat_root_map():
A_view = T.decl_buffer((16,), "float32", data=A.data)
A_view_2 = T.decl_buffer((16,), "float32", data=A_view.data)
for i in range(4):
- A_view_2[i * 4 : i * 4 + 4] = T.broadcast(T.float32(1), 4)
+ A_view_2[T.ramp(i * 4, 1, 4)] = T.broadcast(T.float32(1), 4)
After = transform(Before)
assert tvm.tirx.analysis.verify_well_formed(After)
diff --git a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py
b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py
index 47473737cc..d5b271b468 100644
--- a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py
+++ b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py
@@ -386,7 +386,7 @@ def test_decl_buffer_is_not_vectorized():
"dummy_func", dtype=T.handle("int32").ty
)
A = T.decl_buffer([8], "int32", data=A_data)
- A[0:8] = T.broadcast(42, 8)
+ A[T.ramp(0, 1, 8)] = T.broadcast(42, 8)
After = tvm.tirx.transform.StorageRewrite()(Before)
tvm.ir.assert_structural_equal(After, Before)
diff --git a/tests/python/tirx-transform/test_tir_transform_vectorize.py
b/tests/python/tirx-transform/test_tir_transform_vectorize.py
index 945e510c0f..e7ac072cb7 100644
--- a/tests/python/tirx-transform/test_tir_transform_vectorize.py
+++ b/tests/python/tirx-transform/test_tir_transform_vectorize.py
@@ -79,7 +79,7 @@ def test_vectorize_vector_scalable_error():
@T.prim_func(s_tir=True)
def main(A: T.Buffer((25,), "float32")):
for j in T.vectorized(T.vscale() * 4):
- A[j * 4 : j * 4 + 4] = T.Broadcast(T.float32(1), 4)
+ A[T.ramp(j * 4, 1, 4)] = T.Broadcast(T.float32(1), 4)
error_msg = "Creating scalable vectors from existing vectors is not
supported."
with tvm.target.Target(sve_target):
@@ -106,7 +106,7 @@ def test_vectorize_vector_scalable_error3():
@T.prim_func(s_tir=True)
def main(A: T.Buffer((25,), "float32")):
for j in T.vectorized(4):
- A[j * T.vscale() * 4 : j * T.vscale() * 4 + T.vscale() * 4] =
T.Broadcast(
+ A[T.ramp(j * T.vscale() * 4, 1, T.vscale() * 4)] = T.Broadcast(
T.float32(1), T.vscale() * 4
)
@@ -122,7 +122,7 @@ def test_vectorize_vector_scalable_error4():
@T.prim_func(private=True, s_tir=True)
def main(A: T.Buffer((25,), "float32")):
for j in T.vectorized(T.vscale() * 4):
- A[j * T.vscale() * 4 : j * T.vscale() * 4 + T.vscale() * 4] =
T.Broadcast(
+ A[T.ramp(j * T.vscale() * 4, 1, T.vscale() * 4)] = T.Broadcast(
T.float32(1), T.vscale() * 4
)
diff --git a/tests/python/tirx/test_op_namespace_cleanup.py
b/tests/python/tirx/test_op_namespace_cleanup.py
index 030bc48aa8..09ad39d4a6 100644
--- a/tests/python/tirx/test_op_namespace_cleanup.py
+++ b/tests/python/tirx/test_op_namespace_cleanup.py
@@ -177,7 +177,7 @@ def
test_device_intrinsic_namespaces_are_canonical_and_classified():
T.cuda.elect_sync(),
T.cuda.thread_fence(),
T.nvshmem.fence(),
- T.nki.identity(buffer[0:1], 1),
+ T.nki.identity(buffer[0], 1),
]
expected = [
diff --git a/tests/python/tirx/test_parser_printer.py
b/tests/python/tirx/test_parser_printer.py
index cf950e63e6..c6f6d5c21e 100644
--- a/tests/python/tirx/test_parser_printer.py
+++ b/tests/python/tirx/test_parser_printer.py
@@ -2268,8 +2268,29 @@ def test_buffer_slice_region():
partial = buf[1]
assert isinstance(partial, BufferRegion)
- with pytest.raises(TypeError):
- _ = partial[2]
+
+ narrowed = br[4:12, 2:10]
+ assert isinstance(narrowed, BufferRegion)
+ assert narrowed.buffer.same_as(buf)
+ assert [(int(dim.min), int(dim.extent)) for dim in narrowed.region] == [
+ (36, 8),
+ (2, 8),
+ ]
+
+ chained_load = br[3, 4]
+ assert isinstance(chained_load, tvm.ir.TensorLoad)
+ assert chained_load.source.same_as(buf)
+ assert [int(index) for index in chained_load.indices] == [35, 4]
+
+ point_then_region = br[3]
+ assert isinstance(point_then_region, BufferRegion)
+ assert [(int(dim.min), int(dim.extent)) for dim in
point_then_region.region] == [
+ (35, 1),
+ (0, 32),
+ ]
+
+ with pytest.raises(ValueError, match="non-unit step"):
+ _ = br[::2]
def test_global_call_realizes_buffer_elements():
