You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/09/17 11:11:16 UTC
(tvm) branch main updated: [REFACTOR][IR] Share tensor region expressions across dialects (#20369)
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 73b7960103 [REFACTOR][IR] Share tensor region expressions across dialects (#20369)
73b7960103 is described below
commit 73b7960103bccddbfdd686d34beb1d90ca89143a
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Sep 17 07:11:10 2026 -0400
[REFACTOR][IR] Share tensor region expressions across dialects (#20369)
Represent indexed regions with shared `ir.TensorRegion`, using an
expression source and ranges alongside the inherited type and span. TIRX
retains buffer validation, construction helpers and `BufferRegionType`
subscript semantics.
Migrate scheduling, access analysis, transformations, script
parsing/printing and tile consumers to the shared representation,
preserving source identities and buffer constraints at dialect
boundaries.
---
include/tvm/ir/expr.h | 30 ++
include/tvm/ir/expr_functor.h | 6 +
include/tvm/s_tir/analysis.h | 10 +-
include/tvm/tirx/buffer_region.h | 34 +--
include/tvm/tirx/expr_functor.h | 16 +-
include/tvm/tirx/script/builder/frame.h | 4 +-
include/tvm/tirx/stmt.h | 12 +-
include/tvm/tirx/stmt_functor.h | 5 +-
include/tvm/tirx/tile_primitive.h | 10 +-
python/tvm/backend/cuda/tile_primitive/common.py | 21 +-
.../backend/cuda/tile_primitive/copy/fallback.py | 4 +-
.../cuda/tile_primitive/copy/ld_stmatrix.py | 6 +-
.../tvm/backend/cuda/tile_primitive/copy/utils.py | 4 +-
.../cuda/tile_primitive/copy/vec_auto_gmem_smem.py | 6 +-
.../cuda/tile_primitive/copy/vec_auto_reg.py | 18 +-
.../backend/cuda/tile_primitive/copy/vec_forced.py | 14 +-
.../cuda/tile_primitive/copy_async/dsmem.py | 8 +-
.../cuda/tile_primitive/copy_async/ldgsts.py | 6 +-
.../cuda/tile_primitive/copy_async/tcgen05_cp.py | 8 +-
.../cuda/tile_primitive/copy_async/tcgen05_ldst.py | 6 +-
.../backend/cuda/tile_primitive/copy_async/tma.py | 16 +-
.../cuda/tile_primitive/elementwise/_common.py | 32 +--
.../tile_primitive/elementwise/ops/__init__.py | 11 +-
.../cuda/tile_primitive/elementwise/ops/binary.py | 9 +-
.../cuda/tile_primitive/elementwise/ops/cast.py | 7 +-
.../cuda/tile_primitive/elementwise/ops/fma.py | 7 +-
.../cuda/tile_primitive/elementwise/ops/unary.py | 10 +-
.../backend/cuda/tile_primitive/elementwise/reg.py | 18 +-
.../cuda/tile_primitive/elementwise/smem.py | 16 +-
.../elementwise/vec_emit/binary_f32x2.py | 4 +-
.../elementwise/vec_emit/cast_vec2.py | 4 +-
.../elementwise/vec_emit/fma_f32x2.py | 6 +-
.../cuda/tile_primitive/gemm/mma_m16n8k_.py | 6 +-
.../cuda/tile_primitive/gemm_async/tcgen05.py | 17 +-
.../tvm/backend/cuda/tile_primitive/gemm_utils.py | 6 +-
.../permute_layout/warp_xor_swizzle.py | 11 +-
.../backend/cuda/tile_primitive/reduction/local.py | 21 +-
.../cuda/tile_primitive/reduction/shared.py | 17 +-
.../cuda/tile_primitive/reduction/sm100_packed.py | 15 +-
.../backend/cuda/tile_primitive/reduction/utils.py | 12 +-
.../backend/trn/tile_primitive/binary/default.py | 4 +-
.../tvm/backend/trn/tile_primitive/binary/utils.py | 11 +-
.../trn/tile_primitive/compose_op/binary_chain.py | 11 +-
.../trn/tile_primitive/compose_op/binary_reduce.py | 15 +-
.../trn/tile_primitive/compose_op/unary_reduce.py | 19 +-
.../tvm/backend/trn/tile_primitive/copy/default.py | 18 +-
python/tvm/backend/trn/tile_primitive/dim_utils.py | 16 +-
.../tvm/backend/trn/tile_primitive/gemm/default.py | 18 +-
.../trn/tile_primitive/instruction_generator.py | 78 +++---
.../backend/trn/tile_primitive/private_alloc.py | 12 +-
.../backend/trn/tile_primitive/reduction/utils.py | 8 +-
.../backend/trn/tile_primitive/select/default.py | 19 +-
.../tvm/backend/trn/tile_primitive/unary/utils.py | 23 +-
.../trn/tile_primitive/unary/with_bias_scale.py | 5 +-
python/tvm/ir/__init__.py | 1 +
python/tvm/ir/_overload_prim_expr.py | 7 +-
python/tvm/ir/expr.py | 32 +++
python/tvm/runtime/script_printer.py | 2 +-
python/tvm/s_tir/analysis/__init__.py | 14 +-
python/tvm/s_tir/dlight/adreno/utils.py | 2 +-
.../tvm/s_tir/dlight/analysis/common_analysis.py | 19 +-
python/tvm/s_tir/dlight/analysis/gemv.py | 2 +-
python/tvm/s_tir/dlight/cpu/reduction.py | 2 +-
python/tvm/s_tir/dlight/gpu/gemv.py | 2 +-
python/tvm/s_tir/dlight/gpu/general_reduction.py | 8 +-
python/tvm/s_tir/dlight/gpu/low_batch_gemv.py | 4 +-
python/tvm/s_tir/dlight/gpu/matmul.py | 4 +-
python/tvm/s_tir/schedule/schedule.py | 6 +-
python/tvm/tirx/buffer.py | 4 +-
python/tvm/tirx/op.py | 6 +-
python/tvm/tirx/operator/tile_primitive/ops.py | 2 +-
python/tvm/tirx/script/builder/ir.py | 26 +-
python/tvm/tirx/script/builder/tirx.py | 309 +++++++++++----------
python/tvm/tirx/script/parser/parser.py | 13 +-
python/tvm/tirx/script/tile.py | 9 +-
python/tvm/tirx/stmt.py | 38 ++-
src/ir/expr.cc | 74 +++++
src/ir/expr_functor.cc | 30 ++
src/relax/analysis/layout_transformation.cc | 29 +-
src/relax/analysis/tir_op_pattern_kind.cc | 4 +-
.../transform/lower_global_view_to_local_view.cc | 6 +-
src/relax/transform/fuse_tir.cc | 15 +-
src/relax/transform/split_call_tir_by_pattern.cc | 19 +-
.../transform/split_layout_rewrite_preproc.cc | 5 +-
src/s_tir/analysis/identify_memcpy.cc | 5 +-
.../analysis/sblock_access_region_detector.cc | 65 ++---
.../analysis/sblock_buffer_access_lca_detector.cc | 12 +-
src/s_tir/meta_schedule/postproc/rewrite_layout.cc | 9 +-
.../postproc/rewrite_parallel_vectorize_unroll.cc | 7 +-
src/s_tir/meta_schedule/schedule/cuda/winograd.cc | 8 +-
.../meta_schedule/schedule_rule/auto_inline.cc | 6 +-
.../schedule_rule/multi_level_tiling.cc | 7 +-
.../multi_level_tiling_tensor_core.cc | 7 +-
.../multi_level_tiling_wide_vector.cc | 2 +-
src/s_tir/sblock_scope.cc | 22 +-
src/s_tir/schedule/analysis.h | 8 +-
src/s_tir/schedule/analysis/analysis.cc | 41 +--
src/s_tir/schedule/analysis/reducer.cc | 10 +-
src/s_tir/schedule/ir_comparator.cc | 15 +-
src/s_tir/schedule/ir_comparator.h | 2 +-
.../schedule/primitive/annotate_buffer_access.cc | 8 +-
src/s_tir/schedule/primitive/block_annotate.cc | 8 +-
src/s_tir/schedule/primitive/blockize_tensorize.cc | 34 +--
src/s_tir/schedule/primitive/cache_read_write.cc | 156 ++++++-----
src/s_tir/schedule/primitive/compute_at.cc | 15 +-
src/s_tir/schedule/primitive/compute_inline.cc | 82 +++---
src/s_tir/schedule/primitive/decompose_padding.cc | 21 +-
src/s_tir/schedule/primitive/hide_buffer_access.cc | 2 +-
.../schedule/primitive/layout_transformation.cc | 22 +-
.../schedule/primitive/loop_transformation.cc | 22 +-
src/s_tir/schedule/primitive/pad_einsum.cc | 34 +--
src/s_tir/schedule/primitive/read_write_at.cc | 12 +-
src/s_tir/schedule/primitive/reduction.cc | 34 ++-
src/s_tir/schedule/primitive/rolling_buffer.cc | 37 +--
src/s_tir/schedule/state.cc | 43 +--
src/s_tir/schedule/transform.cc | 81 +++---
src/s_tir/schedule/transform.h | 12 +-
src/s_tir/transform/compact_buffer_region.cc | 60 ++--
src/s_tir/transform/default_gpu_schedule.cc | 4 +-
src/s_tir/transform/inject_double_buffer.cc | 10 +-
src/s_tir/transform/inject_software_pipeline.cc | 86 +++---
src/s_tir/transform/inject_virtual_thread.cc | 7 +-
src/s_tir/transform/lower_async_dma.cc | 12 +-
.../transform/lower_cross_thread_reduction.cc | 30 +-
src/s_tir/transform/lower_match_buffer.cc | 28 +-
.../manifest_shared_memory_local_stage.cc | 2 +-
src/s_tir/transform/memhammer_coalesce.cc | 17 +-
src/s_tir/transform/memhammer_lower_auto_copy.cc | 27 +-
src/s_tir/transform/memhammer_rewrite_rule.h | 36 +--
.../transform/memhammer_tensorcore_rewrite.cc | 4 +-
.../plan_update_buffer_allocation_location.cc | 20 +-
src/s_tir/transform/renew_defs.cc | 13 +-
src/tirx/analysis/verify_well_formed.cc | 4 +-
src/tirx/ir/data_type_rewriter.cc | 21 +-
src/tirx/ir/data_type_rewriter.h | 2 +-
src/tirx/ir/script/script_complete.cc | 6 +-
src/tirx/ir/specialize.cc | 13 +-
src/tirx/ir/stmt.cc | 112 +++-----
src/tirx/ir/stmt_functor.cc | 41 +--
src/tirx/ir/tir_visitor_with_path.cc | 14 +-
src/tirx/ir/tir_visitor_with_path.h | 6 +-
src/tirx/ir/tirx_stmt.cc | 8 +
src/tirx/script/builder/frame.cc | 7 +-
src/tirx/script/builder/ir.cc | 12 +-
src/tirx/script/builder/utils.h | 6 +-
src/tirx/script/printer/buffer.cc | 8 +-
src/tirx/script/printer/stmt.cc | 8 +-
src/tirx/transform/flatten_buffer.cc | 14 +-
src/tirx/transform/ir_utils.cc | 20 +-
src/tirx/transform/vectorize_loop.cc | 2 +-
tests/cpp/ir_functor_test.cc | 4 +-
.../test_s_tir_analysis_identify_memcpy.py | 5 +-
.../s_tir/analysis/test_sblock_access_region.py | 2 +-
.../schedule/test_tir_schedule_transform_layout.py | 6 +-
tests/python/tirx-base/test_tir_constructor.py | 7 +-
tests/python/tirx/test_hint.py | 10 +-
tests/python/tirx/test_parser_printer.py | 19 +-
157 files changed, 1610 insertions(+), 1376 deletions(-)
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index c02b804bab..86cfd41492 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -664,6 +664,36 @@ class Range : public ffi::ObjectRef {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Range, ffi::ObjectRef, RangeNode);
};
+/*! \brief A region of an indexed expression source. */
+class TensorRegionNode : public ExprNode {
+ public:
+ /*! \brief The indexed source expression. */
+ Expr source;
+ /*! \brief The ranges selected from the source. */
+ ffi::Array<Range> region;
+
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ // A region's source and symbolic source type form a definition pattern;
+ // ranges then refer to those same identities.
+ refl::ObjectDef<TensorRegionNode>()
+ .def_ro("source", &TensorRegionNode::source, refl::AttachFieldFlag::SEqHashDefPattern())
+ .def_ro("region", &TensorRegionNode::region);
+ }
+
+ static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.TensorRegion", TensorRegionNode, ExprNode);
+};
+
+/*! \brief Managed reference to TensorRegionNode. */
+class TensorRegion : public Expr {
+ public:
+ TVM_DLL TensorRegion(Expr source, ffi::Array<Range> region, Type ty, Span span = Span());
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TensorRegion, Expr, TensorRegionNode);
+ TVM_DEFINE_OBJECT_REF_COW_METHOD(TensorRegionNode);
+};
+
/*!
* \brief Analyze the runtime effects of an expression.
*
diff --git a/include/tvm/ir/expr_functor.h b/include/tvm/ir/expr_functor.h
index c4ce33a3ea..7af79f7b6c 100644
--- a/include/tvm/ir/expr_functor.h
+++ b/include/tvm/ir/expr_functor.h
@@ -82,6 +82,9 @@ class ExprFunctor<R(const Expr&, Args...)> {
virtual R Dispatch_(const TupleGetItemNode* node, Args... args) {
return DispatchDefault_(node, std::forward<Args>(args)...);
}
+ virtual R Dispatch_(const TensorRegionNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
virtual R Dispatch_(const TensorLoadNode* node, Args... args) {
return DispatchDefault_(node, std::forward<Args>(args)...);
}
@@ -207,6 +210,7 @@ class ExprFunctor<R(const Expr&, Args...)> {
SetDispatch<TSelf, TupleNode>(vtable);
SetDispatch<TSelf, TupleGetItemNode>(vtable);
SetDispatch<TSelf, TensorLoadNode>(vtable);
+ SetDispatch<TSelf, TensorRegionNode>(vtable);
SetDispatch<TSelf, VarNode>(vtable);
SetDispatch<TSelf, GlobalVarNode>(vtable);
SetDispatch<TSelf, CallNode>(vtable);
@@ -306,6 +310,7 @@ class TVM_DLL ExprVisitor : public ObjectVisitor {
virtual ffi::Optional<VisitInterrupt> Visit_(const TupleNode* node);
virtual ffi::Optional<VisitInterrupt> Visit_(const TupleGetItemNode* node);
virtual ffi::Optional<VisitInterrupt> Visit_(const TensorLoadNode* node);
+ virtual ffi::Optional<VisitInterrupt> Visit_(const TensorRegionNode* node);
virtual ffi::Optional<VisitInterrupt> Visit_(const VarNode* node);
virtual ffi::Optional<VisitInterrupt> Visit_(const GlobalVarNode* node);
virtual ffi::Optional<VisitInterrupt> Visit_(const CallNode* node);
@@ -423,6 +428,7 @@ class TVM_DLL ExprMutator : public ObjectMutator {
virtual UnchangedOr<Expr> Mutate_(const TupleNode* node, InplaceMode inplace_mode);
virtual UnchangedOr<Expr> Mutate_(const TupleGetItemNode* node, InplaceMode inplace_mode);
virtual UnchangedOr<PrimExpr> Mutate_(const TensorLoadNode* node, InplaceMode inplace_mode);
+ virtual UnchangedOr<Expr> Mutate_(const TensorRegionNode* node, InplaceMode inplace_mode);
virtual UnchangedOr<Expr> Mutate_(const VarNode* node, InplaceMode inplace_mode);
virtual UnchangedOr<Expr> Mutate_(const GlobalVarNode* node, InplaceMode inplace_mode);
virtual UnchangedOr<Expr> Mutate_(const CallNode* node, InplaceMode inplace_mode);
diff --git a/include/tvm/s_tir/analysis.h b/include/tvm/s_tir/analysis.h
index 9c5640900f..4a7c72d386 100644
--- a/include/tvm/s_tir/analysis.h
+++ b/include/tvm/s_tir/analysis.h
@@ -42,12 +42,12 @@ namespace tirx {
* \param buffer_var_map The outside buffers which may be accessed the block.
* It is a map from buffer var to the buffer.
* \return Array of access regions.
- * There are three arrays of BufferRegion:
+ * There are three arrays of TensorRegion:
* - first: read regions
* - second: write regions
* - third: opaque regions
*/
-TVM_DLL ffi::Array<ffi::Array<BufferRegion>> GetSBlockAccessRegion(
+TVM_DLL ffi::Array<ffi::Array<TensorRegion>> GetSBlockAccessRegion(
const SBlock& block, const ffi::Map<Var, BufferVar>& buffer_var_map);
/*!
@@ -58,7 +58,7 @@ TVM_DLL ffi::Array<ffi::Array<BufferRegion>> GetSBlockAccessRegion(
* It is a map from buffer var to the buffer
* \return An array only consisting of the read regions and write regions of the input block
*/
-TVM_DLL ffi::Array<ffi::Array<BufferRegion>> GetSBlockReadWriteRegion(
+TVM_DLL ffi::Array<ffi::Array<TensorRegion>> GetSBlockReadWriteRegion(
const SBlock& block, const ffi::Map<Var, BufferVar>& buffer_var_map);
/*!
@@ -129,8 +129,8 @@ TVM_DLL bool VerifyGPUCode(const PrimFunc& func, ffi::Map<ffi::String, PrimExpr>
/*! \brief Helper struct for return value of IdentifyMemCpy */
struct MemCpyDetails {
- BufferRegion source;
- BufferRegion dest;
+ TensorRegion source;
+ TensorRegion dest;
};
/*! \brief Identify whether a For loop is semantically equivalent to MemCpy
diff --git a/include/tvm/tirx/buffer_region.h b/include/tvm/tirx/buffer_region.h
index d4bc55abd5..edd70443ca 100644
--- a/include/tvm/tirx/buffer_region.h
+++ b/include/tvm/tirx/buffer_region.h
@@ -45,34 +45,12 @@ class BufferRegionType : public Type {
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::SEqHashDefPattern())
- .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);
-};
+/*! \brief Construct a region with buffer rank validation and BufferRegionType. */
+TVM_DLL TensorRegion BufferRegion(BufferVar buffer, ffi::Array<Range> region, Span span = Span());
+/*! \brief Select the entire buffer. */
+TVM_DLL TensorRegion FullBufferRegion(BufferVar buffer);
+/*! \brief Construct unit or vector-lane ranges from point indices. */
+TVM_DLL TensorRegion BufferRegionFromPoint(BufferVar buffer, ffi::Array<PrimExpr> indices);
} // namespace tirx
} // namespace tvm
diff --git a/include/tvm/tirx/expr_functor.h b/include/tvm/tirx/expr_functor.h
index 93aaba02d5..7b09c6d73a 100644
--- a/include/tvm/tirx/expr_functor.h
+++ b/include/tvm/tirx/expr_functor.h
@@ -34,7 +34,7 @@ namespace tvm {
namespace tirx {
/*!
- * \brief Shared expression dispatch extended with the TIRx BufferRegion hook.
+ * \brief Dialect entry point for shared expression dispatch.
*
* Override Dispatch_ for a node type or DispatchDefault_ for default behavior.
* Core expression hooks, result types, argument forwarding and registered-ancestor
@@ -49,21 +49,16 @@ class ExprFunctor;
template <typename R, typename... Args>
class ExprFunctor<R(const Expr&, Args...)> : public tvm::ExprFunctor<R(const Expr&, Args...)> {
private:
- using TSelf = ExprFunctor<R(const Expr&, Args...)>;
using Parent = tvm::ExprFunctor<R(const Expr&, Args...)>;
public:
using Parent::Dispatch_;
- /*! \brief Construct a functor with the inherited core and TIRx hooks. */
+ /*! \brief Construct a functor with the inherited expression hooks. */
ExprFunctor() : Parent(GlobalVTable()) {}
/*! \brief Destroy through the dialect functor base. */
virtual ~ExprFunctor() = default;
- virtual R Dispatch_(const BufferRegionNode* node, Args... args) {
- return this->DispatchDefault_(node, std::forward<Args>(args)...);
- }
-
protected:
using Parent::SetDispatch;
using typename Parent::VTable;
@@ -71,11 +66,8 @@ class ExprFunctor<R(const Expr&, Args...)> : public tvm::ExprFunctor<R(const Exp
/*! \brief Construct with a finalized table that outlives the functor. */
explicit ExprFunctor(const VTable* vtable) : Parent(vtable) {}
- /*! \brief Initialize inherited dispatch and add the dialect-only hook. */
- static void InitVTable(VTable* vtable) {
- Parent::InitVTable(vtable);
- Parent::template SetDispatch<TSelf, BufferRegionNode>(vtable);
- }
+ /*! \brief Initialize the inherited expression dispatch. */
+ static void InitVTable(VTable* vtable) { Parent::InitVTable(vtable); }
private:
static const VTable* GlobalVTable() {
diff --git a/include/tvm/tirx/script/builder/frame.h b/include/tvm/tirx/script/builder/frame.h
index 3bf3a1a215..38259a1a57 100644
--- a/include/tvm/tirx/script/builder/frame.h
+++ b/include/tvm/tirx/script/builder/frame.h
@@ -144,9 +144,9 @@ class SBlockFrameNode : public TIRFrameNode {
/*! \brief The variables of the block. */
ffi::Array<tvm::tirx::IterVar> iter_vars;
/*! \brief The read buffer regions of the block. */
- ffi::Optional<ffi::Array<tvm::tirx::BufferRegion>> reads;
+ ffi::Optional<ffi::Array<tvm::TensorRegion>> reads;
/*! \brief The write buffer regions of the block. */
- ffi::Optional<ffi::Array<tvm::tirx::BufferRegion>> writes;
+ ffi::Optional<ffi::Array<tvm::TensorRegion>> writes;
/*! \brief The init statement of the bolck. */
ffi::Optional<tvm::tirx::Stmt> init;
/*! \brief The buffer allocated in the block. */
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 4c0e7fd904..d53c58772e 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -781,7 +781,7 @@ class MatchBufferRegionNode : public ffi::Object {
/*! \brief The target buffer. */
BufferVar buffer;
/*! \brief The source buffer region. */
- BufferRegion source;
+ TensorRegion source;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -800,7 +800,7 @@ class MatchBufferRegionNode : public ffi::Object {
*/
class MatchBufferRegion : public ffi::ObjectRef {
public:
- TVM_DLL explicit MatchBufferRegion(BufferVar buffer, BufferRegion source);
+ TVM_DLL explicit MatchBufferRegion(BufferVar buffer, TensorRegion source);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MatchBufferRegion, ffi::ObjectRef,
MatchBufferRegionNode);
@@ -833,9 +833,9 @@ class SBlockNode : public StmtNode {
/*! \brief The variables of the block. */
ffi::Array<IterVar> iter_vars;
/*! \brief The read buffer regions of the block. */
- ffi::Array<BufferRegion> reads;
+ ffi::Array<TensorRegion> reads;
/*! \brief The write buffer regions of the block. */
- ffi::Array<BufferRegion> writes;
+ ffi::Array<TensorRegion> writes;
/*! \brief The name_hint of the block. */
ffi::String name_hint;
/*! \brief The buffer allocated in the block. */
@@ -879,8 +879,8 @@ class SBlockNode : public StmtNode {
class SBlock : public Stmt {
public:
TVM_DLL explicit SBlock(
- ffi::Array<IterVar> iter_vars, ffi::Array<BufferRegion> reads,
- ffi::Array<BufferRegion> writes, ffi::String name_hint, Stmt body,
+ ffi::Array<IterVar> iter_vars, ffi::Array<TensorRegion> reads,
+ ffi::Array<TensorRegion> writes, ffi::String name_hint, Stmt body,
ffi::Optional<Stmt> init = std::nullopt,
ffi::Array<BufferVar> alloc_buffers = ffi::Array<BufferVar>(),
ffi::Array<MatchBufferRegion> match_buffers = ffi::Array<MatchBufferRegion>(),
diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h
index 1209bbd5a1..61c17cf574 100644
--- a/include/tvm/tirx/stmt_functor.h
+++ b/include/tvm/tirx/stmt_functor.h
@@ -176,7 +176,6 @@ class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor {
virtual ffi::Optional<VisitInterrupt> Visit_(const SBlockRealizeNode* op);
virtual ffi::Optional<VisitInterrupt> Visit_(const ScopeIdDefStmtNode* op);
virtual ffi::Optional<VisitInterrupt> Visit_(const TilePrimitiveCallNode* op);
- virtual ffi::Optional<VisitInterrupt> Visit_(const BufferRegionNode* op);
// Preserve TIRx operand traversal where it differs from the shared defaults.
ffi::Optional<VisitInterrupt> Visit_(const VarNode* op) override;
@@ -204,7 +203,7 @@ class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor {
* Base hooks preserve stored types and derived fields. They rewrite structural
* children without re-inferring types or repeating constructor validation,
* including on in-place writes. Passes that change dtypes or index lanes must
- * provide the corresponding TensorLoad, BufferStore, or BufferRegion inference.
+ * provide the corresponding TensorLoad, BufferStore, or TensorRegion inference.
*/
class TVM_DLL StmtExprMutator : public tvm::ExprMutator {
public:
@@ -244,8 +243,6 @@ class TVM_DLL StmtExprMutator : public tvm::ExprMutator {
virtual UnchangedOr<Stmt> Mutate_(const ScopeIdDefStmtNode* op, InplaceMode inplace_mode);
virtual UnchangedOr<Stmt> Mutate_(const TilePrimitiveCallNode* op, InplaceMode inplace_mode);
- virtual UnchangedOr<Expr> Mutate_(const BufferRegionNode* op, InplaceMode inplace_mode);
-
protected:
explicit StmtExprMutator(const VTable* vtable) : tvm::ExprMutator(vtable) {}
static void InitVTable(VTable* vtable);
diff --git a/include/tvm/tirx/tile_primitive.h b/include/tvm/tirx/tile_primitive.h
index 91781a23d4..b0158a25e7 100644
--- a/include/tvm/tirx/tile_primitive.h
+++ b/include/tvm/tirx/tile_primitive.h
@@ -253,28 +253,28 @@ class TilePrimitiveCall : public Stmt {
/*!
* \brief See pesudo code below:
*
- * Tx.cast(BufferRegion dst, BufferRegion src)
+ * Tx.cast(TensorRegion dst, TensorRegion src)
*/
TVM_DLL const Op& cast();
/*!
* \brief See pesudo code below:
*
- * Tx.copy(BufferRegion dst, BufferRegion src)
+ * Tx.copy(TensorRegion dst, TensorRegion src)
*/
TVM_DLL const Op& copy();
/*!
* \brief See pesudo code below:
*
- * Tx.Async.copy(BufferRegion dst, BufferRegion src)
+ * Tx.Async.copy(TensorRegion dst, TensorRegion src)
*/
TVM_DLL const Op& copy_async();
/*!
* \brief See pesudo code below:
*
- * Tx.fill(BufferRegion dst, PrimExpr value)
+ * Tx.fill(TensorRegion dst, PrimExpr value)
*/
TVM_DLL const Op& fill();
@@ -288,7 +288,7 @@ TVM_DLL const Op& gemm();
/*!
* \brief See pesudo code below:
*
- * Tx.gemm_async(BufferRegion C, BufferRegion A, BufferRegion B, bool transA, bool transB,
+ * Tx.gemm_async(TensorRegion C, TensorRegion A, TensorRegion B, bool transA, bool transB,
* bool accum)
*/
TVM_DLL const Op& gemm_async();
diff --git a/python/tvm/backend/cuda/tile_primitive/common.py b/python/tvm/backend/cuda/tile_primitive/common.py
index d229809631..25af4e5453 100644
--- a/python/tvm/backend/cuda/tile_primitive/common.py
+++ b/python/tvm/backend/cuda/tile_primitive/common.py
@@ -23,9 +23,10 @@ import re
from enum import Enum
from tvm.arith.analyzer import Analyzer
+from tvm.ir import TensorRegion
from tvm.runtime import DataType
from tvm.script import tirx as T
-from tvm.tirx import Buffer, BufferRegion, PrimFunc
+from tvm.tirx import Buffer, PrimFunc
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
from tvm.tirx.tile_primitive import TilePrimitiveCall
@@ -37,7 +38,7 @@ def next_power_of_2(x: int) -> int:
return 1 << (x - 1).bit_length()
-def get_st_extent(buffer_region: BufferRegion):
+def get_st_extent(buffer_region: TensorRegion):
"""Get the start and extent of a buffer region."""
region = buffer_region.region
return [r.min for r in region], [r.extent for r in region]
@@ -86,8 +87,8 @@ def validate_copy_op(
) -> bool:
"""Sanity check for copy op"""
dst_buffer_region, src_buffer_region = op_call.args[:2]
- src: Buffer = src_buffer_region.buffer
- dst: Buffer = dst_buffer_region.buffer
+ src: Buffer = src_buffer_region.source
+ dst: Buffer = dst_buffer_region.source
if not (src.layout and dst.layout and src.dtype == dst.dtype):
return False
# Extract regions and validate dimensions
@@ -104,15 +105,15 @@ def validate_copy_op(
def get_vec_len(
- dst_buffer_region: BufferRegion,
- src_buffer_region: BufferRegion,
+ dst_buffer_region: TensorRegion,
+ src_buffer_region: TensorRegion,
vec_candidates: list[int],
thread_cnt=1,
) -> int | None:
"""Get the vector length for the copy operation."""
- dst: Buffer = dst_buffer_region.buffer
- src: Buffer = src_buffer_region.buffer
+ dst: Buffer = dst_buffer_region.source
+ src: Buffer = src_buffer_region.source
# layout=None (flat local buffer) is treated as trivial for vectorization purposes
if not (
(dst.layout is None or dst.layout.is_trivial())
@@ -158,8 +159,8 @@ def copy_vec_load_impl(
threads in a CTA/using a single thread.
"""
dst_buffer_region, src_buffer_region = op_call.args[:2]
- src: Buffer = src_buffer_region.buffer
- dst: Buffer = dst_buffer_region.buffer
+ src: Buffer = src_buffer_region.source
+ dst: Buffer = dst_buffer_region.source
if not (
(src.scope() == "global" and dst.scope().startswith("shared"))
or (src.scope().startswith("shared") and dst.scope() == "global")
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/fallback.py b/python/tvm/backend/cuda/tile_primitive/copy/fallback.py
index 1a78fce1d9..9979fedb97 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/fallback.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/fallback.py
@@ -41,8 +41,8 @@ def _region_st_extent(buffer_region):
def _emit_fallback(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
op_call = TilePrimitiveCall.downcast(op_call)
- src: Buffer = op_call.src.buffer
- dst: Buffer = op_call.dst.buffer
+ src: Buffer = op_call.src.source
+ dst: Buffer = op_call.dst.source
src_st, src_extent = _region_st_extent(op_call.src)
dst_st, dst_extent = _region_st_extent(op_call.dst)
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/ld_stmatrix.py b/python/tvm/backend/cuda/tile_primitive/copy/ld_stmatrix.py
index d446a5de76..65e052cafd 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/ld_stmatrix.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/ld_stmatrix.py
@@ -78,14 +78,14 @@ def _emit(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
# Step 1: identify reg / smem sides and pull their tensor shape + layout.
src_br = op_call.src
dst_br = op_call.dst
- if src_br.buffer.scope() == "local":
+ if src_br.source.scope() == "local":
r_br, s_br = src_br, dst_br
direction = "st" # reg -> smem (stmatrix)
else:
r_br, s_br = dst_br, src_br
direction = "ld" # smem -> reg (ldmatrix)
- r_buf = r_br.buffer
- s_buf = s_br.buffer
+ r_buf = r_br.source
+ s_buf = s_br.source
r_shape = list(r_buf.shape)
r_layout = r_buf.layout
s_shape = list(s_buf.shape)
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/utils.py b/python/tvm/backend/cuda/tile_primitive/copy/utils.py
index cccf8b2b99..2164dec7bb 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/utils.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/utils.py
@@ -48,8 +48,8 @@ def _scope_allowed(
):
op_call = TilePrimitiveCall.downcast(op_call)
dst_buffer_region, src_buffer_region = (op_call.dst, op_call.src)
- src_scope = src_buffer_region.buffer.scope()
- dst_scope = dst_buffer_region.buffer.scope()
+ src_scope = src_buffer_region.source.scope()
+ dst_scope = dst_buffer_region.source.scope()
ok = any(
(
match_scope(src_scope, src_pat) and match_scope(dst_scope, dst_pat)
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_gmem_smem.py b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_gmem_smem.py
index 73074573d5..51d74f4b96 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_gmem_smem.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_gmem_smem.py
@@ -63,7 +63,7 @@ def _divides_thread_cnt(
thread_cnt = _thread_cnt(sctx)
if thread_cnt <= 0:
return False, f"degenerate thread_cnt={thread_cnt} (scope has empty intra)"
- g_br = op_call.src if op_call.src.buffer.scope() == "global" else op_call.dst
+ g_br = op_call.src if op_call.src.source.scope() == "global" else op_call.dst
n_elements = 1
for r in g_br.region:
ext = r.extent
@@ -95,8 +95,8 @@ def _is_gmem_smem(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bo
def _emit_gmem_smem(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
op_call = TilePrimitiveCall.downcast(op_call)
- src: Buffer = op_call.src.buffer
- dst: Buffer = op_call.dst.buffer
+ src: Buffer = op_call.src.source
+ dst: Buffer = op_call.dst.source
if src.scope() == "global":
g_buf, g_br, s_buf, s_br = src, op_call.src, dst, op_call.dst
g_is_src = True
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py
index 0f6fd138b9..5648045be0 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py
@@ -94,8 +94,8 @@ def _r_side_layout_valid(
op_call: TilePrimitiveCall, sctx: DispatchContext
) -> tuple[bool, str | None]:
op_call = TilePrimitiveCall.downcast(op_call)
- src: Buffer = op_call.src.buffer
- dst: Buffer = op_call.dst.buffer
+ src: Buffer = op_call.src.source
+ dst: Buffer = op_call.dst.source
r_buf = src if src.scope() == "local" else dst
layout = r_buf.layout
if layout is None:
@@ -138,8 +138,8 @@ def _s_side_slice_ok(op_call: TilePrimitiveCall) -> tuple[bool, str | None]:
op_call = TilePrimitiveCall.downcast(op_call)
src_br = op_call.src
dst_br = op_call.dst
- s_br = dst_br if src_br.buffer.scope() == "local" else src_br
- s_buf: Buffer = s_br.buffer
+ s_br = dst_br if src_br.source.scope() == "local" else src_br
+ s_buf: Buffer = s_br.source
layout = s_buf.layout
if layout is None:
return False, "S has no layout"
@@ -270,12 +270,12 @@ def _align_layouts(op_call: TilePrimitiveCall, sctx: DispatchContext):
op_call = TilePrimitiveCall.downcast(op_call)
src_br = op_call.src
dst_br = op_call.dst
- if src_br.buffer.scope() == "local":
+ if src_br.source.scope() == "local":
r_br, s_br = src_br, dst_br
else:
r_br, s_br = dst_br, src_br
- r_buf = r_br.buffer
- s_buf = s_br.buffer
+ r_buf = r_br.source
+ s_buf = s_br.source
r_region = [(r.min, r.min + r.extent) for r in r_br.region]
s_region = [(r.min, r.min + r.extent) for r in s_br.region]
with sctx.target:
@@ -525,8 +525,8 @@ def _outer_const_offsets(outer_atoms, flat_idx: int) -> tuple[int, int]:
def _emit_reg(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
op_call = TilePrimitiveCall.downcast(op_call)
- src: Buffer = op_call.src.buffer
- dst: Buffer = op_call.dst.buffer
+ src: Buffer = op_call.src.source
+ dst: Buffer = op_call.dst.source
if src.scope() == "local":
r_buf, s_buf, r_is_src = src, dst, True
else:
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py b/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py
index afd743ece1..fcb9f2f794 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py
@@ -33,23 +33,23 @@ Example::
"""
from tvm.arith.analyzer import Analyzer
+from tvm.ir import TensorRegion
from tvm.runtime import DataType
from tvm.script import tirx as T
from tvm.tirx import Buffer, PrimFunc
from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
from tvm.tirx.operator.tile_primitive.registry import DispatchContext
-from tvm.tirx.stmt import BufferRegion
from tvm.tirx.tile_primitive import TilePrimitiveCall
from ._common import copy_ptx_form, copy_ptx_ld_chain
from .utils import _scope_allowed
-def _region_start(buffer_region: BufferRegion):
+def _region_start(buffer_region: TensorRegion):
return [r.min for r in buffer_region.region]
-def _region_elements(buffer_region: BufferRegion):
+def _region_elements(buffer_region: TensorRegion):
product = 1
for r in buffer_region.region:
product *= r.extent
@@ -106,8 +106,8 @@ def _is_forced_vec_copy(
return False, scope_reason
op_call = TilePrimitiveCall.downcast(op_call)
- src: Buffer = op_call.src.buffer
- dst: Buffer = op_call.dst.buffer
+ src: Buffer = op_call.src.source
+ dst: Buffer = op_call.dst.source
if src.dtype != dst.dtype:
return False, f"dtype mismatch: src={src.dtype}, dst={dst.dtype}"
@@ -137,8 +137,8 @@ def _is_forced_vec_copy(
def _emit_forced_vec_copy(op_call: TilePrimitiveCall, _sctx: DispatchContext, num_bytes: int):
op_call = TilePrimitiveCall.downcast(op_call)
- src: Buffer = op_call.src.buffer
- dst: Buffer = op_call.dst.buffer
+ src: Buffer = op_call.src.source
+ dst: Buffer = op_call.dst.source
src_scope = src.scope()
dst_scope = dst.scope()
src_is_local = src_scope == "local"
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py b/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
index 62ec3261e9..35a0d2570d 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
@@ -39,8 +39,8 @@ from .utils import find_contiguous_region, to_tile_layout
def _is_shared_to_shared(op_call: TilePrimitiveCall) -> bool:
"""Check if both src and dst are in shared memory."""
op_call = TilePrimitiveCall.downcast(op_call)
- src_scope = op_call.src.buffer.scope()
- dst_scope = op_call.dst.buffer.scope()
+ src_scope = op_call.src.source.scope()
+ dst_scope = op_call.dst.source.scope()
return src_scope.startswith("shared") and dst_scope.startswith("shared")
@@ -68,8 +68,8 @@ def copy_dsmem_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFu
# Extract buffer regions
dst_buffer_region = op_call.dst
src_buffer_region = op_call.src
- src_buf: Buffer = src_buffer_region.buffer
- dst_buf: Buffer = dst_buffer_region.buffer
+ src_buf: Buffer = src_buffer_region.source
+ dst_buf: Buffer = dst_buffer_region.source
src_st = [r.min for r in src_buffer_region.region]
src_ext = [r.extent for r in src_buffer_region.region]
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/ldgsts.py b/python/tvm/backend/cuda/tile_primitive/copy_async/ldgsts.py
index e3845f8cfd..42ada04cce 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/ldgsts.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/ldgsts.py
@@ -110,7 +110,7 @@ def _divides_thread_cnt_ldgsts(
thread_cnt = _thread_cnt(sctx)
if thread_cnt <= 0:
return False, f"degenerate thread_cnt={thread_cnt} (scope has empty intra)"
- g_br = op_call.src if op_call.src.buffer.scope() == "global" else op_call.dst
+ g_br = op_call.src if op_call.src.source.scope() == "global" else op_call.dst
n_elements = 1
for r in g_br.region:
ext = r.extent
@@ -142,8 +142,8 @@ def _is_ldgsts(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool,
def _emit_ldgsts(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
op_call = TilePrimitiveCall.downcast(op_call)
- src: Buffer = op_call.src.buffer
- dst: Buffer = op_call.dst.buffer
+ src: Buffer = op_call.src.source
+ dst: Buffer = op_call.dst.source
# Predicate above guarantees src is global, dst is shared.
g_buf, g_br = src, op_call.src
s_buf, s_br = dst, op_call.dst
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
index f385fbff26..578e2dcf41 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
@@ -349,8 +349,8 @@ def _build_plan(op_call: TilePrimitiveCall):
def _plan_for_shape(op_call: TilePrimitiveCall, shape: str, multicast: str):
"""Run A..I for one (shape, multicast); raises ValueError on any mismatch."""
dst_region, src_region = op_call.args[:2]
- s_buf: Buffer = src_region.buffer
- t_buf: Buffer = dst_region.buffer
+ s_buf: Buffer = src_region.source
+ t_buf: Buffer = dst_region.source
dtype = s_buf.dtype
dtype_bits = DataType(dtype).bits
elem_per_128b = 128 // dtype_bits # elements per 16B descriptor unit
@@ -702,8 +702,8 @@ def _validate_smem_tmem_copy(op_call: TilePrimitiveCall, sctx: DispatchContext):
"""Memory-scope envelope only; shape resolution/inference and the detailed
layout validation raise readable ValueErrors in ``_build_plan``."""
dst_region, src_region = op_call.args[:2]
- src: Buffer = src_region.buffer
- dst: Buffer = dst_region.buffer
+ src: Buffer = src_region.source
+ dst: Buffer = dst_region.source
return (
src.scope().startswith("shared")
and dst.scope() == "tmem"
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_ldst.py b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_ldst.py
index b3e7d7e3fb..ff23c11806 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_ldst.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_ldst.py
@@ -225,8 +225,8 @@ def _tmem_window(tmem_buf, tmem_region, atom_kind, frag_rows, analyzer):
def copy_tmem_local_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
op_call = TilePrimitiveCall.downcast(op_call)
dst_buffer_region, src_buffer_region = op_call.dst, op_call.src
- dst: Buffer = dst_buffer_region.buffer
- src: Buffer = src_buffer_region.buffer
+ dst: Buffer = dst_buffer_region.source
+ src: Buffer = src_buffer_region.source
if src.scope() == "tmem" and dst.scope() == "local":
direction = "tmem2local"
@@ -237,7 +237,7 @@ def copy_tmem_local_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> P
else:
raise ValueError(f"Unsupported src scope {src.scope()} and dst scope {dst.scope()}")
- tmem_buf, local_buf = tmem_region.buffer, local_region.buffer
+ tmem_buf, local_buf = tmem_region.source, local_region.source
assert tmem_buf.layout is not None
assert local_buf.layout is not None
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py b/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
index 39a3d84473..e5b5aa677c 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
@@ -1620,8 +1620,8 @@ def _runtime_config(op_call, sctx, direction: str, *, explicit: bool):
def _copy_direction(op_call):
op_call = TilePrimitiveCall.downcast(op_call)
dst_region, src_region = op_call.dst, op_call.src
- src_scope = src_region.buffer.scope()
- dst_scope = dst_region.buffer.scope()
+ src_scope = src_region.source.scope()
+ dst_scope = dst_region.source.scope()
if src_scope == "global" and dst_scope.startswith("shared"):
return "g2s", dst_region, src_region
if src_scope.startswith("shared") and dst_scope == "global":
@@ -1631,8 +1631,8 @@ def _copy_direction(op_call):
def _build_auto_plan(op_call: TilePrimitiveCall, sctx: DispatchContext) -> TMAPlan:
direction, shared_region, global_region = _copy_direction(op_call)
- s_buf = shared_region.buffer
- g_buf = global_region.buffer
+ s_buf = shared_region.source
+ g_buf = global_region.source
if str(s_buf.dtype) != str(g_buf.dtype):
_auto_fail(
"dtype",
@@ -1953,8 +1953,8 @@ def _selector_compatibility(main: TensorMapSpec, candidate: TensorMapSpec, index
def _build_explicit_plan(op_call: TilePrimitiveCall, sctx: DispatchContext):
direction, shared_region, global_region = _copy_direction(op_call)
- s_buf = shared_region.buffer
- g_buf = global_region.buffer
+ s_buf = shared_region.source
+ g_buf = global_region.source
runtime = _runtime_config(op_call, sctx, direction, explicit=True)
gather4 = _normalize_gather4(op_call.config.get("gather4"))
selectors = _normalize_src_selector(op_call.config.get("src_selector"))
@@ -2244,8 +2244,8 @@ def copy_tma_explicit_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) ->
def _validate_tma_copy_op(op_call: TilePrimitiveCall, _sctx: DispatchContext) -> bool:
dst_region, src_region = op_call.args[:2]
- src = src_region.buffer
- dst = dst_region.buffer
+ src = src_region.source
+ dst = dst_region.source
if src.layout is None or dst.layout is None:
return False
src_scope, dst_scope = src.scope(), dst.scope()
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/_common.py b/python/tvm/backend/cuda/tile_primitive/elementwise/_common.py
index 16e400728b..5e668ffddb 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/_common.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/_common.py
@@ -31,9 +31,9 @@ import functools
import operator
from tvm.arith.analyzer import Analyzer
+from tvm.ir import TensorRegion
from tvm.runtime import DataType
from tvm.script import tirx as T
-from tvm.tirx import BufferRegion
from tvm.tirx.layout import Axis, Iter, TileLayout
from ..common import get_indices, get_st_extent
@@ -46,9 +46,9 @@ from ..copy.vec_auto_reg import _all_threads_active, _axis_decl, _compute_perm_r
# -----------------------------------------------------------------------------
# Plan helpers
# -----------------------------------------------------------------------------
-def buffer_regions(plan) -> list[BufferRegion]:
- """All BufferRegion args (dst + buffer-region srcs), in plan order."""
- out: list[BufferRegion] = [plan.dst]
+def buffer_regions(plan) -> list[TensorRegion]:
+ """All TensorRegion args (dst + buffer-region srcs), in plan order."""
+ out: list[TensorRegion] = [plan.dst]
for s in plan.srcs:
if s.buf_region is not None:
out.append(s.buf_region)
@@ -68,10 +68,10 @@ def dtype_bits(dtype) -> int:
def compute_dtype_of(plan) -> str:
"""Widest dtype in bits across dst + buffer/scalar srcs (dst breaks ties)."""
- candidates = [dtype_name(plan.dst.buffer.dtype)]
+ candidates = [dtype_name(plan.dst.source.dtype)]
for s in plan.srcs:
if s.buf_region is not None:
- candidates.append(dtype_name(s.buf_region.buffer.dtype))
+ candidates.append(dtype_name(s.buf_region.source.dtype))
elif s.scalar is not None:
candidates.append(scalar_dtype(s.scalar))
widest = candidates[0]
@@ -96,7 +96,7 @@ def scalar_dtype(scalar) -> str:
return str(dtype)
-def n_elements(buf_region: BufferRegion) -> int:
+def n_elements(buf_region: TensorRegion) -> int:
_, ext = get_st_extent(buf_region)
return functools.reduce(operator.mul, ext, 1)
@@ -104,7 +104,7 @@ def n_elements(buf_region: BufferRegion) -> int:
# -----------------------------------------------------------------------------
# Anchor selection (reg.py)
# -----------------------------------------------------------------------------
-def pick_anchor(plan) -> BufferRegion:
+def pick_anchor(plan) -> TensorRegion:
"""Anchor is always ``plan.dst`` — every operand must have a layout
(enforced by predicate); dst's layout drives iteration. No choice to make.
"""
@@ -118,7 +118,7 @@ def _tensor_shape_of(region) -> tuple[int, ...]:
"""Per-dim region extent (post-slice tensor shape, NOT layout shape).
Accepts either ``[(start, end), ...]`` pairs (as built locally from a
- ``BufferRegion``) or the ``BufferRegion.region`` sequence of ``Range``
+ ``TensorRegion``) or the ``TensorRegion.region`` sequence of ``Range``
objects directly. ``Range.extent`` is already simplified by the
front-end, so we avoid computing ``end - start`` on raw Expr (which
yields an un-simplified ``Sub`` and breaks ``int(...)``).
@@ -205,8 +205,8 @@ def preprocess_operand(op_br, anchor_tshape):
Raises ``ValueError`` if the lift is not broadcast-compatible (caller
should have verified via ``shape_broadcast_compat`` in the predicate).
"""
- op_layout = op_br.buffer.layout
- op_shape = op_br.buffer.shape
+ op_layout = op_br.source.layout
+ op_shape = op_br.source.shape
op_region = [(r.min, r.min + r.extent) for r in op_br.region]
sliced = op_layout.slice(list(op_shape), op_region).canonicalize()
sliced = _extract_tile(sliced, op_region)
@@ -286,8 +286,8 @@ def align_operands_to_anchor(anchor_br, layout_others_br):
Uses ``_align_layouts_no_post_canon`` (not copy's ``align_layouts_raw``
directly) so ``anchor_p.shard`` length matches ``op_seps`` groupings.
"""
- anchor_layout = anchor_br.buffer.layout
- anchor_shape = anchor_br.buffer.shape
+ anchor_layout = anchor_br.source.layout
+ anchor_shape = anchor_br.source.shape
anchor_region = [(r.min, r.min + r.extent) for r in anchor_br.region]
per_op_aligned: dict = {}
@@ -301,8 +301,8 @@ def align_operands_to_anchor(anchor_br, layout_others_br):
return anchor_p, per_op_aligned
for op_br in layout_others_br:
- op_layout = op_br.buffer.layout
- op_shape = op_br.buffer.shape
+ op_layout = op_br.source.layout
+ op_shape = op_br.source.shape
op_region = [(r.min, r.min + r.extent) for r in op_br.region]
r_p, op_p, op_seps = _align_layouts_no_post_canon(
anchor_layout,
@@ -371,7 +371,7 @@ def fetch_src_value(src, fused, dst_indices, dst_start, dst_extent):
idx = _broadcast_indices(dst_indices, dst_start, dst_extent, src_st, src_ext)
else:
idx = get_indices(fused, src_st, src_ext)
- return region.buffer[tuple(idx)]
+ return region.source[tuple(idx)]
def emit_scope_sync(scope_kind: str):
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/__init__.py b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/__init__.py
index 4c2cc78d72..163225f635 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/__init__.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/__init__.py
@@ -30,21 +30,22 @@ from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any
+from tvm.ir import TensorRegion
from tvm.ir.expr import Expr
-from tvm.tirx import BufferRegion, TilePrimitiveCall
+from tvm.tirx import TilePrimitiveCall
@dataclass
class SrcSpec:
"""One operand of an elementwise op.
- Either a ``BufferRegion`` (per-element load) or a scalar ``Expr``.
+ Either a ``TensorRegion`` (per-element load) or a scalar ``Expr``.
``index_fn``, if given, derives per-element indices for broadcasting srcs:
``index_fn(dst_indices, dst_start, dst_extent, src_start, src_extent) -> list[Expr]``
Default is the standard ``get_indices`` over the src's own region.
"""
- buf_region: BufferRegion | None = None
+ buf_region: TensorRegion | None = None
scalar: Expr | None = None
index_fn: Callable | None = None
@@ -54,14 +55,14 @@ class SrcSpec:
@property
def buffer(self):
- return self.buf_region.buffer if self.buf_region is not None else None
+ return self.buf_region.source if self.buf_region is not None else None
@dataclass
class Plan:
"""Parsed elementwise op ready for a schedule to consume."""
- dst: BufferRegion
+ dst: TensorRegion
srcs: list[SrcSpec]
extras: dict[str, Any] = field(default_factory=dict)
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/binary.py b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/binary.py
index 2932b121c5..3b943440c5 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/binary.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/binary.py
@@ -33,8 +33,9 @@ import functools
import operator
from typing import Any
+from tvm.ir import TensorRegion
from tvm.script import tirx as Tx
-from tvm.tirx import BufferRegion, TilePrimitiveCall
+from tvm.tirx import TilePrimitiveCall
from ..vec_emit.binary_f32x2 import BINARY_F32X2_IMPLS
from . import OpSpec, Plan, SrcSpec
@@ -46,12 +47,12 @@ def _parse_binary_for(op_name: str):
"""Build a ``parse(op_call) -> (Plan, msg)`` for a specific binary op."""
def parse(op: TilePrimitiveCall) -> tuple[Plan | None, str | None]:
- _dst: BufferRegion = op.args[0]
+ _dst: TensorRegion = op.args[0]
_src1 = op.args[1]
_src2 = op.args[2]
- s1_scalar = not isinstance(_src1, BufferRegion)
- s2_scalar = not isinstance(_src2, BufferRegion)
+ s1_scalar = not isinstance(_src1, TensorRegion)
+ s2_scalar = not isinstance(_src2, TensorRegion)
if s1_scalar and s2_scalar:
return None, "both inputs are constants"
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/cast.py b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/cast.py
index 7ffb0b72ad..641d0fb8e8 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/cast.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/cast.py
@@ -21,16 +21,17 @@ CUDA intrinsics like ``__float22half2_rn``."""
from __future__ import annotations
-from tvm.tirx import BufferRegion, TilePrimitiveCall
+from tvm.ir import TensorRegion
+from tvm.tirx import TilePrimitiveCall
from ..vec_emit.cast_vec2 import CAST_VEC2_IMPL
from . import OpSpec, Plan, SrcSpec
def _parse_cast(op: TilePrimitiveCall) -> tuple[Plan | None, str | None]:
- _dst: BufferRegion = op.args[0]
+ _dst: TensorRegion = op.args[0]
_src = op.args[1]
- if not isinstance(_src, BufferRegion):
+ if not isinstance(_src, TensorRegion):
return None, "cast src must be a buffer region"
return Plan(dst=_dst, srcs=[SrcSpec(buf_region=_src)], extras={}), None
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/fma.py b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/fma.py
index 5b057f9cdd..8f2a6cfc79 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/fma.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/fma.py
@@ -23,18 +23,19 @@ Attaches ``fma_f32x2`` VecImpl for sm_100+ f32; falls back to scalar
from __future__ import annotations
-from tvm.tirx import BufferRegion, TilePrimitiveCall
+from tvm.ir import TensorRegion
+from tvm.tirx import TilePrimitiveCall
from ..vec_emit.fma_f32x2 import FMA_F32X2_IMPL
from . import OpSpec, Plan, SrcSpec
def _parse_fma(op: TilePrimitiveCall) -> tuple[Plan | None, str | None]:
- _dst: BufferRegion = op.args[0]
+ _dst: TensorRegion = op.args[0]
args = op.args[1:4]
srcs: list[SrcSpec] = []
for a in args:
- if isinstance(a, BufferRegion):
+ if isinstance(a, TensorRegion):
srcs.append(SrcSpec(buf_region=a))
else:
srcs.append(SrcSpec(scalar=a))
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/unary.py b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/unary.py
index be4de45257..ed6d921414 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/unary.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/unary.py
@@ -25,9 +25,9 @@ from __future__ import annotations
from typing import Any
-from tvm.ir import is_prim_expr
+from tvm.ir import TensorRegion, is_prim_expr
from tvm.script import tirx as T
-from tvm.tirx import BufferRegion, TilePrimitiveCall
+from tvm.tirx import TilePrimitiveCall
from tvm.tirx.expr import FloatImm
from .._common import scalar_dtype
@@ -36,13 +36,13 @@ from . import OpSpec, Plan, SrcSpec
def _parse_unary(op: TilePrimitiveCall) -> tuple[Plan | None, str | None]:
"""T.<unary>(dst, src[, bias, scale]) → Plan."""
- _dst: BufferRegion = op.args[0]
+ _dst: TensorRegion = op.args[0]
_src = op.args[1]
_bias = op.args[2] if len(op.args) > 2 else None
_scale = op.args[3] if len(op.args) > 2 else None
srcs: list[SrcSpec] = []
- if isinstance(_src, BufferRegion):
+ if isinstance(_src, TensorRegion):
srcs.append(SrcSpec(buf_region=_src))
elif is_prim_expr(_src):
srcs.append(SrcSpec(scalar=_src))
@@ -53,7 +53,7 @@ def _parse_unary(op: TilePrimitiveCall) -> tuple[Plan | None, str | None]:
"scale": _scale,
"bias_const": _bias if isinstance(_bias, FloatImm) else None,
}
- if isinstance(_bias, BufferRegion):
+ if isinstance(_bias, TensorRegion):
srcs.append(SrcSpec(buf_region=_bias))
extras["has_bias_buf"] = True
else:
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py b/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py
index 324662efec..cec6b6831c 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py
@@ -62,7 +62,7 @@ from .vec_emit import _emit_vec
# Predicate
# -----------------------------------------------------------------------------
def _validate_anchor_layout(anchor_br) -> tuple[bool, str | None]:
- layout = anchor_br.buffer.layout
+ layout = anchor_br.source.layout
if layout.is_swizzle():
return False, "anchor layout is swizzle"
if not isinstance(layout, TileLayout):
@@ -87,7 +87,7 @@ def _validate_scope_level_anchor(anchor_br, sctx: DispatchContext) -> tuple[bool
# Canonicalize the sliced anchor with the target so warp/lane axes fuse.
st, ext = get_st_extent(anchor_br)
- sliced = get_sublayout_from_region(anchor_br.buffer.layout, anchor_br.buffer.shape, st, ext)
+ sliced = get_sublayout_from_region(anchor_br.source.layout, anchor_br.source.shape, st, ext)
with sctx.target:
canon = sliced.canonicalize()
shard = getattr(canon, "shard", None)
@@ -151,7 +151,7 @@ def _check_layout_operands_agree(plan, sctx) -> tuple[bool, str | None]:
for br in layout_brs:
st, ext = get_st_extent(br)
with sctx.target:
- sliced = get_sublayout_from_region(br.buffer.layout, br.buffer.shape, st, ext)
+ sliced = get_sublayout_from_region(br.source.layout, br.source.shape, st, ext)
canon = sliced.canonicalize()
sig = layout_signature(canon)
if sig is None:
@@ -187,9 +187,9 @@ def is_reg_ewise(spec):
if msg is not None or plan is None:
return False, msg
for br in buffer_regions(plan):
- if br.buffer.scope() != "local":
- return False, f"operand scope {br.buffer.scope()} != local"
- if br.buffer.layout is None:
+ if br.source.scope() != "local":
+ return False, f"operand scope {br.source.scope()} != local"
+ if br.source.layout is None:
return False, f"operand {br} has no layout"
if spec.check_extras is not None:
ok2, reason2 = spec.check_extras(plan.extras, compute_dtype_of(plan))
@@ -356,8 +356,8 @@ def _make_views_meta(per_op_carved, per_thread_total):
return {
op_br: T.decl_buffer(
(per_thread_total,),
- op_br.buffer.dtype,
- op_br.buffer.data,
+ op_br.source.dtype,
+ op_br.source.data,
scope="local",
layout=per_op_carved[op_br],
)
@@ -412,7 +412,7 @@ def _emit_induced_scalar(
extras = plan.extras
srcs = plan.srcs
dst_br = plan.dst
- dst_dtype = dst_br.buffer.dtype
+ dst_dtype = dst_br.source.dtype
compute = spec.compute_scalar
@T.prim_func(check_well_formed=False)
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/smem.py b/python/tvm/backend/cuda/tile_primitive/elementwise/smem.py
index 6db3fcb44c..1e950935a4 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/smem.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/smem.py
@@ -72,9 +72,9 @@ def is_smem_ewise(spec):
if msg is not None or plan is None:
return False, msg
for br in buffer_regions(plan):
- if not br.buffer.scope().startswith("shared"):
- return False, f"operand scope {br.buffer.scope()} != shared*"
- if br.buffer.layout is None:
+ if not br.source.scope().startswith("shared"):
+ return False, f"operand scope {br.source.scope()} != shared*"
+ if br.source.layout is None:
return False, "shared operand has no layout"
if spec.check_extras is not None:
ok2, reason2 = spec.check_extras(plan.extras, compute_dtype_of(plan))
@@ -101,10 +101,10 @@ def is_smem_ewise(spec):
def _max_layout_vec(plan, total: int, thread_cnt: int) -> int:
"""Widest vec_chunk dividing all operands' innermost extents AND
``total / thread_cnt``, within dtype-bit candidates ``{128,64,32,16,8}``."""
- max_bits = dtype_bits(plan.dst.buffer.dtype)
+ max_bits = dtype_bits(plan.dst.source.dtype)
for s in plan.srcs:
if s.buf_region is not None:
- max_bits = max(max_bits, dtype_bits(s.buf_region.buffer.dtype))
+ max_bits = max(max_bits, dtype_bits(s.buf_region.source.dtype))
per_thread = total // thread_cnt if thread_cnt > 0 else total
if total % thread_cnt != 0:
return 1
@@ -193,7 +193,7 @@ def _src_lane_indices(src_br, dst_lane_indices, dst_st, dst_ext, vec_chunk, fuse
def _emit_packed(plan, vec_impl, vec_chunk, total, thread_cnt, sctx) -> PrimFunc:
extras = plan.extras
srcs = plan.srcs
- dst_buf = plan.dst.buffer
+ dst_buf = plan.dst.source
dst_st, dst_ext = get_st_extent(plan.dst)
sync = emit_scope_sync(sctx.scope_kind)
n_outer = (total + vec_chunk * thread_cnt - 1) // (vec_chunk * thread_cnt)
@@ -214,7 +214,7 @@ def _emit_packed(plan, vec_impl, vec_chunk, total, thread_cnt, sctx) -> PrimFunc
srcs[i].scalar
if srcs[i].is_scalar
else (
- srcs[i].buf_region.buffer,
+ srcs[i].buf_region.source,
_src_lane_indices(
srcs[i].buf_region,
dst_lane_indices,
@@ -239,7 +239,7 @@ def _emit_packed(plan, vec_impl, vec_chunk, total, thread_cnt, sctx) -> PrimFunc
def _emit_scalar(plan, spec, vec_chunk, total, thread_cnt, sctx) -> PrimFunc:
extras = plan.extras
srcs = plan.srcs
- dst_buf = plan.dst.buffer
+ dst_buf = plan.dst.source
dst_st, dst_ext = get_st_extent(plan.dst)
dst_dtype = dst_buf.dtype
compute = spec.compute_scalar
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/binary_f32x2.py b/python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/binary_f32x2.py
index c25d932b22..581bc40cf9 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/binary_f32x2.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/binary_f32x2.py
@@ -49,7 +49,7 @@ def _f32x2_applies(op_name):
def applies(op_call, sctx, plan):
from ...common import sm_version_ok
- if dtype_name(plan.dst.buffer.dtype) != "float32":
+ if dtype_name(plan.dst.source.dtype) != "float32":
return False, "dst dtype not f32"
if not sm_version_ok(op_call, sctx, min_version=100)[0]:
return False, "sm version < 100"
@@ -60,7 +60,7 @@ def _f32x2_applies(op_name):
if scalar_dtype(s.scalar) != "float32":
return False, "scalar src dtype not f32"
else:
- if dtype_name(s.buf_region.buffer.dtype) != "float32":
+ if dtype_name(s.buf_region.source.dtype) != "float32":
return False, "buffer src dtype not f32"
if s.index_fn is not None:
return False, "broadcasting src not supported by f32x2 packed"
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/cast_vec2.py b/python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/cast_vec2.py
index 802d85919f..85a15c5448 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/cast_vec2.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/cast_vec2.py
@@ -61,8 +61,8 @@ def _cast_vec2_applies(op_call, sctx, plan):
src = plan.srcs[0]
if src.index_fn is not None:
return False, "broadcasting src not supported by cast vec2"
- src_dtype = dtype_name(src.buf_region.buffer.dtype)
- dst_dtype = dtype_name(plan.dst.buffer.dtype)
+ src_dtype = dtype_name(src.buf_region.source.dtype)
+ dst_dtype = dtype_name(plan.dst.source.dtype)
if (src_dtype, dst_dtype) not in _VEC2_CAST_INTRINSICS:
return False, f"no vec2 intrinsic for {src_dtype}->{dst_dtype}"
return True, None
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/fma_f32x2.py b/python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/fma_f32x2.py
index ff095141a8..589a981081 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/fma_f32x2.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/fma_f32x2.py
@@ -33,7 +33,7 @@ from .binary_f32x2 import _lane
def _fma_f32x2_applies(op_call, sctx, plan):
from ...common import sm_version_ok
- if dtype_name(plan.dst.buffer.dtype) != "float32":
+ if dtype_name(plan.dst.source.dtype) != "float32":
return False, "dst dtype not f32"
if not sm_version_ok(op_call, sctx, min_version=100)[0]:
return False, "sm version < 100"
@@ -42,7 +42,7 @@ def _fma_f32x2_applies(op_call, sctx, plan):
a, b, c = plan.srcs
if a.is_scalar:
return False, "fma 'a' must be a buffer (no scalar-a packed FMA)"
- if dtype_name(a.buf_region.buffer.dtype) != "float32":
+ if dtype_name(a.buf_region.source.dtype) != "float32":
return False, "src a dtype not f32"
if a.index_fn is not None:
return False, "broadcasting src a not supported"
@@ -51,7 +51,7 @@ def _fma_f32x2_applies(op_call, sctx, plan):
if scalar_dtype(s.scalar) != "float32":
return False, "scalar b/c dtype not f32"
else:
- if dtype_name(s.buf_region.buffer.dtype) != "float32":
+ if dtype_name(s.buf_region.source.dtype) != "float32":
return False, "buffer b/c dtype not f32"
if s.index_fn is not None:
return False, "broadcasting src b/c not supported"
diff --git a/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py b/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py
index 9bf7f31ce5..c76a6f41b7 100644
--- a/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py
+++ b/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py
@@ -252,8 +252,8 @@ def _full_active_lanes(op: TilePrimitiveCall, sctx: DispatchContext):
def _no_replica(op: TilePrimitiveCall, sctx: DispatchContext):
"""All operand layouts must have no replica (no broadcast/duplicated axes)."""
for region, name in zip(op.args[:4], ("D", "A", "B", "C")):
- if region.buffer.layout.replica:
- return False, f"{name} layout has replica {region.buffer.layout.replica}"
+ if region.source.layout.replica:
+ return False, f"{name} layout has replica {region.source.layout.replica}"
return True
@@ -279,7 +279,7 @@ def gemm_cuda_mma_dispatch(op: TilePrimitiveCall, sctx: DispatchContext) -> Prim
# gemm op args: D = alpha * A @ B + beta * C
# D (args[0]) is the output; C (args[3]) is the beta-accumulator input.
D_region, A_region, B_region, C_region, transpose_A, transpose_B, alpha, beta = op.args
- D, A, B, C = D_region.buffer, A_region.buffer, B_region.buffer, C_region.buffer
+ D, A, B, C = D_region.source, A_region.source, B_region.source, C_region.source
# Pure-register mma path: A/B fragments and C/D accumulators all live in
# registers ("local"). The caller is responsible for staging A/B into
diff --git a/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py b/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
index 8ea5f29166..16a3c739ad 100644
--- a/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
+++ b/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
@@ -27,6 +27,7 @@ import operator
import tvm
from tvm.arith.analyzer import Analyzer
+from tvm.ir import TensorRegion
from tvm.runtime import DataType
from tvm.script import tirx as T
from tvm.tirx import PrimFunc
@@ -391,13 +392,13 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -
op_call = TilePrimitiveCall.downcast(op_call)
is_block_scaled = op_call.is_block_scaled
- C_buffer_region: tvm.tirx.BufferRegion = op_call.output
- A_buffer_region: tvm.tirx.BufferRegion = op_call.lhs
- B_buffer_region: tvm.tirx.BufferRegion = op_call.rhs
+ C_buffer_region: TensorRegion = op_call.output
+ A_buffer_region: TensorRegion = op_call.lhs
+ B_buffer_region: TensorRegion = op_call.rhs
C_buffer, A_buffer, B_buffer = (
- C_buffer_region.buffer,
- A_buffer_region.buffer,
- B_buffer_region.buffer,
+ C_buffer_region.source,
+ A_buffer_region.source,
+ B_buffer_region.source,
)
C_scope, A_scope, B_scope = C_buffer.scope(), A_buffer.scope(), B_buffer.scope()
@@ -465,8 +466,8 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -
if is_block_scaled:
SFA_buffer_region, SFB_buffer_region = op_call.sfa, op_call.sfb
transA, transB, accum = op_call.transA, op_call.transB, op_call.accum
- SFA_buffer: tvm.tirx.Buffer = SFA_buffer_region.buffer
- SFB_buffer: tvm.tirx.Buffer = SFB_buffer_region.buffer
+ SFA_buffer: tvm.tirx.Buffer = SFA_buffer_region.source
+ SFB_buffer: tvm.tirx.Buffer = SFB_buffer_region.source
SFA_scope, SFB_scope = SFA_buffer.scope(), SFB_buffer.scope()
if not (SFA_scope == "tmem" and SFB_scope == "tmem"):
raise ValueError(
diff --git a/python/tvm/backend/cuda/tile_primitive/gemm_utils.py b/python/tvm/backend/cuda/tile_primitive/gemm_utils.py
index 207999cdc2..a7ada173d3 100644
--- a/python/tvm/backend/cuda/tile_primitive/gemm_utils.py
+++ b/python/tvm/backend/cuda/tile_primitive/gemm_utils.py
@@ -26,9 +26,9 @@ from tvm.tirx.tile_primitive import TilePrimitiveCall
def validate_gemm_op(op_call: TilePrimitiveCall, sctx: DispatchContext) -> bool:
"""Sanity check for gemm op"""
C_buffer_region, A_buffer_region, B_buffer_region = op_call.args[:3]
- C: Buffer = C_buffer_region.buffer
- A: Buffer = A_buffer_region.buffer
- B: Buffer = B_buffer_region.buffer
+ C: Buffer = C_buffer_region.source
+ A: Buffer = A_buffer_region.source
+ B: Buffer = B_buffer_region.source
if not (C.layout and A.layout and B.layout and A.dtype == B.dtype):
return False
# Extract regions and validate dimensions
diff --git a/python/tvm/backend/cuda/tile_primitive/permute_layout/warp_xor_swizzle.py b/python/tvm/backend/cuda/tile_primitive/permute_layout/warp_xor_swizzle.py
index 90f037083c..3b179bcede 100644
--- a/python/tvm/backend/cuda/tile_primitive/permute_layout/warp_xor_swizzle.py
+++ b/python/tvm/backend/cuda/tile_primitive/permute_layout/warp_xor_swizzle.py
@@ -73,9 +73,10 @@ from __future__ import annotations
import math
+from tvm.ir import TensorRegion
from tvm.runtime import DataType
from tvm.script import tirx as T
-from tvm.tirx import BufferRegion, IntImm, PrimFunc, is_buffer_var
+from tvm.tirx import IntImm, PrimFunc, is_buffer_var
from tvm.tirx.layout import TileLayout, _flatten_coord
from tvm.tirx.operator.tile_primitive import DispatchContext, fail, register_dispatch
from tvm.tirx.tile_primitive import TilePrimitiveCall
@@ -86,13 +87,15 @@ from ..common import get_indices, get_st_extent
def _as_buffer_and_region(arg):
- """Normalize a Buffer or BufferRegion to (buffer, start_list, extent_list)."""
+ """Normalize a Buffer or TensorRegion to (buffer, start_list, extent_list)."""
if is_buffer_var(arg):
buf = arg
extent = list(buf.ty.shape)
st = [0] * len(extent)
- elif isinstance(arg, BufferRegion):
- buf = arg.buffer
+ elif isinstance(arg, TensorRegion):
+ if not is_buffer_var(arg.source):
+ raise TypeError("permute_layout requires a TensorRegion with a BufferVar source")
+ buf = arg.source
st, extent = get_st_extent(arg)
else:
raise TypeError(f"unexpected permute_layout arg type: {type(arg)}")
diff --git a/python/tvm/backend/cuda/tile_primitive/reduction/local.py b/python/tvm/backend/cuda/tile_primitive/reduction/local.py
index 9738ceb20f..f5629235a4 100644
--- a/python/tvm/backend/cuda/tile_primitive/reduction/local.py
+++ b/python/tvm/backend/cuda/tile_primitive/reduction/local.py
@@ -62,8 +62,9 @@ import operator
from typing import Any
from tvm.arith.analyzer import Analyzer
+from tvm.ir import TensorRegion
from tvm.script import tirx as T
-from tvm.tirx import BufferRegion, PrimFunc
+from tvm.tirx import PrimFunc
from tvm.tirx.layout import TileLayout, laneid
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
from tvm.tirx.operator.tile_primitive.common import ReduceOpType
@@ -154,7 +155,7 @@ def validate_reduction_local(
"""Validate reduction in local memory."""
op = TilePrimitiveCall.downcast(op)
dst_br, src_br = op.output, op.input
- dst, src = dst_br.buffer, src_br.buffer
+ dst, src = dst_br.source, src_br.source
if not (src.scope() == "local" and dst.scope() == "local" and sctx.is_target("cuda")):
return False, "expected local scope and CUDA target"
@@ -224,14 +225,14 @@ def validate_reduction_local(
def _emit_reduction_local_thread_wise(
- dst_br: BufferRegion,
- src_br: BufferRegion,
+ dst_br: TensorRegion,
+ src_br: TensorRegion,
accum: bool,
reduce_op: ReduceOpType,
reduce_dims: list[int],
spatial_dims: list[int],
) -> PrimFunc:
- dst, src = dst_br.buffer, src_br.buffer
+ dst, src = dst_br.source, src_br.source
dtype = src.dtype
src_st, src_extent = get_st_extent(src_br)
dst_st, dst_extent = get_st_extent(dst_br)
@@ -283,8 +284,8 @@ def _emit_reduction_local_thread_wise(
def _emit_reduction_local_view(
- dst_br: BufferRegion,
- src_br: BufferRegion,
+ dst_br: TensorRegion,
+ src_br: TensorRegion,
accum: bool,
reduce_op: ReduceOpType,
config: dict[str, Any],
@@ -294,7 +295,7 @@ def _emit_reduction_local_view(
dst_local_info,
shuffle_masks: list[int],
) -> PrimFunc:
- dst, src = dst_br.buffer, src_br.buffer
+ dst, src = dst_br.source, src_br.source
dtype = src.dtype
op_func = reduce_op_table.get(reduce_op)
@@ -411,8 +412,8 @@ def reduction_local_impl(
dst_br, src_br, accum, op_type, reduce_dims, spatial_dims
)
elif sctx.scope_kind in ["warp", "warpgroup"]:
- src = src_br.buffer
- dst = dst_br.buffer
+ src = src_br.source
+ dst = dst_br.source
if sctx.is_warp:
# --- Try laneid shard->replica shuffle reduce ---
diff --git a/python/tvm/backend/cuda/tile_primitive/reduction/shared.py b/python/tvm/backend/cuda/tile_primitive/reduction/shared.py
index fced38e81a..cee4f7e1cc 100644
--- a/python/tvm/backend/cuda/tile_primitive/reduction/shared.py
+++ b/python/tvm/backend/cuda/tile_primitive/reduction/shared.py
@@ -58,8 +58,9 @@ import math
import operator
from tvm.arith.analyzer import Analyzer
+from tvm.ir import TensorRegion
from tvm.script import tirx as T
-from tvm.tirx import BufferRegion, PrimFunc
+from tvm.tirx import PrimFunc
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
from tvm.tirx.operator.tile_primitive.common import ReduceOpType
from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
@@ -84,7 +85,7 @@ def validate_reduction_shared(
return False, f"unsupported exec_scope {sctx.scope_kind} for shared reduction"
op = TilePrimitiveCall.downcast(op)
- dst, src = op.output.buffer, op.input.buffer
+ dst, src = op.output.source, op.input.source
if not (src.scope().startswith("shared") and dst.scope().startswith("shared")):
return False, "expected shared scope for both src and dst"
if src.dtype != dst.dtype:
@@ -117,8 +118,8 @@ def validate_reduction_shared(
def _emit_reduction_shared_cta(
- dst_br: BufferRegion,
- src_br: BufferRegion,
+ dst_br: TensorRegion,
+ src_br: TensorRegion,
accum: bool,
reduce_op: ReduceOpType,
sctx: DispatchContext,
@@ -138,7 +139,7 @@ def _emit_reduction_shared_cta(
return 1
thread_cnt = get_thread_cnt()
- dst, src = dst_br.buffer, src_br.buffer
+ dst, src = dst_br.source, src_br.source
src_st, src_extent = get_st_extent(src_br)
dst_st, dst_extent = get_st_extent(dst_br)
dtype = src.dtype
@@ -215,15 +216,15 @@ def _emit_reduction_shared_cta(
def _emit_reduction_shared_thread(
- dst_br: BufferRegion,
- src_br: BufferRegion,
+ dst_br: TensorRegion,
+ src_br: TensorRegion,
accum: bool,
reduce_op: ReduceOpType,
sctx: DispatchContext,
reduce_dims: list[int],
spatial_dims: list[int],
) -> PrimFunc:
- dst, src = dst_br.buffer, src_br.buffer
+ dst, src = dst_br.source, src_br.source
src_st, src_extent = get_st_extent(src_br)
dst_st, dst_extent = get_st_extent(dst_br)
dtype = src.dtype
diff --git a/python/tvm/backend/cuda/tile_primitive/reduction/sm100_packed.py b/python/tvm/backend/cuda/tile_primitive/reduction/sm100_packed.py
index dd46f31f9e..bec0bf7774 100644
--- a/python/tvm/backend/cuda/tile_primitive/reduction/sm100_packed.py
+++ b/python/tvm/backend/cuda/tile_primitive/reduction/sm100_packed.py
@@ -45,8 +45,9 @@ With accum=True: accumulator folded into first element/pair of the reduction.
import functools
import operator
+from tvm.ir import TensorRegion
from tvm.script import tirx as T
-from tvm.tirx import BufferRegion, PrimFunc
+from tvm.tirx import PrimFunc
from tvm.tirx.operator.tile_primitive import DispatchContext
from tvm.tirx.operator.tile_primitive.common import ReduceOpType
from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
@@ -65,13 +66,13 @@ from .utils import (
def _emit_reduction_local_thread_packed_add_sum(
- dst_buffer_region: BufferRegion,
- src_buffer_region: BufferRegion,
+ dst_buffer_region: TensorRegion,
+ src_buffer_region: TensorRegion,
accum: bool,
reduce_op: ReduceOpType,
sctx: DispatchContext,
) -> PrimFunc:
- dst, src = dst_buffer_region.buffer, src_buffer_region.buffer
+ dst, src = dst_buffer_region.source, src_buffer_region.source
src_region, dst_region = src_buffer_region.region, dst_buffer_region.region
dtype = src.dtype
@@ -143,13 +144,13 @@ def _emit_reduction_local_thread_packed_add_sum(
def _emit_reduction_local_thread_3input_maxmin(
- dst_buffer_region: BufferRegion,
- src_buffer_region: BufferRegion,
+ dst_buffer_region: TensorRegion,
+ src_buffer_region: TensorRegion,
accum: bool,
reduce_op: ReduceOpType,
sctx: DispatchContext,
) -> PrimFunc:
- dst, src = dst_buffer_region.buffer, src_buffer_region.buffer
+ dst, src = dst_buffer_region.source, src_buffer_region.source
src_region, dst_region = src_buffer_region.region, dst_buffer_region.region
dtype = src.dtype
diff --git a/python/tvm/backend/cuda/tile_primitive/reduction/utils.py b/python/tvm/backend/cuda/tile_primitive/reduction/utils.py
index ffb9b5cf1a..ece1d2734a 100644
--- a/python/tvm/backend/cuda/tile_primitive/reduction/utils.py
+++ b/python/tvm/backend/cuda/tile_primitive/reduction/utils.py
@@ -22,8 +22,8 @@ import math
import operator
from tvm.arith.analyzer import Analyzer
+from tvm.ir import TensorRegion
from tvm.script import tirx as T
-from tvm.tirx import BufferRegion
from tvm.tirx.operator.tile_primitive import DispatchContext
from tvm.tirx.operator.tile_primitive.common import ReduceOpType
from tvm.tirx.tile_primitive import TilePrimitiveCall
@@ -47,7 +47,7 @@ def reduce_default_value_table(dtype):
def _reduction_args(
op: TilePrimitiveCall,
-) -> tuple[BufferRegion, BufferRegion, tuple[int, ...], bool, dict]:
+) -> tuple[TensorRegion, TensorRegion, tuple[int, ...], bool, dict]:
"""Parse ReduceOp -> (dst, src, reduce_axes, accum, config)."""
op = TilePrimitiveCall.downcast(op)
dst = op.output
@@ -63,8 +63,8 @@ def _match_reduction_storage_scope(
) -> tuple[bool, str | None]:
"""Check that dst and src scopes match one of the expected patterns."""
op = TilePrimitiveCall.downcast(op)
- dst_scope = op.output.buffer.scope()
- src_scope = op.input.buffer.scope()
+ dst_scope = op.output.source.scope()
+ src_scope = op.input.source.scope()
ok = any(match_scope(dst_scope, p) and match_scope(src_scope, p) for p in expected_scope)
msg = f"storage scope mismatch: dst {dst_scope}, src {src_scope}; expected {expected_scope}"
@@ -218,7 +218,7 @@ _REDUCE_OP_TO_STR = {ReduceOpType.SUM: "sum", ReduceOpType.MAX: "max", ReduceOpT
def _dtype_ok(op: TilePrimitiveCall, sctx: DispatchContext, expected_dtype: str):
op = TilePrimitiveCall.downcast(op)
- dtype = op.input.buffer.dtype
+ dtype = op.input.source.dtype
ok = dtype == expected_dtype
return (ok, None if ok else f"dtype {dtype} != {expected_dtype}")
@@ -248,7 +248,7 @@ def _src_ndim_ok(op: TilePrimitiveCall, sctx: DispatchContext, expected_ndim: in
def _local_scope_match(op: TilePrimitiveCall, sctx: DispatchContext):
op = TilePrimitiveCall.downcast(op)
- src, dst = op.input.buffer, op.output.buffer
+ src, dst = op.input.source, op.output.source
ok = all(
[
src.scope() == "local",
diff --git a/python/tvm/backend/trn/tile_primitive/binary/default.py b/python/tvm/backend/trn/tile_primitive/binary/default.py
index 998532d939..3a432c3fad 100644
--- a/python/tvm/backend/trn/tile_primitive/binary/default.py
+++ b/python/tvm/backend/trn/tile_primitive/binary/default.py
@@ -50,8 +50,8 @@ def binary_trn(
# Extract buffers and constants
CONST = _src2 if isinstance(_src2, FloatImm) else None
- dst, src1 = _dst.buffer, _src1.buffer
- src2 = None if CONST is not None else _src2.buffer
+ dst, src1 = _dst.source, _src1.source
+ src2 = None if CONST is not None else _src2.source
p_var = T.Var("P", "int32")
b_var = T.Var("B", "int32")
diff --git a/python/tvm/backend/trn/tile_primitive/binary/utils.py b/python/tvm/backend/trn/tile_primitive/binary/utils.py
index e3c35dc603..66453e55ee 100644
--- a/python/tvm/backend/trn/tile_primitive/binary/utils.py
+++ b/python/tvm/backend/trn/tile_primitive/binary/utils.py
@@ -21,7 +21,8 @@ from enum import Enum
from tvm.arith.analyzer import Analyzer
from tvm.backend.trn.layout import is_trainium_layout
-from tvm.tirx import BufferRegion, FloatImm
+from tvm.ir import TensorRegion
+from tvm.tirx import FloatImm
from tvm.tirx.operator.tile_primitive.common import MapOpType
from ..dim_utils import get_ewise_dim_map
@@ -42,8 +43,8 @@ class InstType(Enum):
def try_find_inst_nary(
- _dst: BufferRegion,
- _srcs: list[BufferRegion | FloatImm],
+ _dst: TensorRegion,
+ _srcs: list[TensorRegion | FloatImm],
analyzer: Analyzer,
inst_gen: InstructionGenerator,
allowed_f_dim_dst: tuple[int] | None = None,
@@ -65,8 +66,8 @@ def try_find_inst_nary(
# Extract buffers and validate properties
dst, srcs = (
- _dst.buffer,
- [_src.buffer if isinstance(_src, BufferRegion) else None for _src in _srcs],
+ _dst.source,
+ [_src.source if isinstance(_src, TensorRegion) else None for _src in _srcs],
)
dst_region = _dst.region
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py b/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
index e53651b123..802227c17e 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
@@ -17,8 +17,9 @@
"""Implementation of BinaryChain dispatch."""
+from tvm.ir import TensorRegion
from tvm.script import tirx as T
-from tvm.tirx import BufferRegion, PrimFunc, TilePrimitiveCall
+from tvm.tirx import PrimFunc, TilePrimitiveCall
from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
from tvm.tirx.operator.tile_primitive.ops import BinaryChain
@@ -59,7 +60,7 @@ def binary_chain_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc |
p_var = T.Var("P", "int32")
b_var = T.Var("B", "int32")
f_var = T.Var("F", "int32")
- p_size = output.buffer.ty.layout.size("P")
+ p_size = output.source.ty.layout.size("P")
inst_size_limit = op.config.get("max_inst_size", 512)
inst_repr.bound_inst_size(inst_size_limit, analyzer)
inst_gen.bind_inst_iter(output, p_var, p_size, 1, False)
@@ -67,7 +68,7 @@ def binary_chain_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc |
b_extent = inst_gen.fill_in_block_dim(output, b_var)
# Extract buffers and opcodes
- _src, dst = srcs[0].buffer, output.buffer
+ _src, dst = srcs[0].source, output.source
opcode0, opcode1 = opcode_table[op.op0], opcode_table[op.op1]
# Determine operation function based on instruction type
@@ -81,8 +82,8 @@ def binary_chain_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc |
def get_srcs(inst_gen):
return [
(
- srcs[i].buffer[inst_gen.generate_indices(srcs[i])]
- if isinstance(srcs[i], BufferRegion)
+ srcs[i].source[inst_gen.generate_indices(srcs[i])]
+ if isinstance(srcs[i], TensorRegion)
else srcs[i]
)
for i in range(len(srcs))
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py b/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
index 35edbb3186..165ce1a2ce 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
@@ -17,8 +17,9 @@
"""Implementation of BinaryReduce dispatch."""
+from tvm.ir import TensorRegion
from tvm.script import tirx as T
-from tvm.tirx import BufferRegion, PrimFunc, TilePrimitiveCall
+from tvm.tirx import PrimFunc, TilePrimitiveCall
from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
from tvm.tirx.operator.tile_primitive.ops import BinaryReduce
@@ -42,7 +43,7 @@ def binary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc
analyzer = init_analyzer(sctx)
# Normalize negative axes
- reduce_axes = [i if i >= 0 else len(binary_output.buffer.ty.shape) + i for i in reduce_axes]
+ reduce_axes = [i if i >= 0 else len(binary_output.source.ty.shape) + i for i in reduce_axes]
# Find instruction patterns
inst_gen = InstructionGenerator(
@@ -77,7 +78,7 @@ def binary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc
f_var = T.Var("F", "int32")
reduction_b_var = T.Var("rB", "int32")
spatial_b_var = T.Var("sB", "int32")
- p_size = binary_output.buffer.ty.layout.size("P")
+ p_size = binary_output.source.ty.layout.size("P")
inst_gen.bind_inst_iter(binary_output, p_var, p_size, 1, False)
inst_gen.bind_inst_iter(binary_output, f_var, inst_repr.size, inst_repr.stride, True)
reduction_b_extent = inst_gen.fill_in_block_dim(binary_output, reduction_b_var, reduce_axes)
@@ -88,13 +89,13 @@ def binary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc
)
# Handle source 2 (either buffer region or constant)
- CONST = binary_input2 if not isinstance(binary_input2, BufferRegion) else None
+ CONST = binary_input2 if not isinstance(binary_input2, TensorRegion) else None
# Extract buffers and opcodes
src1, src2 = (
- binary_input1.buffer,
- (binary_input2.buffer if isinstance(binary_input2, BufferRegion) else None),
+ binary_input1.source,
+ (binary_input2.source if isinstance(binary_input2, TensorRegion) else None),
)
- dst1, dst2 = binary_output.buffer, reduce_output.buffer
+ dst1, dst2 = binary_output.source, reduce_output.source
binary_opcode, reduce_opcode = opcode_table[op.binary_op], opcode_table[op.reduce_op]
# Create appropriate implementation based on intermediate buffer requirement
if reduction_b_extent == 1:
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py b/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
index 54dd59fb7d..9f622cf404 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
@@ -17,8 +17,9 @@
"""Implementation of UnaryReduce dispatch."""
+from tvm.ir import TensorRegion
from tvm.script import tirx as T
-from tvm.tirx import BufferRegion, PrimFunc, TilePrimitiveCall
+from tvm.tirx import PrimFunc, TilePrimitiveCall
from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
from tvm.tirx.operator.tile_primitive.ops import UnaryReduce
@@ -42,7 +43,7 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc |
analyzer = init_analyzer(sctx)
# Normalize axes and default values
- reduce_axes = [i if i >= 0 else len(unary_output.buffer.ty.shape) + i for i in op.reduce_axes]
+ reduce_axes = [i if i >= 0 else len(unary_output.source.ty.shape) + i for i in op.reduce_axes]
scale = 1.0 if scale is None else scale
bias = 0.0 if bias is None else bias
@@ -50,7 +51,7 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc |
reduce_dim_map = get_reduction_dim_map(unary_output, reduce_output, reduce_axes, analyzer)
inst_gen.link_buffer_regions(unary_output, reduce_output, reduce_dim_map)
# Find instruction patterns based on bias type
- if isinstance(bias, BufferRegion):
+ if isinstance(bias, TensorRegion):
inst_repr, _, _ = try_find_inst_nary(
unary_output,
[unary_input, bias],
@@ -72,7 +73,7 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc |
f_var = T.Var("F", "int32")
reduction_b_var = T.Var("rB", "int32")
spatial_b_var = T.Var("sB", "int32")
- p_size = unary_output.buffer.ty.layout.size("P")
+ p_size = unary_output.source.ty.layout.size("P")
inst_gen.bind_inst_iter(unary_output, p_var, p_size, 1, False)
inst_gen.bind_inst_iter(unary_output, f_var, inst_repr.size, inst_repr.stride, True)
reduction_b_extent = inst_gen.fill_in_block_dim(unary_output, reduction_b_var, reduce_axes)
@@ -82,14 +83,14 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc |
reduce_output, reduction_b_extent, op.workspace, sctx
)
# Extract buffers and opcodes
- src, dst1, dst2 = unary_input.buffer, unary_output.buffer, reduce_output.buffer
+ src, dst1, dst2 = unary_input.source, unary_output.source, reduce_output.source
unary_opcode = opcode_table[op.unary_op]
reduce_opcode = opcode_table[op.reduce_op]
# Handle bias buffer
bias_buffer = (
- bias.buffer
- if isinstance(bias, BufferRegion)
+ bias.source
+ if isinstance(bias, TensorRegion)
else get_const_bias_tensor(
bias, (p_size, inst_repr.size), dst1.ty.dtype, op.workspace, sctx
)
@@ -110,7 +111,7 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc |
dst_1_indices = T.meta_var(inst_gen.generate_indices(unary_output))
dst_2_indices = T.meta_var(inst_gen.generate_indices(reduce_output))
if inst_gen.make_guard(unary_output):
- if isinstance(bias, BufferRegion):
+ if isinstance(bias, TensorRegion):
src_bias_indices = T.meta_var(inst_gen.generate_indices(bias))
T.evaluate(T.nki.activation_reduce(dst2[tuple(dst_2_indices)], dst1[tuple(dst_1_indices)], src[tuple(src_1_indices)], unary_opcode, reduce_opcode, bias_buffer[tuple(src_bias_indices)], scale)) # noqa: E501
else:
@@ -135,7 +136,7 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc |
src_1_indices = T.meta_var(inst_gen.generate_indices(unary_input))
dst_1_indices = T.meta_var(inst_gen.generate_indices(unary_output))
if inst_gen.make_guard(unary_output):
- if isinstance(bias, BufferRegion):
+ if isinstance(bias, TensorRegion):
src_bias_indices = T.meta_var(inst_gen.generate_indices(bias)) # noqa: E501
T.evaluate(T.nki.activation_reduce(intermediate_buffer[p_loop, reduction_b_loop], dst1[tuple(dst_1_indices)], src[tuple(src_1_indices)], unary_opcode, reduce_opcode, bias_buffer[tuple(src_bias_indices)], scale)) # noqa: E501
else:
diff --git a/python/tvm/backend/trn/tile_primitive/copy/default.py b/python/tvm/backend/trn/tile_primitive/copy/default.py
index 1dd18ed776..e22b3800b6 100644
--- a/python/tvm/backend/trn/tile_primitive/copy/default.py
+++ b/python/tvm/backend/trn/tile_primitive/copy/default.py
@@ -38,7 +38,7 @@ def transpose_schedule(
op: TilePrimitiveCall, inst_gen: InstructionGenerator, sctx: DispatchContext
) -> PrimFunc | None:
dst_region, src_region = op.args
- assert src_region.buffer.scope() != "trn.psum", "Transpose on psum buffer is not supported"
+ assert src_region.source.scope() != "trn.psum", "Transpose on psum buffer is not supported"
inst_repr_dst, inst_repr_src = inst_gen.find_max_inst_size_transpose(dst_region, src_region)
@@ -47,8 +47,8 @@ def transpose_schedule(
dst_f = T.Var("dst_F", "int32")
b_var = T.Var("B", "int32")
extend_b = T.Var("extend_B", "int32")
- p_size = src_region.buffer.ty.layout.size("P")
- lhs_f_size = dst_region.buffer.ty.layout.size("P")
+ p_size = src_region.source.ty.layout.size("P")
+ lhs_f_size = dst_region.source.ty.layout.size("P")
rhs_f_size = p_size
inst_gen.bind_inst_iter(
src_region, lhs_f, inst_repr_src.size, inst_repr_src.stride, is_free_dim=True
@@ -62,7 +62,7 @@ def transpose_schedule(
no_propagate=True,
)
inst_gen.bind_inst_iter(src_region, lhs_p, p_size, 1, is_free_dim=False, no_propagate=True)
- if dst_region.buffer.scope() == "trn.sbuf":
+ if dst_region.source.scope() == "trn.sbuf":
max_extend_num = (
inst_gen.find_max_inst_size_from_one_region(
dst_region, min_stride=inst_repr_dst.stride
@@ -91,7 +91,7 @@ def transpose_schedule(
)
identity_tensor = T.buffer(
(p_size, rhs_f_size),
- src_region.buffer.ty.dtype,
+ src_region.source.ty.dtype,
scope="trn.sbuf",
buffer_name="identity",
)
@@ -110,8 +110,8 @@ def transpose_schedule(
identity_tensor = op.workspace["identity"]
check_workspace_buffer(identity_tensor, (p_size, rhs_f_size), "trn.sbuf")
- dst_buffer = dst_region.buffer
- src_buffer = src_region.buffer
+ dst_buffer = dst_region.source
+ src_buffer = src_region.source
if dst_buffer.scope() == "trn.psum":
@T.prim_func
@@ -196,7 +196,7 @@ def copy_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
fail("requires thread exec_scope for TRN copy")
dst_region, src_region = op.args
- src, dst = src_region.buffer, dst_region.buffer
+ src, dst = src_region.source, dst_region.source
# Check for valid buffer configurations
valid_config = all(
@@ -265,7 +265,7 @@ def copy_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
from_region, _to_region = src_region, dst_region
else:
from_region, _to_region = dst_region, src_region
- p_size = from_region.buffer.ty.layout.size("P")
+ p_size = from_region.source.ty.layout.size("P")
inst_gen.bind_inst_iter(from_region, p_var, p_size, 1, is_free_dim=False)
inst_gen.bind_inst_iter(from_region, f_var, inst.size, inst.stride, is_free_dim=True)
b_extent = inst_gen.fill_in_block_dim(from_region, b_var)
diff --git a/python/tvm/backend/trn/tile_primitive/dim_utils.py b/python/tvm/backend/trn/tile_primitive/dim_utils.py
index ff7064f63d..f1de58c334 100644
--- a/python/tvm/backend/trn/tile_primitive/dim_utils.py
+++ b/python/tvm/backend/trn/tile_primitive/dim_utils.py
@@ -20,8 +20,8 @@
from collections import namedtuple
from tvm.arith.analyzer import Analyzer
+from tvm.ir import TensorRegion
from tvm.script import tirx as T
-from tvm.tirx import BufferRegion
# Represents the part of data iter covered by the buffer region
RangeInfo = namedtuple(
@@ -56,15 +56,15 @@ def normalize_and_group(layout, shape):
def get_ewise_dim_map(
- buffer_region: BufferRegion, second_buffer_region: BufferRegion, analyzer: Analyzer
+ buffer_region: TensorRegion, second_buffer_region: TensorRegion, analyzer: Analyzer
):
"""Get the dimension map between two elementwise buffer regions.
Parameters
----------
- buffer_region : BufferRegion
+ buffer_region : TensorRegion
The first buffer region
- second_buffer_region : BufferRegion
+ second_buffer_region : TensorRegion
The second buffer region
analyzer : Analyzer
The analyzer to use
@@ -108,8 +108,8 @@ def get_ewise_dim_map(
def get_reduction_dim_map(
- src_buffer_region: BufferRegion,
- dst_buffer_region: BufferRegion,
+ src_buffer_region: TensorRegion,
+ dst_buffer_region: TensorRegion,
axes: tuple[int],
analyzer: Analyzer,
):
@@ -117,9 +117,9 @@ def get_reduction_dim_map(
Parameters
----------
- src_buffer_region : BufferRegion
+ src_buffer_region : TensorRegion
The source buffer region
- dst_buffer_region : BufferRegion
+ dst_buffer_region : TensorRegion
The destination buffer region
axes : Tuple[int]
The reduction axes
diff --git a/python/tvm/backend/trn/tile_primitive/gemm/default.py b/python/tvm/backend/trn/tile_primitive/gemm/default.py
index 1a52792fda..fa9a67a7c1 100644
--- a/python/tvm/backend/trn/tile_primitive/gemm/default.py
+++ b/python/tvm/backend/trn/tile_primitive/gemm/default.py
@@ -22,9 +22,9 @@ import operator
from tvm.arith.analyzer import Analyzer
from tvm.backend.trn.layout import is_trainium_layout
-from tvm.ir import assert_structural_equal
+from tvm.ir import TensorRegion, assert_structural_equal
from tvm.script import tirx as T
-from tvm.tirx import BufferRegion, PrimFunc
+from tvm.tirx import PrimFunc
from tvm.tirx.operator.tile_primitive import (
DispatchContext,
fail,
@@ -46,7 +46,7 @@ class OperatorKind:
def get_pf_dim_from_buffer_region(
- buffer_region: BufferRegion,
+ buffer_region: TensorRegion,
analyzer: Analyzer,
operator_kind: OperatorKind,
transposed: bool = False,
@@ -55,13 +55,13 @@ def get_pf_dim_from_buffer_region(
# Find non-unit dimensions
non_unit_dims = [
i
- for i in range(len(buffer_region.buffer.ty.shape))
+ for i in range(len(buffer_region.source.ty.shape))
if not analyzer.can_prove_equal(buffer_region.region[i].extent, 1)
]
assert len(non_unit_dims) == 2, "Only 2D matrix is supported for gemm"
layout, seps = normalize_and_group(
- buffer_region.buffer.ty.layout, buffer_region.buffer.ty.shape
+ buffer_region.source.ty.layout, buffer_region.source.ty.shape
)
# Determine partition and free dimensions based on operator kind
if operator_kind == OperatorKind.A:
@@ -129,10 +129,10 @@ def matmul_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
) = op.args
analyzer = init_analyzer(sctx)
A, B, C, _D = (
- A_buffer_region.buffer,
- B_buffer_region.buffer,
- C_buffer_region.buffer,
- D_buffer_region.buffer,
+ A_buffer_region.source,
+ B_buffer_region.source,
+ C_buffer_region.source,
+ D_buffer_region.source,
)
# Validate alpha, beta
diff --git a/python/tvm/backend/trn/tile_primitive/instruction_generator.py b/python/tvm/backend/trn/tile_primitive/instruction_generator.py
index 06aa94e859..48591d8860 100644
--- a/python/tvm/backend/trn/tile_primitive/instruction_generator.py
+++ b/python/tvm/backend/trn/tile_primitive/instruction_generator.py
@@ -28,9 +28,9 @@ import tvm_ffi
import tvm
from tvm.arith.analyzer import Analyzer
from tvm.backend.trn.layout import is_trainium_layout
-from tvm.ir import Range
+from tvm.ir import Range, TensorRegion
from tvm.script import tirx as T
-from tvm.tirx import BufferRegion, Expr, Var
+from tvm.tirx import BufferRegion, Expr, Var, is_buffer_var
from tvm.tirx.layout import Iter
from .dim_utils import DimensionMapper, RangeInfo, normalize_and_group
@@ -64,14 +64,14 @@ def _replace_vars(expr: Expr, var_map: dict[Var, Expr]) -> Expr:
@dataclass
class InstructionRepr:
- buffer_region: BufferRegion
+ buffer_region: TensorRegion
size: int
stride: int
selected_data_iter_ids: list[int]
def __init__(
self,
- buffer_region: BufferRegion,
+ buffer_region: TensorRegion,
inst_size: int,
inst_stride: int,
selected_data_iter_ids: list[int],
@@ -94,18 +94,22 @@ class InstructionRepr:
class InstructionGenerator:
- def __init__(self, buffer_regions: tuple[BufferRegion], analyzer: Analyzer):
+ def __init__(self, buffer_regions: tuple[TensorRegion], analyzer: Analyzer):
self.buffer_regions = []
self.analyzer = analyzer
self.split_shape_views = {}
self.split_layout_views = {}
self.seps = {}
self.bound_regions = {}
- self.bind_iters: dict[BufferRegion, LogicalIterList] = None
- self.bind_maps: dict[BufferRegion, dict[Var, Expr]] = {}
+ self.bind_iters: dict[TensorRegion, LogicalIterList] = None
+ self.bind_maps: dict[TensorRegion, dict[Var, Expr]] = {}
for buffer_region in buffer_regions:
- if not isinstance(buffer_region, BufferRegion):
+ if not isinstance(buffer_region, TensorRegion):
continue
+ if not is_buffer_var(buffer_region.source):
+ raise TypeError(
+ "Instruction operands require a TensorRegion with a BufferVar source"
+ )
self.buffer_regions.append(buffer_region)
bound_buffer_region = self._bound_buffer_region(buffer_region)
layout, seps = self._get_sub_layout(bound_buffer_region)
@@ -116,7 +120,7 @@ class InstructionGenerator:
self.seps[buffer_region] = seps
self.dim_mapper = DimensionMapper()
- def _bound_buffer_region(self, buffer_region: BufferRegion):
+ def _bound_buffer_region(self, buffer_region: TensorRegion):
region = []
changed = False
for r in buffer_region.region:
@@ -125,14 +129,14 @@ class InstructionGenerator:
changed = True
region.append(Range.from_min_extent(r.min, bound.max_value))
if changed:
- bound_region = BufferRegion(buffer_region.buffer, region)
+ bound_region = BufferRegion(buffer_region.source, region)
self.bound_regions[buffer_region] = bound_region
return bound_region
return buffer_region
- def _get_sub_layout(self, buffer_region: BufferRegion):
- layout = buffer_region.buffer.ty.layout
- layout, seps = normalize_and_group(layout, buffer_region.buffer.ty.shape)
+ def _get_sub_layout(self, buffer_region: TensorRegion):
+ layout = buffer_region.source.ty.layout
+ layout, seps = normalize_and_group(layout, buffer_region.source.ty.shape)
tiled_range_infos_per_dim = []
new_shard = []
new_seps = [0]
@@ -262,7 +266,7 @@ class InstructionGenerator:
return tuple(out)
def _link_buffer_regions(
- self, buffer_region: BufferRegion, to_link: BufferRegion, dim_map: dict[int, int]
+ self, buffer_region: TensorRegion, to_link: TensorRegion, dim_map: dict[int, int]
):
split_shape_view_1 = self.split_shape_views[buffer_region]
split_layout_view_1 = self.split_layout_views[buffer_region]
@@ -294,7 +298,7 @@ class InstructionGenerator:
return {dim_map[i]: i for i in dim_map}
def link_buffer_regions(
- self, buffer_region: BufferRegion, to_link: BufferRegion, dim_map: dict[int, int]
+ self, buffer_region: TensorRegion, to_link: TensorRegion, dim_map: dict[int, int]
):
self.dim_mapper.register_dim_map(buffer_region, to_link, dim_map)
for r in self.buffer_regions:
@@ -308,12 +312,12 @@ class InstructionGenerator:
seps_2 = self.seps[to_link]
for i, j in dim_map.items():
assert seps_1[i + 1] - seps_1[i] == seps_2[j + 1] - seps_2[j], (
- f"The number of data iters at dim {i} of {buffer_region.buffer.name} is not equal to the number of data iters at dim {j} of {to_link.buffer.name}" # noqa: E501
+ f"The number of data iters at dim {i} of {buffer_region.source.name} is not equal to the number of data iters at dim {j} of {to_link.source.name}" # noqa: E501
)
def bind_inst_iter(
self,
- buffer_region: BufferRegion,
+ buffer_region: TensorRegion,
bind: Var,
inst_size: int,
inst_stride: int,
@@ -328,7 +332,7 @@ class InstructionGenerator:
return
self._propagate_bind_iter(buffer_region, logical_iter_list)
- def _propagate_bind_iter(self, buffer_region: BufferRegion, logical_iter_list: LogicalIterList):
+ def _propagate_bind_iter(self, buffer_region: TensorRegion, logical_iter_list: LogicalIterList):
for to_propagate in self.buffer_regions:
if to_propagate == buffer_region:
continue
@@ -345,7 +349,7 @@ class InstructionGenerator:
]
self._add_bind_iter_list(to_propagate, propagated_logical_iter)
- def _add_bind_iter_list(self, buffer_region: BufferRegion, bind_iter_list: LogicalIterList):
+ def _add_bind_iter_list(self, buffer_region: TensorRegion, bind_iter_list: LogicalIterList):
if self.bind_iters is None:
self._init_bind_iters()
seps = self.seps[buffer_region]
@@ -356,12 +360,12 @@ class InstructionGenerator:
)
def fill_in_block_dim(
- self, buffer_region: BufferRegion, bind: Var, dims: list[int] | None = None
+ self, buffer_region: TensorRegion, bind: Var, dims: list[int] | None = None
):
# fixme: be cautious of the min of buffer region. This implementation is not correct.
# we need to first take a view of sub-layout (keep strides, but reduce the extent
# then we analyze the relationship between data iter of sub-layout
- dims = dims or list(range(len(buffer_region.buffer.ty.shape)))
+ dims = dims or list(range(len(buffer_region.source.ty.shape)))
layout = self.split_layout_views[buffer_region]
shards = layout.shard
self._normalize_bind_iters()
@@ -388,12 +392,12 @@ class InstructionGenerator:
else 1
)
assert next_logical_stride % cur == 0, (
- f"Fail to infer block dim for {buffer_region.buffer.name} at dim {i}"
+ f"Fail to infer block dim for {buffer_region.source.name} at dim {i}"
)
gap = next_logical_stride // cur
if is_partition:
assert gap == 1, (
- f"Fail to propagate partition dim. The propagated dim does not cover the whole partition on {buffer_region.buffer.name} at dim {i}" # noqa: E501
+ f"Fail to propagate partition dim. The propagated dim does not cover the whole partition on {buffer_region.source.name} at dim {i}" # noqa: E501
)
elif gap > 1:
new_acc_block_ext = acc_block_ext * gap
@@ -405,7 +409,7 @@ class InstructionGenerator:
self._propagate_bind_iter(buffer_region, logical_iter_list_block)
return acc_block_ext
- def _check_bind_iter_coverage(self, buffer_region: BufferRegion):
+ def _check_bind_iter_coverage(self, buffer_region: TensorRegion):
self._normalize_bind_iters()
seps = self.seps[buffer_region]
iters = self.split_layout_views[buffer_region].shard
@@ -424,20 +428,20 @@ class InstructionGenerator:
next_logical_stride
% (logical_iter_dims[d].logical_stride * logical_iter_dims[d].extent)
== 0
- ), f"Fail to infer block dim for {buffer_region.buffer.name} at dim {i}"
+ ), f"Fail to infer block dim for {buffer_region.source.name} at dim {i}"
gap = next_logical_stride // (
logical_iter_dims[d].logical_stride * logical_iter_dims[d].extent
)
assert gap == 1, "Call fill_in_block_dim() before calling generate_indices()"
- def set_bind_map(self, buffer_region: BufferRegion, bind_map: dict[Var, Expr]):
+ def set_bind_map(self, buffer_region: TensorRegion, bind_map: dict[Var, Expr]):
self.bind_maps[buffer_region] = bind_map
def set_bind_map_all(self, bind_map: dict[Var, Expr]):
for buffer_region in self.buffer_regions:
self.set_bind_map(buffer_region, bind_map)
- def generate_axes(self, buffer_region: BufferRegion) -> list[Expr]:
+ def generate_axes(self, buffer_region: TensorRegion) -> list[Expr]:
self._check_bind_iter_coverage(buffer_region)
layout = self.split_layout_views[buffer_region]
iters = layout.shard
@@ -461,13 +465,13 @@ class InstructionGenerator:
axes.append(index)
return axes
- def generate_indices(self, buffer_region: BufferRegion) -> list[Expr]:
+ def generate_indices(self, buffer_region: TensorRegion) -> list[Expr]:
axes = self.generate_axes(buffer_region)
return [axes[i] + r.min for i, r in enumerate(buffer_region.region)]
def _get_inst_logical_iter_list(
self,
- buffer_region: BufferRegion,
+ buffer_region: TensorRegion,
bind: Var,
stride: int,
size: int,
@@ -518,7 +522,7 @@ class InstructionGenerator:
)
return ret
- def make_guard(self, buffer_region: BufferRegion):
+ def make_guard(self, buffer_region: TensorRegion):
if buffer_region not in self.bound_regions:
return True
bound_region = self.bound_regions[buffer_region]
@@ -561,7 +565,7 @@ class InstructionGenerator:
def find_max_inst_size_from_one_region(
self,
- buffer_region: BufferRegion,
+ buffer_region: TensorRegion,
allowed_f_dim: tuple[int] | None = None,
min_stride: int | None = None,
):
@@ -584,7 +588,7 @@ class InstructionGenerator:
def fit_inst_tile_to_region(
self,
inst_repr: InstructionRepr,
- to_region: BufferRegion,
+ to_region: TensorRegion,
allowed_to_f_dim: tuple[int] | None = None,
broadcast: bool = False,
):
@@ -641,7 +645,7 @@ class InstructionGenerator:
return InstructionRepr(from_region, inst_size, inst_stride_from, idx_list)
def check_partition_dim_match(
- self, buffer_region_1: BufferRegion, buffer_region_2: BufferRegion
+ self, buffer_region_1: TensorRegion, buffer_region_2: TensorRegion
):
dim_map = self.dim_mapper.get_dim_map(buffer_region_1, buffer_region_2)
layout_1 = self.split_layout_views[buffer_region_1]
@@ -666,7 +670,7 @@ class InstructionGenerator:
return True
def find_max_inst_size_transpose(
- self, buffer_region_1: BufferRegion, buffer_region_2: BufferRegion
+ self, buffer_region_1: TensorRegion, buffer_region_2: TensorRegion
):
dim_map = self.dim_mapper.get_dim_map(buffer_region_1, buffer_region_2)
layout_1 = self.split_layout_views[buffer_region_1]
@@ -693,10 +697,10 @@ class InstructionGenerator:
inst_repr_1 = InstructionRepr(buffer_region_1, *self._find_max_linear_inst(indexed_iters_1))
inst_repr_2 = InstructionRepr(buffer_region_2, *self._find_max_linear_inst(indexed_iters_2))
assert inst_repr_1.size == layout_2.size("P"), (
- f"The instruction size of {buffer_region_1.buffer.name} does not match the partition size of {buffer_region_2.buffer.name}" # noqa: E501
+ f"The instruction size of {buffer_region_1.source.name} does not match the partition size of {buffer_region_2.source.name}" # noqa: E501
)
assert inst_repr_2.size == layout_1.size("P"), (
- f"The instruction size of {buffer_region_2.buffer.name} does not match the partition size of {buffer_region_1.buffer.name}" # noqa: E501
+ f"The instruction size of {buffer_region_2.source.name} does not match the partition size of {buffer_region_1.source.name}" # noqa: E501
)
return inst_repr_1, inst_repr_2
@@ -709,7 +713,7 @@ class InstructionGenerator:
indexed_selected_iters = sorted(indexed_selected_iters, key=lambda x: x[1].stride)
iter_idx_to_dim = {
seps[j]: i
- for i in range(len(region.buffer.ty.shape))
+ for i in range(len(region.source.ty.shape))
for j in range(seps[i], seps[i + 1])
}
last_dim = None
diff --git a/python/tvm/backend/trn/tile_primitive/private_alloc.py b/python/tvm/backend/trn/tile_primitive/private_alloc.py
index a0759c6d9c..81ab5827e5 100644
--- a/python/tvm/backend/trn/tile_primitive/private_alloc.py
+++ b/python/tvm/backend/trn/tile_primitive/private_alloc.py
@@ -48,12 +48,12 @@ def _scalar_dtype(scalar) -> str:
def alloc_const_bias_trn(
op: TilePrimitiveCall, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: DispatchContext
) -> dict[str, Any]:
- bias = op.bias if op.bias is not None else FloatImm(op.dsts[0].buffer.ty.dtype, 0.0)
+ bias = op.bias if op.bias is not None else FloatImm(op.dsts[0].source.ty.dtype, 0.0)
if "const_bias" in op.workspace:
return {}
if not isinstance(bias, (FloatImm)):
return {}
- par_size = op.dsts[0].buffer.ty.layout.size("P")
+ par_size = op.dsts[0].source.ty.layout.size("P")
max_inst_size = op.config.get("max_inst_size", 512)
if ("const_bias", bias.value) in buffer_dict:
bias_buffer, bias_init_stmt = buffer_dict[("const_bias", bias.value)]
@@ -104,7 +104,7 @@ def alloc_identity_trn(
) -> dict[str, Any]:
if "identity" in op.workspace:
return {}
- par_size = op.srcs[0].buffer.ty.layout.size("P")
+ par_size = op.srcs[0].source.ty.layout.size("P")
if "identity" in buffer_dict:
identity_buffer, identity_init_stmt = buffer_dict["identity"]
old_shape = identity_buffer.ty.shape
@@ -114,7 +114,7 @@ def alloc_identity_trn(
else:
new_shape = (par_size, par_size)
new_buffer = T.buffer(
- new_shape, dtype=op.srcs[0].buffer.ty.dtype, scope="trn.sbuf", buffer_name="identity"
+ new_shape, dtype=op.srcs[0].source.ty.dtype, scope="trn.sbuf", buffer_name="identity"
)
@T.prim_func
@@ -132,9 +132,9 @@ def alloc_identity_trn(
def alloc_acc_psum_trn(
op: TilePrimitiveCall, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: DispatchContext
) -> dict[str, Any]:
- if "acc_psum" in op.workspace or op.dsts[0].buffer.scope() == "trn.psum":
+ if "acc_psum" in op.workspace or op.dsts[0].source.scope() == "trn.psum":
return {}
- par_size = op.dsts[0].buffer.ty.layout.size("P")
+ par_size = op.dsts[0].source.ty.layout.size("P")
acc_psum = T.buffer(
(8, par_size, 512),
"float32",
diff --git a/python/tvm/backend/trn/tile_primitive/reduction/utils.py b/python/tvm/backend/trn/tile_primitive/reduction/utils.py
index e600ebda6d..2d791908b2 100644
--- a/python/tvm/backend/trn/tile_primitive/reduction/utils.py
+++ b/python/tvm/backend/trn/tile_primitive/reduction/utils.py
@@ -40,7 +40,7 @@ def generate_intermediate_buffer(
Returns:
Tuple[Optional[buffer], int]: The intermediate buffer and reduction factor size.
"""
- intermediate_shape = [dst_buffer_region.buffer.ty.layout.size("P"), rfactor_size]
+ intermediate_shape = [dst_buffer_region.source.ty.layout.size("P"), rfactor_size]
if "partial_reduce" in workspace:
intermediate_buffer = workspace["partial_reduce"]
@@ -51,7 +51,7 @@ def generate_intermediate_buffer(
)
intermediate_buffer = T.buffer(
intermediate_shape,
- dtype=dst_buffer_region.buffer.ty.dtype,
+ dtype=dst_buffer_region.source.ty.dtype,
scope="trn.sbuf",
buffer_name="partial_reduce",
)
@@ -83,8 +83,8 @@ def reduction_trn(
assert reduce_op in reduce_ops, f"Unsupported reduce operation {reduce_op}"
# Extract buffers
- dst = dst_buffer_region.buffer
- src = src_buffer_region.buffer
+ dst = dst_buffer_region.source
+ src = src_buffer_region.source
axes = [i if i >= 0 else len(src.ty.shape) + i for i in axes]
dim_map = get_reduction_dim_map(src_buffer_region, dst_buffer_region, axes, analyzer)
diff --git a/python/tvm/backend/trn/tile_primitive/select/default.py b/python/tvm/backend/trn/tile_primitive/select/default.py
index ca44225ae7..9106f3c8a6 100644
--- a/python/tvm/backend/trn/tile_primitive/select/default.py
+++ b/python/tvm/backend/trn/tile_primitive/select/default.py
@@ -18,8 +18,9 @@
"""Implementation of select schedules."""
from tvm.backend.trn.layout import is_trainium_layout
+from tvm.ir import TensorRegion
from tvm.script import tirx as T
-from tvm.tirx import BufferRegion, FloatImm, PrimFunc, TilePrimitiveCall
+from tvm.tirx import FloatImm, PrimFunc, TilePrimitiveCall
from tvm.tirx.operator.tile_primitive import (
DispatchContext,
fail,
@@ -55,17 +56,17 @@ def select_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
pred = not pred
true_value, false_value = false_value, true_value
- assert isinstance(true_value, BufferRegion), f"{op} expects one of the source to be a buffer"
+ assert isinstance(true_value, TensorRegion), f"{op} expects one of the source to be a buffer"
# Initialize analyzer and validate buffers
analyzer = init_analyzer(sctx)
# Validate buffer layout and scope
buffer_conditions = [
- dst.buffer.ty.layout and true_value.buffer.ty.layout,
- dst.buffer.scope() == "trn.sbuf" and true_value.buffer.scope() == "trn.sbuf",
- is_trainium_layout(true_value.buffer.ty.layout),
- is_trainium_layout(dst.buffer.ty.layout),
+ dst.source.ty.layout and true_value.source.ty.layout,
+ dst.source.scope() == "trn.sbuf" and true_value.source.scope() == "trn.sbuf",
+ is_trainium_layout(true_value.source.ty.layout),
+ is_trainium_layout(dst.source.ty.layout),
]
if not all(buffer_conditions):
@@ -98,14 +99,14 @@ def select_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
p_var = T.Var("p", "int32")
b_var = T.Var("b", "int32")
f_var = T.Var("f", "int32")
- p_size = dst.buffer.ty.layout.size("P")
+ p_size = dst.source.ty.layout.size("P")
inst_gen.bind_inst_iter(dst, f_var, inst_repr.size, inst_repr.stride, True)
inst_gen.bind_inst_iter(dst, p_var, p_size, 1, False)
b_extent = inst_gen.fill_in_block_dim(dst, b_var)
# Get buffer references and guard function
- dst_buffer = dst.buffer
- true_value_buffer = true_value.buffer
+ dst_buffer = dst.source
+ true_value_buffer = true_value.source
# fmt: off
@T.prim_func
diff --git a/python/tvm/backend/trn/tile_primitive/unary/utils.py b/python/tvm/backend/trn/tile_primitive/unary/utils.py
index a982f2fb2f..f2589a7a7d 100644
--- a/python/tvm/backend/trn/tile_primitive/unary/utils.py
+++ b/python/tvm/backend/trn/tile_primitive/unary/utils.py
@@ -19,8 +19,9 @@
from tvm.arith.analyzer import Analyzer
from tvm.backend.trn.layout import is_trainium_layout
+from tvm.ir import TensorRegion
from tvm.script import tirx as T
-from tvm.tirx import BufferRegion, FloatImm
+from tvm.tirx import FloatImm
from tvm.tirx.operator.tile_primitive.common import MapOpType
from ..common import nki_dim
@@ -40,16 +41,16 @@ const_input_ops = [MapOpType.FILL]
def try_find_inst_unary(
- dst_buffer_region: BufferRegion,
- src_buffer_region: BufferRegion,
+ dst_buffer_region: TensorRegion,
+ src_buffer_region: TensorRegion,
analyzer: Analyzer,
inst_gen: InstructionGenerator,
allowed_f_dim_dst: tuple[int] | None = None,
allowed_f_dim_src: tuple[int] | None = None,
):
"""Find instruction parameters for a unary operation."""
- dst = dst_buffer_region.buffer
- src = src_buffer_region.buffer
+ dst = dst_buffer_region.source
+ src = src_buffer_region.source
# Validate buffer layouts and scopes
valid_layout_scope = all(
@@ -136,7 +137,7 @@ def generate_unary_func(
):
"""Generate a function that implements a unary operation."""
# Prepare parameters
- p_size = dst_buffer_region.buffer.ty.layout.size("P")
+ p_size = dst_buffer_region.source.ty.layout.size("P")
# Apply instruction size limits if specified
inst_size_limit = config.get("max_inst_size", 512)
@@ -153,16 +154,16 @@ def generate_unary_func(
opcode = opcode_table.get(unary_op, None)
# Extract buffers
- dst = dst_buffer_region.buffer
- src = _src.buffer if isinstance(_src, BufferRegion) else None
+ dst = dst_buffer_region.source
+ src = _src.source if isinstance(_src, TensorRegion) else None
# Handle bias tensor
if isinstance(bias, FloatImm | float):
bias_buffer = get_const_bias_tensor(
bias, (p_size, inst_repr.size), dst.ty.dtype, workspace, sctx
)
- elif isinstance(bias, BufferRegion):
- bias_buffer = bias.buffer
+ elif isinstance(bias, TensorRegion):
+ bias_buffer = bias.source
# fmt: off
@T.prim_func
@@ -180,7 +181,7 @@ def generate_unary_func(
src_indices = T.meta_var(inst_gen.generate_indices(_src))
if unary_op == MapOpType.RECIPROCAL:
T.evaluate(T.nki.reciprocal(dst[tuple(dst_indices)], src[tuple(src_indices)])) # noqa: E501
- elif isinstance(bias, BufferRegion):
+ elif isinstance(bias, TensorRegion):
bias_indices = T.meta_var(inst_gen.generate_indices(bias))
T.evaluate(T.nki.activation(dst[tuple(dst_indices)], src[tuple(src_indices)], opcode, scale=scale, bias=bias_buffer[tuple(bias_indices)])) # noqa: E501
else:
diff --git a/python/tvm/backend/trn/tile_primitive/unary/with_bias_scale.py b/python/tvm/backend/trn/tile_primitive/unary/with_bias_scale.py
index c75f18544f..6c728c35ad 100644
--- a/python/tvm/backend/trn/tile_primitive/unary/with_bias_scale.py
+++ b/python/tvm/backend/trn/tile_primitive/unary/with_bias_scale.py
@@ -17,7 +17,8 @@
"""Implementation of unary with bias and scale operator dispatches."""
-from tvm.tirx import BufferRegion, PrimFunc
+from tvm.ir import TensorRegion
+from tvm.tirx import PrimFunc
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
from tvm.tirx.operator.tile_primitive.common import MapOpType
from tvm.tirx.tile_primitive import TilePrimitiveCall
@@ -47,7 +48,7 @@ def unary_with_bias_scale_trn(
# Find instruction parameters
inst_gen = InstructionGenerator([dst_buffer_region, src_buffer_region, _bias], analyzer)
- if isinstance(_bias, BufferRegion):
+ if isinstance(_bias, TensorRegion):
inst_repr, _, _ = try_find_inst_nary(
dst_buffer_region,
[src_buffer_region, _bias],
diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py
index af039e0b9c..b8be6269a3 100644
--- a/python/tvm/ir/__init__.py
+++ b/python/tvm/ir/__init__.py
@@ -43,6 +43,7 @@ from .expr import (
OpaqueExpr,
Range,
TensorLoad,
+ TensorRegion,
Tuple,
TupleGetItem,
Var,
diff --git a/python/tvm/ir/_overload_prim_expr.py b/python/tvm/ir/_overload_prim_expr.py
index 8ccd3c46ea..26eaf1332a 100644
--- a/python/tvm/ir/_overload_prim_expr.py
+++ b/python/tvm/ir/_overload_prim_expr.py
@@ -18,7 +18,7 @@
from ..runtime import DataTypeCode, ObjectConvertible, const
from .base import Span
-from .expr import Expr, is_prim_expr
+from .expr import Expr, TensorRegion, is_prim_expr
from .prim import _ffi_api as _prim_ffi_api
from .type import PrimType
@@ -52,10 +52,9 @@ def _dtype_is_float(value):
def _is_scalar_operand(value):
- type_info = getattr(type(value), "__tvm_ffi_type_info__", None)
- if type_info is not None and type_info.type_key == "tirx.BufferRegion":
+ if isinstance(value, TensorRegion):
raise TypeError(
- "BufferRegion is not a primitive operand; construct a BufferLoad explicitly"
+ "TensorRegion is not a primitive operand; construct a TensorLoad explicitly"
)
return isinstance(value, ExprOp | int | float) or is_prim_expr(value)
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 993d13bfc5..5b61fc73ef 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -502,6 +502,38 @@ class Call(_CallableExprWithOp):
self.__init_handle_by_constructor__(_ffi_api.Call, ret_ty, op, args, attrs, ty_args, span)
+@tvm_ffi.register_object("ir.TensorRegion")
+class TensorRegion(Expr, Scriptable):
+ """A region of an arbitrary tensor expression.
+
+ Parameters
+ ----------
+ source : Expr
+ The source expression.
+
+ region : list[Range]
+ The ranges describing the region.
+
+ ty : tvm.ir.Type
+ The result type, including any dialect-specific subscript semantics.
+
+ span : Span | None
+ The location of the expression in the source code.
+ """
+
+ source: Expr
+ region: list["Range"]
+
+ def __init__(
+ self,
+ source: Expr,
+ region: list["Range"],
+ ty: "tvm.ir.Type",
+ span: Span | None = None,
+ ) -> None:
+ self.__init_handle_by_constructor__(_ffi_api.TensorRegion, source, region, ty, span)
+
+
@tvm_ffi.register_object("ir.Var")
class Var(_CallableExprWithOp):
"""A canonical local variable in the IR.
diff --git a/python/tvm/runtime/script_printer.py b/python/tvm/runtime/script_printer.py
index 973069021a..07b0b2d41f 100644
--- a/python/tvm/runtime/script_printer.py
+++ b/python/tvm/runtime/script_printer.py
@@ -195,7 +195,7 @@ class Scriptable:
"""
# Auto-switch to tirx (`T`/`tirx`) flavor only when explicitly
# printing a PrimFunc / IRModule that has no s_tir-tagged content.
- # Free objects (Buffer, BufferRegion, ...) keep the default `T`/`tir`
+ # Free objects (Buffer, buffer-backed TensorRegion, ...) keep the default `T`/`tir`
# flavor -- they have no enclosing function to indicate tirx vs s_tir.
merged_extra: dict = {}
if extra_config is not None:
diff --git a/python/tvm/s_tir/analysis/__init__.py b/python/tvm/s_tir/analysis/__init__.py
index c1d4667e62..07021f092d 100644
--- a/python/tvm/s_tir/analysis/__init__.py
+++ b/python/tvm/s_tir/analysis/__init__.py
@@ -21,9 +21,9 @@
from typing import Optional, Union
import tvm
-from tvm.ir import IRModule
+from tvm.ir import IRModule, TensorRegion
from tvm.tirx.expr import Var
-from tvm.tirx.stmt import SBlock, BufferRegion
+from tvm.tirx.stmt import SBlock
from tvm.tirx import Buffer, Stmt
from tvm.tirx.function import PrimFunc
@@ -32,7 +32,7 @@ from . import _ffi_api
def get_sblock_access_region(
block: SBlock, buffer_var_map: dict[Var, Buffer]
-) -> list[list[BufferRegion]]:
+) -> list[list[TensorRegion]]:
"""Detect which regions of tensors in this block are read or written to.
Regions are sorted by order of appearance in the AST.
@@ -46,8 +46,8 @@ def get_sblock_access_region(
Returns
-------
- result : List[List[BufferRegion]]
- Array of access regions. There are three arrays of BufferRegion:
+ result : List[List[TensorRegion]]
+ Array of access regions. There are three arrays of TensorRegion:
- first: read regions
- second: write regions
- third: opaque regions
@@ -57,7 +57,7 @@ def get_sblock_access_region(
def get_sblock_read_write_region(
block: SBlock, buffer_var_map: dict[Var, Buffer]
-) -> list[list[BufferRegion]]:
+) -> list[list[TensorRegion]]:
"""Auto detect the block read/write region according to its body stmt.
An opaque access will be counted as both a read and a write access
@@ -71,7 +71,7 @@ def get_sblock_read_write_region(
Returns
-------
- result : List[List[BufferRegion]]
+ result : List[List[TensorRegion]]
An array only consisting of the read regions and write regions of the input block
"""
return _ffi_api.GetSBlockReadWriteRegion(block, buffer_var_map) # type: ignore
diff --git a/python/tvm/s_tir/dlight/adreno/utils.py b/python/tvm/s_tir/dlight/adreno/utils.py
index a820cf2b31..40601a1cd0 100644
--- a/python/tvm/s_tir/dlight/adreno/utils.py
+++ b/python/tvm/s_tir/dlight/adreno/utils.py
@@ -38,7 +38,7 @@ def get_texture_storage(block_info: SBlockInfo):
# define it uniformly for all target devices
# spatial_limit = 16384, depth_limit = 2048
# TODO: Check Write Bufs.
- shape = block_info.write_bufs[0].buf_region.buffer.shape
+ shape = block_info.write_bufs[0].buf_region.source.shape
spatial_limit = Target.current().attrs["texture_spatial_limit"]
depth_limit = Target.current().attrs["texture_depth_limit"]
diff --git a/python/tvm/s_tir/dlight/analysis/common_analysis.py b/python/tvm/s_tir/dlight/analysis/common_analysis.py
index 1c9e6edcd9..254c30ee5f 100644
--- a/python/tvm/s_tir/dlight/analysis/common_analysis.py
+++ b/python/tvm/s_tir/dlight/analysis/common_analysis.py
@@ -27,6 +27,7 @@ import tvm_ffi
from tvm_ffi import get_global_func
from tvm import ir, s_tir, tirx
+from tvm.ir import TensorRegion
from tvm.s_tir import Schedule
from tvm.s_tir.schedule import SBlockRV
from tvm.target.target import Target
@@ -79,7 +80,7 @@ BufIndex = list[Index | RemIndex | DivIndex | MergeIndex | None]
class BufferInfo:
"Information about Buffer. Provides useful analysis"
- buf_region: tirx.BufferRegion
+ buf_region: TensorRegion
shape: tuple[int]
assoc_lps: list[s_tir.schedule.LoopRV | None]
assoc_lps_info: list[tirx.For | None]
@@ -88,7 +89,7 @@ class BufferInfo:
self,
sch: s_tir.Schedule,
block_rv: s_tir.schedule.SBlockRV,
- buf_region: tirx.BufferRegion,
+ buf_region: TensorRegion,
lps: list[s_tir.schedule.LoopRV] | None,
):
block = sch.get(block_rv)
@@ -100,7 +101,7 @@ class BufferInfo:
lpvar_lp = dict([loop.loop_var, lp] for loop, lp in zip(loops, lps))
var_lp = dict(zip(iter_vars, [lpvar_lp.get(val, None) for val in iter_values]))
- def extract_index_types(buf: tirx.BufferRegion) -> BufIndex:
+ def extract_index_types(buf: TensorRegion) -> BufIndex:
buf_index = []
for expr in buf.region:
expr = expr.min
@@ -145,10 +146,10 @@ class BufferInfo:
self.buf_region = buf_region
self.assoc_lps = assoc_lps
self.assoc_lps_info = [(sch.get(lp) if lp is not None else None) for lp in assoc_lps]
- self.shape = buf_region.buffer.shape
+ self.shape = buf_region.source.shape
def get_scope(self) -> str:
- return self.buf_region.buffer.scope()
+ return self.buf_region.source.scope()
def get_vecsize(self, buf_index: int = 0, vbits: int = 128):
if self.assoc_lps_info[-1] is None:
@@ -159,7 +160,7 @@ class BufferInfo:
)
vbuf_extent = int(self.shape[-1]) & ~(int(self.shape[-1]) - 1)
- return min(vlp_extent, vbuf_extent, vbits // self.buf_region.buffer.dtype.bits)
+ return min(vlp_extent, vbuf_extent, vbits // self.buf_region.source.dtype.bits)
def __str__(self) -> str:
return f"BufferInfo({self.buf_region})"
@@ -437,7 +438,7 @@ def detect_dominant_read(block: tirx.SBlock) -> tirx.Expr:
num_read_iters = len(tir_vars)
dominant_read = buffer_region
assert dominant_read is not None
- (result,) = dominant_read.buffer.offset_of([e.min for e in dominant_read.region])
+ (result,) = dominant_read.source.offset_of([e.min for e in dominant_read.region])
return result
@@ -447,10 +448,10 @@ def is_broadcast_epilogue(
epilogue: s_tir.schedule.SBlockRV,
) -> bool:
"""Check if the epilogue block is a broadcast pattern"""
- write_buffers = {r.buffer for r in sch.get(block).writes}
+ write_buffers = {r.source for r in sch.get(block).writes}
epilogue_iters = {i.var: i for i in sch.get(epilogue).iter_vars if i.dom != 1}
for buffer_region in sch.get(epilogue).reads:
- if buffer_region.buffer not in write_buffers:
+ if buffer_region.source not in write_buffers:
continue
tir_vars = collect_block_iter_vars_used_in_access_region(
sch.get(epilogue), buffer_region.region
diff --git a/python/tvm/s_tir/dlight/analysis/gemv.py b/python/tvm/s_tir/dlight/analysis/gemv.py
index 11e280e642..6c31c2057f 100644
--- a/python/tvm/s_tir/dlight/analysis/gemv.py
+++ b/python/tvm/s_tir/dlight/analysis/gemv.py
@@ -93,7 +93,7 @@ def is_gemv(sch: s_tir.Schedule, block_info: SBlockInfo) -> list[tirx.Buffer] |
iter_num = len(block_stmt.iter_vars)
ret = [
- read.buffer
+ read.source
for read in block_stmt.reads
if len(collect_block_iter_vars_used_in_access_region(block_stmt, read.region)) < iter_num
and len(collect_block_iter_vars_used_in_access_region(block_stmt, read.region)) > 0
diff --git a/python/tvm/s_tir/dlight/cpu/reduction.py b/python/tvm/s_tir/dlight/cpu/reduction.py
index bbce215098..050e0c6a1b 100644
--- a/python/tvm/s_tir/dlight/cpu/reduction.py
+++ b/python/tvm/s_tir/dlight/cpu/reduction.py
@@ -85,7 +85,7 @@ class Reduction(CPUScheduleRule):
# Infer dtype from the last block's write buffer.
last_block_stmt = sch.get(block_infos[-1].block_rv)
- dtype_bits = last_block_stmt.writes[0].buffer.dtype.bits if last_block_stmt.writes else 32
+ dtype_bits = last_block_stmt.writes[0].source.dtype.bits if last_block_stmt.writes else 32
# Determine vector lanes from target VLEN.
vlen_bits = llvm_get_vector_width(target)
diff --git a/python/tvm/s_tir/dlight/gpu/gemv.py b/python/tvm/s_tir/dlight/gpu/gemv.py
index 1a71fa697e..808df099ae 100644
--- a/python/tvm/s_tir/dlight/gpu/gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/gemv.py
@@ -172,7 +172,7 @@ class GEMV(GPUScheduleRule):
fused_load = cache_loops[0]
else:
fused_load = sch.fuse(*cache_loops[-2:])
- aq_vec_len = max(1, VEC_LOAD // get_bytes(sch.get(Aq_local).reads[0].buffer.dtype))
+ aq_vec_len = max(1, VEC_LOAD // get_bytes(sch.get(Aq_local).reads[0].source.dtype))
fused_load, vec_load = sch.split(
fused_load, factors=[None, aq_vec_len], preserve_unit_iters=True
)
diff --git a/python/tvm/s_tir/dlight/gpu/general_reduction.py b/python/tvm/s_tir/dlight/gpu/general_reduction.py
index 235cbd34ef..22f1fde590 100644
--- a/python/tvm/s_tir/dlight/gpu/general_reduction.py
+++ b/python/tvm/s_tir/dlight/gpu/general_reduction.py
@@ -80,13 +80,13 @@ class GeneralReduction(GPUScheduleRule):
continue
other_block_buffers = [
- region.buffer
+ region.source
for other_index, other_block in enumerate(blocks)
if other_index != block_index
for region in (*other_block.reads, *other_block.writes)
]
for buffer_index, write in enumerate(block.writes):
- buffer = write.buffer
+ buffer = write.source
is_allocated = any(buffer.same_as(other) for other in alloc_buffers)
is_cross_block = any(buffer.same_as(other) for other in other_block_buffers)
if buffer.scope() == "global" and is_allocated and not is_cross_block:
@@ -152,7 +152,7 @@ class GeneralReduction(GPUScheduleRule):
reduced_buffers = []
for block_info in block_infos[:-1]:
for buffer_write in sch.get(block_info.block_rv).writes:
- reduced_buffers.append(buffer_write.buffer)
+ reduced_buffers.append(buffer_write.source)
spatial_block = sch.get(block_infos[-1].block_rv)
spatial_loops = set()
@@ -166,7 +166,7 @@ class GeneralReduction(GPUScheduleRule):
spatial_loops.add(block_var_to_loop_var[e])
for buffer_read in spatial_block.reads:
- buffer = buffer_read.buffer
+ buffer = buffer_read.source
if buffer in reduced_buffers:
for read_range in buffer_read.region:
tvm_ffi.structural_walk(
diff --git a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
index b64235224d..c997e32355 100644
--- a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
@@ -112,7 +112,7 @@ def is_gemv(sch: s_tir.Schedule, block_info: SBlockInfo) -> list[tirx.Buffer] |
if symbolic_iter_var.iter_type != tirx.stmt.IterVar.DataPar:
return None
ret = [
- read.buffer
+ read.source
for read in block_stmt.reads
if len(
collect_block_iter_vars_used_in_access_region(block_stmt, read.region) & const_iter_vars
@@ -139,7 +139,7 @@ def detect_dominant_read(block: tirx.SBlock, const_iter_vars: set[tirx.Var]) ->
num_read_iters = len(tir_vars)
dominant_read = buffer_region
assert dominant_read is not None
- (result,) = dominant_read.buffer.offset_of([e.min for e in dominant_read.region])
+ (result,) = dominant_read.source.offset_of([e.min for e in dominant_read.region])
return result
diff --git a/python/tvm/s_tir/dlight/gpu/matmul.py b/python/tvm/s_tir/dlight/gpu/matmul.py
index 7c5fd2290b..48141e182f 100644
--- a/python/tvm/s_tir/dlight/gpu/matmul.py
+++ b/python/tvm/s_tir/dlight/gpu/matmul.py
@@ -331,8 +331,8 @@ def get_in_out_dtypes(block: tirx.SBlock) -> tuple[str]:
Detect In/Out data types for the given block based on the analysis if read/write buffers.
"""
assert len(block.reads) > 0 and len(block.writes) > 0
- in_dtype = block.reads[0].buffer.dtype
- out_dtype = block.writes[0].buffer.dtype
+ in_dtype = block.reads[0].source.dtype
+ out_dtype = block.writes[0].source.dtype
return (in_dtype, out_dtype)
diff --git a/python/tvm/s_tir/schedule/schedule.py b/python/tvm/s_tir/schedule/schedule.py
index 36e590b492..e97195ad54 100644
--- a/python/tvm/s_tir/schedule/schedule.py
+++ b/python/tvm/s_tir/schedule/schedule.py
@@ -3252,9 +3252,9 @@ class Schedule(Object):
def iter_buffers():
for i, read in enumerate(block_obj.reads):
- yield "read", i, read.buffer
+ yield "read", i, read.source
for i, write in enumerate(block_obj.writes):
- yield "write", i, write.buffer
+ yield "write", i, write.source
if isinstance(buffer, int):
buffer = (required_buffer_type, buffer)
@@ -3295,7 +3295,7 @@ class Schedule(Object):
f"Block {block_name} has only "
f"{len(buffer_list)} {buffer_index_type} buffers."
)
- buffer_obj = buffer_list[buffer_index].buffer
+ buffer_obj = buffer_list[buffer_index].source
else:
raise TypeError(f"Invalid type for argument 'buffer': {type(buffer)}")
diff --git a/python/tvm/tirx/buffer.py b/python/tvm/tirx/buffer.py
index dc62b35f4c..10415d0ffa 100644
--- a/python/tvm/tirx/buffer.py
+++ b/python/tvm/tirx/buffer.py
@@ -402,7 +402,7 @@ class _BufferMethods:
"""Numpy-style view indexer: ``buf.sub[2, 4:8, ::4]``.
Unlike plain ``buf[...]`` (BufferLoad for scalar indices, extent-1
- BufferRegion dims for tile-primitive operands), ``sub`` follows numpy
+ TensorRegion dims for tile-primitive operands), ``sub`` follows numpy
basic-indexing semantics as a *view constructor*: an integer index
removes the dim (``select``), ``a:b`` narrows it, and ``a::s`` takes
every s-th element (requires the extent divisible by ``s`` and
@@ -450,7 +450,7 @@ class _BufferMethods:
the chunk index (int / Expr) and **narrows that dim** to the chunk's
``[c*E//n : (c+1)*E//n)`` range — the dim is kept at ``E // n``, no
dimension is added; an unchunked dim's pick is a normal index (``:`` /
- int / slice). The result is the *same BufferRegion* as the hand-written
+ int / slice). The result is the *same TensorRegion* as the hand-written
slice — one line instead of the ``c*k : (c+1)*k`` arithmetic::
X[.., c * k : (c + 1) * k, ..] # before (k = E // n)
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index 1b321c9c28..0494f89b46 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -25,7 +25,7 @@ from tvm_ffi import Array
import tvm
import tvm.ir.prim._ffi_api as _prim_ffi_api
from tvm import tirx
-from tvm.ir import Call, Expr, ExprWithOp, Op, PointerType, PrimType, TensorLoad
+from tvm.ir import Call, Expr, ExprWithOp, Op, PointerType, PrimType, TensorLoad, TensorRegion
from tvm.ir.base import Span
from tvm.ir.prim import clz as clz
from tvm.ir.prim import max_value, min_value
@@ -103,9 +103,9 @@ def _canonical_device_intrin_name(func_name: str) -> str:
def _reject_buffer_region(value, api_name):
"""Reject region metadata where a call argument must denote a runtime value."""
- if isinstance(value, tirx.BufferRegion):
+ if isinstance(value, TensorRegion):
raise TypeError(
- f"tirx.{api_name} does not accept BufferRegion arguments; "
+ f"tirx.{api_name} does not accept TensorRegion arguments; "
"construct a BufferLoad with explicit indices"
)
return value
diff --git a/python/tvm/tirx/operator/tile_primitive/ops.py b/python/tvm/tirx/operator/tile_primitive/ops.py
index e4be09f45f..87ff93fef3 100644
--- a/python/tvm/tirx/operator/tile_primitive/ops.py
+++ b/python/tvm/tirx/operator/tile_primitive/ops.py
@@ -169,7 +169,7 @@ class FMA(TilePrimitiveCall):
fma(output, input, scale, bias)
- scale and bias can each be either a BufferRegion or a Expr scalar.
+ scale and bias can each be either a TensorRegion or an Expr scalar.
"""
op = get_tirx_op("fma")
diff --git a/python/tvm/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py
index 3cd5d45e88..f9be5ec39d 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/python/tvm/tirx/script/builder/ir.py
@@ -25,6 +25,8 @@ from functools import partial
from numbers import Integral
from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, Union
+from tvm.ir import TensorRegion
+
# isort: off
from typing import Literal
@@ -45,7 +47,7 @@ from tvm.target import Target
# pylint: disable=unused-import
from tvm.target.codegen import llvm_lookup_intrinsic_id
-from tvm.tirx import Buffer, BufferRegion, Expr, IndexMap, is_buffer_var, type_annotation
+from tvm.tirx import Buffer, Expr, IndexMap, is_buffer_var, type_annotation
from tvm.tirx import op as _tir_op
from tvm.tirx.exec_scope import ExecScope, ScopeIdDef, Var
@@ -462,7 +464,7 @@ def Tuple(*fields: Type) -> Type: # pylint: disable=invalid-name
def match_buffer(
- param: Var | TensorLoad | BufferRegion,
+ param: Var | TensorLoad | TensorRegion,
shape: list[Expr] | tuple[Expr] | Expr | Integral = None,
dtype: str = "float32",
data: Var = None,
@@ -493,7 +495,7 @@ def match_buffer(
Parameters
----------
- param : Union[Var, TensorLoad, BufferRegion]
+ param : Union[Var, TensorLoad, TensorRegion]
The parameter of the PrimFunc to match.
shape : Union[List[Expr], Tuple[Expr], Expr, Integral]
@@ -528,9 +530,11 @@ def match_buffer(
res : Buffer
The matched buffer.
"""
+ if isinstance(param, TensorRegion) and not is_buffer_var(param.source):
+ raise TypeError("match_buffer requires a TensorRegion with a BufferVar source")
if shape is None:
- if isinstance(param, BufferRegion):
- dtype = param.buffer.ty.dtype
+ if isinstance(param, TensorRegion):
+ dtype = param.source.ty.dtype
shape = [region.extent for region in param.region]
else:
raise ValueError("Shape must be specified when binding input param")
@@ -769,12 +773,12 @@ def where(predicate: Expr | int) -> None:
_ffi_api.Where(predicate) # type: ignore[attr-defined] # pylint: disable=no-member
-def reads(*buffer_slices: list[BufferRegion | TensorLoad]) -> None:
+def reads(*buffer_slices: list[TensorRegion | TensorLoad]) -> None:
"""The block buffer region reading statement.
Parameters
----------
- buffer_slices : List[Union[BufferRegion, TensorLoad]]
+ buffer_slices : List[Union[TensorRegion, TensorLoad]]
The array of buffer regions to read.
"""
if len(buffer_slices) == 1:
@@ -789,12 +793,12 @@ def reads(*buffer_slices: list[BufferRegion | TensorLoad]) -> None:
_ffi_api.Reads(buffer_slices) # type: ignore[attr-defined] # pylint: disable=no-member
-def writes(*buffer_slices: list[BufferRegion | TensorLoad]) -> None:
+def writes(*buffer_slices: list[TensorRegion | TensorLoad]) -> None:
"""The block buffer region writing statement.
Parameters
----------
- buffer_slices : List[Union[BufferRegion, TensorLoad]]
+ buffer_slices : List[Union[TensorRegion, TensorLoad]]
The array of buffer regions to write.
"""
if len(buffer_slices) == 1:
@@ -2390,9 +2394,9 @@ def evaluate(value: Expr) -> None:
value = StringImm(value)
if isinstance(value, bool):
value = IntImm("bool", value)
- if isinstance(value, tir.BufferRegion):
+ if isinstance(value, TensorRegion):
raise TypeError(
- "T.evaluate does not accept BufferRegion values; "
+ "T.evaluate does not accept TensorRegion 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/script/builder/tirx.py b/python/tvm/tirx/script/builder/tirx.py
index 76d181621c..e8cba7d76b 100644
--- a/python/tvm/tirx/script/builder/tirx.py
+++ b/python/tvm/tirx/script/builder/tirx.py
@@ -21,8 +21,8 @@ from collections.abc import Callable
import tvm
import tvm.tirx.operator as tirx_op
-from tvm.ir import Op
-from tvm.tirx import Buffer, BufferRegion, Expr, LambdaExpr, buffer_data, is_buffer_var
+from tvm.ir import Op, TensorRegion
+from tvm.tirx import Buffer, Expr, LambdaExpr, buffer_data, is_buffer_var
from tvm.tirx.exec_scope import _SCOPE_KIND_TO_NAME, ExecScope
from tvm.tirx.expr import FloatImm, IntImm
from tvm.tirx.lang.alloc_pool import SMEMPool, TMEMPool
@@ -123,13 +123,14 @@ thread = ScopeNamespace("thread", "thread")
def _is_buffer_or_region(x):
- return is_buffer_var(x) or isinstance(x, BufferRegion)
+ return is_buffer_var(x) or isinstance(x, TensorRegion)
-def _to_region(buffer: BufferRegion | Buffer):
+def _to_region(buffer: TensorRegion | Buffer):
if is_buffer_var(buffer):
return buffer[tuple(slice(None) for _ in buffer.ty.shape)]
- assert isinstance(buffer, BufferRegion)
+ if not isinstance(buffer, TensorRegion) or not is_buffer_var(buffer.source):
+ raise TypeError("Tile operands require a Buffer or a TensorRegion with a BufferVar source")
return buffer
@@ -144,8 +145,8 @@ f_insert = _ffi_api.TilePrimitiveCall # pylint: disable=no-member
@ScopedOp
def zero(
- dst: BufferRegion | Buffer,
- src: BufferRegion | Buffer | None = None,
+ dst: TensorRegion | Buffer,
+ src: TensorRegion | Buffer | None = None,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
@@ -155,11 +156,11 @@ def zero(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region for zero result.
When src is omitted, also used as the source (in-place).
- src : Union[BufferRegion, Buffer], optional
+ src : Union[TensorRegion, Buffer], optional
The source buffer region. If omitted, dst is used (in-place).
workspace : Optional[Dict[str, Buffer]]
@@ -179,9 +180,9 @@ def zero(
@ScopedOp
def sqrt(
- dst: BufferRegion | Buffer,
- src: BufferRegion | Buffer | None = None,
- bias: BufferRegion | Buffer | FloatImm | None = None,
+ dst: TensorRegion | Buffer,
+ src: TensorRegion | Buffer | None = None,
+ bias: TensorRegion | Buffer | FloatImm | None = None,
scale: FloatImm | None = None,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
@@ -194,14 +195,14 @@ def sqrt(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region for sqrt result.
When src is omitted, also used as the source (in-place).
- src : Union[BufferRegion, Buffer], optional
+ src : Union[TensorRegion, Buffer], optional
The source buffer region. If omitted, dst is used (in-place).
- bias : Optional[Union[BufferRegion, Buffer, FloatImm]]
+ bias : Optional[Union[TensorRegion, Buffer, FloatImm]]
The bias of the sqrt src. Only supported on Trn.
scale : Optional[FloatImm]
@@ -240,9 +241,9 @@ def sqrt(
@ScopedOp
def add(
- dst: BufferRegion | Buffer,
- src1: BufferRegion | Buffer | FloatImm,
- src2: BufferRegion | Buffer | FloatImm,
+ dst: TensorRegion | Buffer,
+ src1: TensorRegion | Buffer | FloatImm,
+ src2: TensorRegion | Buffer | FloatImm,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
@@ -252,13 +253,13 @@ def add(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region for add result.
- src1 : Union[BufferRegion, Buffer, FloatImm]
+ src1 : Union[TensorRegion, Buffer, FloatImm]
The source buffer region 1, or float.
- src2 : Union[BufferRegion, Buffer, FloatImm]
+ src2 : Union[TensorRegion, Buffer, FloatImm]
The source buffer region 2, or float.
workspace : Optional[Dict[str, Buffer]]
@@ -281,9 +282,9 @@ def add(
@ScopedOp
def sub(
- dst: BufferRegion | Buffer,
- src1: BufferRegion | Buffer,
- src2: BufferRegion | Buffer | FloatImm,
+ dst: TensorRegion | Buffer,
+ src1: TensorRegion | Buffer,
+ src2: TensorRegion | Buffer | FloatImm,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
@@ -293,13 +294,13 @@ def sub(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region for sub result.
- src1 : Union[BufferRegion, Buffer]
+ src1 : Union[TensorRegion, Buffer]
The source buffer region 1.
- src2 : Union[BufferRegion, Buffer, FloatImm]
+ src2 : Union[TensorRegion, Buffer, FloatImm]
The source buffer region 2, or float.
workspace : Dict[str, Buffer]
@@ -322,9 +323,9 @@ def sub(
@ScopedOp
def mul(
- dst: BufferRegion | Buffer,
- src1: BufferRegion | Buffer | FloatImm,
- src2: BufferRegion | Buffer | FloatImm,
+ dst: TensorRegion | Buffer,
+ src1: TensorRegion | Buffer | FloatImm,
+ src2: TensorRegion | Buffer | FloatImm,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
@@ -334,13 +335,13 @@ def mul(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region for mul result.
- src1 : Union[BufferRegion, Buffer, FloatImm]
+ src1 : Union[TensorRegion, Buffer, FloatImm]
The source buffer region 1, or float.
- src2 : Union[BufferRegion, Buffer, FloatImm]
+ src2 : Union[TensorRegion, Buffer, FloatImm]
The source buffer region 2, or float.
workspace : Dict[str, Buffer]
@@ -363,9 +364,9 @@ def mul(
@ScopedOp
def fdiv(
- dst: BufferRegion | Buffer,
- src1: BufferRegion | Buffer,
- src2: BufferRegion | Buffer | FloatImm,
+ dst: TensorRegion | Buffer,
+ src1: TensorRegion | Buffer,
+ src2: TensorRegion | Buffer | FloatImm,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
@@ -375,13 +376,13 @@ def fdiv(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region for div result.
- src1 : Union[BufferRegion, Buffer]
+ src1 : Union[TensorRegion, Buffer]
The source buffer region 1.
- src2 : Union[BufferRegion, Buffer, FloatImm]
+ src2 : Union[TensorRegion, Buffer, FloatImm]
The source buffer region 2, or float.
workspace : Optional[Dict[str, Buffer]]
@@ -403,10 +404,10 @@ def fdiv(
@ScopedOp
def fma(
- dst: BufferRegion | Buffer,
- src: BufferRegion | Buffer,
- scale: BufferRegion | Buffer | Expr,
- bias: BufferRegion | Buffer | Expr,
+ dst: TensorRegion | Buffer,
+ src: TensorRegion | Buffer,
+ scale: TensorRegion | Buffer | Expr,
+ bias: TensorRegion | Buffer | Expr,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
@@ -416,16 +417,16 @@ def fma(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region.
- src : Union[BufferRegion, Buffer]
+ src : Union[TensorRegion, Buffer]
The input buffer region.
- scale : Union[BufferRegion, Buffer, Expr]
+ scale : Union[TensorRegion, Buffer, Expr]
The scale factor (buffer region or scalar).
- bias : Union[BufferRegion, Buffer, Expr]
+ bias : Union[TensorRegion, Buffer, Expr]
The bias term (buffer region or scalar).
workspace : Optional[Dict[str, Buffer]]
@@ -516,7 +517,7 @@ def _check_copy_regions_match(dst, src, config, name, dispatch=None):
def _payload_bits(region):
nonlocal analyzer
- value = tvm.DataType(region.buffer.dtype).bits
+ value = tvm.DataType(region.source.dtype).bits
for axis in region.region:
value = value * axis.extent
if analyzer is None:
@@ -583,8 +584,8 @@ def _check_copy_regions_match(dst, src, config, name, dispatch=None):
@ScopedOp
def copy(
- dst: BufferRegion | Buffer,
- src: BufferRegion | Buffer,
+ dst: TensorRegion | Buffer,
+ src: TensorRegion | Buffer,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
@@ -594,10 +595,10 @@ def copy(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region.
- src : Union[BufferRegion, Buffer]
+ src : Union[TensorRegion, Buffer]
The source buffer region.
workspace : Optional[Dict[str, Buffer]]
@@ -616,8 +617,8 @@ def copy(
@ScopedOp
def copy_async(
- dst: BufferRegion | Buffer,
- src: BufferRegion | Buffer,
+ dst: TensorRegion | Buffer,
+ src: TensorRegion | Buffer,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
@@ -638,11 +639,11 @@ def copy_async(
@ScopedOp
def gemm_async(
- C: BufferRegion | Buffer,
- A: BufferRegion | Buffer,
- B: BufferRegion | Buffer,
- SFA: BufferRegion | Buffer | None = None,
- SFB: BufferRegion | Buffer | None = None,
+ C: TensorRegion | Buffer,
+ A: TensorRegion | Buffer,
+ B: TensorRegion | Buffer,
+ SFA: TensorRegion | Buffer | None = None,
+ SFB: TensorRegion | Buffer | None = None,
transA: bool = False,
transB: bool = False,
accum: bool = False,
@@ -655,19 +656,19 @@ def gemm_async(
Parameters
----------
- C : Union[BufferRegion, Buffer]
+ C : Union[TensorRegion, Buffer]
The buffer of matrix C.
- A : Union[BufferRegion, Buffer]
+ A : Union[TensorRegion, Buffer]
The buffer of matrix A.
- B : Union[BufferRegion, Buffer]
+ B : Union[TensorRegion, Buffer]
The buffer of matrix B.
- SFA : Optional[Union[BufferRegion, Buffer]]
+ SFA : Optional[Union[TensorRegion, Buffer]]
The scale factor buffer for matrix A (block-scaled MMA only).
- SFB : Optional[Union[BufferRegion, Buffer]]
+ SFB : Optional[Union[TensorRegion, Buffer]]
The scale factor buffer for matrix B (block-scaled MMA only).
transA : bool
@@ -728,7 +729,7 @@ def gemm_async(
@ScopedOp
def fill(
- dst: BufferRegion | Buffer,
+ dst: TensorRegion | Buffer,
value: Expr,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
@@ -739,7 +740,7 @@ def fill(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region.
value : Expr
@@ -759,10 +760,10 @@ def fill(
@ScopedOp
def gemm(
- D: BufferRegion | Buffer,
- A: BufferRegion | Buffer,
- B: BufferRegion | Buffer,
- C: BufferRegion | Buffer,
+ D: TensorRegion | Buffer,
+ A: TensorRegion | Buffer,
+ B: TensorRegion | Buffer,
+ C: TensorRegion | Buffer,
transpose_A: bool = False,
transpose_B: bool = False,
alpha: Expr = 1.0,
@@ -778,16 +779,16 @@ def gemm(
Parameters
----------
- D : Union[BufferRegion, Buffer]
+ D : Union[TensorRegion, Buffer]
The buffer of matrix D.
- A : Union[BufferRegion, Buffer]
+ A : Union[TensorRegion, Buffer]
The buffer of matrix A.
- B : Union[BufferRegion, Buffer]
+ B : Union[TensorRegion, Buffer]
The buffer of matrix B.
- C : Union[BufferRegion, Buffer]
+ C : Union[TensorRegion, Buffer]
The buffer of matrix C.
transpose_A : bool
@@ -832,8 +833,8 @@ def gemm(
@ScopedOp
def sum(
- dst: BufferRegion | Buffer,
- src: BufferRegion | Buffer,
+ dst: TensorRegion | Buffer,
+ src: TensorRegion | Buffer,
axes: int | tuple[int] = -1,
accum: bool = False,
workspace: dict[str, Buffer] | None = None,
@@ -846,10 +847,10 @@ def sum(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region for sum result.
- src : Union[BufferRegion, Buffer]
+ src : Union[TensorRegion, Buffer]
The source buffer region.
axes : Union[int, Tuple[int]]
@@ -964,8 +965,8 @@ def min(
@ScopedOp
def reciprocal(
- dst: BufferRegion | Buffer,
- src: BufferRegion | Buffer | None = None,
+ dst: TensorRegion | Buffer,
+ src: TensorRegion | Buffer | None = None,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
@@ -975,11 +976,11 @@ def reciprocal(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region for reciprocal result.
When src is omitted, also used as the source (in-place).
- src : Union[BufferRegion, Buffer], optional
+ src : Union[TensorRegion, Buffer], optional
The source buffer region. If omitted, dst is used (in-place).
workspace : Optional[Dict[str, Buffer]]
@@ -1006,8 +1007,8 @@ def reciprocal(
@ScopedOp
def silu(
- dst: BufferRegion | Buffer,
- src: BufferRegion | Buffer,
+ dst: TensorRegion | Buffer,
+ src: TensorRegion | Buffer,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
@@ -1017,10 +1018,10 @@ def silu(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region for SiLU result.
- src : Union[BufferRegion, Buffer]
+ src : Union[TensorRegion, Buffer]
The source buffer region.
workspace : Optional[Dict[str, Buffer]]
@@ -1043,7 +1044,7 @@ def silu(
@ScopedOp
def memset(
- dst: BufferRegion | Buffer,
+ dst: TensorRegion | Buffer,
value: Expr,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
@@ -1054,7 +1055,7 @@ def memset(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region for memset.
value : Expr
@@ -1076,9 +1077,9 @@ def memset(
@ScopedOp
def maximum(
- dst: BufferRegion | Buffer,
- src1: BufferRegion | Buffer | FloatImm,
- src2: BufferRegion | Buffer | FloatImm,
+ dst: TensorRegion | Buffer,
+ src1: TensorRegion | Buffer | FloatImm,
+ src2: TensorRegion | Buffer | FloatImm,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
@@ -1088,13 +1089,13 @@ def maximum(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region for maximum result.
- src1 : Union[BufferRegion, Buffer, FloatImm]
+ src1 : Union[TensorRegion, Buffer, FloatImm]
The source buffer region 1, or float.
- src2 : Union[BufferRegion, Buffer, FloatImm]
+ src2 : Union[TensorRegion, Buffer, FloatImm]
The source buffer region 2, or float.
workspace : Dict[str, Buffer]
@@ -1117,9 +1118,9 @@ def maximum(
@ScopedOp
def minimum(
- dst: BufferRegion | Buffer,
- src1: BufferRegion | Buffer | FloatImm,
- src2: BufferRegion | Buffer | FloatImm,
+ dst: TensorRegion | Buffer,
+ src1: TensorRegion | Buffer | FloatImm,
+ src2: TensorRegion | Buffer | FloatImm,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
@@ -1129,13 +1130,13 @@ def minimum(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region for minimum result.
- src1 : Union[BufferRegion, Buffer, FloatImm]
+ src1 : Union[TensorRegion, Buffer, FloatImm]
The source buffer region 1, or float.
- src2 : Union[BufferRegion, Buffer, FloatImm]
+ src2 : Union[TensorRegion, Buffer, FloatImm]
The source buffer region 2, or float.
workspace : Dict[str, Buffer]
@@ -1158,9 +1159,9 @@ def minimum(
@ScopedOp
def exp(
- dst: BufferRegion | Buffer,
- src: BufferRegion | Buffer | None = None,
- bias: BufferRegion | Buffer | FloatImm | None = None,
+ dst: TensorRegion | Buffer,
+ src: TensorRegion | Buffer | None = None,
+ bias: TensorRegion | Buffer | FloatImm | None = None,
scale: FloatImm | None = None,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
@@ -1171,14 +1172,14 @@ def exp(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region for exp result.
When src is omitted, also used as the source (in-place).
- src : Union[BufferRegion, Buffer], optional
+ src : Union[TensorRegion, Buffer], optional
The source buffer region. If omitted, dst is used (in-place).
- bias : Optional[Union[BufferRegion, Buffer, FloatImm]]
+ bias : Optional[Union[TensorRegion, Buffer, FloatImm]]
The bias of the exp src. Only supported on Trn.
scale : Optional[FloatImm]
@@ -1217,9 +1218,9 @@ def exp(
@ScopedOp
def exp2(
- dst: BufferRegion | Buffer,
- src: BufferRegion | Buffer | None = None,
- bias: BufferRegion | Buffer | FloatImm | None = None,
+ dst: TensorRegion | Buffer,
+ src: TensorRegion | Buffer | None = None,
+ bias: TensorRegion | Buffer | FloatImm | None = None,
scale: FloatImm | None = None,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
@@ -1230,14 +1231,14 @@ def exp2(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region for exp2 result.
When src is omitted, also used as the source (in-place).
- src : Union[BufferRegion, Buffer], optional
+ src : Union[TensorRegion, Buffer], optional
The source buffer region. If omitted, dst is used (in-place).
- bias : Optional[Union[BufferRegion, Buffer, FloatImm]]
+ bias : Optional[Union[TensorRegion, Buffer, FloatImm]]
The bias of the exp2 src.
scale : Optional[FloatImm]
@@ -1276,9 +1277,9 @@ def exp2(
@ScopedOp
def log2(
- dst: BufferRegion | Buffer,
- src: BufferRegion | Buffer | None = None,
- bias: BufferRegion | Buffer | FloatImm | None = None,
+ dst: TensorRegion | Buffer,
+ src: TensorRegion | Buffer | None = None,
+ bias: TensorRegion | Buffer | FloatImm | None = None,
scale: FloatImm | None = None,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
@@ -1337,10 +1338,10 @@ def compose_op(
@ScopedOp
def binary_reduce(
- binary_output: BufferRegion | Buffer,
- reduce_output: BufferRegion | Buffer,
- binary_input1: BufferRegion | Buffer | FloatImm,
- binary_input2: BufferRegion | Buffer | FloatImm,
+ binary_output: TensorRegion | Buffer,
+ reduce_output: TensorRegion | Buffer,
+ binary_input1: TensorRegion | Buffer | FloatImm,
+ binary_input2: TensorRegion | Buffer | FloatImm,
binary_op: str | Op,
reduce_op: str | Op,
reduce_axes: int | tuple[int] = -1,
@@ -1353,16 +1354,16 @@ def binary_reduce(
Parameters
----------
- binary_output : Union[BufferRegion, Buffer]
+ binary_output : Union[TensorRegion, Buffer]
The destination buffer region for binary operation result.
- reduce_output : Union[BufferRegion, Buffer]
+ reduce_output : Union[TensorRegion, Buffer]
The destination buffer region for reduction result.
- binary_input1 : Union[BufferRegion, Buffer, FloatImm]
+ binary_input1 : Union[TensorRegion, Buffer, FloatImm]
The first source input for binary operation.
- binary_input2 : Union[BufferRegion, Buffer, FloatImm]
+ binary_input2 : Union[TensorRegion, Buffer, FloatImm]
The second source input for binary operation.
binary_op : Union[str, Op]
@@ -1415,12 +1416,12 @@ def binary_reduce(
@ScopedOp
def unary_reduce(
- unary_output: BufferRegion | Buffer,
- reduce_output: BufferRegion | Buffer,
- unary_input: BufferRegion | Buffer,
+ unary_output: TensorRegion | Buffer,
+ reduce_output: TensorRegion | Buffer,
+ unary_input: TensorRegion | Buffer,
unary_op: str | Op,
reduce_op: str | Op,
- bias: BufferRegion | Buffer | FloatImm | None = None,
+ bias: TensorRegion | Buffer | FloatImm | None = None,
scale: FloatImm | None = None,
reduce_axes: int | tuple[int] = -1,
workspace: dict[str, Buffer] | None = None,
@@ -1432,13 +1433,13 @@ def unary_reduce(
Parameters
----------
- unary_output : Union[BufferRegion, Buffer]
+ unary_output : Union[TensorRegion, Buffer]
The destination buffer region for unary operation result.
- reduce_output : Union[BufferRegion, Buffer]
+ reduce_output : Union[TensorRegion, Buffer]
The destination buffer region for reduction result.
- unary_input : Union[BufferRegion, Buffer]
+ unary_input : Union[TensorRegion, Buffer]
The source input for unary operation.
unary_op : Union[str, Op]
@@ -1447,7 +1448,7 @@ def unary_reduce(
reduce_op : Union[str, Op]
The reduction operation to perform.
- bias : Optional[Union[BufferRegion, Buffer, FloatImm]]
+ bias : Optional[Union[TensorRegion, Buffer, FloatImm]]
The bias to apply before unary operation.
scale : Optional[FloatImm]
@@ -1499,10 +1500,10 @@ def unary_reduce(
@ScopedOp
def binary_chain(
- output: BufferRegion | Buffer,
- data: BufferRegion | Buffer,
- operand0: BufferRegion | Buffer | FloatImm,
- operand1: BufferRegion | Buffer | FloatImm,
+ output: TensorRegion | Buffer,
+ data: TensorRegion | Buffer,
+ operand0: TensorRegion | Buffer | FloatImm,
+ operand1: TensorRegion | Buffer | FloatImm,
op0: str | Op,
op1: str | Op,
reverse1: bool = False,
@@ -1520,16 +1521,16 @@ def binary_chain(
Parameters
----------
- output : Union[BufferRegion, Buffer]
+ output : Union[TensorRegion, Buffer]
The destination buffer region for the result.
- data : Union[BufferRegion, Buffer]
+ data : Union[TensorRegion, Buffer]
The input data to operate on.
- operand0 : Union[BufferRegion, Buffer, FloatImm]
+ operand0 : Union[TensorRegion, Buffer, FloatImm]
The first operand to combine with data.
- operand1 : Union[BufferRegion, Buffer, FloatImm]
+ operand1 : Union[TensorRegion, Buffer, FloatImm]
The second operand to use in chained operation.
op0 : Union[str, Op]
@@ -1582,8 +1583,8 @@ def binary_chain(
@ScopedOp
def reduce_negate(
- output: BufferRegion | Buffer,
- input: BufferRegion | Buffer,
+ output: TensorRegion | Buffer,
+ input: TensorRegion | Buffer,
reduce_op: str | Op,
reduce_axes: int | tuple[int] = -1,
accum: bool = False,
@@ -1596,10 +1597,10 @@ def reduce_negate(
Parameters
----------
- output : Union[BufferRegion, Buffer]
+ output : Union[TensorRegion, Buffer]
The destination buffer region for the negated reduction result.
- input : Union[BufferRegion, Buffer]
+ input : Union[TensorRegion, Buffer]
The input buffer region to reduce.
reduce_axes : Union[int, Tuple[int]]
@@ -1644,9 +1645,9 @@ def reduce_negate(
@ScopedOp
def select(
- dst: BufferRegion | Buffer,
- true_value: BufferRegion | Buffer | FloatImm,
- false_value: BufferRegion | Buffer | FloatImm,
+ dst: TensorRegion | Buffer,
+ true_value: TensorRegion | Buffer | FloatImm,
+ false_value: TensorRegion | Buffer | FloatImm,
pred: LambdaExpr | Callable[..., Expr],
scope: ExecScope | None = None,
):
@@ -1654,13 +1655,13 @@ def select(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
The destination buffer region for the result.
- true_value : Union[BufferRegion, Buffer, FloatImm]
+ true_value : Union[TensorRegion, Buffer, FloatImm]
The value to select if the predicate is true.
- false_value : Union[BufferRegion, Buffer, FloatImm]
+ false_value : Union[TensorRegion, Buffer, FloatImm]
The value to select if the predicate is false.
pred : Union[LambdaExpr, Callable[..., Expr]]
@@ -1711,8 +1712,8 @@ def reshape(buffer: Buffer, shape: list[Expr]):
@ScopedOp
def permute_layout(
- dst: BufferRegion | Buffer,
- src: BufferRegion | Buffer,
+ dst: TensorRegion | Buffer,
+ src: TensorRegion | Buffer,
workspace: dict[str, Buffer] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
@@ -1726,9 +1727,9 @@ def permute_layout(
Parameters
----------
- dst : Union[BufferRegion, Buffer]
+ dst : Union[TensorRegion, Buffer]
Destination view (carries the target layout).
- src : Union[BufferRegion, Buffer]
+ src : Union[TensorRegion, Buffer]
Source view (carries the current layout).
workspace : Dict[str, Buffer]
Optional workspace for the operator.
@@ -1736,7 +1737,7 @@ def permute_layout(
Force a specific dispatch variant by name.
"""
- # Promote Buffer to BufferRegion covering the full extent, matching the
+ # Promote Buffer to TensorRegion covering the full extent, matching the
# convention used by ``Tx.<dynamic>`` fallback registration.
def _to_region(b):
if is_buffer_var(b):
diff --git a/python/tvm/tirx/script/parser/parser.py b/python/tvm/tirx/script/parser/parser.py
index e9c9aafc43..753e46c413 100644
--- a/python/tvm/tirx/script/parser/parser.py
+++ b/python/tvm/tirx/script/parser/parser.py
@@ -23,7 +23,7 @@ from functools import partial
from typing import Any, TypeVar
import tvm
-from tvm.ir import Expr, GlobalVar, PointerType, PrimType, TensorLoad
+from tvm.ir import Expr, GlobalVar, PointerType, PrimType, TensorLoad, TensorRegion
from tvm.script.ir_builder import ir as I
from tvm.script.ir_builder.base import IRBuilder
from tvm.script.ir_builder.base import IRBuilderFrame as Frame
@@ -32,21 +32,22 @@ from tvm.script.parser.core.doc import from_doc
from tvm.tirx import Buffer, IterVar, Layout, buffer_data, is_buffer_var
from tvm.tirx.script import builder as T
from tvm.tirx.script.builder.ir import name_meta_class_value
-from tvm.tirx.stmt import BufferRegion
from .entry import _OptionalAnnotation, inline
from .entry import constexpr as _constexpr_sentinel
-def slice_buffer_from_region(br: BufferRegion) -> Buffer:
- """Create a matched DeclBuffer from a BufferRegion.
+def slice_buffer_from_region(br: TensorRegion) -> Buffer:
+ """Create a matched DeclBuffer from a TensorRegion.
Slices the layout (if present) or computes elem_offset for the sub-region,
producing a DeclBuffer that views the same underlying data.
"""
import functools # pylint: disable=import-outside-toplevel
- buf = br.buffer
+ if not is_buffer_var(br.source):
+ raise TypeError("A matched buffer region requires a BufferVar source")
+ buf = br.source
region = br.region
new_shape = [r.extent for r in region]
sliced_layout = None
@@ -244,7 +245,7 @@ def bind_assign_value(
for i, v in enumerate(value):
bind_assign_value(self, node, f"{var_name}_{i}", v)
return value
- elif isinstance(value, BufferRegion):
+ elif isinstance(value, TensorRegion):
return value
elif isinstance(value, Frame):
value.add_callback(partial(value.__exit__, None, None, None))
diff --git a/python/tvm/tirx/script/tile.py b/python/tvm/tirx/script/tile.py
index 4aaf79113b..ac9a29c478 100644
--- a/python/tvm/tirx/script/tile.py
+++ b/python/tvm/tirx/script/tile.py
@@ -18,7 +18,8 @@
import functools
-from tvm.tirx import BufferRegion, is_buffer_var
+from tvm.ir import TensorRegion
+from tvm.tirx import is_buffer_var
from .builder import tirx as _builder
@@ -30,10 +31,12 @@ def _get_arg(args, kwargs, index, name):
def _require_buffer_arg(op_name, arg_name, value):
- if not (is_buffer_var(value) or isinstance(value, BufferRegion)):
+ if not (
+ is_buffer_var(value) or (isinstance(value, TensorRegion) and is_buffer_var(value.source))
+ ):
raise TypeError(
f"Tx.{op_name} is tile-only and expects `{arg_name}` to be a Buffer "
- f"or BufferRegion; use T.{op_name} for expression/builtin calls"
+ f"or TensorRegion with a BufferVar source; use T.{op_name} for expression/builtin calls"
)
diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py
index 46c7931c58..5d15b227af 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, Type
+from tvm.ir import Expr, Range, Span, TensorRegion, Type
from tvm.runtime import Object, Scriptable, const
from . import _ffi_api
@@ -608,30 +608,24 @@ class Evaluate(Stmt):
@tvm_ffi.register_object("tirx.BufferRegionType")
class BufferRegionType(Type):
- """The structural type of a :class:`BufferRegion` expression."""
+ """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
-@tvm_ffi.register_object("tirx.BufferRegion")
-class BufferRegion(Expr, Scriptable):
- """BufferRegion node.
+def BufferRegion(buffer: Buffer, region: list[Range]) -> TensorRegion:
+ """Construct a buffer-backed tensor region with TIRX subscript semantics.
Parameters
----------
buffer : Buffer
- The buffer of the buffer region
+ The source buffer.
region : List[Range]
- The region array of the buffer region
+ The ranges, with one entry for each buffer dimension.
"""
-
- buffer: Buffer
- region: list[Range]
-
- def __init__(self, buffer: Buffer, region: list[Range]) -> None:
- self.__init_handle_by_constructor__(_ffi_api.BufferRegion, buffer, region) # type: ignore
+ return _ffi_api.BufferRegion(buffer, region)
@tvm_ffi.register_object("tirx.MatchBufferRegion")
@@ -643,14 +637,14 @@ class MatchBufferRegion(Object, Scriptable):
buffer : Buffer
The target buffer
- source : BufferRegion
+ source : TensorRegion
The region of source buffer
"""
buffer: Buffer
- source: BufferRegion
+ source: TensorRegion
- def __init__(self, buffer: Buffer, source: BufferRegion) -> None:
+ def __init__(self, buffer: Buffer, source: TensorRegion) -> None:
self.__init_handle_by_constructor__(
_ffi_api.MatchBufferRegion,
buffer,
@@ -667,10 +661,10 @@ class SBlock(Stmt):
iter_vars : List[IterVar]
The block Variable.
- reads : List[BufferRegion]
+ reads : List[TensorRegion]
The read buffer regions of the block.
- writes: List[BufferRegion]
+ writes: List[TensorRegion]
The write buffer regions of the block.
name_hint: str
@@ -696,8 +690,8 @@ class SBlock(Stmt):
"""
iter_vars: list[IterVar]
- reads: list[BufferRegion]
- writes: list[BufferRegion]
+ reads: list[TensorRegion]
+ writes: list[TensorRegion]
name_hint: str
body: Stmt
init: Stmt | None
@@ -709,8 +703,8 @@ class SBlock(Stmt):
def __init__(
self,
iter_vars: list[IterVar],
- reads: list[BufferRegion],
- writes: list[BufferRegion],
+ reads: list[TensorRegion],
+ writes: list[TensorRegion],
name_hint: str,
body: Stmt,
init: Stmt | None = None,
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index cf2ffbeb26..6f1f224a2a 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -126,6 +126,57 @@ TVMFFIAny TensorLoadMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
return ffi::Unchanged().CopyToTVMFFIAny();
}
+TVMFFIAny TensorRegionVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+ const TensorRegionNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorRegionNode>(value);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->source));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->region));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny TensorRegionMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+ const TensorRegionNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorRegionNode>(value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty_u,
+ mutator->MutateExpected(self->ty));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_source_u,
+ mutator->MutateExpected(self->source));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<Range>>, mapped_region_u,
+ mutator->MutateExpected(self->region));
+ if (mapped_ty_u.UnchangedOrSameAs(self->ty) && mapped_source_u.UnchangedOrSameAs(self->source) &&
+ mapped_region_u.UnchangedOrSameAs(self->region)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::ObjectPtr<TensorRegionNode> copy = ffi::make_object<TensorRegionNode>(*self);
+ if (!mapped_ty_u.IsUnchanged()) copy->ty = std::move(mapped_ty_u).ValueUnchecked();
+ if (!mapped_source_u.IsUnchanged()) copy->source = std::move(mapped_source_u).ValueUnchecked();
+ if (!mapped_region_u.IsUnchanged()) copy->region = std::move(mapped_region_u).ValueUnchecked();
+ return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny TensorRegionMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ TensorRegionNode* self = const_cast<TensorRegionNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorRegionNode>(value));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty_u,
+ mutator->MutateExpected(self->ty, ffi::InplaceMode::kAllow));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ ffi::UnchangedOr<Expr>, mapped_source_u,
+ mutator->MutateExpected(self->source, ffi::InplaceMode::kAllow));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ ffi::UnchangedOr<ffi::Array<Range>>, mapped_region_u,
+ mutator->MutateExpected(self->region, ffi::InplaceMode::kAllow));
+ if (mapped_ty_u.UnchangedOrSameAs(self->ty) && mapped_source_u.UnchangedOrSameAs(self->source) &&
+ mapped_region_u.UnchangedOrSameAs(self->region)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ if (!mapped_ty_u.IsUnchanged()) self->ty = std::move(mapped_ty_u).ValueUnchecked();
+ if (!mapped_source_u.IsUnchanged()) self->source = std::move(mapped_source_u).ValueUnchecked();
+ if (!mapped_region_u.IsUnchanged()) self->region = std::move(mapped_region_u).ValueUnchecked();
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
TVMFFIAny TupleVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const TupleNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TupleNode>(value);
@@ -537,6 +588,29 @@ TVM_FFI_STATIC_INIT_BLOCK() {
reinterpret_cast<void*>(&TensorLoadMaybeInplaceMutate));
}
+TensorRegion::TensorRegion(Expr source, ffi::Array<Range> region, Type ty, Span span) {
+ auto node = ffi::make_object<TensorRegionNode>();
+ node->source = std::move(source);
+ node->region = std::move(region);
+ node->ty = std::move(ty);
+ node->span = std::move(span);
+ data_ = std::move(node);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ TensorRegionNode::RegisterReflection();
+ refl::TypeAttrDef<TensorRegionNode>()
+ .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&TensorRegionVisit))
+ .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&TensorRegionMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&TensorRegionMaybeInplaceMutate));
+ refl::GlobalDef().def("ir.TensorRegion",
+ [](Expr source, ffi::Array<Range> region, Type ty, Span span) {
+ return TensorRegion(source, region, ty, span);
+ });
+}
+
// Tuple
Tuple::Tuple(ffi::Array<Expr> fields, Span span) {
ffi::Optional<Type> tuple_ty = [&]() -> ffi::Optional<Type> {
diff --git a/src/ir/expr_functor.cc b/src/ir/expr_functor.cc
index 7d12787bf5..f00e5d92d1 100644
--- a/src/ir/expr_functor.cc
+++ b/src/ir/expr_functor.cc
@@ -26,6 +26,7 @@ void ExprVisitor::InitVTable(VTable* vtable) {
SetDispatch<ExprVisitor, TupleNode>(vtable);
SetDispatch<ExprVisitor, TupleGetItemNode>(vtable);
SetDispatch<ExprVisitor, TensorLoadNode>(vtable);
+ SetDispatch<ExprVisitor, TensorRegionNode>(vtable);
SetDispatch<ExprVisitor, VarNode>(vtable);
SetDispatch<ExprVisitor, GlobalVarNode>(vtable);
SetDispatch<ExprVisitor, CallNode>(vtable);
@@ -84,6 +85,13 @@ ffi::Optional<VisitInterrupt> ExprVisitor::Visit_(const TensorLoadNode* node) {
return std::nullopt;
}
+ffi::Optional<VisitInterrupt> ExprVisitor::Visit_(const TensorRegionNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(node->ty));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(node->source));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(node->region));
+ return std::nullopt;
+}
+
ffi::Optional<VisitInterrupt> ExprVisitor::Visit_(const VarNode* node) {
// Primitive types have no children; dynamic type fields are visited.
if (!node->ty.as<PrimTypeNode>()) {
@@ -281,6 +289,7 @@ void ExprMutator::InitVTable(VTable* vtable) {
SetDispatch<ExprMutator, TupleNode>(vtable);
SetDispatch<ExprMutator, TupleGetItemNode>(vtable);
SetDispatch<ExprMutator, TensorLoadNode>(vtable);
+ SetDispatch<ExprMutator, TensorRegionNode>(vtable);
SetDispatch<ExprMutator, VarNode>(vtable);
SetDispatch<ExprMutator, GlobalVarNode>(vtable);
SetDispatch<ExprMutator, CallNode>(vtable);
@@ -383,6 +392,27 @@ UnchangedOr<PrimExpr> ExprMutator::Mutate_(const TensorLoadNode* node, InplaceMo
return PrimExpr(std::move(copy));
}
+UnchangedOr<Expr> ExprMutator::Mutate_(const TensorRegionNode* node, InplaceMode inplace_mode) {
+ auto ty_u = Mutate(node->ty, inplace_mode).as_or_throw<UnchangedOr<Type>>();
+ auto source_u = Mutate(node->source, inplace_mode);
+ auto region_u = Mutate(node->region, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<Range>>>();
+ if (ty_u.UnchangedOrSameAs(node->ty) && source_u.UnchangedOrSameAs(node->source) &&
+ region_u.UnchangedOrSameAs(node->region))
+ return ffi::Unchanged();
+ if (inplace_mode == InplaceMode::kAllow) {
+ auto* writable = const_cast<TensorRegionNode*>(node);
+ if (!ty_u.IsUnchanged()) writable->ty = std::move(ty_u).ValueUnchecked();
+ if (!source_u.IsUnchanged()) writable->source = std::move(source_u).ValueUnchecked();
+ if (!region_u.IsUnchanged()) writable->region = std::move(region_u).ValueUnchecked();
+ return ffi::Unchanged();
+ }
+ auto copy = ffi::make_object<TensorRegionNode>(*node);
+ if (!ty_u.IsUnchanged()) copy->ty = std::move(ty_u).ValueUnchecked();
+ if (!source_u.IsUnchanged()) copy->source = std::move(source_u).ValueUnchecked();
+ if (!region_u.IsUnchanged()) copy->region = std::move(region_u).ValueUnchecked();
+ return Expr(std::move(copy));
+}
+
UnchangedOr<Expr> ExprMutator::Mutate_(const GlobalVarNode* node, InplaceMode inplace_mode) {
// Registry atoms and constant leaves do not descend into metadata or types.
return ffi::Unchanged();
diff --git a/src/relax/analysis/layout_transformation.cc b/src/relax/analysis/layout_transformation.cc
index 7c8309c8d8..92cf537bee 100644
--- a/src/relax/analysis/layout_transformation.cc
+++ b/src/relax/analysis/layout_transformation.cc
@@ -336,7 +336,7 @@ class BlockAnalyzer : public StmtExprVisitor {
void Analyze() {
const auto& block = block_;
TVM_FFI_ICHECK(block_->writes.size() == 1);
- auto write_buffer = block_->writes[0]->buffer;
+ auto write_buffer = block_->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>();
ComputeBlockSpatialDomain();
@@ -402,7 +402,8 @@ class BlockAnalyzer : public StmtExprVisitor {
// Infer read buffer transformations from write buffer transformation.
for (const auto& r : block->reads) {
- SpatialLayout read_spatial_layout = get_spatial_layout(r->buffer);
+ SpatialLayout read_spatial_layout =
+ get_spatial_layout(r->source.as_or_throw<tvm::tirx::BufferVar>());
if (read_spatial_layout.empty()) continue;
if (!IsSequentialAccess(read_spatial_layout, iter_var_to_block_index)) continue;
@@ -410,14 +411,19 @@ class BlockAnalyzer : public StmtExprVisitor {
write_spatial_layout, write_transformation_, read_spatial_layout);
if (!maybe_read_transformation.has_value()) continue;
IndexMap read_transformation = maybe_read_transformation.value();
- if (buffer_transformation_cache_.count(r->buffer) != 0) {
- if (!AreIdenticalTransforms(read_transformation, buffer_transformation_cache_[r->buffer]))
- LOG(WARNING) << "[LayoutInference] Buffer: " << r->buffer
- << " has conflicting transform proposals -- (preferred) "
- << buffer_transformation_cache_[r->buffer] << " vs. " << read_transformation;
+ if (buffer_transformation_cache_.count(r->source.as_or_throw<tvm::tirx::BufferVar>()) != 0) {
+ if (!AreIdenticalTransforms(
+ read_transformation,
+ buffer_transformation_cache_[r->source.as_or_throw<tvm::tirx::BufferVar>()]))
+ LOG(WARNING)
+ << "[LayoutInference] Buffer: " << r->source.as_or_throw<tvm::tirx::BufferVar>()
+ << " has conflicting transform proposals -- (preferred) "
+ << buffer_transformation_cache_[r->source.as_or_throw<tvm::tirx::BufferVar>()]
+ << " vs. " << read_transformation;
continue;
}
- read_buffer_transformations_.Set(r->buffer, read_transformation);
+ read_buffer_transformations_.Set(r->source.as_or_throw<tvm::tirx::BufferVar>(),
+ read_transformation);
}
}
@@ -492,11 +498,12 @@ class BlockAnalyzer : public StmtExprVisitor {
if (!access_info.IsValid()) return std::nullopt;
// Only single write buffer is supported for each block.
- if (!op->buffer.same_as(block_->writes[0]->buffer)) {
+ if (!op->buffer.same_as(block_->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>())) {
access_info.Invalidate();
LOG(WARNING) << "[LayoutInference] Exactly one write buffer is supported for layout "
"inference, found two: "
- << op->buffer << " and " << block_->writes[0]->buffer;
+ << op->buffer << " and "
+ << block_->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>();
can_transform_block_ = false;
return std::nullopt;
}
@@ -601,7 +608,7 @@ class PrimFuncAnalyzer : public StmtExprVisitor {
SBlock block = ffi::GetRef<SBlock>(op);
// Get block write buffer transformation.
if (block->writes.size() != 1) return std::nullopt;
- auto write_buffer = block->writes[0]->buffer;
+ auto write_buffer = block->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>();
block_to_buffer_[block].push_back(write_buffer);
auto block_analyzer = ffi::make_object<BlockAnalyzer>(
block, buffer_transformation_cache_, buffer_transformation_cache_[write_buffer]);
diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc
index b7904282fd..21a4e45bc6 100644
--- a/src/relax/analysis/tir_op_pattern_kind.cc
+++ b/src/relax/analysis/tir_op_pattern_kind.cc
@@ -51,8 +51,8 @@ class PatternKindAnalyzer : public StmtExprVisitor {
private:
bool IsOutputBlock(const SBlockNode* block) {
- for (const BufferRegion& write_region : block->writes) {
- if (param_buffers_.count(write_region->buffer)) {
+ for (const TensorRegion& write_region : block->writes) {
+ if (param_buffers_.count(write_region->source.as_or_throw<tvm::tirx::BufferVar>())) {
return true;
}
}
diff --git a/src/relax/distributed/transform/lower_global_view_to_local_view.cc b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
index 6f4278c09f..42efd11af5 100644
--- a/src/relax/distributed/transform/lower_global_view_to_local_view.cc
+++ b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
@@ -68,7 +68,7 @@ class DistSBlockInfoCollector : public StmtExprVisitor {
for (const auto& iter_var : op->iter_vars) {
if (iter_var->iter_type == kCommReduce) {
TVM_FFI_ICHECK(op->writes.size() == 1);
- reduce_buffer_ = op->writes[0]->buffer;
+ reduce_buffer_ = op->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>();
}
}
return StmtExprVisitor::Visit_(op);
@@ -188,10 +188,10 @@ class DistributedBufferCompactor : public StmtExprMutator {
ffi::ObjectPtrEqual>& buffer_access_indices) {
std::vector<BufferVar> buffers;
for (const auto& read : block->reads) {
- buffers.push_back(read->buffer);
+ buffers.push_back(read->source.as_or_throw<tvm::tirx::BufferVar>());
}
for (const auto& write : block->writes) {
- buffers.push_back(write->buffer);
+ buffers.push_back(write->source.as_or_throw<tvm::tirx::BufferVar>());
}
ffi::Map<Var, Range> iter_var_range;
for (const auto& iter_var : block->iter_vars) {
diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc
index 9602c4fadb..859512a11a 100644
--- a/src/relax/transform/fuse_tir.cc
+++ b/src/relax/transform/fuse_tir.cc
@@ -196,8 +196,8 @@ class FuseTIRBufferSubstitutor : public StmtExprMutator {
SBlock block = StmtExprMutator::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
.as_or_throw<SBlock>();
- ffi::Array<BufferRegion> reads = UnionAccessRegion(block->reads);
- ffi::Array<BufferRegion> writes = UnionAccessRegion(block->writes);
+ ffi::Array<TensorRegion> reads = UnionAccessRegion(block->reads);
+ ffi::Array<TensorRegion> writes = UnionAccessRegion(block->writes);
if (!reads.same_as(block->reads) || !writes.same_as(block->writes)) {
auto* n = block.CopyOnWrite();
n->reads = std::move(reads);
@@ -206,19 +206,20 @@ class FuseTIRBufferSubstitutor : public StmtExprMutator {
return block;
}
- ffi::Array<tirx::BufferRegion> UnionAccessRegion(const ffi::Array<BufferRegion>& regions) const {
+ ffi::Array<tvm::TensorRegion> UnionAccessRegion(const ffi::Array<TensorRegion>& regions) const {
// For now we only allow buffers to access the same elements.
// e.g. `[A[vi, vj], A[vi, vj]]` is a legal pattern but need to union to `A[vi, vj]`
// However, `A[vi, vj], A[vi, vj + 1]` is not allow for now.
// Note: the order of return region should remain the same as the first occurrence of the region
- ffi::Array<BufferRegion> ret;
+ ffi::Array<TensorRegion> ret;
std::unordered_map<const VarNode*, ffi::Array<Range>> buffer_region_set;
- for (const BufferRegion& region : regions) {
- auto it = buffer_region_set.find(region->buffer.get());
+ for (const TensorRegion& region : regions) {
+ auto it = buffer_region_set.find(region->source.as_or_throw<tvm::tirx::BufferVar>().get());
if (it == buffer_region_set.end()) {
ret.push_back(region);
- buffer_region_set[region->buffer.get()] = region->region;
+ buffer_region_set[region->source.as_or_throw<tvm::tirx::BufferVar>().get()] =
+ region->region;
}
}
diff --git a/src/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc
index 9bc1b14948..5f24da7cec 100644
--- a/src/relax/transform/split_call_tir_by_pattern.cc
+++ b/src/relax/transform/split_call_tir_by_pattern.cc
@@ -346,8 +346,9 @@ class ForMatcher : public TensorizeComparator {
return equal;
}
- bool CompareBufferRegion(const BufferRegion& lhs, const BufferRegion& rhs) {
- if (!CompareBuffer(lhs->buffer, rhs->buffer)) {
+ bool CompareBufferRegion(const TensorRegion& lhs, const TensorRegion& rhs) {
+ if (!CompareBuffer(lhs->source.as_or_throw<tvm::tirx::BufferVar>(),
+ rhs->source.as_or_throw<tvm::tirx::BufferVar>())) {
return false;
}
return CompareArray(lhs->region, rhs->region, &ForMatcher::CompareRange);
@@ -487,24 +488,24 @@ class FunctionPartitioner : public StmtExprVisitor {
}
for (const auto& read : op->reads) {
if (is_matching_) {
- input1.insert(read->buffer);
+ input1.insert(read->source.as_or_throw<tvm::tirx::BufferVar>());
} else {
- input2.insert(read->buffer);
+ input2.insert(read->source.as_or_throw<tvm::tirx::BufferVar>());
}
}
for (const auto& write : op->writes) {
if (is_matching_) {
- allocs1.insert(write->buffer);
- } else if (allocs1.count(write->buffer)) {
+ allocs1.insert(write->source.as_or_throw<tvm::tirx::BufferVar>());
+ } else if (allocs1.count(write->source.as_or_throw<tvm::tirx::BufferVar>())) {
fail = true;
return std::nullopt;
} else {
- allocs2.insert(write->buffer);
+ allocs2.insert(write->source.as_or_throw<tvm::tirx::BufferVar>());
}
if (is_matching_) {
- intermediate_buffer = write->buffer;
+ intermediate_buffer = write->source.as_or_throw<tvm::tirx::BufferVar>();
} else {
- input2.insert(write->buffer);
+ input2.insert(write->source.as_or_throw<tvm::tirx::BufferVar>());
}
}
block_partition.Set(ffi::GetRef<SBlock>(op), is_matching_);
diff --git a/src/relax/transform/split_layout_rewrite_preproc.cc b/src/relax/transform/split_layout_rewrite_preproc.cc
index 53466dbf03..845900ecc5 100644
--- a/src/relax/transform/split_layout_rewrite_preproc.cc
+++ b/src/relax/transform/split_layout_rewrite_preproc.cc
@@ -196,7 +196,7 @@ class SplitPrimFuncLayoutRewrite : public StmtExprMutator {
<< "There should be no alloc buffer in the layout rewrite";
TVM_FFI_ICHECK(op->match_buffers.empty())
<< "There should be no match buffer in the layout rewrite";
- const BufferVar& preproc_buffer = op->reads[0]->buffer;
+ const BufferVar& preproc_buffer = op->reads[0]->source.as_or_throw<tvm::tirx::BufferVar>();
int buffer_index = -1;
for (size_t i = 0; i < original_func_->params.size(); ++i) {
BufferVar buffer = original_func_->params[i].as_or_throw<tvm::tirx::BufferVar>();
@@ -208,7 +208,8 @@ class SplitPrimFuncLayoutRewrite : public StmtExprMutator {
TVM_FFI_ICHECK(buffer_index != -1)
<< "The preproc buffer is not found in the original primfunc.";
rewrite_infos_.push_back(
- RewriteInfo{buffer_index, op->reads[0]->buffer, op->writes[0]->buffer});
+ RewriteInfo{buffer_index, op->reads[0]->source.as_or_throw<tvm::tirx::BufferVar>(),
+ op->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>()});
auto new_annotations = op->annotations;
new_annotations.erase(s_tir::attr::meta_schedule_layout_rewrite_preproc);
diff --git a/src/s_tir/analysis/identify_memcpy.cc b/src/s_tir/analysis/identify_memcpy.cc
index 83f343c76b..f965745fa6 100644
--- a/src/s_tir/analysis/identify_memcpy.cc
+++ b/src/s_tir/analysis/identify_memcpy.cc
@@ -273,10 +273,11 @@ std::variant<MemCpyDetails, std::string> IdentifyMemCpyImpl(const For& loop,
}
}
- BufferRegion src_region(
+ TensorRegion src_region = BufferRegion(
load->source.as_or_throw<tvm::tirx::BufferVar>(),
DomainTouched(loop, load->source.as_or_throw<tvm::tirx::BufferVar>(), true, true));
- BufferRegion dst_region(store->buffer, DomainTouched(loop, store->buffer, true, true));
+ TensorRegion dst_region =
+ BufferRegion(store->buffer, DomainTouched(loop, store->buffer, true, true));
return MemCpyDetails{src_region, dst_region};
}
diff --git a/src/s_tir/analysis/sblock_access_region_detector.cc b/src/s_tir/analysis/sblock_access_region_detector.cc
index f15382fdf3..574401da95 100644
--- a/src/s_tir/analysis/sblock_access_region_detector.cc
+++ b/src/s_tir/analysis/sblock_access_region_detector.cc
@@ -56,17 +56,17 @@ class BlockReadWriteDetector : public StmtExprVisitor {
}
/*! \brief Return read regions of the block */
- ffi::Array<BufferRegion> CollectReads(
+ ffi::Array<TensorRegion> CollectReads(
const std::unordered_set<const VarNode*>* excluded_buffers = nullptr);
/*! \brief Return write regions of the block */
- ffi::Array<BufferRegion> CollectWrites(
+ ffi::Array<TensorRegion> CollectWrites(
const std::unordered_set<const VarNode*>* excluded_buffers = nullptr);
/*!
* \brief Return opaque buffer regions of the block
* \note The buffer accessed by load/store or call with buffer.data will
* be marked as opaque.
*/
- ffi::Array<BufferRegion> CollectOpaques();
+ ffi::Array<TensorRegion> CollectOpaques();
/*! \brief overload operator() to make sure it accepts a block node */
void operator()(const Stmt& stmt);
@@ -109,7 +109,7 @@ class BlockReadWriteDetector : public StmtExprVisitor {
BufferVar buffer, std::vector<arith::IntSet> region);
/*! \brief Helper function to collect access regions. */
- ffi::Array<BufferRegion> CollectRegions(
+ ffi::Array<TensorRegion> CollectRegions(
const std::vector<BufferVar>& buffers,
const std::vector<std::vector<tvm::arith::IntSet>>& regions,
const std::unordered_set<const VarNode*>* excluded_buffers = nullptr);
@@ -125,7 +125,8 @@ class BlockReadWriteDetector : public StmtExprVisitor {
arith::IntSet RelaxAccessIndex(const PrimExpr& index);
// Declared regions carry bounds, not opaque runtime accesses.
- ffi::Optional<VisitInterrupt> Visit_(const BufferRegionNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorRegionNode* op) final {
+ if (!op->source.as<BufferVar>()) return StmtExprVisitor::Visit_(op);
for (const Range& range : op->region) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(range->min));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(range->extent));
@@ -150,7 +151,7 @@ void BlockReadWriteDetector::operator()(const Stmt& stmt) {
<< "Only visiting Blocks is allowed, but got " << stmt->GetTypeKey();
for (const MatchBufferRegion& match_buffer : block->match_buffers) {
const Var target_var = match_buffer->buffer.var();
- const Var source_var = match_buffer->source->buffer.var();
+ const Var source_var = match_buffer->source->source.as_or_throw<tvm::tirx::BufferVar>().var();
if (buffer_var_map_.find(source_var) != buffer_var_map_.end()) {
match_buffers_[target_var.get()] = match_buffer;
buffer_var_map_.Set(target_var, match_buffer->buffer);
@@ -159,17 +160,17 @@ void BlockReadWriteDetector::operator()(const Stmt& stmt) {
StmtExprVisitor::Visit(stmt);
}
-ffi::Array<BufferRegion> BlockReadWriteDetector::CollectReads(
+ffi::Array<TensorRegion> BlockReadWriteDetector::CollectReads(
const std::unordered_set<const VarNode*>* excluded_buffers) {
return CollectRegions(read_buffers_, read_regions_, excluded_buffers);
}
-ffi::Array<BufferRegion> BlockReadWriteDetector::CollectWrites(
+ffi::Array<TensorRegion> BlockReadWriteDetector::CollectWrites(
const std::unordered_set<const VarNode*>* excluded_buffers) {
return CollectRegions(writes_buffers_, write_regions_, excluded_buffers);
}
-ffi::Array<BufferRegion> BlockReadWriteDetector::CollectOpaques() {
+ffi::Array<TensorRegion> BlockReadWriteDetector::CollectOpaques() {
return CollectRegions(opaque_buffers_, opaque_regions_);
}
@@ -294,7 +295,7 @@ ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const CallNode* op)
auto it = buffer_var_map_.find(ffi::GetRef<Var>(buffer_var));
if (it != buffer_var_map_.end()) {
const BufferVar& buffer = (*it).second;
- const BufferRegion buffer_region = BufferRegion::FullRegion(buffer);
+ const TensorRegion buffer_region = FullBufferRegion(buffer);
const ffi::Array<Range>& region = buffer_region->region;
std::vector<arith::IntSet> int_set;
int_set.reserve(region.size());
@@ -383,7 +384,8 @@ ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const SBlockRealize
relaxed_region.push_back(
arith::EvalSet(arith::IntSet::FromRange(Range::FromMinExtent(min, extent)), dom_map_));
}
- Update(&read_buffers_, &read_regions_, read->buffer, relaxed_region);
+ Update(&read_buffers_, &read_regions_, read->source.as_or_throw<tvm::tirx::BufferVar>(),
+ relaxed_region);
}
for (const auto& write : op->block->writes) {
std::vector<arith::IntSet> relaxed_region;
@@ -395,7 +397,8 @@ ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const SBlockRealize
relaxed_region.push_back(
arith::EvalSet(arith::IntSet::FromRange(Range::FromMinExtent(min, extent)), dom_map_));
}
- Update(&writes_buffers_, &write_regions_, write->buffer, relaxed_region);
+ Update(&writes_buffers_, &write_regions_, write->source.as_or_throw<tvm::tirx::BufferVar>(),
+ relaxed_region);
}
return std::nullopt;
}
@@ -430,7 +433,7 @@ void BlockReadWriteDetector::Update(std::vector<BufferVar>* buffers,
auto it = match_buffers_.find(buffer.get());
if (it != match_buffers_.end()) {
const MatchBufferRegion& match_buffer = it->second;
- buffer = match_buffer->source->buffer;
+ buffer = match_buffer->source->source.as_or_throw<tvm::tirx::BufferVar>();
region = ConvertMatchedRegion(match_buffer, std::move(region));
}
TVM_FFI_ICHECK_EQ(buffers->size(), regions->size())
@@ -448,12 +451,12 @@ void BlockReadWriteDetector::Update(std::vector<BufferVar>* buffers,
regions->push_back(std::move(region));
}
-ffi::Array<BufferRegion> BlockReadWriteDetector::CollectRegions(
+ffi::Array<TensorRegion> BlockReadWriteDetector::CollectRegions(
const std::vector<BufferVar>& buffers,
const std::vector<std::vector<tvm::arith::IntSet>>& regions,
const std::unordered_set<const VarNode*>* excluded_buffers) {
TVM_FFI_ICHECK_EQ(buffers.size(), regions.size());
- ffi::Array<BufferRegion> res;
+ ffi::Array<TensorRegion> res;
res.reserve(buffers.size());
for (size_t i = 0; i < regions.size(); ++i) {
if (excluded_buffers != nullptr && excluded_buffers->count(buffers[i].get())) {
@@ -480,7 +483,7 @@ void BlockReadWriteDetector::UpdateOpaque(const Var& buffer_var) {
auto it = buffer_var_map_.find(buffer_var);
if (it != buffer_var_map_.end()) {
const BufferVar& buffer = (*it).second;
- const BufferRegion buffer_region = BufferRegion::FullRegion(buffer);
+ const TensorRegion buffer_region = FullBufferRegion(buffer);
const ffi::Array<Range>& region = buffer_region->region;
std::vector<arith::IntSet> int_set;
int_set.reserve(region.size());
@@ -491,40 +494,40 @@ void BlockReadWriteDetector::UpdateOpaque(const Var& buffer_var) {
}
}
-ffi::Array<ffi::Array<BufferRegion>> GetSBlockAccessRegion(
+ffi::Array<ffi::Array<TensorRegion>> GetSBlockAccessRegion(
const SBlock& block, const ffi::Map<Var, BufferVar>& buffer_var_map) {
auto detector = ffi::make_object<BlockReadWriteDetector>(buffer_var_map);
detector->operator()(block);
- ffi::Array<BufferRegion> writes = detector->CollectWrites();
+ ffi::Array<TensorRegion> writes = detector->CollectWrites();
std::unordered_set<const VarNode*> excluded_buffers;
// exclude write buffers from read regions for reductions if init block is defined.
if (block->init.has_value()) {
- for (const BufferRegion& write_access : writes) {
- excluded_buffers.insert(write_access->buffer.get());
+ for (const TensorRegion& write_access : writes) {
+ excluded_buffers.insert(write_access->source.as_or_throw<tvm::tirx::BufferVar>().get());
}
}
- ffi::Array<BufferRegion> reads = detector->CollectReads(&excluded_buffers);
- ffi::Array<BufferRegion> opaques = detector->CollectOpaques();
+ ffi::Array<TensorRegion> reads = detector->CollectReads(&excluded_buffers);
+ ffi::Array<TensorRegion> opaques = detector->CollectOpaques();
return {reads, writes, opaques};
}
-ffi::Array<ffi::Array<BufferRegion>> GetSBlockReadWriteRegion(
+ffi::Array<ffi::Array<TensorRegion>> GetSBlockReadWriteRegion(
const SBlock& block, const ffi::Map<Var, BufferVar>& buffer_var_map) {
auto detector = ffi::make_object<BlockReadWriteDetector>(buffer_var_map);
detector->operator()(block);
- ffi::Array<BufferRegion> opaques = detector->CollectOpaques();
+ ffi::Array<TensorRegion> opaques = detector->CollectOpaques();
std::unordered_set<const VarNode*> excluded_buffers;
- for (const BufferRegion& opaque_access : opaques) {
- excluded_buffers.insert(opaque_access->buffer.get());
+ for (const TensorRegion& opaque_access : opaques) {
+ excluded_buffers.insert(opaque_access->source.as_or_throw<tvm::tirx::BufferVar>().get());
}
- ffi::Array<BufferRegion> writes = detector->CollectWrites(&excluded_buffers);
+ ffi::Array<TensorRegion> writes = detector->CollectWrites(&excluded_buffers);
if (block->init.has_value()) {
- for (const BufferRegion& write_access : writes) {
- excluded_buffers.insert(write_access->buffer.get());
+ for (const TensorRegion& write_access : writes) {
+ excluded_buffers.insert(write_access->source.as_or_throw<tvm::tirx::BufferVar>().get());
}
}
- ffi::Array<BufferRegion> reads = detector->CollectReads(&excluded_buffers);
- for (const BufferRegion& opaque_access : opaques) {
+ ffi::Array<TensorRegion> reads = detector->CollectReads(&excluded_buffers);
+ for (const TensorRegion& opaque_access : opaques) {
reads.push_back(opaque_access);
writes.push_back(opaque_access);
}
diff --git a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc
index 3e636aa180..568a330dd5 100644
--- a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc
+++ b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc
@@ -135,7 +135,8 @@ class LCADetector : public StmtExprVisitor {
// Update match_buffers
for (const MatchBufferRegion& match_buffer : block->match_buffers) {
- UpdateBufferLCA(match_buffer->source->buffer.get(), ancestor_scopes_.back());
+ UpdateBufferLCA(match_buffer->source->source.as_or_throw<tvm::tirx::BufferVar>().get(),
+ ancestor_scopes_.back());
match_buffers_.insert(match_buffer->buffer.get());
}
@@ -193,7 +194,7 @@ class LCADetector : public StmtExprVisitor {
} else {
opaque_var_scope[iter_var->var.get()] = scope;
for (const auto& write : block->writes) {
- UpdateBufferLCA(write->buffer.get(), scope);
+ UpdateBufferLCA(write->source.as_or_throw<tvm::tirx::BufferVar>().get(), scope);
}
}
}
@@ -202,9 +203,9 @@ class LCADetector : public StmtExprVisitor {
// function to update lca scope of the buffer with loop carried dependent buffer accesses.
// the result scope should be above all loop scopes the accessed opaque block iter vars
// relate to, which is record in `itervar_to_dom_scope`.
- auto do_update = [this, &opaque_var_scope, highest_reduce_scope](const BufferRegion& region,
+ auto do_update = [this, &opaque_var_scope, highest_reduce_scope](const TensorRegion& region,
bool is_reduce_write = false) {
- const BufferVar& buffer = region->buffer;
+ const BufferVar& buffer = region->source.as_or_throw<tvm::tirx::BufferVar>();
const ScopeInfo* scope = ancestor_scopes_.back();
auto handle_itervar = [&opaque_var_scope,
@@ -266,7 +267,8 @@ class LCADetector : public StmtExprVisitor {
}
// Declared regions carry bounds, not opaque runtime accesses.
- ffi::Optional<VisitInterrupt> Visit_(const BufferRegionNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorRegionNode* op) final {
+ if (!op->source.as<BufferVar>()) return StmtExprVisitor::Visit_(op);
for (const Range& range : op->region) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(range->min));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(range->extent));
diff --git a/src/s_tir/meta_schedule/postproc/rewrite_layout.cc b/src/s_tir/meta_schedule/postproc/rewrite_layout.cc
index 84bf296fe5..b47fc6c65c 100644
--- a/src/s_tir/meta_schedule/postproc/rewrite_layout.cc
+++ b/src/s_tir/meta_schedule/postproc/rewrite_layout.cc
@@ -95,7 +95,7 @@ class BufferReadPosCollector : public StmtExprVisitor {
static int GetReadBufferIndex(const SBlock& block, const BufferVar& buffer) {
for (size_t i = 0; i < block->reads.size(); i++) {
- if (block->reads[i]->buffer.same_as(buffer)) {
+ if (block->reads[i]->source.as_or_throw<tvm::tirx::BufferVar>().same_as(buffer)) {
return i;
}
}
@@ -192,9 +192,9 @@ std::vector<std::string> GetCacheReadChain(const BufferVar& buf, const PrimFuncN
ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op) final {
// Check if this block is doing cache_read or a similar operation that consumes cur_buffer_.
if (!op->init && op->reads.size() == 1 && op->writes.size() == 1 &&
- op->reads[0]->buffer.get() == cur_buffer_) {
+ op->reads[0]->source.as_or_throw<tvm::tirx::BufferVar>().get() == cur_buffer_) {
cache_read_chain.push_back(op->name_hint);
- cur_buffer_ = op->writes[0]->buffer.get();
+ cur_buffer_ = op->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>().get();
}
return StmtExprVisitor::Visit_(op);
}
@@ -244,7 +244,8 @@ bool RewriteLayout(const Schedule& sch) {
// in cache_read_chain corresponds to that buffer.
SBlock cache_read_block = sch->Get(sch->GetSBlock(cache_read_chain.back(), func_name));
TVM_FFI_ICHECK_EQ(cache_read_block->writes.size(), 1);
- auto tup_opt = GetSuggestedIndexMap(cache_read_block->writes[0]->buffer, prim_func);
+ auto tup_opt = GetSuggestedIndexMap(
+ cache_read_block->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>(), prim_func);
if (tup_opt == std::nullopt) continue;
auto [anchor_block, buffer_index, index_map] = *tup_opt;
diff --git a/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc b/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc
index 99836eb42b..eeed8dff57 100644
--- a/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc
+++ b/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc
@@ -208,7 +208,7 @@ void AdjustParallelVectorize(const Schedule& sch, const SBlockRV& block_rv,
}
// check the maximal number of axes that are vectorizable (contiguous memory access)
SBlockRealize realize = GetSBlockRealize(sch->state(), block_sref);
- ffi::Array<BufferRegion> buffer_access(realize->block->reads);
+ ffi::Array<TensorRegion> buffer_access(realize->block->reads);
buffer_access.insert(buffer_access.end(), realize->block->writes.begin(),
realize->block->writes.end());
std::unordered_map<const VarNode*, PrimExpr> binding_map;
@@ -224,7 +224,7 @@ void AdjustParallelVectorize(const Schedule& sch, const SBlockRV& block_rv,
int max_fusible = INT32_MAX;
// for each block read/write, get the strides of the loop vars and find the fusible
// (vectorizable) axes
- for (const BufferRegion& access : buffer_access) {
+ for (const TensorRegion& access : buffer_access) {
int fusible = 0;
bool can_analyze_contiguous_access = true;
std::vector<int64_t> strides;
@@ -242,7 +242,8 @@ void AdjustParallelVectorize(const Schedule& sch, const SBlockRV& block_rv,
stride = coef * buffer_stride;
break;
}
- const auto* shape = access->buffer->shape[i].as<IntImmNode>();
+ const auto* shape =
+ access->source.as_or_throw<tvm::tirx::BufferVar>()->shape[i].as<IntImmNode>();
if (shape == nullptr) {
can_analyze_contiguous_access = false;
break;
diff --git a/src/s_tir/meta_schedule/schedule/cuda/winograd.cc b/src/s_tir/meta_schedule/schedule/cuda/winograd.cc
index 250ca202f0..49960b0b57 100644
--- a/src/s_tir/meta_schedule/schedule/cuda/winograd.cc
+++ b/src/s_tir/meta_schedule/schedule/cuda/winograd.cc
@@ -144,8 +144,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
[](Schedule sch, SBlockRV inverse) -> ffi::Array<Schedule> {
GetWinogradProducerAndInlineConst(sch, inverse);
// loops on top of the inverse block: [CO, P, tile_size, tile_size, alpha, alpha]
- int64_t tile_size =
- sch->Get(inverse)->writes[0]->buffer->shape[2].as_or_throw<IntImm>()->value;
+ int64_t tile_size = sch->Get(inverse)
+ ->writes[0]
+ ->source.as_or_throw<tvm::tirx::BufferVar>()
+ ->shape[2]
+ .as_or_throw<IntImm>()
+ ->value;
LoopRV outer{ffi::UnsafeInit()};
{
SBlockRV output = sch->GetConsumers(inverse)[0];
diff --git a/src/s_tir/meta_schedule/schedule_rule/auto_inline.cc b/src/s_tir/meta_schedule/schedule_rule/auto_inline.cc
index 06acc411f4..e53f55b665 100644
--- a/src/s_tir/meta_schedule/schedule_rule/auto_inline.cc
+++ b/src/s_tir/meta_schedule/schedule_rule/auto_inline.cc
@@ -145,8 +145,8 @@ inline InlineType AutoInlineNode::CheckInline(const s_tir::Schedule& sch,
}
// Cond 5. The mapping from read indices to write indices are injective and ordered
if (!is_pure_sptial && (require_injective || require_ordered)) {
- const BufferRegion& write_region = block->writes[0];
- for (const BufferRegion& read_region : block->reads) {
+ const TensorRegion& write_region = block->writes[0];
+ for (const TensorRegion& read_region : block->reads) {
bool injective, ordered;
auto _ = std::ignore;
std::tie(/*exists=*/_, /*surjective=*/_, injective, ordered, /*no_const_read=*/_,
@@ -244,7 +244,7 @@ class InlineConstantScalarsNode : public ScheduleRuleNode {
// }
auto block = sch->Get(block_rv);
if (block->reads.size() == 0 && block->writes.size() == 1 &&
- block->writes[0]->buffer->shape.size() == 0) {
+ block->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>()->shape.size() == 0) {
auto sref = sch->GetSRef(block_rv);
if (!s_tir::IsOutputBlock(sch->state(), sref,
s_tir::GetScopeRoot(sch->state(), sref, true))) {
diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc
index 50cffabb86..0750a6ddad 100644
--- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc
+++ b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc
@@ -35,11 +35,11 @@ using namespace tvm::tirx;
std::vector<int> GetReadBufferNDims(const StmtSRef& block_sref) {
const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref);
- const VarNode* write_buffer = block->writes[0]->buffer.get();
+ const VarNode* write_buffer = block->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>().get();
int n = block->reads.size();
std::vector<int> results(n, -1);
for (int i = 0; i < n; ++i) {
- const VarNode* read_buffer = block->reads[i]->buffer.get();
+ const VarNode* read_buffer = block->reads[i]->source.as_or_throw<tvm::tirx::BufferVar>().get();
if (read_buffer != write_buffer) {
results[i] = GetBufferVar(read_buffer)->shape.size();
}
@@ -369,7 +369,8 @@ void MultiLevelTilingNode::AnnotateCooperativeFetching(Schedule* sch,
// Filter out invalid vector lanes according to the data type.
const tirx::SBlockNode* block_node = (*sch)->GetSRef(block)->StmtAs<tirx::SBlockNode>();
TVM_FFI_ICHECK_EQ(block_node->writes.size(), 1);
- const DLDataType dtype = block_node->writes[0]->buffer->dtype->dtype;
+ const DLDataType dtype =
+ block_node->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>()->dtype->dtype;
std::function<bool(int)> f_filter = nullptr;
if (dtype == DLDataType{kDLFloat, 32, 1}) {
f_filter = [&](int vector_len) { return vector_len <= 4; };
diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
index de20749a15..a6c21c955c 100644
--- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
+++ b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
@@ -489,7 +489,10 @@ std::vector<State> MultiLevelTilingTensorCoreNode::TransformIntermediateOutputLa
auto warp_num_frag_n = f_get_inner_tile_product(-1);
Schedule& sch = state->sch;
- int buffer_ndim = static_cast<int>(sch->Get(state->block_rv)->writes[0]->buffer->shape.size());
+ int buffer_ndim = static_cast<int>(sch->Get(state->block_rv)
+ ->writes[0]
+ ->source.as_or_throw<tvm::tirx::BufferVar>()
+ ->shape.size());
// The dimension of the buffer should be larger or same as that of the tensor intrin.
TVM_FFI_ICHECK_GE(buffer_ndim, 2);
int num_higher_dims = buffer_ndim - 2;
@@ -861,7 +864,7 @@ ffi::Optional<LoopRV> MultiLevelTilingTensorCoreNode::TransformWithTensorIntrin(
visited_buffers.insert(lhs_buffer);
// Refresh block pointer (block sref is not invalidated)
block = TVM_SREF_TO_SBLOCK(block_sref);
- const tirx::BufferRegion& reindexed_buffer_region = s_tir::GetNthAccessBufferRegion(
+ const tvm::TensorRegion& reindexed_buffer_region = s_tir::GetNthAccessBufferRegion(
state->sch->state(), ffi::GetRef<tirx::SBlock>(block), buffer_index, index_type);
auto sub_index_map = f_get_sub_index_map(lhs_buffer, reindexed_buffer_region->region);
buffer_sub_index_map.Set(lhs_buffer, sub_index_map);
diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc
index e11ad4ded1..e863621729 100644
--- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc
+++ b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc
@@ -71,7 +71,7 @@ MultiLevelTilingWideVectorNode::SplitLoop(const Schedule& sch, SBlockRV block_rv
const tirx::SBlockRealize block_realize = s_tir::GetSBlockRealize(sch->state(), block_sref);
TVM_FFI_ICHECK(block_node && block_node->writes.size() == 1);
- const auto out_dtype = block_node->writes[0]->buffer->dtype;
+ const auto out_dtype = block_node->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>()->dtype;
const int vec_len = vector_length_in_bits / out_dtype.bits();
// Determine if this loop is over the innermost axis of the output buffer.
diff --git a/src/s_tir/sblock_scope.cc b/src/s_tir/sblock_scope.cc
index 7076df6b5a..f4f2f0c815 100644
--- a/src/s_tir/sblock_scope.cc
+++ b/src/s_tir/sblock_scope.cc
@@ -87,15 +87,17 @@ SBlockScope::SBlockScope(const ffi::Array<StmtSRef>& child_block_srefs) {
for (const StmtSRef& child_block_sref : child_block_srefs) {
const SBlockNode* child_block = TVM_SREF_TO_SBLOCK(child_block_sref);
// Step 1. Update `buffer_readers` and `buffer_writers` for each buffer
- for (const BufferRegion& region : child_block->reads) {
- buffer_readers[region->buffer].push_back(child_block_sref);
+ for (const TensorRegion& region : child_block->reads) {
+ buffer_readers[region->source.as_or_throw<tvm::tirx::BufferVar>()].push_back(
+ child_block_sref);
}
- for (const BufferRegion& region : child_block->writes) {
- buffer_writers[region->buffer].push_back(child_block_sref);
+ for (const TensorRegion& region : child_block->writes) {
+ buffer_writers[region->source.as_or_throw<tvm::tirx::BufferVar>()].push_back(
+ child_block_sref);
}
// Step 2. Update RAW dependency
- for (const BufferRegion& region : child_block->reads) {
- auto it = buffer_writers.find(region->buffer);
+ for (const TensorRegion& region : child_block->reads) {
+ auto it = buffer_writers.find(region->source.as_or_throw<tvm::tirx::BufferVar>());
if (it != buffer_writers.end()) {
for (const StmtSRef& from : it->second) {
AddDependency(n.get(), from, child_block_sref, DepKind::kRAW);
@@ -103,8 +105,8 @@ SBlockScope::SBlockScope(const ffi::Array<StmtSRef>& child_block_srefs) {
}
}
// Step 3. Update WAW dependency
- for (const BufferRegion& region : child_block->writes) {
- auto it = buffer_writers.find(region->buffer);
+ for (const TensorRegion& region : child_block->writes) {
+ auto it = buffer_writers.find(region->source.as_or_throw<tvm::tirx::BufferVar>());
if (it != buffer_writers.end()) {
for (const StmtSRef& from : it->second) {
AddDependency(n.get(), from, child_block_sref, DepKind::kWAW);
@@ -112,8 +114,8 @@ SBlockScope::SBlockScope(const ffi::Array<StmtSRef>& child_block_srefs) {
}
}
// Step 4. Update WAR dependency
- for (const BufferRegion& region : child_block->writes) {
- auto it = buffer_readers.find(region->buffer);
+ for (const TensorRegion& region : child_block->writes) {
+ auto it = buffer_readers.find(region->source.as_or_throw<tvm::tirx::BufferVar>());
if (it != buffer_readers.end()) {
for (const StmtSRef& from : it->second) {
AddDependency(n.get(), from, child_block_sref, DepKind::kWAR);
diff --git a/src/s_tir/schedule/analysis.h b/src/s_tir/schedule/analysis.h
index 8dcc5b3787..54a77e93d1 100644
--- a/src/s_tir/schedule/analysis.h
+++ b/src/s_tir/schedule/analysis.h
@@ -464,7 +464,7 @@ BufferVar GetNthAccessBuffer(const ScheduleState& self, const SBlock& block, int
* \return The n-th read/write region of the block.
* \throw ScheduleError If the buffer index is out of bound.
*/
-BufferRegion GetNthAccessBufferRegion(const ScheduleState& self, const SBlock& block, int n,
+TensorRegion GetNthAccessBufferRegion(const ScheduleState& self, const SBlock& block, int n,
BufferIndexType index_type);
/*!
@@ -649,7 +649,7 @@ std::tuple</*exists=*/bool,
/*ordered=*/bool,
/*no_const_read=*/bool,
/*no_shift_read=*/bool>
-AnalyzeReadWritePattern(const BufferRegion& read_region, const BufferRegion& write_region);
+AnalyzeReadWritePattern(const TensorRegion& read_region, const TensorRegion& write_region);
/*!
* \brief Check if the block is a data parallel block, i.e. all the block vars are data parallel
@@ -705,7 +705,7 @@ bool NeedsRFactorOrCrossThreadReduction(const s_tir::ScheduleState& self, //
* \param dom_high_exclusive The highest node in the sref tree path
* \return An n-dimensional integer set
*/
-ffi::Array<arith::IntSet> AnalyzeRegionUpperBound(const BufferRegion& region,
+ffi::Array<arith::IntSet> AnalyzeRegionUpperBound(const TensorRegion& region,
const PrimExpr& predicate,
const StmtSRef& dom_low_inclusive,
const StmtSRef& dom_high_exclusive,
@@ -721,7 +721,7 @@ ffi::Array<arith::IntSet> AnalyzeRegionUpperBound(const BufferRegion& region,
* \param analyzer The analyzer
* \return An n-dimensional integer set
*/
-ffi::Array<arith::IntSet> AnalyzeRegionLowerBound(const BufferRegion& region,
+ffi::Array<arith::IntSet> AnalyzeRegionLowerBound(const TensorRegion& region,
const PrimExpr& predicate,
const StmtSRef& dom_low_inclusive,
const StmtSRef& dom_high_exclusive,
diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc
index 409fc2338a..6d695274b5 100644
--- a/src/s_tir/schedule/analysis/analysis.cc
+++ b/src/s_tir/schedule/analysis/analysis.cc
@@ -213,9 +213,9 @@ bool IsDominantBlock(const ScheduleState& self, const StmtSRef& scope_root_sref,
}
// Check whether the input block is the only writer of its outputs
const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref);
- for (const BufferRegion& write_region : block->writes) {
- if (buffer_writers.count(write_region->buffer)) {
- if (buffer_writers.at(write_region->buffer).size() != 1) {
+ for (const TensorRegion& write_region : block->writes) {
+ if (buffer_writers.count(write_region->source.as_or_throw<tvm::tirx::BufferVar>())) {
+ if (buffer_writers.at(write_region->source.as_or_throw<tvm::tirx::BufferVar>()).size() != 1) {
return false;
}
}
@@ -249,11 +249,11 @@ int CheckCompleteBlockErrorCode(const ScheduleState& self, const StmtSRef& block
// Cond 3. No overlap between the buffers the block reads and writes
std::unordered_set<const VarNode*> written_buffers;
written_buffers.reserve(block->writes.size());
- for (const BufferRegion& write : block->writes) {
- written_buffers.insert(write->buffer.get());
+ for (const TensorRegion& write : block->writes) {
+ written_buffers.insert(write->source.as_or_throw<tvm::tirx::BufferVar>().get());
}
- for (const BufferRegion& read : block->reads) {
- if (written_buffers.count(read->buffer.get())) {
+ for (const TensorRegion& read : block->reads) {
+ if (written_buffers.count(read->source.as_or_throw<tvm::tirx::BufferVar>().get())) {
return 3;
}
}
@@ -504,8 +504,8 @@ bool IsOutputBlock(const ScheduleState& self, const StmtSRef& block_sref,
for (const BufferVar& buffer : scope_root->alloc_buffers) {
scope_allocated.insert(buffer.get());
}
- for (const BufferRegion& buffer_region : block->writes) {
- if (!scope_allocated.count(buffer_region->buffer.get())) {
+ for (const TensorRegion& buffer_region : block->writes) {
+ if (!scope_allocated.count(buffer_region->source.as_or_throw<tvm::tirx::BufferVar>().get())) {
return true;
}
}
@@ -552,8 +552,8 @@ bool IsWriteCache(const StmtSRef& block_sref) {
if (block->writes.size() != 1) {
return false;
}
- const BufferRegion& write_region = block->writes[0];
- for (const BufferRegion& read_region : block->reads) {
+ const TensorRegion& write_region = block->writes[0];
+ for (const TensorRegion& read_region : block->reads) {
auto [exists, surjective, injective, ordered, no_const_read, no_shift_read] =
AnalyzeReadWritePattern(read_region, write_region);
// https://gcc.gnu.org/bugzilla/show_bug.cgi?id=81767
@@ -1252,7 +1252,7 @@ ProducerConsumerSplit ProducerConsumerSplit::Find(
/******** Block-buffer relation ********/
-BufferRegion GetNthAccessBufferRegion(const ScheduleState& self, const SBlock& block, int n,
+TensorRegion GetNthAccessBufferRegion(const ScheduleState& self, const SBlock& block, int n,
BufferIndexType index_type) {
class BufferIndexOutOfRangeError : public ScheduleErrorContextObj {
public:
@@ -1298,7 +1298,7 @@ BufferRegion GetNthAccessBufferRegion(const ScheduleState& self, const SBlock& b
BufferIndexType index_type_;
};
- const ffi::Array<BufferRegion>& access_region =
+ const ffi::Array<TensorRegion>& access_region =
index_type == BufferIndexType::kWrite ? block->writes : block->reads;
if (n < 0 || static_cast<int>(access_region.size()) <= n) {
@@ -1309,7 +1309,8 @@ BufferRegion GetNthAccessBufferRegion(const ScheduleState& self, const SBlock& b
BufferVar GetNthAccessBuffer(const ScheduleState& self, const SBlock& block, int n,
BufferIndexType index_type) {
- return GetNthAccessBufferRegion(self, block, n, index_type)->buffer;
+ return GetNthAccessBufferRegion(self, block, n, index_type)
+ ->source.as_or_throw<tvm::tirx::BufferVar>();
}
std::pair<ffi::Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref,
@@ -1416,11 +1417,11 @@ std::tuple</*exists=*/bool,
/*ordered=*/bool,
/*no_const_read=*/bool,
/*no_shift_read=*/bool>
-AnalyzeReadWritePattern(const BufferRegion& read_region, const BufferRegion& write_region) {
+AnalyzeReadWritePattern(const TensorRegion& read_region, const TensorRegion& write_region) {
static constexpr const std::tuple<bool, bool, bool, bool, bool, bool> kNotExist =
std::make_tuple(false, false, false, false, false, false);
// Step 1. Extract the write indices
- int w_dim = write_region->buffer->shape.size();
+ int w_dim = write_region->source.as_or_throw<tvm::tirx::BufferVar>()->shape.size();
std::unordered_map<const VarNode*, int> var2idx;
var2idx.reserve(w_dim);
for (int i = 0; i < w_dim; ++i) {
@@ -1437,7 +1438,7 @@ AnalyzeReadWritePattern(const BufferRegion& read_region, const BufferRegion& wri
// Step 2. Map each read index to a write index
bool no_const_read = true;
bool no_shift_read = true;
- int r_dim = read_region->buffer->shape.size();
+ int r_dim = read_region->source.as_or_throw<tvm::tirx::BufferVar>()->shape.size();
std::vector<int> mapped(r_dim, -1);
for (int i = 0; i < r_dim; ++i) {
const Range& dom = read_region->region[i];
@@ -1565,7 +1566,7 @@ bool NeedsMultiLevelTiling(const ScheduleState& self, const StmtSRef& block_sref
!IsTrivialBinding(self, block_sref)) {
return false;
}
- const VarNode* write_buffer = block->writes[0]->buffer.get();
+ const VarNode* write_buffer = block->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>().get();
// Step 1. Sort out spatial block variables. Skip the block iters of domain [0, 1), since such
// block iters distracts the following check of the unused block iters.
std::vector<const VarNode*> spatial_block_vars;
@@ -1584,8 +1585,8 @@ bool NeedsMultiLevelTiling(const ScheduleState& self, const StmtSRef& block_sref
int total_unused_block_vars = 0;
std::unordered_set<const VarNode*> read_buffers;
read_buffers.reserve(block->reads.size());
- for (const BufferRegion& buffer_region : block->reads) {
- const VarNode* buffer = buffer_region->buffer.get();
+ for (const TensorRegion& buffer_region : block->reads) {
+ const VarNode* buffer = buffer_region->source.as_or_throw<tvm::tirx::BufferVar>().get();
const ffi::Array<Range>& regions = buffer_region->region;
// Step 2.1. Duplication of read buffers are not allowed
if (read_buffers.insert(buffer).second == false) {
diff --git a/src/s_tir/schedule/analysis/reducer.cc b/src/s_tir/schedule/analysis/reducer.cc
index 07e9b5cea3..9de832d805 100644
--- a/src/s_tir/schedule/analysis/reducer.cc
+++ b/src/s_tir/schedule/analysis/reducer.cc
@@ -561,8 +561,8 @@ bool ReductionIterNotIndexOutputBuffer(const SBlock& block) {
// Step 2. Check if the reduction block iters are used to index the output buffer.
std::unordered_set<const VarNode*> buffer_written;
buffer_written.reserve(block->writes.size());
- for (const BufferRegion& write_region : block->writes) {
- buffer_written.insert(write_region->buffer.get());
+ for (const TensorRegion& write_region : block->writes) {
+ buffer_written.insert(write_region->source.as_or_throw<tvm::tirx::BufferVar>().get());
}
std::unordered_set<const VarNode*> buffer_allocated;
@@ -582,11 +582,13 @@ bool ReductionIterNotIndexOutputBuffer(const SBlock& block) {
std::unordered_map<const VarNode*, const VarNode*> match_buffer_sources;
for (const MatchBufferRegion& region : block->match_buffers) {
- match_buffer_sources[region->buffer.get()] = region->source->buffer.get();
+ match_buffer_sources[region->buffer.get()] =
+ region->source->source.as_or_throw<tvm::tirx::BufferVar>().get();
}
auto visit_block = [&](const SBlock& nested_block) -> ffi::Expected<ffi::WalkResult> {
for (const MatchBufferRegion& region : nested_block->match_buffers) {
- match_buffer_sources[region->buffer.get()] = region->source->buffer.get();
+ match_buffer_sources[region->buffer.get()] =
+ region->source->source.as_or_throw<tvm::tirx::BufferVar>().get();
}
return ffi::WalkResult::Advance();
};
diff --git a/src/s_tir/schedule/ir_comparator.cc b/src/s_tir/schedule/ir_comparator.cc
index 447bcd5053..81da621e37 100644
--- a/src/s_tir/schedule/ir_comparator.cc
+++ b/src/s_tir/schedule/ir_comparator.cc
@@ -517,12 +517,14 @@ bool TensorizeComparator::CompareBuffer(const BufferVar& lhs, const BufferVar& r
return equal;
}
-bool TensorizeComparator::CompareBufferRegion(const BufferRegion& lhs, const BufferRegion& rhs) {
- if (!CompareBuffer(lhs->buffer, rhs->buffer)) {
+bool TensorizeComparator::CompareBufferRegion(const TensorRegion& lhs, const TensorRegion& rhs) {
+ if (!CompareBuffer(lhs->source.as_or_throw<tvm::tirx::BufferVar>(),
+ rhs->source.as_or_throw<tvm::tirx::BufferVar>())) {
if (assert_mode_) {
std::ostringstream os;
- os << "CompareBufferRegion returning false due to buffer mismatch: lhs->buffer="
- << lhs->buffer << " vs rhs->buffer=" << rhs->buffer;
+ os << "CompareBufferRegion returning false due to buffer mismatch: lhs->source="
+ << lhs->source.as_or_throw<tvm::tirx::BufferVar>()
+ << " vs rhs->source=" << rhs->source.as_or_throw<tvm::tirx::BufferVar>();
EmitError(os.str());
}
return false;
@@ -540,7 +542,7 @@ bool TensorizeComparator::CompareBufferRegion(const BufferRegion& lhs, const Buf
return false;
}
- auto it = buffer_indices_.find(lhs->buffer);
+ auto it = buffer_indices_.find(lhs->source.as_or_throw<tvm::tirx::BufferVar>());
if (it == buffer_indices_.end()) {
// Update base indices for the buffer, this can only happen if it is visiting the scope block.
TVM_FFI_ICHECK(is_scope_block);
@@ -574,7 +576,8 @@ bool TensorizeComparator::CompareBufferRegion(const BufferRegion& lhs, const Buf
return false;
}
}
- buffer_indices_.emplace(lhs->buffer, std::move(indices_base));
+ buffer_indices_.emplace(lhs->source.as_or_throw<tvm::tirx::BufferVar>(),
+ std::move(indices_base));
} else {
// Check the base indices are consistent.
const std::vector<PrimExpr>& indices_base = it->second;
diff --git a/src/s_tir/schedule/ir_comparator.h b/src/s_tir/schedule/ir_comparator.h
index 3d84e01b94..28c8332341 100644
--- a/src/s_tir/schedule/ir_comparator.h
+++ b/src/s_tir/schedule/ir_comparator.h
@@ -90,7 +90,7 @@ class TensorizeComparator : public ExprComparator, public StmtComparator {
bool DefEqual(const Var& lhs, const Var& rhs);
bool CompareExpr(const Expr& lhs, const Expr& rhs);
virtual bool CompareBuffer(const BufferVar& lhs, const BufferVar& rhs);
- bool CompareBufferRegion(const BufferRegion& lhs, const BufferRegion& rhs);
+ bool CompareBufferRegion(const TensorRegion& lhs, const TensorRegion& rhs);
bool CompareAnnotation(const std::pair<ffi::String, ffi::Any>& lhs,
const std::pair<ffi::String, ffi::Any>& rhs);
bool CompareAnnotationMap(const ffi::Map<ffi::String, ffi::Any>& lhs,
diff --git a/src/s_tir/schedule/primitive/annotate_buffer_access.cc b/src/s_tir/schedule/primitive/annotate_buffer_access.cc
index 9b785c90fa..72ca2d407b 100644
--- a/src/s_tir/schedule/primitive/annotate_buffer_access.cc
+++ b/src/s_tir/schedule/primitive/annotate_buffer_access.cc
@@ -30,7 +30,7 @@ class AnnotateRegionRewriter : public StmtExprMutator {
using StmtExprMutator::Mutate;
using StmtExprMutator::Mutate_;
- AnnotateRegionRewriter(BufferVar buffer, int buffer_index, BufferRegion new_region,
+ AnnotateRegionRewriter(BufferVar buffer, int buffer_index, TensorRegion new_region,
BufferIndexType buffer_index_type)
: buffer_(buffer),
buffer_index_(buffer_index),
@@ -42,7 +42,7 @@ class AnnotateRegionRewriter : public StmtExprMutator {
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
.as_or_throw<SBlock>();
- ffi::Array<BufferRegion> regions =
+ ffi::Array<TensorRegion> regions =
buffer_index_type_ == BufferIndexType::kWrite ? block->writes : block->reads;
TVM_FFI_ICHECK_GE(buffer_index_, 0) << "Buffer index must be non-negative";
TVM_FFI_ICHECK_LT(buffer_index_, static_cast<int>(regions.size()))
@@ -86,7 +86,7 @@ class AnnotateRegionRewriter : public StmtExprMutator {
private:
BufferVar buffer_;
int buffer_index_;
- BufferRegion new_region_;
+ TensorRegion new_region_;
BufferIndexType buffer_index_type_;
};
@@ -110,7 +110,7 @@ void AnnotateBufferAccess(ScheduleState self, const StmtSRef& block_sref, int bu
new_indices[i], analyzer->Simplify(new_indices[i + 1] - new_indices[i])));
}
- BufferRegion new_region(buffer, new_ranges);
+ TensorRegion new_region = BufferRegion(buffer, new_ranges);
auto mutator =
ffi::make_object<AnnotateRegionRewriter>(buffer, buffer_index, new_region, buffer_index_type);
diff --git a/src/s_tir/schedule/primitive/block_annotate.cc b/src/s_tir/schedule/primitive/block_annotate.cc
index c50b28b8d7..396051b25c 100644
--- a/src/s_tir/schedule/primitive/block_annotate.cc
+++ b/src/s_tir/schedule/primitive/block_annotate.cc
@@ -215,7 +215,9 @@ class StorageScopeMutator : public ReplaceBufferMutator {
private:
MatchBufferRegion VisitMatchBufferRegion(const MatchBufferRegion& match_buffer) final {
- if (auto replacement = VarRemapGet(match_buffer->source->buffer).as<BufferVar>()) {
+ if (auto replacement =
+ VarRemapGet(match_buffer->source->source.as_or_throw<tvm::tirx::BufferVar>())
+ .as<BufferVar>()) {
BufferVar new_target_buffer = WithScope(match_buffer->buffer, replacement.value().scope());
VarRemapSet(match_buffer->buffer, new_target_buffer);
return MatchBufferRegion(new_target_buffer,
@@ -320,7 +322,9 @@ class DTypeMutator : public ReplaceBufferMutator {
private:
MatchBufferRegion VisitMatchBufferRegion(const MatchBufferRegion& match_buffer) final {
- if (auto replacement = VarRemapGet(match_buffer->source->buffer).as<BufferVar>()) {
+ if (auto replacement =
+ VarRemapGet(match_buffer->source->source.as_or_throw<tvm::tirx::BufferVar>())
+ .as<BufferVar>()) {
BufferVar new_target_buffer = WithDType(match_buffer->buffer, replacement.value()->dtype);
VarRemapSet(match_buffer->buffer, new_target_buffer);
return MatchBufferRegion(new_target_buffer,
diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc b/src/s_tir/schedule/primitive/blockize_tensorize.cc
index acf1cb6456..830909374d 100644
--- a/src/s_tir/schedule/primitive/blockize_tensorize.cc
+++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc
@@ -300,7 +300,7 @@ SBlockRealize GenerateInner(bool is_write_reduction,
n->iter_vars = iter_vars;
n->init = std::nullopt;
if (is_write_reduction) {
- ffi::Array<BufferRegion> reads;
+ ffi::Array<TensorRegion> reads;
reads.reserve(block->writes.size() + block->reads.size());
reads.insert(reads.end(), block->writes.begin(), block->writes.end());
reads.insert(reads.end(), block->reads.begin(), block->reads.end());
@@ -448,12 +448,12 @@ Stmt ReplaceAndSimplify(const Stmt& stmt, const ffi::Map<Var, PrimExpr>& sub,
* \param dom_map The variables to be relaxed
* \return The relaxed regions
*/
-ffi::Array<BufferRegion> EvalSetRegions(const ffi::Array<BufferRegion>& regions,
+ffi::Array<TensorRegion> EvalSetRegions(const ffi::Array<TensorRegion>& regions,
const ffi::Map<Var, arith::IntSet>& dom_map) {
- ffi::Array<BufferRegion> results;
+ ffi::Array<TensorRegion> results;
results.reserve(regions.size());
- for (const BufferRegion& buffer_region : regions) {
- const BufferVar& buffer = buffer_region->buffer;
+ for (const TensorRegion& buffer_region : regions) {
+ const BufferVar& buffer = buffer_region->source.as_or_throw<tvm::tirx::BufferVar>();
ffi::Array<arith::IntSet> relaxed = arith::EvalSet(buffer_region->region, dom_map);
TVM_FFI_ICHECK_EQ(relaxed.size(), buffer->shape.size());
int ndim = buffer->shape.size();
@@ -472,11 +472,11 @@ ffi::Array<BufferRegion> EvalSetRegions(const ffi::Array<BufferRegion>& regions,
* \param regions The input regions for the union.
* \return The union regions
*/
-ffi::Array<BufferRegion> UnionRegions(const ffi::Array<BufferRegion>& regions) {
+ffi::Array<TensorRegion> UnionRegions(const ffi::Array<TensorRegion>& regions) {
typedef std::vector<ffi::Array<arith::IntSet>> ranges_t;
std::unordered_map<BufferVar, ranges_t, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> intset_map;
- for (const BufferRegion& buffer_region : regions) {
- const BufferVar& buffer = buffer_region->buffer;
+ for (const TensorRegion& buffer_region : regions) {
+ const BufferVar& buffer = buffer_region->source.as_or_throw<tvm::tirx::BufferVar>();
if (intset_map.find(buffer) == intset_map.end()) {
intset_map[buffer] = {buffer->shape.size(), ffi::Array<arith::IntSet>()};
}
@@ -486,7 +486,7 @@ ffi::Array<BufferRegion> UnionRegions(const ffi::Array<BufferRegion>& regions) {
intset_map[buffer][dim].push_back(arith::IntSet::FromRange(buffer_region->region[dim]));
}
}
- ffi::Array<BufferRegion> results;
+ ffi::Array<TensorRegion> results;
for (const auto& it : intset_map) {
const BufferVar& buffer = it.first;
ffi::Array<Range> regions;
@@ -606,8 +606,8 @@ SBlockRealize BlockizeBlocks(const ScheduleState& self, const ffi::Array<StmtSRe
PrimExpr outer_predicate{nullptr};
ffi::Array<IterVar> outer_iter_vars{nullptr};
ffi::Array<PrimExpr> outer_bindings{nullptr};
- ffi::Array<BufferRegion> read_regions;
- ffi::Array<BufferRegion> write_regions;
+ ffi::Array<TensorRegion> read_regions;
+ ffi::Array<TensorRegion> write_regions;
std::string outer_block_name = "outer_";
ffi::Map<Var, Var> loop_var_subst;
arith::Analyzer analyzer;
@@ -821,8 +821,8 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int
PrimFunc intrin_impl = DeepCopy(intrin->impl);
int index_dtype_bits = -1;
- auto f_update_max_dtype_bits_from_region = [&](const ffi::Array<BufferRegion>& buffer_regions) {
- for (const BufferRegion& buffer_region : buffer_regions) {
+ auto f_update_max_dtype_bits_from_region = [&](const ffi::Array<TensorRegion>& buffer_regions) {
+ for (const TensorRegion& buffer_region : buffer_regions) {
for (const auto& range : buffer_region->region) {
index_dtype_bits = std::max(index_dtype_bits, range->min.ty().bits());
}
@@ -857,11 +857,11 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int
std::unordered_map<BufferVar, ffi::Array<Range>, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
impl2region;
SBlock impl_block = intrin_impl->body.as_or_throw<SBlockRealize>()->block;
- for (const BufferRegion& read : impl_block->reads) {
- impl2region.emplace(read->buffer, read->region);
+ for (const TensorRegion& read : impl_block->reads) {
+ impl2region.emplace(read->source.as_or_throw<tvm::tirx::BufferVar>(), read->region);
}
- for (const BufferRegion& write : impl_block->writes) {
- impl2region.emplace(write->buffer, write->region);
+ for (const TensorRegion& write : impl_block->writes) {
+ impl2region.emplace(write->source.as_or_throw<tvm::tirx::BufferVar>(), write->region);
}
// Step 4: Create MatchBufferRegion for the params of the impl function of the tensor
// intrin to make them subregions of the buffer in the original IR.
diff --git a/src/s_tir/schedule/primitive/cache_read_write.cc b/src/s_tir/schedule/primitive/cache_read_write.cc
index 555d6f5eca..417b7c6380 100644
--- a/src/s_tir/schedule/primitive/cache_read_write.cc
+++ b/src/s_tir/schedule/primitive/cache_read_write.cc
@@ -91,15 +91,15 @@ struct CacheStageInfo {
/*! \brief A set of blocks that will consume the new cache. */
std::unordered_set<StmtSRef, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> consumer_blocks;
/*! \brief cache region for the buffer to be cached */
- BufferRegion cache_region;
+ TensorRegion cache_region;
};
/*! \brief Return the buffer region related with the buffer */
-ffi::Optional<BufferRegion> GetBufferRegionFromBuffer(
- const ffi::Array<BufferRegion>& buffer_regions, const BufferVar& buffer) {
- ffi::Optional<BufferRegion> res = std::nullopt;
+ffi::Optional<TensorRegion> GetBufferRegionFromBuffer(
+ const ffi::Array<TensorRegion>& buffer_regions, const BufferVar& buffer) {
+ ffi::Optional<TensorRegion> res = std::nullopt;
for (const auto& region : buffer_regions) {
- if (region->buffer.same_as(buffer)) {
+ if (region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(buffer)) {
TVM_FFI_ICHECK(!res.has_value());
res = region;
}
@@ -122,7 +122,7 @@ struct ReindexCacheStageInfo : CacheStageInfo {
* reindex_cache_read/write. */
class NotSinglePointAccess : public ScheduleErrorContextObj {
public:
- explicit NotSinglePointAccess(IRModule mod, SBlock block, BufferRegion cache_region,
+ explicit NotSinglePointAccess(IRModule mod, SBlock block, TensorRegion cache_region,
bool is_cache_read)
: mod_(std::move(mod)), block_(std::move(block)), cache_region_(cache_region) {
primitive_name_ = is_cache_read ? "reindex_cache_read" : "reindex_cache_write";
@@ -146,7 +146,7 @@ class NotSinglePointAccess : public ScheduleErrorContextObj {
private:
IRModule mod_;
SBlock block_;
- BufferRegion cache_region_;
+ TensorRegion cache_region_;
ffi::String primitive_name_;
};
@@ -159,7 +159,7 @@ class NotSinglePointAccess : public ScheduleErrorContextObj {
* \returns A block indicating the body of the loop nesting.
*/
template <bool is_cache_read>
-SBlock MakeReindexCacheStage(const BufferRegion& cache_region, ReindexCacheStageInfo* info,
+SBlock MakeReindexCacheStage(const TensorRegion& cache_region, ReindexCacheStageInfo* info,
const ffi::String& storage_scope) {
// loop variables
std::vector<PrimVar> loop_vars;
@@ -217,7 +217,8 @@ SBlock MakeReindexCacheStage(const BufferRegion& cache_region, ReindexCacheStage
/*iter_vars*/ std::move(block_vars),
/*reads=*/{BufferRegion(info->read_buffer, read_access_region)},
/*writes=*/{BufferRegion(info->write_buffer, write_access_region)},
- /*name_hint*/ cache_region->buffer.name() + "_" + storage_scope,
+ /*name_hint*/ cache_region->source.as_or_throw<tvm::tirx::BufferVar>().name() + "_" +
+ storage_scope,
/*body=*/
BufferStore(info->write_buffer, BufferLoad(info->read_buffer, read_access_indices),
write_access_indices),
@@ -253,7 +254,7 @@ SBlock MakeReindexCacheStage(const BufferRegion& cache_region, ReindexCacheStage
* full region or compact region.
* \returns A block indicating the body of the loop nesting.
*/
-SBlock MakeCacheStage(const BufferRegion& cache_region, CacheStageInfo* info,
+SBlock MakeCacheStage(const TensorRegion& cache_region, CacheStageInfo* info,
const ffi::String& storage_scope, bool cache_full_region = true) {
// loop variables
std::vector<PrimVar> loop_vars;
@@ -274,11 +275,13 @@ SBlock MakeCacheStage(const BufferRegion& cache_region, CacheStageInfo* info,
ffi::Array<PrimExpr> read_access_indices;
ffi::Array<PrimExpr> write_access_indices;
// Create block vars, block's accessed region and accessing indices
- for (int i = 0; i < static_cast<int>(cache_region->buffer->shape.size()); ++i) {
+ for (int i = 0;
+ i < static_cast<int>(cache_region->source.as_or_throw<tvm::tirx::BufferVar>()->shape.size());
+ ++i) {
Range axis_range = cache_region->region[i];
PrimVar var("v" + std::to_string(read_access_indices.size()), axis_range->extent.ty());
if (cache_full_region) {
- PrimExpr dim = cache_region->buffer->shape[i];
+ PrimExpr dim = cache_region->source.as_or_throw<tvm::tirx::BufferVar>()->shape[i];
block_vars.push_back(IterVar(/*dom=*/Range::FromMinExtent(IntImm(dim.ty(), 0), dim),
/*var=*/var,
/*IterVarType=*/kDataPar));
@@ -291,7 +294,7 @@ SBlock MakeCacheStage(const BufferRegion& cache_region, CacheStageInfo* info,
/*dom=*/Range::FromMinExtent(IntImm(axis_range->extent.ty(), 0), axis_range->extent),
/*var=*/var,
/*IterVarType=*/kDataPar));
- if (cache_region->buffer.same_as(info->read_buffer)) {
+ if (cache_region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(info->read_buffer)) {
// cache_read
read_access_indices.push_back(axis_range->min + var);
read_access_region.push_back(
@@ -317,7 +320,8 @@ SBlock MakeCacheStage(const BufferRegion& cache_region, CacheStageInfo* info,
/*iter_vars=*/std::move(block_vars),
/*reads=*/{BufferRegion(info->read_buffer, read_access_region)},
/*writes=*/{BufferRegion(info->write_buffer, write_access_region)},
- /*name_hint=*/cache_region->buffer.name() + "_" + storage_scope,
+ /*name_hint=*/cache_region->source.as_or_throw<tvm::tirx::BufferVar>().name() + "_" +
+ storage_scope,
/*body=*/
BufferStore(info->write_buffer, BufferLoad(info->read_buffer, read_access_indices),
write_access_indices),
@@ -422,8 +426,8 @@ SBlock MakeReIndexStage(const SBlock& block, CacheStageInfo* info,
// Create the body block
SBlock new_block(
/*iter_vars=*/new_block_iters,
- /*reads=*/{BufferRegion::FromPoint(info->read_buffer, src_indices)},
- /*writes=*/{BufferRegion::FromPoint(info->write_buffer, dst_indices)},
+ /*reads=*/{BufferRegionFromPoint(info->read_buffer, src_indices)},
+ /*writes=*/{BufferRegionFromPoint(info->write_buffer, dst_indices)},
/*name_hint=*/info->write_buffer.name() + "_reindex",
/*body=*/
BufferStore(info->write_buffer, BufferLoad(info->read_buffer, src_indices), dst_indices));
@@ -597,8 +601,8 @@ static PrimExpr CollectNestedBlockPredicates(const Stmt& body, const BufferVar&
const SBlockNode* block = realize->block.get();
const auto& regions = (index_type_ == BufferIndexType::kRead) ? block->reads : block->writes;
bool accesses_buffer = false;
- for (const BufferRegion& region : regions) {
- if (region->buffer.same_as(buffer_)) {
+ for (const TensorRegion& region : regions) {
+ if (region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(buffer_)) {
accesses_buffer = true;
break;
}
@@ -654,13 +658,13 @@ static PrimExpr CollectNestedBlockPredicates(const Stmt& body, const BufferVar&
* with the block's own predicate before relaxation. Defaults to true (no effect).
* \return The relaxed buffer region.
*/
-BufferRegion RelaxBufferRegion(ScheduleState self, const BufferRegion& buffer_region,
+TensorRegion RelaxBufferRegion(ScheduleState self, const TensorRegion& buffer_region,
const StmtSRef& block_sref, const StmtSRef& dom_low_inclusive,
const StmtSRef& dom_high_exclusive,
PrimExpr extra_predicate = IntImm::Bool(true)) {
SBlockRealize realize = GetSBlockRealize(self, block_sref);
ffi::Map<Var, PrimExpr> binding = GetBindings(realize);
- const BufferVar& buffer = buffer_region->buffer;
+ const BufferVar& buffer = buffer_region->source.as_or_throw<tvm::tirx::BufferVar>();
arith::Analyzer analyzer;
auto f_substitute = [&binding](const Var& var) -> ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
if (auto repl = binding.Get(var)) return ffi::Any(*std::move(repl));
@@ -673,7 +677,7 @@ BufferRegion RelaxBufferRegion(ScheduleState self, const BufferRegion& buffer_re
.as_or_throw<PrimExpr>();
return Range::FromMinExtent(min, extent);
});
- BufferRegion subst_region = BufferRegion(buffer, mapped_region);
+ TensorRegion subst_region = BufferRegion(buffer, mapped_region);
ffi::Array<arith::IntSet> int_sets = AnalyzeRegionUpperBound(
/*region=*/subst_region,
/*predicate=*/
@@ -984,14 +988,14 @@ class CacheReadRewriter : public StmtExprMutator {
return ret;
};
- update_access_regions = [this, update_region](ffi::Array<BufferRegion> regions) {
+ update_access_regions = [this, update_region](ffi::Array<TensorRegion> regions) {
if (cache_full_region_) {
return ReplaceBuffer(std::move(regions), info_->read_buffer, info_->write_buffer);
}
- ffi::Array<BufferRegion> ret;
- for (const BufferRegion& region : regions) {
- if (region->buffer.same_as(info_->read_buffer)) {
+ ffi::Array<TensorRegion> ret;
+ for (const TensorRegion& region : regions) {
+ if (region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(info_->read_buffer)) {
ret.push_back(BufferRegion(info_->write_buffer,
update_region(region->region, info_->cache_region->region)));
} else {
@@ -1007,7 +1011,8 @@ class CacheReadRewriter : public StmtExprMutator {
ffi::Array<MatchBufferRegion> ret;
for (const MatchBufferRegion& match_buffer : match_buffers) {
- if (match_buffer->source->buffer.same_as(info_->read_buffer)) {
+ if (match_buffer->source->source.as_or_throw<tvm::tirx::BufferVar>().same_as(
+ info_->read_buffer)) {
ret.push_back(MatchBufferRegion(
match_buffer->buffer,
BufferRegion(info_->write_buffer, update_region(match_buffer->source->region,
@@ -1084,7 +1089,7 @@ class CacheReadRewriter : public StmtExprMutator {
// Only make this change if the block is one of the specified consumers.
if (is_consumer) {
// Use the updated block stmt
- ffi::Array<BufferRegion> reads = update_access_regions(stmt->reads);
+ ffi::Array<TensorRegion> reads = update_access_regions(stmt->reads);
ffi::Array<MatchBufferRegion> match_buffers = update_match_buffers(stmt->match_buffers);
if (!reads.same_as(stmt->reads) || !match_buffers.same_as(stmt->match_buffers)) {
ffi::ObjectPtr<SBlockNode> n = ffi::make_object<SBlockNode>(*stmt.as<SBlockNode>());
@@ -1146,10 +1151,13 @@ class CacheReadRewriter : public StmtExprMutator {
return node;
}
- UnchangedOr<Expr> Mutate_(const BufferRegionNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Expr> Mutate_(const TensorRegionNode* op, InplaceMode inplace_mode) final {
+ if (!op->source.as<BufferVar>()) {
+ return StmtExprMutator::Mutate_(op, inplace_mode);
+ }
auto region = Mutate(op->region).as_or_throw<UnchangedOr<ffi::Array<Range>>>();
if (region.UnchangedOrSameAs(op->region)) return ffi::Unchanged();
- BufferRegion node = ffi::GetRef<BufferRegion>(op);
+ TensorRegion node = ffi::GetRef<TensorRegion>(op);
node.CopyOnWrite()->region = std::move(region).ValueUnchecked();
return node;
}
@@ -1161,7 +1169,7 @@ class CacheReadRewriter : public StmtExprMutator {
/*! \brief Whether the most recently visited block is a specified consumer. */
bool current_block_consumes;
/*! \brief function to update read/write region of block being cache read.*/
- std::function<ffi::Array<BufferRegion>(ffi::Array<BufferRegion>)> update_access_regions;
+ std::function<ffi::Array<TensorRegion>(ffi::Array<TensorRegion>)> update_access_regions;
/*! \brief function to update match buffers of block being cache read.*/
std::function<ffi::Array<MatchBufferRegion>(ffi::Array<MatchBufferRegion>)> update_match_buffers;
/*!
@@ -1196,10 +1204,10 @@ class ReindexCacheReadRewriter : public CacheReadRewriter {
explicit ReindexCacheReadRewriter(const StmtSRef& scope_sref, ReindexCacheStageInfo* info)
: CacheReadRewriter(scope_sref, info) {
new_indices_ = info->indices;
- update_access_regions = [&](ffi::Array<BufferRegion> reads) {
- ffi::Array<BufferRegion> new_reads;
- for (const BufferRegion& buf_region : reads) {
- if (buf_region->buffer.same_as(info_->read_buffer)) {
+ update_access_regions = [&](ffi::Array<TensorRegion> reads) {
+ ffi::Array<TensorRegion> new_reads;
+ for (const TensorRegion& buf_region : reads) {
+ if (buf_region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(info_->read_buffer)) {
Region region;
for (const PrimExpr index : new_indices_) {
region.push_back(Range::FromMinExtent(index, IntImm::Int32(1)));
@@ -1214,8 +1222,8 @@ class ReindexCacheReadRewriter : public CacheReadRewriter {
update_match_buffers = [&](const ffi::Array<MatchBufferRegion> match_buffers) {
ffi::Array<MatchBufferRegion> new_match_buffers;
for (const MatchBufferRegion& match_buffer_region : match_buffers) {
- BufferRegion source = match_buffer_region->source;
- if (source->buffer.same_as(info_->read_buffer)) {
+ TensorRegion source = match_buffer_region->source;
+ if (source->source.as_or_throw<tvm::tirx::BufferVar>().same_as(info_->read_buffer)) {
Region region;
for (const PrimExpr index : new_indices_) {
region.push_back(Range::FromMinExtent(index, IntImm::Int32(1)));
@@ -1289,14 +1297,14 @@ class CacheWriteRewriter : public StmtExprMutator {
return ret;
};
- update_access_regions = [this, update_region](ffi::Array<BufferRegion> regions) {
+ update_access_regions = [this, update_region](ffi::Array<TensorRegion> regions) {
if (cache_full_region_) {
return ReplaceBuffer(regions, info_->write_buffer, info_->read_buffer);
}
- ffi::Array<BufferRegion> ret;
- for (const BufferRegion& region : regions) {
- if (region->buffer.same_as(info_->write_buffer)) {
+ ffi::Array<TensorRegion> ret;
+ for (const TensorRegion& region : regions) {
+ if (region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(info_->write_buffer)) {
ret.push_back(BufferRegion(info_->read_buffer,
update_region(region->region, info_->cache_region->region)));
} else {
@@ -1312,7 +1320,8 @@ class CacheWriteRewriter : public StmtExprMutator {
ffi::Array<MatchBufferRegion> ret;
for (const MatchBufferRegion& match_buffer : match_buffers) {
- if (match_buffer->source->buffer.same_as(info_->write_buffer)) {
+ if (match_buffer->source->source.as_or_throw<tvm::tirx::BufferVar>().same_as(
+ info_->write_buffer)) {
ret.push_back(MatchBufferRegion(
match_buffer->buffer,
BufferRegion(info_->read_buffer, update_region(match_buffer->source->region,
@@ -1348,8 +1357,8 @@ class CacheWriteRewriter : public StmtExprMutator {
const SBlockNode* consumer_node = TVM_SREF_TO_SBLOCK(consumer_sref);
SBlock consumer_block = ffi::GetRef<SBlock>(consumer_node);
if (old_stmt.same_as(consumer_block)) {
- ffi::Array<BufferRegion> writes = update_access_regions(block->writes);
- ffi::Array<BufferRegion> reads = update_access_regions(block->reads);
+ ffi::Array<TensorRegion> writes = update_access_regions(block->writes);
+ ffi::Array<TensorRegion> reads = update_access_regions(block->reads);
ffi::Array<MatchBufferRegion> match_buffers = update_match_buffers(block->match_buffers);
if (!writes.same_as(block->writes) || !reads.same_as(block->reads) ||
!match_buffers.same_as(block->match_buffers)) {
@@ -1485,10 +1494,13 @@ class CacheWriteRewriter : public StmtExprMutator {
return node;
}
- UnchangedOr<Expr> Mutate_(const BufferRegionNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Expr> Mutate_(const TensorRegionNode* op, InplaceMode inplace_mode) final {
+ if (!op->source.as<BufferVar>()) {
+ return StmtExprMutator::Mutate_(op, inplace_mode);
+ }
auto region = Mutate(op->region).as_or_throw<UnchangedOr<ffi::Array<Range>>>();
if (region.UnchangedOrSameAs(op->region)) return ffi::Unchanged();
- BufferRegion node = ffi::GetRef<BufferRegion>(op);
+ TensorRegion node = ffi::GetRef<TensorRegion>(op);
node.CopyOnWrite()->region = std::move(region).ValueUnchecked();
return node;
}
@@ -1502,7 +1514,7 @@ class CacheWriteRewriter : public StmtExprMutator {
/*! \brief Whether the current node is under the given block. */
bool under_writer_block_{false};
/*! \brief function to update read/write region of block being cache write.*/
- std::function<ffi::Array<BufferRegion>(ffi::Array<BufferRegion>)> update_access_regions;
+ std::function<ffi::Array<TensorRegion>(ffi::Array<TensorRegion>)> update_access_regions;
/*! \brief function to update match buffers of block being cache write.*/
std::function<ffi::Array<MatchBufferRegion>(ffi::Array<MatchBufferRegion>)> update_match_buffers;
/*!
@@ -1541,10 +1553,10 @@ class ReindexCacheWriteRewriter : public CacheWriteRewriter {
ReindexCacheStageInfo* info)
: CacheWriteRewriter(scope_sref, writer_block_sref, info) {
new_indices_ = info->indices;
- update_access_regions = [&](ffi::Array<BufferRegion> reads) {
- ffi::Array<BufferRegion> new_reads;
- for (const BufferRegion& buf_region : reads) {
- if (buf_region->buffer.same_as(info_->write_buffer)) {
+ update_access_regions = [&](ffi::Array<TensorRegion> reads) {
+ ffi::Array<TensorRegion> new_reads;
+ for (const TensorRegion& buf_region : reads) {
+ if (buf_region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(info_->write_buffer)) {
Region region;
for (const PrimExpr index : new_indices_) {
region.push_back(Range::FromMinExtent(index, IntImm::Int32(1)));
@@ -1559,8 +1571,8 @@ class ReindexCacheWriteRewriter : public CacheWriteRewriter {
update_match_buffers = [&](const ffi::Array<MatchBufferRegion> match_buffers) {
ffi::Array<MatchBufferRegion> new_match_buffers;
for (const MatchBufferRegion& match_buffer_region : match_buffers) {
- BufferRegion source = match_buffer_region->source;
- if (source->buffer.same_as(info_->write_buffer)) {
+ TensorRegion source = match_buffer_region->source;
+ if (source->source.as_or_throw<tvm::tirx::BufferVar>().same_as(info_->write_buffer)) {
Region region;
for (const PrimExpr index : new_indices_) {
region.push_back(Range::FromMinExtent(index, IntImm::Int32(1)));
@@ -1826,11 +1838,11 @@ class ReIndexRewriter : public StmtExprMutator {
.as_or_throw<SBlock>();
// Update block reads/writes to use the intermediate reindex buffer
auto writes =
- ReplaceBufferRegion(block->writes, old_buffer_, BufferRegion{new_buffer_, region_});
+ ReplaceBufferRegion(block->writes, old_buffer_, BufferRegion(new_buffer_, region_));
auto reads =
- ReplaceBufferRegion(block->reads, old_buffer_, BufferRegion{new_buffer_, region_});
+ ReplaceBufferRegion(block->reads, old_buffer_, BufferRegion(new_buffer_, region_));
auto match_buffers = ReplaceBufferRegion(block->match_buffers, old_buffer_,
- BufferRegion{new_buffer_, region_});
+ BufferRegion(new_buffer_, region_));
if (!writes.same_as(block->writes) || !reads.same_as(block->reads) ||
!match_buffers.same_as(block->match_buffers)) {
ffi::ObjectPtr<SBlockNode> n = ffi::make_object<SBlockNode>(*stmt.as<SBlockNode>());
@@ -1921,8 +1933,8 @@ The region cover property require to hold for every of its child blocks
for (const auto& child_block_sref : GetChildBlocks(self, scope_root)) {
const SBlockNode* child_block = TVM_SREF_TO_SBLOCK(child_block_sref);
- for (const BufferRegion& region : child_block->reads) {
- if (region->buffer.same_as(read_buffer)) {
+ for (const TensorRegion& region : child_block->reads) {
+ if (region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(read_buffer)) {
if (!self->block_info.at(child_block_sref).region_cover) {
const SBlockNode* block = TVM_SREF_TO_SBLOCK(scope_root);
throw MakeScheduleError<NotRegionCoverError>(self->mod, ffi::GetRef<SBlock>(block));
@@ -1972,14 +1984,14 @@ StmtSRef CacheRead(ScheduleState self, const StmtSRef& block_sref, int read_buff
}
// Step 3. Update cache stage info.
- BufferRegion cache_region{nullptr};
+ TensorRegion cache_region{nullptr};
if (ffi::Optional<StmtSRef> _write_block_sref =
GetOnlyWriteBlock(self, scope_sref, read_buffer)) {
// Case 1. The buffer is written inside the block.
StmtSRef write_block_sref = _write_block_sref.value();
const SBlockNode* write_block = TVM_SREF_TO_SBLOCK(write_block_sref);
// Find the producing region
- BufferRegion region = GetBufferRegionFromBuffer(write_block->writes, read_buffer).value();
+ TensorRegion region = GetBufferRegionFromBuffer(write_block->writes, read_buffer).value();
StmtSRef parent_sref = ffi::GetRef<StmtSRef>(write_block_sref->parent);
// Detect insert position
@@ -1995,7 +2007,7 @@ StmtSRef CacheRead(ScheduleState self, const StmtSRef& block_sref, int read_buff
// block's read region under that predicate to get a tighter cache allocation.
// Without a nested predicate we fall back to scope_block->reads (which preserves the
// original buffer's dtype in its extents, e.g. int64 shapes).
- ffi::Optional<BufferRegion> read_region_opt =
+ ffi::Optional<TensorRegion> read_region_opt =
GetBufferRegionFromBuffer(block->reads, read_buffer);
PrimExpr nested_pred = read_region_opt ? CollectNestedBlockPredicates(block->body, read_buffer,
BufferIndexType::kRead)
@@ -2004,11 +2016,11 @@ StmtSRef CacheRead(ScheduleState self, const StmtSRef& block_sref, int read_buff
StmtSRef parent_sref = ffi::GetRef<StmtSRef>(block_sref->parent);
cache_region = RelaxBufferRegion(self, read_region_opt.value(), block_sref, parent_sref,
scope_sref, nested_pred);
- } else if (ffi::Optional<BufferRegion> scope_region =
+ } else if (ffi::Optional<TensorRegion> scope_region =
GetBufferRegionFromBuffer(scope_block->reads, read_buffer)) {
cache_region = scope_region.value();
} else {
- cache_region = BufferRegion::FullRegion(read_buffer);
+ cache_region = FullBufferRegion(read_buffer);
}
}
@@ -2085,7 +2097,7 @@ StmtSRef CacheWrite(ScheduleState self, const StmtSRef& block_sref, int write_bu
TVM_FFI_ICHECK_EQ(block_sref.get(), only_write_block.value().get());
// Step 4. Find the producing region and insert position
- BufferRegion region = GetBufferRegionFromBuffer(block->writes, write_buffer).value();
+ TensorRegion region = GetBufferRegionFromBuffer(block->writes, write_buffer).value();
// Detect insert position
CacheLocDetector::Detect</*is_cache_read=*/false>(self, block_sref, scope_sref, &info);
// Collect predicates from any nested blocks that gate the actual write (e.g. T.where on an
@@ -2093,7 +2105,7 @@ StmtSRef CacheWrite(ScheduleState self, const StmtSRef& block_sref, int write_bu
// is restricted by a nested predicate, so we OR them together for a tighter region estimate.
PrimExpr nested_write_pred =
CollectNestedBlockPredicates(block->body, write_buffer, BufferIndexType::kWrite);
- BufferRegion cache_region;
+ TensorRegion cache_region;
if (block_sref->parent != nullptr) {
StmtSRef parent_sref = ffi::GetRef<StmtSRef>(block_sref->parent);
cache_region =
@@ -2205,7 +2217,7 @@ template <bool is_cache_read>
void CollectReindexCacheStageInfoAndCreateBuffer(
ReindexCacheStageInfo* info, const IRModule& mod, const StmtSRef& block_sref,
const ffi::String& storage_scope, const IndexMap& index_map, const SBlock& block,
- const SBlockRealize& realize, const BufferVar& old_buffer, const BufferRegion& cache_region) {
+ const SBlockRealize& realize, const BufferVar& old_buffer, const TensorRegion& cache_region) {
arith::Analyzer analyzer;
ffi::Array<PrimExpr> block_iter_vars, block_shape;
for (const IterVar& iter_var : block->iter_vars) {
@@ -2272,7 +2284,7 @@ void CollectReindexCacheStageInfoAndCreateBuffer(
/*! \brief Check whether given cache_region is a single point access. */
template <bool is_cache_read>
-void CheckSinglePoint(ScheduleState self, const SBlock& block, const BufferRegion& cache_region) {
+void CheckSinglePoint(ScheduleState self, const SBlock& block, const TensorRegion& cache_region) {
bool single_point = true;
for (const Range& range : cache_region->region) {
const auto* ext_int = range->extent.as<IntImmNode>();
@@ -2314,10 +2326,10 @@ StmtSRef ReindexCacheRead(ScheduleState self, const StmtSRef& block_sref, int re
info.consumer_blocks.insert(block_sref);
// Step 3. Update cache stage info.
- ffi::Optional<BufferRegion> maybe_region = GetBufferRegionFromBuffer(block->reads, read_buffer);
+ ffi::Optional<TensorRegion> maybe_region = GetBufferRegionFromBuffer(block->reads, read_buffer);
TVM_FFI_ICHECK(maybe_region.has_value())
<< read_buffer << " should appear in the block's read region: " << block->reads;
- BufferRegion cache_region = maybe_region.value();
+ TensorRegion cache_region = maybe_region.value();
if (ffi::Optional<StmtSRef> _write_block_sref =
GetOnlyWriteBlock(self, scope_sref, read_buffer)) {
// Case 1. The buffer is written inside the block.
@@ -2390,13 +2402,13 @@ StmtSRef ReindexCacheWrite(ScheduleState self, const StmtSRef& block_sref, int w
TVM_FFI_ICHECK_EQ(block_sref.get(), only_write_block.value().get());
// Step 4. Find the producing region and insert position
- ffi::Optional<BufferRegion> maybe_region = GetBufferRegionFromBuffer(block->writes, write_buffer);
+ ffi::Optional<TensorRegion> maybe_region = GetBufferRegionFromBuffer(block->writes, write_buffer);
TVM_FFI_ICHECK(maybe_region.has_value())
<< write_buffer << " should appear in the block's write region";
StmtSRef parent_sref = ffi::GetRef<StmtSRef>(block_sref->parent);
// Detect insert position
CacheLocDetector::Detect</*is_cache_read=*/false>(self, block_sref, scope_sref, &info);
- BufferRegion cache_region = maybe_region.value();
+ TensorRegion cache_region = maybe_region.value();
CollectReindexCacheStageInfoAndCreateBuffer</*is_cache_read=*/false>(
&info, self->mod, block_sref, storage_scope, index_map, block, realize, write_buffer,
@@ -2464,8 +2476,8 @@ ffi::Array<StmtSRef> CacheInplace(ScheduleState self, const StmtSRef& block_sref
// Check 4. Check if target block both read & write target buffer.
const SBlockNode* rw_block = TVM_SREF_TO_SBLOCK(block_sref);
- ffi::Optional<BufferRegion> read_region = GetBufferRegionFromBuffer(rw_block->reads, buffer);
- ffi::Optional<BufferRegion> write_region = GetBufferRegionFromBuffer(rw_block->writes, buffer);
+ ffi::Optional<TensorRegion> read_region = GetBufferRegionFromBuffer(rw_block->reads, buffer);
+ ffi::Optional<TensorRegion> write_region = GetBufferRegionFromBuffer(rw_block->writes, buffer);
if (!read_region.has_value() || !write_region.has_value()) {
throw MakeScheduleError<NotReadWriteError>(self->mod, ffi::GetRef<SBlock>(rw_block), buffer);
}
diff --git a/src/s_tir/schedule/primitive/compute_at.cc b/src/s_tir/schedule/primitive/compute_at.cc
index 7f8b752fd8..6b8d920c45 100644
--- a/src/s_tir/schedule/primitive/compute_at.cc
+++ b/src/s_tir/schedule/primitive/compute_at.cc
@@ -372,7 +372,7 @@ class ScopeReconstructor : public StmtExprMutator {
*/
template <bool relax_storage_scope>
void RelaxBufferRegions(const ffi::Map<Var, PrimExpr>& binding,
- const ffi::Array<BufferRegion>& buffer_regions,
+ const ffi::Array<TensorRegion>& buffer_regions,
const StmtSRef& relax_path_low_inclusive,
const StmtSRef& relax_path_high_exclusive,
std::unordered_map<const VarNode*, std::vector<NDIntSet>>* relaxed) {
@@ -385,8 +385,8 @@ void RelaxBufferRegions(const ffi::Map<Var, PrimExpr>& binding,
return ffi::Unchanged();
};
// Enumerate every buffer region
- for (const BufferRegion& buffer_region : buffer_regions) {
- const BufferVar& buffer = buffer_region->buffer;
+ for (const TensorRegion& buffer_region : buffer_regions) {
+ const BufferVar& buffer = buffer_region->source.as_or_throw<tvm::tirx::BufferVar>();
const ffi::Array<Range>& region = buffer_region->region;
// Skip the buffer regions we are not interested in
auto it = relaxed->find(buffer.get());
@@ -490,7 +490,7 @@ std::pair<Var, BlockVarDomainInfo> SolveBlockVarDomain(const arith::IntSet& prov
}
}
TVM_FFI_CHECK(var.has_value(), ValueError)
- << "BufferRegion pattern match failed: " << provided_min;
+ << "TensorRegion pattern match failed: " << provided_min;
return {var.value(), BlockVarDomainInfo{var_dom, var_bound}};
}
@@ -691,11 +691,12 @@ void CalculateProvidedRequiredRegions(
std::unordered_map<const VarNode*, std::vector<NDIntSet>>* provided_regions,
std::unordered_map<const VarNode*, std::vector<NDIntSet>>* required_regions) {
// Step 1. Calculate the region provided by a single execution instance of `block`
- const ffi::Array<BufferRegion>& provided_buffers = is_compute_at ? block->writes : block->reads;
+ const ffi::Array<TensorRegion>& provided_buffers = is_compute_at ? block->writes : block->reads;
provided_regions->reserve(provided_buffers.size());
required_regions->reserve(provided_buffers.size());
- for (const BufferRegion& provided_buffer_region : provided_buffers) {
- const VarNode* buffer = provided_buffer_region->buffer.get();
+ for (const TensorRegion& provided_buffer_region : provided_buffers) {
+ const VarNode* buffer =
+ provided_buffer_region->source.as_or_throw<tvm::tirx::BufferVar>().get();
const ffi::Array<Range>& region = provided_buffer_region->region;
(*provided_regions)[buffer].push_back(support::NDIntSetFromRegion(region));
(*required_regions)[buffer].clear();
diff --git a/src/s_tir/schedule/primitive/compute_inline.cc b/src/s_tir/schedule/primitive/compute_inline.cc
index cb3af7717d..1340a49a78 100644
--- a/src/s_tir/schedule/primitive/compute_inline.cc
+++ b/src/s_tir/schedule/primitive/compute_inline.cc
@@ -100,8 +100,8 @@ class NotSingleReadWriteBuffer : public ScheduleErrorContextObj {
ffi::ObjectPtrEqual>& buffer_writers =
self->block_info.at(scope_root_sref).scope->buffer_writers;
const VarNode* read_buffer = nullptr;
- for (const BufferRegion& read_region : block->reads) {
- const VarNode* buffer = read_region->buffer.get();
+ for (const TensorRegion& read_region : block->reads) {
+ const VarNode* buffer = read_region->source.as_or_throw<tvm::tirx::BufferVar>().get();
if (buffer == read_buffer) {
continue;
}
@@ -122,7 +122,7 @@ class NotSingleReadWriteBuffer : public ScheduleErrorContextObj {
if (block->writes.size() != 1) {
throw MakeScheduleError<NotSingleReadWriteBuffer>(self->mod, false, block);
}
- return block->writes[0]->buffer;
+ return block->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>();
}
};
@@ -224,7 +224,7 @@ class NonSingleProducerError : public ScheduleErrorContextObj {
// leaf block
producer_across_scope_.pop_back();
for (const auto& write : node->writes) {
- if (write->buffer.same_as(buffer_)) {
+ if (write->source.as_or_throw<tvm::tirx::BufferVar>().same_as(buffer_)) {
// Check if the producer block is a complete block
StmtSRef producer_block_sref = self_->stmt2ref.at(node);
if (!IsCompleteBlock(self_, producer_block_sref, scope_root_sref_)) {
@@ -323,10 +323,13 @@ class BaseInliner : public StmtExprMutator {
}
protected:
- UnchangedOr<Expr> Mutate_(const BufferRegionNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Expr> Mutate_(const TensorRegionNode* op, InplaceMode inplace_mode) final {
+ if (!op->source.as<BufferVar>()) {
+ return StmtExprMutator::Mutate_(op, inplace_mode);
+ }
auto region = Mutate(op->region).as_or_throw<UnchangedOr<ffi::Array<Range>>>();
if (region.UnchangedOrSameAs(op->region)) return ffi::Unchanged();
- BufferRegion node = ffi::GetRef<BufferRegion>(op);
+ TensorRegion node = ffi::GetRef<TensorRegion>(op);
node.CopyOnWrite()->region = std::move(region).ValueUnchecked();
return node;
}
@@ -393,12 +396,12 @@ class BaseInliner : public StmtExprMutator {
* \param block The block whose signature to be added
*/
void AddBuffersInBlockSignature(const SBlockNode* block) {
- for (const BufferRegion& buffer_region : block->reads) {
- const BufferVar& buffer = buffer_region->buffer;
+ for (const TensorRegion& buffer_region : block->reads) {
+ const BufferVar& buffer = buffer_region->source.as_or_throw<tvm::tirx::BufferVar>();
buffer_var_map_.Set(buffer.var(), buffer);
}
- for (const BufferRegion& buffer_region : block->writes) {
- const BufferVar& buffer = buffer_region->buffer;
+ for (const TensorRegion& buffer_region : block->writes) {
+ const BufferVar& buffer = buffer_region->source.as_or_throw<tvm::tirx::BufferVar>();
buffer_var_map_.Set(buffer.var(), buffer);
}
for (const BufferVar& buffer : block->alloc_buffers) {
@@ -429,14 +432,14 @@ class BaseInliner : public StmtExprMutator {
alloc_buffers = std::move(block->alloc_buffers);
}
// Step 2. Update `BlockNode::reads` and `BlockNode::writes`
- ffi::Array<BufferRegion> reads = std::move(block->reads);
- ffi::Array<BufferRegion> writes = std::move(block->writes);
- auto f_access_inline_buffer = [this](const BufferRegion& access) {
- return access->buffer.same_as(this->inlined_buffer_);
+ ffi::Array<TensorRegion> reads = std::move(block->reads);
+ ffi::Array<TensorRegion> writes = std::move(block->writes);
+ auto f_access_inline_buffer = [this](const TensorRegion& access) {
+ return access->source.as_or_throw<tvm::tirx::BufferVar>().same_as(this->inlined_buffer_);
};
if (!is_scope_root && (std::any_of(reads.begin(), reads.end(), f_access_inline_buffer) ||
std::any_of(writes.begin(), writes.end(), f_access_inline_buffer))) {
- ffi::Array<ffi::Array<BufferRegion>> inspected =
+ ffi::Array<ffi::Array<TensorRegion>> inspected =
GetSBlockReadWriteRegion(block, buffer_var_map_);
reads = inspected[0];
writes = inspected[1];
@@ -467,7 +470,8 @@ class BaseInliner : public StmtExprMutator {
*/
void CheckMatchBufferRegion(const SBlockNode* block) {
for (const MatchBufferRegion& match_buffer_region : block->match_buffers) {
- const BufferVar& matched = match_buffer_region->source->buffer;
+ const BufferVar& matched =
+ match_buffer_region->source->source.as_or_throw<tvm::tirx::BufferVar>();
if (matched.same_as(inlined_buffer_)) {
this->has_opaque_access = true;
}
@@ -1145,9 +1149,9 @@ class ReductionEpilogueFuser : public BaseInliner {
nullptr}; // The reduction buffer load in epilogue expression
BufferVar epilogue_output_buffer_{nullptr}; // Output buffer D
ffi::Array<PrimExpr> epilogue_output_indices_{nullptr}; // Indices of D[vi, vj]
- BufferRegion epilogue_output_region_{nullptr}; // Write region of D
+ TensorRegion epilogue_output_region_{nullptr}; // Write region of D
BufferVar epilogue_addend_buffer_{nullptr}; // Additional buffer (e.g., bias buffer C)
- BufferRegion epilogue_addend_region_{nullptr}; // Read region of additional buffer
+ TensorRegion epilogue_addend_region_{nullptr}; // Read region of additional buffer
};
bool ReductionEpilogueFuser::BodyPatternAllowFusion(const SBlockRealize& epilogue_block_realize) {
@@ -1296,8 +1300,8 @@ void ReductionEpilogueFuser::ExtractEpilogueInfo() {
epilogue_output_indices_ = inlined_store_->indices;
// Extract epilogue output region from epilogue block writes
- for (const BufferRegion& write : epilogue_block_->writes) {
- if (write->buffer.same_as(epilogue_output_buffer_)) {
+ for (const TensorRegion& write : epilogue_block_->writes) {
+ if (write->source.as_or_throw<tvm::tirx::BufferVar>().same_as(epilogue_output_buffer_)) {
epilogue_output_region_ = write;
break;
}
@@ -1327,8 +1331,8 @@ void ReductionEpilogueFuser::ExtractEpilogueInfo() {
const VarNode* first_buffer = *extractor->other_buffers.begin();
epilogue_addend_buffer_ = BufferVar(ffi::GetRef<Var>(first_buffer));
// Find the read region from epilogue block reads
- for (const BufferRegion& read : epilogue_block_->reads) {
- if (read->buffer.get() == first_buffer) {
+ for (const TensorRegion& read : epilogue_block_->reads) {
+ if (read->source.as_or_throw<tvm::tirx::BufferVar>().get() == first_buffer) {
epilogue_addend_region_ = read;
break;
}
@@ -1601,9 +1605,9 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock(
new_block->body = replacer->Mutate(reduction_block->body).ValueOrUnchanged(reduction_block->body);
// 4. Update write regions
- ffi::Array<BufferRegion> new_writes;
- for (const BufferRegion& write : reduction_block->writes) {
- if (write->buffer.same_as(inlined_buffer_)) {
+ ffi::Array<TensorRegion> new_writes;
+ for (const TensorRegion& write : reduction_block->writes) {
+ if (write->source.as_or_throw<tvm::tirx::BufferVar>().same_as(inlined_buffer_)) {
ffi::Array<Range> mapped_region = write->region.Map([&f_substitute](const Range& range) {
PrimExpr min = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(range->min, f_substitute)
.as_or_throw<PrimExpr>();
@@ -1619,12 +1623,12 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock(
new_block->writes = new_writes;
// 5. Update read regions: add all buffers from epilogue expression (except reduction buffer)
- ffi::Array<BufferRegion> new_reads;
+ ffi::Array<TensorRegion> new_reads;
std::unordered_set<const VarNode*> read_bufs;
// Add all non-reduction buffers from epilogue expression
- for (const BufferRegion& read : epilogue_block_->reads) {
- if (!read->buffer.same_as(inlined_buffer_)) {
+ for (const TensorRegion& read : epilogue_block_->reads) {
+ if (!read->source.as_or_throw<tvm::tirx::BufferVar>().same_as(inlined_buffer_)) {
ffi::Array<Range> mapped_region = read->region.Map([&f_substitute](const Range& range) {
PrimExpr min = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(range->min, f_substitute)
.as_or_throw<PrimExpr>();
@@ -1632,18 +1636,20 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock(
.as_or_throw<PrimExpr>();
return Range::FromMinExtent(min, extent);
});
- new_reads.push_back(BufferRegion(read->buffer, mapped_region));
- read_bufs.insert(read->buffer.get());
+ new_reads.push_back(
+ BufferRegion(read->source.as_or_throw<tvm::tirx::BufferVar>(), mapped_region));
+ read_bufs.insert(read->source.as_or_throw<tvm::tirx::BufferVar>().get());
}
}
// Add existing read regions from reduction block (A, B, etc.)
- for (const BufferRegion& read : reduction_block->reads) {
- if (!read->buffer.same_as(inlined_buffer_)) {
+ for (const TensorRegion& read : reduction_block->reads) {
+ if (!read->source.as_or_throw<tvm::tirx::BufferVar>().same_as(inlined_buffer_)) {
// Only add non-temp buffers that haven't been added yet
- if (read_bufs.find(read->buffer.get()) == read_bufs.end()) {
+ if (read_bufs.find(read->source.as_or_throw<tvm::tirx::BufferVar>().get()) ==
+ read_bufs.end()) {
new_reads.push_back(read);
- read_bufs.insert(read->buffer.get());
+ read_bufs.insert(read->source.as_or_throw<tvm::tirx::BufferVar>().get());
}
}
}
@@ -1688,16 +1694,16 @@ static bool CheckBufferStillUsed(const SBlock& scope_root, const BufferVar& buff
}
// Check reads
- for (const BufferRegion& read : block->reads) {
- if (read->buffer.same_as(buffer_)) {
+ for (const TensorRegion& read : block->reads) {
+ if (read->source.as_or_throw<tvm::tirx::BufferVar>().same_as(buffer_)) {
found_usage_ = true;
return std::nullopt;
}
}
// Check writes
- for (const BufferRegion& write : block->writes) {
- if (write->buffer.same_as(buffer_)) {
+ for (const TensorRegion& write : block->writes) {
+ if (write->source.as_or_throw<tvm::tirx::BufferVar>().same_as(buffer_)) {
found_usage_ = true;
return std::nullopt;
}
diff --git a/src/s_tir/schedule/primitive/decompose_padding.cc b/src/s_tir/schedule/primitive/decompose_padding.cc
index ccc0c12de2..6a4d03ab24 100644
--- a/src/s_tir/schedule/primitive/decompose_padding.cc
+++ b/src/s_tir/schedule/primitive/decompose_padding.cc
@@ -239,10 +239,11 @@ static std::pair<Stmt, SBlockRealize> CreateConstBlock(const SBlockRealizeNode*
// create new write region
TVM_FFI_ICHECK_EQ(block->writes.size(), 1U);
- BufferRegion write_region = BufferRegion(
- block->writes[0]->buffer, block->writes[0]->region.Map([rewrite_expr](const Range& r) {
- return Range::FromMinExtent(rewrite_expr(r->min), rewrite_expr(r->extent));
- }));
+ TensorRegion write_region =
+ BufferRegion(block->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>(),
+ block->writes[0]->region.Map([rewrite_expr](const Range& r) {
+ return Range::FromMinExtent(rewrite_expr(r->min), rewrite_expr(r->extent));
+ }));
// create block to fill const pad values
BufferStore store = block->body.as_or_throw<BufferStore>();
@@ -345,12 +346,14 @@ static std::pair<Stmt, SBlockRealize> CreateInBoundBlock(const SBlockRealizeNode
};
// create new read/write region for in-bound accesses
- ffi::Array<BufferRegion> reads, writes;
- for (const BufferRegion& read : block->reads) {
- reads.push_back(BufferRegion(read->buffer, rewrite_region(read->region)));
+ ffi::Array<TensorRegion> reads, writes;
+ for (const TensorRegion& read : block->reads) {
+ reads.push_back(BufferRegion(read->source.as_or_throw<tvm::tirx::BufferVar>(),
+ rewrite_region(read->region)));
}
- for (const BufferRegion& write : block->writes) {
- writes.push_back(BufferRegion(write->buffer, rewrite_region(write->region)));
+ for (const TensorRegion& write : block->writes) {
+ writes.push_back(BufferRegion(write->source.as_or_throw<tvm::tirx::BufferVar>(),
+ rewrite_region(write->region)));
}
// create new block realize node
diff --git a/src/s_tir/schedule/primitive/hide_buffer_access.cc b/src/s_tir/schedule/primitive/hide_buffer_access.cc
index 605069b543..43c0322f28 100644
--- a/src/s_tir/schedule/primitive/hide_buffer_access.cc
+++ b/src/s_tir/schedule/primitive/hide_buffer_access.cc
@@ -113,7 +113,7 @@ void UnsafeHideBufferAccess(ScheduleState self, const StmtSRef& block_sref,
/* Step 0: Collect new buffer access regions. */
- ffi::Array<BufferRegion> reads, writes;
+ ffi::Array<TensorRegion> reads, writes;
if (buf_type == "read") {
for (size_t i = 0; i < block->reads.size(); ++i) {
diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc b/src/s_tir/schedule/primitive/layout_transformation.cc
index 1cff59598f..72475f2465 100644
--- a/src/s_tir/schedule/primitive/layout_transformation.cc
+++ b/src/s_tir/schedule/primitive/layout_transformation.cc
@@ -532,7 +532,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
std::stringstream block_name;
block_name << "buffer_" << new_buffer.name() << "_assumptions";
- auto read_region = BufferRegion::FromPoint(new_buffer, indices);
+ auto read_region = BufferRegionFromPoint(new_buffer, indices);
stmt = SBlockRealize(iter_values, IntImm::Bool(true),
SBlock(iter_vars, {read_region}, {}, block_name.str(), stmt));
@@ -625,7 +625,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
std::stringstream block_name;
block_name << "buffer_" << new_buffer.name() << "_padding";
- auto write_region = BufferRegion::FromPoint(new_buffer, indices);
+ auto write_region = BufferRegionFromPoint(new_buffer, indices);
stmt = SBlockRealize(iter_values, padding_predicate,
SBlock(iter_vars, {}, {write_region}, block_name.str(), stmt));
@@ -915,17 +915,17 @@ class TransformLayoutRewriter : public tirx::IRMutatorWithAnalyzer {
return buffer_store;
}
- void RewriteAccessRegion(ffi::Array<BufferRegion>* old_access_regions,
- const ffi::Array<BufferRegion>& infered_access_regions) {
- auto fmutate = [this, &infered_access_regions](const BufferRegion& buffer_region) {
- if (buffer_region->buffer.same_as(old_buffer_)) {
+ void RewriteAccessRegion(ffi::Array<TensorRegion>* old_access_regions,
+ const ffi::Array<TensorRegion>& infered_access_regions) {
+ auto fmutate = [this, &infered_access_regions](const TensorRegion& buffer_region) {
+ if (buffer_region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(old_buffer_)) {
TVM_FFI_ICHECK(infered_access_regions.size() == 1);
- BufferRegion result = infered_access_regions[0];
+ TensorRegion result = infered_access_regions[0];
// The inferred region may reference old_buffer_ (e.g. when resolved
// through match_buffer source). Ensure we use new_buffer_ instead.
- if (result->buffer.same_as(old_buffer_)) {
+ if (result->source.as_or_throw<tvm::tirx::BufferVar>().same_as(old_buffer_)) {
auto* n = result.CopyOnWrite();
- n->buffer = new_buffer_;
+ n->source = new_buffer_;
}
return result;
}
@@ -957,10 +957,10 @@ class TransformLayoutRewriter : public tirx::IRMutatorWithAnalyzer {
RewriteAccessRegion(&n->writes, infered_access_regions[1]);
// Update match_buffers whose source references old_buffer_
n->match_buffers.MutateByApply([this](const MatchBufferRegion& match_buf) {
- if (match_buf->source->buffer.same_as(old_buffer_)) {
+ if (match_buf->source->source.as_or_throw<tvm::tirx::BufferVar>().same_as(old_buffer_)) {
auto new_source = match_buf->source;
auto* source_n = new_source.CopyOnWrite();
- source_n->buffer = new_buffer_;
+ source_n->source = new_buffer_;
auto new_match = match_buf;
new_match.CopyOnWrite()->source = new_source;
return new_match;
diff --git a/src/s_tir/schedule/primitive/loop_transformation.cc b/src/s_tir/schedule/primitive/loop_transformation.cc
index 0044577f80..91e165d54e 100644
--- a/src/s_tir/schedule/primitive/loop_transformation.cc
+++ b/src/s_tir/schedule/primitive/loop_transformation.cc
@@ -547,17 +547,19 @@ class BufferIndicesMapExtractor : public StmtExprVisitor {
ffi::Map<ffi::String, ffi::Array<ffi::String>> buffer_indices_map;
};
-ffi::Array<BufferRegion> MutateBufferRegion(
+ffi::Array<TensorRegion> MutateBufferRegion(
ffi::Map<ffi::String, ffi::Array<ffi::String>> buffer_indices_map,
- ffi::Map<ffi::String, Range> index_range_map, ffi::Array<BufferRegion> region_arr) {
- // Update the region with new Ranges and return new BufferRegion
- ffi::Array<BufferRegion> new_region_arr =
- MutateArray(region_arr, [&buffer_indices_map, &index_range_map](const BufferRegion& region) {
- BufferRegion new_region = region;
- auto it = buffer_indices_map.find(new_region->buffer.name());
+ ffi::Map<ffi::String, Range> index_range_map, ffi::Array<TensorRegion> region_arr) {
+ // Update the region with new Ranges and return new TensorRegion
+ ffi::Array<TensorRegion> new_region_arr =
+ MutateArray(region_arr, [&buffer_indices_map, &index_range_map](const TensorRegion& region) {
+ TensorRegion new_region = region;
+ auto it =
+ buffer_indices_map.find(new_region->source.as_or_throw<tvm::tirx::BufferVar>().name());
if (it == buffer_indices_map.end()) return new_region;
- ffi::Array<ffi::String> old_indices = buffer_indices_map[new_region->buffer.name()];
+ ffi::Array<ffi::String> old_indices =
+ buffer_indices_map[new_region->source.as_or_throw<tvm::tirx::BufferVar>().name()];
ffi::Array<Range> new_ranges;
for (size_t i = 0; i < old_indices.size(); i++) {
new_ranges.push_back(index_range_map[old_indices[i]]);
@@ -624,13 +626,13 @@ class BlockMutator : public StmtExprMutator {
// Get the (BufferVar, indices) map
ffi::Map<ffi::String, ffi::Array<ffi::String>> buffer_indices_map =
BufferIndicesMapExtractor::Extract(new_loop_var_, new_block);
- ffi::Array<BufferRegion> new_writes =
+ ffi::Array<TensorRegion> new_writes =
MutateBufferRegion(buffer_indices_map, index_range_map, new_block->writes);
if (!new_block->writes.same_as(new_writes)) {
// Update the writes with new_writes
new_block.CopyOnWrite()->writes = std::move(new_writes);
}
- ffi::Array<BufferRegion> new_reads =
+ ffi::Array<TensorRegion> new_reads =
MutateBufferRegion(buffer_indices_map, index_range_map, new_block->reads);
if (!new_block->reads.same_as(new_reads)) {
// Update the reads with new_reads
diff --git a/src/s_tir/schedule/primitive/pad_einsum.cc b/src/s_tir/schedule/primitive/pad_einsum.cc
index 7ad9fe9333..05de3ac334 100644
--- a/src/s_tir/schedule/primitive/pad_einsum.cc
+++ b/src/s_tir/schedule/primitive/pad_einsum.cc
@@ -48,7 +48,7 @@ ffi::Optional<ffi::Array<Var>> CheckTrivialBufferIndices(
return indices;
}
-ffi::Optional<ffi::Array<Var>> CheckTrivialBufferAccess(const BufferRegion& buffer_region) {
+ffi::Optional<ffi::Array<Var>> CheckTrivialBufferAccess(const TensorRegion& buffer_region) {
ffi::Array<Var> indices;
indices.reserve(buffer_region->region.size());
for (const Range& range : buffer_region->region) {
@@ -137,10 +137,10 @@ struct BufferPadding {
BufferVar buffer;
BufferVar padded_buffer;
- static BufferPadding FromBufferRegion(const BufferRegion& buffer_region,
+ static BufferPadding FromBufferRegion(const TensorRegion& buffer_region,
const ffi::Map<Var, PrimExpr>& iter_extents) {
BufferPadding result;
- result.buffer = buffer_region->buffer;
+ result.buffer = buffer_region->source.as_or_throw<tvm::tirx::BufferVar>();
ffi::Array<PrimExpr> shape;
shape.reserve(buffer_region->region.size());
int ndim = buffer_region->region.size();
@@ -152,7 +152,7 @@ struct BufferPadding {
} else if (ffi::Optional<PrimExpr> extent = iter_extents.Get(pos.as_or_throw<Var>())) {
shape.push_back(extent.value());
} else {
- shape.push_back(buffer_region->buffer->shape[i]);
+ shape.push_back(buffer_region->source.as_or_throw<tvm::tirx::BufferVar>()->shape[i]);
}
}
result.padded_buffer = decl_buffer(shape, result.buffer->dtype, result.buffer.name() + "_pad",
@@ -196,8 +196,8 @@ struct BufferPadding {
} else {
body = BufferStore(buffer, BufferLoad(padded_buffer, indices), indices);
}
- BufferRegion read_region(buffer, instance_dom);
- BufferRegion write_region(padded_buffer, instance_dom);
+ TensorRegion read_region = BufferRegion(buffer, instance_dom);
+ TensorRegion write_region = BufferRegion(padded_buffer, instance_dom);
if (!is_read) {
std::swap(read_region, write_region);
}
@@ -221,7 +221,7 @@ Einsum ExtractEinsum(const ScheduleState& self, const SBlock& block) {
std::unordered_set<const VarNode*> buffer_used;
int n_reads = block->reads.size();
for (int i = 0; i < n_reads; ++i) {
- const BufferVar& buffer = block->reads[i]->buffer;
+ const BufferVar& buffer = block->reads[i]->source.as_or_throw<tvm::tirx::BufferVar>();
if (buffer_used.count(buffer.get()) != 0) {
throw MakeScheduleError<NonEinsumError>(self->mod, block);
}
@@ -235,7 +235,7 @@ Einsum ExtractEinsum(const ScheduleState& self, const SBlock& block) {
}
int n_writes = block->writes.size();
for (int i = 0; i < n_writes; ++i) {
- const BufferVar& buffer = block->writes[i]->buffer;
+ const BufferVar& buffer = block->writes[i]->source.as_or_throw<tvm::tirx::BufferVar>();
if (buffer_used.count(buffer.get()) != 0) {
throw MakeScheduleError<NonEinsumError>(self->mod, block);
}
@@ -322,19 +322,21 @@ class PadEinsumBufferReplacer : public StmtExprMutator {
iter_vars.push_back(iter_var);
}
}
- ffi::Array<BufferRegion> reads;
+ ffi::Array<TensorRegion> reads;
reads.reserve(block->reads.size());
- for (const BufferRegion& read : block->reads) {
- if (ffi::Optional<BufferVar> buffer = VarRemapGet(read->buffer).as<BufferVar>()) {
+ for (const TensorRegion& read : block->reads) {
+ if (ffi::Optional<BufferVar> buffer =
+ VarRemapGet(read->source.as_or_throw<tvm::tirx::BufferVar>()).as<BufferVar>()) {
reads.push_back(BufferRegion(buffer.value(), read->region));
} else {
reads.push_back(read);
}
}
- ffi::Array<BufferRegion> writes;
+ ffi::Array<TensorRegion> writes;
writes.reserve(block->writes.size());
- for (const BufferRegion& write : block->writes) {
- if (ffi::Optional<BufferVar> buffer = VarRemapGet(write->buffer).as<BufferVar>()) {
+ for (const TensorRegion& write : block->writes) {
+ if (ffi::Optional<BufferVar> buffer =
+ VarRemapGet(write->source.as_or_throw<tvm::tirx::BufferVar>()).as<BufferVar>()) {
writes.push_back(BufferRegion(buffer.value(), write->region));
} else {
writes.push_back(write);
@@ -451,7 +453,7 @@ void PadEinsum(ScheduleState self, const StmtSRef& block_sref, const ffi::Array<
ffi::Array<Stmt> write_blocks;
ffi::Array<SBlock> new_copy_blocks;
ffi::Array<BufferVar> alloc_buffers;
- for (const BufferRegion& buffer_region : block->reads) {
+ for (const TensorRegion& buffer_region : block->reads) {
if (f_needs_padding(buffer_region->region)) {
BufferPadding bp =
BufferPadding::FromBufferRegion(buffer_region, replacer->iter2padded_extents);
@@ -460,7 +462,7 @@ void PadEinsum(ScheduleState self, const StmtSRef& block_sref, const ffi::Array<
alloc_buffers.push_back(bp.padded_buffer);
}
}
- for (const BufferRegion& buffer_region : block->writes) {
+ for (const TensorRegion& buffer_region : block->writes) {
if (f_needs_padding(buffer_region->region)) {
BufferPadding bp =
BufferPadding::FromBufferRegion(buffer_region, replacer->iter2padded_extents);
diff --git a/src/s_tir/schedule/primitive/read_write_at.cc b/src/s_tir/schedule/primitive/read_write_at.cc
index bbabd40c9d..a6329e4db5 100644
--- a/src/s_tir/schedule/primitive/read_write_at.cc
+++ b/src/s_tir/schedule/primitive/read_write_at.cc
@@ -32,16 +32,16 @@ using namespace tvm::tirx;
using support::NDIntSet;
-bool HasBuffer(const ffi::Array<BufferRegion>& buffer_regions, const BufferVar& buffer) {
- for (const BufferRegion& buffer_region : buffer_regions) {
- if (buffer_region->buffer.same_as(buffer)) {
+bool HasBuffer(const ffi::Array<TensorRegion>& buffer_regions, const BufferVar& buffer) {
+ for (const TensorRegion& buffer_region : buffer_regions) {
+ if (buffer_region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(buffer)) {
return true;
}
}
return false;
}
-void RelaxBufferRegions(const ffi::Array<BufferRegion>& buffer_regions,
+void RelaxBufferRegions(const ffi::Array<TensorRegion>& buffer_regions,
const BufferVar& buffer, //
const ffi::Map<Var, arith::IntSet>& var_dom, //
const ffi::Map<Var, PrimExpr>& bindings, //
@@ -50,8 +50,8 @@ void RelaxBufferRegions(const ffi::Array<BufferRegion>& buffer_regions,
if (auto repl = bindings.Get(var)) return ffi::Any(*std::move(repl));
return ffi::Unchanged();
};
- for (const BufferRegion& buffer_region : buffer_regions) {
- if (buffer_region->buffer.same_as(buffer)) {
+ for (const TensorRegion& buffer_region : buffer_regions) {
+ if (buffer_region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(buffer)) {
ffi::Array<Range> mapped_region =
buffer_region->region.Map([&f_substitute](const Range& range) {
PrimExpr min = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(range->min, f_substitute)
diff --git a/src/s_tir/schedule/primitive/reduction.cc b/src/s_tir/schedule/primitive/reduction.cc
index ddfc8f3aca..cf6f246df2 100644
--- a/src/s_tir/schedule/primitive/reduction.cc
+++ b/src/s_tir/schedule/primitive/reduction.cc
@@ -83,17 +83,18 @@ class DecomposeReductionBlockReplacer : public StmtExprMutator {
p_new_block->name_hint = p_new_block->name_hint + "_update";
p_new_block->init = std::nullopt;
// Add write regions back to read regions in update block.
- ffi::Array<BufferRegion> new_reads;
+ ffi::Array<TensorRegion> new_reads;
std::unordered_set<const VarNode*> read_bufs;
- for (const BufferRegion& read_access : block->reads) {
- read_bufs.insert(read_access->buffer.get());
+ for (const TensorRegion& read_access : block->reads) {
+ read_bufs.insert(read_access->source.as_or_throw<tvm::tirx::BufferVar>().get());
}
- for (const BufferRegion& write_access : block->writes) {
- if (read_bufs.find(write_access->buffer.get()) == read_bufs.end()) {
+ for (const TensorRegion& write_access : block->writes) {
+ if (read_bufs.find(write_access->source.as_or_throw<tvm::tirx::BufferVar>().get()) ==
+ read_bufs.end()) {
new_reads.push_back(write_access);
}
}
- for (const BufferRegion& read_access : block->reads) {
+ for (const TensorRegion& read_access : block->reads) {
new_reads.push_back(read_access);
}
p_new_block->reads = new_reads;
@@ -256,7 +257,7 @@ StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref,
init_block->body =
ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(block->init.value(), map_block_var)
.as_or_throw<Stmt>();
- for (const BufferRegion& write : block->writes) {
+ for (const TensorRegion& write : block->writes) {
ffi::Array<Range> mapped_region = write->region.Map([&map_block_var](const Range& range) {
PrimExpr min = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(range->min, map_block_var)
.as_or_throw<PrimExpr>();
@@ -264,7 +265,8 @@ StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref,
.as_or_throw<PrimExpr>();
return Range::FromMinExtent(min, extent);
});
- init_block->writes.push_back(BufferRegion(write->buffer, mapped_region));
+ init_block->writes.push_back(
+ BufferRegion(write->source.as_or_throw<tvm::tirx::BufferVar>(), mapped_region));
}
// Step 3. Scan loops not higher than the specified loop above the reduction block.
// If the loop is used in the init block binding, then it is chosen.
@@ -932,9 +934,9 @@ class BaseBlockCreator {
/*! \brief THe RHS values of the reduction in this block */
ffi::Array<PrimExpr> update_rhs_;
/*! \brief The read regions of the new created block */
- ffi::Array<BufferRegion> read_regions_;
+ ffi::Array<TensorRegion> read_regions_;
/*! \brief The write regions of the new created block */
- ffi::Array<BufferRegion> write_regions_;
+ ffi::Array<TensorRegion> write_regions_;
};
/*!
@@ -1057,7 +1059,7 @@ class RFactorBlockCreator : public BaseBlockCreator {
return ffi::Unchanged();
};
read_regions_.reserve(old_block->reads.size());
- for (const BufferRegion& read_region : old_block->reads) {
+ for (const TensorRegion& read_region : old_block->reads) {
ffi::Array<Range> region = read_region->region.Map([&map_block_var](const Range& range) {
PrimExpr min = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(range->min, map_block_var)
.as_or_throw<PrimExpr>();
@@ -1066,15 +1068,17 @@ class RFactorBlockCreator : public BaseBlockCreator {
.as_or_throw<PrimExpr>();
return Range::FromMinExtent(min, extent);
});
- read_regions_.push_back(BufferRegion(read_region->buffer, region));
+ read_regions_.push_back(
+ BufferRegion(read_region->source.as_or_throw<tvm::tirx::BufferVar>(), region));
}
write_regions_.reserve(old_block->writes.size());
- for (const BufferRegion& write_region : old_block->writes) {
+ for (const TensorRegion& write_region : old_block->writes) {
ffi::Array<Range> region = write_region->region;
region.insert(
region.begin() + factor_axis_,
Range::FromMinExtent(additional_iter_->var, IntImm(additional_iter_->var.ty(), 1)));
- ffi::Optional<BufferVar> rf_buffer = buffer_map.Get(write_region->buffer);
+ ffi::Optional<BufferVar> rf_buffer =
+ buffer_map.Get(write_region->source.as_or_throw<tvm::tirx::BufferVar>());
TVM_FFI_ICHECK(rf_buffer.has_value());
region.MutateByApply([&map_block_var](const Range& range) {
PrimExpr min = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(range->min, map_block_var)
@@ -1175,7 +1179,7 @@ class WriteBackBlockCreator : public BaseBlockCreator {
}
void CreateRegion(const ffi::Array<PrimExpr>& buf_loads, bool is_read) {
- ffi::Array<BufferRegion>& buf_regions = is_read ? read_regions_ : write_regions_;
+ ffi::Array<TensorRegion>& buf_regions = is_read ? read_regions_ : write_regions_;
for (const PrimExpr& expr : buf_loads) {
const auto* buf_load = expr.as<TensorLoadNode>();
TVM_FFI_ICHECK(buf_load != nullptr);
diff --git a/src/s_tir/schedule/primitive/rolling_buffer.cc b/src/s_tir/schedule/primitive/rolling_buffer.cc
index b555432c90..4d83af50b7 100644
--- a/src/s_tir/schedule/primitive/rolling_buffer.cc
+++ b/src/s_tir/schedule/primitive/rolling_buffer.cc
@@ -42,7 +42,7 @@ struct RollingBufferInfo {
ffi::Map<SBlock, SBlock> block_reuse;
};
-BufferRegion GetRelaxedBufferRegion(const SBlockRealize& realize, const BufferRegion& buffer_region,
+TensorRegion GetRelaxedBufferRegion(const SBlockRealize& realize, const TensorRegion& buffer_region,
const ffi::Map<Var, arith::IntSet>& dom_map) {
ffi::Map<Var, PrimExpr> bindings = GetBindings(realize);
auto f_substitute = [&bindings](const Var& var) -> ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
@@ -60,10 +60,10 @@ BufferRegion GetRelaxedBufferRegion(const SBlockRealize& realize, const BufferRe
Region relaxed_region;
relaxed_region.reserve(relaxed_intsets.size());
for (size_t i = 0; i < relaxed_intsets.size(); ++i) {
- relaxed_region.push_back(
- relaxed_intsets[i].CoverRange(Range::FromMinExtent(0, buffer_region->buffer->shape[i])));
+ relaxed_region.push_back(relaxed_intsets[i].CoverRange(Range::FromMinExtent(
+ 0, buffer_region->source.as_or_throw<tvm::tirx::BufferVar>()->shape[i])));
}
- return BufferRegion(buffer_region->buffer, relaxed_region);
+ return BufferRegion(buffer_region->source.as_or_throw<tvm::tirx::BufferVar>(), relaxed_region);
}
class RollingBufferDependencyError : public ScheduleErrorContextObj {
@@ -115,7 +115,7 @@ class RollingBufferDependencyError : public ScheduleErrorContextObj {
class RollingBufferMatchError : public ScheduleErrorContextObj {
public:
- RollingBufferMatchError(IRModule mod, SBlock block, BufferRegion buffer_region)
+ RollingBufferMatchError(IRModule mod, SBlock block, TensorRegion buffer_region)
: mod_(mod), block_(block), buffer_region_(buffer_region) {}
ffi::String FastErrorString() const final {
return "ScheduleError: rolling_buffer expect the buffer region to have at least one dimention"
@@ -123,8 +123,8 @@ class RollingBufferMatchError : public ScheduleErrorContextObj {
}
ffi::String DetailRenderTemplate() const final {
std::ostringstream os;
- os << "The target buffer " << buffer_region_->buffer.name() << " with region "
- << buffer_region_->region
+ os << "The target buffer " << buffer_region_->source.as_or_throw<tvm::tirx::BufferVar>().name()
+ << " with region " << buffer_region_->region
<< " should have at least one dimension range that matches a rolling pattern "
"such as hh.outer * stride + hh.inner. ";
return os.str();
@@ -136,7 +136,7 @@ class RollingBufferMatchError : public ScheduleErrorContextObj {
private:
IRModule mod_;
SBlock block_;
- BufferRegion buffer_region_;
+ TensorRegion buffer_region_;
};
class RollingBufferInsertionError : public ScheduleErrorContextObj {
@@ -168,7 +168,7 @@ class RollingBufferInfoCollector {
public:
static RollingBufferInfo CheckAndGetRollingBufferInfo(const IRModule& mod,
const StmtSRef& block_sref,
- const BufferRegion& buffer_region) {
+ const TensorRegion& buffer_region) {
RollingBufferInfoCollector collector;
if (!collector.MatchRollingBuffer(block_sref, buffer_region)) {
const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref);
@@ -179,8 +179,8 @@ class RollingBufferInfoCollector {
}
private:
- bool MatchRollingBuffer(const StmtSRef& block_sref, const BufferRegion& buffer_region) {
- const BufferVar& buffer = buffer_region->buffer;
+ bool MatchRollingBuffer(const StmtSRef& block_sref, const TensorRegion& buffer_region) {
+ const BufferVar& buffer = buffer_region->source.as_or_throw<tvm::tirx::BufferVar>();
const Region& region = buffer_region->region;
std::vector<ffi::Optional<Var>> bound_iter_vars;
@@ -286,10 +286,10 @@ class RollingBufferRewriter : public StmtExprMutator {
: scope_sref_(scope_sref), info_(info) {}
private:
- void RewriteAccessRegion(ffi::Array<BufferRegion>* old_access_regions,
- const ffi::Array<BufferRegion>& infered_access_regions) {
- auto fmutate = [this, &infered_access_regions](const BufferRegion& buffer_region) {
- if (buffer_region->buffer.same_as(info_->old_buffer)) {
+ void RewriteAccessRegion(ffi::Array<TensorRegion>* old_access_regions,
+ const ffi::Array<TensorRegion>& infered_access_regions) {
+ auto fmutate = [this, &infered_access_regions](const TensorRegion& buffer_region) {
+ if (buffer_region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(info_->old_buffer)) {
TVM_FFI_ICHECK(infered_access_regions.size() == 1);
return infered_access_regions[0];
}
@@ -439,7 +439,7 @@ void RollingBuffer(ScheduleState self, const StmtSRef& block_sref, int write_buf
const SBlock& block = realize->block;
// Step 1. Checking index, getting the target buffer region and the parent scope.
- const BufferRegion& buffer_region =
+ const TensorRegion& buffer_region =
GetNthAccessBufferRegion(self, block, write_buffer_index, BufferIndexType::kWrite);
StmtSRef scope_root_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false);
// Step 2. Check if the target block is not an output block and has only RAW dependencies.
@@ -452,7 +452,8 @@ void RollingBuffer(ScheduleState self, const StmtSRef& block_sref, int write_buf
consumers_sref.push_back(block_sref);
StmtSRef lca = GetSRefLowestCommonAncestor(consumers_sref);
if (!lca->StmtAs<ForNode>()) {
- throw MakeScheduleError<RollingBufferInsertionError>(self->mod, buffer_region->buffer, block);
+ throw MakeScheduleError<RollingBufferInsertionError>(
+ self->mod, buffer_region->source.as_or_throw<tvm::tirx::BufferVar>(), block);
}
for (auto it = loop_srefs.rbegin(); it != loop_srefs.rend(); ++it) {
@@ -465,7 +466,7 @@ void RollingBuffer(ScheduleState self, const StmtSRef& block_sref, int write_buf
Range range = Range::FromMinExtent(cur_loop->min, cur_loop->extent);
dom_map.Set(cur_loop->loop_var, arith::IntSet::FromRange(range));
}
- BufferRegion relaxed_region = GetRelaxedBufferRegion(realize, buffer_region, dom_map);
+ TensorRegion relaxed_region = GetRelaxedBufferRegion(realize, buffer_region, dom_map);
// Step 4. Find a valid rolling axis and collect bound overlaps on the target buffer.
RollingBufferInfo info = RollingBufferInfoCollector::CheckAndGetRollingBufferInfo(
diff --git a/src/s_tir/schedule/state.cc b/src/s_tir/schedule/state.cc
index 6cd6a5ccca..ab5e6b370b 100644
--- a/src/s_tir/schedule/state.cc
+++ b/src/s_tir/schedule/state.cc
@@ -42,7 +42,7 @@ using SMap = std::unordered_map<K, V, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>;
* \param dom_high_exclusive The highest node in the sref tree path
* \return An n-dimensional integer set
*/
-ffi::Array<arith::IntSet> AnalyzeRegionUpperBound(const BufferRegion& region, //
+ffi::Array<arith::IntSet> AnalyzeRegionUpperBound(const TensorRegion& region, //
const PrimExpr& predicate, //
const StmtSRef& dom_low_inclusive, //
const StmtSRef& dom_high_exclusive, //
@@ -50,7 +50,8 @@ ffi::Array<arith::IntSet> AnalyzeRegionUpperBound(const BufferRegion& region,
ffi::Map<Var, Range> var_dom = LoopDomainOfSRefTreePath(
/*low_inclusive=*/dom_low_inclusive,
/*high_exclusive=*/dom_high_exclusive,
- /*extra_relax_scope=*/runtime::StorageScope::Create(region->buffer.scope()));
+ /*extra_relax_scope=*/
+ runtime::StorageScope::Create(region->source.as_or_throw<tvm::tirx::BufferVar>().scope()));
arith::Analyzer analyzer_ref = ffi::GetRef<arith::Analyzer>(analyzer);
return EstimateRegionUpperBound(
/*region=*/region->region,
@@ -68,7 +69,7 @@ ffi::Array<arith::IntSet> AnalyzeRegionUpperBound(const BufferRegion& region,
* \param analyzer The analyzer
* \return An n-dimensional integer set
*/
-ffi::Array<arith::IntSet> AnalyzeRegionLowerBound(const BufferRegion& region, //
+ffi::Array<arith::IntSet> AnalyzeRegionLowerBound(const TensorRegion& region, //
const PrimExpr& predicate, //
const StmtSRef& dom_low_inclusive, //
const StmtSRef& dom_high_exclusive, //
@@ -76,7 +77,8 @@ ffi::Array<arith::IntSet> AnalyzeRegionLowerBound(const BufferRegion& region,
ffi::Map<Var, Range> var_dom = LoopDomainOfSRefTreePath(
/*low_inclusive=*/dom_low_inclusive,
/*high_exclusive=*/dom_high_exclusive,
- /*extra_relax_scope=*/runtime::StorageScope::Create(region->buffer.scope()));
+ /*extra_relax_scope=*/
+ runtime::StorageScope::Create(region->source.as_or_throw<tvm::tirx::BufferVar>().scope()));
arith::Analyzer analyzer_ref = ffi::GetRef<arith::Analyzer>(analyzer);
if (ffi::Optional<ffi::Array<arith::IntSet>> result = EstimateRegionLowerBound(
/*region=*/region->region,
@@ -84,7 +86,8 @@ ffi::Array<arith::IntSet> AnalyzeRegionLowerBound(const BufferRegion& region,
/*predicate=*/predicate, /*analyzer=*/analyzer_ref)) {
return result.value();
}
- return ffi::Array<arith::IntSet>(region->buffer->shape.size(), arith::IntSet::Nothing());
+ return ffi::Array<arith::IntSet>(region->source.as_or_throw<tvm::tirx::BufferVar>()->shape.size(),
+ arith::IntSet::Nothing());
}
/*!
@@ -214,8 +217,8 @@ class SBlockInfoCollector : public StmtExprVisitor {
const StmtSRefNode* limit = scope_root->parent;
bool stage_pipeline = true;
// Step 1. Unbind the read/write regions of each child block
- std::unordered_map<const StmtSRefNode*, ffi::Array<BufferRegion>> block_reads_unbound;
- std::unordered_map<const StmtSRefNode*, ffi::Array<BufferRegion>> block_writes_unbound;
+ std::unordered_map<const StmtSRefNode*, ffi::Array<TensorRegion>> block_reads_unbound;
+ std::unordered_map<const StmtSRefNode*, ffi::Array<TensorRegion>> block_writes_unbound;
block_reads_unbound.reserve(child_block_srefs.size());
block_writes_unbound.reserve(child_block_srefs.size());
for (const StmtSRef& block_sref : child_block_srefs) {
@@ -226,9 +229,9 @@ class SBlockInfoCollector : public StmtExprVisitor {
return ffi::Unchanged();
};
// Step 1.1. Unbind read regions
- ffi::Array<BufferRegion> reads;
+ ffi::Array<TensorRegion> reads;
reads.reserve(block->reads.size());
- for (const BufferRegion& region : block->reads) {
+ for (const TensorRegion& region : block->reads) {
ffi::Array<Range> mapped_region = region->region.Map([&f_substitute](const Range& range) {
PrimExpr min = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(range->min, f_substitute)
.as_or_throw<PrimExpr>();
@@ -237,13 +240,14 @@ class SBlockInfoCollector : public StmtExprVisitor {
.as_or_throw<PrimExpr>();
return Range::FromMinExtent(min, extent);
});
- reads.push_back(BufferRegion(region->buffer, mapped_region));
+ reads.push_back(
+ BufferRegion(region->source.as_or_throw<tvm::tirx::BufferVar>(), mapped_region));
}
block_reads_unbound.emplace(block_sref.get(), std::move(reads));
// Step 1.2. Unbind write regions
- ffi::Array<BufferRegion> writes;
+ ffi::Array<TensorRegion> writes;
writes.reserve(block->writes.size());
- for (const BufferRegion& region : block->writes) {
+ for (const TensorRegion& region : block->writes) {
ffi::Array<Range> mapped_region = region->region.Map([&f_substitute](const Range& range) {
PrimExpr min = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(range->min, f_substitute)
.as_or_throw<PrimExpr>();
@@ -252,7 +256,8 @@ class SBlockInfoCollector : public StmtExprVisitor {
.as_or_throw<PrimExpr>();
return Range::FromMinExtent(min, extent);
});
- writes.push_back(BufferRegion(region->buffer, mapped_region));
+ writes.push_back(
+ BufferRegion(region->source.as_or_throw<tvm::tirx::BufferVar>(), mapped_region));
}
block_writes_unbound.emplace(block_sref.get(), std::move(writes));
}
@@ -304,16 +309,16 @@ class SBlockInfoCollector : public StmtExprVisitor {
ffi::ObjectPtrEqual>
touched_regions;
// Step 2.3.1. Find all the regions read by the consumer that we care about
- for (const BufferRegion& region : block_reads_unbound.at(consumer_block_sref.get())) {
- BufferVar buffer = region->buffer;
+ for (const TensorRegion& region : block_reads_unbound.at(consumer_block_sref.get())) {
+ BufferVar buffer = region->source.as_or_throw<tvm::tirx::BufferVar>();
touched_regions[buffer] = {};
}
// Step 2.3.2. Find all the regions written by each producer
for (const StmtSRefNode* producer_block_sref : producer_block_srefs) {
const SBlockRealize& producer_realize = block2realize_.at(producer_block_sref->stmt);
StmtSRef parent_sref = ffi::GetRef<StmtSRef>(producer_block_sref->parent);
- for (const BufferRegion& region : block_writes_unbound.at(producer_block_sref)) {
- BufferVar buffer = region->buffer;
+ for (const TensorRegion& region : block_writes_unbound.at(producer_block_sref)) {
+ BufferVar buffer = region->source.as_or_throw<tvm::tirx::BufferVar>();
auto it = touched_regions.find(buffer);
// Skip the regions that is not read by the consumer
if (it != touched_regions.end()) {
@@ -334,8 +339,8 @@ class SBlockInfoCollector : public StmtExprVisitor {
// Step 2.3.3. For each buffer, check the region cover property
{
StmtSRef parent_sref = ffi::GetRef<StmtSRef>(consumer_block_sref->parent);
- for (const BufferRegion& region : block_reads_unbound.at(consumer_block_sref.get())) {
- BufferVar buffer = region->buffer;
+ for (const TensorRegion& region : block_reads_unbound.at(consumer_block_sref.get())) {
+ BufferVar buffer = region->source.as_or_throw<tvm::tirx::BufferVar>();
const std::vector<ffi::Array<arith::IntSet>>& touched_region =
touched_regions.at(buffer);
if (!touched_region.empty()) {
diff --git a/src/s_tir/schedule/transform.cc b/src/s_tir/schedule/transform.cc
index ee4ebabe38..bd39c56fb2 100644
--- a/src/s_tir/schedule/transform.cc
+++ b/src/s_tir/schedule/transform.cc
@@ -56,26 +56,26 @@ BufferVar WithDType(const BufferVar& buffer, PrimType dtype) {
return BufferVar(buffer.name(), new_type, buffer.span());
}
-ffi::Array<BufferRegion> ReplaceBuffer(ffi::Array<BufferRegion> regions, const BufferVar& source,
+ffi::Array<TensorRegion> ReplaceBuffer(ffi::Array<TensorRegion> regions, const BufferVar& source,
const BufferVar& target) {
- regions.MutateByApply([&source, &target](BufferRegion region) -> BufferRegion {
- if (region->buffer.same_as(source)) {
- ffi::ObjectPtr<BufferRegionNode> n = ffi::make_object<BufferRegionNode>(*region.get());
- n->buffer = target;
- return BufferRegion(n);
+ regions.MutateByApply([&source, &target](TensorRegion region) -> TensorRegion {
+ if (region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(source)) {
+ ffi::ObjectPtr<TensorRegionNode> n = ffi::make_object<TensorRegionNode>(*region.get());
+ n->source = target;
+ return TensorRegion(n);
}
return region;
});
return regions;
}
-ffi::Array<BufferRegion> ReplaceBuffer(ffi::Array<BufferRegion> regions,
+ffi::Array<TensorRegion> ReplaceBuffer(ffi::Array<TensorRegion> regions,
const ffi::Map<BufferVar, BufferVar>& buffer_map) {
- regions.MutateByApply([&buffer_map](BufferRegion region) -> BufferRegion {
- if (buffer_map.count(region->buffer)) {
- ffi::ObjectPtr<BufferRegionNode> n = ffi::make_object<BufferRegionNode>(*region.get());
- n->buffer = buffer_map[region->buffer];
- return BufferRegion(n);
+ regions.MutateByApply([&buffer_map](TensorRegion region) -> TensorRegion {
+ if (buffer_map.count(region->source.as_or_throw<tvm::tirx::BufferVar>())) {
+ ffi::ObjectPtr<TensorRegionNode> n = ffi::make_object<TensorRegionNode>(*region.get());
+ n->source = buffer_map[region->source.as_or_throw<tvm::tirx::BufferVar>()];
+ return TensorRegion(n);
}
return region;
});
@@ -86,7 +86,7 @@ ffi::Array<MatchBufferRegion> ReplaceBuffer(ffi::Array<MatchBufferRegion> match_
const BufferVar& source, const BufferVar& target) {
match_buffers.MutateByApply(
[&source, &target](MatchBufferRegion match_buffer) -> MatchBufferRegion {
- if (match_buffer->source->buffer.same_as(source)) {
+ if (match_buffer->source->source.as_or_throw<tvm::tirx::BufferVar>().same_as(source)) {
ffi::ObjectPtr<MatchBufferRegionNode> n =
ffi::make_object<MatchBufferRegionNode>(*match_buffer.get());
n->source = BufferRegion(target, n->source->region);
@@ -97,11 +97,11 @@ ffi::Array<MatchBufferRegion> ReplaceBuffer(ffi::Array<MatchBufferRegion> match_
return match_buffers;
}
-ffi::Array<BufferRegion> ReplaceBufferRegion(ffi::Array<BufferRegion> regions,
+ffi::Array<TensorRegion> ReplaceBufferRegion(ffi::Array<TensorRegion> regions,
const BufferVar& source_buffer,
- const BufferRegion& target) {
- regions.MutateByApply([&source_buffer, &target](const BufferRegion& region) -> BufferRegion {
- if (region->buffer.same_as(source_buffer)) {
+ const TensorRegion& target) {
+ regions.MutateByApply([&source_buffer, &target](const TensorRegion& region) -> TensorRegion {
+ if (region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(source_buffer)) {
return target;
}
return region;
@@ -111,17 +111,17 @@ ffi::Array<BufferRegion> ReplaceBufferRegion(ffi::Array<BufferRegion> regions,
ffi::Array<MatchBufferRegion> ReplaceBufferRegion(ffi::Array<MatchBufferRegion> match_buffers,
const BufferVar& source_buffer,
- const BufferRegion& target) {
- match_buffers.MutateByApply(
- [&source_buffer, &target](const MatchBufferRegion& match_buffer) -> MatchBufferRegion {
- if (match_buffer->source->buffer.same_as(source_buffer)) {
- ffi::ObjectPtr<MatchBufferRegionNode> n =
- ffi::make_object<MatchBufferRegionNode>(*match_buffer.get());
- n->source = target;
- return MatchBufferRegion(n);
- }
- return match_buffer;
- });
+ const TensorRegion& target) {
+ match_buffers.MutateByApply([&source_buffer, &target](
+ const MatchBufferRegion& match_buffer) -> MatchBufferRegion {
+ if (match_buffer->source->source.as_or_throw<tvm::tirx::BufferVar>().same_as(source_buffer)) {
+ ffi::ObjectPtr<MatchBufferRegionNode> n =
+ ffi::make_object<MatchBufferRegionNode>(*match_buffer.get());
+ n->source = target;
+ return MatchBufferRegion(n);
+ }
+ return match_buffer;
+ });
return match_buffers;
}
@@ -160,7 +160,9 @@ UnchangedOr<Expr> ReplaceBufferMutator::Mutate_(const CallNode* op, InplaceMode
MatchBufferRegion ReplaceBufferMutator::VisitMatchBufferRegion(
const MatchBufferRegion& match_buffer) {
- if (auto replacement = VarRemapGet(match_buffer->source->buffer).as<BufferVar>()) {
+ if (auto replacement =
+ VarRemapGet(match_buffer->source->source.as_or_throw<tvm::tirx::BufferVar>())
+ .as<BufferVar>()) {
return MatchBufferRegion(match_buffer->buffer,
BufferRegion(replacement.value(), match_buffer->source->region));
} else {
@@ -176,7 +178,7 @@ UnchangedOr<Stmt> ReplaceBufferMutator::Mutate_(const SBlockNode* block, Inplace
auto f_mutate_match_buffer = [this](const MatchBufferRegion& match_buffer) {
return this->VisitMatchBufferRegion(match_buffer);
};
- auto f_mutate_read_write_region = [this](const BufferRegion& buffer_region) {
+ auto f_mutate_read_write_region = [this](const TensorRegion& buffer_region) {
auto region = MutateArray(buffer_region->region, [this](const Range& range) {
auto min_result = Mutate(range->min, InplaceMode::kDisallow);
bool min_unchanged = min_result.UnchangedOrSameAs(range->min);
@@ -191,10 +193,12 @@ UnchangedOr<Stmt> ReplaceBufferMutator::Mutate_(const SBlockNode* block, Inplace
}
});
- BufferVar buf =
- VarRemapGet(buffer_region->buffer).as<BufferVar>().value_or(buffer_region->buffer);
+ BufferVar buf = VarRemapGet(buffer_region->source.as_or_throw<tvm::tirx::BufferVar>())
+ .as<BufferVar>()
+ .value_or(buffer_region->source.as_or_throw<tvm::tirx::BufferVar>());
- if (buf.same_as(buffer_region->buffer) && region.same_as(buffer_region->region)) {
+ if (buf.same_as(buffer_region->source.as_or_throw<tvm::tirx::BufferVar>()) &&
+ region.same_as(buffer_region->region)) {
return buffer_region;
} else {
return BufferRegion(buf, region);
@@ -207,8 +211,8 @@ UnchangedOr<Stmt> ReplaceBufferMutator::Mutate_(const SBlockNode* block, Inplace
// Step 1. Mutate `match_buffers`. If an old buffer appears as a source of MatchBufferRegion,
ffi::Array<MatchBufferRegion> match_buffers = block->match_buffers.Map(f_mutate_match_buffer);
// Step 2. Mutate the read/write region.
- ffi::Array<BufferRegion> reads = block->reads.Map(f_mutate_read_write_region);
- ffi::Array<BufferRegion> writes = block->writes.Map(f_mutate_read_write_region);
+ ffi::Array<TensorRegion> reads = block->reads.Map(f_mutate_read_write_region);
+ ffi::Array<TensorRegion> writes = block->writes.Map(f_mutate_read_write_region);
// Step 3. Mutate `alloc_buffers` for the old buffer allocated in this block.
ffi::Array<BufferVar> alloc_buffers = block->alloc_buffers.Map(f_mutate_alloc_buffers);
// Step 4. Recursively mutate the block.
@@ -438,8 +442,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
/******** BlockBufferAccessSimplifier ********/
void BlockBufferAccessSimplifier::SimplifyAccessRegion(
- ffi::Array<BufferRegion>* old_access_regions) {
- auto fmutate = [this](const BufferRegion& buffer_region) {
+ ffi::Array<TensorRegion>* old_access_regions) {
+ auto fmutate = [this](const TensorRegion& buffer_region) {
ffi::Array<Range> new_buffer_region;
ffi::Array<PrimExpr> simplified_min;
for (const auto& range : buffer_region->region) {
@@ -452,7 +456,8 @@ void BlockBufferAccessSimplifier::SimplifyAccessRegion(
PrimExpr extent = analyzer_->Simplify(buffer_region->region[i]->extent);
new_buffer_region.push_back(Range::FromMinExtent(min, extent));
}
- return BufferRegion(buffer_region->buffer, new_buffer_region);
+ return BufferRegion(buffer_region->source.as_or_throw<tvm::tirx::BufferVar>(),
+ new_buffer_region);
};
(*old_access_regions).MutateByApply(fmutate);
}
diff --git a/src/s_tir/schedule/transform.h b/src/s_tir/schedule/transform.h
index 31e90c6845..72c76daaa9 100644
--- a/src/s_tir/schedule/transform.h
+++ b/src/s_tir/schedule/transform.h
@@ -71,7 +71,7 @@ BufferVar WithDType(const BufferVar& buffer, PrimType dtype);
* \param target The buffer to be replaced to
* \return The new sequence of regions after replacement
*/
-ffi::Array<BufferRegion> ReplaceBuffer(ffi::Array<BufferRegion> regions, const BufferVar& source,
+ffi::Array<TensorRegion> ReplaceBuffer(ffi::Array<TensorRegion> regions, const BufferVar& source,
const BufferVar& target);
/*!
@@ -80,7 +80,7 @@ ffi::Array<BufferRegion> ReplaceBuffer(ffi::Array<BufferRegion> regions, const B
* \param buffer_map The mapping from old buffers to new buffers
* \return The new sequence of regions after replacement
*/
-ffi::Array<BufferRegion> ReplaceBuffer(ffi::Array<BufferRegion> regions,
+ffi::Array<TensorRegion> ReplaceBuffer(ffi::Array<TensorRegion> regions,
const ffi::Map<BufferVar, BufferVar>& buffer_map);
/*!
@@ -100,9 +100,9 @@ ffi::Array<MatchBufferRegion> ReplaceBuffer(ffi::Array<MatchBufferRegion> match_
* \param target The buffer region to be replaced to
* \return The new sequence of regions after replacement
*/
-ffi::Array<BufferRegion> ReplaceBufferRegion(ffi::Array<BufferRegion> regions,
+ffi::Array<TensorRegion> ReplaceBufferRegion(ffi::Array<TensorRegion> regions,
const BufferVar& source_buffer,
- const BufferRegion& target);
+ const TensorRegion& target);
/*!
* \brief Replaces the buffer region within the specific sequence of match_buffers
@@ -113,7 +113,7 @@ ffi::Array<BufferRegion> ReplaceBufferRegion(ffi::Array<BufferRegion> regions,
*/
ffi::Array<MatchBufferRegion> ReplaceBufferRegion(ffi::Array<MatchBufferRegion> match_buffers,
const BufferVar& source_buffer,
- const BufferRegion& target);
+ const TensorRegion& target);
/*!
* \brief A helper mutator which recursively replaces the old buffer with the new buffer and
@@ -231,7 +231,7 @@ class BlockBufferAccessSimplifier : public tirx::IRMutatorWithAnalyzer {
: IRMutatorWithAnalyzer(analyzer) {}
private:
- void SimplifyAccessRegion(ffi::Array<BufferRegion>* old_access_regions);
+ void SimplifyAccessRegion(ffi::Array<TensorRegion>* old_access_regions);
void SimplifyBufferIndices(ffi::Array<PrimExpr>* indices);
UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final;
diff --git a/src/s_tir/transform/compact_buffer_region.cc b/src/s_tir/transform/compact_buffer_region.cc
index 007bd21b35..bd5e5042e4 100644
--- a/src/s_tir/transform/compact_buffer_region.cc
+++ b/src/s_tir/transform/compact_buffer_region.cc
@@ -93,7 +93,8 @@ class Var2BufferCollector : public StmtExprVisitor {
}
for (const MatchBufferRegion& region : op->match_buffers) {
var2buffer_[region->buffer.var()].insert(region->buffer);
- var2buffer_[region->source->buffer.var()].insert(region->source->buffer);
+ var2buffer_[region->source->source.as_or_throw<tvm::tirx::BufferVar>().var()].insert(
+ region->source->source.as_or_throw<tvm::tirx::BufferVar>());
}
return StmtExprVisitor::Visit_(op);
}
@@ -149,7 +150,8 @@ class BufferAccessRegionCollector : public StmtExprVisitor {
/**************** Visitor overload ****************/
// Declared regions carry bounds, not opaque runtime accesses.
- ffi::Optional<VisitInterrupt> Visit_(const BufferRegionNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorRegionNode* op) final {
+ if (!op->source.as<BufferVar>()) return StmtExprVisitor::Visit_(op);
for (const Range& range : op->region) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(range->min));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(range->extent));
@@ -158,7 +160,7 @@ class BufferAccessRegionCollector : public StmtExprVisitor {
}
ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
- VisitBufferAccess(BufferRegion::FromPoint(op->buffer, op->indices));
+ VisitBufferAccess(BufferRegionFromPoint(op->buffer, op->indices));
return Visit(op->value);
}
@@ -168,7 +170,7 @@ class BufferAccessRegionCollector : public StmtExprVisitor {
if (explicit_it != explicit_access_annotations_.end()) {
VisitBufferAccess(explicit_it->second);
} else {
- VisitBufferAccess(BufferRegion::FromPoint(buffer, op->indices));
+ VisitBufferAccess(BufferRegionFromPoint(buffer, op->indices));
}
for (const auto& index : op->indices) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(index));
@@ -267,14 +269,14 @@ class BufferAccessRegionCollector : public StmtExprVisitor {
TVM_FFI_ICHECK(!op->init.has_value());
TVM_FFI_ICHECK_EQ(op->iter_vars.size(), 0) << "CompactBufferRegion only works on opaque blocks";
// Step 1. Record and update current read/write region annotations
- std::unordered_map<BufferVar, std::vector<BufferRegion>, ffi::ObjectPtrHash,
+ std::unordered_map<BufferVar, std::vector<TensorRegion>, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>
cur_access_annotations;
- for (const BufferRegion& region : op->reads) {
- cur_access_annotations[region->buffer].push_back(region);
+ for (const TensorRegion& region : op->reads) {
+ cur_access_annotations[region->source.as_or_throw<tvm::tirx::BufferVar>()].push_back(region);
}
- for (const BufferRegion& region : op->writes) {
- cur_access_annotations[region->buffer].push_back(region);
+ for (const TensorRegion& region : op->writes) {
+ cur_access_annotations[region->source.as_or_throw<tvm::tirx::BufferVar>()].push_back(region);
}
for (auto& p : cur_access_annotations) {
auto& regions = access_annotations_[p.first];
@@ -289,10 +291,12 @@ class BufferAccessRegionCollector : public StmtExprVisitor {
for (int64_t index : buffer_indices) {
int buffer_index = static_cast<int>(index);
if (buffer_index >= 0 && buffer_index < static_cast<int>(op->reads.size())) {
- const BufferRegion& explicit_region = index_type == BufferIndexType::kRead
+ const TensorRegion& explicit_region = index_type == BufferIndexType::kRead
? op->reads[buffer_index]
: op->writes[buffer_index];
- explicit_access_annotations_[explicit_region->buffer] = explicit_region;
+ explicit_access_annotations_[explicit_region->source
+ .as_or_throw<tvm::tirx::BufferVar>()] =
+ explicit_region;
}
}
}
@@ -372,8 +376,8 @@ class BufferAccessRegionCollector : public StmtExprVisitor {
buffer_scope_depth_.insert(it, {buffer_data, ancestor_iters_.size()});
}
- void VisitBufferAccess(const BufferRegion& buffer_region) {
- const BufferVar& buffer = buffer_region->buffer;
+ void VisitBufferAccess(const TensorRegion& buffer_region) {
+ const BufferVar& buffer = buffer_region->source.as_or_throw<tvm::tirx::BufferVar>();
auto it = buffer_scope_depth_.find(buffer.var());
if (it != buffer_scope_depth_.end()) {
size_t n_ancestor_loops = it->second;
@@ -430,11 +434,11 @@ class BufferAccessRegionCollector : public StmtExprVisitor {
auto annotation_it = access_annotations_.find(buffer);
if (annotation_it != access_annotations_.end()) {
// opaque buffer has explicit accessed region annotations
- for (const BufferRegion& region : annotation_it->second) {
+ for (const TensorRegion& region : annotation_it->second) {
VisitBufferAccess(region);
}
} else {
- VisitBufferAccess(BufferRegion::FullRegion(buffer));
+ VisitBufferAccess(FullBufferRegion(buffer));
}
}
}
@@ -568,10 +572,10 @@ class BufferAccessRegionCollector : public StmtExprVisitor {
buffer_access_region_;
/*! \brief The map from BufferVar to it's access regions annotated by current block. */
- std::unordered_map<BufferVar, std::vector<BufferRegion>, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
+ std::unordered_map<BufferVar, std::vector<TensorRegion>, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
access_annotations_;
/*! \brief The map from BufferVar to its explicit access region annotated by the block. */
- std::unordered_map<BufferVar, BufferRegion, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
+ std::unordered_map<BufferVar, TensorRegion, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
explicit_access_annotations_;
};
@@ -712,13 +716,15 @@ class BufferCompactor : public StmtExprMutator {
*region = std::move(new_region);
}
- void RewriteBufferRegions(ffi::Array<BufferRegion>* regions) const {
- ffi::Array<BufferRegion> new_regions;
+ void RewriteBufferRegions(ffi::Array<TensorRegion>* regions) const {
+ ffi::Array<TensorRegion> new_regions;
new_regions.reserve(regions->size());
for (const auto& region : *regions) {
- BufferRegion buffer_region = region;
- BufferRegionNode* p = buffer_region.CopyOnWrite();
- RewriteBufferRegion(&p->buffer, &p->region);
+ TensorRegion buffer_region = region;
+ TensorRegionNode* p = buffer_region.CopyOnWrite();
+ BufferVar source = p->source.as_or_throw<tvm::tirx::BufferVar>();
+ RewriteBufferRegion(&source, &p->region);
+ p->source = source;
new_regions.push_back(buffer_region);
}
*regions = std::move(new_regions);
@@ -728,10 +734,12 @@ class BufferCompactor : public StmtExprMutator {
ffi::Array<MatchBufferRegion> result;
result.reserve(match_buffers->size());
for (const auto& match_buffer : *match_buffers) {
- const BufferRegion& buffer_region = match_buffer->source;
- auto p = ffi::make_object<BufferRegionNode>(*buffer_region.get());
- RewriteBufferRegion(&p->buffer, &p->region);
- result.push_back(MatchBufferRegion(match_buffer->buffer, BufferRegion(p)));
+ const TensorRegion& buffer_region = match_buffer->source;
+ auto p = ffi::make_object<TensorRegionNode>(*buffer_region.get());
+ BufferVar source = p->source.as_or_throw<tvm::tirx::BufferVar>();
+ RewriteBufferRegion(&source, &p->region);
+ p->source = source;
+ result.push_back(MatchBufferRegion(match_buffer->buffer, TensorRegion(p)));
}
*match_buffers = std::move(result);
}
diff --git a/src/s_tir/transform/default_gpu_schedule.cc b/src/s_tir/transform/default_gpu_schedule.cc
index c67014a9a5..f9623e2aed 100644
--- a/src/s_tir/transform/default_gpu_schedule.cc
+++ b/src/s_tir/transform/default_gpu_schedule.cc
@@ -145,8 +145,8 @@ tirx::PrimFunc WrapBareSBlockBody(const tirx::PrimFunc& func) {
tirx::Stmt for_stmt =
tirx::For(loop_var.as_or_throw<PrimVar>(), zero, one, tirx::ForKind::kSerial, inner_realize);
tirx::SBlock root_block(/*iter_vars=*/ffi::Array<tirx::IterVar>{},
- /*reads=*/ffi::Array<tirx::BufferRegion>{},
- /*writes=*/ffi::Array<tirx::BufferRegion>{},
+ /*reads=*/ffi::Array<tvm::TensorRegion>{},
+ /*writes=*/ffi::Array<tvm::TensorRegion>{},
/*name_hint=*/"root", /*body=*/for_stmt);
tirx::SBlockRealize root_realize(/*iter_values=*/ffi::Array<tvm::PrimExpr>{},
/*predicate=*/IntImm::Bool(true), root_block);
diff --git a/src/s_tir/transform/inject_double_buffer.cc b/src/s_tir/transform/inject_double_buffer.cc
index 666c4db7c0..2100a82780 100644
--- a/src/s_tir/transform/inject_double_buffer.cc
+++ b/src/s_tir/transform/inject_double_buffer.cc
@@ -109,7 +109,8 @@ class DoubleBufferDetector : public StmtExprVisitor {
}
// Declared regions carry bounds, not opaque runtime accesses.
- ffi::Optional<VisitInterrupt> Visit_(const BufferRegionNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorRegionNode* op) final {
+ if (!op->source.as<BufferVar>()) return StmtExprVisitor::Visit_(op);
for (const Range& range : op->region) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(range->min));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(range->extent));
@@ -341,10 +342,13 @@ class DoubleBufferInjector : public StmtExprMutator {
return buf;
}
- UnchangedOr<Expr> Mutate_(const BufferRegionNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Expr> Mutate_(const TensorRegionNode* op, InplaceMode inplace_mode) final {
+ if (!op->source.as<BufferVar>()) {
+ return StmtExprMutator::Mutate_(op, inplace_mode);
+ }
auto region = Mutate(op->region).as_or_throw<UnchangedOr<ffi::Array<Range>>>();
if (region.UnchangedOrSameAs(op->region)) return ffi::Unchanged();
- BufferRegion node = ffi::GetRef<BufferRegion>(op);
+ TensorRegion node = ffi::GetRef<TensorRegion>(op);
node.CopyOnWrite()->region = std::move(region).ValueUnchecked();
return node;
}
diff --git a/src/s_tir/transform/inject_software_pipeline.cc b/src/s_tir/transform/inject_software_pipeline.cc
index ba86936954..9f6eb256f9 100644
--- a/src/s_tir/transform/inject_software_pipeline.cc
+++ b/src/s_tir/transform/inject_software_pipeline.cc
@@ -80,7 +80,7 @@ SBlock MakeSBlock(const Stmt& body, const ffi::Map<Var, BufferVar>& buffer_data_
}
}
SBlock block(/*iter_vars=*/{}, /*reads=*/{}, /*writes=*/{}, /*name_hint=*/"", /*body*/ body);
- ffi::Array<ffi::Array<BufferRegion>> access =
+ ffi::Array<ffi::Array<TensorRegion>> access =
GetSBlockReadWriteRegion(block, buffer_data_to_buffer);
SBlockNode* n = block.CopyOnWrite();
n->reads = access[0];
@@ -265,8 +265,9 @@ class PipelineBodyRewriter : public StmtExprMutator {
}
private:
- BufferRegion RewritePipelineBufferRegion(const BufferRegion& buffer_region) {
- if (auto replacement = VarRemapGet(buffer_region->buffer).as<BufferVar>()) {
+ TensorRegion RewritePipelineBufferRegion(const TensorRegion& buffer_region) {
+ if (auto replacement = VarRemapGet(buffer_region->source.as_or_throw<tvm::tirx::BufferVar>())
+ .as<BufferVar>()) {
Region new_region = buffer_region->region;
BufferVar new_buffer = replacement.value();
// For pipeline buffers, relax the access region of the first dimension to full extent
@@ -296,8 +297,11 @@ class PipelineBodyRewriter : public StmtExprMutator {
return block;
}
- UnchangedOr<Expr> Mutate_(const BufferRegionNode* op, InplaceMode inplace_mode) final {
- BufferRegion region = RewritePipelineBufferRegion(ffi::GetRef<BufferRegion>(op));
+ UnchangedOr<Expr> Mutate_(const TensorRegionNode* op, InplaceMode inplace_mode) final {
+ if (!op->source.as<BufferVar>()) {
+ return StmtExprMutator::Mutate_(op, inplace_mode);
+ }
+ TensorRegion region = RewritePipelineBufferRegion(ffi::GetRef<TensorRegion>(op));
return StmtExprMutator::Mutate_(region.get(), InplaceMode::kDisallow).ValueOrUnchanged(region);
}
@@ -452,11 +456,11 @@ class PipelineRewriter : public StmtExprMutator {
int stage = pair.second.stage;
max_stage_ = std::max(max_stage_, stage);
- for (const BufferRegion& write : block->writes) {
- if (!infos.count(write->buffer)) {
- infos.emplace(write->buffer, BufferAccessInfo{});
+ for (const TensorRegion& write : block->writes) {
+ if (!infos.count(write->source.as_or_throw<tvm::tirx::BufferVar>())) {
+ infos.emplace(write->source.as_or_throw<tvm::tirx::BufferVar>(), BufferAccessInfo{});
}
- auto& info = infos.at(write->buffer);
+ auto& info = infos.at(write->source.as_or_throw<tvm::tirx::BufferVar>());
if (info.def == -1) {
info.def = stage;
} else {
@@ -464,11 +468,11 @@ class PipelineRewriter : public StmtExprMutator {
}
}
- for (const BufferRegion& read : block->reads) {
- if (!infos.count(read->buffer)) {
- infos.emplace(read->buffer, BufferAccessInfo{});
+ for (const TensorRegion& read : block->reads) {
+ if (!infos.count(read->source.as_or_throw<tvm::tirx::BufferVar>())) {
+ infos.emplace(read->source.as_or_throw<tvm::tirx::BufferVar>(), BufferAccessInfo{});
}
- auto& info = infos.at(read->buffer);
+ auto& info = infos.at(read->source.as_or_throw<tvm::tirx::BufferVar>());
info.use = std::max(info.use, stage);
}
}
@@ -529,10 +533,11 @@ class PipelineRewriter : public StmtExprMutator {
const SBlock& writer_block = pair1.first;
const auto& writer_info = pair1.second;
- auto it1 = std::find_if(writer_block->writes.begin(), writer_block->writes.end(),
- [&](const BufferRegion& buffer_region) {
- return buffer_region->buffer.same_as(buffer);
- });
+ auto it1 = std::find_if(
+ writer_block->writes.begin(), writer_block->writes.end(),
+ [&](const TensorRegion& buffer_region) {
+ return buffer_region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(buffer);
+ });
if (it1 == writer_block->writes.end()) {
continue;
}
@@ -540,10 +545,11 @@ class PipelineRewriter : public StmtExprMutator {
for (const auto& pair2 : pipeline_info_) {
const SBlock& reader_block = pair2.first;
const auto& reader_info = pair2.second;
- auto it2 = std::find_if(reader_block->reads.begin(), reader_block->reads.end(),
- [&](const BufferRegion& buffer_region) {
- return buffer_region->buffer.same_as(buffer);
- });
+ auto it2 = std::find_if(
+ reader_block->reads.begin(), reader_block->reads.end(),
+ [&](const TensorRegion& buffer_region) {
+ return buffer_region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(buffer);
+ });
if (it2 == reader_block->reads.end()) {
continue;
}
@@ -653,15 +659,18 @@ class PipelineRewriter : public StmtExprMutator {
if (new_blocks[i].is_async) {
// Record the fact that we have encountered these write buffers.
for (auto write_region : new_blocks[i].block->writes) {
- (*async_states_local)[new_blocks[i].stage].seen.insert(write_region->buffer.get());
+ (*async_states_local)[new_blocks[i].stage].seen.insert(
+ write_region->source.as_or_throw<tvm::tirx::BufferVar>().get());
}
}
int producer_stage_idx = -1;
for (auto read_region : new_blocks[i].block->reads) {
for (auto kv : async_states) {
- if (kv.first <= new_blocks[i].stage && kv.second.writes(read_region->buffer)) {
- // Found an earlier stage where read_region->buffer was asynchronously written
+ if (kv.first <= new_blocks[i].stage &&
+ kv.second.writes(read_region->source.as_or_throw<tvm::tirx::BufferVar>())) {
+ // Found an earlier stage where read_region->source.as_or_throw<tvm::tirx::BufferVar>()
+ // was asynchronously written
TVM_FFI_ICHECK(producer_stage_idx == -1 || producer_stage_idx == kv.first)
<< "A dependency on multiple async stages is not supported";
producer_stage_idx = kv.first;
@@ -727,11 +736,15 @@ class PipelineRewriter : public StmtExprMutator {
std::vector<bool> need_wait_count(num_commit_group, true);
for (auto read_region : new_blocks[i].block->reads) {
- if (!async_states[producer_stage_idx].writes(read_region->buffer)) continue;
- auto commit_group_id = buffer_to_commit_group.at(read_region->buffer.get());
+ if (!async_states[producer_stage_idx].writes(
+ read_region->source.as_or_throw<tvm::tirx::BufferVar>()))
+ continue;
+ auto commit_group_id = buffer_to_commit_group.at(
+ read_region->source.as_or_throw<tvm::tirx::BufferVar>().get());
if (!need_wait_count[commit_group_id]) continue;
- if (!dep_local_state.seen.count(read_region->buffer.get())) {
+ if (!dep_local_state.seen.count(
+ read_region->source.as_or_throw<tvm::tirx::BufferVar>().get())) {
// Multiple async producers interleaved: The most recent async write is from the
// previous iteration. This is the B_shared case above.
producer_head_per_commit.push_back(dep_local_state.producer_head.value() - 1);
@@ -963,8 +976,10 @@ class PipelineRewriter : public StmtExprMutator {
}
for (auto write_region : new_block->writes) {
- async_states[stage].dst_buffers.insert(write_region->buffer.get());
- buffer_to_commit_group[write_region->buffer.get()] = commit_group_id;
+ async_states[stage].dst_buffers.insert(
+ write_region->source.as_or_throw<tvm::tirx::BufferVar>().get());
+ buffer_to_commit_group[write_region->source.as_or_throw<tvm::tirx::BufferVar>().get()] =
+ commit_group_id;
}
local_state.producer_head = normalized_access_index;
@@ -985,7 +1000,8 @@ class PipelineRewriter : public StmtExprMutator {
for (auto read_region : new_block->reads) {
for (auto kv : async_states) {
int producer_stage_id = kv.first;
- if (producer_stage_id <= stage && kv.second.writes(read_region->buffer)) {
+ if (producer_stage_id <= stage &&
+ kv.second.writes(read_region->source.as_or_throw<tvm::tirx::BufferVar>())) {
async_states_local[producer_stage_id].consumed = true;
}
}
@@ -1065,8 +1081,8 @@ void BuildDependencyGraph(const ffi::Array<SBlock>& blocks,
std::unordered_map<Var, ffi::Array<SBlock>> buffer_writers;
for (const SBlock& block : blocks) {
- for (const BufferRegion& read : block->reads) {
- auto it = buffer_writers.find(read->buffer.var());
+ for (const TensorRegion& read : block->reads) {
+ auto it = buffer_writers.find(read->source.as_or_throw<tvm::tirx::BufferVar>().var());
if (it != buffer_writers.end()) {
for (const SBlock& writer : it->second) {
if (dep_src2dst != nullptr) {
@@ -1078,8 +1094,8 @@ void BuildDependencyGraph(const ffi::Array<SBlock>& blocks,
}
}
}
- for (const BufferRegion& write : block->writes) {
- buffer_writers[write->buffer.var()].push_back(block);
+ for (const TensorRegion& write : block->writes) {
+ buffer_writers[write->source.as_or_throw<tvm::tirx::BufferVar>().var()].push_back(block);
}
}
}
@@ -1291,7 +1307,7 @@ class PipelineInjector : public StmtExprMutator {
ValueError)
<< "Index of the buffer exceeds the size of the write regions of the block. ("
<< buffer_index << " vs. " << op->writes.size() << ")";
- double_buffers.insert(op->writes[buffer_index]->buffer);
+ double_buffers.insert(op->writes[buffer_index]->source.as_or_throw<tvm::tirx::BufferVar>());
}
SBlock block = StmtExprMutator::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
diff --git a/src/s_tir/transform/inject_virtual_thread.cc b/src/s_tir/transform/inject_virtual_thread.cc
index 1815973ea6..6397723504 100644
--- a/src/s_tir/transform/inject_virtual_thread.cc
+++ b/src/s_tir/transform/inject_virtual_thread.cc
@@ -260,10 +260,13 @@ class VTInjector : public tirx::IRMutatorWithAnalyzer {
return result;
}
// Variable
- UnchangedOr<Expr> Mutate_(const BufferRegionNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Expr> Mutate_(const TensorRegionNode* op, InplaceMode inplace_mode) final {
+ if (!op->source.as<BufferVar>()) {
+ return StmtExprMutator::Mutate_(op, inplace_mode);
+ }
auto region = Mutate(op->region).as_or_throw<UnchangedOr<ffi::Array<Range>>>();
if (region.UnchangedOrSameAs(op->region)) return ffi::Unchanged();
- BufferRegion node = ffi::GetRef<BufferRegion>(op);
+ TensorRegion node = ffi::GetRef<TensorRegion>(op);
node.CopyOnWrite()->region = std::move(region).ValueUnchecked();
return node;
}
diff --git a/src/s_tir/transform/lower_async_dma.cc b/src/s_tir/transform/lower_async_dma.cc
index 75fba27ea4..12fdd9e13a 100644
--- a/src/s_tir/transform/lower_async_dma.cc
+++ b/src/s_tir/transform/lower_async_dma.cc
@@ -76,17 +76,17 @@ class AsyncDMALowerer : public tirx::IRMutatorWithAnalyzer {
tvm::PrimExpr dst_min = mem_copy->dest->region[0]->min;
tvm::PrimExpr dst_extent = mem_copy->dest->region[0]->extent;
- auto src = BufferLoad(mem_copy->source->buffer, {src_min});
- auto dst = BufferLoad(mem_copy->dest->buffer, {dst_min});
+ auto src = BufferLoad(mem_copy->source->source.as_or_throw<BufferVar>(), {src_min});
+ auto dst = BufferLoad(mem_copy->dest->source.as_or_throw<BufferVar>(), {dst_min});
PrimExpr dst_nbytes = dst_extent * static_cast<int>(src.ty().StorageBytes());
return Evaluate(
Call(PrimType::Int(32), tirx::builtin::dma_copy(),
ffi::Array<Expr>{
PrimExpr(async_queue_id_.value()),
- Call(mem_copy->dest->buffer.DataPointerType(), tirx::builtin::address_of(),
- ffi::Array<Expr>{dst}, Attrs(), {}, Span()),
- Call(mem_copy->source->buffer.DataPointerType(), tirx::builtin::address_of(),
- ffi::Array<Expr>{src}, Attrs(), {}, Span()),
+ Call(mem_copy->dest->source.as_or_throw<BufferVar>().DataPointerType(),
+ tirx::builtin::address_of(), ffi::Array<Expr>{dst}, Attrs(), {}, Span()),
+ Call(mem_copy->source->source.as_or_throw<BufferVar>().DataPointerType(),
+ tirx::builtin::address_of(), ffi::Array<Expr>{src}, Attrs(), {}, Span()),
dst_nbytes, PrimExpr(dma_bypass_cache_)},
Attrs(), {}, Span())
.as_or_throw<PrimExpr>());
diff --git a/src/s_tir/transform/lower_cross_thread_reduction.cc b/src/s_tir/transform/lower_cross_thread_reduction.cc
index 606cda0c20..4926d9161f 100644
--- a/src/s_tir/transform/lower_cross_thread_reduction.cc
+++ b/src/s_tir/transform/lower_cross_thread_reduction.cc
@@ -83,16 +83,17 @@ bool IsDominantBlock(const SBlock& scope_block, const SBlock& block) {
// Step 1. Count the number of writers for each buffer written by the scope block.
std::unordered_map<const VarNode*, int> buffer_writer_cnt;
auto walk_fn = [&buffer_writer_cnt](const SBlock& block) -> ffi::Expected<ffi::WalkResult> {
- for (const BufferRegion& buffer_region : block->writes) {
- ++buffer_writer_cnt[buffer_region->buffer.get()];
+ for (const TensorRegion& buffer_region : block->writes) {
+ ++buffer_writer_cnt[buffer_region->source.as_or_throw<tvm::tirx::BufferVar>().get()];
}
return ffi::WalkResult::Skip();
};
ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(scope_block->body, walk_fn);
// Step 2. Check whether `block` is the only writer of its outputs.
- for (const BufferRegion& buffer_region : block->writes) {
- TVM_FFI_ICHECK(buffer_writer_cnt.count(buffer_region->buffer.get()));
- if (buffer_writer_cnt[buffer_region->buffer.get()] != 1) {
+ for (const TensorRegion& buffer_region : block->writes) {
+ TVM_FFI_ICHECK(
+ buffer_writer_cnt.count(buffer_region->source.as_or_throw<tvm::tirx::BufferVar>().get()));
+ if (buffer_writer_cnt[buffer_region->source.as_or_throw<tvm::tirx::BufferVar>().get()] != 1) {
return false;
}
}
@@ -333,7 +334,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize,
const SBlockNode* block = realize->block.get();
auto f_create_buffer_regions = [](ffi::Array<BufferVar> buffers) {
- ffi::Array<BufferRegion> regions;
+ ffi::Array<TensorRegion> regions;
regions.reserve(buffers.size());
for (const BufferVar& buffer : buffers) {
regions.push_back(BufferRegion(buffer, {Range::FromMinExtent(0, 1)}));
@@ -341,8 +342,8 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize,
return regions;
};
- ffi::Array<BufferRegion> ct_buffer_regions = f_create_buffer_regions(ct_buffers);
- ffi::Optional<ffi::Array<BufferRegion>> it_buffer_regions = std::nullopt;
+ ffi::Array<TensorRegion> ct_buffer_regions = f_create_buffer_regions(ct_buffers);
+ ffi::Optional<ffi::Array<TensorRegion>> it_buffer_regions = std::nullopt;
if (it_buffers.has_value()) {
it_buffer_regions = f_create_buffer_regions(it_buffers.value());
}
@@ -423,7 +424,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize,
// Step 3.2. Create the block and the block-realize.
ffi::Array<IterVar> iter_vars{nullptr};
ffi::Array<PrimExpr> bindings{nullptr};
- ffi::Array<BufferRegion> reads{nullptr};
+ ffi::Array<TensorRegion> reads{nullptr};
if (it_buffers.has_value()) {
iter_vars = ffi::Array<IterVar>{};
bindings = ffi::Array<PrimExpr>{};
@@ -477,7 +478,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize,
}
}
ffi::Array<Stmt> wb_updates;
- ffi::Array<BufferRegion> wb_regions;
+ ffi::Array<TensorRegion> wb_regions;
wb_updates.reserve(n_buffers);
wb_regions.reserve(n_buffers);
int n_dim = static_cast<int>(old_wb_indices.size());
@@ -632,8 +633,8 @@ class CrossThreadReductionTransformer : public StmtExprMutator {
SBlock block = realize->block;
// If the block writes to local memory, no rewrite is needed.
- for (BufferRegion write_region : block->writes) {
- if (write_region->buffer.scope() == "local") {
+ for (TensorRegion write_region : block->writes) {
+ if (write_region->source.as_or_throw<tvm::tirx::BufferVar>().scope() == "local") {
return {};
}
}
@@ -641,8 +642,9 @@ class CrossThreadReductionTransformer : public StmtExprMutator {
// Find out the reduction threads for the read-buffers which are produced by
// cross-thread reduction.
std::unordered_map<ThreadScope, Range, ThreadScopeHash, ThreadScopeEqual> thread2range;
- for (BufferRegion read_region : block->reads) {
- auto buf_it = crt_buf2threads_.find(read_region->buffer.get());
+ for (TensorRegion read_region : block->reads) {
+ auto buf_it =
+ crt_buf2threads_.find(read_region->source.as_or_throw<tvm::tirx::BufferVar>().get());
if (buf_it == crt_buf2threads_.end()) {
continue;
}
diff --git a/src/s_tir/transform/lower_match_buffer.cc b/src/s_tir/transform/lower_match_buffer.cc
index 5cb43aaba5..61196b2b1d 100644
--- a/src/s_tir/transform/lower_match_buffer.cc
+++ b/src/s_tir/transform/lower_match_buffer.cc
@@ -78,9 +78,9 @@ class MatchBufferLower : public StmtExprMutator {
}
op = stmt.as<SBlockNode>();
TVM_FFI_ICHECK(op != nullptr);
- ffi::Array<BufferRegion> reads =
+ ffi::Array<TensorRegion> reads =
op->reads.Map(std::bind(&MatchBufferLower::VisitBufferRegion, this, std::placeholders::_1));
- ffi::Array<BufferRegion> writes = op->writes.Map(
+ ffi::Array<TensorRegion> writes = op->writes.Map(
std::bind(&MatchBufferLower::VisitBufferRegion, this, std::placeholders::_1));
if (reads.same_as(op->reads) && writes.same_as(op->writes) && op->match_buffers.empty()) {
@@ -114,7 +114,7 @@ class MatchBufferLower : public StmtExprMutator {
var.has_value() && var.value()->ty.as<BufferTypeNode>()) {
auto it = match_buffers_.find(BufferVar(var.value()));
if (it != match_buffers_.end()) {
- return (*it).second->buffer.data();
+ return (*it).second->source.as_or_throw<tvm::tirx::BufferVar>().data();
}
}
}
@@ -136,11 +136,11 @@ class MatchBufferLower : public StmtExprMutator {
return stmt;
} else {
const BufferVar& buffer = (*it).first;
- const BufferRegion& source = (*it).second;
+ const TensorRegion& source = (*it).second;
auto* n = stmt.CopyOnWrite();
n->indices = ConvertIndices(MatchBufferRegion(buffer, source), op->indices);
- n->buffer = source->buffer;
+ n->buffer = source->source.as_or_throw<tvm::tirx::BufferVar>();
return stmt;
}
}
@@ -158,29 +158,29 @@ class MatchBufferLower : public StmtExprMutator {
return expr;
} else {
const BufferVar& buffer = (*it).first;
- const BufferRegion& source = (*it).second;
+ const TensorRegion& source = (*it).second;
ffi::Array<PrimExpr> indices = ConvertIndices(MatchBufferRegion(buffer, source), op->indices);
- return BufferLoad(source->buffer, indices);
+ return BufferLoad(source->source.as_or_throw<tvm::tirx::BufferVar>(), indices);
}
}
- BufferRegion VisitBufferRegion(const BufferRegion& buffer_region) {
- const BufferVar& buffer = buffer_region->buffer;
+ TensorRegion VisitBufferRegion(const TensorRegion& buffer_region) {
+ const BufferVar& buffer = buffer_region->source.as_or_throw<tvm::tirx::BufferVar>();
auto it = match_buffers_.find(buffer);
if (it == match_buffers_.end()) {
return buffer_region;
} else {
- const BufferRegion& source = (*it).second;
+ const TensorRegion& source = (*it).second;
Region region = ConvertRegion(MatchBufferRegion(buffer, source), buffer_region->region);
- return BufferRegion(source->buffer, std::move(region));
+ return BufferRegion(source->source.as_or_throw<tvm::tirx::BufferVar>(), std::move(region));
}
}
void CheckAndUpdateVarMap(const MatchBufferRegion& match_buffer) {
// Step.1. Check
const BufferVar& buffer = match_buffer->buffer;
- const BufferRegion& source = VisitBufferRegion(match_buffer->source);
- const BufferVar& source_buffer = source->buffer;
+ const TensorRegion& source = VisitBufferRegion(match_buffer->source);
+ const BufferVar& source_buffer = source->source.as_or_throw<tvm::tirx::BufferVar>();
// Step.1.1. Check scope & dtype
TVM_FFI_ICHECK_EQ(buffer.scope(), source_buffer.scope())
@@ -310,7 +310,7 @@ class MatchBufferLower : public StmtExprMutator {
}
/*! \brief BufferVar region mapping. */
- ffi::Map<BufferVar, BufferRegion> match_buffers_;
+ ffi::Map<BufferVar, TensorRegion> match_buffers_;
/*! \brief The analyzer */
arith::Analyzer analyzer_;
};
diff --git a/src/s_tir/transform/manifest_shared_memory_local_stage.cc b/src/s_tir/transform/manifest_shared_memory_local_stage.cc
index f49f557bae..f7484ce295 100644
--- a/src/s_tir/transform/manifest_shared_memory_local_stage.cc
+++ b/src/s_tir/transform/manifest_shared_memory_local_stage.cc
@@ -128,7 +128,7 @@ class IntermediateStageRewriter {
Stmt local_stage = BufferStore(new_buffer, store->value, local_stage_indices);
// Step 1: Make block and block realize
- BufferRegion write_buffer_region = BufferRegion::FromPoint(new_buffer, local_stage_indices);
+ TensorRegion write_buffer_region = BufferRegionFromPoint(new_buffer, local_stage_indices);
local_stage =
SBlock(/*iter_vars=*/{}, /*reads=*/block->reads, /*writes=*/{write_buffer_region}, "",
/*body=*/std::move(local_stage));
diff --git a/src/s_tir/transform/memhammer_coalesce.cc b/src/s_tir/transform/memhammer_coalesce.cc
index 4a0263ff83..5a7897aca4 100644
--- a/src/s_tir/transform/memhammer_coalesce.cc
+++ b/src/s_tir/transform/memhammer_coalesce.cc
@@ -169,10 +169,11 @@ ffi::Array<PrimExpr> GetMapping(const Stmt& stmt, const ConstraintSet& constrain
body = loop->body;
}
const BufferStoreNode* buf_store = TVM_TYPE_AS(body, BufferStoreNode);
- BufferRegion write_region = constraints.write_region;
+ TensorRegion write_region = constraints.write_region;
const ffi::Array<PrimExpr>& write_index = buf_store->indices;
- TVM_FFI_ICHECK(write_region->region.size() == write_index.size() &&
- write_region->buffer.same_as(buf_store->buffer));
+ TVM_FFI_ICHECK(
+ write_region->region.size() == write_index.size() &&
+ write_region->source.as_or_throw<tvm::tirx::BufferVar>().same_as(buf_store->buffer));
ffi::Array<PrimExpr> result;
arith::Analyzer analyzer;
for (int i = 0; i < static_cast<int>(write_region->region.size()); i++) {
@@ -204,8 +205,8 @@ Stmt InverseMapping::Rewrite(const Stmt& stmt, const ConstraintSet& constraints,
ffi::Map<Var, PrimExpr> inverse_mapping =
arith::InverseAffineIterMap(iter_map->indices, loop_vars);
// Step 3. Generate new body
- BufferRegion read_region = constraints.read_region;
- BufferRegion write_region = constraints.write_region;
+ TensorRegion read_region = constraints.read_region;
+ TensorRegion write_region = constraints.write_region;
ffi::Array<PrimExpr> write_index;
ffi::Array<PrimExpr> read_index;
ffi::Array<PrimVar> new_loop_vars;
@@ -237,8 +238,10 @@ Stmt InverseMapping::Rewrite(const Stmt& stmt, const ConstraintSet& constraints,
read_index.push_back(read_region->region[i]->min + inverse);
}
}
- TensorLoad new_buf_load = BufferLoad(read_region->buffer, read_index);
- BufferStore new_buf_store = BufferStore(write_region->buffer, new_buf_load, write_index);
+ TensorLoad new_buf_load =
+ BufferLoad(read_region->source.as_or_throw<tvm::tirx::BufferVar>(), read_index);
+ BufferStore new_buf_store = BufferStore(write_region->source.as_or_throw<tvm::tirx::BufferVar>(),
+ new_buf_load, write_index);
Stmt ret = new_buf_store;
// Step 3.3 construct loop body
for (int i = static_cast<int>(new_loop_vars.size()) - 1; i >= 0; i--) {
diff --git a/src/s_tir/transform/memhammer_lower_auto_copy.cc b/src/s_tir/transform/memhammer_lower_auto_copy.cc
index 13b3f963a4..ba3b2b0108 100644
--- a/src/s_tir/transform/memhammer_lower_auto_copy.cc
+++ b/src/s_tir/transform/memhammer_lower_auto_copy.cc
@@ -231,9 +231,10 @@ class AutoPadder {
// after mutation. Otherwise we just return the original block.
bool changed = false;
// Step 1. Mutate the read region.
- ffi::Array<BufferRegion> reads;
- for (const BufferRegion& read : op->reads) {
- if (auto replacement = VarRemapGet(read->buffer).as<BufferVar>()) {
+ ffi::Array<TensorRegion> reads;
+ for (const TensorRegion& read : op->reads) {
+ if (auto replacement =
+ VarRemapGet(read->source.as_or_throw<tvm::tirx::BufferVar>()).as<BufferVar>()) {
changed = true;
reads.push_back(BufferRegion(replacement.value(), read->region));
} else {
@@ -241,9 +242,10 @@ class AutoPadder {
}
}
// Step 2. Mutate the write region.
- ffi::Array<BufferRegion> writes;
- for (const BufferRegion& write : op->writes) {
- if (auto replacement = VarRemapGet(write->buffer).as<BufferVar>()) {
+ ffi::Array<TensorRegion> writes;
+ for (const TensorRegion& write : op->writes) {
+ if (auto replacement =
+ VarRemapGet(write->source.as_or_throw<tvm::tirx::BufferVar>()).as<BufferVar>()) {
changed = true;
writes.push_back(BufferRegion(replacement.value(), write->region));
} else {
@@ -254,7 +256,9 @@ class AutoPadder {
// MatchBufferRegion, the storage scope of the target buffer also needs to be set.
ffi::Array<MatchBufferRegion> match_buffers;
for (const MatchBufferRegion& match_buffer : op->match_buffers) {
- if (auto replacement = VarRemapGet(match_buffer->source->buffer).as<BufferVar>()) {
+ if (auto replacement =
+ VarRemapGet(match_buffer->source->source.as_or_throw<tvm::tirx::BufferVar>())
+ .as<BufferVar>()) {
changed = true;
BufferVar new_buffer = replacement.value();
match_buffers.push_back(MatchBufferRegion(
@@ -612,7 +616,7 @@ class AutoPadder {
if (call->op.same_as(tvm_load_matrix_sync_op) ||
call->op.same_as(tvm_store_matrix_sync_op)) {
for (const MatchBufferRegion& r : op->match_buffers) {
- BufferVar src_buffer = r->source->buffer;
+ BufferVar src_buffer = r->source->source.as_or_throw<tvm::tirx::BufferVar>();
runtime::StorageScope scope = runtime::StorageScope::Create(src_buffer.scope());
if (scope.rank == runtime::StorageRank::kShared) {
Region region = r->source->region;
@@ -731,11 +735,12 @@ class AutoCopyMutator : public StmtExprMutator {
TVM_FFI_ICHECK_EQ(block->writes.size(), 1);
TVM_FFI_ICHECK_GE(block->reads.size(), 1);
- BufferRegion target_read = block->reads[0];
+ TensorRegion target_read = block->reads[0];
if (block->reads.size() > 1) {
bool found = false;
for (size_t i = 0; i < block->reads.size(); i++) {
- if (block->reads[i]->buffer.scope() == "wmma.accumulator") {
+ if (block->reads[i]->source.as_or_throw<tvm::tirx::BufferVar>().scope() ==
+ "wmma.accumulator") {
found = true;
target_read = block->reads[i];
}
@@ -743,7 +748,7 @@ class AutoCopyMutator : public StmtExprMutator {
TVM_FFI_ICHECK(found) << "Multiple buffer read";
}
- int data_bits = target_read->buffer->dtype.bits();
+ int data_bits = target_read->source.as_or_throw<tvm::tirx::BufferVar>()->dtype.bits();
ConstraintSet constraints(this->thread_extent_, //
this->outer_loops_, //
target_read, //
diff --git a/src/s_tir/transform/memhammer_rewrite_rule.h b/src/s_tir/transform/memhammer_rewrite_rule.h
index 64806b02c3..3ae73d4119 100644
--- a/src/s_tir/transform/memhammer_rewrite_rule.h
+++ b/src/s_tir/transform/memhammer_rewrite_rule.h
@@ -43,9 +43,9 @@ struct ConstraintSet {
/*! \brief The outer loops surrounding the data copy */
ffi::Array<For> outer_loops;
/*! \brief The read region of the data copy */
- BufferRegion read_region;
+ TensorRegion read_region;
/*! \brief The write region of the data copy */
- BufferRegion write_region;
+ TensorRegion write_region;
/*! \brief The dtype size in bits */
int data_bits;
/*! \brief Whether to insert a local stage in the data copy */
@@ -55,8 +55,8 @@ struct ConstraintSet {
explicit ConstraintSet(ffi::Map<ffi::String, int64_t> thread_extent, //
ffi::Array<For> outer_loops, //
- BufferRegion read_region, //
- BufferRegion write_region, //
+ TensorRegion read_region, //
+ TensorRegion write_region, //
int data_bits, //
const ffi::Map<ffi::String, ffi::Any>& ann)
: thread_extent(thread_extent),
@@ -135,8 +135,8 @@ class CoalescedAccess : public RewriteRule {
CoalescedAccess() = default;
Stmt Rewrite(const Stmt& stmt, const ConstraintSet& constraints, OutputSet* output) const final;
bool CanApply(const Stmt& stmt, const ConstraintSet& constraints) const final {
- BufferVar src_buffer = constraints.read_region->buffer;
- BufferVar tgt_buffer = constraints.write_region->buffer;
+ BufferVar src_buffer = constraints.read_region->source.as_or_throw<tvm::tirx::BufferVar>();
+ BufferVar tgt_buffer = constraints.write_region->source.as_or_throw<tvm::tirx::BufferVar>();
return IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kGlobal,
runtime::StorageRank::kShared) ||
IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kShared,
@@ -152,8 +152,8 @@ class InverseMapping : public RewriteRule {
InverseMapping() = default;
Stmt Rewrite(const Stmt& stmt, const ConstraintSet& constraints, OutputSet* output) const final;
bool CanApply(const Stmt& stmt, const ConstraintSet& constraints) const final {
- BufferVar src_buffer = constraints.read_region->buffer;
- BufferVar tgt_buffer = constraints.write_region->buffer;
+ BufferVar src_buffer = constraints.read_region->source.as_or_throw<tvm::tirx::BufferVar>();
+ BufferVar tgt_buffer = constraints.write_region->source.as_or_throw<tvm::tirx::BufferVar>();
return IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kShared,
runtime::StorageRank::kGlobal);
}
@@ -167,8 +167,8 @@ class CreateLocalStage : public RewriteRule {
CreateLocalStage() = default;
Stmt Rewrite(const Stmt& stmt, const ConstraintSet& constraints, OutputSet* output) const final;
bool CanApply(const Stmt& stmt, const ConstraintSet& constraints) const final {
- BufferVar src_buffer = constraints.read_region->buffer;
- BufferVar tgt_buffer = constraints.write_region->buffer;
+ BufferVar src_buffer = constraints.read_region->source.as_or_throw<tvm::tirx::BufferVar>();
+ BufferVar tgt_buffer = constraints.write_region->source.as_or_throw<tvm::tirx::BufferVar>();
return IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kGlobal,
runtime::StorageRank::kShared) &&
is_one(constraints.add_local_stage);
@@ -184,8 +184,8 @@ class WmmaToGlobal : public RewriteRule {
WmmaToGlobal() = default;
Stmt Rewrite(const Stmt& stmt, const ConstraintSet& constraints, OutputSet* output) const final;
bool CanApply(const Stmt& stmt, const ConstraintSet& constraints) const final {
- BufferVar src_buffer = constraints.read_region->buffer;
- BufferVar tgt_buffer = constraints.write_region->buffer;
+ BufferVar src_buffer = constraints.read_region->source.as_or_throw<tvm::tirx::BufferVar>();
+ BufferVar tgt_buffer = constraints.write_region->source.as_or_throw<tvm::tirx::BufferVar>();
return IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kWMMAAccumulator,
runtime::StorageRank::kGlobal);
}
@@ -200,8 +200,8 @@ class MmaToGlobal : public RewriteRule {
MmaToGlobal() = default;
Stmt Rewrite(const Stmt& stmt, const ConstraintSet& constraints, OutputSet* output) const final;
bool CanApply(const Stmt& stmt, const ConstraintSet& constraints) const final {
- BufferVar src_buffer = constraints.read_region->buffer;
- BufferVar tgt_buffer = constraints.write_region->buffer;
+ BufferVar src_buffer = constraints.read_region->source.as_or_throw<tvm::tirx::BufferVar>();
+ BufferVar tgt_buffer = constraints.write_region->source.as_or_throw<tvm::tirx::BufferVar>();
return IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kMMAMatrixC,
runtime::StorageRank::kGlobal);
}
@@ -215,8 +215,8 @@ class SharedToWmma : public RewriteRule {
SharedToWmma() = default;
Stmt Rewrite(const Stmt& stmt, const ConstraintSet& constraints, OutputSet* output) const final;
bool CanApply(const Stmt& stmt, const ConstraintSet& constraints) const final {
- BufferVar src_buffer = constraints.read_region->buffer;
- BufferVar tgt_buffer = constraints.write_region->buffer;
+ BufferVar src_buffer = constraints.read_region->source.as_or_throw<tvm::tirx::BufferVar>();
+ BufferVar tgt_buffer = constraints.write_region->source.as_or_throw<tvm::tirx::BufferVar>();
return IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kShared,
runtime::StorageRank::kWMMAMatrixA) ||
IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kShared,
@@ -232,8 +232,8 @@ class WmmaToShared : public RewriteRule {
WmmaToShared() = default;
Stmt Rewrite(const Stmt& stmt, const ConstraintSet& constraints, OutputSet* output) const final;
bool CanApply(const Stmt& stmt, const ConstraintSet& constraints) const final {
- BufferVar src_buffer = constraints.read_region->buffer;
- BufferVar tgt_buffer = constraints.write_region->buffer;
+ BufferVar src_buffer = constraints.read_region->source.as_or_throw<tvm::tirx::BufferVar>();
+ BufferVar tgt_buffer = constraints.write_region->source.as_or_throw<tvm::tirx::BufferVar>();
return IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kWMMAAccumulator,
runtime::StorageRank::kShared);
}
diff --git a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
index b46ef0ab9b..7363572751 100644
--- a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
+++ b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
@@ -310,14 +310,14 @@ Stmt RewriteWmmaStore(Stmt stmt) {
Stmt SharedToWmma::Rewrite(const Stmt& stmt, const ConstraintSet& constraints,
OutputSet* output) const {
Stmt after_tiling = TileWmmaBlock(stmt).first;
- output->padding_min.Set(constraints.read_region->buffer, 8);
+ output->padding_min.Set(constraints.read_region->source.as_or_throw<tvm::tirx::BufferVar>(), 8);
return RewriteWmmaLoad(after_tiling);
}
Stmt WmmaToShared::Rewrite(const Stmt& stmt, const ConstraintSet& constraints,
OutputSet* output) const {
Stmt after_tiling = TileWmmaBlock(stmt).first;
- output->padding_min.Set(constraints.write_region->buffer, 8);
+ output->padding_min.Set(constraints.write_region->source.as_or_throw<tvm::tirx::BufferVar>(), 8);
return RewriteWmmaStore(after_tiling);
}
diff --git a/src/s_tir/transform/plan_update_buffer_allocation_location.cc b/src/s_tir/transform/plan_update_buffer_allocation_location.cc
index 7d4e2cacbb..f49ba03a24 100644
--- a/src/s_tir/transform/plan_update_buffer_allocation_location.cc
+++ b/src/s_tir/transform/plan_update_buffer_allocation_location.cc
@@ -80,8 +80,9 @@ class BufferAllocateOrderCollector : public StmtExprVisitor {
// Also visit match_buffers to collect buffers that only appear in read and match_buffer
// regions.
for (const auto& region : op->match_buffers) {
- if (!find(region->source->buffer)) {
- buffer_alloc_recorder_.push_back(region->source->buffer);
+ if (!find(region->source->source.as_or_throw<tvm::tirx::BufferVar>())) {
+ buffer_alloc_recorder_.push_back(
+ region->source->source.as_or_throw<tvm::tirx::BufferVar>());
}
}
@@ -183,7 +184,7 @@ class BufferAllocationLocator : public StmtExprMutator {
}
for (const MatchBufferRegion match_buffer : op->match_buffers) {
const Var target_var = match_buffer->buffer.var();
- const Var source_var = match_buffer->source->buffer.var();
+ const Var source_var = match_buffer->source->source.as_or_throw<tvm::tirx::BufferVar>().var();
TVM_FFI_ICHECK(buffer_data_to_buffer_.count(source_var));
buffer_data_to_buffer_.Set(target_var, match_buffer->buffer);
}
@@ -224,7 +225,7 @@ class BufferAllocationLocator : public StmtExprMutator {
/*init=*/std::nullopt,
/*alloc_buffers=*/alloc_buffers);
SBlockNode* n = opaque_block.CopyOnWrite();
- ffi::Array<ffi::Array<BufferRegion>> access =
+ ffi::Array<ffi::Array<TensorRegion>> access =
GetSBlockReadWriteRegion(opaque_block, buffer_data_to_buffer_);
n->reads = access[0];
n->writes = access[1];
@@ -232,11 +233,12 @@ class BufferAllocationLocator : public StmtExprMutator {
return realize;
}
- ffi::Array<BufferRegion> RemoveRedundantBufferRegion(
- const ffi::Array<BufferRegion>& region) const {
- ffi::Array<BufferRegion> result;
- for (const BufferRegion& buffer_region : region) {
- if (buffer_data_to_buffer_.count(buffer_region->buffer.var())) {
+ ffi::Array<TensorRegion> RemoveRedundantBufferRegion(
+ const ffi::Array<TensorRegion>& region) const {
+ ffi::Array<TensorRegion> result;
+ for (const TensorRegion& buffer_region : region) {
+ if (buffer_data_to_buffer_.count(
+ buffer_region->source.as_or_throw<tvm::tirx::BufferVar>().var())) {
result.push_back(buffer_region);
}
}
diff --git a/src/s_tir/transform/renew_defs.cc b/src/s_tir/transform/renew_defs.cc
index 70b5ff844e..9ddc21bdb9 100644
--- a/src/s_tir/transform/renew_defs.cc
+++ b/src/s_tir/transform/renew_defs.cc
@@ -126,9 +126,9 @@ class RenewDefMutator : public StmtExprMutator {
Stmt body = this->Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body);
// Step 4. Revisit access region
- ffi::Array<BufferRegion> reads =
+ ffi::Array<TensorRegion> reads =
op->reads.Map(std::bind(&RenewDefMutator::VisitBufferRegion, this, std::placeholders::_1));
- ffi::Array<BufferRegion> writes =
+ ffi::Array<TensorRegion> writes =
op->writes.Map(std::bind(&RenewDefMutator::VisitBufferRegion, this, std::placeholders::_1));
// Step 5. Regenerate block. Since the defs are changed, we need to create a new block
@@ -231,7 +231,7 @@ class RenewDefMutator : public StmtExprMutator {
MatchBufferRegion VisitMatchBuffer(const MatchBufferRegion& match_buffer) {
BufferVar buffer = DefineBuffer(match_buffer->buffer);
- BufferRegion region = VisitBufferRegion(match_buffer->source);
+ TensorRegion region = VisitBufferRegion(match_buffer->source);
return MatchBufferRegion(std::move(buffer), std::move(region));
}
@@ -249,11 +249,12 @@ class RenewDefMutator : public StmtExprMutator {
}
}
- BufferRegion VisitBufferRegion(const BufferRegion& buffer_region) {
- BufferVar buffer = UseOrRemapBuffer(buffer_region->buffer);
+ TensorRegion VisitBufferRegion(const TensorRegion& buffer_region) {
+ BufferVar buffer = UseOrRemapBuffer(buffer_region->source.as_or_throw<tvm::tirx::BufferVar>());
ffi::Array<Range> region = buffer_region->region.Map(
std::bind(&RenewDefMutator::VisitRange, this, std::placeholders::_1));
- if (buffer.same_as(buffer_region->buffer) && region.same_as(buffer_region->region)) {
+ if (buffer.same_as(buffer_region->source.as_or_throw<tvm::tirx::BufferVar>()) &&
+ region.same_as(buffer_region->region)) {
return buffer_region;
} else {
return BufferRegion(std::move(buffer), std::move(region));
diff --git a/src/tirx/analysis/verify_well_formed.cc b/src/tirx/analysis/verify_well_formed.cc
index 9241047757..48a2d2d897 100644
--- a/src/tirx/analysis/verify_well_formed.cc
+++ b/src/tirx/analysis/verify_well_formed.cc
@@ -106,7 +106,7 @@ class BlockVarAccessVerifier : public StmtExprVisitor {
// Step 0. Skip block iter var's domain
// Step 1. Visit read/write regions
- auto fvisit_buffer_region = [this](const BufferRegion& s) -> ffi::Optional<VisitInterrupt> {
+ auto fvisit_buffer_region = [this](const TensorRegion& s) -> ffi::Optional<VisitInterrupt> {
for (const auto& range : s->region) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(range->min));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(range->extent));
@@ -246,7 +246,7 @@ class UndefinedVarVerifier : public Verifier<UndefinedVarVerifier> {
* - SBlock::alloc_buffers
* - SBlock::match_buffers
*
- * it must not appear in a BufferLoad, BufferStore, or BufferRegion outside that declaration's
+ * it must not appear in a BufferLoad, BufferStore, or TensorRegion outside that declaration's
* scope.
*
* All buffers that appear in TensorLoad or BufferStore must have a prior declaration.
diff --git a/src/tirx/ir/data_type_rewriter.cc b/src/tirx/ir/data_type_rewriter.cc
index d27daa68b9..2b7504513d 100644
--- a/src/tirx/ir/data_type_rewriter.cc
+++ b/src/tirx/ir/data_type_rewriter.cc
@@ -435,14 +435,14 @@ UnchangedOr<Stmt> IndexDataTypeRewriter::Mutate_(const SBlockNode* op, InplaceMo
.as_or_throw<UnchangedOr<BufferVar>>()
.ValueOrUnchanged(match->buffer);
});
- BufferRegion source = VisitBufferRegion(match->source);
+ TensorRegion source = VisitBufferRegion(match->source);
if (buffer.same_as(match->buffer) && source.same_as(match->source)) return match;
return MatchBufferRegion(buffer, source);
});
- ffi::Array<BufferRegion> new_reads = op->reads.Map(
- [this](const BufferRegion& buffer_region) { return this->VisitBufferRegion(buffer_region); });
- ffi::Array<BufferRegion> new_writes = op->writes.Map(
- [this](const BufferRegion& buffer_region) { return this->VisitBufferRegion(buffer_region); });
+ ffi::Array<TensorRegion> new_reads = op->reads.Map(
+ [this](const TensorRegion& buffer_region) { return this->VisitBufferRegion(buffer_region); });
+ ffi::Array<TensorRegion> new_writes = op->writes.Map(
+ [this](const TensorRegion& buffer_region) { return this->VisitBufferRegion(buffer_region); });
ffi::Array<IterVar> new_iter_vars =
op->iter_vars.Map([this](const IterVar& iter_var) { return this->VisitIterVar(iter_var); });
ffi::Optional<Stmt> new_init = std::nullopt;
@@ -528,10 +528,11 @@ IterVar IndexDataTypeRewriter::VisitIterVar(const IterVar& iter_var) {
return iter_var;
}
-BufferRegion IndexDataTypeRewriter::VisitBufferRegion(const BufferRegion& buffer_region) {
- BufferVar remapped_buffer = Mutate(buffer_region->buffer, InplaceMode::kDisallow)
- .as_or_throw<UnchangedOr<BufferVar>>()
- .ValueOrUnchanged(buffer_region->buffer);
+TensorRegion IndexDataTypeRewriter::VisitBufferRegion(const TensorRegion& buffer_region) {
+ BufferVar remapped_buffer =
+ Mutate(buffer_region->source.as_or_throw<BufferVar>(), InplaceMode::kDisallow)
+ .as_or_throw<UnchangedOr<BufferVar>>()
+ .ValueOrUnchanged(buffer_region->source.as_or_throw<BufferVar>());
bool is_enabled = is_enabled_;
is_enabled_ = true;
@@ -542,7 +543,7 @@ BufferRegion IndexDataTypeRewriter::VisitBufferRegion(const BufferRegion& buffer
});
is_enabled_ = is_enabled;
- if (!remapped_buffer.same_as(buffer_region->buffer) ||
+ if (!remapped_buffer.same_as(buffer_region->source.as_or_throw<BufferVar>()) ||
!new_region.same_as(buffer_region->region)) {
return BufferRegion(remapped_buffer, new_region);
} else {
diff --git a/src/tirx/ir/data_type_rewriter.h b/src/tirx/ir/data_type_rewriter.h
index ad6eea46e5..c2cca01ed8 100644
--- a/src/tirx/ir/data_type_rewriter.h
+++ b/src/tirx/ir/data_type_rewriter.h
@@ -126,7 +126,7 @@ class IndexDataTypeRewriter : public DataTypeLegalizer {
ffi::Map<ffi::String, ffi::Any> VisitBlockAnnotations(
const ffi::Map<ffi::String, ffi::Any>& annotations);
- BufferRegion VisitBufferRegion(const BufferRegion& region);
+ TensorRegion VisitBufferRegion(const TensorRegion& region);
IterVar VisitIterVar(const IterVar& iter_var);
// indicator of index expr to rewrite
bool is_enabled_{false};
diff --git a/src/tirx/ir/script/script_complete.cc b/src/tirx/ir/script/script_complete.cc
index 1da981116d..5d55087aa3 100644
--- a/src/tirx/ir/script/script_complete.cc
+++ b/src/tirx/ir/script/script_complete.cc
@@ -93,9 +93,9 @@ class ScriptCompleter : public StmtExprMutator {
// ignore root block or blocks which already has reads/writes regions
if (mask != 0 && s_tir_) {
auto access_region = GetSBlockAccessRegion(block, *buffer_var_map_);
- const ffi::Array<BufferRegion>& reads = access_region[0];
- const ffi::Array<BufferRegion>& writes = access_region[1];
- const ffi::Array<BufferRegion>& opaque = access_region[2];
+ const ffi::Array<TensorRegion>& reads = access_region[0];
+ const ffi::Array<TensorRegion>& writes = access_region[1];
+ const ffi::Array<TensorRegion>& opaque = access_region[2];
TVM_FFI_CHECK(opaque.empty(), ValueError)
<< "Can not auto detect buffer access region from tirx.Load, tirx.Store or "
"direct access by buffer data. Please annotation the access region manually";
diff --git a/src/tirx/ir/specialize.cc b/src/tirx/ir/specialize.cc
index e9ef406d5f..1455517378 100644
--- a/src/tirx/ir/specialize.cc
+++ b/src/tirx/ir/specialize.cc
@@ -189,10 +189,10 @@ class PrimFuncSpecializer : public StmtExprMutator {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(iter->dom->min));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(iter->dom->extent));
}
- for (const BufferRegion& region : op->reads) {
+ for (const TensorRegion& region : op->reads) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(region));
}
- for (const BufferRegion& region : op->writes) {
+ for (const TensorRegion& region : op->writes) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(region));
}
for (const MatchBufferRegion& match : op->match_buffers) {
@@ -238,11 +238,14 @@ class PrimFuncSpecializer : public StmtExprMutator {
return load;
}
- UnchangedOr<Expr> Mutate_(const BufferRegionNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Expr> Mutate_(const TensorRegionNode* op, InplaceMode inplace_mode) final {
auto result = StmtExprMutator::Mutate_(op, InplaceMode::kDisallow);
if (result.UnchangedOrSameAs(ffi::GetRef<Expr>(op))) return ffi::Unchanged();
- auto region = std::move(result).ValueUnchecked().as_or_throw<BufferRegion>();
- return BufferRegion(region->buffer, region->region);
+ auto region = std::move(result).ValueUnchecked().as_or_throw<TensorRegion>();
+ if (auto buffer = region->source.as<BufferVar>()) {
+ return BufferRegion(buffer.value(), region->region, region->span);
+ }
+ return region;
}
DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::AddNode, add);
diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc
index 2dd2c5a3f3..b873b1e0a6 100644
--- a/src/tirx/ir/stmt.cc
+++ b/src/tirx/ir/stmt.cc
@@ -650,7 +650,7 @@ TVMFFIAny BufferStoreMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
}
ffi::ObjectRef RealizeBufferRegionSubscript(Expr value, SubscriptSlice slice, Span span) {
- BufferRegion source = value.as_or_throw<BufferRegion>();
+ TensorRegion source = value.as_or_throw<TensorRegion>();
TVM_FFI_CHECK_LE(slice.size(), source->region.size(), IndexError)
<< "Too many indices for a " << source->region.size() << "-dimensional buffer region";
@@ -661,7 +661,7 @@ ffi::ObjectRef RealizeBufferRegionSubscript(Expr value, SubscriptSlice slice, Sp
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";
+ << "TensorRegion slices with a non-unit step are not supported";
}
}
@@ -671,7 +671,7 @@ ffi::ObjectRef RealizeBufferRegionSubscript(Expr value, SubscriptSlice slice, Sp
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);
+ return BufferLoad(source->source.as_or_throw<BufferVar>(), indices, span);
}
arith::Analyzer analyzer;
@@ -696,7 +696,7 @@ ffi::ObjectRef RealizeBufferRegionSubscript(Expr value, SubscriptSlice slice, Sp
for (size_t i = slice.size(); i < source->region.size(); ++i) {
region.push_back(source->region[i]);
}
- return BufferRegion(source->buffer, region, span);
+ return BufferRegion(source->source.as_or_throw<BufferVar>(), region, span);
}
TVMFFIAny BufferRegionTypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
@@ -711,50 +711,6 @@ TVMFFIAny BufferRegionTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::An
return ffi::Unchanged().CopyToTVMFFIAny();
}
-TVMFFIAny BufferRegionVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
- const BufferRegionNode* self =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BufferRegionNode>(value);
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->buffer));
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->region));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
-}
-
-TVMFFIAny BufferRegionMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
- const BufferRegionNode* self =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BufferRegionNode>(value);
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<BufferVar>, mapped_buffer,
- mutator->MutateExpected(self->buffer));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<Range>>, mapped_region,
- mutator->MutateExpected(self->region));
- if (mapped_buffer.UnchangedOrSameAs(self->buffer) &&
- mapped_region.UnchangedOrSameAs(self->region)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
- }
- ffi::ObjectPtr<BufferRegionNode> copy = ffi::make_object<BufferRegionNode>(*self);
- copy->buffer = std::move(mapped_buffer).ValueOrUnchanged(std::move(copy->buffer));
- copy->region = std::move(mapped_region).ValueOrUnchanged(std::move(copy->region));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
-}
-
-TVMFFIAny BufferRegionMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
- BufferRegionNode* self = const_cast<BufferRegionNode*>(
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BufferRegionNode>(value));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
- ffi::UnchangedOr<BufferVar>, mapped_buffer,
- mutator->MutateExpected(self->buffer, ffi::InplaceMode::kAllow));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
- ffi::UnchangedOr<ffi::Array<Range>>, mapped_region,
- mutator->MutateExpected(self->region, ffi::InplaceMode::kAllow));
- if (mapped_buffer.UnchangedOrSameAs(self->buffer) &&
- mapped_region.UnchangedOrSameAs(self->region)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
- }
- if (!mapped_buffer.IsUnchanged()) self->buffer = std::move(mapped_buffer).ValueUnchecked();
- if (!mapped_region.IsUnchanged()) self->region = std::move(mapped_region).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
-}
-
TVMFFIAny MatchBufferRegionVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const MatchBufferRegionNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const MatchBufferRegionNode>(
@@ -773,7 +729,7 @@ TVMFFIAny MatchBufferRegionMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyVi
mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() {
return mutator->MutateExpected(self->buffer);
}));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<BufferRegion>, mapped_source,
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<TensorRegion>, mapped_source,
mutator->MutateExpected(self->source));
if (mapped_buffer.UnchangedOrSameAs(self->buffer) &&
mapped_source.UnchangedOrSameAs(self->source)) {
@@ -796,7 +752,7 @@ TVMFFIAny MatchBufferRegionMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator
ffi::InplaceMode::kAllow);
}));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
- ffi::UnchangedOr<BufferRegion>, mapped_source,
+ ffi::UnchangedOr<TensorRegion>, mapped_source,
mutator->MutateExpected(self->source, ffi::InplaceMode::kAllow));
if (mapped_buffer.UnchangedOrSameAs(self->buffer) &&
mapped_source.UnchangedOrSameAs(self->source)) {
@@ -829,9 +785,9 @@ TVMFFIAny SBlockMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) n
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SBlockNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<IterVar>>, mapped_iter_vars,
mutator->MutateExpected(self->iter_vars));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<BufferRegion>>, mapped_reads,
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<TensorRegion>>, mapped_reads,
mutator->MutateExpected(self->reads));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<BufferRegion>>, mapped_writes,
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<TensorRegion>>, mapped_writes,
mutator->MutateExpected(self->writes));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<BufferVar>>, mapped_alloc_buffers,
mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() {
@@ -878,10 +834,10 @@ TVMFFIAny SBlockMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
ffi::UnchangedOr<ffi::Array<IterVar>>, mapped_iter_vars,
mutator->MutateExpected(self->iter_vars, ffi::InplaceMode::kAllow));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<BufferRegion>>, mapped_reads,
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<TensorRegion>>, mapped_reads,
mutator->MutateExpected(self->reads, ffi::InplaceMode::kAllow));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
- ffi::UnchangedOr<ffi::Array<BufferRegion>>, mapped_writes,
+ ffi::UnchangedOr<ffi::Array<TensorRegion>>, mapped_writes,
mutator->MutateExpected(self->writes, ffi::InplaceMode::kAllow));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<BufferVar>>, mapped_alloc_buffers,
mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() {
@@ -1534,7 +1490,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Span span) { return BufferStore(buffer, value, indices, span); });
}
-// BufferRegion
+// TensorRegion
BufferRegionType::BufferRegionType() : Type(ffi::UnsafeInit{}) {
static ffi::ObjectPtr<BufferRegionTypeNode> singleton = ffi::make_object<BufferRegionTypeNode>();
data_ = singleton;
@@ -1553,33 +1509,20 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("tirx.BufferRegionType", []() { return BufferRegionType(); });
}
-BufferRegion::BufferRegion(BufferVar buffer, ffi::Array<Range> region, Span span) {
+TensorRegion 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);
+ << "Buffer rank and region dimension mismatch";
+ return TensorRegion(std::move(buffer), std::move(region), BufferRegionType(), std::move(span));
}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- BufferRegionNode::RegisterReflection();
- refl::TypeAttrDef<BufferRegionNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BufferRegionVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BufferRegionMutate))
- .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BufferRegionMaybeInplaceMutate));
-
refl::GlobalDef().def("tirx.BufferRegion", [](BufferVar buffer, ffi::Array<Range> region) {
return BufferRegion(buffer, region);
});
}
-BufferRegion BufferRegion::FullRegion(BufferVar buffer) {
+TensorRegion FullBufferRegion(BufferVar buffer) {
ffi::Array<Range> region;
for (PrimExpr extent : buffer->shape) {
region.push_back(Range::FromMinExtent(0, extent));
@@ -1587,7 +1530,7 @@ BufferRegion BufferRegion::FullRegion(BufferVar buffer) {
return BufferRegion(buffer, region);
}
-BufferRegion BufferRegion::FromPoint(BufferVar buffer, ffi::Array<PrimExpr> indices) {
+TensorRegion BufferRegionFromPoint(BufferVar buffer, ffi::Array<PrimExpr> indices) {
ffi::Array<Range> region;
for (const PrimExpr& index : indices) {
if (const prim::RampNode* ramp_index = index.as<prim::RampNode>()) {
@@ -1601,8 +1544,10 @@ BufferRegion BufferRegion::FromPoint(BufferVar buffer, ffi::Array<PrimExpr> indi
}
// MatchBufferRegion
-MatchBufferRegion::MatchBufferRegion(BufferVar buffer, BufferRegion source) {
- const BufferVar& source_buffer = source->buffer;
+MatchBufferRegion::MatchBufferRegion(BufferVar buffer, TensorRegion source) {
+ const BufferVar& source_buffer = source->source.as_or_throw<BufferVar>();
+ TVM_FFI_ICHECK_EQ(source_buffer->shape.size(), source->region.size())
+ << "MatchBufferRegion source must match its buffer rank";
arith::Analyzer analyzer;
// Check scope and dtype
TVM_FFI_ICHECK_EQ(buffer.scope(), source_buffer.scope())
@@ -1654,17 +1599,24 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
reinterpret_cast<void*>(&MatchBufferRegionMaybeInplaceMutate));
- refl::GlobalDef().def("tirx.MatchBufferRegion", [](BufferVar buffer, BufferRegion source) {
+ refl::GlobalDef().def("tirx.MatchBufferRegion", [](BufferVar buffer, TensorRegion source) {
return MatchBufferRegion(buffer, source);
});
}
// Block
-SBlock::SBlock(ffi::Array<IterVar> iter_vars, ffi::Array<BufferRegion> reads,
- ffi::Array<BufferRegion> writes, ffi::String name_hint, Stmt body,
+SBlock::SBlock(ffi::Array<IterVar> iter_vars, ffi::Array<TensorRegion> reads,
+ ffi::Array<TensorRegion> writes, ffi::String name_hint, Stmt body,
ffi::Optional<Stmt> init, ffi::Array<BufferVar> alloc_buffers,
ffi::Array<MatchBufferRegion> match_buffers, ffi::Map<ffi::String, Any> annotations,
Span span) {
+ for (const auto& regions : {reads, writes}) {
+ for (const TensorRegion& region : regions) {
+ const auto buffer = region->source.as_or_throw<BufferVar>();
+ TVM_FFI_ICHECK_EQ(buffer->shape.size(), region->region.size())
+ << "SBlock region must match its buffer rank";
+ }
+ }
ffi::ObjectPtr<SBlockNode> node = ffi::make_object<SBlockNode>();
node->iter_vars = std::move(iter_vars);
node->reads = std::move(reads);
@@ -1704,8 +1656,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
reinterpret_cast<void*>(&SBlockMaybeInplaceMutate));
refl::GlobalDef().def("tirx.SBlock",
- [](ffi::Array<IterVar> iter_vars, ffi::Array<BufferRegion> reads,
- ffi::Array<BufferRegion> writes, ffi::String name_hint, Stmt body,
+ [](ffi::Array<IterVar> iter_vars, ffi::Array<TensorRegion> reads,
+ ffi::Array<TensorRegion> writes, ffi::String name_hint, Stmt body,
ffi::Optional<Stmt> init, ffi::Array<BufferVar> alloc_buffers,
ffi::Array<MatchBufferRegion> match_buffers,
ffi::Map<ffi::String, Any> annotations, Span span) {
diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc
index 3dba665700..3bce83fe86 100644
--- a/src/tirx/ir/stmt_functor.cc
+++ b/src/tirx/ir/stmt_functor.cc
@@ -60,7 +60,6 @@ void StmtExprVisitor::InitVTable(VTable* vtable) {
SetDispatch<StmtExprVisitor, SBlockRealizeNode>(vtable);
SetDispatch<StmtExprVisitor, ScopeIdDefStmtNode>(vtable);
SetDispatch<StmtExprVisitor, TilePrimitiveCallNode>(vtable);
- SetDispatch<StmtExprVisitor, BufferRegionNode>(vtable);
}
ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const VarNode* op) { return std::nullopt; }
@@ -83,15 +82,6 @@ ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const TensorLoadNode* op)
return std::nullopt;
}
-ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const BufferRegionNode* op) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->buffer));
- for (const auto& range : op->region) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(range->min));
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(range->extent));
- }
- return std::nullopt;
-}
-
ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const TupleNode* op) {
for (const auto& child : op->fields) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(child));
@@ -249,10 +239,10 @@ ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const SBlockNode* op) {
this->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() { return this->Visit(buf); }));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(VisitBufferMetadata(buf));
}
- for (const BufferRegion& region : op->reads) {
+ for (const TensorRegion& region : op->reads) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(region));
}
- for (const BufferRegion& region : op->writes) {
+ for (const TensorRegion& region : op->writes) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(region));
}
for (const MatchBufferRegion& match_buffer_region : op->match_buffers) {
@@ -295,7 +285,7 @@ ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const TilePrimitiveCallNod
std::function<ffi::Optional<VisitInterrupt>(const ffi::Any&)> fvisit;
fvisit = [this, &fvisit](const ffi::Any& e) -> ffi::Optional<VisitInterrupt> {
if (e == nullptr) return std::nullopt;
- if (auto buffer_region = e.as<BufferRegion>()) {
+ if (auto buffer_region = e.as<TensorRegion>()) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(buffer_region.value()));
} else if (auto var = e.as<Var>(); var && var.value()->ty.as<BufferTypeNode>()) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(BufferVar(var.value())));
@@ -339,7 +329,6 @@ void StmtExprMutator::InitVTable(VTable* vtable) {
SetDispatch<StmtExprMutator, SBlockRealizeNode>(vtable);
SetDispatch<StmtExprMutator, ScopeIdDefStmtNode>(vtable);
SetDispatch<StmtExprMutator, TilePrimitiveCallNode>(vtable);
- SetDispatch<StmtExprMutator, BufferRegionNode>(vtable);
}
UnchangedOr<Stmt> StmtExprMutator::Mutate_(const BindNode* op, InplaceMode inplace_mode) {
@@ -597,9 +586,9 @@ UnchangedOr<Stmt> StmtExprMutator::Mutate_(const SBlockNode* op, InplaceMode inp
auto alloc_buffers = WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] {
return Mutate(op->alloc_buffers, inplace_mode);
}).as_or_throw<UnchangedOr<ffi::Array<BufferVar>>>();
- auto reads = Mutate(op->reads, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<BufferRegion>>>();
+ auto reads = Mutate(op->reads, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<TensorRegion>>>();
auto writes =
- Mutate(op->writes, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<BufferRegion>>>();
+ Mutate(op->writes, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<TensorRegion>>>();
auto match_buffers = Mutate(op->match_buffers, inplace_mode)
.as_or_throw<UnchangedOr<ffi::Array<MatchBufferRegion>>>();
auto init = Mutate(op->init, inplace_mode).as_or_throw<UnchangedOr<ffi::Optional<Stmt>>>();
@@ -633,24 +622,6 @@ UnchangedOr<Stmt> StmtExprMutator::Mutate_(const SBlockNode* op, InplaceMode inp
return Stmt(std::move(copy));
}
-UnchangedOr<Expr> StmtExprMutator::Mutate_(const BufferRegionNode* op, InplaceMode inplace_mode) {
- auto buffer = Mutate(op->buffer, inplace_mode).as_or_throw<UnchangedOr<BufferVar>>();
- auto region = Mutate(op->region, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<Range>>>();
- if (buffer.UnchangedOrSameAs(op->buffer) && region.UnchangedOrSameAs(op->region)) {
- return ffi::Unchanged();
- }
- if (inplace_mode == InplaceMode::kAllow) {
- auto* writable = const_cast<BufferRegionNode*>(op);
- if (!buffer.IsUnchanged()) writable->buffer = std::move(buffer).ValueUnchecked();
- if (!region.IsUnchanged()) writable->region = std::move(region).ValueUnchecked();
- return ffi::Unchanged();
- }
- auto copy = ffi::make_object<BufferRegionNode>(*op);
- if (!buffer.IsUnchanged()) copy->buffer = std::move(buffer).ValueUnchecked();
- if (!region.IsUnchanged()) copy->region = std::move(region).ValueUnchecked();
- return Expr(std::move(copy));
-}
-
UnchangedOr<Stmt> StmtExprMutator::Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode) {
return detail::MutateSeqStmt(op, inplace_mode, [this](ffi::AnyView element, InplaceMode mode) {
return Mutate(element, mode).as_or_throw<UnchangedOr<Stmt>>();
@@ -709,7 +680,7 @@ UnchangedOr<Stmt> StmtExprMutator::Mutate_(const TilePrimitiveCallNode* op,
InplaceMode inplace_mode) {
std::function<UnchangedOr<ffi::Any>(ffi::AnyView, InplaceMode)> mutate_arg;
mutate_arg = [&](ffi::AnyView value, InplaceMode mode) -> UnchangedOr<ffi::Any> {
- if (value.as<BufferRegionNode>()) {
+ if (value.as<TensorRegionNode>()) {
return Mutate(value, mode);
}
if (const auto* var = value.as<VarNode>(); var && var->ty.as<BufferTypeNode>()) {
diff --git a/src/tirx/ir/tir_visitor_with_path.cc b/src/tirx/ir/tir_visitor_with_path.cc
index d0e86972f0..53fe95f12c 100644
--- a/src/tirx/ir/tir_visitor_with_path.cc
+++ b/src/tirx/ir/tir_visitor_with_path.cc
@@ -150,8 +150,12 @@ void TIRVisitorWithPath::VisitBufferDef(const BufferVar& buffer, AccessPath path
// are allocated in a different scope than where they are used.
void TIRVisitorWithPath::VisitBufferUse(const BufferVar& buffer, AccessPath path) {}
-void TIRVisitorWithPath::Visit(const BufferRegion& region, AccessPath path) {
- VisitBufferUse(region->buffer, path->Attr("buffer"));
+void TIRVisitorWithPath::Visit(const TensorRegion& region, AccessPath path) {
+ if (auto buffer = region->source.as<BufferVar>()) {
+ VisitBufferUse(buffer.value(), path->Attr("source"));
+ } else {
+ Visit(region->source, path->Attr("source"));
+ }
Visit(region->region, path->Attr("region"));
}
@@ -320,7 +324,7 @@ void TIRVisitorWithPath::VisitStmt_(const tirx::TilePrimitiveCallNode* op, Acces
if (op->args[i] == nullptr) {
continue;
}
- if (auto buf_region = op->args[i].as<BufferRegion>()) {
+ if (auto buf_region = op->args[i].as<TensorRegion>()) {
Visit(buf_region.value(), path->Attr("args")->ArrayItem(i));
} else if (auto expr = op->args[i].as<PrimExpr>()) {
Visit(expr.value(), path->Attr("args")->ArrayItem(i));
@@ -357,8 +361,8 @@ void TIRVisitorWithPath::Dispatch_(const TensorLoadNode* op, AccessPath path) {
Visit(op->indices, path->Attr("indices"));
}
-void TIRVisitorWithPath::Dispatch_(const BufferRegionNode* op, AccessPath path) {
- Visit(ffi::GetRef<BufferRegion>(op), path);
+void TIRVisitorWithPath::Dispatch_(const TensorRegionNode* op, AccessPath path) {
+ Visit(ffi::GetRef<TensorRegion>(op), path);
}
void TIRVisitorWithPath::Dispatch_(const OpaqueExprNode* op, AccessPath path) {}
diff --git a/src/tirx/ir/tir_visitor_with_path.h b/src/tirx/ir/tir_visitor_with_path.h
index c1bcc07a8c..ef54a7c35e 100644
--- a/src/tirx/ir/tir_visitor_with_path.h
+++ b/src/tirx/ir/tir_visitor_with_path.h
@@ -67,7 +67,7 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const Expr&, ffi::reflecti
Dispatch_(tuple, path);
} else if (auto* tuple_get_item = obj.as<TupleGetItemNode>()) {
Dispatch_(tuple_get_item, path);
- } else if (auto* buffer_region = obj.as<BufferRegionNode>()) {
+ } else if (auto* buffer_region = obj.as<TensorRegionNode>()) {
Dispatch_(buffer_region, path);
} else if (obj.as<OpaqueExprNode>()) {
Dispatch(obj, path);
@@ -92,7 +92,7 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const Expr&, ffi::reflecti
virtual void Visit(const PrimFunc& obj, ffi::reflection::AccessPath path);
virtual void Visit(const GlobalVar& obj, ffi::reflection::AccessPath path) {}
virtual void Visit(const Range& obj, ffi::reflection::AccessPath path);
- virtual void Visit(const BufferRegion& obj, ffi::reflection::AccessPath path);
+ virtual void Visit(const TensorRegion& obj, ffi::reflection::AccessPath path);
virtual void Visit(const MatchBufferRegion& obj, ffi::reflection::AccessPath path);
virtual void Visit(const IterVar& obj, ffi::reflection::AccessPath path);
@@ -155,7 +155,7 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const Expr&, ffi::reflecti
using ExprFunctor::Dispatch;
void Dispatch_(const VarNode* op, ffi::reflection::AccessPath path) override;
void Dispatch_(const TensorLoadNode* op, ffi::reflection::AccessPath path) override;
- void Dispatch_(const BufferRegionNode* op, ffi::reflection::AccessPath path) override;
+ void Dispatch_(const TensorRegionNode* op, ffi::reflection::AccessPath path) override;
void Dispatch_(const OpaqueExprNode* op, ffi::reflection::AccessPath path) override;
void Dispatch_(const TupleNode* op, ffi::reflection::AccessPath path) override;
void Dispatch_(const TupleGetItemNode* op, ffi::reflection::AccessPath path) override;
diff --git a/src/tirx/ir/tirx_stmt.cc b/src/tirx/ir/tirx_stmt.cc
index ef92f11639..be57749739 100644
--- a/src/tirx/ir/tirx_stmt.cc
+++ b/src/tirx/ir/tirx_stmt.cc
@@ -43,6 +43,14 @@ TilePrimitiveCall::TilePrimitiveCall(tvm::Op op, ffi::Array<ffi::Any> args,
static const auto& category_map = Op::GetAttrMap<TIRxOpCategory>("TIRxOpCategory");
TVM_FFI_ICHECK(category_map.get(op, ffi::String("")) == "tile_primitive")
<< "Only tile primitive ops can be used in tirx::TilePrimitiveCall";
+ ffi::StructuralVisit(args,
+ [](const TensorRegionNode* region,
+ ffi::StructuralVisitorObj*) -> ffi::Optional<ffi::VisitInterrupt> {
+ const auto buffer = region->source.as_or_throw<BufferVar>();
+ TVM_FFI_ICHECK_EQ(buffer->shape.size(), region->region.size())
+ << "Tile region must match its buffer rank";
+ return std::nullopt;
+ });
ffi::ObjectPtr<TilePrimitiveCallNode> n = ffi::make_object<TilePrimitiveCallNode>(
std::move(op), std::move(args), std::move(workspace), std::move(config), std::move(dispatch),
std::move(scope));
diff --git a/src/tirx/script/builder/frame.cc b/src/tirx/script/builder/frame.cc
index dcd198f7ed..031984e83d 100644
--- a/src/tirx/script/builder/frame.cc
+++ b/src/tirx/script/builder/frame.cc
@@ -213,10 +213,9 @@ void SBlockFrameNode::ExitWithScope() {
if (int detect_access = (!reads.has_value()) | (!writes.has_value() << 1)) {
attrs.Set("tirx.script_parsing_detect_access", tvm::IntImm::Int64(detect_access));
}
- tvm::tirx::SBlock block(iter_vars, reads.value_or(ffi::Array<tvm::tirx::BufferRegion>()),
- writes.value_or(ffi::Array<tvm::tirx::BufferRegion>()), name,
- AsStmt(stmts), init, tir_alloc_buffers, match_buffers, attrs,
- tvm::Span());
+ tvm::tirx::SBlock block(iter_vars, reads.value_or(ffi::Array<tvm::TensorRegion>()),
+ writes.value_or(ffi::Array<tvm::TensorRegion>()), name, AsStmt(stmts),
+ init, tir_alloc_buffers, match_buffers, attrs, tvm::Span());
if (no_realize) {
TVM_FFI_CHECK(iter_values.empty(), ValueError)
<< "Block bindings are not allowed when `no_realize=True`";
diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc
index 8c57c0b4c4..53bab5646b 100644
--- a/src/tirx/script/builder/ir.cc
+++ b/src/tirx/script/builder/ir.cc
@@ -159,10 +159,10 @@ BufferVar MatchBuffer(ffi::ObjectRef param, ffi::Array<PrimExpr> shape, PrimType
SBlockFrame frame = FindSBlockFrame("T.match_buffer");
frame->match_buffers.push_back(tvm::tirx::MatchBufferRegion(
buffer, BufferRegionFromLoad(ffi::GetRef<tvm::TensorLoad>(buffer_load))));
- } else if (const auto* buffer_region = param.as<tvm::tirx::BufferRegionNode>()) {
+ } else if (const auto* buffer_region = param.as<tvm::TensorRegionNode>()) {
SBlockFrame frame = FindSBlockFrame("T.match_buffer");
frame->match_buffers.push_back(
- tvm::tirx::MatchBufferRegion(buffer, ffi::GetRef<tvm::tirx::BufferRegion>(buffer_region)));
+ tvm::tirx::MatchBufferRegion(buffer, ffi::GetRef<tvm::TensorRegion>(buffer_region)));
} else {
TVM_FFI_THROW(InternalError) << "ValueError: Unexpected type for TIR MatchBuffer.";
}
@@ -292,9 +292,9 @@ void Reads(ffi::Array<ffi::ObjectRef> buffer_slices) {
TVM_FFI_THROW(InternalError)
<< "ValueError: Duplicate read region declaration, previous one is " << frame->reads;
}
- ffi::Array<BufferRegion> reads;
+ ffi::Array<TensorRegion> reads;
for (const ffi::ObjectRef& obj : buffer_slices) {
- if (auto buffer_region = obj.as<BufferRegion>()) {
+ if (auto buffer_region = obj.as<TensorRegion>()) {
reads.push_back(buffer_region.value());
} else if (auto buffer_load = obj.as<TensorLoad>()) {
reads.push_back(BufferRegionFromLoad(buffer_load.value()));
@@ -312,9 +312,9 @@ void Writes(ffi::Array<ffi::ObjectRef> buffer_slices) {
TVM_FFI_THROW(InternalError)
<< "ValueError: Duplicate write region declaration, previous one is " << frame->writes;
}
- ffi::Array<BufferRegion> writes;
+ ffi::Array<TensorRegion> writes;
for (const ffi::ObjectRef& obj : buffer_slices) {
- if (auto buffer_region = obj.as<BufferRegion>()) {
+ if (auto buffer_region = obj.as<TensorRegion>()) {
writes.push_back(buffer_region.value());
} else if (auto buffer_load = obj.as<TensorLoad>()) {
writes.push_back(BufferRegionFromLoad(buffer_load.value()));
diff --git a/src/tirx/script/builder/utils.h b/src/tirx/script/builder/utils.h
index 4175e80dfa..340e4261dd 100644
--- a/src/tirx/script/builder/utils.h
+++ b/src/tirx/script/builder/utils.h
@@ -127,11 +127,11 @@ inline IfFrame FindIfFrame(const ffi::String& method) {
}
/*!
- * \brief Convert TensorLoad to BufferRegion.
+ * \brief Convert TensorLoad to TensorRegion.
* \param buffer_load The BufferLoad.
- * \return The converted BufferRegion.
+ * \return The converted TensorRegion.
*/
-inline tvm::tirx::BufferRegion BufferRegionFromLoad(tvm::TensorLoad buffer_load) {
+inline tvm::TensorRegion BufferRegionFromLoad(tvm::TensorLoad buffer_load) {
ffi::Array<Range> ranges;
for (const PrimExpr& index : buffer_load->indices) {
ranges.push_back(Range::FromMinExtent(index, IntImm(index.ty(), 1)));
diff --git a/src/tirx/script/printer/buffer.cc b/src/tirx/script/printer/buffer.cc
index cf9ba72937..a69b018f57 100644
--- a/src/tirx/script/printer/buffer.cc
+++ b/src/tirx/script/printer/buffer.cc
@@ -393,9 +393,9 @@ ffi::Array<Doc> BufferSlices(const ffi::Array<Range>& region, const AccessPath&
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<tirx::BufferRegion>(
- "", [](tirx::BufferRegion buffer_region, AccessPath p, IRDocsifier d) -> Doc {
- ExprDoc prefix = d->AsDoc<ExprDoc>(buffer_region->buffer, p->Attr("buffer"));
+ IRDocsifier::vtable().set_dispatch<tvm::TensorRegion>(
+ "", [](tvm::TensorRegion buffer_region, AccessPath p, IRDocsifier d) -> Doc {
+ ExprDoc prefix = d->AsDoc<ExprDoc>(buffer_region->source, p->Attr("source"));
return prefix[BufferSlices(buffer_region->region, p->Attr("region"), d)];
});
}
@@ -586,7 +586,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_SCRIPT_REPR(tirx::BufferRegionNode, ReprPrintTIR);
+TVM_SCRIPT_REPR(tvm::TensorRegionNode, ReprPrintTIR);
TVM_SCRIPT_REPR(TensorLoadNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::BufferStoreNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::BufferTypeNode, ReprPrintTIR);
diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc
index 7e1e26b581..989bb2beaf 100644
--- a/src/tirx/script/printer/stmt.cc
+++ b/src/tirx/script/printer/stmt.cc
@@ -123,10 +123,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// and args[0]/args[1] refer to the same buffer region, collapse to 1 arg
bool inplace_unary = false;
if (n_args == 2) {
- auto dst_opt = op_call->args[0].as<tirx::BufferRegion>();
- auto src_opt = op_call->args[1].as<tirx::BufferRegion>();
+ auto dst_opt = op_call->args[0].as<tvm::TensorRegion>();
+ auto src_opt = op_call->args[1].as<tvm::TensorRegion>();
if (dst_opt.has_value() && src_opt.has_value() &&
- dst_opt.value()->buffer.same_as(src_opt.value()->buffer) &&
+ dst_opt.value()->source.same_as(src_opt.value()->source) &&
StructuralEqual()(dst_opt.value()->region, src_opt.value()->region)) {
inplace_unary = true;
}
@@ -346,7 +346,7 @@ ffi::Optional<ExprDoc> TryDeclBufferSugarWithParent(const tirx::BufferVar& child
// NOTE: an earlier sugar printed rank-preserving aliases with a different
// elem_offset as ``parent[slices]``. That print is not roundtrippable: it
- // reparses as a BufferRegion, not a Buffer, so any later Buffer use of the
+ // reparses as a TensorRegion, not a Buffer, so any later Buffer use of the
// alias (stores, views) breaks. Such aliases now print as plain
// T.decl_buffer, which reparses exactly.
diff --git a/src/tirx/transform/flatten_buffer.cc b/src/tirx/transform/flatten_buffer.cc
index 96119a39e9..7baad9cc3a 100644
--- a/src/tirx/transform/flatten_buffer.cc
+++ b/src/tirx/transform/flatten_buffer.cc
@@ -191,14 +191,14 @@ class BufferFlattener : public IRMutatorWithAnalyzer {
block.CopyOnWrite()->alloc_buffers = alloc_buffers;
}
- ffi::Array<BufferRegion> reads = op->reads;
- reads.MutateByApply([this](BufferRegion region) { return MutateBufferRegion(region); });
+ ffi::Array<TensorRegion> reads = op->reads;
+ reads.MutateByApply([this](TensorRegion region) { return MutateBufferRegion(region); });
if (!reads.same_as(op->reads)) {
block.CopyOnWrite()->reads = reads;
}
- ffi::Array<BufferRegion> writes = op->writes;
- writes.MutateByApply([this](BufferRegion region) { return MutateBufferRegion(region); });
+ ffi::Array<TensorRegion> writes = op->writes;
+ writes.MutateByApply([this](TensorRegion region) { return MutateBufferRegion(region); });
if (!writes.same_as(op->writes)) {
block.CopyOnWrite()->writes = writes;
}
@@ -323,9 +323,9 @@ class BufferFlattener : public IRMutatorWithAnalyzer {
return BufferLoad(info.flattened, FoldIndices(info, node->indices), node->span);
}
- BufferRegion MutateBufferRegion(BufferRegion region) {
- const FlatInfo& info = Lookup(region->buffer);
- if (info.flattened.same_as(region->buffer)) {
+ TensorRegion MutateBufferRegion(TensorRegion region) {
+ const FlatInfo& info = Lookup(region->source.as_or_throw<tvm::tirx::BufferVar>());
+ if (info.flattened.same_as(region->source.as_or_throw<tvm::tirx::BufferVar>())) {
return region;
}
diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc
index a05311ca5c..121b12ba8a 100644
--- a/src/tirx/transform/ir_utils.cc
+++ b/src/tirx/transform/ir_utils.cc
@@ -270,9 +270,9 @@ class IRConvertSSA final : public StmtExprMutator {
}
return iter_var;
});
- ffi::Array<BufferRegion> reads =
+ ffi::Array<TensorRegion> reads =
block->reads.Map([&](const auto& region) { return VisitBufferAccess(region); });
- ffi::Array<BufferRegion> writes =
+ ffi::Array<TensorRegion> writes =
block->writes.Map([&](const auto& region) { return VisitBufferAccess(region); });
if (!reads.same_as(block->reads) || !writes.same_as(block->writes) ||
@@ -301,6 +301,15 @@ class IRConvertSSA final : public StmtExprMutator {
return node;
}
+ TensorRegion VisitBufferAccess(TensorRegion node) {
+ BufferVar buffer = node->source.as_or_throw<BufferVar>();
+ BufferVar new_buf = GetRemappedBuffer(buffer);
+ if (!new_buf.same_as(buffer)) {
+ node.CopyOnWrite()->source = new_buf;
+ }
+ return node;
+ }
+
TensorLoad VisitBufferAccess(TensorLoad node) {
BufferVar buffer = node->source.as_or_throw<BufferVar>();
BufferVar new_buf = GetRemappedBuffer(buffer);
@@ -778,7 +787,7 @@ ffi::Array<PrimExpr> GetBufferAllocationShape(const BufferVar& buffer) {
ffi::Array<PrimExpr> ConvertIndices(const MatchBufferRegion& match_buffer,
const ffi::Array<PrimExpr>& indices) {
const BufferVar& target = match_buffer->buffer;
- const BufferRegion& source = match_buffer->source;
+ const TensorRegion& source = match_buffer->source;
TVM_FFI_ICHECK_EQ(indices.size(), target->shape.size());
arith::Analyzer analyzer;
@@ -800,7 +809,7 @@ ffi::Array<PrimExpr> ConvertIndices(const MatchBufferRegion& match_buffer,
Region ConvertRegion(const MatchBufferRegion& match_buffer, const Region& region) {
const BufferVar& target = match_buffer->buffer;
- const BufferRegion& source = match_buffer->source;
+ const TensorRegion& source = match_buffer->source;
TVM_FFI_ICHECK_EQ(region.size(), target->shape.size());
arith::Analyzer analyzer;
@@ -847,7 +856,8 @@ class StorageAlignCollector : public StmtExprVisitor {
auto storage_align_annotation = (*it).second.as_or_throw<StorageAlignAnnotation>();
for (const auto& storage_align_tuple : storage_align_annotation) {
int buffer_index = storage_align_tuple.get<0>();
- const BufferVar& buffer = op->writes[buffer_index]->buffer;
+ const BufferVar& buffer =
+ op->writes[buffer_index]->source.as_or_throw<tvm::tirx::BufferVar>();
storage_align_[buffer.var()].push_back(storage_align_tuple);
}
}
diff --git a/src/tirx/transform/vectorize_loop.cc b/src/tirx/transform/vectorize_loop.cc
index b93e2ccd52..854f3e75bc 100644
--- a/src/tirx/transform/vectorize_loop.cc
+++ b/src/tirx/transform/vectorize_loop.cc
@@ -493,7 +493,7 @@ class Vectorizer : public StmtExprMutator {
TVM_FFI_UNREACHABLE();
}
- UnchangedOr<Expr> Mutate_(const BufferRegionNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Expr> Mutate_(const TensorRegionNode* op, InplaceMode inplace_mode) final {
TVM_FFI_THROW(InternalError) << "Do not have a default for " << op->GetTypeKey();
TVM_FFI_UNREACHABLE();
}
diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc
index 7960f4b4e5..abbdf56035 100644
--- a/tests/cpp/ir_functor_test.cc
+++ b/tests/cpp/ir_functor_test.cc
@@ -231,7 +231,7 @@ TEST(IRF, StmtVisitor) {
tirx::Var buf_var("b", PointerType(dtype));
BufferVar buffer = decl_buffer({16});
body = SeqStmt({DeclBuffer(buffer, buf_var), std::move(body)});
- BufferRegion buffer_region(buffer, {Range::FromMinExtent(x + 1, 1)});
+ TensorRegion buffer_region = BufferRegion(buffer, {Range::FromMinExtent(x + 1, 1)});
MatchBufferRegion match_buffer_region(decl_buffer({1}), buffer_region);
// construct block and block_realize
@@ -359,7 +359,7 @@ TEST(IRF, StmtExprMutator) {
Stmt alloc = fmakealloc();
// body is: DeclBuffer, AllocBuffer, Evaluate
Stmt body = SeqStmt({decl, alloc, eval_body});
- BufferRegion buffer_region(buffer, {Range::FromMinExtent(x + 1, 1)});
+ TensorRegion buffer_region = BufferRegion(buffer, {Range::FromMinExtent(x + 1, 1)});
MatchBufferRegion match_buffer_region(decl_buffer({1}), buffer_region);
// construct block and block_realize
SBlock block = SBlock({}, {buffer_region}, {buffer_region}, "block", body, body, {},
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py b/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py
index acd33f9814..7da451701d 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py
@@ -22,9 +22,10 @@ import pytest
import tvm
import tvm.testing
+from tvm.ir import TensorRegion
from tvm.script import ir as I
from tvm.script import tirx as T
-from tvm.tirx import BufferRegion, StringImm
+from tvm.tirx import StringImm
identify_memcpy = tvm.s_tir.analysis._ffi_api._identify_memcpy
@@ -34,7 +35,7 @@ def _check_memcpy_results(func, expected):
results = identify_memcpy(func.body)
if isinstance(expected, str) or (
- isinstance(expected, tuple) and isinstance(expected[0], BufferRegion)
+ isinstance(expected, tuple) and isinstance(expected[0], TensorRegion)
):
expected = [expected]
diff --git a/tests/python/s_tir/analysis/test_sblock_access_region.py b/tests/python/s_tir/analysis/test_sblock_access_region.py
index 7218fad53f..268e621fed 100644
--- a/tests/python/s_tir/analysis/test_sblock_access_region.py
+++ b/tests/python/s_tir/analysis/test_sblock_access_region.py
@@ -363,7 +363,7 @@ def test_access_of_padding_pattern():
buffer_var_map = {buf: buf for buf in alloc_buffers}
def do_compare_buffer_region(region, expect):
- assert region.buffer == expect.buffer
+ assert region.source == expect.source
analyzer = tvm.arith.Analyzer()
for observed_range, expected_range in zip(region.region, expect.region):
analyzer.can_prove_equal(observed_range.min, expected_range.min)
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py
index 9cd1907cff..701aaa08bf 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py
@@ -310,7 +310,7 @@ def test_simplify():
)
actual_block = sch.get(block_outer)
actual = tvm.tirx.PrimFunc(
- [actual_block.reads[0].buffer, actual_block.writes[0].buffer],
+ [actual_block.reads[0].source, actual_block.writes[0].source],
sch.get(sch.get_loops(block_outer)[0]),
)
tvm.ir.assert_structural_equal(expected, actual)
@@ -1023,7 +1023,7 @@ def test_padded_transform_repeated_buffer_element():
B[vi] = A[vi // 4, vi % 4]
sch = tvm.s_tir.Schedule(Before)
- A = sch.get(sch.get_sblock("block")).reads[0].buffer
+ A = sch.get(sch.get_sblock("block")).reads[0].source
sch.transform_layout(
"block",
"A",
@@ -1052,7 +1052,7 @@ def test_pad_value_may_not_reference_other_buffer():
B[vi] = A[vi]
sch = tvm.s_tir.Schedule(Before)
- A = sch.get(sch.get_sblock("block")).reads[0].buffer
+ A = sch.get(sch.get_sblock("block")).reads[0].source
other = tirx.decl_buffer(1, A.ty.dtype, name="other")
with pytest.raises(tvm.s_tir.schedule.schedule.ScheduleError):
sch.transform_layout(
diff --git a/tests/python/tirx-base/test_tir_constructor.py b/tests/python/tirx-base/test_tir_constructor.py
index 88d0f36964..cca9089609 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.ir import TensorRegion
from tvm.ir.prim import expr_deep_equal
from tvm.script import tirx as T
@@ -230,14 +231,14 @@ def test_buffer_region_call_wrappers_reject():
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 isinstance(lhs, TensorRegion)
+ assert isinstance(rhs, TensorRegion)
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"):
+ with pytest.raises(TypeError, match="construct a TensorLoad explicitly"):
tvm.tirx.IterVar((0, 4), "i", tvm.tirx.IterVar.DataPar) + int_region
diff --git a/tests/python/tirx/test_hint.py b/tests/python/tirx/test_hint.py
index 6186c471e3..13c6d5fc21 100644
--- a/tests/python/tirx/test_hint.py
+++ b/tests/python/tirx/test_hint.py
@@ -21,7 +21,7 @@ import tvm_ffi
import tvm
import tvm.script
import tvm.testing
-from tvm.ir import assert_structural_equal
+from tvm.ir import TensorRegion, assert_structural_equal
from tvm.script import tirx as T
from tvm.tirx import AttrStmt
@@ -217,9 +217,8 @@ def test_hint_no_message():
# Should have "access" key but no "message" key
assert "access" in stmt.node
assert "message" not in stmt.node
- from tvm.tirx import BufferRegion
- assert isinstance(stmt.node["access"], BufferRegion)
+ assert isinstance(stmt.node["access"], TensorRegion)
found[0] = True
tvm_ffi.structural_walk(func.body, visit)
@@ -245,11 +244,10 @@ def test_hint_access_buffer_region():
assert isinstance(stmt.node, tvm.ir.Map)
assert str(stmt.node["message"]) == "partition"
assert "access" in stmt.node
- from tvm.tirx import BufferRegion
- assert isinstance(stmt.node["access"], BufferRegion)
+ assert isinstance(stmt.node["access"], TensorRegion)
br = stmt.node["access"]
- assert br.buffer.name == "A"
+ assert br.source.name == "A"
assert len(br.region) == 2
found[0] = True
diff --git a/tests/python/tirx/test_parser_printer.py b/tests/python/tirx/test_parser_printer.py
index 98e613d22c..4fc72f0414 100644
--- a/tests/python/tirx/test_parser_printer.py
+++ b/tests/python/tirx/test_parser_printer.py
@@ -22,7 +22,7 @@ import tvm_ffi
import tvm
import tvm.script
import tvm.testing
-from tvm.ir import PointerType, PrimType, assert_structural_equal
+from tvm.ir import PointerType, PrimType, TensorRegion, assert_structural_equal
from tvm.script import ir as I
from tvm.script import tirx as T
from tvm.script.tirx import tile as Tx
@@ -2191,8 +2191,6 @@ def test_buffer_chunk_ir():
into n equal chunks. chunk(spec)[picks] is the exact same BufferRegion as
the hand-written a*k:(a+1)*k slice — no reshape, no extra dim."""
- from tvm.tirx.stmt import BufferRegion
-
compose = T.ComposeLayout(3, 3, 3, T.TileLayout(T.S[(4, 512) : (512, 1)]))
A = tvm.tirx.decl_buffer(
(4, 8, 16), "float16", layout=tvm.tirx.layout.TileLayout(T.S[(4, 8, 16) : (128, 16, 1)])
@@ -2202,7 +2200,7 @@ def test_buffer_chunk_ir():
# chunk((None, None, 2))[:, :, 1] narrows dim 2 (extent 16) to chunk 1 of 2
# → [8:16] (k = 16 // 2 = 8); rank preserved, dims 0/1 pass through as ':'.
reg = A.chunk((None, None, 2))[:, :, 1]
- assert isinstance(reg, BufferRegion)
+ assert isinstance(reg, TensorRegion)
assert len(reg.region) == 3 # rank-preserving: no extra extent-1 chunk dim
assert (int(reg.region[2].min), int(reg.region[2].extent)) == (8, 8)
assert_structural_equal(reg, A[:, :, 8:16])
@@ -2265,12 +2263,11 @@ def test_buffer_view_dtype_ir():
def test_buffer_slice_region():
"""Verify A[slice] returns BufferRegion (not DeclBuffer)."""
- from tvm.tirx.stmt import BufferRegion
buf = tvm.tirx.decl_buffer((128, 64), "float16")
br = buf[32:64, 0:32]
- assert isinstance(br, BufferRegion)
- assert br.buffer.same_as(buf)
+ assert isinstance(br, TensorRegion)
+ assert br.source.same_as(buf)
assert int(br.region[0].extent) == 32
assert int(br.region[1].extent) == 32
@@ -2278,11 +2275,11 @@ def test_buffer_slice_region():
assert isinstance(load, tvm.ir.TensorLoad)
partial = buf[1]
- assert isinstance(partial, BufferRegion)
+ assert isinstance(partial, TensorRegion)
narrowed = br[4:12, 2:10]
- assert isinstance(narrowed, BufferRegion)
- assert narrowed.buffer.same_as(buf)
+ assert isinstance(narrowed, TensorRegion)
+ assert narrowed.source.same_as(buf)
assert [(int(dim.min), int(dim.extent)) for dim in narrowed.region] == [
(36, 8),
(2, 8),
@@ -2294,7 +2291,7 @@ def test_buffer_slice_region():
assert [int(index) for index in chained_load.indices] == [35, 4]
point_then_region = br[3]
- assert isinstance(point_then_region, BufferRegion)
+ assert isinstance(point_then_region, TensorRegion)
assert [(int(dim.min), int(dim.extent)) for dim in point_then_region.region] == [
(35, 1),
(0, 32),