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/10/07 15:43:51 UTC
(tvm) branch main updated: [REFACTOR][TIRx] Rename BufferStore to TensorStore (#20590)
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 b250f7f386 [REFACTOR][TIRx] Rename BufferStore to TensorStore (#20590)
b250f7f386 is described below
commit b250f7f3867beee05eacd53738793ad124821fb0
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Wed Oct 7 11:43:44 2026 -0400
[REFACTOR][TIRx] Rename BufferStore to TensorStore (#20590)
---
docs/arch/codegen.rst | 2 +-
docs/tirx/arch/lowering_pipeline.rst | 2 +-
include/tvm/relax/distributed/axis_group_graph.h | 2 +-
include/tvm/s_tir/analysis.h | 2 +-
.../tvm/s_tir/meta_schedule/feature_extractor.h | 4 +-
include/tvm/s_tir/schedule/schedule.h | 4 +-
include/tvm/tirx/script/ir_builder/ir.h | 2 +-
include/tvm/tirx/stmt.h | 30 ++---
include/tvm/tirx/stmt_functor.h | 10 +-
include/tvm/tirx/transform.h | 4 +-
include/tvm/tirx/type.h | 2 +-
python/tvm/backend/cuda/lang/clc.py | 2 +-
.../backend/cuda/tile_primitive/copy/fallback.py | 4 +-
python/tvm/relax/transform/legalize_ops/grad.py | 6 +-
python/tvm/s_tir/analysis/__init__.py | 2 +-
python/tvm/s_tir/dlight/analysis/gemv.py | 12 +-
python/tvm/s_tir/dlight/gpu/low_batch_gemv.py | 12 +-
python/tvm/s_tir/dlight/gpu/reduction.py | 12 +-
python/tvm/s_tir/dlight/gpu/rmsnorm.py | 6 +-
python/tvm/s_tir/dlight/gpu/transpose.py | 2 +-
.../feature_extractor/per_store_feature.py | 8 +-
python/tvm/s_tir/schedule/schedule.py | 6 +-
python/tvm/s_tir/tensor_intrin/cuda.py | 4 +-
python/tvm/tirx/__init__.py | 2 +-
.../tvm/tirx/script/ir_builder/parser_protocol.py | 12 +-
python/tvm/tirx/stmt.py | 10 +-
python/tvm/tirx/transform/transform.py | 4 +-
python/tvm/topi/gpu/scan.py | 32 +++---
python/tvm/topi/gpu/scatter_elements.py | 14 +--
python/tvm/topi/gpu/scatter_nd.py | 12 +-
python/tvm/topi/gpu/sort.py | 82 +++++++-------
python/tvm/topi/index_put.py | 6 +-
python/tvm/topi/scan.py | 8 +-
python/tvm/topi/scatter.py | 12 +-
python/tvm/topi/scatter_elements.py | 14 +--
python/tvm/topi/searchsorted.py | 10 +-
python/tvm/topi/signal.py | 24 ++--
python/tvm/topi/sparse_reshape.py | 38 +++----
python/tvm/topi/unique.py | 4 +-
python/tvm/topi/vision/nms.py | 122 ++++++++++-----------
python/tvm/topi/vision/nms_util.py | 14 +--
src/backend/cuda/transforms/lower_iket.cc | 8 +-
src/backend/opencl/runtime/opencl_common.h | 2 +-
src/backend/trn/codegen/codegen_trn.cc | 2 +-
src/backend/trn/codegen/codegen_trn.h | 2 +-
src/backend/trn/transform/lower_trainium_layout.cc | 10 +-
src/backend/vulkan/codegen/codegen_spirv.cc | 2 +-
src/backend/vulkan/codegen/codegen_spirv.h | 2 +-
src/backend/webgpu/codegen/codegen_webgpu.cc | 4 +-
src/backend/webgpu/codegen/codegen_webgpu.h | 2 +-
src/relax/analysis/layout_transformation.cc | 2 +-
src/relax/analysis/tir_op_pattern_kind.cc | 42 +++----
src/relax/backend/vm/vm_shape_lower.cc | 2 +-
.../transform/lower_global_view_to_local_view.cc | 2 +-
src/relax/transform/split_call_tir_by_pattern.cc | 8 +-
src/s_tir/analysis/domain_touched.cc | 2 +-
src/s_tir/analysis/estimate_flops.cc | 2 +-
src/s_tir/analysis/identify_memcpy.cc | 10 +-
src/s_tir/analysis/is_pure_function.cc | 2 +-
src/s_tir/analysis/oob_checker.cc | 2 +-
.../analysis/sblock_access_region_detector.cc | 4 +-
.../analysis/sblock_buffer_access_lca_detector.cc | 2 +-
src/s_tir/analysis/verify_gpu_code.cc | 2 +-
src/s_tir/backend/adreno/texture_flatten.cc | 4 +-
.../feature_extractor/per_store_feature.cc | 16 +--
.../postproc/disallow_async_strided_mem_copy.cc | 2 +-
.../postproc/rewrite_cooperative_fetch.cc | 2 +-
src/s_tir/schedule/analysis.h | 14 +--
src/s_tir/schedule/analysis/reducer.cc | 60 +++++-----
src/s_tir/schedule/ir_comparator.cc | 12 +-
src/s_tir/schedule/ir_comparator.h | 4 +-
src/s_tir/schedule/primitive.h | 4 +-
src/s_tir/schedule/primitive/block_annotate.cc | 6 +-
src/s_tir/schedule/primitive/cache_index.cc | 6 +-
src/s_tir/schedule/primitive/cache_read_write.cc | 32 +++---
src/s_tir/schedule/primitive/compute_inline.cc | 64 +++++------
src/s_tir/schedule/primitive/decompose_padding.cc | 10 +-
.../schedule/primitive/layout_transformation.cc | 46 ++++----
.../schedule/primitive/loop_transformation.cc | 2 +-
src/s_tir/schedule/primitive/pad_einsum.cc | 16 +--
src/s_tir/schedule/primitive/read_write_at.cc | 12 +-
src/s_tir/schedule/primitive/reduction.cc | 28 ++---
src/s_tir/schedule/primitive/rolling_buffer.cc | 8 +-
src/s_tir/schedule/transform.cc | 6 +-
src/s_tir/schedule/transform.h | 2 +-
src/s_tir/transform/compact_buffer_region.cc | 12 +-
src/s_tir/transform/inject_permuted_layout.cc | 4 +-
src/s_tir/transform/inject_ptx_async_copy.cc | 8 +-
src/s_tir/transform/inject_software_pipeline.cc | 6 +-
src/s_tir/transform/inject_virtual_thread.cc | 12 +-
.../transform/lower_cross_thread_reduction.cc | 14 +--
src/s_tir/transform/lower_match_buffer.cc | 8 +-
.../manifest_shared_memory_local_stage.cc | 16 +--
src/s_tir/transform/memhammer_coalesce.cc | 4 +-
.../transform/memhammer_intermediate_stage.cc | 8 +-
src/s_tir/transform/memhammer_lower_auto_copy.cc | 10 +-
.../transform/memhammer_tensorcore_rewrite.cc | 12 +-
.../transform/merge_shared_memory_allocations.cc | 6 +-
.../plan_update_buffer_allocation_location.cc | 2 +-
src/s_tir/transform/remove_store_undef.cc | 16 +--
.../remove_weight_layout_rewrite_block.cc | 4 +-
src/s_tir/transform/storage_access.cc | 2 +-
src/s_tir/transform/storage_access.h | 2 +-
src/s_tir/transform/transform_mma_buffer_layout.cc | 6 +-
.../transform/using_assume_to_reduce_branches.cc | 18 +--
src/script/printer/script_printer.cc | 2 +-
src/target/llvm/codegen_llvm.cc | 2 +-
src/target/llvm/codegen_llvm.h | 4 +-
src/target/source/codegen_c.cc | 2 +-
src/target/source/codegen_c.h | 2 +-
src/te/operation/create_primfunc.cc | 8 +-
src/tirx/analysis/verify_memory.cc | 2 +-
src/tirx/analysis/verify_well_formed.h | 4 +-
src/tirx/ir/buffer.cc | 2 +-
src/tirx/ir/data_type_rewriter.cc | 4 +-
src/tirx/ir/data_type_rewriter.h | 2 +-
src/tirx/ir/stmt.cc | 44 ++++----
src/tirx/ir/stmt_functor.cc | 12 +-
src/tirx/ir/tir_visitor_with_path.cc | 2 +-
src/tirx/ir/tir_visitor_with_path.h | 4 +-
src/tirx/script/ir_builder/ir.cc | 12 +-
src/tirx/script/printer/buffer.cc | 8 +-
src/tirx/script/printer/stmt.cc | 4 +-
src/tirx/transform/flatten_buffer.cc | 4 +-
src/tirx/transform/lower_thread_allreduce.h | 26 ++---
src/tirx/transform/lower_tirx_cleanup.cc | 6 +-
src/tirx/transform/lower_tvm_builtin.cc | 2 +-
src/tirx/transform/lower_warp_memory.cc | 6 +-
src/tirx/transform/narrow_datatype.cc | 2 +-
src/tirx/transform/remove_no_op.cc | 6 +-
src/tirx/transform/stmt_simplify.cc | 6 +-
src/tirx/transform/stmt_simplify.h | 2 +-
src/tirx/transform/storage_rewrite.cc | 36 +++---
src/tirx/transform/unroll_loop.cc | 2 +-
src/tirx/transform/unsupported_dtype_legalize.cc | 8 +-
src/tirx/transform/vectorize_loop.cc | 20 ++--
tests/cpp/ir_functor_test.cc | 4 +-
tests/python/codegen/test_target_codegen.py | 4 +-
tests/python/codegen/test_target_codegen_llvm.py | 2 +-
tests/python/codegen/test_target_codegen_vulkan.py | 2 +-
.../python/relax/test_runtime_builtin_rnn_state.py | 4 +-
.../test_transform_annotate_tir_op_pattern.py | 6 +-
tests/python/runtime/test_runtime_module_load.py | 2 +-
.../test_s_tir_analysis_identify_memcpy.py | 4 +-
.../test_s_tir_analysis_verify_block_scope.py | 4 +-
...schedule_feature_extractor_per_store_feature.py | 12 +-
.../s_tir/schedule/test_tir_schedule_rfactor.py | 42 +++----
.../s_tir/script/test_s_tir_script_basic_usage.py | 2 +-
.../s_tir/script/test_s_tir_script_blocks.py | 4 +-
.../script/test_s_tir_script_error_handling.py | 2 +-
.../s_tir/script/test_s_tir_script_printer.py | 2 +-
tests/python/s_tir/test_s_tir_renew_defs.py | 6 +-
...s_tir_transform_lower_cross_thread_reduction.py | 6 +-
.../test_s_tir_transform_mma_buffer_layout.py | 2 +-
.../transform/test_s_tir_transform_remove_undef.py | 2 +-
tests/python/tirx-base/test_tir_constructor.py | 4 +-
tests/python/tirx-base/test_tir_nodes.py | 16 +--
.../test_tir_transform_common_subexpr_elim.py | 8 +-
.../test_tir_transform_lower_intrin.py | 2 +-
.../test_tir_transform_lower_tvm_builtin.py | 4 +-
...est_tir_transform_pointer_value_type_rewrite.py | 10 +-
.../test_tir_transform_remove_no_op.py | 2 +-
.../tirx-transform/test_tir_transform_simplify.py | 10 +-
.../test_tir_transform_unroll_loop.py | 2 +-
.../tirx/script/test_tirx_script_error_handling.py | 12 +-
.../tirx/script/test_tirx_script_ir_builder.py | 22 ++--
.../script/test_tirx_script_meta_programming.py | 4 +-
.../python/tirx/script/test_tirx_script_printer.py | 4 +-
.../script/test_tirx_script_source_locations.py | 2 +-
tests/python/tirx/test_op.py | 4 +-
.../transform/test_transform_flatten_buffer.py | 4 +-
.../tirx/transform/test_transform_lower_tirx.py | 4 +-
172 files changed, 837 insertions(+), 837 deletions(-)
diff --git a/docs/arch/codegen.rst b/docs/arch/codegen.rst
index 47ee76d869..7f8eed6c37 100644
--- a/docs/arch/codegen.rst
+++ b/docs/arch/codegen.rst
@@ -151,7 +151,7 @@ backend (x86, ARM, NVPTX, AMDGPU, etc.).
arithmetic ops → LLVM binary instructions, ``TensorLoad`` → load with pointer arithmetic,
``Cast`` → LLVM type conversions, ``Call`` → intrinsic or extern function calls.
- **Statements** (``VisitStmt_``) emit LLVM IR side effects:
- ``BufferStore`` → store instructions, ``For`` → loop basic blocks with branches,
+ ``TensorStore`` → store instructions, ``For`` → loop basic blocks with branches,
``IfThenElse`` → conditional branches, ``tirx.alloc_tensor`` calls → stack or heap allocation.
The key methods on ``CodeGenLLVM`` are:
diff --git a/docs/tirx/arch/lowering_pipeline.rst b/docs/tirx/arch/lowering_pipeline.rst
index 9f4350c79c..a55b2121ec 100644
--- a/docs/tirx/arch/lowering_pipeline.rst
+++ b/docs/tirx/arch/lowering_pipeline.rst
@@ -64,7 +64,7 @@ The ``tirx_pipeline`` module pass applies this exact sequence (a few are gated b
- lowers remaining opaque constructs to lower-level TIRx forms
* - 4
- ``FlattenBuffer``
- - flattens multi-dimensional ``TensorLoad`` / ``BufferStore`` to 1-D
+ - flattens multi-dimensional ``TensorLoad`` / ``TensorStore`` to 1-D
* - 5
- ``BF16ComputeLegalize``
- rewrites ``bfloat16`` compute to a legal (f32-up-cast) form
diff --git a/include/tvm/relax/distributed/axis_group_graph.h b/include/tvm/relax/distributed/axis_group_graph.h
index c1e045576e..90e5d5ef4d 100644
--- a/include/tvm/relax/distributed/axis_group_graph.h
+++ b/include/tvm/relax/distributed/axis_group_graph.h
@@ -119,7 +119,7 @@ class BufferAxisGraphExtractor : public s_tir::StmtExprVisitor {
}
private:
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op));
buffer_access_indices_.push_back({op->buffer, op->indices});
diff --git a/include/tvm/s_tir/analysis.h b/include/tvm/s_tir/analysis.h
index f9e53c4528..80b75f8b0f 100644
--- a/include/tvm/s_tir/analysis.h
+++ b/include/tvm/s_tir/analysis.h
@@ -64,7 +64,7 @@ TVM_DLL ffi::Array<ffi::Array<TensorRegion>> GetSBlockReadWriteRegion(
/*!
* \brief Detect the lowest common ancestor(LCA) of buffer access, including both high-level
- * access(TensorLoad, BufferStore) and low-level access(Load, Store and opaque access).
+ * access(TensorLoad, TensorStore) and low-level access(Load, Store and opaque access).
* The LCA may be a For loop or a Block.
* \param func The PrimFunc to be detected.
* \return The Map from buffer to the LCA of all access to it. The lca is function root if the
diff --git a/include/tvm/s_tir/meta_schedule/feature_extractor.h b/include/tvm/s_tir/meta_schedule/feature_extractor.h
index 70b9977136..e86b63cb4f 100644
--- a/include/tvm/s_tir/meta_schedule/feature_extractor.h
+++ b/include/tvm/s_tir/meta_schedule/feature_extractor.h
@@ -90,8 +90,8 @@ class PyFeatureExtractorNode : public FeatureExtractorNode {
class FeatureExtractor : public ffi::ObjectRef {
public:
/*!
- * \brief Create a feature extractor that extracts features from each BufferStore
- * \param buffers_per_store The number of buffers in each BufferStore; Pad or truncate if
+ * \brief Create a feature extractor that extracts features from each TensorStore
+ * \param buffers_per_store The number of buffers in each TensorStore; Pad or truncate if
* necessary.
* \param arith_intensity_curve_num_samples The number of samples used in the arithmetic intensity
* curve.
diff --git a/include/tvm/s_tir/schedule/schedule.h b/include/tvm/s_tir/schedule/schedule.h
index ef0183b331..d149445d90 100644
--- a/include/tvm/s_tir/schedule/schedule.h
+++ b/include/tvm/s_tir/schedule/schedule.h
@@ -593,7 +593,7 @@ class ScheduleNode : public ffi::Object {
* \brief Inline a block into its consumer(s). It requires:
* 1) The block is a complete non-root block, which only produces one buffer
* 2) The block must not be the only leaf in the scope.
- * 3) The body of the block must be a BufferStore statement in the form of,
+ * 3) The body of the block must be a TensorStore statement in the form of,
* A[i, j, k, ...] = ...
* where the indices of the LHS are all distinct atomic variables,
* and no variables other than those indexing variables are allowed in the statement.
@@ -605,7 +605,7 @@ class ScheduleNode : public ffi::Object {
* 1) The block is a complete non-root block, which only produces and consumers one buffer
* 2) The block must not be the only leaf in the scope.
* 3) The only producer of the block is a read-after-write producer and a complete non-root block
- * 4) The body of the block must be a BufferStore statement in the form of,
+ * 4) The body of the block must be a TensorStore statement in the form of,
* B[f(i, j, k, ...)] = g(i, j, k, A[i, j, k, ...] ...)
* where the indices of each `TensorLoad` on the RHS are all distinct atomic variables,
* and no variables other than those indexing variables are allowed in the statement.
diff --git a/include/tvm/tirx/script/ir_builder/ir.h b/include/tvm/tirx/script/ir_builder/ir.h
index 0263aa709e..8288c4babb 100644
--- a/include/tvm/tirx/script/ir_builder/ir.h
+++ b/include/tvm/tirx/script/ir_builder/ir.h
@@ -333,7 +333,7 @@ RegionFrame Region(Op op, ffi::Array<Expr> args,
* \param indices The indices location to be stored.
* \return The same statement that was added to the parent frame.
*/
-tvm::tirx::Stmt BufferStore(TensorVar buffer, PrimExpr value, ffi::Array<PrimExpr> indices);
+tvm::tirx::Stmt TensorStore(TensorVar buffer, PrimExpr value, ffi::Array<PrimExpr> indices);
/*!
* \brief Evaluate the input expression.
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 315e01e638..2327bdcd1a 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -239,11 +239,11 @@ class AssertStmt : public Stmt {
* \endcode
* \sa MakeTensorLoad
*/
-class BufferStoreNode : public StmtNode {
+class TensorStoreNode : public StmtNode {
public:
- explicit BufferStoreNode(ffi::UnsafeInit tag) : buffer(tag), value(tag) {}
+ explicit TensorStoreNode(ffi::UnsafeInit tag) : buffer(tag), value(tag) {}
- BufferStoreNode(TensorVar buffer, PrimExpr value)
+ TensorStoreNode(TensorVar buffer, PrimExpr value)
: buffer(std::move(buffer)), value(std::move(value)) {}
/*! \brief The buffer variable. */
@@ -255,27 +255,27 @@ class BufferStoreNode : public StmtNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<BufferStoreNode>()
- .def_ro("buffer", &BufferStoreNode::buffer)
- .def_ro("value", &BufferStoreNode::value)
- .def_ro("indices", &BufferStoreNode::indices);
+ refl::ObjectDef<TensorStoreNode>()
+ .def_ro("buffer", &TensorStoreNode::buffer)
+ .def_ro("value", &TensorStoreNode::value)
+ .def_ro("indices", &TensorStoreNode::indices);
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferStore", BufferStoreNode, StmtNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.TensorStore", TensorStoreNode, StmtNode);
};
/*!
- * \brief Managed reference to BufferStoreNode.
- * \sa BufferStoreNode
+ * \brief Managed reference to TensorStoreNode.
+ * \sa TensorStoreNode
*/
-class BufferStore : public Stmt {
+class TensorStore : public Stmt {
public:
- TVM_DLL explicit BufferStore(TensorVar buffer, PrimExpr value, ffi::Array<PrimExpr> indices,
+ TVM_DLL explicit TensorStore(TensorVar buffer, PrimExpr value, ffi::Array<PrimExpr> indices,
Span span = Span());
- explicit BufferStore(ffi::ObjectPtr<BufferStoreNode> node) : Stmt(std::move(node)) {}
+ explicit TensorStore(ffi::ObjectPtr<TensorStoreNode> node) : Stmt(std::move(node)) {}
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BufferStore, Stmt, BufferStoreNode);
- TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferStoreNode);
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorStore, Stmt, TensorStoreNode);
+ TVM_DEFINE_OBJECT_REF_COW_METHOD(TensorStoreNode);
};
/*!
diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h
index 6762f53769..430c4f97d9 100644
--- a/include/tvm/tirx/stmt_functor.h
+++ b/include/tvm/tirx/stmt_functor.h
@@ -98,7 +98,7 @@ class StmtFunctor<R(const Stmt&, Args...)> {
virtual R Dispatch_(const ContinueNode* node, Args... args) {
return DispatchDefault_(node, std::forward<Args>(args)...);
}
- virtual R Dispatch_(const BufferStoreNode* node, Args... args) {
+ virtual R Dispatch_(const TensorStoreNode* node, Args... args) {
return DispatchDefault_(node, std::forward<Args>(args)...);
}
virtual R Dispatch_(const AssertStmtNode* node, Args... args) {
@@ -141,7 +141,7 @@ class StmtFunctor<R(const Stmt&, Args...)> {
SetDispatch<TSelf, ReturnNode>(vtable);
SetDispatch<TSelf, BreakNode>(vtable);
SetDispatch<TSelf, ContinueNode>(vtable);
- SetDispatch<TSelf, BufferStoreNode>(vtable);
+ SetDispatch<TSelf, TensorStoreNode>(vtable);
SetDispatch<TSelf, AssertStmtNode>(vtable);
SetDispatch<TSelf, SeqStmtNode>(vtable);
SetDispatch<TSelf, EvaluateNode>(vtable);
@@ -198,7 +198,7 @@ class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor {
virtual ffi::Optional<VisitInterrupt> Visit_(const ReturnNode* op);
virtual ffi::Optional<VisitInterrupt> Visit_(const BreakNode* op);
virtual ffi::Optional<VisitInterrupt> Visit_(const ContinueNode* op);
- virtual ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op);
+ virtual ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op);
virtual ffi::Optional<VisitInterrupt> Visit_(const AssertStmtNode* op);
virtual ffi::Optional<VisitInterrupt> Visit_(const SeqStmtNode* op);
virtual ffi::Optional<VisitInterrupt> Visit_(const EvaluateNode* op);
@@ -228,7 +228,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 TensorRegion inference.
+ * provide the corresponding TensorLoad, TensorStore, or TensorRegion inference.
*/
class TVM_DLL StmtExprMutator : public tvm::ExprMutator {
public:
@@ -266,7 +266,7 @@ class TVM_DLL StmtExprMutator : public tvm::ExprMutator {
virtual UnchangedOr<Stmt> Mutate_(const ReturnNode* op, InplaceMode inplace_mode);
virtual UnchangedOr<Stmt> Mutate_(const BreakNode* op, InplaceMode inplace_mode);
virtual UnchangedOr<Stmt> Mutate_(const ContinueNode* op, InplaceMode inplace_mode);
- virtual UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode);
+ virtual UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode);
virtual UnchangedOr<Stmt> Mutate_(const AssertStmtNode* op, InplaceMode inplace_mode);
virtual UnchangedOr<Stmt> Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode);
virtual UnchangedOr<Stmt> Mutate_(const EvaluateNode* op, InplaceMode inplace_mode);
diff --git a/include/tvm/tirx/transform.h b/include/tvm/tirx/transform.h
index 14542f5460..38b08a82ff 100644
--- a/include/tvm/tirx/transform.h
+++ b/include/tvm/tirx/transform.h
@@ -280,8 +280,8 @@ TVM_DLL Pass InlinePrivateFunctions();
TVM_DLL Pass PointerValueTypeRewrite();
/*!
- * \brief Flatten the multi-dimensional TensorLoad and BufferStore to single dimensional
- * TensorLoad/BufferStore for the TIR not contains opaque block.
+ * \brief Flatten the multi-dimensional TensorLoad and TensorStore to single dimensional
+ * TensorLoad/TensorStore for the TIR not contains opaque block.
* \return The pass.
*/
TVM_DLL Pass FlattenBuffer();
diff --git a/include/tvm/tirx/type.h b/include/tvm/tirx/type.h
index d0ba52231f..c2208d5f2a 100644
--- a/include/tvm/tirx/type.h
+++ b/include/tvm/tirx/type.h
@@ -74,7 +74,7 @@ class TensorTypeNode : public TypeNode {
/*! \brief The type of the buffer prior to flattening
*
* This contains the shape as it is accessed by
- * TensorLoad/BufferStore nodes, and used by the low-level code
+ * TensorLoad/TensorStore nodes, and used by the low-level code
* generators.
*/
ffi::Array<PrimExpr> shape;
diff --git a/python/tvm/backend/cuda/lang/clc.py b/python/tvm/backend/cuda/lang/clc.py
index 65577963a7..670bc0c503 100644
--- a/python/tvm/backend/cuda/lang/clc.py
+++ b/python/tvm/backend/cuda/lang/clc.py
@@ -41,7 +41,7 @@ def query_cancel_first_ctaid_x(first_ctaid_x, handle, *, use_ld_acquire=True):
T.ptx[f"ld{'.acquire.cta' if use_ld_acquire else ''}.shared.b128"](response, handle)
T.ptx.clusterlaunchcontrol.query_cancel.is_canceled.pred.b128(canceled, response)
- T.buffer_store(first_ctaid_x.source, T.uint32(0xFFFFFFFF), first_ctaid_x.indices)
+ T.tensor_store(first_ctaid_x.source, T.uint32(0xFFFFFFFF), first_ctaid_x.indices)
T.ptx.clusterlaunchcontrol.query_cancel.get_first_ctaid__x.b32.b128(
first_ctaid_x, response, pred=canceled
)
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/fallback.py b/python/tvm/backend/cuda/tile_primitive/copy/fallback.py
index 7df2a215e0..b67c4b2227 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/fallback.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/fallback.py
@@ -76,11 +76,11 @@ def _emit_fallback(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFun
return coord
if not copy_extents:
- T.buffer_store(dst_buf, src_buf[tuple(src_st)], dst_st)
+ T.tensor_store(dst_buf, src_buf[tuple(src_st)], dst_st)
return
with T.grid(*copy_extents) as lvs:
- T.buffer_store(dst_buf, src_buf[tuple(_src_coord(lvs))], _dst_coord(lvs))
+ T.tensor_store(dst_buf, src_buf[tuple(_src_coord(lvs))], _dst_coord(lvs))
scope_kind = sctx.scope_kind
diff --git a/python/tvm/relax/transform/legalize_ops/grad.py b/python/tvm/relax/transform/legalize_ops/grad.py
index a37829e3a9..14910147c9 100644
--- a/python/tvm/relax/transform/legalize_ops/grad.py
+++ b/python/tvm/relax/transform/legalize_ops/grad.py
@@ -180,7 +180,7 @@ def _grad_take_backward(bb: BlockBuilder, call: Call) -> Expr:
with IRBuilder() as ib:
# Init loop (zero-fill output buffer)
with T.serial(fused_shape) as i:
- T.buffer_store(out, tirx.const(0, dtype=x_ptr.dtype), T.tensor_indices(out, i))
+ T.tensor_store(out, tirx.const(0, dtype=x_ptr.dtype), T.tensor_indices(out, i))
# Accumulation loop
if axis is not None:
@@ -210,7 +210,7 @@ def _grad_take_backward(bb: BlockBuilder, call: Call) -> Expr:
+ loop_l * fused_output_grad_shape_nxt
+ j
)
- T.buffer_store(
+ T.tensor_store(
out,
out[T.tensor_indices(out, out_idx)]
+ grad[T.tensor_indices(grad, grad_idx)],
@@ -218,7 +218,7 @@ def _grad_take_backward(bb: BlockBuilder, call: Call) -> Expr:
)
else:
with T.serial(indices_len) as loop_l:
- T.buffer_store(
+ T.tensor_store(
out,
out[T.tensor_indices(out, (idx[T.tensor_indices(idx, loop_l)]))]
+ grad[T.tensor_indices(grad, loop_l)],
diff --git a/python/tvm/s_tir/analysis/__init__.py b/python/tvm/s_tir/analysis/__init__.py
index 018c5778e5..bc907f9576 100644
--- a/python/tvm/s_tir/analysis/__init__.py
+++ b/python/tvm/s_tir/analysis/__init__.py
@@ -79,7 +79,7 @@ def get_sblock_read_write_region(
def detect_buffer_access_lca(func: PrimFunc) -> dict[Var, Stmt]:
"""Detect the lowest common ancestor(LCA) of buffer access, including both high-level
- access (TensorLoad, BufferStore) and low-level access (TensorLoad, BufferStore and opaque
+ access (TensorLoad, TensorStore) and low-level access (TensorLoad, TensorStore and opaque
access).
The LCA may be a For loop or a Block.
diff --git a/python/tvm/s_tir/dlight/analysis/gemv.py b/python/tvm/s_tir/dlight/analysis/gemv.py
index 44327a3029..123dd96cfe 100644
--- a/python/tvm/s_tir/dlight/analysis/gemv.py
+++ b/python/tvm/s_tir/dlight/analysis/gemv.py
@@ -45,18 +45,18 @@ def get_reduction_expr(block: s_tir.SBlock) -> tirx.Expr | None:
The reduction expression (`Y`) if detected, otherwise None.
"""
- buffer_store = block.body
- if not isinstance(buffer_store, tirx.BufferStore):
+ tensor_store = block.body
+ if not isinstance(tensor_store, tirx.TensorStore):
return None
- if not isinstance(buffer_store.value, tirx.Add):
+ if not isinstance(tensor_store.value, tirx.Add):
return None
if not tvm_ffi.structural_equal(
- buffer_store.value.a,
- tirx.TensorLoad(buffer_store.buffer, block.body.indices),
+ tensor_store.value.a,
+ tirx.TensorLoad(tensor_store.buffer, block.body.indices),
map_free_vars=True,
):
return None
- return buffer_store.value.b
+ return tensor_store.value.b
def is_gemv(sch: s_tir.Schedule, block_info: SBlockInfo) -> list[tirx.Var] | None:
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 5e9b0ecaf4..b1385d8a57 100644
--- a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
@@ -39,18 +39,18 @@ from .base import GPUScheduleRule
def _get_reduction_expr(block: s_tir.SBlock) -> tirx.Expr | None:
# Detect and return `Y` in `X[...] = X[...] + Y`
- buffer_store = block.body
- if not isinstance(buffer_store, tirx.BufferStore):
+ tensor_store = block.body
+ if not isinstance(tensor_store, tirx.TensorStore):
return None
- if not isinstance(buffer_store.value, tirx.Add):
+ if not isinstance(tensor_store.value, tirx.Add):
return None
if not tvm_ffi.structural_equal(
- buffer_store.value.a,
- tirx.TensorLoad(buffer_store.buffer, block.body.indices),
+ tensor_store.value.a,
+ tirx.TensorLoad(tensor_store.buffer, block.body.indices),
map_free_vars=True,
):
return None
- return buffer_store.value.b
+ return tensor_store.value.b
def _has_pad_einsum_compatible_access(block: s_tir.SBlock) -> bool:
diff --git a/python/tvm/s_tir/dlight/gpu/reduction.py b/python/tvm/s_tir/dlight/gpu/reduction.py
index 1f27e8a174..6b50ab5ba7 100644
--- a/python/tvm/s_tir/dlight/gpu/reduction.py
+++ b/python/tvm/s_tir/dlight/gpu/reduction.py
@@ -36,18 +36,18 @@ from .base import GPUScheduleRule
def _get_reduction_expr(block: s_tir.SBlock) -> tirx.Expr | None:
# Detect and return `Y` in `X[...] = X[...] + Y`
- buffer_store = block.body
- if not isinstance(buffer_store, tirx.BufferStore):
+ tensor_store = block.body
+ if not isinstance(tensor_store, tirx.TensorStore):
return None
- if not isinstance(buffer_store.value, tirx.Add):
+ if not isinstance(tensor_store.value, tirx.Add):
return None
if not tvm_ffi.structural_equal(
- buffer_store.value.a,
- tirx.TensorLoad(buffer_store.buffer, block.body.indices),
+ tensor_store.value.a,
+ tirx.TensorLoad(tensor_store.buffer, block.body.indices),
map_free_vars=True,
):
return None
- return buffer_store.value.b
+ return tensor_store.value.b
def _has_reduction_loop(block_info):
diff --git a/python/tvm/s_tir/dlight/gpu/rmsnorm.py b/python/tvm/s_tir/dlight/gpu/rmsnorm.py
index 8245c6930a..4cf0ee3cc3 100644
--- a/python/tvm/s_tir/dlight/gpu/rmsnorm.py
+++ b/python/tvm/s_tir/dlight/gpu/rmsnorm.py
@@ -22,7 +22,7 @@ from tvm import tirx
from tvm.ir import Call, TensorLoad
from tvm.s_tir import SBlock
from tvm.target import Target
-from tvm.tirx import BufferStore
+from tvm.tirx import TensorStore
from tvm.tirx.expr import Cast
from ..base import ScheduleRule
@@ -32,7 +32,7 @@ def identify_cast_or_load_block(block: SBlock) -> bool:
if len(block.reads) != 1 or len(block.writes) != 1:
return False
- if not isinstance(block.body, BufferStore):
+ if not isinstance(block.body, TensorStore):
return False
store = block.body
@@ -61,7 +61,7 @@ def identify_rsqrt_block(block: SBlock) -> bool:
if len(block.reads) != 1 or len(block.writes) != 1:
return False
- if not isinstance(block.body, BufferStore):
+ if not isinstance(block.body, TensorStore):
return False
store = block.body
diff --git a/python/tvm/s_tir/dlight/gpu/transpose.py b/python/tvm/s_tir/dlight/gpu/transpose.py
index 412b31f1f0..850e631dd3 100644
--- a/python/tvm/s_tir/dlight/gpu/transpose.py
+++ b/python/tvm/s_tir/dlight/gpu/transpose.py
@@ -32,7 +32,7 @@ class Transpose(GPUScheduleRule):
def is_transpose(self, sch: Schedule, block_rv: SBlockRV):
block = sch.get(block_rv)
- if isinstance(block.body, tirx.BufferStore):
+ if isinstance(block.body, tirx.TensorStore):
rhs = block.body.value
if isinstance(rhs, TensorLoad):
lhs_indices = block.body.indices
diff --git a/python/tvm/s_tir/meta_schedule/feature_extractor/per_store_feature.py b/python/tvm/s_tir/meta_schedule/feature_extractor/per_store_feature.py
index 4cf085e900..9f1130bc9e 100644
--- a/python/tvm/s_tir/meta_schedule/feature_extractor/per_store_feature.py
+++ b/python/tvm/s_tir/meta_schedule/feature_extractor/per_store_feature.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name
-"""We extract one feature vector per BufferStoreNode statement in a TIR Stmt,
+"""We extract one feature vector per TensorStoreNode statement in a TIR Stmt,
so we call this feature as "per-store" feature.
"""
@@ -27,12 +27,12 @@ from .feature_extractor import FeatureExtractor
@register_object("s_tir.meta_schedule.PerStoreFeature")
class PerStoreFeature(FeatureExtractor):
- """PerStoreFeature extracts one feature vector per BufferStoreNode
+ """PerStoreFeature extracts one feature vector per TensorStoreNode
Parameters
----------
buffers_per_store : int
- The number of buffers in each BufferStore; Pad or truncate if necessary.
+ The number of buffers in each TensorStore; Pad or truncate if necessary.
arith_intensity_curve_num_samples : int
The number of samples used in the arithmetic intensity curve.
cache_line_bytes : int
@@ -42,7 +42,7 @@ class PerStoreFeature(FeatureExtractor):
"""
buffers_per_store: int
- """The number of buffers in each BufferStore; Pad or truncate if necessary."""
+ """The number of buffers in each TensorStore; Pad or truncate if necessary."""
arith_intensity_curve_num_samples: int # pylint: disable=invalid-name
"""The number of samples used in the arithmetic intensity curve."""
cache_line_bytes: int
diff --git a/python/tvm/s_tir/schedule/schedule.py b/python/tvm/s_tir/schedule/schedule.py
index 4b09fb0fde..0c3ec79b7d 100644
--- a/python/tvm/s_tir/schedule/schedule.py
+++ b/python/tvm/s_tir/schedule/schedule.py
@@ -2240,7 +2240,7 @@ class Schedule(Object):
2) The block must not be the only leaf in the scope.
- 3) The body of the block must be a BufferStore statement in
+ 3) The body of the block must be a TensorStore statement in
the form of, ``A[i, j, k, ...] = ...`` where the indices of
the LHS are all distinct atomic variables, and no variables
other than those indexing variables are allowed in the
@@ -2309,7 +2309,7 @@ class Schedule(Object):
3) The only producer of the block is a read-after-write producer and a
complete non-root block
- 4) The body of the block must be a BufferStore statement in the form of,
+ 4) The body of the block must be a TensorStore statement in the form of,
``B[f(i, j, k, ...)] = g(i, j, k, A[i, j, k, ...] ...)`` where the
indices of each `TensorLoad` on the RHS are all distinct atomic
variables, and no variables other than those indexing variables are
@@ -2641,7 +2641,7 @@ class Schedule(Object):
7) An unary extent loop that is not bound to any reduction or data parallel variables in
the block binding should not appear under some reduction loop;
8) The reduction block should write to only one buffer, and its init and body are both
- simple `BufferStore`s, and the pattern is registered as an associative reducer.
+ simple `TensorStore`s, and the pattern is registered as an associative reducer.
The pre-defined patterns include: plus, multiplication, min and max;
9) Each of the loops on top of the block cannot be bound to a data parallel and a
reduction block binding at the same time;
diff --git a/python/tvm/s_tir/tensor_intrin/cuda.py b/python/tvm/s_tir/tensor_intrin/cuda.py
index dc536592f9..5173f90073 100644
--- a/python/tvm/s_tir/tensor_intrin/cuda.py
+++ b/python/tvm/s_tir/tensor_intrin/cuda.py
@@ -699,12 +699,12 @@ def get_mma_intrin_group(
Whether the input matrix B is transposed.
not_use_mma_store_intrinic : bool
- Whether to not use the mma_store intrinsic. If True, use BufferStore stmts to store the
+ Whether to not use the mma_store intrinsic. If True, use TensorStore stmts to store the
result of mma. Otherwise, use mma_store intrinsic.
This is because if we use mma_store intrinsic, during swizzling shared memory visits, our
rearrangement scheme will involve areas accessed by different mma_store calls. This makes
- swizzling quite complex. But BufferStore will not face this problem.
+ swizzling quite complex. But TensorStore will not face this problem.
store_to_smem_dtype : Optional[Literal["float16", "float32", "int32"]]
The dtype that we use to store from register to shared memory. By default it is out_dtype.
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index 63dd38366c..8e17805fe3 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -48,7 +48,7 @@ from .stmt import Stmt, Bind, AssertStmt, ForKind, For, While, Return, Break, Co
# Legacy alias: LetStmt was folded into the body-less Bind statement.
LetStmt = Bind
-from .stmt import BufferStore, RegionStmt
+from .stmt import TensorStore, RegionStmt
from .stmt import SeqStmt
from .stmt import IfThenElse, Evaluate, stmt_seq, stmt_list
diff --git a/python/tvm/tirx/script/ir_builder/parser_protocol.py b/python/tvm/tirx/script/ir_builder/parser_protocol.py
index 75e9e93589..557233da43 100644
--- a/python/tvm/tirx/script/ir_builder/parser_protocol.py
+++ b/python/tvm/tirx/script/ir_builder/parser_protocol.py
@@ -325,14 +325,14 @@ def set_mutable_cell_(
Updates emit a scalar buffer store. Targets must denote scalar storage.
"""
if isinstance(target, _ir.TensorLoad):
- return _base.at_(span, buffer_store(target.source, value, list(target.indices)))
+ return _base.at_(span, tensor_store(target.source, value, list(target.indices)))
elif (
_tir.is_tensor_var(target)
and len(target.ty.shape) == 1
and isinstance(target.ty.shape[0], _tir.IntImm)
and target.ty.shape[0].value == 1
):
- return _base.at_(span, buffer_store(target, value, [0]))
+ return _base.at_(span, tensor_store(target, value, [0]))
else:
raise TypeError("A mutable assignment requires scalar storage")
@@ -386,7 +386,7 @@ def setitem_(
target: Any, key: Any, value: Any, *, span: _Span = None
) -> _base.AlreadyEmitted[_tir.Stmt]:
"""Implements :func:`tvm.script.ir_builder.parser_protocol.setitem_`."""
- return _base.at_(span, buffer_store(target, value, key))
+ return _base.at_(span, tensor_store(target, value, key))
def setattr_(
@@ -437,7 +437,7 @@ def bind( # pylint: disable=invalid-name
return _ffi_api.Bind(value, type_annotation, var) # type: ignore[attr-defined] # pylint: disable=no-member
-def buffer_store(
+def tensor_store(
buffer: Var, # pylint: disable=redefined-outer-name
value: Expr,
indices: list[Expr | slice],
@@ -481,7 +481,7 @@ def buffer_store(
expr_indices.append(index)
if isinstance(value, bool) and buffer.ty.dtype == "bool":
value = IntImm("bool", value)
- return AlreadyEmitted(_ffi_api.BufferStore(buffer, value, expr_indices))
+ return AlreadyEmitted(_ffi_api.TensorStore(buffer, value, expr_indices))
def evaluate(value: Expr) -> AlreadyEmitted[_tir.Stmt]:
@@ -987,7 +987,6 @@ __all__ = [
"bind",
"bind_",
"break_",
- "buffer_store",
"call_global_var_",
"check_well_formed_",
"compute_scope",
@@ -1028,6 +1027,7 @@ __all__ = [
"set_mutable_cell_",
"setattr_",
"setitem_",
+ "tensor_store",
"then_",
"thread_binding",
"unpack_",
diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py
index fff7db4697..b42c19b2b0 100644
--- a/python/tvm/tirx/stmt.py
+++ b/python/tvm/tirx/stmt.py
@@ -22,8 +22,8 @@ Each statement node have subfields that can be visited from python side.
x = tvm.tirx.Var("n", "int32")
buffer = tvm.tirx.decl_tensor((16,), "float32")
- st = tvm.tirx.stmt.BufferStore(buffer, 1, (x,))
- assert isinstance(st, tvm.tirx.stmt.BufferStore)
+ st = tvm.tirx.stmt.TensorStore(buffer, 1, (x,))
+ assert isinstance(st, tvm.tirx.stmt.TensorStore)
assert(st.buffer == buffer)
"""
@@ -233,8 +233,8 @@ class While(Stmt):
self.__init_handle_by_constructor__(_ffi_api.While, condition, body, span) # type: ignore
-@tvm_ffi.register_object("tirx.BufferStore")
-class BufferStore(Stmt):
+@tvm_ffi.register_object("tirx.TensorStore")
+class TensorStore(Stmt):
"""Store a value into a tensor variable.
Parameters
@@ -265,7 +265,7 @@ class BufferStore(Stmt):
span: Span | None = None,
) -> None:
self.__init_handle_by_constructor__(
- _ffi_api.BufferStore,
+ _ffi_api.TensorStore,
buffer,
value,
indices,
diff --git a/python/tvm/tirx/transform/transform.py b/python/tvm/tirx/transform/transform.py
index 3bf78b0edb..aa1dfdf31b 100644
--- a/python/tvm/tirx/transform/transform.py
+++ b/python/tvm/tirx/transform/transform.py
@@ -459,8 +459,8 @@ class HoistExpressionConfig(_ffi.Object):
def FlattenBuffer():
- """Flatten the multi-dimensional TensorLoad and BufferStore to single dimensional
- TensorLoad/BufferStore for the TIR not contains opaque block.
+ """Flatten the multi-dimensional TensorLoad and TensorStore to single dimensional
+ TensorLoad/TensorStore for the TIR not contains opaque block.
Returns
-------
diff --git a/python/tvm/topi/gpu/scan.py b/python/tvm/topi/gpu/scan.py
index 74e5b84984..5cda86f64e 100644
--- a/python/tvm/topi/gpu/scan.py
+++ b/python/tvm/topi/gpu/scan.py
@@ -93,7 +93,7 @@ def exclusive_scan_ir(data, output, reduction=None, binop=operator.add, identity
with T.if_(bx < batch_size):
with T.then_():
if reduction is not None:
- T.buffer_store(
+ T.tensor_store(
reduction,
cast(identity_value, out_dtype),
T.tensor_indices(reduction, bx),
@@ -115,7 +115,7 @@ def exclusive_scan_ir(data, output, reduction=None, binop=operator.add, identity
tid = tvm.tirx.indexmod(bx, blocks_per_batch) * nthread_tx + tx
with T.if_(tid < scan_axis_size):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
output,
cast(
data[T.tensor_indices(data, batch * scan_axis_size + tid)],
@@ -146,15 +146,15 @@ def exclusive_scan_ir(data, output, reduction=None, binop=operator.add, identity
start = start_buf
middle = middle_buf
end = end_buf
- T.buffer_store(start, width * tid, T.tensor_indices(start, 0))
+ T.tensor_store(start, width * tid, T.tensor_indices(start, 0))
with T.if_(start[T.tensor_indices(start, 0)] < scan_axis_size):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
middle,
start[T.tensor_indices(start, 0)] + tvm.tirx.indexdiv(width, 2),
T.tensor_indices(middle, 0),
)
- T.buffer_store(
+ T.tensor_store(
end,
tvm.te.min(
start[T.tensor_indices(start, 0)] + width, scan_axis_size
@@ -163,7 +163,7 @@ def exclusive_scan_ir(data, output, reduction=None, binop=operator.add, identity
)
with T.if_(middle[T.tensor_indices(middle, 0)] < scan_axis_size):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
output,
binop(
output[
@@ -202,14 +202,14 @@ def exclusive_scan_ir(data, output, reduction=None, binop=operator.add, identity
with T.if_(bx < batch_size):
with T.then_():
if reduction is not None:
- T.buffer_store(
+ T.tensor_store(
reduction,
output[
T.tensor_indices(output, ((bx + 1) * scan_axis_size - 1))
],
T.tensor_indices(reduction, bx),
)
- T.buffer_store(
+ T.tensor_store(
output,
cast(identity_value, out_dtype),
T.tensor_indices(output, ((bx + 1) * scan_axis_size - 1)),
@@ -235,17 +235,17 @@ def exclusive_scan_ir(data, output, reduction=None, binop=operator.add, identity
middle = middle_buf
end = end_buf
tmp = tmp_buf
- T.buffer_store(start, width * tid, T.tensor_indices(start, 0))
+ T.tensor_store(start, width * tid, T.tensor_indices(start, 0))
with T.if_(
tvm.tirx.all(start[T.tensor_indices(start, 0)] < scan_axis_size)
):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
middle,
start[T.tensor_indices(start, 0)] + tvm.tirx.indexdiv(width, 2),
T.tensor_indices(middle, 0),
)
- T.buffer_store(
+ T.tensor_store(
end,
tvm.tirx.min(
start[T.tensor_indices(start, 0)] + width, scan_axis_size
@@ -254,7 +254,7 @@ def exclusive_scan_ir(data, output, reduction=None, binop=operator.add, identity
)
with T.if_(middle[T.tensor_indices(middle, 0)] < scan_axis_size):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
tmp,
output[
T.tensor_indices(
@@ -268,7 +268,7 @@ def exclusive_scan_ir(data, output, reduction=None, binop=operator.add, identity
],
T.tensor_indices(tmp, 0),
)
- T.buffer_store(
+ T.tensor_store(
output,
output[
T.tensor_indices(
@@ -289,7 +289,7 @@ def exclusive_scan_ir(data, output, reduction=None, binop=operator.add, identity
),
),
)
- T.buffer_store(
+ T.tensor_store(
output,
binop(
output[
@@ -368,7 +368,7 @@ def get_reduction_from_exclusive_scan(data, ex_scan_output, binop=operator.add):
with T.then_():
with T.if_(scan_axis_size > 0):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
reduction,
binop(
data_ex_scan[
@@ -386,7 +386,7 @@ def get_reduction_from_exclusive_scan(data, ex_scan_output, binop=operator.add):
T.tensor_indices(reduction, tid),
)
with T.else_():
- T.buffer_store(
+ T.tensor_store(
reduction,
cast(0, reduction_buf.dtype),
T.tensor_indices(reduction, tid),
diff --git a/python/tvm/topi/gpu/scatter_elements.py b/python/tvm/topi/gpu/scatter_elements.py
index 514e84f732..e8362149dd 100644
--- a/python/tvm/topi/gpu/scatter_elements.py
+++ b/python/tvm/topi/gpu/scatter_elements.py
@@ -80,7 +80,7 @@ def scatter_elements(data, indices, updates, axis=0, reduction="update"):
tid = bx_init * max_threads + tx_init
with T.if_(tid < full_range):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
out,
data[T.tensor_indices(data, tid)],
T.tensor_indices(out, tid),
@@ -114,38 +114,38 @@ def scatter_elements(data, indices, updates, axis=0, reduction="update"):
return ib.get()
def update_func(dst_ptr, dst_index, update):
- T.buffer_store(dst_ptr, update, T.tensor_indices(dst_ptr, dst_index))
+ T.tensor_store(dst_ptr, update, T.tensor_indices(dst_ptr, dst_index))
def add_func(dst_ptr, dst_index, update):
- T.buffer_store(
+ T.tensor_store(
dst_ptr,
dst_ptr[T.tensor_indices(dst_ptr, dst_index)] + (update),
T.tensor_indices(dst_ptr, dst_index),
)
def mul_func(dst_ptr, dst_index, update):
- T.buffer_store(
+ T.tensor_store(
dst_ptr,
dst_ptr[T.tensor_indices(dst_ptr, dst_index)] * (update),
T.tensor_indices(dst_ptr, dst_index),
)
def mean_func(dst_ptr, dst_index, update):
- T.buffer_store(
+ T.tensor_store(
dst_ptr,
(dst_ptr[T.tensor_indices(dst_ptr, dst_index)] + update) / 2,
T.tensor_indices(dst_ptr, dst_index),
)
def min_func(dst_ptr, dst_index, update):
- T.buffer_store(
+ T.tensor_store(
dst_ptr,
tirx.min(dst_ptr[T.tensor_indices(dst_ptr, dst_index)], update),
T.tensor_indices(dst_ptr, dst_index),
)
def max_func(dst_ptr, dst_index, update):
- T.buffer_store(
+ T.tensor_store(
dst_ptr,
tirx.max(dst_ptr[T.tensor_indices(dst_ptr, dst_index)], update),
T.tensor_indices(dst_ptr, dst_index),
diff --git a/python/tvm/topi/gpu/scatter_nd.py b/python/tvm/topi/gpu/scatter_nd.py
index f9db776650..2a1879faf1 100644
--- a/python/tvm/topi/gpu/scatter_nd.py
+++ b/python/tvm/topi/gpu/scatter_nd.py
@@ -68,7 +68,7 @@ def scatter_nd(data, indices, updates, mode):
tid = bx_init * max_threads + tx_init
with T.if_(tid < fused_shape):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
out,
data[T.tensor_indices(data, tid)],
T.tensor_indices(out, tid),
@@ -100,7 +100,7 @@ def scatter_nd(data, indices, updates, mode):
)
offset *= data_ptr.shape[l]
if mode == "update":
- T.buffer_store(
+ T.tensor_store(
out,
updates[
T.tensor_indices(updates, (i * fused_updates_dimension + j))
@@ -108,7 +108,7 @@ def scatter_nd(data, indices, updates, mode):
T.tensor_indices(out, index),
)
elif mode == "add":
- T.buffer_store(
+ T.tensor_store(
out,
out[T.tensor_indices(out, index)]
+ (
@@ -121,7 +121,7 @@ def scatter_nd(data, indices, updates, mode):
T.tensor_indices(out, index),
)
elif mode == "mul":
- T.buffer_store(
+ T.tensor_store(
out,
out[T.tensor_indices(out, index)]
* (
@@ -134,7 +134,7 @@ def scatter_nd(data, indices, updates, mode):
T.tensor_indices(out, index),
)
elif mode == "min":
- T.buffer_store(
+ T.tensor_store(
out,
tirx.min(
out[T.tensor_indices(out, index)],
@@ -147,7 +147,7 @@ def scatter_nd(data, indices, updates, mode):
T.tensor_indices(out, index),
)
elif mode == "max":
- T.buffer_store(
+ T.tensor_store(
out,
tirx.max(
out[T.tensor_indices(out, index)],
diff --git a/python/tvm/topi/gpu/sort.py b/python/tvm/topi/gpu/sort.py
index c072581049..03eea6d4c5 100644
--- a/python/tvm/topi/gpu/sort.py
+++ b/python/tvm/topi/gpu/sort.py
@@ -66,13 +66,13 @@ def _sort_init(shape, axis, keys_in, keys_out, values_out=None, value_init_func=
idx = (by_val * shape[axis] + tid) * axis_mul_after + bz
with T.if_(tid < shape[axis]):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
keys_out,
keys_in[T.tensor_indices(keys_in, idx)],
T.tensor_indices(keys_out, idx),
)
if values_out is not None:
- T.buffer_store(
+ T.tensor_store(
values_out, value_init_func(idx, tid), T.tensor_indices(values_out, idx)
)
@@ -124,7 +124,7 @@ def _odd_even_sort(
with T.serial(0, 2) as n:
with T.if_((tid + n + start) < size):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
tmp_keys_swap,
keys[
T.tensor_indices(keys, (base_idx + (tid + n + start) * axis_mul_after))
@@ -132,7 +132,7 @@ def _odd_even_sort(
[tid + n],
)
if values_swap is not None:
- T.buffer_store(
+ T.tensor_store(
tmp_values_swap,
values[
T.tensor_indices(
@@ -151,25 +151,25 @@ def _odd_even_sort(
n = idxm(tid + k, 2)
with T.if_(tid + n < current_sort_num - 1):
with T.then_():
- T.buffer_store(temp_cond1, tmp_keys_swap[tid + n], [0])
- T.buffer_store(temp_cond2, tmp_keys_swap[tid + n + 1], [0])
+ T.tensor_store(temp_cond1, tmp_keys_swap[tid + n], [0])
+ T.tensor_store(temp_cond2, tmp_keys_swap[tid + n + 1], [0])
if is_ascend:
cond = temp_cond1[0] > temp_cond2[0]
else:
cond = temp_cond1[0] < temp_cond2[0]
with T.if_(cond):
with T.then_():
- T.buffer_store(temp_keys, tmp_keys_swap[tid + n], [0])
- T.buffer_store(tmp_keys_swap, tmp_keys_swap[tid + n + 1], [tid + n])
- T.buffer_store(tmp_keys_swap, temp_keys[0], [tid + n + 1])
+ T.tensor_store(temp_keys, tmp_keys_swap[tid + n], [0])
+ T.tensor_store(tmp_keys_swap, tmp_keys_swap[tid + n + 1], [tid + n])
+ T.tensor_store(tmp_keys_swap, temp_keys[0], [tid + n + 1])
if values_swap is not None:
- T.buffer_store(temp_values, tmp_values_swap[tid + n], [0])
- T.buffer_store(
+ T.tensor_store(temp_values, tmp_values_swap[tid + n], [0])
+ T.tensor_store(
tmp_values_swap,
tmp_values_swap[tid + n + 1],
[tid + n],
)
- T.buffer_store(tmp_values_swap, temp_values[0], [tid + n + 1])
+ T.tensor_store(tmp_values_swap, temp_values[0], [tid + n + 1])
T.evaluate(
tvm.ir.Call("tirx.tvm_storage_sync", [tvm.ir.StringImm("shared")], ty="void")
)
@@ -179,15 +179,15 @@ def _odd_even_sort(
with T.if_(tid + n + start < size):
with T.then_():
out_idx = base_idx + (tid + n + start) * axis_mul_after
- T.buffer_store(keys, tmp_keys_swap[tid + n], T.tensor_indices(keys, out_idx))
- T.buffer_store(
+ T.tensor_store(keys, tmp_keys_swap[tid + n], T.tensor_indices(keys, out_idx))
+ T.tensor_store(
keys_swap, tmp_keys_swap[tid + n], T.tensor_indices(keys_swap, out_idx)
)
if values_swap is not None:
- T.buffer_store(
+ T.tensor_store(
values, tmp_values_swap[tid + n], T.tensor_indices(values, out_idx)
)
- T.buffer_store(
+ T.tensor_store(
values_swap,
tmp_values_swap[tid + n],
T.tensor_indices(values_swap, out_idx),
@@ -253,20 +253,20 @@ def _sort_common(
max_val = tvm.te.max(0, diag - bCount)
min_val = tvm.te.min(diag, aCount)
if is_webgpu:
- T.buffer_store(first, cast(max_val, target_dtype), T.tensor_indices(first, 0))
- T.buffer_store(last, cast(min_val, target_dtype), T.tensor_indices(last, 0))
+ T.tensor_store(first, cast(max_val, target_dtype), T.tensor_indices(first, 0))
+ T.tensor_store(last, cast(min_val, target_dtype), T.tensor_indices(last, 0))
else:
- T.buffer_store(first, max_val, T.tensor_indices(first, 0))
- T.buffer_store(last, min_val, T.tensor_indices(last, 0))
+ T.tensor_store(first, max_val, T.tensor_indices(first, 0))
+ T.tensor_store(last, min_val, T.tensor_indices(last, 0))
with T.while_(first[T.tensor_indices(first, 0)] < last[T.tensor_indices(last, 0)]):
mid = (first[T.tensor_indices(first, 0)] + last[T.tensor_indices(last, 0)]) >> 1
a = source[T.tensor_indices(source, (base_idx + (aStart + mid)))]
b = source[T.tensor_indices(source, (base_idx + (bStart + diag - 1 - mid)))]
with T.if_(compare(a, b)):
with T.then_():
- T.buffer_store(first, mid + 1, T.tensor_indices(first, 0))
+ T.tensor_store(first, mid + 1, T.tensor_indices(first, 0))
with T.else_():
- T.buffer_store(last, mid, T.tensor_indices(last, 0))
+ T.tensor_store(last, mid, T.tensor_indices(last, 0))
def serial_merge(
source,
@@ -289,11 +289,11 @@ def _sort_common(
i_val = aStart + first[T.tensor_indices(first, 0)]
j_val = bStart + diag - last[T.tensor_indices(last, 0)]
if is_webgpu:
- T.buffer_store(i_buf, cast(i_val, target_dtype), T.tensor_indices(i_buf, 0))
- T.buffer_store(j_buf, cast(j_val, target_dtype), T.tensor_indices(j_buf, 0))
+ T.tensor_store(i_buf, cast(i_val, target_dtype), T.tensor_indices(i_buf, 0))
+ T.tensor_store(j_buf, cast(j_val, target_dtype), T.tensor_indices(j_buf, 0))
else:
- T.buffer_store(i_buf, i_val, T.tensor_indices(i_buf, 0))
- T.buffer_store(j_buf, j_val, T.tensor_indices(j_buf, 0))
+ T.tensor_store(i_buf, i_val, T.tensor_indices(i_buf, 0))
+ T.tensor_store(j_buf, j_val, T.tensor_indices(j_buf, 0))
with T.serial(0, tvm.te.min(aCount + bCount - diag, step_count)) as count:
i_idx = base_idx + i_buf[T.tensor_indices(i_buf, 0)]
@@ -314,35 +314,35 @@ def _sort_common(
)
):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
dest,
source[T.tensor_indices(source, i_idx)],
T.tensor_indices(dest, k_idx),
)
if values is not None:
- T.buffer_store(
+ T.tensor_store(
dest_idx,
source_idx[T.tensor_indices(source_idx, i_idx)],
T.tensor_indices(dest_idx, k_idx),
)
- T.buffer_store(
+ T.tensor_store(
i_buf,
i_buf[T.tensor_indices(i_buf, 0)] + 1,
T.tensor_indices(i_buf, 0),
)
with T.else_():
- T.buffer_store(
+ T.tensor_store(
dest,
source[T.tensor_indices(source, j_idx)],
T.tensor_indices(dest, k_idx),
)
if values is not None:
- T.buffer_store(
+ T.tensor_store(
dest_idx,
source_idx[T.tensor_indices(source_idx, j_idx)],
T.tensor_indices(dest_idx, k_idx),
)
- T.buffer_store(
+ T.tensor_store(
j_buf,
j_buf[T.tensor_indices(j_buf, 0)] + 1,
T.tensor_indices(j_buf, 0),
@@ -350,35 +350,35 @@ def _sort_common(
with T.else_():
with T.if_(i_buf[T.tensor_indices(i_buf, 0)] < aStart + aCount):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
dest,
source[T.tensor_indices(source, i_idx)],
T.tensor_indices(dest, k_idx),
)
if values is not None:
- T.buffer_store(
+ T.tensor_store(
dest_idx,
source_idx[T.tensor_indices(source_idx, i_idx)],
T.tensor_indices(dest_idx, k_idx),
)
- T.buffer_store(
+ T.tensor_store(
i_buf,
i_buf[T.tensor_indices(i_buf, 0)] + 1,
T.tensor_indices(i_buf, 0),
)
with T.else_():
- T.buffer_store(
+ T.tensor_store(
dest,
source[T.tensor_indices(source, j_idx)],
T.tensor_indices(dest, k_idx),
)
if values is not None:
- T.buffer_store(
+ T.tensor_store(
dest_idx,
source_idx[T.tensor_indices(source_idx, j_idx)],
T.tensor_indices(dest_idx, k_idx),
)
- T.buffer_store(
+ T.tensor_store(
j_buf,
j_buf[T.tensor_indices(j_buf, 0)] + 1,
T.tensor_indices(j_buf, 0),
@@ -654,13 +654,13 @@ def _sort_common(
idx = by2 * size + tid
with T.if_(tid < size):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
keys,
keys_swap[T.tensor_indices(keys_swap, idx)],
T.tensor_indices(keys, idx),
)
if values is not None:
- T.buffer_store(
+ T.tensor_store(
values,
values_swap[T.tensor_indices(values_swap, idx)],
T.tensor_indices(values, idx),
@@ -1231,7 +1231,7 @@ def searchsorted(sorted_sequence, values, right=False, out_dtype="int64"):
sequence_id = tid // values_shape[-1]
sequence_offset = sequence_id * search_range
- T.buffer_store(
+ T.tensor_store(
indices_ptr,
binary_search(
sequence_offset,
diff --git a/python/tvm/topi/index_put.py b/python/tvm/topi/index_put.py
index 1c4f6c4f28..e5399e1fc4 100644
--- a/python/tvm/topi/index_put.py
+++ b/python/tvm/topi/index_put.py
@@ -96,7 +96,7 @@ def index_put(data, indices, values, accumulate=False):
with IRBuilder() as ib:
with T.parallel(0, full_range) as i:
- T.buffer_store(out, data[T.tensor_indices(data, i)], T.tensor_indices(out, i))
+ T.tensor_store(out, data[T.tensor_indices(data, i)], T.tensor_indices(out, i))
with T.parallel(0, index_len) as k:
# Decompose k into multi-dimensional broadcast index
@@ -140,10 +140,10 @@ def index_put(data, indices, values, accumulate=False):
return ib.get()
def update_func(dst_ptr, dst_index, update):
- T.buffer_store(dst_ptr, update, T.tensor_indices(dst_ptr, dst_index))
+ T.tensor_store(dst_ptr, update, T.tensor_indices(dst_ptr, dst_index))
def add_func(dst_ptr, dst_index, update):
- T.buffer_store(
+ T.tensor_store(
dst_ptr,
dst_ptr[T.tensor_indices(dst_ptr, dst_index)] + (update),
T.tensor_indices(dst_ptr, dst_index),
diff --git a/python/tvm/topi/scan.py b/python/tvm/topi/scan.py
index 07e7e33c86..02b5e43453 100644
--- a/python/tvm/topi/scan.py
+++ b/python/tvm/topi/scan.py
@@ -122,11 +122,11 @@ def scanop(
j = fused % axis_mul_after
base_idx = i * cumsum_axis_len * axis_mul_after + j
if exclusive:
- T.buffer_store(
+ T.tensor_store(
out_buf, cast(identity_value, dtype), T.tensor_indices(out_buf, base_idx)
)
else:
- T.buffer_store(
+ T.tensor_store(
out_buf,
maybe_cast(data_buf[T.tensor_indices(data_buf, base_idx)]),
T.tensor_indices(out_buf, base_idx),
@@ -136,7 +136,7 @@ def scanop(
cur_idx = base_idx + k * axis_mul_after
prev_idx = base_idx + (k - 1) * axis_mul_after
if exclusive:
- T.buffer_store(
+ T.tensor_store(
out_buf,
binop(
out_buf[T.tensor_indices(out_buf, prev_idx)],
@@ -145,7 +145,7 @@ def scanop(
T.tensor_indices(out_buf, cur_idx),
)
else:
- T.buffer_store(
+ T.tensor_store(
out_buf,
binop(
out_buf[T.tensor_indices(out_buf, prev_idx)],
diff --git a/python/tvm/topi/scatter.py b/python/tvm/topi/scatter.py
index 08cb1e5758..2b7ac9f899 100644
--- a/python/tvm/topi/scatter.py
+++ b/python/tvm/topi/scatter.py
@@ -118,7 +118,7 @@ def scatter_nd(data, indices, updates, mode):
with IRBuilder() as ib:
with T.serial(0, fused_shape) as i:
- T.buffer_store(out, data[T.tensor_indices(data, i)], T.tensor_indices(out, i))
+ T.tensor_store(out, data[T.tensor_indices(data, i)], T.tensor_indices(out, i))
with T.serial(0, fused_indices_dimension) as i:
with T.parallel(0, fused_updates_dimension) as j:
@@ -134,27 +134,27 @@ def scatter_nd(data, indices, updates, mode):
)
offset *= data_ptr.shape[l]
if mode == "update":
- T.buffer_store(
+ T.tensor_store(
out,
updates[T.tensor_indices(updates, i * fused_updates_dimension + j)],
T.tensor_indices(out, index),
)
elif mode == "add":
- T.buffer_store(
+ T.tensor_store(
out,
out[T.tensor_indices(out, index)]
+ (updates[T.tensor_indices(updates, i * fused_updates_dimension + j)]),
T.tensor_indices(out, index),
)
elif mode == "mul":
- T.buffer_store(
+ T.tensor_store(
out,
out[T.tensor_indices(out, index)]
* (updates[T.tensor_indices(updates, i * fused_updates_dimension + j)]),
T.tensor_indices(out, index),
)
elif mode == "min":
- T.buffer_store(
+ T.tensor_store(
out,
tirx.min(
out[T.tensor_indices(out, index)],
@@ -163,7 +163,7 @@ def scatter_nd(data, indices, updates, mode):
T.tensor_indices(out, index),
)
elif mode == "max":
- T.buffer_store(
+ T.tensor_store(
out,
tirx.max(
out[T.tensor_indices(out, index)],
diff --git a/python/tvm/topi/scatter_elements.py b/python/tvm/topi/scatter_elements.py
index 4f88dcc4aa..73e0715d6d 100644
--- a/python/tvm/topi/scatter_elements.py
+++ b/python/tvm/topi/scatter_elements.py
@@ -106,7 +106,7 @@ def scatter_elements(data, indices, updates, axis=0, reduction="update"):
# Copy initial input data to output
with IRBuilder() as ib:
with T.parallel(0, full_range) as i:
- T.buffer_store(out, data[T.tensor_indices(data, i)], T.tensor_indices(out, i))
+ T.tensor_store(out, data[T.tensor_indices(data, i)], T.tensor_indices(out, i))
with T.parallel(0, ind_before_axis_range * ind_after_axis_range) as fused:
i = fused // ind_after_axis_range
@@ -126,38 +126,38 @@ def scatter_elements(data, indices, updates, axis=0, reduction="update"):
return ib.get()
def update_func(dst_ptr, dst_index, update):
- T.buffer_store(dst_ptr, update, T.tensor_indices(dst_ptr, dst_index))
+ T.tensor_store(dst_ptr, update, T.tensor_indices(dst_ptr, dst_index))
def add_func(dst_ptr, dst_index, update):
- T.buffer_store(
+ T.tensor_store(
dst_ptr,
dst_ptr[T.tensor_indices(dst_ptr, dst_index)] + (update),
T.tensor_indices(dst_ptr, dst_index),
)
def mul_func(dst_ptr, dst_index, update):
- T.buffer_store(
+ T.tensor_store(
dst_ptr,
dst_ptr[T.tensor_indices(dst_ptr, dst_index)] * (update),
T.tensor_indices(dst_ptr, dst_index),
)
def mean_func(dst_ptr, dst_index, update):
- T.buffer_store(
+ T.tensor_store(
dst_ptr,
(dst_ptr[T.tensor_indices(dst_ptr, dst_index)] + update) / 2,
T.tensor_indices(dst_ptr, dst_index),
)
def min_func(dst_ptr, dst_index, update):
- T.buffer_store(
+ T.tensor_store(
dst_ptr,
tirx.min(dst_ptr[T.tensor_indices(dst_ptr, dst_index)], update),
T.tensor_indices(dst_ptr, dst_index),
)
def max_func(dst_ptr, dst_index, update):
- T.buffer_store(
+ T.tensor_store(
dst_ptr,
tirx.max(dst_ptr[T.tensor_indices(dst_ptr, dst_index)], update),
T.tensor_indices(dst_ptr, dst_index),
diff --git a/python/tvm/topi/searchsorted.py b/python/tvm/topi/searchsorted.py
index 3a117f83dc..d6d11db654 100644
--- a/python/tvm/topi/searchsorted.py
+++ b/python/tvm/topi/searchsorted.py
@@ -43,8 +43,8 @@ def binary_search(sequence_offset, search_range, sorted_sequence, value, right,
lo = lo_buf
hi = hi_buf
- T.buffer_store(lo, cast(0, out_dtype), T.tensor_indices(lo, 0))
- T.buffer_store(hi, cast(search_range, out_dtype), T.tensor_indices(hi, 0))
+ T.tensor_store(lo, cast(0, out_dtype), T.tensor_indices(lo, 0))
+ T.tensor_store(hi, cast(search_range, out_dtype), T.tensor_indices(hi, 0))
# Reference: pytorch/aten/src/ATen/native/cuda/Bucketization.cu
def condition(current_val, target_val):
@@ -62,9 +62,9 @@ def binary_search(sequence_offset, search_range, sorted_sequence, value, right,
)
):
with T.then_():
- T.buffer_store(lo, mid + 1, T.tensor_indices(lo, 0))
+ T.tensor_store(lo, mid + 1, T.tensor_indices(lo, 0))
with T.else_():
- T.buffer_store(hi, mid, T.tensor_indices(hi, 0))
+ T.tensor_store(hi, mid, T.tensor_indices(hi, 0))
return lo[T.tensor_indices(lo, 0)]
@@ -115,7 +115,7 @@ def searchsorted(sorted_sequence, values, right=False, out_dtype="int64"):
sequence_id = i // values_shape[-1]
sequence_offset = sequence_id * search_range
- T.buffer_store(
+ T.tensor_store(
indices,
binary_search(
sequence_offset,
diff --git a/python/tvm/topi/signal.py b/python/tvm/topi/signal.py
index d305357ef2..d2ea464ac5 100644
--- a/python/tvm/topi/signal.py
+++ b/python/tvm/topi/signal.py
@@ -90,18 +90,18 @@ def stft(
with col_loop(0, output_ptr.shape[2]) as col:
batch = tirx.floordiv(batch_row, output_ptr.shape[1])
row = tirx.floormod(batch_row, output_ptr.shape[1])
- T.buffer_store(
+ T.tensor_store(
output,
tirx.Cast(data_ptr.dtype, 0),
(batch, row, col, 0),
)
- T.buffer_store(
+ T.tensor_store(
output,
tirx.Cast(data_ptr.dtype, 0),
(batch, row, col, 1),
)
with T.serial(0, win_length) as wlen:
- T.buffer_store(
+ T.tensor_store(
output,
output[(batch, row, col, 0)]
+ (
@@ -111,7 +111,7 @@ def stft(
),
(batch, row, col, 0),
)
- T.buffer_store(
+ T.tensor_store(
output,
output[(batch, row, col, 1)]
- (
@@ -123,13 +123,13 @@ def stft(
)
with T.if_(normalized):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
output,
output[(batch, row, col, 0)]
/ (tirx.sqrt(tirx.const(n_fft, "float32"))),
(batch, row, col, 0),
)
- T.buffer_store(
+ T.tensor_store(
output,
output[(batch, row, col, 1)]
/ (tirx.sqrt(tirx.const(n_fft, "float32"))),
@@ -210,12 +210,12 @@ def dft(
base_idx = i * n_fft
with T.serial(0, n_fft) as n:
n_idx = base_idx + n
- T.buffer_store(
+ T.tensor_store(
re_output_ptr,
tirx.Cast(re_output_ptr.dtype, 0),
T.tensor_indices(re_output_ptr, n_idx),
)
- T.buffer_store(
+ T.tensor_store(
im_output_ptr,
tirx.Cast(im_output_ptr.dtype, 0),
T.tensor_indices(im_output_ptr, n_idx),
@@ -226,7 +226,7 @@ def dft(
w = _w * k
cos_w = tirx.Cast(re_output_ptr.dtype, tirx.cos(w))
sin_w = tirx.Cast(re_output_ptr.dtype, tirx.sin(w))
- T.buffer_store(
+ T.tensor_store(
re_output_ptr,
re_output_ptr[T.tensor_indices(re_output_ptr, n_idx)]
+ (
@@ -235,7 +235,7 @@ def dft(
),
T.tensor_indices(re_output_ptr, n_idx),
)
- T.buffer_store(
+ T.tensor_store(
im_output_ptr,
im_output_ptr[T.tensor_indices(im_output_ptr, n_idx)]
+ (
@@ -245,13 +245,13 @@ def dft(
T.tensor_indices(im_output_ptr, n_idx),
)
- T.buffer_store(
+ T.tensor_store(
re_output_ptr,
re_output_ptr[T.tensor_indices(re_output_ptr, n_idx)]
* (tirx.Cast(re_output_ptr.dtype, factor)),
T.tensor_indices(re_output_ptr, n_idx),
)
- T.buffer_store(
+ T.tensor_store(
im_output_ptr,
im_output_ptr[T.tensor_indices(im_output_ptr, n_idx)]
* (tirx.Cast(im_output_ptr.dtype, factor)),
diff --git a/python/tvm/topi/sparse_reshape.py b/python/tvm/topi/sparse_reshape.py
index 1024ec5760..eb483f80e8 100644
--- a/python/tvm/topi/sparse_reshape.py
+++ b/python/tvm/topi/sparse_reshape.py
@@ -104,34 +104,34 @@ def sparse_reshape(
equal_shape_buf = T.alloc_tensor([1], "bool", scope="local")
equal_shape = equal_shape_buf
- T.buffer_store(
+ T.tensor_store(
total_ele,
prev_shape[T.tensor_indices(prev_shape, 0)],
T.tensor_indices(total_ele, 0),
)
# Cumulative Reverse Exclusive Multiply
- T.buffer_store(
+ T.tensor_store(
multipliers,
Cast(new_shape_ptr.dtype, 1),
T.tensor_indices(multipliers, prev_shape_size - 1),
)
with T.serial(0, prev_shape_size - 1) as i_:
i = i_ + 1
- T.buffer_store(
+ T.tensor_store(
multipliers,
prev_shape[T.tensor_indices(prev_shape, prev_shape_size - i)]
* multipliers[T.tensor_indices(multipliers, prev_shape_size - i)],
T.tensor_indices(multipliers, prev_shape_size - 1 - i),
)
- T.buffer_store(
+ T.tensor_store(
total_ele,
total_ele[T.tensor_indices(total_ele, 0)]
* (prev_shape[T.tensor_indices(prev_shape, prev_shape_size - i)]),
T.tensor_indices(total_ele, 0),
)
- T.buffer_store(
+ T.tensor_store(
division_total_ele,
Cast(new_shape_ptr.dtype, 1),
T.tensor_indices(division_total_ele, 0),
@@ -139,7 +139,7 @@ def sparse_reshape(
with T.serial(0, new_shape_size) as i:
with T.if_(new_shape[T.tensor_indices(new_shape, i)] != -1):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
division_total_ele,
division_total_ele[T.tensor_indices(division_total_ele, 0)]
* (new_shape[T.tensor_indices(new_shape, i)]),
@@ -150,7 +150,7 @@ def sparse_reshape(
with T.serial(0, new_shape_size) as i:
with T.if_(new_shape[T.tensor_indices(new_shape, i)] == -1):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
out_new_shape,
Cast(
new_shape_ptr.dtype,
@@ -162,14 +162,14 @@ def sparse_reshape(
T.tensor_indices(out_new_shape, i),
)
with T.else_():
- T.buffer_store(
+ T.tensor_store(
out_new_shape,
new_shape[T.tensor_indices(new_shape, i)],
T.tensor_indices(out_new_shape, i),
)
# Check if prev_shape and new_shape are equal
- T.buffer_store(equal_shape, True, T.tensor_indices(equal_shape, 0))
+ T.tensor_store(equal_shape, True, T.tensor_indices(equal_shape, 0))
with T.if_(prev_shape_size == new_shape_size):
with T.then_():
with T.serial(0, prev_shape_size) as i:
@@ -178,16 +178,16 @@ def sparse_reshape(
!= out_new_shape[T.tensor_indices(out_new_shape, i)]
):
with T.then_():
- T.buffer_store(equal_shape, False, T.tensor_indices(equal_shape, 0))
+ T.tensor_store(equal_shape, False, T.tensor_indices(equal_shape, 0))
with T.else_():
- T.buffer_store(equal_shape, False, T.tensor_indices(equal_shape, 0))
+ T.tensor_store(equal_shape, False, T.tensor_indices(equal_shape, 0))
# Return same inputs if shapes are equal
with T.if_(equal_shape[T.tensor_indices(equal_shape, 0)]):
with T.then_():
with T.parallel(0, sparse_indices_ptr.shape[0]) as i:
with T.serial(0, sparse_indices_ptr.shape[1]) as j:
- T.buffer_store(
+ T.tensor_store(
new_sparse_indices,
sparse_indices[(i, j)],
(i, j),
@@ -195,14 +195,14 @@ def sparse_reshape(
# Else compute new_sparse_indices
with T.else_():
- T.buffer_store(
+ T.tensor_store(
dividers,
Cast(new_shape_ptr.dtype, 1),
T.tensor_indices(dividers, new_shape_size - 1),
)
with T.serial(0, new_shape_size - 1) as i_:
i = i_ + 1
- T.buffer_store(
+ T.tensor_store(
dividers,
dividers[T.tensor_indices(dividers, new_shape_size - i)]
* out_new_shape[T.tensor_indices(out_new_shape, new_shape_size - i)],
@@ -210,13 +210,13 @@ def sparse_reshape(
)
with T.parallel(0, sparse_indices_ptr.shape[0]) as i:
- T.buffer_store(
+ T.tensor_store(
flattened_indices,
Cast(new_shape_ptr.dtype, 0),
T.tensor_indices(flattened_indices, i),
)
with T.serial(0, sparse_indices_ptr.shape[1]) as j:
- T.buffer_store(
+ T.tensor_store(
flattened_indices,
flattened_indices[T.tensor_indices(flattened_indices, i)]
+ (
@@ -231,14 +231,14 @@ def sparse_reshape(
[1], new_shape_ptr.dtype, scope="local"
)
current_element = current_element_buf
- T.buffer_store(
+ T.tensor_store(
current_element,
flattened_indices[T.tensor_indices(flattened_indices, i)],
T.tensor_indices(current_element, 0),
)
with T.serial(0, new_sparse_indices_ptr.shape[1]) as j:
- T.buffer_store(
+ T.tensor_store(
new_sparse_indices,
Cast(
sparse_indices_ptr.dtype,
@@ -249,7 +249,7 @@ def sparse_reshape(
),
(i, j),
)
- T.buffer_store(
+ T.tensor_store(
current_element,
floormod(
current_element[T.tensor_indices(current_element, 0)],
diff --git a/python/tvm/topi/unique.py b/python/tvm/topi/unique.py
index 474181794a..250f11c790 100644
--- a/python/tvm/topi/unique.py
+++ b/python/tvm/topi/unique.py
@@ -46,9 +46,9 @@ def _calc_adjacent_diff_ir(data, output, binop=tirx.Sub):
with T.parallel(0, data.shape[0]) as i:
with T.if_(i == 0):
with T.then_():
- T.buffer_store(output_ptr, 0, T.tensor_indices(output_ptr, 0))
+ T.tensor_store(output_ptr, 0, T.tensor_indices(output_ptr, 0))
with T.else_():
- T.buffer_store(
+ T.tensor_store(
output_ptr,
tirx.Cast(
output.dtype,
diff --git a/python/tvm/topi/vision/nms.py b/python/tvm/topi/vision/nms.py
index 89b24f4031..1427c536bf 100644
--- a/python/tvm/topi/vision/nms.py
+++ b/python/tvm/topi/vision/nms.py
@@ -46,7 +46,7 @@ def _get_valid_counts_ir(
with IRBuilder() as ib:
with T.parallel(0, batch_size) as i:
- T.buffer_store(valid_count, T.int32(0), T.tensor_indices(valid_count, i))
+ T.tensor_store(valid_count, T.int32(0), T.tensor_indices(valid_count, i))
with T.serial(0, num_anchors) as j:
score = data[(i, j, score_index)]
@@ -59,25 +59,25 @@ def _get_valid_counts_ir(
with T.then_():
cur = valid_count[T.tensor_indices(valid_count, i)]
with T.serial(0, box_data_length) as k:
- T.buffer_store(
+ T.tensor_store(
out_tensor,
data[(i, j, k)],
(i, cur, k),
)
- T.buffer_store(out_indices, j, (i, cur))
- T.buffer_store(valid_count, cur + 1, T.tensor_indices(valid_count, i))
+ T.tensor_store(out_indices, j, (i, cur))
+ T.tensor_store(valid_count, cur + 1, T.tensor_indices(valid_count, i))
# Fill remaining slots with -1
with T.serial(0, num_anchors) as j:
with T.if_(j >= valid_count[T.tensor_indices(valid_count, i)]):
with T.then_():
with T.serial(0, box_data_length) as k:
- T.buffer_store(
+ T.tensor_store(
out_tensor,
tvm.tirx.Cast(data.dtype, T.float32(-1.0)),
(i, j, k),
)
- T.buffer_store(out_indices, T.int32(-1), (i, j))
+ T.tensor_store(out_indices, T.int32(-1), (i, j))
return ib.get()
@@ -214,7 +214,7 @@ def _classic_nms_ir(
# Step 1: Reorder data by sorted score
nkeep_buf = T.alloc_tensor((1,), "int32", scope="local")
nkeep_local = nkeep_buf
- T.buffer_store(
+ T.tensor_store(
nkeep_local,
valid_count[T.tensor_indices(valid_count, i)],
T.tensor_indices(nkeep_local, 0),
@@ -223,7 +223,7 @@ def _classic_nms_ir(
tvm.tirx.all(top_k > 0, top_k < nkeep_local[T.tensor_indices(nkeep_local, 0)])
):
with T.then_():
- T.buffer_store(nkeep_local, top_k, T.tensor_indices(nkeep_local, 0))
+ T.tensor_store(nkeep_local, top_k, T.tensor_indices(nkeep_local, 0))
# Copy sorted boxes to output
with T.serial(0, num_anchors) as j:
@@ -231,29 +231,29 @@ def _classic_nms_ir(
with T.then_():
src_idx = sorted_index[(i, j)]
with T.serial(0, box_data_length) as k:
- T.buffer_store(
+ T.tensor_store(
out_data,
data[(i, src_idx, k)],
(i, j, k),
)
- T.buffer_store(
+ T.tensor_store(
out_box_indices,
sorted_index[(i, j)],
(i, j),
)
with T.else_():
with T.serial(0, box_data_length) as k:
- T.buffer_store(
+ T.tensor_store(
out_data,
tvm.tirx.Cast(data.dtype, T.float32(-1.0)),
(i, j, k),
)
- T.buffer_store(out_box_indices, T.int32(-1), (i, j))
+ T.tensor_store(out_box_indices, T.int32(-1), (i, j))
# Step 2: Apply NMS - greedy suppression
num_valid_boxes_buf = T.alloc_tensor((1,), "int32", scope="local")
num_valid_boxes = num_valid_boxes_buf
- T.buffer_store(num_valid_boxes, T.int32(0), T.tensor_indices(num_valid_boxes, 0))
+ T.tensor_store(num_valid_boxes, T.int32(0), T.tensor_indices(num_valid_boxes, 0))
best_idx_buf = T.alloc_tensor((1,), "int32", scope="local")
best_idx = best_idx_buf
best_score_buf = T.alloc_tensor((1,), data.dtype, scope="local")
@@ -320,8 +320,8 @@ def _classic_nms_ir(
)
):
with T.then_():
- T.buffer_store(best_idx, T.int32(-1), T.tensor_indices(best_idx, 0))
- T.buffer_store(best_score, thresh, T.tensor_indices(best_score, 0))
+ T.tensor_store(best_idx, T.int32(-1), T.tensor_indices(best_idx, 0))
+ T.tensor_store(best_score, thresh, T.tensor_indices(best_score, 0))
with T.serial(0, nkeep_local[T.tensor_indices(nkeep_local, 0)]) as j:
with T.if_(
@@ -333,8 +333,8 @@ def _classic_nms_ir(
)
):
with T.then_():
- T.buffer_store(best_idx, j, T.tensor_indices(best_idx, 0))
- T.buffer_store(
+ T.tensor_store(best_idx, j, T.tensor_indices(best_idx, 0))
+ T.tensor_store(
best_score,
out_data[(i, j, score_index)],
T.tensor_indices(best_score, 0),
@@ -347,7 +347,7 @@ def _classic_nms_ir(
!= num_valid_boxes[T.tensor_indices(num_valid_boxes, 0)]
):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
tmp_idx,
out_box_indices[
(
@@ -359,7 +359,7 @@ def _classic_nms_ir(
],
T.tensor_indices(tmp_idx, 0),
)
- T.buffer_store(
+ T.tensor_store(
out_box_indices,
out_box_indices[
(
@@ -374,14 +374,14 @@ def _classic_nms_ir(
],
),
)
- T.buffer_store(
+ T.tensor_store(
out_box_indices,
tmp_idx[T.tensor_indices(tmp_idx, 0)],
(i, best_idx[T.tensor_indices(best_idx, 0)]),
)
with T.serial(0, box_data_length) as k:
- T.buffer_store(
+ T.tensor_store(
tmp_val,
out_data[
(
@@ -396,7 +396,7 @@ def _classic_nms_ir(
],
T.tensor_indices(tmp_val, 0),
)
- T.buffer_store(
+ T.tensor_store(
out_data,
out_data[
(
@@ -413,7 +413,7 @@ def _classic_nms_ir(
k,
),
)
- T.buffer_store(
+ T.tensor_store(
out_data,
tmp_val[T.tensor_indices(tmp_val, 0)],
(
@@ -471,14 +471,14 @@ def _classic_nms_ir(
with T.if_(iou >= iou_threshold):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
out_box_indices,
T.int32(-1),
(i, j),
)
with T.if_(iou < iou_threshold):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
out_data,
out_data[(i, j, score_index)]
* tvm.tirx.exp(
@@ -491,20 +491,20 @@ def _classic_nms_ir(
<= thresh
):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
out_box_indices,
T.int32(-1),
(i, j),
)
- T.buffer_store(
+ T.tensor_store(
num_valid_boxes,
num_valid_boxes[T.tensor_indices(num_valid_boxes, 0)] + 1,
T.tensor_indices(num_valid_boxes, 0),
)
if return_indices:
- T.buffer_store(
+ T.tensor_store(
out_valid_box_count,
num_valid_boxes[T.tensor_indices(num_valid_boxes, 0)],
(i, 0),
@@ -514,7 +514,7 @@ def _classic_nms_ir(
with T.if_(j < num_valid_boxes[T.tensor_indices(num_valid_boxes, 0)]):
with T.then_():
orig_idx = out_box_indices[(i, j)]
- T.buffer_store(
+ T.tensor_store(
out_box_indices,
indices[(i, orig_idx)],
(i, j),
@@ -522,12 +522,12 @@ def _classic_nms_ir(
with T.if_(j >= num_valid_boxes[T.tensor_indices(num_valid_boxes, 0)]):
with T.then_():
with T.serial(0, box_data_length) as k:
- T.buffer_store(
+ T.tensor_store(
out_data,
tvm.tirx.Cast(data.dtype, T.float32(-1.0)),
(i, j, k),
)
- T.buffer_store(
+ T.tensor_store(
out_box_indices,
T.int32(-1),
(i, j),
@@ -537,7 +537,7 @@ def _classic_nms_ir(
with T.if_(j >= num_valid_boxes[T.tensor_indices(num_valid_boxes, 0)]):
with T.then_():
with T.serial(0, box_data_length) as k:
- T.buffer_store(
+ T.tensor_store(
out_data,
tvm.tirx.Cast(data.dtype, T.float32(-1.0)),
(i, j, k),
@@ -556,7 +556,7 @@ def _classic_nms_ir(
)
):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
num_valid_boxes,
num_valid_boxes[T.tensor_indices(num_valid_boxes, 0)] + 1,
T.tensor_indices(num_valid_boxes, 0),
@@ -587,14 +587,14 @@ def _classic_nms_ir(
with T.if_(iou >= iou_threshold):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
out_data,
tvm.tirx.Cast(
data.dtype, T.float32(-1.0)
),
(i, k, score_index),
)
- T.buffer_store(
+ T.tensor_store(
out_box_indices,
T.int32(-1),
(i, k),
@@ -602,12 +602,12 @@ def _classic_nms_ir(
with T.else_():
with T.serial(0, box_data_length) as k:
- T.buffer_store(
+ T.tensor_store(
out_data,
tvm.tirx.Cast(data.dtype, T.float32(-1.0)),
(i, j, k),
)
- T.buffer_store(
+ T.tensor_store(
out_box_indices,
T.int32(-1),
(i, j),
@@ -616,24 +616,24 @@ def _classic_nms_ir(
if return_indices:
valid_idx_buf = T.alloc_tensor((1,), "int32", scope="local")
valid_idx = valid_idx_buf
- T.buffer_store(valid_idx, T.int32(0), T.tensor_indices(valid_idx, 0))
+ T.tensor_store(valid_idx, T.int32(0), T.tensor_indices(valid_idx, 0))
with T.serial(0, num_anchors) as j:
with T.if_(out_box_indices[(i, j)] >= 0):
with T.then_():
orig_idx = out_box_indices[(i, j)]
- T.buffer_store(
+ T.tensor_store(
out_box_indices,
indices[(i, orig_idx)],
(i, valid_idx[T.tensor_indices(valid_idx, 0)]),
)
- T.buffer_store(
+ T.tensor_store(
valid_idx,
valid_idx[T.tensor_indices(valid_idx, 0)] + 1,
T.tensor_indices(valid_idx, 0),
)
- T.buffer_store(
+ T.tensor_store(
out_valid_box_count,
valid_idx[T.tensor_indices(valid_idx, 0)],
(i, 0),
@@ -642,7 +642,7 @@ def _classic_nms_ir(
with T.serial(0, num_anchors) as j:
with T.if_(j >= valid_idx[T.tensor_indices(valid_idx, 0)]):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
out_box_indices,
T.int32(-1),
(i, j),
@@ -853,7 +853,7 @@ def _rearrange_out(data, batch_size, num_anchors, box_data_length, score_index):
with T.parallel(0, batch_size) as i:
valid_idx_buf = T.alloc_tensor((1,), "int32", scope="local")
valid_idx = valid_idx_buf
- T.buffer_store(valid_idx, T.int32(0), T.tensor_indices(valid_idx, 0))
+ T.tensor_store(valid_idx, T.int32(0), T.tensor_indices(valid_idx, 0))
with T.serial(0, num_anchors) as j:
with T.if_(
@@ -861,12 +861,12 @@ def _rearrange_out(data, batch_size, num_anchors, box_data_length, score_index):
):
with T.then_():
with T.serial(0, box_data_length) as k:
- T.buffer_store(
+ T.tensor_store(
out,
data[(i, j, k)],
(i, valid_idx[T.tensor_indices(valid_idx, 0)], k),
)
- T.buffer_store(
+ T.tensor_store(
valid_idx,
valid_idx[T.tensor_indices(valid_idx, 0)] + 1,
T.tensor_indices(valid_idx, 0),
@@ -876,7 +876,7 @@ def _rearrange_out(data, batch_size, num_anchors, box_data_length, score_index):
with T.if_(j >= valid_idx[T.tensor_indices(valid_idx, 0)]):
with T.then_():
with T.serial(0, box_data_length) as k:
- T.buffer_store(
+ T.tensor_store(
out,
tvm.tirx.Cast(data.dtype, T.float32(-1.0)),
(i, j, k),
@@ -915,7 +915,7 @@ def _nms_loop(
on_new_valid_box_func(
0, num_valid_boxes_local[T.tensor_indices(num_valid_boxes_local, 0)], i, j
)
- T.buffer_store(
+ T.tensor_store(
num_valid_boxes_local,
num_valid_boxes_local[T.tensor_indices(num_valid_boxes_local, 0)] + 1,
T.tensor_indices(num_valid_boxes_local, 0),
@@ -938,7 +938,7 @@ def _nms_loop(
with T.if_(iou >= iou_threshold):
with T.then_():
- T.buffer_store(out_scores, T.float32(-1.0), (i, k))
+ T.tensor_store(out_scores, T.float32(-1.0), (i, k))
on_new_invalidated_box_func(i, k)
with T.serial(0, batch_size) as i:
@@ -948,7 +948,7 @@ def _nms_loop(
with T.then_():
num_valid_boxes_local_buf = T.alloc_tensor((1,), "int32", scope="local")
num_valid_boxes_local = num_valid_boxes_local_buf
- T.buffer_store(
+ T.tensor_store(
num_valid_boxes_local, T.int32(0), T.tensor_indices(num_valid_boxes_local, 0)
)
@@ -968,14 +968,14 @@ def _nms_loop(
else:
nms_inner_loop(i, j, nkeep, num_valid_boxes_local)
- T.buffer_store(
+ T.tensor_store(
num_valid_boxes,
num_valid_boxes_local[T.tensor_indices(num_valid_boxes_local, 0)],
T.tensor_indices(num_valid_boxes, i),
)
with T.else_():
- T.buffer_store(num_valid_boxes, T.int32(0), T.tensor_indices(num_valid_boxes, i))
+ T.tensor_store(num_valid_boxes, T.int32(0), T.tensor_indices(num_valid_boxes, i))
def _get_valid_box_count(scores, score_threshold):
@@ -1067,7 +1067,7 @@ def _collect_selected_indices_ir(
max_output_rows = batch_classes * 10
with T.serial(0, max_output_rows) as init_i:
with T.serial(0, 3) as init_j: # 3 columns
- T.buffer_store(out, cast(0, "int64"), (init_i, init_j))
+ T.tensor_store(out, cast(0, "int64"), (init_i, init_j))
with T.parallel(0, batch_classes) as i:
i_64 = cast(i, "int64")
@@ -1088,9 +1088,9 @@ def _collect_selected_indices_ir(
limit = num_detections[i]
with T.serial(0, limit) as j:
- T.buffer_store(out, batch_id, (row_offsets[i] + j, 0))
- T.buffer_store(out, class_id, (row_offsets[i] + j, 1))
- T.buffer_store(
+ T.tensor_store(out, batch_id, (row_offsets[i] + j, 0))
+ T.tensor_store(out, class_id, (row_offsets[i] + j, 1))
+ T.tensor_store(
out,
cast(selected_indices[i, j], "int64"),
(row_offsets[i] + j, 2),
@@ -1123,17 +1123,17 @@ def _collect_selected_indices_and_scores_ir(
with T.if_(j < num_detections[batch_id, class_id]):
with T.then_():
offset = row_offsets[batch_id, class_id] + j
- T.buffer_store(
+ T.tensor_store(
collected_indices,
class_id,
(batch_id, offset, 0),
)
- T.buffer_store(
+ T.tensor_store(
collected_indices,
cast(selected_indices[i, j], "int64"),
(batch_id, offset, 1),
)
- T.buffer_store(
+ T.tensor_store(
collected_scores,
selected_scores[i, j],
(batch_id, offset),
@@ -1146,17 +1146,17 @@ def _collect_selected_indices_and_scores_ir(
+ j
- num_detections[batch_id, class_id]
)
- T.buffer_store(
+ T.tensor_store(
collected_indices,
zero,
(batch_id, offset, 0),
)
- T.buffer_store(
+ T.tensor_store(
collected_indices,
zero,
(batch_id, offset, 1),
)
- T.buffer_store(
+ T.tensor_store(
collected_scores,
T.float32(0.0),
(batch_id, offset),
diff --git a/python/tvm/topi/vision/nms_util.py b/python/tvm/topi/vision/nms_util.py
index 1bac2c73da..3227479ff5 100644
--- a/python/tvm/topi/vision/nms_util.py
+++ b/python/tvm/topi/vision/nms_util.py
@@ -71,16 +71,16 @@ def binary_search(y, num_boxes, scores, score_threshold, out):
hi_buf = T.decl_tensor([1], "int32", scope="local")
lo = lo_buf
hi = hi_buf
- T.buffer_store(lo, T.int32(0), T.tensor_indices(lo, 0))
- T.buffer_store(hi, tvm.tirx.Cast("int32", num_boxes), T.tensor_indices(hi, 0))
+ T.tensor_store(lo, T.int32(0), T.tensor_indices(lo, 0))
+ T.tensor_store(hi, tvm.tirx.Cast("int32", num_boxes), T.tensor_indices(hi, 0))
with T.while_(lo[T.tensor_indices(lo, 0)] < hi[T.tensor_indices(hi, 0)]):
mid = (hi[T.tensor_indices(hi, 0)] + lo[T.tensor_indices(lo, 0)]) >> 1
with T.if_(scores[y, mid] > score_threshold):
with T.then_():
- T.buffer_store(lo, mid + 1, T.tensor_indices(lo, 0))
+ T.tensor_store(lo, mid + 1, T.tensor_indices(lo, 0))
with T.else_():
- T.buffer_store(hi, mid, T.tensor_indices(hi, 0))
- T.buffer_store(out, lo[T.tensor_indices(lo, 0)], T.tensor_indices(out, y))
+ T.tensor_store(hi, mid, T.tensor_indices(hi, 0))
+ T.tensor_store(out, lo[T.tensor_indices(lo, 0)], T.tensor_indices(out, y))
def _estimate_max_detections(batch_class, input_image_size=None):
@@ -338,13 +338,13 @@ def _all_class_nms_ir(
def on_new_valid_box(tid, num_current_valid_box, i, j):
with T.if_(tid + 0 == 0):
with T.then_():
- T.buffer_store(
+ T.tensor_store(
box_indices,
sorted_indices[i, j],
(i, num_current_valid_box),
)
if selected_scores is not None:
- T.buffer_store(
+ T.tensor_store(
selected_scores,
sorted_scores[i, j],
(i, num_current_valid_box),
diff --git a/src/backend/cuda/transforms/lower_iket.cc b/src/backend/cuda/transforms/lower_iket.cc
index 136b6e2664..3a94284601 100644
--- a/src/backend/cuda/transforms/lower_iket.cc
+++ b/src/backend/cuda/transforms/lower_iket.cc
@@ -386,7 +386,7 @@ class TokenBufferCollector : public StmtExprVisitor {
bool changed{false};
private:
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* store) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* store) final {
bool is_token_value = false;
if (const auto* call = store->value.as<CallNode>()) {
is_token_value = IsTokenProducer(call);
@@ -421,7 +421,7 @@ class TokenDeclarationCollector : public StmtExprVisitor {
bool changed{false};
private:
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* store) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* store) final {
std::set<DeclarationKey> possible;
if (const auto* call = store->value.as<CallNode>(); call && IsTokenProducer(call)) {
possible.insert(DeclarationKey{DeclarationKind::kRange, GetName(call)});
@@ -514,7 +514,7 @@ class TokenVerifier : public StmtExprVisitor {
explicit TokenVerifier(const TokenBufferSet& token_buffers) : token_buffers_(token_buffers) {}
private:
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* store) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* store) final {
if (!token_buffers_.count(store->buffer.get())) {
return StmtExprVisitor::Visit_(store);
}
@@ -604,7 +604,7 @@ class StripIket : public StmtExprMutator {
return StmtExprMutator::Mutate_(alloc, inplace_mode);
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* store, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* store, InplaceMode inplace_mode) final {
if (token_buffers_.count(store->buffer.get())) return Evaluate(0);
return StmtExprMutator::Mutate_(store, inplace_mode);
}
diff --git a/src/backend/opencl/runtime/opencl_common.h b/src/backend/opencl/runtime/opencl_common.h
index 6954d607cf..f7564f7e40 100644
--- a/src/backend/opencl/runtime/opencl_common.h
+++ b/src/backend/opencl/runtime/opencl_common.h
@@ -421,7 +421,7 @@ class OpenCLThreadEntry {
/*! \brief OpenCL runtime buffer structure with tracked memory layout
TODO(tvm-team): Uncouple use of storage scope and data layout by using the transform_layout
schedule primitive to express the desired texture layout. This will require supporting Nd
- indices in TensorLoad and BufferStore in CodegenOpenCL, and ensuring Nd allocations for
+ indices in TensorLoad and TensorStore in CodegenOpenCL, and ensuring Nd allocations for
texture are correctly routed to the AllocateTexture packed function in the OpenCL DeviceAPI.
*/
struct BufferDescriptor {
diff --git a/src/backend/trn/codegen/codegen_trn.cc b/src/backend/trn/codegen/codegen_trn.cc
index 4c8ac7e61d..b33badb6ba 100644
--- a/src/backend/trn/codegen/codegen_trn.cc
+++ b/src/backend/trn/codegen/codegen_trn.cc
@@ -351,7 +351,7 @@ std::string CodeGenTrainium::PrintIndices(const Array<PrimExpr>& indices) {
return os.str();
}
-void CodeGenTrainium::Dispatch_(const BufferStoreNode* op) {
+void CodeGenTrainium::Dispatch_(const TensorStoreNode* op) {
LOG(FATAL) << "Trainium codegen does not support buffer store";
}
diff --git a/src/backend/trn/codegen/codegen_trn.h b/src/backend/trn/codegen/codegen_trn.h
index 45a1f04ce4..2ea5b32409 100644
--- a/src/backend/trn/codegen/codegen_trn.h
+++ b/src/backend/trn/codegen/codegen_trn.h
@@ -62,7 +62,7 @@ class CodeGenTrainium final : public CodeGenC {
void DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call);
void Dispatch_(const RegionStmtNode* op) final; // NOLINT(*)
void Dispatch_(const ForNode* op) final; // NOLINT(*)
- void Dispatch_(const BufferStoreNode* op) final; // NOLINT(*)=
+ void Dispatch_(const TensorStoreNode* op) final; // NOLINT(*)=
void Dispatch_(const EvaluateNode* op) final; // NOLINT(*)
std::string PrintIndices(const ffi::Array<PrimExpr>& indices); // NOLINT(*)
void Dispatch_(const TensorLoadNode* op, std::ostream& os) final; // NOLINT(*)
diff --git a/src/backend/trn/transform/lower_trainium_layout.cc b/src/backend/trn/transform/lower_trainium_layout.cc
index 822943f94f..ea3c8a71ed 100644
--- a/src/backend/trn/transform/lower_trainium_layout.cc
+++ b/src/backend/trn/transform/lower_trainium_layout.cc
@@ -214,12 +214,12 @@ class TrainiumLayoutApplier : public tirx::IRMutatorWithAnalyzer {
return flattened;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
// Index conversion needs the original logical layout after the parent remaps the buffer.
TensorVar logical_buffer = op->buffer;
- BufferStore store = StmtExprMutator::Mutate_(op, inplace_mode)
+ TensorStore store = StmtExprMutator::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<BufferStore>();
+ .as_or_throw<TensorStore>();
PrimType store_value_ty = op->value.ty();
bool store_returns_bool = store_value_ty.MatchesCode(DLDataTypeCode::kDLBool);
store = VisitBufferAccess(store, logical_buffer);
@@ -289,7 +289,7 @@ class TrainiumLayoutApplier : public tirx::IRMutatorWithAnalyzer {
return res;
}
if (tile_layout && tile_layout->HasThreadAxis()) {
- LOG(FATAL) << "Cannot lower direct TensorLoad/BufferStore on a buffer with thread-axis "
+ LOG(FATAL) << "Cannot lower direct TensorLoad/TensorStore on a buffer with thread-axis "
<< "layout: unable to verify that the coordinate matches the current thread. "
<< "Use .view() + .local() to decompose thread and memory axes.";
}
@@ -302,7 +302,7 @@ class TrainiumLayoutApplier : public tirx::IRMutatorWithAnalyzer {
return {analyzer_->Simplify(flattened_indices[0])};
}
- BufferStore VisitBufferAccess(BufferStore node, const TensorVar& logical_buffer) {
+ TensorStore VisitBufferAccess(TensorStore node, const TensorVar& logical_buffer) {
TVM_FFI_ICHECK(logical_buffer.defined());
if (!logical_buffer->layout.has_value()) {
return node;
diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc b/src/backend/vulkan/codegen/codegen_spirv.cc
index b38b219bad..9f2c295f91 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.cc
+++ b/src/backend/vulkan/codegen/codegen_spirv.cc
@@ -712,7 +712,7 @@ spirv::Value CodeGenSPIRV::Dispatch_(const prim::ShuffleNode* op) {
return element;
}
-void CodeGenSPIRV::Dispatch_(const BufferStoreNode* op) {
+void CodeGenSPIRV::Dispatch_(const TensorStoreNode* op) {
TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "SPIR-V codegen expects flat memory buffers";
Var buffer_var = op->buffer.var();
PrimExpr prim_index = op->indices[0];
diff --git a/src/backend/vulkan/codegen/codegen_spirv.h b/src/backend/vulkan/codegen/codegen_spirv.h
index 867fac6d5c..67208de6ed 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.h
+++ b/src/backend/vulkan/codegen/codegen_spirv.h
@@ -114,7 +114,7 @@ class CodeGenSPIRV : public tirx::ExprFunctor<spirv::Value(const Expr&)>,
spirv::Value Dispatch_(const TensorLoadNode* op) override;
spirv::Value Dispatch_(const prim::ShuffleNode* op) override;
// stmt
- void Dispatch_(const BufferStoreNode* op) override;
+ void Dispatch_(const TensorStoreNode* op) override;
void Dispatch_(const ForNode* op) override;
void Dispatch_(const WhileNode* op) override;
void Dispatch_(const IfThenElseNode* op) override;
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc b/src/backend/webgpu/codegen/codegen_webgpu.cc
index 86eeb8e127..85a1a7f5d0 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -116,7 +116,7 @@ class WebGPUWorkgroupInfoCollector : public StmtExprVisitor {
return std::nullopt;
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
info_.write_access_set.insert(ResolveBuffer(op->buffer.var()));
@@ -705,7 +705,7 @@ void CodeGenWebGPU::Dispatch_(const BindNode* op) {
}
}
-void CodeGenWebGPU::Dispatch_(const BufferStoreNode* op) {
+void CodeGenWebGPU::Dispatch_(const TensorStoreNode* op) {
TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "Store to non-flat memory not supported.";
PrimType value_ty = op->value.ty();
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.h b/src/backend/webgpu/codegen/codegen_webgpu.h
index 1c3bfd2caa..6eccd1a6e8 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.h
+++ b/src/backend/webgpu/codegen/codegen_webgpu.h
@@ -81,7 +81,7 @@ class CodeGenWebGPU final : public CodeGenC {
// stmt printing
void Dispatch_(const RegionStmtNode* op) final;
void Dispatch_(const BindNode* op) final;
- void Dispatch_(const BufferStoreNode* op) final;
+ void Dispatch_(const TensorStoreNode* op) final;
void Dispatch_(const ForNode* op) final;
void DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call);
void Dispatch_(const AssertStmtNode* op) final;
diff --git a/src/relax/analysis/layout_transformation.cc b/src/relax/analysis/layout_transformation.cc
index 7c5b058eec..d1268743fa 100644
--- a/src/relax/analysis/layout_transformation.cc
+++ b/src/relax/analysis/layout_transformation.cc
@@ -491,7 +491,7 @@ class BlockAnalyzer : public s_tir::StmtExprVisitor {
return std::nullopt;
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op));
BufferAccessInfo& access_info = buffer_access_info_[op->buffer];
diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc
index 034e01c099..13f0a9bdd3 100644
--- a/src/relax/analysis/tir_op_pattern_kind.cc
+++ b/src/relax/analysis/tir_op_pattern_kind.cc
@@ -60,14 +60,14 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor {
return false;
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
// We only support one buffer store in a block (usually generated by TE compute)
// If we have already seen buffer store in the current block, classify as Opaque.
if (store_.has_value() && !IsSameArray(op->indices, store_.value()->indices)) {
kind_ = kOpaque;
return std::nullopt;
}
- store_ = ffi::GetRef<BufferStore>(op);
+ store_ = ffi::GetRef<TensorStore>(op);
return s_tir::StmtExprVisitor::Visit_(op);
}
@@ -95,7 +95,7 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor {
return std::nullopt;
}
- BufferStore store = store_.value();
+ TensorStore store = store_.value();
// Step 3. Checking load store indices pattern
OpPatternKind index_pair_pattern = kElemWise;
@@ -188,7 +188,7 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor {
* It's elemwise pattern iff load indices and store indices are the same.
* E.g A[i, j] = B[i, j]
*/
- static bool IsElemwisePattern(const BufferStore& store, const TensorLoad& load) {
+ static bool IsElemwisePattern(const TensorStore& store, const TensorLoad& load) {
return IsSameArray(store->indices, load->indices);
}
@@ -199,7 +199,7 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor {
* A[i, j] = B[i, k] is not broadcast since `k` are not in the store indices.
* A[i, j] = B[j, i] is not broadcast the load indices are not in the same order as store's
*/
- static bool IsBroadcastPattern(const BufferStore& store, const TensorLoad& load) {
+ static bool IsBroadcastPattern(const TensorStore& store, const TensorLoad& load) {
size_t ndim_load_buf = load->source.as_or_throw<tvm::tirx::TensorVar>()->shape.size();
size_t ndim_store_buf = store->buffer->shape.size();
@@ -231,7 +231,7 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor {
* E.g. A[i, j] = B[j, i] is injective.
* A[i, j] = B[i - j] is injective since the load index vars are only i, j
*/
- static bool IsInjectivePattern(const BufferStore& store, const TensorLoad& load) {
+ static bool IsInjectivePattern(const TensorStore& store, const TensorLoad& load) {
std::unordered_set<const tirx::VarNode*> vars;
for (const PrimExpr& store_index : store->indices) {
if (auto var = store_index.as<PrimVar>()) {
@@ -259,7 +259,7 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor {
* E.g. Store = A[i, j] and Load = B[i, j, k] allow data reuse.
* Store = A[i, j] and Load = B[i, j + k] allow data reuse.
*/
- static bool IsAllowReusePattern(const BufferStore& store, const TensorLoad& load) {
+ static bool IsAllowReusePattern(const TensorStore& store, const TensorLoad& load) {
std::unordered_set<const tirx::VarNode*> vars;
for (const PrimExpr& index : store->indices) {
if (auto var = index.as<PrimVar>()) {
@@ -293,7 +293,7 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor {
/*! \brief Checking if the stmt is multiply add. E.g. C[i, j] += A[i, k] * B[j, k] */
static bool IsFMA(const Stmt& body) {
- if (const auto* store = body.as<BufferStoreNode>()) {
+ if (const auto* store = body.as<TensorStoreNode>()) {
if (const auto* add = RemoveCast(store->value).as<prim::AddNode>()) {
if (const auto* mul = RemoveCast(add->b).as<prim::MulNode>()) {
const auto* store_lhs = RemoveCast(add->a).as<TensorLoadNode>();
@@ -307,9 +307,9 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor {
if (!lhs || !rhs) {
return false;
}
- return IsAllowReusePattern(ffi::GetRef<BufferStore>(store),
+ return IsAllowReusePattern(ffi::GetRef<TensorStore>(store),
ffi::GetRef<TensorLoad>(lhs)) &&
- IsAllowReusePattern(ffi::GetRef<BufferStore>(store), ffi::GetRef<TensorLoad>(rhs));
+ IsAllowReusePattern(ffi::GetRef<TensorStore>(store), ffi::GetRef<TensorLoad>(rhs));
}
}
}
@@ -349,10 +349,10 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor {
private:
/*!
- * \brief The BufferStore node in the current block.
- * \note We only support one BufferStore node in a block (usually generated by TE compute)
+ * \brief The TensorStore node in the current block.
+ * \note We only support one TensorStore node in a block (usually generated by TE compute)
*/
- ffi::Optional<BufferStore> store_;
+ ffi::Optional<TensorStore> store_;
/*! \brief The TensorLoad nodes in the current block. */
ffi::Array<TensorLoad> loads_;
/*! \brief The result of op pattern. */
@@ -434,18 +434,18 @@ bool HasReshapePattern(const PrimFunc& func) {
// Step 1. Get the load/store pattern of the block body.
// To detect the reshape pattern, we require the block body to be a
- // BufferStore, which has a TensorLoad as value.
- const auto* buffer_store = block->body.as<BufferStoreNode>();
- if (buffer_store == nullptr) {
+ // TensorStore, which has a TensorLoad as value.
+ const auto* tensor_store = block->body.as<TensorStoreNode>();
+ if (tensor_store == nullptr) {
return std::nullopt;
}
- const auto* buffer_load = buffer_store->value.as<TensorLoadNode>();
+ const auto* buffer_load = tensor_store->value.as<TensorLoadNode>();
if (buffer_load == nullptr) {
return std::nullopt;
}
// Further, we require the buffer being stored and being loaded to
// match the parameter of the PrimFunc, namely `dst_buffer_` and `src_buffer_`.
- if (!(buffer_store->buffer.same_as(dst_buffer_) &&
+ if (!(tensor_store->buffer.same_as(dst_buffer_) &&
buffer_load->source.as_or_throw<tvm::tirx::TensorVar>().same_as(src_buffer_))) {
return std::nullopt;
}
@@ -490,11 +490,11 @@ bool HasReshapePattern(const PrimFunc& func) {
ffi::Array<PrimExpr> nontrivial_indices{nullptr};
ffi::Optional<TensorVar> nontrivial_buffer;
- if (f_is_trivial_indices(dst_buffer_, buffer_store->indices)) {
+ if (f_is_trivial_indices(dst_buffer_, tensor_store->indices)) {
nontrivial_indices = buffer_load->indices;
nontrivial_buffer = src_buffer_;
} else if (f_is_trivial_indices(src_buffer_, buffer_load->indices)) {
- nontrivial_indices = buffer_store->indices;
+ nontrivial_indices = tensor_store->indices;
nontrivial_buffer = dst_buffer_;
}
@@ -545,7 +545,7 @@ bool HasReshapePattern(const PrimFunc& func) {
// Apply check 2 as followup when check 1 is not satisfied.
// Calculate the flattened access index according to the load/store pattern.
PrimExpr src_idx = f_calc_flattened_idx(src_buffer_, buffer_load->indices);
- PrimExpr dst_idx = f_calc_flattened_idx(dst_buffer_, buffer_store->indices);
+ PrimExpr dst_idx = f_calc_flattened_idx(dst_buffer_, tensor_store->indices);
// Check if we can prove the equality of flattened indices.
if (ana_->CanProveEqual(src_idx, dst_idx)) {
this->is_reshape_ = true;
diff --git a/src/relax/backend/vm/vm_shape_lower.cc b/src/relax/backend/vm/vm_shape_lower.cc
index 4a3cbf859e..43ec5f7be4 100644
--- a/src/relax/backend/vm/vm_shape_lower.cc
+++ b/src/relax/backend/vm/vm_shape_lower.cc
@@ -735,7 +735,7 @@ class VMShapeLowerMutator
PrimExpr value = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(slot->expr, f_substitute)
.as_or_throw<PrimExpr>();
seq.push_back(
- tirx::BufferStore(buffer, value, {IntImm(tvm::PrimType(ShapeDType()), slot->index)}));
+ tirx::TensorStore(buffer, value, {IntImm(tvm::PrimType(ShapeDType()), slot->index)}));
}
tirx::Stmt body = tirx::SeqStmt::Flatten(seq);
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 d8652468a8..5701d2c964 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
@@ -55,7 +55,7 @@ class DistBufferReplacer : public s_tir::StmtExprMutator {
class DistSBlockInfoCollector : public s_tir::StmtExprVisitor {
private:
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
buffer_access_indices[op->buffer].push_back(op->indices);
return s_tir::StmtExprVisitor::Visit_(op);
}
diff --git a/src/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc
index 46915b52da..124c3b5701 100644
--- a/src/relax/transform/split_call_tir_by_pattern.cc
+++ b/src/relax/transform/split_call_tir_by_pattern.cc
@@ -294,8 +294,8 @@ class ForMatcher : public TensorizeComparator {
if (!CompareArray(op->reads, rhs->reads, &ForMatcher::CompareBufferRegion)) {
return false;
}
- // The body of the block has to be BufferStore
- if (!op->body->IsInstance<BufferStoreNode>() || !rhs->body->IsInstance<BufferStoreNode>()) {
+ // The body of the block has to be TensorStore
+ if (!op->body->IsInstance<TensorStoreNode>() || !rhs->body->IsInstance<TensorStoreNode>()) {
return false;
}
// Handle init block
@@ -321,8 +321,8 @@ class ForMatcher : public TensorizeComparator {
return Dispatch(op->block, rhs->block);
}
- bool Dispatch_(const BufferStoreNode* op, const Stmt& other) {
- const auto* rhs = other.as<BufferStoreNode>();
+ bool Dispatch_(const TensorStoreNode* op, const Stmt& other) {
+ const auto* rhs = other.as<TensorStoreNode>();
return CompareBufferAccess(op, rhs) && Dispatch(op->value, rhs->value);
}
diff --git a/src/s_tir/analysis/domain_touched.cc b/src/s_tir/analysis/domain_touched.cc
index 5c3ff3e3d5..6e9559192a 100644
--- a/src/s_tir/analysis/domain_touched.cc
+++ b/src/s_tir/analysis/domain_touched.cc
@@ -110,7 +110,7 @@ class BufferTouchedDomain final : public s_tir::IRVisitorWithAnalyzer {
return Parent::Visit_(op);
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
// Record store-exclusive buffer access
Touch(&std::get<StoreAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
// Record load-store inclusive buffer access
diff --git a/src/s_tir/analysis/estimate_flops.cc b/src/s_tir/analysis/estimate_flops.cc
index 5ac83f32ec..e707f39b7b 100644
--- a/src/s_tir/analysis/estimate_flops.cc
+++ b/src/s_tir/analysis/estimate_flops.cc
@@ -171,7 +171,7 @@ class FlopEstimator : private tirx::ExprFunctor<TResult(const Expr& n)>,
for (const Expr& arg : op->args) result += Dispatch(arg);
return result;
}
- TResult Dispatch_(const BufferStoreNode* store) override { return Dispatch(store->value); }
+ TResult Dispatch_(const TensorStoreNode* store) override { return Dispatch(store->value); }
TResult Dispatch_(const SBlockRealizeNode* block) override {
return Dispatch(block->block->body);
}
diff --git a/src/s_tir/analysis/identify_memcpy.cc b/src/s_tir/analysis/identify_memcpy.cc
index b4b267860e..aead644f74 100644
--- a/src/s_tir/analysis/identify_memcpy.cc
+++ b/src/s_tir/analysis/identify_memcpy.cc
@@ -62,13 +62,13 @@ std::variant<MemCpyDetails, std::string> IdentifyMemCpyImpl(const For& loop,
stmt = for_node->body;
}
- BufferStore store{ffi::UnsafeInit{}};
- if (auto opt = stmt.as<BufferStore>()) {
+ TensorStore store{ffi::UnsafeInit{}};
+ if (auto opt = stmt.as<TensorStore>()) {
store = opt.value();
} else {
return static_cast<const std::stringstream&>(
std::stringstream()
- << "Expected innermost loop to have BufferStore body, but instead found " << stmt)
+ << "Expected innermost loop to have TensorStore body, but instead found " << stmt)
.str();
}
@@ -78,12 +78,12 @@ std::variant<MemCpyDetails, std::string> IdentifyMemCpyImpl(const For& loop,
} else {
return static_cast<const std::stringstream&>(
std::stringstream()
- << "Expected BufferStore's value to be TensorLoad, but instead found "
+ << "Expected TensorStore's value to be TensorLoad, but instead found "
<< store->value)
.str();
}
- // Now, we have a BufferStore whose value is a TensorLoad. Because
+ // Now, we have a TensorStore whose value is a TensorLoad. Because
// non-flat physical indices are target-dependent, only handle cases
// where the buffer will be flattened to a 1-d physical buffer.
ffi::Array<PrimExpr> flattened_dst = store->buffer.OffsetOf(store->indices);
diff --git a/src/s_tir/analysis/is_pure_function.cc b/src/s_tir/analysis/is_pure_function.cc
index 2cbc34e9a6..778cc03c22 100644
--- a/src/s_tir/analysis/is_pure_function.cc
+++ b/src/s_tir/analysis/is_pure_function.cc
@@ -62,7 +62,7 @@ class PurityChecker : TIRVisitorWithPath {
TIRVisitorWithPath::Dispatch_(op, path);
}
- void Dispatch_(const BufferStoreNode* op, ffi::reflection::AccessPath path) override {
+ void Dispatch_(const TensorStoreNode* op, ffi::reflection::AccessPath path) override {
TIRVisitorWithPath::Dispatch_(op, path);
if (!internal_allocations_.count(op->buffer.var())) {
diff --git a/src/s_tir/analysis/oob_checker.cc b/src/s_tir/analysis/oob_checker.cc
index 60bf827a93..fc46ca83ae 100644
--- a/src/s_tir/analysis/oob_checker.cc
+++ b/src/s_tir/analysis/oob_checker.cc
@@ -73,7 +73,7 @@ class OOBCheckerVisitor final : public s_tir::IRVisitorWithAnalyzer {
public:
using s_tir::IRVisitorWithAnalyzer::Visit_;
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* node) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* node) final {
for (size_t i = 0; i < node->buffer->shape.size(); i++) {
CheckBounds(node, node->buffer, i);
}
diff --git a/src/s_tir/analysis/sblock_access_region_detector.cc b/src/s_tir/analysis/sblock_access_region_detector.cc
index 9f74a03674..7feabf2f38 100644
--- a/src/s_tir/analysis/sblock_access_region_detector.cc
+++ b/src/s_tir/analysis/sblock_access_region_detector.cc
@@ -140,7 +140,7 @@ class BlockReadWriteDetector : public s_tir::StmtExprVisitor {
ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) override;
ffi::Optional<VisitInterrupt> Visit_(const IfThenElseNode* op) override;
ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockRealizeNode* op) override;
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) override;
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) override;
ffi::Optional<VisitInterrupt> Visit_(const BindNode* op) override;
ffi::Optional<VisitInterrupt> Visit_(const TensorLoadNode* op) override;
ffi::Optional<VisitInterrupt> Visit_(const VarNode* op) override;
@@ -338,7 +338,7 @@ ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const CallNode* op)
return s_tir::StmtExprVisitor::Visit_(op);
}
-ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const BufferStoreNode* op) {
+ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const TensorStoreNode* op) {
auto f_substitute = [this](const Var& var) -> ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
if (auto it = let_bindings_.find(var.get()); it != let_bindings_.end()) {
return ffi::Any(it->second);
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 8e46aa2c7d..6732c88d13 100644
--- a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc
+++ b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc
@@ -293,7 +293,7 @@ class LCADetector : public s_tir::StmtExprVisitor {
return std::nullopt;
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
UpdateBufferLCA(op->buffer.get(), ancestor_scopes_.back());
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->value));
for (const auto& index : op->indices) {
diff --git a/src/s_tir/analysis/verify_gpu_code.cc b/src/s_tir/analysis/verify_gpu_code.cc
index 28628de8aa..40e5c19b2a 100644
--- a/src/s_tir/analysis/verify_gpu_code.cc
+++ b/src/s_tir/analysis/verify_gpu_code.cc
@@ -273,7 +273,7 @@ class GPUCodeVerifier : public StmtExprVisitor {
return StmtExprVisitor::Visit_(op);
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) {
PrimType value_ty = op->value.ty();
if (value_ty.IsFixedLengthVector()) {
if (ElementBytes(value_ty) > max_vector_bytes_) {
diff --git a/src/s_tir/backend/adreno/texture_flatten.cc b/src/s_tir/backend/adreno/texture_flatten.cc
index f25cfb56ec..e911cfc886 100644
--- a/src/s_tir/backend/adreno/texture_flatten.cc
+++ b/src/s_tir/backend/adreno/texture_flatten.cc
@@ -93,9 +93,9 @@ class TextureFlattener : public TextureLoweringBase {
explicit TextureFlattener(const ffi::Array<Var>& params, IRVisitorWithAnalyzer* bound_analyzer)
: TextureLoweringBase(params, bound_analyzer) {}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
Stmt stmt = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- op = stmt.as<BufferStoreNode>();
+ op = stmt.as<TensorStoreNode>();
std::string storage_scope = GetStorageScope(op->buffer);
// Lower to two dimensional access
if (IsTextureStorage(storage_scope)) {
diff --git a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
index 3920f592a6..5a684c7b7a 100644
--- a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
+++ b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
@@ -530,7 +530,7 @@ struct Feature {
static constexpr int64_t kCount = 16;
ArithOps() = default;
- ArithOps(const BufferStoreNode* store, int64_t prod_loop_extent);
+ ArithOps(const TensorStoreNode* store, int64_t prod_loop_extent);
void Export(std::vector<double>* v) const {
double vs[] = {
@@ -598,7 +598,7 @@ struct Feature {
static constexpr int64_t kCount = ArithOps::kCount + ForKindFeature::kCount * 3 + 8;
- explicit Feature(const BufferStoreNode* store, const LoopNest& loop_nest, bool is_gpu)
+ explicit Feature(const TensorStoreNode* store, const LoopNest& loop_nest, bool is_gpu)
: arith_ops(store, loop_nest.prod),
vectorize(loop_nest.vectorize),
unroll(loop_nest.unroll),
@@ -630,7 +630,7 @@ struct Feature {
}
};
-Feature::ArithOps::ArithOps(const BufferStoreNode* store, int64_t prod_loop_extent) {
+Feature::ArithOps::ArithOps(const TensorStoreNode* store, int64_t prod_loop_extent) {
class ArithOpCounter : public StmtExprVisitor {
public:
using StmtExprVisitor::Visit_;
@@ -868,11 +868,11 @@ struct Feature {
}
}
- explicit Feature(const BufferStoreNode* store, const LoopNest& loop_nest,
+ explicit Feature(const TensorStoreNode* store, const LoopNest& loop_nest,
int64_t cache_line_bytes, IntVec* for_touched_bytes,
ForBufferMap<IntVec>* buffer_touched_under_loop, sym::AnalyzerObj* analyzer);
- void Init(const BufferStoreNode* store, int n_loops);
+ void Init(const TensorStoreNode* store, int n_loops);
void SetRegion(const LoopNest& loop_nest, //
IntVec* for_touched_bytes, //
@@ -882,7 +882,7 @@ struct Feature {
std::vector<SubFeature> sub_features;
};
-void Feature::Init(const BufferStoreNode* store, int n_loops) {
+void Feature::Init(const TensorStoreNode* store, int n_loops) {
struct Info {
AccessType access_type = AccessType::kUnknownRW;
std::vector<MultiIndex> multi_indices;
@@ -1120,7 +1120,7 @@ void Feature::SubFeature::SetFeature(const LoopNest& loop_nest, int64_t cache_li
this->unique_lines_d_reuse_ct = this->unique_lines / proxy_reuse_ct;
}
-Feature::Feature(const BufferStoreNode* store, const LoopNest& loop_nest, int64_t cache_line_bytes,
+Feature::Feature(const TensorStoreNode* store, const LoopNest& loop_nest, int64_t cache_line_bytes,
IntVec* for_touched_bytes, ForBufferMap<IntVec>* buffer_touched_under_loop,
sym::AnalyzerObj* analyzer) {
int n_loops = loop_nest.loops.size();
@@ -1445,7 +1445,7 @@ class PerStoreFeatureCollector : public StmtExprVisitor {
return std::nullopt;
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* store) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* store) final {
if (store->value->IsInstance<IntImmNode>() || store->value->IsInstance<FloatImmNode>()) {
return std::nullopt;
}
diff --git a/src/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.cc b/src/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.cc
index fd488714c0..499f5b4bf8 100644
--- a/src/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.cc
+++ b/src/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.cc
@@ -61,7 +61,7 @@ struct AsyncStridedMemCopyFinder : public StmtExprVisitor {
return result;
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
if (!found_ && in_async_copy_) {
if (const auto* load = op->value.as<TensorLoadNode>()) {
// Inspect each copy, including copies grouped under one commit or predicate.
diff --git a/src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc b/src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc
index b2da7cd70d..44ad95e909 100644
--- a/src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc
+++ b/src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc
@@ -93,7 +93,7 @@ bool ParseWarpExecutionAnn(const Schedule& sch, const Instruction& inst) {
size_t GetMaxUsedDtypeBytes(SBlock block) {
size_t max_bytes = 1;
- auto visit_store = [&](const tirx::BufferStore& store) -> ffi::Expected<ffi::WalkResult> {
+ auto visit_store = [&](const tirx::TensorStore& store) -> ffi::Expected<ffi::WalkResult> {
max_bytes = std::max(max_bytes, store->value.ty().StorageBytes());
return ffi::WalkResult::Advance();
};
diff --git a/src/s_tir/schedule/analysis.h b/src/s_tir/schedule/analysis.h
index ddf77ceba1..25fdfa0f9e 100644
--- a/src/s_tir/schedule/analysis.h
+++ b/src/s_tir/schedule/analysis.h
@@ -481,13 +481,13 @@ std::pair<ffi::Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& b
/******** Reduction SBlock Related ********/
/*!
- * \brief Get the init values and the BufferStore updates from the input reduction block
+ * \brief Get the init values and the TensorStore updates from the input reduction block
* \param self The schedule state, used for error reporting
- * \param block The block from which the init values and BufferStore updates are extracted from
- * \return The extracted init values and BufferStore updates
+ * \param block The block from which the init values and TensorStore updates are extracted from
+ * \return The extracted init values and TensorStore updates
* \throw ScheduleError If rfactor or cross-thread reduction cannot be applied to the block
*/
-std::pair<ffi::Array<PrimExpr>, ffi::Array<BufferStore>> GetInitValuesAndUpdatesFromReductionBlock(
+std::pair<ffi::Array<PrimExpr>, ffi::Array<TensorStore>> GetInitValuesAndUpdatesFromReductionBlock(
const ffi::Optional<ScheduleState>& self, SBlock block);
/*!
@@ -519,7 +519,7 @@ bool ReductionIterNotIndexOutputBuffer(const SBlock& block);
*/
std::tuple<te::CommReducer, ffi::Array<PrimExpr>, ffi::Array<PrimExpr>> GetReducerAndCombinerLhsRhs(
const ffi::Optional<ScheduleState>& self, const ffi::Array<PrimExpr>& identities,
- const ffi::Array<BufferStore>& combiners);
+ const ffi::Array<TensorStore>& combiners);
/******** Commutative Reducer ********/
@@ -532,7 +532,7 @@ std::vector<ffi::TypedFunction<ffi::Optional<te::CommReducer>(ffi::Array<PrimExp
GetReducerGetters();
/*!
- * \brief Given the input identities and the combiner BufferStores of a reduction, extract the
+ * \brief Given the input identities and the combiner TensorStores of a reduction, extract the
* corresponding commutative reducer, LHS values and RHS values, if possible.
* \param identities The identities of the reduction
* \param combiners The combiners of the reduction
@@ -542,7 +542,7 @@ GetReducerGetters();
* \return A boolean indicating whether a corresponding commutative reducer is found
*/
bool FromIdentityCombiner(const ffi::Array<PrimExpr>& identities,
- const ffi::Array<BufferStore>& combiners, te::CommReducer* result_reducer,
+ const ffi::Array<TensorStore>& combiners, te::CommReducer* result_reducer,
ffi::Array<PrimExpr>* lhs, ffi::Array<PrimExpr>* rhs);
/******** Misc ********/
diff --git a/src/s_tir/schedule/analysis/reducer.cc b/src/s_tir/schedule/analysis/reducer.cc
index 638dd1d9f6..0220cd7576 100644
--- a/src/s_tir/schedule/analysis/reducer.cc
+++ b/src/s_tir/schedule/analysis/reducer.cc
@@ -318,18 +318,18 @@ class PatternMatcher : public StmtExprVisitor {
static const char* kRFactorCrossThreadReductionApplicableBlockDef =
R"(Definition of a reduction block that is applicable by RFactor and Cross-Thread Reduction:
-1) The block init should be a single BufferStore or a SeqStmt of BufferStores
+1) The block init should be a single TensorStore or a SeqStmt of TensorStores
2) The buffers initialized in the block init should be all different
-3) The number of consecutive Binds in the block body (if any) should equal the number of BufferStores in the block init
+3) The number of consecutive Binds in the block body (if any) should equal the number of TensorStores in the block init
4) The variables of the Binds in the block body should be all different
-5) The statement after the innermost Bind should be a single BufferStore or a SeqStmt of BufferStores
-6) The number of BufferStores under the block body should equal the number of BufferStores in the block init, and thereby equal the number of Binds above
-7) The variables bound by the Binds in the block body must all directly serve as values of the BufferStores inside, and the stored values of the BufferStores can only be those variables
-8) The variables stored by the BufferStores in the block body should be all different
-9) The buffers written by the BufferStores in the block body should be all different
+5) The statement after the innermost Bind should be a single TensorStore or a SeqStmt of TensorStores
+6) The number of TensorStores under the block body should equal the number of TensorStores in the block init, and thereby equal the number of Binds above
+7) The variables bound by the Binds in the block body must all directly serve as values of the TensorStores inside, and the stored values of the TensorStores can only be those variables
+8) The variables stored by the TensorStores in the block body should be all different
+9) The buffers written by the TensorStores in the block body should be all different
10) The buffers initialized in the block init and written in the block body should match
11) The buffers written by the block should have same shape
-12) The indices of all BufferStores in the reduction block should be the same)";
+12) The indices of all TensorStores in the reduction block should be the same)";
void ErrorRFactorCrossThreadReductionNotApplicable(const ffi::Optional<ScheduleState>& self,
SBlock block, int violated_cond) {
@@ -371,8 +371,8 @@ void ErrorRFactorCrossThreadReductionNotApplicable(const ffi::Optional<ScheduleS
}
/*!
- * \brief Extract the BufferStores, which serve as the reduction updates, from the given Bind nodes
- * and the BufferStores inside. And meanwhile set the buffer order of the reduction
+ * \brief Extract the TensorStores, which serve as the reduction updates, from the given Bind nodes
+ * and the TensorStores inside. And meanwhile set the buffer order of the reduction
* \param self The schedule state, used for error reporting
* \param block The reduction block, used for error reporting
* \param let The Bind nodes from which the reduction updates are extracted
@@ -383,7 +383,7 @@ void ErrorRFactorCrossThreadReductionNotApplicable(const ffi::Optional<ScheduleS
*/
void ExtractReductionUpdates(const ffi::Optional<ScheduleState>& self, SBlock block,
const ffi::Array<Stmt>& stmts, int n_buffers,
- ffi::Array<BufferStore>* updates,
+ ffi::Array<TensorStore>* updates,
std::unordered_map<const VarNode*, int>* buf2index) {
std::unordered_map<const VarNode*, int> var2index;
ffi::Array<PrimExpr> let_values;
@@ -414,7 +414,7 @@ void ExtractReductionUpdates(const ffi::Optional<ScheduleState>& self, SBlock bl
ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/3);
}
- // The remaining stmts after the Bind nodes should be BufferStores.
+ // The remaining stmts after the Bind nodes should be TensorStores.
// Collect them into a sequence.
ffi::Array<Stmt> seq;
for (int i = n_buffers; i < static_cast<int>(stmts.size()); ++i) {
@@ -425,10 +425,10 @@ void ExtractReductionUpdates(const ffi::Optional<ScheduleState>& self, SBlock bl
}
// Step 2.
- // - Create BufferStores according to the variables being stored.
+ // - Create TensorStores according to the variables being stored.
// - Construct the mapping from reduction buffers to the index.
for (const Stmt& stmt : seq) {
- const auto* buf_store = stmt.as<BufferStoreNode>();
+ const auto* buf_store = stmt.as<TensorStoreNode>();
if (buf_store == nullptr) {
ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/5);
}
@@ -445,7 +445,7 @@ void ExtractReductionUpdates(const ffi::Optional<ScheduleState>& self, SBlock bl
if ((*updates)[idx].defined()) {
ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/8);
}
- updates->Set(idx, BufferStore(buf_store->buffer, let_values[idx], buf_store->indices));
+ updates->Set(idx, TensorStore(buf_store->buffer, let_values[idx], buf_store->indices));
auto insert_result = buf2index->insert(std::make_pair(buf_store->buffer.get(), idx));
if (!insert_result.second) {
ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/9);
@@ -456,18 +456,18 @@ void ExtractReductionUpdates(const ffi::Optional<ScheduleState>& self, SBlock bl
}
}
-std::pair<ffi::Array<PrimExpr>, ffi::Array<BufferStore>> GetInitValuesAndUpdatesFromReductionBlock(
+std::pair<ffi::Array<PrimExpr>, ffi::Array<TensorStore>> GetInitValuesAndUpdatesFromReductionBlock(
const ffi::Optional<ScheduleState>& self, SBlock block) {
- ffi::Array<BufferStore> inits;
- ffi::Array<BufferStore> updates;
+ ffi::Array<TensorStore> inits;
+ ffi::Array<TensorStore> updates;
- // Step 1. Extract the BufferStores serving as block inits.
- if (auto init = block->init.as<BufferStore>()) {
+ // Step 1. Extract the TensorStores serving as block inits.
+ if (auto init = block->init.as<TensorStore>()) {
inits.push_back(init.value());
} else if (const auto* seq_init = block->init.as<SeqStmtNode>()) {
std::unordered_set<const VarNode*> init_buffers;
for (const Stmt& stmt : seq_init->seq) {
- auto init = stmt.as<BufferStore>();
+ auto init = stmt.as<TensorStore>();
if (!init) {
ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/1);
}
@@ -481,11 +481,11 @@ std::pair<ffi::Array<PrimExpr>, ffi::Array<BufferStore>> GetInitValuesAndUpdates
ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/1);
}
- // Step 2. Extract the block updates, in the form of BufferStores.
+ // Step 2. Extract the block updates, in the form of TensorStores.
int n_buffers = inits.size();
std::unordered_map<const VarNode*, int> buf2index;
- if (const auto* update = block->body.as<BufferStoreNode>()) {
- updates.push_back(ffi::GetRef<BufferStore>(update));
+ if (const auto* update = block->body.as<TensorStoreNode>()) {
+ updates.push_back(ffi::GetRef<TensorStore>(update));
buf2index[update->buffer.get()] = 0;
} else if (const auto* seq = block->body.as<SeqStmtNode>()) {
ExtractReductionUpdates(self, block, seq->seq, n_buffers, &updates, &buf2index);
@@ -500,7 +500,7 @@ std::pair<ffi::Array<PrimExpr>, ffi::Array<BufferStore>> GetInitValuesAndUpdates
init_values.resize(n_buffers);
// - Check all buffers have the same shape
- // - Check all indices of the BufferStores are the same
+ // - Check all indices of the TensorStores are the same
// - Check buffers written in the block init and the block body can match
// - Check buffers do not duplicate
const ffi::Array<PrimExpr>& expected_shape = updates[0]->buffer->shape;
@@ -605,7 +605,7 @@ bool ReductionIterNotIndexOutputBuffer(const SBlock& block) {
}
return ffi::WalkResult::Advance();
};
- auto visit_store = [&](const BufferStore& store) -> ffi::Expected<ffi::WalkResult> {
+ auto visit_store = [&](const TensorStore& store) -> ffi::Expected<ffi::WalkResult> {
bool write_is_covered_by_match_buffer =
match_buffer_sources.count(store->buffer.get()) &&
buffer_written.count(match_buffer_sources.find(store->buffer.get())->second);
@@ -630,7 +630,7 @@ bool ReductionIterNotIndexOutputBuffer(const SBlock& block) {
class NoMatchedReducerError : public ScheduleErrorContextObj {
public:
explicit NoMatchedReducerError(IRModule mod, ffi::Array<PrimExpr> identities,
- ffi::Array<BufferStore> combiners)
+ ffi::Array<TensorStore> combiners)
: mod_(std::move(mod)),
identities_(std::move(identities)),
combiners_(std::move(combiners)) {}
@@ -653,12 +653,12 @@ class NoMatchedReducerError : public ScheduleErrorContextObj {
IRModule mod_;
ffi::Array<PrimExpr> identities_;
- ffi::Array<BufferStore> combiners_;
+ ffi::Array<TensorStore> combiners_;
};
std::tuple<te::CommReducer, ffi::Array<PrimExpr>, ffi::Array<PrimExpr>> GetReducerAndCombinerLhsRhs(
const ffi::Optional<ScheduleState>& self, const ffi::Array<PrimExpr>& identities,
- const ffi::Array<BufferStore>& combiners) {
+ const ffi::Array<TensorStore>& combiners) {
te::CommReducer reducer{nullptr};
ffi::Array<PrimExpr> combiner_lhs, combiner_rhs;
bool matched =
@@ -714,7 +714,7 @@ bool MatchReducer(const te::CommReducer& reducer, const ffi::Array<PrimExpr>& id
}
bool FromIdentityCombiner(const ffi::Array<PrimExpr>& identities,
- const ffi::Array<BufferStore>& combiners, te::CommReducer* result_reducer,
+ const ffi::Array<TensorStore>& combiners, te::CommReducer* result_reducer,
ffi::Array<PrimExpr>* lhs, ffi::Array<PrimExpr>* rhs) {
int n = identities.size();
ffi::Array<TensorLoad> buf_loads;
diff --git a/src/s_tir/schedule/ir_comparator.cc b/src/s_tir/schedule/ir_comparator.cc
index b4247f101d..0c04060e45 100644
--- a/src/s_tir/schedule/ir_comparator.cc
+++ b/src/s_tir/schedule/ir_comparator.cc
@@ -247,8 +247,8 @@ bool TensorizeComparator::Dispatch_(const SeqStmtNode* op, const Stmt& other) {
&TensorizeComparator::Dispatch));
}
-bool TensorizeComparator::Dispatch_(const BufferStoreNode* op, const Stmt& other) {
- const auto* rhs = other.as<BufferStoreNode>();
+bool TensorizeComparator::Dispatch_(const TensorStoreNode* op, const Stmt& other) {
+ const auto* rhs = other.as<TensorStoreNode>();
return CompareBufferAccess(op, rhs) && Dispatch(op->value, rhs->value);
}
@@ -653,8 +653,8 @@ bool TensorizeComparator::CompareBufferRegion(const TensorRegion& lhs, const Ten
return true;
}
-// Comparator for BufferStoreNode and TensorLoadNode
-inline TensorVar GetBufferAccessBuffer(const BufferStoreNode* op) { return op->buffer; }
+// Comparator for TensorStoreNode and TensorLoadNode
+inline TensorVar GetBufferAccessBuffer(const TensorStoreNode* op) { return op->buffer; }
inline TensorVar GetBufferAccessBuffer(const TensorLoadNode* op) {
return op->source.as_or_throw<tvm::tirx::TensorVar>();
}
@@ -805,8 +805,8 @@ bool AutoTensorizeComparator::CompareBuffer(const TensorVar& lhs, const TensorVa
return equal;
}
-bool AutoTensorizeComparator::Dispatch_(const BufferStoreNode* op, const Stmt& other) {
- const auto* rhs = other.as<BufferStoreNode>();
+bool AutoTensorizeComparator::Dispatch_(const TensorStoreNode* op, const Stmt& other) {
+ const auto* rhs = other.as<TensorStoreNode>();
return CompareBufferAccess(op, rhs) && Dispatch(op->value, rhs->value);
}
diff --git a/src/s_tir/schedule/ir_comparator.h b/src/s_tir/schedule/ir_comparator.h
index 7d6bc9db87..64cfb3bd0c 100644
--- a/src/s_tir/schedule/ir_comparator.h
+++ b/src/s_tir/schedule/ir_comparator.h
@@ -53,7 +53,7 @@ class TensorizeComparator : public ExprComparator, public StmtComparator {
bool Dispatch_(const CallNode* op, const PrimExpr& other) override;
bool Dispatch_(const ForNode* op, const Stmt& other) override;
bool Dispatch_(const SeqStmtNode* op, const Stmt& other) override;
- bool Dispatch_(const BufferStoreNode* op, const Stmt& other) override;
+ bool Dispatch_(const TensorStoreNode* op, const Stmt& other) override;
bool Dispatch_(const SBlockRealizeNode* op, const Stmt& other) override;
bool Dispatch_(const SBlockNode* op, const Stmt& other) override;
@@ -149,7 +149,7 @@ class AutoTensorizeComparator : public TensorizeComparator {
bool DispatchDefault_(const ffi::Object* op, const Stmt& other) override;
bool Dispatch_(const SBlockNode* op, const Stmt& other) override;
- bool Dispatch_(const BufferStoreNode* op, const Stmt& other) override;
+ bool Dispatch_(const TensorStoreNode* op, const Stmt& other) override;
bool Dispatch_(const TensorLoadNode* op, const PrimExpr& other) override;
diff --git a/src/s_tir/schedule/primitive.h b/src/s_tir/schedule/primitive.h
index c1dd9adc19..4a984663d8 100644
--- a/src/s_tir/schedule/primitive.h
+++ b/src/s_tir/schedule/primitive.h
@@ -492,7 +492,7 @@ TVM_DLL void ReverseComputeAt(ScheduleState self, const StmtSRef& block_sref,
* \brief Inline a block into its consumer(s). It requires:
* 1) The block is a complete non-root block, which only produces one buffer
* 2) The block must not be the only leaf in the scope.
- * 3) The body of the block must be a BufferStore statement in the form of,
+ * 3) The body of the block must be a TensorStore statement in the form of,
* A[i, j, k, ...] = ...
* where the indices of the LHS are all distinct atomic variables,
* and no variables other than those indexing variables are allowed in the statement.
@@ -505,7 +505,7 @@ TVM_DLL void ComputeInline(ScheduleState self, const StmtSRef& block_sref);
* 1) The block is a complete non-root block, which only produces and consumers one buffer
* 2) The block must not be the only leaf in the scope.
* 3) The only producer of the block is a read-after-write producer and a complete non-root block
- * 4) The body of the block must be a BufferStore statement in the form of,
+ * 4) The body of the block must be a TensorStore statement in the form of,
* B[f(i, j, k, ...)] = g(i, j, k, A[i, j, k, ...] ...)
* where the indices of each `TensorLoad` on the RHS are all distinct atomic variables,
* and no variables other than those indexing variables are allowed in the statement.
diff --git a/src/s_tir/schedule/primitive/block_annotate.cc b/src/s_tir/schedule/primitive/block_annotate.cc
index 0f0013b32e..1debadc292 100644
--- a/src/s_tir/schedule/primitive/block_annotate.cc
+++ b/src/s_tir/schedule/primitive/block_annotate.cc
@@ -334,11 +334,11 @@ class DTypeMutator : public ReplaceBufferMutator {
}
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
TensorVar original_buffer = op->buffer;
- BufferStore node = StmtExprMutator::Mutate_(op, inplace_mode)
+ TensorStore node = StmtExprMutator::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<BufferStore>();
+ .as_or_throw<TensorStore>();
if (auto replacement = VarRemapGet(original_buffer).as<TensorVar>()) {
node.CopyOnWrite()->buffer = replacement.value();
node.CopyOnWrite()->value = Cast(tgt_dtype_, node->value);
diff --git a/src/s_tir/schedule/primitive/cache_index.cc b/src/s_tir/schedule/primitive/cache_index.cc
index 5cb0203f61..cfac3a4ac3 100644
--- a/src/s_tir/schedule/primitive/cache_index.cc
+++ b/src/s_tir/schedule/primitive/cache_index.cc
@@ -151,7 +151,7 @@ class IndexInfoCollector : public StmtExprVisitor {
return std::nullopt;
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* store) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* store) final {
// Only analyze the cache candidate for stores in target block
if (visiting_target_sblock) {
auto IsEligibleComputation = [](const PrimExpr& expr) {
@@ -344,7 +344,7 @@ ffi::Array<SBlock> MakeIndexCacheStage(IndexInfo* info, const ffi::String& stora
/*writes=*/{BufferRegion(info->cache_buffer[expr_index], access_region)},
/*name_hint=*/"index_" + std::to_string(expr_index),
/*body=*/
- BufferStore(info->cache_buffer[expr_index], new_expr, access_indices),
+ TensorStore(info->cache_buffer[expr_index], new_expr, access_indices),
/*init=*/std::nullopt,
/*alloc_buffers=*/{},
/*match_buffers=*/{},
@@ -442,7 +442,7 @@ class CacheIndexRewriter : public StmtExprMutator {
return stmt;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* store, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* store, InplaceMode inplace_mode) final {
Stmt ret_stmt =
StmtExprMutator::Mutate_(store, inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(store));
// Replace common sub expr for target block, with cached buffer load
diff --git a/src/s_tir/schedule/primitive/cache_read_write.cc b/src/s_tir/schedule/primitive/cache_read_write.cc
index 899df685f6..f40ee32ad0 100644
--- a/src/s_tir/schedule/primitive/cache_read_write.cc
+++ b/src/s_tir/schedule/primitive/cache_read_write.cc
@@ -221,7 +221,7 @@ SBlock MakeReindexCacheStage(const TensorRegion& cache_region, ReindexCacheStage
/*name_hint*/ cache_region->source.as_or_throw<tvm::tirx::TensorVar>().name() + "_" +
storage_scope,
/*body=*/
- BufferStore(info->write_buffer, MakeTensorLoad(info->read_buffer, read_access_indices),
+ TensorStore(info->write_buffer, MakeTensorLoad(info->read_buffer, read_access_indices),
write_access_indices),
/*init=*/std::nullopt,
/*alloc_buffers=*/{},
@@ -324,7 +324,7 @@ SBlock MakeCacheStage(const TensorRegion& cache_region, CacheStageInfo* info,
/*name_hint=*/cache_region->source.as_or_throw<tvm::tirx::TensorVar>().name() + "_" +
storage_scope,
/*body=*/
- BufferStore(info->write_buffer, MakeTensorLoad(info->read_buffer, read_access_indices),
+ TensorStore(info->write_buffer, MakeTensorLoad(info->read_buffer, read_access_indices),
write_access_indices),
/*init=*/std::nullopt,
/*alloc_buffers=*/{},
@@ -431,7 +431,7 @@ SBlock MakeReIndexStage(const SBlock& block, CacheStageInfo* info,
/*writes=*/{BufferRegionFromPoint(info->write_buffer, dst_indices)},
/*name_hint=*/info->write_buffer.name() + "_reindex",
/*body=*/
- BufferStore(info->write_buffer, MakeTensorLoad(info->read_buffer, src_indices), dst_indices));
+ TensorStore(info->write_buffer, MakeTensorLoad(info->read_buffer, src_indices), dst_indices));
// Step 4: Create surrounding loops
@@ -1434,18 +1434,18 @@ class CacheWriteRewriter : public StmtExprMutator {
return ret;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* store, InplaceMode inplace_mode) override {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* store, InplaceMode inplace_mode) override {
bool rewrite_buffer = store->buffer.same_as(info_->write_buffer);
auto value = Mutate(store->value);
auto indices = Mutate(store->indices).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
- BufferStore stmt = ffi::GetRef<BufferStore>(store);
+ TensorStore stmt = ffi::GetRef<TensorStore>(store);
if (!value.UnchangedOrSameAs(store->value) || !indices.UnchangedOrSameAs(store->indices)) {
auto* n = stmt.CopyOnWrite();
n->value = std::move(value).ValueOrUnchanged(store->value);
n->indices = std::move(indices).ValueOrUnchanged(store->indices);
}
if (rewrite_buffer) {
- BufferStoreNode* n = stmt.CopyOnWrite();
+ TensorStoreNode* n = stmt.CopyOnWrite();
n->buffer = info_->read_buffer;
if (!cache_full_region_) {
n->indices = RewriteIndices(n->indices);
@@ -1589,18 +1589,18 @@ class ReindexCacheWriteRewriter : public CacheWriteRewriter {
}
private:
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* store, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* store, InplaceMode inplace_mode) final {
bool rewrite_buffer = store->buffer.same_as(info_->write_buffer);
auto value = Mutate(store->value);
auto indices = Mutate(store->indices).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
- BufferStore stmt = ffi::GetRef<BufferStore>(store);
+ TensorStore stmt = ffi::GetRef<TensorStore>(store);
if (!value.UnchangedOrSameAs(store->value) || !indices.UnchangedOrSameAs(store->indices)) {
auto* n = stmt.CopyOnWrite();
n->value = std::move(value).ValueOrUnchanged(store->value);
n->indices = std::move(indices).ValueOrUnchanged(store->indices);
}
if (rewrite_buffer) {
- BufferStoreNode* n = stmt.CopyOnWrite();
+ TensorStoreNode* n = stmt.CopyOnWrite();
n->buffer = info_->read_buffer;
n->indices = new_indices_;
return stmt;
@@ -1684,14 +1684,14 @@ class InvalidBufferAccessError : public ScheduleErrorContextObj {
InvalidBufferAccessError(IRModule mod, TensorVar buffer, SBlock block, ErrorKind kind)
: mod_(std::move(mod)), buffer_(std::move(buffer)), block_(std::move(block)), kind_(kind) {}
ffi::String FastErrorString() const final {
- return "ScheduleError: The target buffer should be accessed via TensorLoad or BufferStore. The "
+ return "ScheduleError: The target buffer should be accessed via TensorLoad or TensorStore. The "
"indices should be the same if there are multiple accesses to the target buffer.";
}
ffi::String DetailRenderTemplate() const final {
std::ostringstream os;
os << "The target buffer " << buffer_.name()
- << " should be accessed in the leaf block {0} via TensorLoad or BufferStore. The indices "
+ << " should be accessed in the leaf block {0} via TensorLoad or TensorStore. The indices "
"should be the same if there are multiple accesses to the target buffer. ";
if (kind_ == ErrorKind::kNoAccess) {
os << "No buffer accesses found.";
@@ -1747,7 +1747,7 @@ class ReIndexCollector : public StmtExprVisitor {
throw MakeScheduleError<NotLeafBlockError>(mod_, block_);
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* store) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* store) final {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(store->value));
for (const PrimExpr& index : store->indices) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(index));
@@ -1877,16 +1877,16 @@ class ReIndexRewriter : public StmtExprMutator {
}
return node;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
auto value = Mutate(op->value);
auto indices = Mutate(op->indices).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
- BufferStore buffer_store = ffi::GetRef<BufferStore>(op);
+ TensorStore tensor_store = ffi::GetRef<TensorStore>(op);
if (!value.UnchangedOrSameAs(op->value) || !indices.UnchangedOrSameAs(op->indices)) {
- auto* n = buffer_store.CopyOnWrite();
+ auto* n = tensor_store.CopyOnWrite();
n->value = std::move(value).ValueOrUnchanged(op->value);
n->indices = std::move(indices).ValueOrUnchanged(op->indices);
}
- return VisitBufferAccess(std::move(buffer_store));
+ return VisitBufferAccess(std::move(tensor_store));
}
UnchangedOr<PrimExpr> Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final {
diff --git a/src/s_tir/schedule/primitive/compute_inline.cc b/src/s_tir/schedule/primitive/compute_inline.cc
index 695239097d..98af21ab54 100644
--- a/src/s_tir/schedule/primitive/compute_inline.cc
+++ b/src/s_tir/schedule/primitive/compute_inline.cc
@@ -317,7 +317,7 @@ class BaseInliner : public StmtExprMutator {
explicit BaseInliner(const TensorVar& inlined_buffer, const SBlock& inlined_block,
const StmtSRef& scope_root_sref)
: inlined_buffer_(inlined_buffer),
- inlined_store_(inlined_block->body.as<BufferStoreNode>()),
+ inlined_store_(inlined_block->body.as<TensorStoreNode>()),
scope_root_sref_(scope_root_sref) {
AddBuffersInBlockSignature(inlined_block.get());
}
@@ -343,10 +343,10 @@ class BaseInliner : public StmtExprMutator {
return node;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) override {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) override {
auto value = Mutate(op->value);
auto indices = Mutate(op->indices).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
- BufferStore node = ffi::GetRef<BufferStore>(op);
+ TensorStore node = ffi::GetRef<TensorStore>(op);
if (!value.UnchangedOrSameAs(op->value) || !indices.UnchangedOrSameAs(op->indices)) {
auto* n = node.CopyOnWrite();
n->value = std::move(value).ValueOrUnchanged(op->value);
@@ -482,7 +482,7 @@ class BaseInliner : public StmtExprMutator {
/*! \brief The buffer to be inlined */
TensorVar inlined_buffer_;
/*! \brief The body of the block to be inlined */
- const BufferStoreNode* inlined_store_{nullptr};
+ const TensorStoreNode* inlined_store_{nullptr};
/*! \brief The scope root */
StmtSRef scope_root_sref_{nullptr};
/*! \brief Maps a buffer's data field to itself */
@@ -566,7 +566,7 @@ class ComputeInliner : public BaseInliner {
/*analyzer=*/analyzer_,
/*simplify_trivial_iterators=*/false);
if (!res->errors.empty()) {
- // Failure: indices of BufferStore are not bijective affine
+ // Failure: indices of TensorStore are not bijective affine
return false;
}
idx_vars_.clear();
@@ -626,7 +626,7 @@ class ComputeInliner : public BaseInliner {
* \brief Helper to inline the consumer block into its producer
* The derived class implements the following functionalities:
* 1) Analyze the consumer block to determine the remapping of index variables
- * 2) Substitute `BufferStore` of the buffer to be inlined,
+ * 2) Substitute `TensorStore` of the buffer to be inlined,
* replacing it with direct writing to the buffer that consumer writes
*/
class ReverseComputeInliner : public BaseInliner {
@@ -710,7 +710,7 @@ class ReverseComputeInliner : public BaseInliner {
return false;
}
if (inlined_store_ == nullptr) {
- // Failure: block body is not BufferStore
+ // Failure: block body is not TensorStore
return false;
}
std::vector<const TensorLoadNode*> loads = ExtractTensorLoad(inlined_buffer_, inlined_store_);
@@ -748,24 +748,24 @@ class ReverseComputeInliner : public BaseInliner {
return false;
}
- const BufferStoreNode* producer_store = nullptr;
+ const TensorStoreNode* producer_store = nullptr;
if (const auto* producer_if = producer_block_->body.as<tirx::IfThenElseNode>()) {
if (producer_if->else_case.has_value()) {
return false;
}
- producer_store = producer_if->then_case.as<BufferStoreNode>();
+ producer_store = producer_if->then_case.as<TensorStoreNode>();
} else {
- producer_store = producer_block_->body.as<BufferStoreNode>();
+ producer_store = producer_block_->body.as<TensorStoreNode>();
if (producer_block_->annotations.count(s_tir::attr::auto_copy) != 0) {
const ForNode* producer_inner_loop = producer_block_->body.as<ForNode>();
while (producer_inner_loop->body.as<ForNode>()) {
producer_inner_loop = producer_inner_loop->body.as<ForNode>();
}
- producer_store = producer_inner_loop->body.as<BufferStoreNode>();
+ producer_store = producer_inner_loop->body.as<TensorStoreNode>();
}
}
if (producer_store == nullptr) {
- // Failure: producer block body is not BufferStore
+ // Failure: producer block body is not TensorStore
return false;
}
CreateInverseMapping(producer_store->indices);
@@ -844,10 +844,10 @@ class ReverseComputeInliner : public BaseInliner {
return tgt_block_realize;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* _store, InplaceMode inplace_mode) final {
- BufferStore store = BaseInliner::Mutate_(_store, inplace_mode)
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* _store, InplaceMode inplace_mode) final {
+ TensorStore store = BaseInliner::Mutate_(_store, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(_store))
- .as_or_throw<BufferStore>();
+ .as_or_throw<TensorStore>();
if (!store->buffer.same_as(inlined_buffer_)) {
return store;
}
@@ -881,9 +881,9 @@ class ReverseComputeInliner : public BaseInliner {
/*!
* \brief Apply the inverse of `buffer_load_iter_map_` to producer indices. Seed the inherited
- * remapping environment with the result. It will be later used to transform the BufferStore
+ * remapping environment with the result. It will be later used to transform the TensorStore
* indices of the producer.
- * \param producer_indices The BufferStore indices of the producer.
+ * \param producer_indices The TensorStore indices of the producer.
*/
void CreateInverseMapping(const ffi::Array<PrimExpr> producer_indices) {
auto inverse_iter_map = sym::InverseAffineIterMap(buffer_load_iter_map_, producer_indices);
@@ -892,25 +892,25 @@ class ReverseComputeInliner : public BaseInliner {
}
}
- Stmt ReplaceInlinedBuffer(BufferStore producer) {
+ Stmt ReplaceInlinedBuffer(TensorStore producer) {
// "producer->value" may contain the buffer that is inlined in cases of reduction,
// so we need to resolve the recursion first
producer_rhs_ = ffi::make_object<RecursionResolver>(this)
->Mutate(producer->value)
.ValueOrUnchanged(producer->value);
return ffi::make_object<Substituter>(this)
- ->Mutate(ffi::GetRef<BufferStore>(inlined_store_))
- .ValueOrUnchanged(ffi::GetRef<BufferStore>(inlined_store_));
+ ->Mutate(ffi::GetRef<TensorStore>(inlined_store_))
+ .ValueOrUnchanged(ffi::GetRef<TensorStore>(inlined_store_));
}
/*!
* \brief Extracts expressions that loads a specific buffer
* \param buffer The buffer to be loaded from
- * \param from The BufferStore statement to be extracted from
+ * \param from The TensorStore statement to be extracted from
* \return A list of `TensorLoad` expressions
*/
static std::vector<const TensorLoadNode*> ExtractTensorLoad(const TensorVar& buffer,
- const BufferStoreNode* from) {
+ const TensorStoreNode* from) {
struct Extractor : public StmtExprVisitor {
using StmtExprVisitor::Visit_;
@@ -950,7 +950,7 @@ class ReverseComputeInliner : public BaseInliner {
return true;
}
- /*! \brief The RHS value of the producer's BufferStore statement */
+ /*! \brief The RHS value of the producer's TensorStore statement */
PrimExpr producer_rhs_{ffi::UnsafeInit{}};
/*! \brief The indices of the consumer's TensorLoad */
ffi::Array<PrimExpr> buffer_load_indices_;
@@ -1113,9 +1113,9 @@ class ReductionEpilogueFuser : public BaseInliner {
private:
bool IsReductionBlock(const SBlockNode* block);
void ExtractEpilogueInfo();
- // Helper function to extract TensorLoad nodes from BufferStore
+ // Helper function to extract TensorLoad nodes from TensorStore
static std::vector<const TensorLoadNode*> ExtractTensorLoad(const TensorVar& buffer,
- const BufferStoreNode* from) {
+ const TensorStoreNode* from) {
struct Extractor : public StmtExprVisitor {
using StmtExprVisitor::Visit_;
@@ -1161,9 +1161,9 @@ bool ReductionEpilogueFuser::BodyPatternAllowFusion(const SBlockRealize& epilogu
return false;
}
- // 2. Check if epilogue body is BufferStore
+ // 2. Check if epilogue body is TensorStore
if (inlined_store_ == nullptr) {
- // Failure: epilogue block body is not BufferStore
+ // Failure: epilogue block body is not TensorStore
return false;
}
@@ -1419,7 +1419,7 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock(
return ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(index, f_substitute)
.as_or_throw<PrimExpr>();
});
- BufferStore new_init_store = BufferStore(epilogue_output_buffer_, init_epilogue, init_indices);
+ TensorStore new_init_store = TensorStore(epilogue_output_buffer_, init_epilogue, init_indices);
new_block->init = new_init_store;
// 3. Generalized update transformation: apply epilogue expression with reduction buffer replaced
@@ -1439,10 +1439,10 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock(
epilogue_expression_(epilogue_expr),
var_map_(var_map) {}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
- BufferStore store = StmtExprMutator::Mutate_(op, inplace_mode)
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
+ TensorStore store = StmtExprMutator::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<BufferStore>();
+ .as_or_throw<TensorStore>();
if (store->buffer.same_as(old_buffer_)) {
// Replace old_buffer_ in store->value with new_buffer_ to get the reduction update
// expression This ensures store->value references new_buffer_ instead of old_buffer_
@@ -1571,7 +1571,7 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock(
new_value = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(new_value, f_substitute)
.as_or_throw<PrimExpr>();
- return BufferStore(new_buffer_, new_value, store->indices);
+ return TensorStore(new_buffer_, new_value, store->indices);
}
return store;
}
diff --git a/src/s_tir/schedule/primitive/decompose_padding.cc b/src/s_tir/schedule/primitive/decompose_padding.cc
index cfeea4e9fb..a5fd78a11a 100644
--- a/src/s_tir/schedule/primitive/decompose_padding.cc
+++ b/src/s_tir/schedule/primitive/decompose_padding.cc
@@ -96,14 +96,14 @@ class PaddingInfoAnalyzer {
Var block_var = block->iter_vars[i]->var;
iter_values.insert_or_assign(block_var.get(), realize->iter_values[i]);
}
- const BufferStoreNode* store = block->body.as<BufferStoreNode>();
+ const TensorStoreNode* store = block->body.as<TensorStoreNode>();
if (!store) {
- SetError("Block body expect a BufferStore to the write buffer");
+ SetError("Block body expect a TensorStore to the write buffer");
return false;
}
const CallNode* if_then_else = store->value.as<CallNode>();
if (!if_then_else || !if_then_else->op.same_as(prim::builtin::if_then_else())) {
- SetError("Value of BufferStore expect to be constrained by a padding predicate");
+ SetError("Value of TensorStore expect to be constrained by a padding predicate");
return false;
}
auto f_substitute =
@@ -245,7 +245,7 @@ static std::pair<Stmt, SBlockRealize> CreateConstBlock(const SBlockRealizeNode*
}));
// create block to fill const pad values
- BufferStore store = block->body.as_or_throw<BufferStore>();
+ TensorStore store = block->body.as_or_throw<TensorStore>();
store.CopyOnWrite()->value = info.pad_value;
store.CopyOnWrite()->indices = store->indices.Map(rewrite_expr);
SBlock new_block(/*iter_vars=*/new_iter_vars, /*reads=*/{}, /*writes=*/{write_region},
@@ -356,7 +356,7 @@ static std::pair<Stmt, SBlockRealize> CreateInBoundBlock(const SBlockRealizeNode
}
// create new block realize node
- BufferStore store = block->body.as_or_throw<BufferStore>();
+ TensorStore store = block->body.as_or_throw<TensorStore>();
store.CopyOnWrite()->value = rewrite_expr(info.in_bound_value);
store.CopyOnWrite()->indices = store->indices.Map(rewrite_expr);
SBlock new_block(/*iter_vars=*/new_iter_vars, /*reads=*/reads, /*writes=*/writes,
diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc b/src/s_tir/schedule/primitive/layout_transformation.cc
index c17887ba15..39326950cf 100644
--- a/src/s_tir/schedule/primitive/layout_transformation.cc
+++ b/src/s_tir/schedule/primitive/layout_transformation.cc
@@ -41,13 +41,13 @@ using namespace tvm::tirx;
*
* There are four ways that transformation may be handled. Each
* updates the buffer shape and the indices used to acces the buffer
- * in BufferStore/TensorLoad nodes, but differ in how they handle the
+ * in TensorStore/TensorLoad nodes, but differ in how they handle the
* `pad_value`. In order of preference, the different strategies are
* as follows:
*
* 1. NoPaddingRequired. The transformation does not introduce
* padding, so only local changes to update the indices of
- * TensorLoad/BufferStore nodes are required. No blocks are added,
+ * TensorLoad/TensorStore nodes are required. No blocks are added,
* removed, or replaced.
*
* 2. ProloguePlan. The transformation introduces padding, but the
@@ -115,8 +115,8 @@ class TransformLayoutPlanner : public StmtExprVisitor {
private:
struct WriteInfo {
- // The BufferStore object
- BufferStore store;
+ // The TensorStore object
+ TensorStore store;
// The block realize that contains the store, if any.
ffi::Optional<SBlockRealize> innermost_block_realize;
@@ -152,7 +152,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
return StmtExprVisitor::Visit_(op);
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) override {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) override {
if (!op->buffer.same_as(old_buffer_)) {
return std::nullopt;
}
@@ -170,7 +170,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
}
}
- WriteInfo write_info{ffi::GetRef<BufferStore>(op)};
+ WriteInfo write_info{ffi::GetRef<TensorStore>(op)};
if (loop_dependency_range) {
size_t i = loop_dependency_range.value().first;
size_t j = loop_dependency_range.value().second;
@@ -216,7 +216,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
write_info_.push_back(write_info);
// Don't need to continue recursing, as the entire goal was to
- // find the BufferStore.
+ // find the TensorStore.
return std::nullopt;
}
@@ -237,12 +237,12 @@ class TransformLayoutPlanner : public StmtExprVisitor {
return prev;
}
- class BufferStoreReplacer : public StmtExprMutator {
+ class TensorStoreReplacer : public StmtExprMutator {
public:
using StmtExprMutator::Mutate;
using StmtExprMutator::Mutate_;
- BufferStoreReplacer(const WriteInfo& info, const TensorVar& new_buffer,
+ TensorStoreReplacer(const WriteInfo& info, const TensorVar& new_buffer,
PrimExpr padding_predicate, const IndexMap& inverse,
const ffi::Optional<IndexMap>& pad_value,
ffi::Map<SBlock, SBlock>* new_block_to_old, sym::AnalyzerObj* analyzer)
@@ -367,7 +367,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
this->new_iter_values = new_iter_values;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
bool can_replace = [&]() -> bool {
if (!op->buffer.same_as(info.store->buffer)) {
return false;
@@ -385,14 +385,14 @@ class TransformLayoutPlanner : public StmtExprVisitor {
return true;
}();
- BufferStore store = ffi::GetRef<BufferStore>(op);
+ TensorStore store = ffi::GetRef<TensorStore>(op);
if (can_replace) {
ffi::Array<PrimExpr> new_index_exprs =
new_indices.Map([](const Var& var) { return var.as_or_throw<PrimExpr>(); });
PrimExpr pad_value_at_index =
pad_value.value()->MapIndices(new_index_exprs, ffi::GetRef<sym::Analyzer>(analyzer))[0];
store =
- BufferStore(new_buffer, if_then_else(padding_predicate, pad_value_at_index, op->value),
+ TensorStore(new_buffer, if_then_else(padding_predicate, pad_value_at_index, op->value),
new_index_exprs);
} else {
all_stores_replaced = false;
@@ -562,7 +562,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
return std::nullopt;
}
- auto replacer = ffi::make_object<BufferStoreReplacer>(
+ auto replacer = ffi::make_object<TensorStoreReplacer>(
info, new_buffer, padding_predicate, inverse, pad_value, &new_block_to_old, analyzer);
Stmt stmt = replacer->Mutate(info.dependent_loopnest.back()->body)
.ValueOrUnchanged(info.dependent_loopnest.back()->body);
@@ -622,7 +622,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
PrimExpr pad_value_at_index =
pad_value.value()->MapIndices(indices, ffi::GetRef<sym::Analyzer>(analyzer))[0];
- Stmt stmt = BufferStore(new_buffer, pad_value_at_index, indices);
+ Stmt stmt = TensorStore(new_buffer, pad_value_at_index, indices);
std::stringstream block_name;
block_name << "buffer_" << new_buffer.name() << "_padding";
@@ -723,7 +723,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
std::vector<BindVariableDefinition> bound_vars_;
};
- /*! \brief Collected information about each BufferStore */
+ /*! \brief Collected information about each TensorStore */
std::vector<WriteInfo> write_info_;
/*! \brief The loop iterators surrounding the current node
@@ -743,7 +743,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
/*! \brief The variable mappings that are currently in-scope
*
- * Used to determine whether the indices of a BufferStore are a
+ * Used to determine whether the indices of a TensorStore are a
* row-major traversal, even if they are rebound in let/block
* mappings.
*/
@@ -906,15 +906,15 @@ class TransformLayoutRewriter : public s_tir::IRMutatorWithAnalyzer {
return buffer_load;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
- BufferStore buffer_store = Parent::Mutate_(op, inplace_mode)
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
+ TensorStore tensor_store = Parent::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<BufferStore>();
- if (buffer_store->buffer.same_as(old_buffer_)) {
- auto* n = buffer_store.CopyOnWrite();
+ .as_or_throw<TensorStore>();
+ if (tensor_store->buffer.same_as(old_buffer_)) {
+ auto* n = tensor_store.CopyOnWrite();
RewriteBufferAccess(&n->buffer, &n->indices);
}
- return buffer_store;
+ return tensor_store;
}
void RewriteAccessRegion(ffi::Array<TensorRegion>* old_access_regions,
@@ -1297,7 +1297,7 @@ void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_
new_buffer_type->shape = index_map->MapShape(old_buffer->shape, analyzer);
TensorVar new_buffer = RebuildTensorVar(old_buffer, std::move(new_buffer_type));
- // Step 3: Rewrite TensorLoad/BufferStore access indices, block read/write regions, and block
+ // Step 3: Rewrite TensorLoad/TensorStore access indices, block read/write regions, and block
// alloc_buffers.
auto [new_stmt, block_sref_reuse] =
TransformLayoutRewriter::Rewrite(ffi::GetRef<SBlock>(scope_block), old_buffer, new_buffer,
diff --git a/src/s_tir/schedule/primitive/loop_transformation.cc b/src/s_tir/schedule/primitive/loop_transformation.cc
index 40169d1869..6fba2d1250 100644
--- a/src/s_tir/schedule/primitive/loop_transformation.cc
+++ b/src/s_tir/schedule/primitive/loop_transformation.cc
@@ -509,7 +509,7 @@ class BufferIndicesMapExtractor : public StmtExprVisitor {
}
private:
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* store) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* store) final {
ffi::Array<Var> indices;
bool check_ = false;
for (size_t i = 0; i < store->indices.size(); i++) {
diff --git a/src/s_tir/schedule/primitive/pad_einsum.cc b/src/s_tir/schedule/primitive/pad_einsum.cc
index 8e8798fbf8..6168901350 100644
--- a/src/s_tir/schedule/primitive/pad_einsum.cc
+++ b/src/s_tir/schedule/primitive/pad_einsum.cc
@@ -30,7 +30,7 @@ using namespace tvm::tirx;
/*!
* \brief Check if buffer indices are all Vars and expr
- * \param buffer_access The TensorLoad or BufferStore
+ * \param buffer_access The TensorLoad or TensorStore
* \return The indices if the indices are all Vars, otherwise std::nullopt
*/
ffi::Optional<ffi::Array<Var>> CheckTrivialBufferIndices(
@@ -184,10 +184,10 @@ struct BufferPadding {
}
}
PrimExpr rhs = MakeTensorLoad(buffer, indices);
- return BufferStore(padded_buffer,
+ return TensorStore(padded_buffer,
if_then_else(predicate, rhs, prim::MakeConst(rhs.ty(), 0)), indices);
} else {
- return BufferStore(buffer, MakeTensorLoad(padded_buffer, indices), indices);
+ return TensorStore(buffer, MakeTensorLoad(padded_buffer, indices), indices);
}
}();
TensorRegion read_region = BufferRegion(buffer, instance_dom);
@@ -282,7 +282,7 @@ class InvalidProducerError : public ScheduleErrorContextObj {
ffi::String DetailRenderTemplate() const final {
std::ostringstream os;
os << "The producer block {0} cannot be padded. It should write to a single buffer and the "
- "body should be a BufferStore.";
+ "body should be a TensorStore.";
return os.str();
}
@@ -355,12 +355,12 @@ class PadEinsumBufferReplacer : public StmtExprMutator {
return new_for;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* old_store_ptr, InplaceMode inplace_mode) final {
- BufferStore store = StmtExprMutator::Mutate_(old_store_ptr, inplace_mode)
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* old_store_ptr, InplaceMode inplace_mode) final {
+ TensorStore store = StmtExprMutator::Mutate_(old_store_ptr, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(old_store_ptr))
- .as_or_throw<BufferStore>();
+ .as_or_throw<TensorStore>();
if (ffi::Optional<TensorVar> buffer = VarRemapGet(store->buffer).as<TensorVar>()) {
- return BufferStore(buffer.value(), store->value, store->indices);
+ return TensorStore(buffer.value(), store->value, store->indices);
} else {
return store;
}
diff --git a/src/s_tir/schedule/primitive/read_write_at.cc b/src/s_tir/schedule/primitive/read_write_at.cc
index 50059c13aa..e14179881b 100644
--- a/src/s_tir/schedule/primitive/read_write_at.cc
+++ b/src/s_tir/schedule/primitive/read_write_at.cc
@@ -116,14 +116,14 @@ class ReadWriteAtBufferReplacer : public StmtExprMutator {
: src_(src), dst_(dst), block_sref_reuse_(block_sref_reuse) {}
private:
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* _store, InplaceMode inplace_mode) final {
- BufferStore store = StmtExprMutator::Mutate_(_store, inplace_mode)
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* _store, InplaceMode inplace_mode) final {
+ TensorStore store = StmtExprMutator::Mutate_(_store, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(_store))
- .as_or_throw<BufferStore>();
+ .as_or_throw<TensorStore>();
if (store->buffer.same_as(src_)) {
- ffi::ObjectPtr<BufferStoreNode> new_store = ffi::make_object<BufferStoreNode>(*store.get());
+ ffi::ObjectPtr<TensorStoreNode> new_store = ffi::make_object<TensorStoreNode>(*store.get());
new_store->buffer = dst_;
- return BufferStore(new_store);
+ return TensorStore(new_store);
}
return store;
}
@@ -334,7 +334,7 @@ struct ReadWriteAtImpl {
indices.push_back(domain[i]->min + loop_vars[i].as_or_throw<PrimExpr>());
}
Stmt stmt =
- BufferStore(copy_to, /*value=*/MakeTensorLoad(copy_from, indices), /*indices=*/indices);
+ TensorStore(copy_to, /*value=*/MakeTensorLoad(copy_from, indices), /*indices=*/indices);
for (int i = n - 1; i >= 0; --i) {
stmt = For(loop_vars[i].as_or_throw<PrimVar>(), IntImm::Int32(0), domain[i]->extent,
ForKind::kSerial, stmt);
diff --git a/src/s_tir/schedule/primitive/reduction.cc b/src/s_tir/schedule/primitive/reduction.cc
index e25e8f5673..c299ed6f54 100644
--- a/src/s_tir/schedule/primitive/reduction.cc
+++ b/src/s_tir/schedule/primitive/reduction.cc
@@ -729,17 +729,17 @@ std::unordered_map<const VarNode*, For> GetLoopVar2LoopMap(const ffi::Array<For>
/*!
* \brief Create the intermediate rfactor buffers, which the rfactor block writes to and the
* write-back block reads from
- * \param buf_stores The BufferStores of the original block, where the rfactor buffers will be
+ * \param buf_stores The TensorStores of the original block, where the rfactor buffers will be
* created from
* \param factor_axis The `factor_axis` parameter of rfactor
* \param rf_loop The rfactor loop
* \return The new created intermediate rfactor buffer
*/
-ffi::Array<TensorVar> CreateRFactorBuffers(const ffi::Array<BufferStore>& buf_stores,
+ffi::Array<TensorVar> CreateRFactorBuffers(const ffi::Array<TensorStore>& buf_stores,
int factor_axis, const ForNode* rf_loop) {
ffi::Array<TensorVar> rf_buffers;
rf_buffers.reserve(buf_stores.size());
- for (const BufferStore& buf_store : buf_stores) {
+ for (const TensorStore& buf_store : buf_stores) {
TensorVar buffer = buf_store->buffer;
ffi::Array<PrimExpr> rf_shape = buffer->shape;
rf_shape.insert(rf_shape.begin() + factor_axis, rf_loop->extent);
@@ -762,7 +762,7 @@ ffi::Array<TensorVar> CreateRFactorBuffers(const ffi::Array<BufferStore>& buf_st
class BaseBlockCreator {
public:
explicit BaseBlockCreator(SBlockRealize old_block_realize, For rf_loop,
- ffi::Array<BufferStore> old_reduction_updates, te::CommReducer reducer,
+ ffi::Array<TensorStore> old_reduction_updates, te::CommReducer reducer,
ffi::Array<TensorVar> rf_buffers, bool is_rf_block)
: old_block_realize_(std::move(old_block_realize)),
rf_loop_(std::move(rf_loop)),
@@ -841,15 +841,15 @@ class BaseBlockCreator {
// buffers.
if (!has_reduce_iter) {
for (int i = 0; i < n_buffers_; ++i) {
- buf_stores.push_back(BufferStore(update_buffers_[i], update_rhs_[i], update_indices_[i]));
+ buf_stores.push_back(TensorStore(update_buffers_[i], update_rhs_[i], update_indices_[i]));
}
return n_buffers_ > 1 ? SeqStmt(buf_stores) : buf_stores[0];
}
- // Case 2. If the reduction is for single buffer, the block body is a single BufferStore.
+ // Case 2. If the reduction is for single buffer, the block body is a single TensorStore.
ffi::Array<PrimExpr> stored_values = (*reducer_.get())(update_lhs_, update_rhs_);
if (n_buffers_ == 1) {
- return BufferStore(update_buffers_[0], stored_values[0], update_indices_[0]);
+ return TensorStore(update_buffers_[0], stored_values[0], update_indices_[0]);
}
// Case 3. In case the reduction is for multiple buffers, we should create the reduction with
@@ -860,7 +860,7 @@ class BaseBlockCreator {
Var var("v_" + update_buffers_[i].name(), stored_values[i].ty());
let_vars.push_back(var);
buf_stores.push_back(
- BufferStore(update_buffers_[i], var.as_or_throw<PrimExpr>(), update_indices_[i]));
+ TensorStore(update_buffers_[i], var.as_or_throw<PrimExpr>(), update_indices_[i]));
}
ffi::Array<Stmt> stmts;
for (int i = 0; i < n_buffers_; ++i) {
@@ -881,7 +881,7 @@ class BaseBlockCreator {
inits.reserve(n_buffers_);
for (int i = 0; i < n_buffers_; ++i) {
inits.push_back(
- BufferStore(update_buffers_[i], reducer_->identity_element[i], update_indices_[i]));
+ TensorStore(update_buffers_[i], reducer_->identity_element[i], update_indices_[i]));
}
return n_buffers_ > 1 ? SeqStmt(inits) : inits[0];
}
@@ -901,8 +901,8 @@ class BaseBlockCreator {
int n_block_iters_;
/*! \brief The rfactor loop */
For rf_loop_;
- /*! \brief The update BufferStores of the old block */
- ffi::Array<BufferStore> old_reduction_updates_;
+ /*! \brief The update TensorStores of the old block */
+ ffi::Array<TensorStore> old_reduction_updates_;
/*! \brief The matched commutative reducer */
te::CommReducer reducer_;
/*! \brief The intermediate rfactor buffers */
@@ -960,7 +960,7 @@ class BaseBlockCreator {
class RFactorBlockCreator : public BaseBlockCreator {
public:
explicit RFactorBlockCreator(SBlockRealize old_block_realize, For rf_loop,
- ffi::Array<BufferStore> old_reduction_updates,
+ ffi::Array<TensorStore> old_reduction_updates,
te::CommReducer reducer, ffi::Array<TensorVar> rf_buffers,
std::unordered_map<const VarNode*, For> loop_vars2loop,
int factor_axis, ffi::Array<PrimExpr> combiner_rhs)
@@ -1117,7 +1117,7 @@ class RFactorBlockCreator : public BaseBlockCreator {
class WriteBackBlockCreator : public BaseBlockCreator {
public:
explicit WriteBackBlockCreator(SBlockRealize old_block_realize, For rf_loop,
- ffi::Array<BufferStore> old_reduction_updates,
+ ffi::Array<TensorStore> old_reduction_updates,
te::CommReducer reducer, ffi::Array<TensorVar> rf_buffers,
IterVar rf_additional_iter, ffi::Array<PrimExpr> combiner_lhs,
ffi::Array<PrimExpr> rf_buf_access_indices)
@@ -1407,7 +1407,7 @@ StmtSRef RFactor(ScheduleState self, const StmtSRef& rf_loop_sref, int factor_ax
// reduction combiner. The lhs will be used when constructing the write-back block, and the rhs
// will be used when constructing the rfactor block.
ffi::Array<PrimExpr> init_values{nullptr};
- ffi::Array<BufferStore> updates{nullptr};
+ ffi::Array<TensorStore> updates{nullptr};
te::CommReducer reducer{nullptr};
ffi::Array<PrimExpr> combiner_lhs{nullptr};
ffi::Array<PrimExpr> combiner_rhs{nullptr};
diff --git a/src/s_tir/schedule/primitive/rolling_buffer.cc b/src/s_tir/schedule/primitive/rolling_buffer.cc
index 7f2f721cf9..3ad08ef07f 100644
--- a/src/s_tir/schedule/primitive/rolling_buffer.cc
+++ b/src/s_tir/schedule/primitive/rolling_buffer.cc
@@ -390,12 +390,12 @@ class RollingBufferRewriter : public StmtExprMutator {
return stmt;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
- BufferStore stmt = StmtExprMutator::Mutate_(op, inplace_mode)
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
+ TensorStore stmt = StmtExprMutator::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<BufferStore>();
+ .as_or_throw<TensorStore>();
if (stmt->buffer.same_as(info_->old_buffer)) {
- BufferStoreNode* n = stmt.CopyOnWrite();
+ TensorStoreNode* n = stmt.CopyOnWrite();
RewriteBufferAccess(&n->buffer, &n->indices);
// Need to add predicate to the current block to avoid recomputing elements.
rewrite_block_predicate_ = true;
diff --git a/src/s_tir/schedule/transform.cc b/src/s_tir/schedule/transform.cc
index 22c935166f..458c936d16 100644
--- a/src/s_tir/schedule/transform.cc
+++ b/src/s_tir/schedule/transform.cc
@@ -479,11 +479,11 @@ UnchangedOr<Stmt> BlockBufferAccessSimplifier::Mutate_(const SBlockNode* op,
return block;
}
-UnchangedOr<Stmt> BlockBufferAccessSimplifier::Mutate_(const BufferStoreNode* op,
+UnchangedOr<Stmt> BlockBufferAccessSimplifier::Mutate_(const TensorStoreNode* op,
InplaceMode inplace_mode) {
- BufferStore node = s_tir::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
+ TensorStore node = s_tir::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<BufferStore>();
+ .as_or_throw<TensorStore>();
SimplifyBufferIndices(&node.CopyOnWrite()->indices);
return node;
}
diff --git a/src/s_tir/schedule/transform.h b/src/s_tir/schedule/transform.h
index a79b937411..2a2eeb0f93 100644
--- a/src/s_tir/schedule/transform.h
+++ b/src/s_tir/schedule/transform.h
@@ -235,7 +235,7 @@ class BlockBufferAccessSimplifier : public s_tir::IRMutatorWithAnalyzer {
void SimplifyBufferIndices(ffi::Array<PrimExpr>* indices);
UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final;
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final;
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final;
UnchangedOr<PrimExpr> Mutate_(const TensorLoadNode* 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 246b01bf13..eed566fc1e 100644
--- a/src/s_tir/transform/compact_buffer_region.cc
+++ b/src/s_tir/transform/compact_buffer_region.cc
@@ -78,7 +78,7 @@ class Var2BufferCollector : public StmtExprVisitor {
var2buffer_;
private:
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
var2buffer_[op->buffer.var()].insert(op->buffer);
return StmtExprVisitor::Visit_(op);
}
@@ -158,7 +158,7 @@ class BufferAccessRegionCollector : public StmtExprVisitor {
return std::nullopt;
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
VisitBufferAccess(BufferRegionFromPoint(op->buffer, op->indices));
return Visit(op->value);
}
@@ -617,12 +617,12 @@ class BufferCompactor : public StmtExprMutator {
explicit BufferCompactor(std::unordered_map<Var, BufferAllocInfo> buffer_info)
: buffer_info_(std::move(buffer_info)) {}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* _op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* _op, InplaceMode inplace_mode) final {
TensorVar original_buffer = _op->buffer;
- BufferStore store = StmtExprMutator::Mutate_(_op, inplace_mode)
+ TensorStore store = StmtExprMutator::Mutate_(_op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(_op))
- .as_or_throw<BufferStore>();
- BufferStoreNode* op = store.CopyOnWrite();
+ .as_or_throw<TensorStore>();
+ TensorStoreNode* op = store.CopyOnWrite();
RewriteBufferAccess(original_buffer, &op->buffer, &op->indices);
return store;
}
diff --git a/src/s_tir/transform/inject_permuted_layout.cc b/src/s_tir/transform/inject_permuted_layout.cc
index ed156e51e5..146abe7520 100644
--- a/src/s_tir/transform/inject_permuted_layout.cc
+++ b/src/s_tir/transform/inject_permuted_layout.cc
@@ -205,13 +205,13 @@ class PermutedLayoutInjector : public IRMutatorWithAnalyzer {
return indices;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
// Rewrite write from global to shared.dyn or shared
// We assume the shape of the shared memory is [..., row_size, col_size],
// where row_size is divisible by 64, or divisible by 32 and col_size is divisible by 2.
auto store = IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<BufferStore>();
+ .as_or_throw<TensorStore>();
if (!permute_ || store->buffer->shape.size() < 2) {
return store;
diff --git a/src/s_tir/transform/inject_ptx_async_copy.cc b/src/s_tir/transform/inject_ptx_async_copy.cc
index 7d1937d98b..5a062b44e9 100644
--- a/src/s_tir/transform/inject_ptx_async_copy.cc
+++ b/src/s_tir/transform/inject_ptx_async_copy.cc
@@ -61,7 +61,7 @@ class PTXAsyncCopyInjector : public StmtExprMutator {
return StmtExprMutator::Mutate_(region, inplace_mode);
}
- Stmt InjectPTX(const TensorLoadNode* load, const BufferStoreNode* store, bool predicated = false,
+ Stmt InjectPTX(const TensorLoadNode* load, const TensorStoreNode* store, bool predicated = false,
ffi::Optional<PrimExpr> predicate_value = std::nullopt) {
TVM_FFI_ICHECK(!predicated || predicate_value.has_value());
if (load->source.as_or_throw<tvm::tirx::TensorVar>().scope() == "global") {
@@ -86,10 +86,10 @@ class PTXAsyncCopyInjector : public StmtExprMutator {
// is a byte buffer generated by merging dynamic shared memory.
TVM_FFI_ICHECK(store->buffer.scope() == "shared.dyn");
TVM_FFI_ICHECK((dst_elem_type.value() == PrimType::UInt(8)));
- // BufferStore/Load have the "pointer reinterpret" semantics according to their
+ // TensorStore/Load have the "pointer reinterpret" semantics according to their
// "value" dtype. Their "indices" are supposed to be applied after such pointer cast,
// for example: ((*float16)(byte_buffer))[buffer->indices] = fp16_value;
- // To replace BufferStore/Load with cp.async, we need to multiply the store index by
+ // To replace TensorStore/Load with cp.async, we need to multiply the store index by
// the byte size of the "value" dtype, to get the correct offset into the byte buffer.
index_factor = static_cast<int>(src_elem_type.value().StorageBytes());
}
@@ -184,7 +184,7 @@ class PTXAsyncCopyInjector : public StmtExprMutator {
.ValueOrUnchanged(ffi::GetRef<Stmt>(store));
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* store, InplaceMode inplace_mode) {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* store, InplaceMode inplace_mode) {
if (in_async && (store->buffer.scope() == "shared" || store->buffer.scope() == "shared.dyn")) {
if (auto* load = store->value.as<TensorLoadNode>()) {
return InjectPTX(load, store);
diff --git a/src/s_tir/transform/inject_software_pipeline.cc b/src/s_tir/transform/inject_software_pipeline.cc
index fa88a2b74f..e43c4f2dc0 100644
--- a/src/s_tir/transform/inject_software_pipeline.cc
+++ b/src/s_tir/transform/inject_software_pipeline.cc
@@ -301,11 +301,11 @@ class PipelineBodyRewriter : public StmtExprMutator {
return StmtExprMutator::Mutate_(region.get(), InplaceMode::kDisallow).ValueOrUnchanged(region);
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
auto replacement = VarRemapGet(op->buffer).as<TensorVar>();
- BufferStore store = StmtExprMutator::Mutate_(op, inplace_mode)
+ TensorStore store = StmtExprMutator::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<BufferStore>();
+ .as_or_throw<TensorStore>();
if (!replacement) {
return store;
}
diff --git a/src/s_tir/transform/inject_virtual_thread.cc b/src/s_tir/transform/inject_virtual_thread.cc
index 1dc7740c7c..0be162bbdf 100644
--- a/src/s_tir/transform/inject_virtual_thread.cc
+++ b/src/s_tir/transform/inject_virtual_thread.cc
@@ -164,7 +164,7 @@ class VarTouchedAnalysis : public StmtExprVisitor {
return std::nullopt;
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
expr_touched_->Reset(false);
expr_touched_->Visit(op->value);
for (const auto& index : op->indices) {
@@ -327,7 +327,7 @@ class VTInjector : public s_tir::IRMutatorWithAnalyzer {
args.push_back(predicate);
return Call(op->ty, op->op, args, op->attrs, op->ty_args, op->span);
}
- BufferStore access = VisitBufferAccess(BufferStore(buffer, value.value(), indices, op->span));
+ TensorStore access = VisitBufferAccess(TensorStore(buffer, value.value(), indices, op->span));
ffi::Array<Expr> args{access->buffer.var(), access->value};
for (const PrimExpr& index : access->indices) args.push_back(index);
args.push_back(predicate);
@@ -377,11 +377,11 @@ class VTInjector : public s_tir::IRMutatorWithAnalyzer {
}
return VisitBufferAccess(std::move(node));
}
- // BufferStore
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ // TensorStore
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
auto value = Mutate(op->value);
auto indices = Mutate(op->indices).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
- BufferStore node = ffi::GetRef<BufferStore>(op);
+ TensorStore node = ffi::GetRef<TensorStore>(op);
if (!value.UnchangedOrSameAs(op->value) || !indices.UnchangedOrSameAs(op->indices)) {
auto* n = node.CopyOnWrite();
n->value = std::move(value).ValueOrUnchanged(op->value);
@@ -706,7 +706,7 @@ class VTInjector : public s_tir::IRMutatorWithAnalyzer {
*
* Maps from the buffer_var of an allocate node to the original
* extent of the allocation. Used when rewriting the indices of
- * TensorLoad/BufferStore.
+ * TensorLoad/TensorStore.
*/
std::unordered_map<const VarNode*, PrimExpr> alloc_remap_;
/*! \brief Map of buffers that are modified.
diff --git a/src/s_tir/transform/lower_cross_thread_reduction.cc b/src/s_tir/transform/lower_cross_thread_reduction.cc
index 176b2a6efa..08cd13e1c0 100644
--- a/src/s_tir/transform/lower_cross_thread_reduction.cc
+++ b/src/s_tir/transform/lower_cross_thread_reduction.cc
@@ -193,12 +193,12 @@ class BufferReplacer : public StmtExprMutator {
return replacement ? MakeTensorLoad(replacement.value(), {0}) : ffi::GetRef<TensorLoad>(load);
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* store, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* store, InplaceMode inplace_mode) final {
if (auto replacement = VarRemapGet(store->buffer).as<TensorVar>()) {
PrimExpr value = StmtExprMutator::Mutate(ffi::AnyView(store->value), inplace_mode)
.ValueOrUnchanged(store->value)
.as_or_throw<PrimExpr>();
- return BufferStore(replacement.value(), std::move(value), {0});
+ return TensorStore(replacement.value(), std::move(value), {0});
} else {
return StmtExprMutator::Mutate_(store, inplace_mode);
}
@@ -374,7 +374,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize,
inits.reserve(n_buffers);
for (int i = 0; i < n_buffers; ++i) {
inits.push_back(
- BufferStore(it_buffers.value()[i], reducer->identity_element[i], {IntImm::Int32(0)}));
+ TensorStore(it_buffers.value()[i], reducer->identity_element[i], {IntImm::Int32(0)}));
}
stmts.push_back(SBlockRealize(/*iter_values=*/{},
/*predicate=*/IntImm::Bool(true),
@@ -516,7 +516,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize,
.as_or_throw<PrimExpr>());
}
for (int i = 0; i < n_buffers; ++i) {
- wb_updates.push_back(BufferStore(
+ wb_updates.push_back(TensorStore(
wb_buffers[i], MakeTensorLoad(ct_buffers[i], {IntImm::Int32(0)}), wb_indices));
wb_regions.push_back(BufferRegion(wb_buffers[i], region));
}
@@ -733,11 +733,11 @@ class CrossThreadReductionTransformer : public StmtExprMutator {
}
}
- // Condition 3. Get the identity values of the block init and the BufferStore block combiner
+ // Condition 3. Get the identity values of the block init and the TensorStore block combiner
// updates of the reduction. Extract the commutative reducer, combiner lhs and combiner rhs from
// the reduction identities and the reduction combiner.
ffi::Array<PrimExpr> init_values{nullptr};
- ffi::Array<BufferStore> updates{nullptr};
+ ffi::Array<TensorStore> updates{nullptr};
te::CommReducer reducer{nullptr};
ffi::Array<PrimExpr> combiner_lhs{nullptr};
ffi::Array<PrimExpr> combiner_rhs{nullptr};
@@ -750,7 +750,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator {
int is_local_buf = -1;
ffi::Array<TensorVar> reduction_buffers;
reduction_buffers.reserve(updates.size());
- for (const BufferStore& buf_store : updates) {
+ for (const TensorStore& buf_store : updates) {
reduction_buffers.push_back(buf_store->buffer);
if (buf_store->buffer.scope() == "local") {
TVM_FFI_CHECK_NE(is_local_buf, 0, ValueError)
diff --git a/src/s_tir/transform/lower_match_buffer.cc b/src/s_tir/transform/lower_match_buffer.cc
index c6a121a7aa..d609e94706 100644
--- a/src/s_tir/transform/lower_match_buffer.cc
+++ b/src/s_tir/transform/lower_match_buffer.cc
@@ -122,13 +122,13 @@ class MatchBufferLower : public StmtExprMutator {
return StmtExprMutator::Mutate_(op, inplace_mode);
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
// Save the original buffer before base class mutation may remap it
TensorVar orig_buffer = op->buffer;
- BufferStore stmt = StmtExprMutator::Mutate_(op, inplace_mode)
+ TensorStore stmt = StmtExprMutator::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<BufferStore>();
- op = stmt.as<BufferStoreNode>();
+ .as_or_throw<TensorStore>();
+ op = stmt.as<TensorStoreNode>();
TVM_FFI_ICHECK(op != nullptr);
// Look up using original buffer (before the inherited Var environment may have remapped it)
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 061bf64fbe..d307a3dc8d 100644
--- a/src/s_tir/transform/manifest_shared_memory_local_stage.cc
+++ b/src/s_tir/transform/manifest_shared_memory_local_stage.cc
@@ -55,11 +55,11 @@ class IntermediateStageRewriter {
: ancestor_loop_or_blocks_(ancestor_loop_or_blocks) {}
std::tuple<TensorVar, TensorVar, SBlock, Stmt> Rewrite(const SBlockNode* block) {
- const BufferStoreNode* store = block->body.as<BufferStoreNode>();
+ const TensorStoreNode* store = block->body.as<TensorStoreNode>();
TVM_FFI_CHECK(store != nullptr && runtime::StorageScope::Create(store->buffer.scope()).rank ==
runtime::StorageRank::kShared,
ValueError)
- << "Expect the body of the block to be BufferStore to shared memory.";
+ << "Expect the body of the block to be TensorStore to shared memory.";
const TensorVar& target_buffer = store->buffer;
@@ -74,10 +74,10 @@ class IntermediateStageRewriter {
// Step 3: Create TensorLoad from the intermediate buffer
TensorLoad new_buffer_load = MakeTensorLoad(new_buffer, buffer_indices);
- BufferStore new_buffer_store = block->body.as_or_throw<BufferStore>();
- new_buffer_store.CopyOnWrite()->value = new_buffer_load;
+ TensorStore new_tensor_store = block->body.as_or_throw<TensorStore>();
+ new_tensor_store.CopyOnWrite()->value = new_buffer_load;
SBlock new_block = ffi::GetRef<SBlock>(block);
- new_block.CopyOnWrite()->body = std::move(new_buffer_store);
+ new_block.CopyOnWrite()->body = std::move(new_tensor_store);
return {target_buffer, new_buffer, new_block, local_stage};
}
@@ -123,9 +123,9 @@ class IntermediateStageRewriter {
/*! \brief Create the intermediate stage. */
Stmt MakeLocalStage(const SBlockNode* block, const TensorVar& new_buffer,
ffi::Array<PrimExpr> local_stage_indices,
- std::vector<const ForNode*> relaxed_loops, const BufferStoreNode* store) {
- // Step 0: Create the body of the local stage, which is BufferStore to the intermediate buffer.
- Stmt local_stage = BufferStore(new_buffer, store->value, local_stage_indices);
+ std::vector<const ForNode*> relaxed_loops, const TensorStoreNode* store) {
+ // Step 0: Create the body of the local stage, which is TensorStore to the intermediate buffer.
+ Stmt local_stage = TensorStore(new_buffer, store->value, local_stage_indices);
// Step 1: Make block and block realize
TensorRegion write_buffer_region = BufferRegionFromPoint(new_buffer, local_stage_indices);
diff --git a/src/s_tir/transform/memhammer_coalesce.cc b/src/s_tir/transform/memhammer_coalesce.cc
index 63814e66cf..ec061b3522 100644
--- a/src/s_tir/transform/memhammer_coalesce.cc
+++ b/src/s_tir/transform/memhammer_coalesce.cc
@@ -165,7 +165,7 @@ ffi::Array<PrimExpr> GetMapping(const Stmt& stmt, const ConstraintSet& constrain
while (const ForNode* loop = body.as<ForNode>()) {
body = loop->body;
}
- const BufferStoreNode* buf_store = TVM_TYPE_AS(body, BufferStoreNode);
+ const TensorStoreNode* buf_store = TVM_TYPE_AS(body, TensorStoreNode);
TensorRegion write_region = constraints.write_region;
const ffi::Array<PrimExpr>& write_index = buf_store->indices;
TVM_FFI_ICHECK(
@@ -236,7 +236,7 @@ Stmt InverseMapping::Rewrite(const Stmt& stmt, const ConstraintSet& constraints,
}
TensorLoad new_buf_load =
MakeTensorLoad(read_region->source.as_or_throw<tvm::tirx::TensorVar>(), read_index);
- BufferStore new_buf_store = BufferStore(write_region->source.as_or_throw<tvm::tirx::TensorVar>(),
+ TensorStore new_buf_store = TensorStore(write_region->source.as_or_throw<tvm::tirx::TensorVar>(),
new_buf_load, write_index);
Stmt ret = new_buf_store;
// Step 3.3 construct loop body
diff --git a/src/s_tir/transform/memhammer_intermediate_stage.cc b/src/s_tir/transform/memhammer_intermediate_stage.cc
index ad7be357ba..4ae5f0d613 100644
--- a/src/s_tir/transform/memhammer_intermediate_stage.cc
+++ b/src/s_tir/transform/memhammer_intermediate_stage.cc
@@ -313,7 +313,7 @@ std::pair<Stmt, SeqStmt> InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S
TVM_FFI_ICHECK(target_buffer_load);
}
- const BufferStoreNode* buf_store = TVM_TYPE_AS(body, BufferStoreNode);
+ const TensorStoreNode* buf_store = TVM_TYPE_AS(body, TensorStoreNode);
ffi::Array<PrimExpr> cache_indices;
ffi::Array<PrimExpr> new_shape;
bool use_rank_promotion = false;
@@ -401,7 +401,7 @@ std::pair<Stmt, SeqStmt> InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S
ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(generate_body, map_var).as_or_throw<Stmt>();
} else {
generate_body =
- BufferStore(new_buffer,
+ TensorStore(new_buffer,
ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(buf_store->value, map_var)
.as_or_throw<PrimExpr>(),
subst_cache_indices);
@@ -438,9 +438,9 @@ std::pair<Stmt, SeqStmt> InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S
if (is_write_cache) {
TensorLoad new_buffer_load = MakeTensorLoad(new_buffer, cache_indices);
rewrite_body =
- BufferStore(new_buffer, ffi::GetRef<TensorLoad>(target_buffer_load), cache_indices);
+ TensorStore(new_buffer, ffi::GetRef<TensorLoad>(target_buffer_load), cache_indices);
} else {
- rewrite_body = BufferStore(buf_store->buffer, MakeTensorLoad(new_buffer, cache_indices),
+ rewrite_body = TensorStore(buf_store->buffer, MakeTensorLoad(new_buffer, cache_indices),
buf_store->indices);
}
if (predicate.has_value()) {
diff --git a/src/s_tir/transform/memhammer_lower_auto_copy.cc b/src/s_tir/transform/memhammer_lower_auto_copy.cc
index e77944878f..337b97c0bd 100644
--- a/src/s_tir/transform/memhammer_lower_auto_copy.cc
+++ b/src/s_tir/transform/memhammer_lower_auto_copy.cc
@@ -218,11 +218,11 @@ class AutoPadder {
return load;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* _op, InplaceMode inplace_mode) final {
- BufferStore store = StmtExprMutator::Mutate_(_op, inplace_mode)
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* _op, InplaceMode inplace_mode) final {
+ TensorStore store = StmtExprMutator::Mutate_(_op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(_op))
- .as_or_throw<BufferStore>();
- BufferStoreNode* op = store.CopyOnWrite();
+ .as_or_throw<TensorStore>();
+ TensorStoreNode* op = store.CopyOnWrite();
if (auto replacement = VarRemapGet(op->buffer).as<TensorVar>()) {
op->buffer = replacement.value();
}
@@ -542,7 +542,7 @@ class AutoPadder {
* The iteration space would be {{0, 1}, {0, 4, ..., 60}}.
* \param op the buffer store
*/
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
runtime::StorageScope scope = runtime::StorageScope::Create(op->buffer.scope());
if (scope.rank == runtime::StorageRank::kShared) {
ffi::Array<PrimExpr> substitued_indices;
diff --git a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
index 28e06129da..aeb560b835 100644
--- a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
+++ b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
@@ -136,8 +136,8 @@ Stmt RewriteWmmaLoad(Stmt stmt) {
{loops[n - 1]->loop_var, IntSet::FromMinExtent(loops[n - 1]->min, loops[n - 1]->extent)},
{loops[n - 2]->loop_var, IntSet::FromMinExtent(loops[n - 2]->min, loops[n - 2]->extent)},
};
- // TODO(tian): the assumption that the RHS of BufferStore is TensorLoad may not be accurate
- const BufferStoreNode* buf_store = TVM_TYPE_AS(body, BufferStoreNode);
+ // TODO(tian): the assumption that the RHS of TensorStore is TensorLoad may not be accurate
+ const TensorStoreNode* buf_store = TVM_TYPE_AS(body, TensorStoreNode);
const TensorLoadNode* buf_load = TVM_TYPE_AS(buf_store->value, TensorLoadNode);
TensorVar src_buffer = buf_load->source.as_or_throw<tvm::tirx::TensorVar>();
@@ -233,8 +233,8 @@ Stmt RewriteWmmaStore(Stmt stmt) {
{loops[n - 1]->loop_var, IntSet::FromMinExtent(loops[n - 1]->min, loops[n - 1]->extent)},
{loops[n - 2]->loop_var, IntSet::FromMinExtent(loops[n - 2]->min, loops[n - 2]->extent)},
};
- // TODO(tian): the assumption that the RHS of BufferStore is TensorLoad may not be accurate
- const BufferStoreNode* buf_store = TVM_TYPE_AS(body, BufferStoreNode);
+ // TODO(tian): the assumption that the RHS of TensorStore is TensorLoad may not be accurate
+ const TensorStoreNode* buf_store = TVM_TYPE_AS(body, TensorStoreNode);
const TensorLoadNode* buf_load = nullptr;
auto walk_fn = [&](const TensorLoad& load) -> ffi::Expected<ffi::WalkResult> {
if (load->source.as_or_throw<tvm::tirx::TensorVar>().scope() == "wmma.accumulator") {
@@ -451,7 +451,7 @@ Stmt RewriteMmaStore(Stmt stmt) {
};
// Step 2. Find matrixC buffer
- const BufferStoreNode* buf_store = TVM_TYPE_AS(body, BufferStoreNode);
+ const TensorStoreNode* buf_store = TVM_TYPE_AS(body, TensorStoreNode);
const TensorLoadNode* buf_load = nullptr;
auto walk_fn = [&](const TensorLoad& load) -> ffi::Expected<ffi::WalkResult> {
if (load->source.as_or_throw<tvm::tirx::TensorVar>().scope() == "m16n8k8.matrixC") {
@@ -506,7 +506,7 @@ Stmt RewriteMmaStore(Stmt stmt) {
{StringImm("threadIdx.x"), IntImm::Int32(32)}, {tx}, DictAttrs(), /*body=*/
For(vec.as_or_throw<PrimVar>(), 0, 2, ForKind::kVectorized,
/*body=*/
- BufferStore(new_tgt_buffer,
+ TensorStore(new_tgt_buffer,
MakeTensorLoad(new_src_buffer, {floordiv(tx, 4),
floormod(tx, 4) * 2 + vec}),
{floordiv(tx, 4), floormod(tx, 4) * 2 + vec}))),
diff --git a/src/s_tir/transform/merge_shared_memory_allocations.cc b/src/s_tir/transform/merge_shared_memory_allocations.cc
index 010b685e52..1b539c50cc 100644
--- a/src/s_tir/transform/merge_shared_memory_allocations.cc
+++ b/src/s_tir/transform/merge_shared_memory_allocations.cc
@@ -193,7 +193,7 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor {
return StmtExprVisitor::Visit_(op);
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
scope_.push_back(StmtEntry());
// visit subexpr
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->value));
@@ -598,10 +598,10 @@ class SharedMemoryRewriter : public StmtExprMutator {
return VisitBufferAccess(std::move(node));
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
auto node = StmtExprMutator::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<BufferStore>();
+ .as_or_throw<TensorStore>();
return VisitBufferAccess(std::move(node));
}
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 c56a3837fc..1185eca2a2 100644
--- a/src/s_tir/transform/plan_update_buffer_allocation_location.cc
+++ b/src/s_tir/transform/plan_update_buffer_allocation_location.cc
@@ -98,7 +98,7 @@ class BufferAllocateOrderCollector : public StmtExprVisitor {
return StmtExprVisitor::Visit_(op);
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
if (!find(op->buffer)) {
buffer_alloc_recorder_.push_back(op->buffer);
}
diff --git a/src/s_tir/transform/remove_store_undef.cc b/src/s_tir/transform/remove_store_undef.cc
index d1c8b5bf78..053941cb0c 100644
--- a/src/s_tir/transform/remove_store_undef.cc
+++ b/src/s_tir/transform/remove_store_undef.cc
@@ -37,7 +37,7 @@ namespace s_tir {
using namespace tvm::tirx;
struct UndefInfo {
- std::unordered_set<const BufferStoreNode*> undef_stores;
+ std::unordered_set<const TensorStoreNode*> undef_stores;
std::unordered_set<const VarNode*> undef_bind_vars;
};
@@ -53,7 +53,7 @@ class StoreUndefLocator : public StmtExprVisitor {
StoreUndefLocator() = default;
private:
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
// Check the value for undef.
bool stash_undef = false;
std::swap(has_undef_, stash_undef);
@@ -62,7 +62,7 @@ class StoreUndefLocator : public StmtExprVisitor {
if (stash_undef) {
auto value = op->value.as<PrimExpr>();
TVM_FFI_ICHECK(value && SideEffect(value.value()) <= CallEffectKind::kReadState)
- << "Error: T.undef() used in BufferStore expressions "
+ << "Error: T.undef() used in TensorStore expressions "
<< "must not have other side effects";
undef_stores_.insert(op);
}
@@ -125,10 +125,10 @@ class StoreUndefLocator : public StmtExprVisitor {
bool has_undef_{false};
std::unordered_set<const VarNode*> var_bindings_with_undef_;
- std::unordered_set<const BufferStoreNode*> undef_stores_;
+ std::unordered_set<const TensorStoreNode*> undef_stores_;
};
-// Remove BufferStores whose value depends on T.undef, and also
+// Remove TensorStores whose value depends on T.undef, and also
// remove Bind nodes whose value contains undef. Undef in buffer
// indices is already caught eagerly in the locator phase.
class StoreUndefRemover : public StmtExprMutator {
@@ -150,7 +150,7 @@ class StoreUndefRemover : public StmtExprMutator {
: stores_to_remove_(info.undef_stores), bind_vars_to_remove_(info.undef_bind_vars) {}
private:
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
if (stores_to_remove_.count(op)) {
return Evaluate(0);
} else {
@@ -166,7 +166,7 @@ class StoreUndefRemover : public StmtExprMutator {
}
}
- const std::unordered_set<const BufferStoreNode*>& stores_to_remove_;
+ const std::unordered_set<const TensorStoreNode*>& stores_to_remove_;
const std::unordered_set<const VarNode*>& bind_vars_to_remove_;
};
@@ -207,7 +207,7 @@ Pass ValidateAllUndefRemoved() {
if (!f->body.has_value()) return f;
bool contains_undef = ContainsUndefChecker::Check(f->body.value());
TVM_FFI_ICHECK(!contains_undef)
- << "Expected removal of BufferStore containing tirx::builtin::undef() "
+ << "Expected removal of TensorStore containing tirx::builtin::undef() "
<< "to remove all instances of tirx::builtin::undef(). "
<< "Instead, result was"
<< "\n"
diff --git a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc
index aa25b8e055..596c265aff 100644
--- a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc
+++ b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc
@@ -86,8 +86,8 @@ class RemoveLayoutRewriteBlock : public StmtExprMutator {
TVM_FFI_ICHECK(block->alloc_buffers.empty());
TVM_FFI_ICHECK(block->match_buffers.empty());
- // Step 1. Checking the body is a BufferStore
- const auto* store = block->body.as<BufferStoreNode>();
+ // Step 1. Checking the body is a TensorStore
+ const auto* store = block->body.as<TensorStoreNode>();
TVM_FFI_ICHECK(store);
// Step 2. Checking the rhs of buffer store is a TensorLoad
diff --git a/src/s_tir/transform/storage_access.cc b/src/s_tir/transform/storage_access.cc
index 3b48b36399..8800b81ae4 100644
--- a/src/s_tir/transform/storage_access.cc
+++ b/src/s_tir/transform/storage_access.cc
@@ -71,7 +71,7 @@ ffi::Optional<VisitInterrupt> StorageAccessVisitor::Visit_(const TensorLoadNode*
return StmtExprVisitor::Visit_(op);
}
-ffi::Optional<VisitInterrupt> StorageAccessVisitor::Visit_(const BufferStoreNode* op) {
+ffi::Optional<VisitInterrupt> StorageAccessVisitor::Visit_(const TensorStoreNode* op) {
allow_append_ = true;
TVM_FFI_ICHECK_EQ(curr_stmt_.access.size(), 0U);
curr_stmt_.stmt = op;
diff --git a/src/s_tir/transform/storage_access.h b/src/s_tir/transform/storage_access.h
index 795b70f2f4..9e30eae631 100644
--- a/src/s_tir/transform/storage_access.h
+++ b/src/s_tir/transform/storage_access.h
@@ -83,7 +83,7 @@ class StorageAccessVisitor : public StmtExprVisitor {
};
// override visitor pattern
ffi::Optional<VisitInterrupt> Visit_(const TensorLoadNode* op) final;
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final;
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final;
ffi::Optional<VisitInterrupt> Visit_(const EvaluateNode* op) final;
ffi::Optional<VisitInterrupt> Visit_(const BindNode* op) final;
ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op) final;
diff --git a/src/s_tir/transform/transform_mma_buffer_layout.cc b/src/s_tir/transform/transform_mma_buffer_layout.cc
index 4b7c590ca3..3df6cc116b 100644
--- a/src/s_tir/transform/transform_mma_buffer_layout.cc
+++ b/src/s_tir/transform/transform_mma_buffer_layout.cc
@@ -129,12 +129,12 @@ class MmaBufferLayoutTransformer : public StmtExprMutator {
return block;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) {
TensorVar original_buffer = op->buffer;
auto value = Mutate(op->value, inplace_mode);
auto indices =
Mutate(op->indices, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
- BufferStore store = ffi::GetRef<BufferStore>(op);
+ TensorStore store = ffi::GetRef<TensorStore>(op);
if (!value.UnchangedOrSameAs(op->value) || !indices.UnchangedOrSameAs(op->indices)) {
auto* n = store.CopyOnWrite();
n->value = std::move(value).ValueOrUnchanged(op->value);
@@ -154,7 +154,7 @@ class MmaBufferLayoutTransformer : public StmtExprMutator {
TVM_FFI_ICHECK(false)
<< "TransformMmaBufferLayout requires " << original_buffer.scope()
<< " buffers to be accessed through opaque ldmatrix/mma_sync operations, but found "
- "an explicit BufferStore.";
+ "an explicit TensorStore.";
}
}
return store;
diff --git a/src/s_tir/transform/using_assume_to_reduce_branches.cc b/src/s_tir/transform/using_assume_to_reduce_branches.cc
index 1aae149d7b..0c82c76bde 100644
--- a/src/s_tir/transform/using_assume_to_reduce_branches.cc
+++ b/src/s_tir/transform/using_assume_to_reduce_branches.cc
@@ -220,10 +220,10 @@ class ParseAssumeAndOvercompute : public IRMutatorWithAnalyzer {
return ffi::Unchanged();
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
- BufferStore store = Parent::Mutate_(op, inplace_mode)
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
+ TensorStore store = Parent::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<BufferStore>();
+ .as_or_throw<TensorStore>();
if (!op->unique()) inplace_mode = InplaceMode::kDisallow;
// Eliminate the builtin if_then_else statement
@@ -265,23 +265,23 @@ class ParseAssumeAndOvercompute : public IRMutatorWithAnalyzer {
if (ffi::StructuralEqual()(then_clause_in_then_context, else_clause_in_then_context)) {
PrimExpr value = analyzer_->Simplify(else_clause);
if (inplace_mode == InplaceMode::kAllow) {
- const_cast<BufferStoreNode*>(op)->value = std::move(value);
+ const_cast<TensorStoreNode*>(op)->value = std::move(value);
return ffi::Unchanged();
} else {
- auto copy = ffi::make_object<BufferStoreNode>(*op);
+ auto copy = ffi::make_object<TensorStoreNode>(*op);
copy->value = std::move(value);
- return BufferStore(std::move(copy));
+ return TensorStore(std::move(copy));
}
} else if (ffi::StructuralEqual()(then_clause_in_else_context,
else_clause_in_else_context)) {
PrimExpr value = analyzer_->Simplify(then_clause);
if (inplace_mode == InplaceMode::kAllow) {
- const_cast<BufferStoreNode*>(op)->value = std::move(value);
+ const_cast<TensorStoreNode*>(op)->value = std::move(value);
return ffi::Unchanged();
} else {
- auto copy = ffi::make_object<BufferStoreNode>(*op);
+ auto copy = ffi::make_object<TensorStoreNode>(*op);
copy->value = std::move(value);
- return BufferStore(std::move(copy));
+ return TensorStore(std::move(copy));
}
} else {
return Parent::Mutate_(op, inplace_mode);
diff --git a/src/script/printer/script_printer.cc b/src/script/printer/script_printer.cc
index 66db749495..1e21537ff9 100644
--- a/src/script/printer/script_printer.cc
+++ b/src/script/printer/script_printer.cc
@@ -126,7 +126,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
RegisterScriptRepr<tirx::RegionStmtNode>();
RegisterScriptRepr<tirx::BindNode>();
RegisterScriptRepr<tirx::BreakNode>();
- RegisterScriptRepr<tirx::BufferStoreNode>();
+ RegisterScriptRepr<tirx::TensorStoreNode>();
RegisterScriptRepr<tirx::TensorTypeNode>();
RegisterScriptRepr<tirx::ComposeLayoutNode>();
RegisterScriptRepr<tirx::ContinueNode>();
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index c4e896ea63..ae85afe9db 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -2054,7 +2054,7 @@ llvm::Value* CodeGenLLVM::Dispatch_(const prim::BroadcastNode* op) {
return builder_->CreateShuffleVector(value, undef, mask);
}
-void CodeGenLLVM::Dispatch_(const BufferStoreNode* op) {
+void CodeGenLLVM::Dispatch_(const TensorStoreNode* op) {
EmitDebugLocation(op);
PrimType value_dtype = PrimType(op->value.ty()->dtype);
Var buffer_var = op->buffer.var();
diff --git a/src/target/llvm/codegen_llvm.h b/src/target/llvm/codegen_llvm.h
index 199953add3..ef4e723e6f 100644
--- a/src/target/llvm/codegen_llvm.h
+++ b/src/target/llvm/codegen_llvm.h
@@ -237,7 +237,7 @@ class CodeGenLLVM : public tirx::ExprFunctor<llvm::Value*(const Expr&)>,
llvm::Value* Dispatch_(const prim::ShuffleNode* op) override;
llvm::Value* Dispatch_(const prim::BroadcastNode* op) override;
// stmt
- void Dispatch_(const BufferStoreNode* op) override;
+ void Dispatch_(const TensorStoreNode* op) override;
void Dispatch_(const ForNode* op) override;
void Dispatch_(const WhileNode* op) override;
void Dispatch_(const ReturnNode* op) override;
@@ -350,7 +350,7 @@ class CodeGenLLVM : public tirx::ExprFunctor<llvm::Value*(const Expr&)>,
* number of lanes being accessed.
*
* \param value_dtype The datatype to be read from (TensorLoad) or
- * written to (BufferStore) the buffer.
+ * written to (TensorStore) the buffer.
*
* \param make_instruction A callback function that generates that
* actual call.
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index 11ec216935..7cb95263f6 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -1050,7 +1050,7 @@ void CodeGenC::Dispatch_(const TensorLoadNode* op, std::ostream& os) { // NOLIN
}
}
-void CodeGenC::Dispatch_(const BufferStoreNode* op) {
+void CodeGenC::Dispatch_(const TensorStoreNode* op) {
TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "Store to non-flat memory not supported.";
PrimType value_ty = op->value.ty();
diff --git a/src/target/source/codegen_c.h b/src/target/source/codegen_c.h
index 96af65a167..5d3e768745 100644
--- a/src/target/source/codegen_c.h
+++ b/src/target/source/codegen_c.h
@@ -210,7 +210,7 @@ class CodeGenC : public tirx::ExprFunctor<void(const Expr&, std::ostream&)>,
void Dispatch_(const StringImmNode* op, std::ostream& os) override; // NOLINT(*)
// statment
void Dispatch_(const BindNode* op) override;
- void Dispatch_(const BufferStoreNode* op) override;
+ void Dispatch_(const TensorStoreNode* op) override;
void Dispatch_(const ForNode* op) override;
void Dispatch_(const WhileNode* op) override;
void Dispatch_(const ReturnNode* op) override;
diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc
index deba53108a..a1a5c62a35 100644
--- a/src/te/operation/create_primfunc.cc
+++ b/src/te/operation/create_primfunc.cc
@@ -452,7 +452,7 @@ Stmt GenerateInitStmt(const ffi::Array<PrimExpr>& indices, const ffi::Array<Tens
for (int i = 0; i < n_buffers; ++i) {
const TensorVar& buffer = buffers[i];
PrimExpr identity = f_transform_and_remap(reduce->combiner->identity_element[i]);
- init_stmts.push_back(BufferStore(buffer, identity, indices));
+ init_stmts.push_back(TensorStore(buffer, identity, indices));
}
return SeqStmt::Flatten(init_stmts);
}
@@ -505,7 +505,7 @@ Stmt GenerateBodyStmt(const ffi::Array<PrimExpr>& indices, const ffi::Array<Tens
temp_vars.reserve(n_buffers);
body_stmts.reserve(n_buffers);
- // - When there is only one buffer, we directly create a BufferStore which stores "combiner(lhs,
+ // - When there is only one buffer, we directly create a TensorStore which stores "combiner(lhs,
// rhs)" into the target buffer position.
// - In case there are multiple buffers, to avoid incorrect results, we create some intermediate
// variables and use Bind nodes to bind the variables with "combiner(lhs, rhs)". After that,
@@ -520,7 +520,7 @@ Stmt GenerateBodyStmt(const ffi::Array<PrimExpr>& indices, const ffi::Array<Tens
PrimExpr combined = reduce->combiner.get()->operator()(lhs, rhs)[i];
return f_transform_and_remap(combined);
}();
- body_stmts.push_back(BufferStore(buffer, value, indices));
+ body_stmts.push_back(TensorStore(buffer, value, indices));
}
Stmt body = SeqStmt::Flatten(body_stmts);
if (n_buffers > 1) {
@@ -538,7 +538,7 @@ Stmt GenerateBodyStmt(const ffi::Array<PrimExpr>& indices, const ffi::Array<Tens
// Case 2. Data parallel compute
TVM_FFI_ICHECK_EQ(buffers.size(), 1);
const PrimExpr& compute_body = f_transform_and_remap(expr_body);
- return BufferStore(buffers[0], analyzer->Simplify(compute_body), indices);
+ return TensorStore(buffers[0], analyzer->Simplify(compute_body), indices);
}
}
diff --git a/src/tirx/analysis/verify_memory.cc b/src/tirx/analysis/verify_memory.cc
index e3c25db845..d195d73593 100644
--- a/src/tirx/analysis/verify_memory.cc
+++ b/src/tirx/analysis/verify_memory.cc
@@ -98,7 +98,7 @@ class MemoryAccessVerifier final : public StmtExprVisitor {
return StmtExprVisitor::Visit_(op);
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
HandleLoadStoreToVariable(op->buffer.var());
return StmtExprVisitor::Visit_(op);
}
diff --git a/src/tirx/analysis/verify_well_formed.h b/src/tirx/analysis/verify_well_formed.h
index 27e92b2615..95e54b0b4a 100644
--- a/src/tirx/analysis/verify_well_formed.h
+++ b/src/tirx/analysis/verify_well_formed.h
@@ -108,10 +108,10 @@ class UndefinedVarVerifier : public Verifier<UndefinedVarVerifier<PathVisitor>,
* - DeclTensor statement
* - Dialect-specific definitions exposed by PathVisitor
*
- * it must not appear in a TensorLoad, BufferStore, or BufferRegion outside that declaration's
+ * it must not appear in a TensorLoad, TensorStore, or BufferRegion outside that declaration's
* scope.
*
- * All buffers that appear in TensorLoad or BufferStore must have a prior declaration.
+ * All buffers that appear in TensorLoad or TensorStore must have a prior declaration.
*/
template <typename PathVisitor>
class UndefinedBufferVerifier : public Verifier<UndefinedBufferVerifier<PathVisitor>, PathVisitor> {
diff --git a/src/tirx/ir/buffer.cc b/src/tirx/ir/buffer.cc
index 374c1ff29f..30de2b1357 100644
--- a/src/tirx/ir/buffer.cc
+++ b/src/tirx/ir/buffer.cc
@@ -555,7 +555,7 @@ Stmt TensorVar::vstore(ffi::Array<PrimExpr> begin, PrimExpr value) const {
indices.Set(indices.size() - 1, prim::Ramp(base, 1, factor));
}
}
- return BufferStore(*this, value, indices);
+ return TensorStore(*this, value, indices);
}
ffi::String TensorVar::scope() const { return (*this)->storage_scope; }
diff --git a/src/tirx/ir/data_type_rewriter.cc b/src/tirx/ir/data_type_rewriter.cc
index a5837f8dbc..9927e98fc0 100644
--- a/src/tirx/ir/data_type_rewriter.cc
+++ b/src/tirx/ir/data_type_rewriter.cc
@@ -344,9 +344,9 @@ UnchangedOr<ffi::Any> IndexDataTypeRewriter::Mutate(ffi::AnyView value, InplaceM
return result;
}
-UnchangedOr<Stmt> IndexDataTypeRewriter::Mutate_(const BufferStoreNode* op,
+UnchangedOr<Stmt> IndexDataTypeRewriter::Mutate_(const TensorStoreNode* op,
InplaceMode inplace_mode) {
- BufferStore store = ffi::GetRef<BufferStore>(op);
+ TensorStore store = ffi::GetRef<TensorStore>(op);
TensorVar new_buffer = Mutate(op->buffer, inplace_mode)
.as_or_throw<UnchangedOr<TensorVar>>()
diff --git a/src/tirx/ir/data_type_rewriter.h b/src/tirx/ir/data_type_rewriter.h
index 6763bed0cf..0fe7227e0d 100644
--- a/src/tirx/ir/data_type_rewriter.h
+++ b/src/tirx/ir/data_type_rewriter.h
@@ -107,7 +107,7 @@ class IndexDataTypeRewriter : public DataTypeLegalizer {
explicit IndexDataTypeRewriter(const VTable* vtable) : DataTypeLegalizer(vtable) {}
using Parent = DataTypeLegalizer;
UnchangedOr<ffi::Any> Mutate(ffi::AnyView value, InplaceMode inplace_mode) override;
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) override;
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) override;
UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) override;
UnchangedOr<PrimExpr> Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) override;
ffi::Array<PrimExpr> VisitIndices(const ffi::Array<PrimExpr>& indices, InplaceMode inplace_mode);
diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc
index e34780039e..d52e65af06 100644
--- a/src/tirx/ir/stmt.cc
+++ b/src/tirx/ir/stmt.cc
@@ -604,20 +604,20 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> EvaluateMaybeInplaceMut
return ffi::Unchanged();
}
-TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> BufferStoreVisit(
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> TensorStoreVisit(
ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
- const BufferStoreNode* self =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BufferStoreNode>(value);
+ const TensorStoreNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorStoreNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->buffer));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->value));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->indices));
return std::nullopt;
}
-TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BufferStoreMutate(
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorStoreMutate(
ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
- const BufferStoreNode* self =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BufferStoreNode>(value);
+ const TensorStoreNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorStoreNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<TensorVar>, mapped_buffer,
mutator->MutateExpected(self->buffer));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_value,
@@ -629,17 +629,17 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BufferStoreMutate(
mapped_indices.UnchangedOrSameAs(self->indices)) {
return ffi::Unchanged();
}
- ffi::ObjectPtr<BufferStoreNode> copy = ffi::make_object<BufferStoreNode>(*self);
+ ffi::ObjectPtr<TensorStoreNode> copy = ffi::make_object<TensorStoreNode>(*self);
copy->buffer = std::move(mapped_buffer).ValueOrUnchanged(std::move(copy->buffer));
copy->value = std::move(mapped_value).ValueOrUnchanged(std::move(copy->value));
copy->indices = std::move(mapped_indices).ValueOrUnchanged(std::move(copy->indices));
return ffi::Any(std::move(copy));
}
-TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BufferStoreMaybeInplaceMutate(
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorStoreMaybeInplaceMutate(
ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
- BufferStoreNode* self = const_cast<BufferStoreNode*>(
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BufferStoreNode>(value));
+ TensorStoreNode* self = const_cast<TensorStoreNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorStoreNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
ffi::UnchangedOr<TensorVar>, mapped_buffer,
mutator->MutateExpected(self->buffer, ffi::InplaceMode::kAllow));
@@ -1132,12 +1132,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
[](Expr value, Span span) { return Evaluate(value, span); });
}
-// BufferStore
+// TensorStore
TVM_FFI_INLINE int GetLanesOrVScaleFactor(const PrimType& ty) {
return ty.IsScalableVector() ? ty.VScaleFactor() : ty.lanes();
}
-BufferStore::BufferStore(TensorVar buffer, PrimExpr value, ffi::Array<PrimExpr> indices, Span span)
+TensorStore::TensorStore(TensorVar buffer, PrimExpr value, ffi::Array<PrimExpr> indices, Span span)
: Stmt(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK_EQ(buffer->shape.size(), indices.size())
<< "TensorVar " << buffer.name() << " is " << buffer->shape.size()
@@ -1179,15 +1179,15 @@ BufferStore::BufferStore(TensorVar buffer, PrimExpr value, ffi::Array<PrimExpr>
buffer_dtype = buffer->dtype.WithLanes(buffer_lanes * index_lanes);
}
if (buffer_dtype != value_ty) {
- TVM_FFI_THROW(TypeError) << "dtype mismatch on BufferStore: " //
+ TVM_FFI_THROW(TypeError) << "dtype mismatch on TensorStore: " //
<< "buffer's dtype is `" << buffer->dtype //
<< "`, the lanes of indexing are: `" << index_lanes //
<< "`, the scalability is: `" << buffer_dtype.IsScalableVector()
<< "`, but RHS's dtype is `" << value_ty << "`";
}
- ffi::ObjectPtr<BufferStoreNode> node =
- ffi::make_object<BufferStoreNode>(std::move(buffer), std::move(value));
+ ffi::ObjectPtr<TensorStoreNode> node =
+ ffi::make_object<TensorStoreNode>(std::move(buffer), std::move(value));
node->indices = std::move(indices);
node->span = std::move(span);
data_ = std::move(node);
@@ -1195,18 +1195,18 @@ BufferStore::BufferStore(TensorVar buffer, PrimExpr value, ffi::Array<PrimExpr>
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- BufferStoreNode::RegisterReflection();
- refl::TypeAttrDef<BufferStoreNode>()
+ TensorStoreNode::RegisterReflection();
+ refl::TypeAttrDef<TensorStoreNode>()
.attr(refl::type_attr::kStructuralVisit,
- ffi::FStructuralVisit::FromNative<&BufferStoreVisit>())
+ ffi::FStructuralVisit::FromNative<&TensorStoreVisit>())
.attr(refl::type_attr::kStructuralMutate,
- ffi::FStructuralMutate::FromNative<&BufferStoreMutate>())
+ ffi::FStructuralMutate::FromNative<&TensorStoreMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- ffi::FStructuralMutate::FromNative<&BufferStoreMaybeInplaceMutate>());
+ ffi::FStructuralMutate::FromNative<&TensorStoreMaybeInplaceMutate>());
- refl::GlobalDef().def("tirx.BufferStore",
+ refl::GlobalDef().def("tirx.TensorStore",
[](TensorVar buffer, PrimExpr value, ffi::Array<PrimExpr> indices,
- Span span) { return BufferStore(buffer, value, indices, span); });
+ Span span) { return TensorStore(buffer, value, indices, span); });
}
// ScopeIdDefStmt
diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc
index 530dde0e2e..99021ff084 100644
--- a/src/tirx/ir/stmt_functor.cc
+++ b/src/tirx/ir/stmt_functor.cc
@@ -50,7 +50,7 @@ void StmtExprVisitor::InitVTable(VTable* vtable) {
SetDispatch<StmtExprVisitor, ReturnNode>(vtable);
SetDispatch<StmtExprVisitor, BreakNode>(vtable);
SetDispatch<StmtExprVisitor, ContinueNode>(vtable);
- SetDispatch<StmtExprVisitor, BufferStoreNode>(vtable);
+ SetDispatch<StmtExprVisitor, TensorStoreNode>(vtable);
SetDispatch<StmtExprVisitor, AssertStmtNode>(vtable);
SetDispatch<StmtExprVisitor, SeqStmtNode>(vtable);
SetDispatch<StmtExprVisitor, EvaluateNode>(vtable);
@@ -161,7 +161,7 @@ ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const ReturnNode* op) {
return this->Visit(op->value);
}
-ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const BufferStoreNode* op) {
+ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const TensorStoreNode* op) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->buffer));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->value));
for (const auto& child : op->indices) {
@@ -232,7 +232,7 @@ void StmtExprMutator::InitVTable(VTable* vtable) {
SetDispatch<StmtExprMutator, ReturnNode>(vtable);
SetDispatch<StmtExprMutator, BreakNode>(vtable);
SetDispatch<StmtExprMutator, ContinueNode>(vtable);
- SetDispatch<StmtExprMutator, BufferStoreNode>(vtable);
+ SetDispatch<StmtExprMutator, TensorStoreNode>(vtable);
SetDispatch<StmtExprMutator, AssertStmtNode>(vtable);
SetDispatch<StmtExprMutator, SeqStmtNode>(vtable);
SetDispatch<StmtExprMutator, EvaluateNode>(vtable);
@@ -417,7 +417,7 @@ UnchangedOr<Stmt> StmtExprMutator::Mutate_(const ContinueNode* op, InplaceMode i
return ffi::Unchanged();
}
-UnchangedOr<Stmt> StmtExprMutator::Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) {
+UnchangedOr<Stmt> StmtExprMutator::Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) {
auto buffer = Mutate(op->buffer, inplace_mode).as_or_throw<UnchangedOr<TensorVar>>();
auto value = Mutate(op->value, inplace_mode);
auto indices = Mutate(op->indices, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
@@ -425,13 +425,13 @@ UnchangedOr<Stmt> StmtExprMutator::Mutate_(const BufferStoreNode* op, InplaceMod
indices.UnchangedOrSameAs(op->indices))
return ffi::Unchanged();
if (inplace_mode == InplaceMode::kAllow) {
- auto* writable = const_cast<BufferStoreNode*>(op);
+ auto* writable = const_cast<TensorStoreNode*>(op);
if (!buffer.IsUnchanged()) writable->buffer = std::move(buffer).ValueUnchecked();
if (!value.IsUnchanged()) writable->value = std::move(value).ValueUnchecked();
if (!indices.IsUnchanged()) writable->indices = std::move(indices).ValueUnchecked();
return ffi::Unchanged();
}
- auto copy = ffi::make_object<BufferStoreNode>(*op);
+ auto copy = ffi::make_object<TensorStoreNode>(*op);
if (!buffer.IsUnchanged()) copy->buffer = std::move(buffer).ValueUnchecked();
if (!value.IsUnchanged()) copy->value = std::move(value).ValueUnchecked();
if (!indices.IsUnchanged()) copy->indices = std::move(indices).ValueUnchecked();
diff --git a/src/tirx/ir/tir_visitor_with_path.cc b/src/tirx/ir/tir_visitor_with_path.cc
index 53c8352e9e..65b1b11bd9 100644
--- a/src/tirx/ir/tir_visitor_with_path.cc
+++ b/src/tirx/ir/tir_visitor_with_path.cc
@@ -203,7 +203,7 @@ void TIRVisitorWithPath::Dispatch_(const BreakNode* op, AccessPath path) {}
void TIRVisitorWithPath::Dispatch_(const ContinueNode* op, AccessPath path) {}
-void TIRVisitorWithPath::Dispatch_(const BufferStoreNode* op, AccessPath path) {
+void TIRVisitorWithPath::Dispatch_(const TensorStoreNode* op, AccessPath path) {
Visit(op->value, path->Attr("value"));
VisitBufferUse(op->buffer, path->Attr("buffer"));
Visit(op->indices, path->Attr("indices"));
diff --git a/src/tirx/ir/tir_visitor_with_path.h b/src/tirx/ir/tir_visitor_with_path.h
index 964f66e222..3d646a9186 100644
--- a/src/tirx/ir/tir_visitor_with_path.h
+++ b/src/tirx/ir/tir_visitor_with_path.h
@@ -98,7 +98,7 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const Expr&, ffi::reflecti
Dispatch(obj, path);
}
- // Visit a buffer at a use site (TensorLoad, BufferStore, reads/writes).
+ // Visit a buffer at a use site (TensorLoad, TensorStore, reads/writes).
// By default, does not re-visit buffer fields (shape, strides, elem_offset),
// as those are visited at the definition site via EnterDef.
virtual void VisitBufferUse(const TensorVar& obj, ffi::reflection::AccessPath path);
@@ -148,7 +148,7 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const Expr&, ffi::reflecti
void Dispatch_(const ReturnNode* op, ffi::reflection::AccessPath path) override;
void Dispatch_(const BreakNode* op, ffi::reflection::AccessPath path) override;
void Dispatch_(const ContinueNode* op, ffi::reflection::AccessPath path) override;
- void Dispatch_(const BufferStoreNode* op, ffi::reflection::AccessPath path) override;
+ void Dispatch_(const TensorStoreNode* op, ffi::reflection::AccessPath path) override;
void Dispatch_(const AssertStmtNode* op, ffi::reflection::AccessPath path) override;
void Dispatch_(const SeqStmtNode* op, ffi::reflection::AccessPath path) override;
void Dispatch_(const EvaluateNode* op, ffi::reflection::AccessPath path) override;
diff --git a/src/tirx/script/ir_builder/ir.cc b/src/tirx/script/ir_builder/ir.cc
index e386b70926..7fe6aa1b0e 100644
--- a/src/tirx/script/ir_builder/ir.cc
+++ b/src/tirx/script/ir_builder/ir.cc
@@ -464,7 +464,7 @@ ElseFrame Else() {
return ElseFrame(n);
}
-tvm::tirx::Stmt BufferStore(TensorVar buffer, PrimExpr value, ffi::Array<PrimExpr> indices) {
+tvm::tirx::Stmt TensorStore(TensorVar buffer, PrimExpr value, ffi::Array<PrimExpr> indices) {
PrimType buffer_dtype = buffer->dtype;
PrimType index_ty = indices.empty() ? PrimType::Int(32) : indices.back().ty();
bool is_index_scalable = !indices.empty() && index_ty.IsScalableVector();
@@ -506,13 +506,13 @@ tvm::tirx::Stmt BufferStore(TensorVar buffer, PrimExpr value, ffi::Array<PrimExp
}
if (!lanes_match) {
- TVM_FFI_THROW(InternalError) << "TypeError: Incompatible types in BufferStore"
+ TVM_FFI_THROW(InternalError) << "TypeError: Incompatible types in TensorStore"
<< ": LHS is `" << lhs_dtype << "`, RHS is `" << rhs_dtype
<< "`, indexing lanes: " << index_lanes;
}
value = tvm::prim::cast(lhs_dtype, value);
}
- tvm::tirx::Stmt store = tvm::tirx::BufferStore(buffer, value, indices);
+ tvm::tirx::Stmt store = tvm::tirx::TensorStore(buffer, value, indices);
if (lhs_dtype != rhs_dtype) {
if (lhs_dtype.code() != rhs_dtype.code()) {
if ((lhs_dtype.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) &&
@@ -522,10 +522,10 @@ tvm::tirx::Stmt BufferStore(TensorVar buffer, PrimExpr value, ffi::Array<PrimExp
if (ffi::Optional<PrimFuncFrame> frame = IRBuilder::Current()->FindFrame<PrimFuncFrame>()) {
kernel_name = frame.value()->name.value_or("<anonymous>");
}
- LOG(WARNING) << "Casting in BufferStore may lose precision"
+ LOG(WARNING) << "Casting in TensorStore may lose precision"
<< ": LHS is `" << lhs_dtype << "`, RHS is `" << rhs_dtype
<< "`, indexing lanes: " << index_lanes << ", kernel: `" << kernel_name << "`"
- << "\nBufferStore:\n"
+ << "\nTensorStore:\n"
<< store;
}
}
@@ -697,7 +697,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("script.ir_builder.tirx.Else", Else)
.def("script.ir_builder.tirx.DeclTensor", DeclTensor)
.def("script.ir_builder.tirx.Region", Region)
- .def("script.ir_builder.tirx.BufferStore", BufferStore)
+ .def("script.ir_builder.tirx.TensorStore", TensorStore)
.def("script.ir_builder.tirx.Evaluate", Evaluate)
.def("script.ir_builder.tirx.Ptr", Ptr);
}
diff --git a/src/tirx/script/printer/buffer.cc b/src/tirx/script/printer/buffer.cc
index ae4ea66216..7198fcbd24 100644
--- a/src/tirx/script/printer/buffer.cc
+++ b/src/tirx/script/printer/buffer.cc
@@ -305,10 +305,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
kDocTranslate, FDocTranslate::FromNative<&TensorRegionDocTranslate>());
}
-ffi::Optional<ExprDoc> BufferStoreDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ffi::Optional<ExprDoc> TensorStoreDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
const ffi::Object* destination) {
const auto* store =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const tirx::BufferStoreNode>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const tirx::TensorStoreNode>(
input);
TVM_FFI_CHECK(destination == nullptr, TypeError)
<< "printer statement-only node cannot fulfill a destination";
@@ -322,8 +322,8 @@ ffi::Optional<ExprDoc> BufferStoreDocTranslate(DocTranslatorObj* d, ffi::AnyView
}
TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<tirx::BufferStoreNode>().attr(
- kDocTranslate, FDocTranslate::FromNative<&BufferStoreDocTranslate>());
+ ffi::reflection::TypeAttrDef<tirx::TensorStoreNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&TensorStoreDocTranslate>());
}
ffi::Optional<ExprDoc> TIRxTensorLoadDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc
index c0e8494e68..a0c632b77d 100644
--- a/src/tirx/script/printer/stmt.cc
+++ b/src/tirx/script/printer/stmt.cc
@@ -348,7 +348,7 @@ ffi::Optional<ExprDoc> SeqStmtDocTranslate(DocTranslatorObj* d, ffi::AnyView inp
if (i + 1 == stmt->seq.size()) continue;
const auto* alloc = stmt->seq[i].as<tirx::BindNode>();
const auto* allocation = alloc ? alloc->value.as<CallNode>() : nullptr;
- const auto* store = stmt->seq[i + 1].as<tirx::BufferStoreNode>();
+ const auto* store = stmt->seq[i + 1].as<tirx::TensorStoreNode>();
auto docs = d->CurrentScopeDocs();
if (!allocation || !allocation->op.same_as(tirx::builtin::alloc_tensor()) || !store ||
!alloc->var.same_as(store->buffer) || docs.empty())
@@ -373,7 +373,7 @@ ffi::Optional<ExprDoc> SeqStmtDocTranslate(DocTranslatorObj* d, ffi::AnyView inp
// Preserve both statement origins using ordinary annotation/assignment
// occurrences while retaining the value's more precise child origin.
d->RecordOrigin(scalar.value()->annotation.value(), ffi::GetRef<tirx::Bind>(alloc));
- d->RecordOrigin(scalar.value(), ffi::GetRef<tirx::BufferStore>(store));
+ d->RecordOrigin(scalar.value(), ffi::GetRef<tirx::TensorStore>(store));
docs.pop_back();
}
}
diff --git a/src/tirx/transform/flatten_buffer.cc b/src/tirx/transform/flatten_buffer.cc
index 51d7591ab0..fed7ba1aa4 100644
--- a/src/tirx/transform/flatten_buffer.cc
+++ b/src/tirx/transform/flatten_buffer.cc
@@ -235,12 +235,12 @@ class BufferFlattener : public IRMutatorWithAnalyzer {
op->span);
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
// The buffer and its indices must be flattened together by VisitBufferAccess.
auto value = Mutate(op->value, inplace_mode);
auto indices =
Mutate(op->indices, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
- BufferStore store = ffi::GetRef<BufferStore>(op);
+ TensorStore store = ffi::GetRef<TensorStore>(op);
if (!value.UnchangedOrSameAs(op->value) || !indices.UnchangedOrSameAs(op->indices)) {
auto* n = store.CopyOnWrite();
n->value = std::move(value).ValueOrUnchanged(op->value);
diff --git a/src/tirx/transform/lower_thread_allreduce.h b/src/tirx/transform/lower_thread_allreduce.h
index 1166e51f86..27cd434287 100644
--- a/src/tirx/transform/lower_thread_allreduce.h
+++ b/src/tirx/transform/lower_thread_allreduce.h
@@ -180,11 +180,11 @@ class ThreadAllreduceBuilder final : public DialectMutator {
return load;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
const VarNode* allocation = GetAllocationKey(op->buffer.get());
- BufferStore store = DialectMutator::Mutate_(op, inplace_mode)
+ TensorStore store = DialectMutator::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .template as_or_throw<BufferStore>();
+ .template as_or_throw<TensorStore>();
if (auto it = load_remap_.find(allocation); it != load_remap_.end()) {
const auto* replacement = it->second.template as<TensorLoadNode>();
@@ -391,7 +391,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
TVM_FFI_ICHECK_EQ(val.ty(), dtypes[i]);
PrimExpr splat = WarpShuffle(tirx::builtin::tvm_warp_shuffle(), new_alloc_bufs.back(),
val, reduce_extent * group_index);
- seq.push_back(BufferStore(buf, splat, {zero_index}));
+ seq.push_back(TensorStore(buf, splat, {zero_index}));
}
} else {
int n_warps = reduce_extent / warp_size_;
@@ -421,7 +421,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
new_alloc_bufs.push_back(reduce_results[i]
.as_or_throw<TensorLoad>()
->source.as_or_throw<tvm::tirx::TensorVar>());
- write_staging_buf.push_back(BufferStore(
+ write_staging_buf.push_back(TensorStore(
/*buffer=*/staging_shared_bufs[i],
/*value=*/reduce_results[i],
/*indices=*/{group_index * n_warps + floordiv(reduce_index, warp_size_)}));
@@ -453,7 +453,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
/*shape=*/{IntImm(reduce_index.ty(), group_extent)},
/*dtype=*/buffers[i]->dtype, /*name=*/"red_result", /*storage_scope=*/"shared");
write_result.push_back(
- BufferStore(broadcast_shared_buf, reduce_results[i], {group_index}));
+ TensorStore(broadcast_shared_buf, reduce_results[i], {group_index}));
// Update `reduce_results`, pointing to the value loaded from the shared memory buffer.
reduce_results[i] = MakeTensorLoad(broadcast_shared_buf, {group_index});
}
@@ -483,7 +483,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
// special case, no reduction is needed.
std::vector<Stmt> stores;
for (size_t i = 0; i < size; ++i) {
- stores.push_back(BufferStore(buffers[i], values[i], {0}));
+ stores.push_back(TensorStore(buffers[i], values[i], {0}));
}
return SeqStmt::Flatten(stores);
}
@@ -493,7 +493,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
for (size_t idx = 0; idx < size; ++idx) {
shared_bufs.push_back(decl_tensor({IntImm(group_index.ty(), group_extent * reduce_extent)},
dtypes[idx], "red_buf" + std::to_string(idx), "shared"));
- seq.emplace_back(BufferStore(shared_bufs[idx], values[idx],
+ seq.emplace_back(TensorStore(shared_bufs[idx], values[idx],
{BufIndex(reduce_index, group_index, reduce_extent)}));
}
seq.emplace_back(SyncThread("shared"));
@@ -557,7 +557,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
for (int idx = 0; idx < n_buffers; ++idx) {
shared_bufs.push_back(
decl_tensor(shape, dtypes[idx], "red_buf" + std::to_string(idx), "local"));
- load_values.push_back(BufferStore(shared_bufs[idx], src_values[idx], zero_indices));
+ load_values.push_back(TensorStore(shared_bufs[idx], src_values[idx], zero_indices));
// Uses a local variable to store the shuffled data. Later
// on, an allocation will be built for this local variable.
@@ -576,7 +576,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
ffi::Optional<TensorVar> mask_buffer;
if (need_warp_shuffle_mask_) {
mask_buffer = decl_tensor(shape, mask.ty(), "mask", "local");
- seq->emplace_back(BufferStore(mask_buffer.value(), mask, zero_indices));
+ seq->emplace_back(TensorStore(mask_buffer.value(), mask, zero_indices));
// Push the buffer description. Later this will have an
// allocation built for it.
local_bufs.push_back(mask_buffer.value());
@@ -611,7 +611,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
PrimExpr other =
WarpShuffle(tirx::builtin::tvm_warp_shuffle_down(), mask_buffer, val, offset);
TensorVar local_buf = local_bufs[i];
- Stmt s = BufferStore(local_buf, other, zero_indices);
+ Stmt s = TensorStore(local_buf, other, zero_indices);
seq->push_back(s);
TensorLoad load = MakeTensorLoad(local_buf, zero_indices);
@@ -627,7 +627,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
stores.reserve(n_buffers);
for (int i = 0; i < n_buffers; ++i) {
TensorVar buf = shared_bufs[i];
- stores.push_back(BufferStore(buf, ret[i], zero_indices));
+ stores.push_back(TensorStore(buf, ret[i], zero_indices));
}
// During the sub-warp reduction, values from inactive threads could be read,
@@ -688,7 +688,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
std::vector<Stmt> stores;
stores.reserve(size);
for (size_t i = 0; i < size; ++i) {
- stores.push_back(BufferStore(shared_bufs[i], ret[i], {buf_index}));
+ stores.push_back(TensorStore(shared_bufs[i], ret[i], {buf_index}));
}
return SeqStmt::Flatten(stores);
};
diff --git a/src/tirx/transform/lower_tirx_cleanup.cc b/src/tirx/transform/lower_tirx_cleanup.cc
index 308d7cdd9c..17e4e1011b 100644
--- a/src/tirx/transform/lower_tirx_cleanup.cc
+++ b/src/tirx/transform/lower_tirx_cleanup.cc
@@ -251,12 +251,12 @@ class LayoutApplier : public IRMutatorWithAnalyzer {
return flattened;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
// Preserve the logical buffer until VisitBufferAccess linearizes its indices.
auto value = Mutate(op->value, inplace_mode);
auto indices =
Mutate(op->indices, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
- BufferStore store = ffi::GetRef<BufferStore>(op);
+ TensorStore store = ffi::GetRef<TensorStore>(op);
if (!value.UnchangedOrSameAs(op->value) || !indices.UnchangedOrSameAs(op->indices)) {
auto* n = store.CopyOnWrite();
n->value = std::move(value).ValueOrUnchanged(op->value);
@@ -312,7 +312,7 @@ class LayoutApplier : public IRMutatorWithAnalyzer {
return res;
}
if (auto tile = buffer->layout.value().as<TileLayoutNode>(); tile && tile->HasThreadAxis()) {
- LOG(FATAL) << "Cannot lower direct TensorLoad/BufferStore on a buffer with thread-axis "
+ LOG(FATAL) << "Cannot lower direct TensorLoad/TensorStore on a buffer with thread-axis "
<< "layout: unable to verify that the coordinate matches the current thread. "
<< "Use .view() + .local() to decompose thread and memory axes.";
}
diff --git a/src/tirx/transform/lower_tvm_builtin.cc b/src/tirx/transform/lower_tvm_builtin.cc
index 92ed136afc..8a5c89d7de 100644
--- a/src/tirx/transform/lower_tvm_builtin.cc
+++ b/src/tirx/transform/lower_tvm_builtin.cc
@@ -492,7 +492,7 @@ class BuiltinLower : public StmtExprMutator {
// no need to perform any store for a scalar shape
for (size_t i = 0; i < op->args.size(); ++i) {
prep_seq.emplace_back(
- BufferStore(scope.stack_shape.value(),
+ TensorStore(scope.stack_shape.value(),
prim::cast(PrimType::Int(64), op->args[i].as_or_throw<PrimExpr>()),
{ConstInt32(stack_begin + i)}));
}
diff --git a/src/tirx/transform/lower_warp_memory.cc b/src/tirx/transform/lower_warp_memory.cc
index d6905b7bdd..ec397077cc 100644
--- a/src/tirx/transform/lower_warp_memory.cc
+++ b/src/tirx/transform/lower_warp_memory.cc
@@ -226,7 +226,7 @@ class WarpStoreCoeffFinder : public StmtExprVisitor {
return StmtExprVisitor::Visit_(op);
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
if (op->buffer.get() != buffer_) {
return StmtExprVisitor::Visit_(op);
}
@@ -469,12 +469,12 @@ class WarpAccessRewriter : public StmtExprMutator {
return StmtExprMutator::Mutate_(op, inplace_mode);
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) override {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) override {
// The source is a memory access, not a direct address use checked by the Var hook.
auto value = Mutate(op->value, inplace_mode);
auto indices =
Mutate(op->indices, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
- BufferStore store = ffi::GetRef<BufferStore>(op);
+ TensorStore store = ffi::GetRef<TensorStore>(op);
if (!value.UnchangedOrSameAs(op->value) || !indices.UnchangedOrSameAs(op->indices)) {
auto* n = store.CopyOnWrite();
n->value = std::move(value).ValueOrUnchanged(op->value);
diff --git a/src/tirx/transform/narrow_datatype.cc b/src/tirx/transform/narrow_datatype.cc
index 98c4c6f457..abdbf9f88d 100644
--- a/src/tirx/transform/narrow_datatype.cc
+++ b/src/tirx/transform/narrow_datatype.cc
@@ -38,7 +38,7 @@ namespace tvm {
namespace tirx {
using namespace tvm::prim;
-// This pass narrows indexing expressions (like BufferStoreNode::indices)
+// This pass narrows indexing expressions (like TensorStoreNode::indices)
// that trivially fit into i32/i16 (denoted by `target_bits_`) to
// i32/i16. Considering that i32/i16 indices may be more
// efficient on some backends (while i64 may be more efficient
diff --git a/src/tirx/transform/remove_no_op.cc b/src/tirx/transform/remove_no_op.cc
index 902c41292a..5fb43c9842 100644
--- a/src/tirx/transform/remove_no_op.cc
+++ b/src/tirx/transform/remove_no_op.cc
@@ -151,11 +151,11 @@ class NoOpRemover : public IRMutatorWithAnalyzer {
return value.as<CallNode>() != nullptr;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
- BufferStore store = ffi::GetRef<BufferStore>(op);
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
+ TensorStore store = ffi::GetRef<TensorStore>(op);
// Helper function that returns a statement containing only the
- // side effects of evaluating this BufferStore, but not the store
+ // side effects of evaluating this TensorStore, but not the store
// itself.
auto only_side_effects = [&]() {
ffi::Array<Stmt> statements;
diff --git a/src/tirx/transform/stmt_simplify.cc b/src/tirx/transform/stmt_simplify.cc
index 1aa9547afd..b613c3bcc2 100644
--- a/src/tirx/transform/stmt_simplify.cc
+++ b/src/tirx/transform/stmt_simplify.cc
@@ -210,10 +210,10 @@ UnchangedOr<Stmt> StmtSimplifier::Mutate_(const IfThenElseNode* op, InplaceMode
}
}
-UnchangedOr<Stmt> StmtSimplifier::Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) {
- BufferStore store = Parent::Mutate_(op, inplace_mode)
+UnchangedOr<Stmt> StmtSimplifier::Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) {
+ TensorStore store = Parent::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<BufferStore>();
+ .as_or_throw<TensorStore>();
if (const TensorLoadNode* load = store->value.as<TensorLoadNode>()) {
TensorVar buffer = load->source.as_or_throw<tvm::tirx::TensorVar>();
if (buffer.same_as(store->buffer) && ArrayDeepEqual(load->indices, store->indices) &&
diff --git a/src/tirx/transform/stmt_simplify.h b/src/tirx/transform/stmt_simplify.h
index 37694283e1..e5a3366cea 100644
--- a/src/tirx/transform/stmt_simplify.h
+++ b/src/tirx/transform/stmt_simplify.h
@@ -75,7 +75,7 @@ class StmtSimplifier : public IRMutatorWithAnalyzer {
UnchangedOr<Stmt> Mutate_(const IfThenElseNode* op, InplaceMode inplace_mode) override;
// eliminate useless stores
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) override;
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) override;
private:
bool ArrayDeepEqual(const ffi::Array<PrimExpr>& lhs, const ffi::Array<PrimExpr>& rhs);
diff --git a/src/tirx/transform/storage_rewrite.cc b/src/tirx/transform/storage_rewrite.cc
index 42cf133f4f..4e7f38b3ea 100644
--- a/src/tirx/transform/storage_rewrite.cc
+++ b/src/tirx/transform/storage_rewrite.cc
@@ -149,7 +149,7 @@ class LinearAccessPatternFinder final : public StmtExprVisitor {
return std::nullopt;
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
scope_.push_back(StmtEntry());
// visit subexpr
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->value));
@@ -352,8 +352,8 @@ class InplaceOpVerifier : public StmtExprVisitor {
Visit_(static_cast<const IfThenElseNode*>(stmt));
} else if (stmt->IsInstance<WhileNode>()) {
Visit_(static_cast<const WhileNode*>(stmt));
- } else if (stmt->IsInstance<BufferStoreNode>()) {
- Visit_(static_cast<const BufferStoreNode*>(stmt));
+ } else if (stmt->IsInstance<TensorStoreNode>()) {
+ Visit_(static_cast<const TensorStoreNode*>(stmt));
} else {
return false;
}
@@ -377,7 +377,7 @@ class InplaceOpVerifier : public StmtExprVisitor {
return std::nullopt;
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
++mem_nest_;
for (const auto& index : op->indices) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(index));
@@ -456,7 +456,7 @@ class InplaceOpVerifier : public StmtExprVisitor {
// it is not safe to inplace when there is nested load like A[B[i]]
int mem_nest_{0};
// The current store to be inspected
- const BufferStoreNode* store_{nullptr};
+ const TensorStoreNode* store_{nullptr};
};
/* \brief Rewrite and merge memory allocation.
@@ -554,15 +554,15 @@ class StoragePlanRewriter : public StmtExprMutator {
return remapped;
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
// Access remapping needs the original buffer to preserve its view and element offset.
auto value = Mutate(op->value, inplace_mode);
auto indices =
Mutate(op->indices, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
- BufferStore node =
+ TensorStore node =
value.UnchangedOrSameAs(op->value) && indices.UnchangedOrSameAs(op->indices)
- ? ffi::GetRef<BufferStore>(op)
- : BufferStore(op->buffer, std::move(value).ValueOrUnchanged(op->value),
+ ? ffi::GetRef<TensorStore>(op)
+ : TensorStore(op->buffer, std::move(value).ValueOrUnchanged(op->value),
std::move(indices).ValueOrUnchanged(op->indices), op->span);
return VisitBufferAccess(std::move(node));
}
@@ -623,7 +623,7 @@ class StoragePlanRewriter : public StmtExprMutator {
.as_or_throw<Expr>());
return Call(access->ty, op->op, args, op->attrs, op->ty_args, op->span);
} else {
- BufferStore access(buffer, value.value(), indices, op->span);
+ TensorStore access(buffer, value.value(), indices, op->span);
access = VisitBufferAccess(std::move(access));
ffi::Array<Expr> args{access->buffer.var(), access->value};
for (const PrimExpr& index : access->indices) args.push_back(index);
@@ -1465,7 +1465,7 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
return StmtExprVisitor::Visit_(op);
}
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
OnArrayAccess(op->value.ty(), op->buffer.get(), op->indices, /*is_buffer_load=*/false);
return StmtExprVisitor::Visit_(op);
}
@@ -1809,7 +1809,7 @@ class VectorTypeRewriter : public StmtExprMutator {
}
/*!
- * \brief Mutator for TensorLoad or BufferStore.
+ * \brief Mutator for TensorLoad or TensorStore.
* \return The rewritten node and the shuffle index. (Only for TensorLoad) When the shuffle index
* is non-negative, the caller should generate Shuffle to extract the element from the vector.
*/
@@ -1912,7 +1912,7 @@ class VectorTypeRewriter : public StmtExprMutator {
std::move(indices).ValueUnchecked(), op->span);
auto [modified, shuffle_index] = VisitBufferAccess(node);
- // Not needed for BufferStoreNode, so we can't just call
+ // Not needed for TensorStoreNode, so we can't just call
// LegalizeDtype() in VisitBufferAccess.
if (node.same_as(modified)) {
return node;
@@ -1925,14 +1925,14 @@ class VectorTypeRewriter : public StmtExprMutator {
}
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
auto value = Mutate(op->value, inplace_mode);
auto indices =
Mutate(op->indices, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
- BufferStore node =
+ TensorStore node =
value.UnchangedOrSameAs(op->value) && indices.UnchangedOrSameAs(op->indices)
- ? ffi::GetRef<BufferStore>(op)
- : BufferStore(op->buffer, std::move(value).ValueOrUnchanged(op->value),
+ ? ffi::GetRef<TensorStore>(op)
+ : TensorStore(op->buffer, std::move(value).ValueOrUnchanged(op->value),
std::move(indices).ValueOrUnchanged(op->indices), op->span);
auto [modified, shuffle_index] = VisitBufferAccess(std::move(node));
TVM_FFI_ICHECK(shuffle_index < 0);
@@ -1965,7 +1965,7 @@ class VectorTypeRewriter : public StmtExprMutator {
indices.push_back(this->Mutate(op->args[i].as_or_throw<PrimExpr>())
.ValueOrUnchanged(op->args[i].as_or_throw<PrimExpr>()));
}
- BufferStore access(buffer, value, indices, op->span);
+ TensorStore access(buffer, value, indices, op->span);
auto [modified, shuffle_index] = VisitBufferAccess(std::move(access));
TVM_FFI_ICHECK_LT(shuffle_index, 0)
<< "A masked vector store cannot be rewritten into a scalar shuffle.";
diff --git a/src/tirx/transform/unroll_loop.cc b/src/tirx/transform/unroll_loop.cc
index 8535def0aa..1b74900615 100644
--- a/src/tirx/transform/unroll_loop.cc
+++ b/src/tirx/transform/unroll_loop.cc
@@ -210,7 +210,7 @@ class LoopUnroller : public StmtExprMutator {
return ffi::Unchanged();
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
++step_count_;
if (unroll_local_access_) {
auto storage_scope = runtime::StorageScope::Create(op->buffer.scope());
diff --git a/src/tirx/transform/unsupported_dtype_legalize.cc b/src/tirx/transform/unsupported_dtype_legalize.cc
index 54b4dd1767..003034988f 100644
--- a/src/tirx/transform/unsupported_dtype_legalize.cc
+++ b/src/tirx/transform/unsupported_dtype_legalize.cc
@@ -372,7 +372,7 @@ class ComputeLegalizer : public StmtExprMutator {
}
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
auto value_result = this->Mutate(op->value, inplace_mode);
bool value_unchanged = value_result.UnchangedOrSameAs(op->value);
PrimExpr value = std::move(value_result).ValueOrUnchanged(op->value);
@@ -394,7 +394,7 @@ class ComputeLegalizer : public StmtExprMutator {
TVM_FFI_ICHECK(MatchType(value.ty()));
value = DTypeConversion(value, storage_dtype);
}
- return BufferStore(new_buf, value, indices);
+ return TensorStore(new_buf, value, indices);
}
}
@@ -565,7 +565,7 @@ class StorageLegalizer : public StmtExprMutator {
}
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
PrimExpr value =
this->ChangeToUInt(Mutate(op->value, inplace_mode).ValueOrUnchanged(op->value));
TensorVar new_buf = GetRemappedBuffer(op->buffer);
@@ -578,7 +578,7 @@ class StorageLegalizer : public StmtExprMutator {
if (MatchType(op->value.ty())) {
TVM_FFI_ICHECK(new_buf->dtype.MatchesCode(DLDataTypeCode::kDLUInt));
}
- return BufferStore(new_buf, value, indices);
+ return TensorStore(new_buf, value, indices);
}
}
diff --git a/src/tirx/transform/vectorize_loop.cc b/src/tirx/transform/vectorize_loop.cc
index 62c57bed4d..c4133da2c6 100644
--- a/src/tirx/transform/vectorize_loop.cc
+++ b/src/tirx/transform/vectorize_loop.cc
@@ -223,10 +223,10 @@ class TryPredicateBufferAccesses : public StmtExprMutator {
return TryPredicateBufferAccess(load);
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
auto store = StmtExprMutator::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<BufferStore>();
+ .as_or_throw<TensorStore>();
return TryPredicateBufferAccess(store);
}
@@ -299,7 +299,7 @@ class TryPredicateBufferAccesses : public StmtExprMutator {
return load;
}
- Stmt TryPredicateBufferAccess(BufferStore store) {
+ Stmt TryPredicateBufferAccess(TensorStore store) {
if (auto mask = GetLaneMask(store->indices)) {
ffi::Array<Expr> args{store->buffer.var(), store->value};
for (const PrimExpr& index : store->indices) args.push_back(index);
@@ -350,14 +350,14 @@ class VecAllocAccess : public StmtExprMutator {
return UpdateBufferAccess(load);
}
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
auto value = Mutate(op->value, inplace_mode);
auto indices =
Mutate(op->indices, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
- BufferStore store =
+ TensorStore store =
value.UnchangedOrSameAs(op->value) && indices.UnchangedOrSameAs(op->indices)
- ? ffi::GetRef<BufferStore>(op)
- : BufferStore(op->buffer, std::move(value).ValueOrUnchanged(op->value),
+ ? ffi::GetRef<TensorStore>(op)
+ : TensorStore(op->buffer, std::move(value).ValueOrUnchanged(op->value),
std::move(indices).ValueOrUnchanged(op->indices), op->span);
return UpdateBufferAccess(store);
}
@@ -1003,9 +1003,9 @@ class Vectorizer : public StmtExprMutator {
return vectors[0];
}
}
- // BufferStore
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final {
- auto store = ffi::GetRef<BufferStore>(op);
+ // TensorStore
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final {
+ auto store = ffi::GetRef<TensorStore>(op);
auto fmutate = [this](const PrimExpr& index) {
return this->Mutate(index).ValueOrUnchanged(index);
diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc
index 5516ab339e..45f5679740 100644
--- a/tests/cpp/ir_functor_test.cc
+++ b/tests/cpp/ir_functor_test.cc
@@ -701,7 +701,7 @@ TEST(IRF, StructuralMapBufferDefinition) {
tirx::Var y = x.CopyWithSuffix("subst");
PrimVar m("m", PrimType::Int(32));
TensorVar buffer = fmakebuffer();
- Stmt store = BufferStore(buffer, FloatImm(dtype, 0), {IntImm::Int32(0)});
+ Stmt store = TensorStore(buffer, FloatImm(dtype, 0), {IntImm::Int32(0)});
Stmt decl =
SeqStmt({Bind(buffer, Call(buffer.type(), tvm::tirx::builtin::decl_tensor(),
{x, tvm::Tuple(buffer->shape), DataTypeImm(buffer->dtype->dtype),
@@ -723,7 +723,7 @@ TEST(IRF, StructuralMapBufferDefinition) {
TVM_FFI_ICHECK(decl_call->args[0].same_as(y));
TVM_FFI_ICHECK(decl_node->var.as_or_throw<TensorVar>()->shape[0].same_as(m));
TVM_FFI_ICHECK(!decl_node->var.same_as(buffer));
- auto* store_node = seq_node->seq[1].as<BufferStoreNode>();
+ auto* store_node = seq_node->seq[1].as<TensorStoreNode>();
TVM_FFI_ICHECK(store_node != nullptr);
TVM_FFI_ICHECK(store_node->buffer.same_as(decl_node->var));
}
diff --git a/tests/python/codegen/test_target_codegen.py b/tests/python/codegen/test_target_codegen.py
index 6c354e8414..580b4ee432 100644
--- a/tests/python/codegen/test_target_codegen.py
+++ b/tests/python/codegen/test_target_codegen.py
@@ -23,7 +23,7 @@ import tvm.testing
from tvm.script import tirx as T
-def test_buffer_store_predicate_not_supported():
+def test_tensor_store_predicate_not_supported():
target = "c"
@T.prim_func
@@ -55,7 +55,7 @@ def test_buffer_store_predicate_not_supported():
pytest.param({"kind": "vulkan", "from_device": 0}, marks=pytest.mark.gpu),
],
)
-def test_buffer_store_predicate_not_supported_gpu(target):
+def test_tensor_store_predicate_not_supported_gpu(target):
if not tvm.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
diff --git a/tests/python/codegen/test_target_codegen_llvm.py b/tests/python/codegen/test_target_codegen_llvm.py
index c3fa1331df..6620d5afb9 100644
--- a/tests/python/codegen/test_target_codegen_llvm.py
+++ b/tests/python/codegen/test_target_codegen_llvm.py
@@ -1129,7 +1129,7 @@ def test_invalid_volatile_masked_decl_buffer_load():
tvm.compile(Module)
-def test_invalid_volatile_masked_buffer_store():
+def test_invalid_volatile_masked_tensor_store():
@I.ir_module
class Module:
@T.prim_func
diff --git a/tests/python/codegen/test_target_codegen_vulkan.py b/tests/python/codegen/test_target_codegen_vulkan.py
index b77009b9e6..6ccc1c6307 100644
--- a/tests/python/codegen/test_target_codegen_vulkan.py
+++ b/tests/python/codegen/test_target_codegen_vulkan.py
@@ -247,7 +247,7 @@ def test_vulkan_constant_passing(vulkan_parameter_impl, vulkan_parameter_dtype):
v_i = i_0 * 64 + i_1
with T_builder.if_(v_i < n_var):
with T_builder.then_():
- T_builder.buffer_store(B, scalar_sum + A[v_i], [v_i])
+ T_builder.tensor_store(B, scalar_sum + A[v_i], [v_i])
mod = ib.get()
f_add = tvm.compile(mod, target=target)
diff --git a/tests/python/relax/test_runtime_builtin_rnn_state.py b/tests/python/relax/test_runtime_builtin_rnn_state.py
index 923460831e..c8ad71b252 100644
--- a/tests/python/relax/test_runtime_builtin_rnn_state.py
+++ b/tests/python/relax/test_runtime_builtin_rnn_state.py
@@ -231,7 +231,7 @@ def rnn_state_get(
# The following line is equivalent to:
# `output[vi, *vs] = storage[seq_id, history_id, *vs]`
# However, unpacking operator in subscript requires Python 3.11 or newer
- T.buffer_store(
+ T.tensor_store(
output, T.TensorLoad(storage, [seq_id, history_id, *vs]), [vi, *vs]
)
# fmt: on
@@ -264,7 +264,7 @@ def rnn_state_set(
# The following line is equivalent to:
# `storage[seq_id, history_id, *vs] = data[vi, *vs]`
# However, unpacking operator in subscript requires Python 3.11 or newer
- T.buffer_store(
+ T.tensor_store(
storage, T.TensorLoad(data, [vi, *vs]), [seq_id, history_id, *vs]
)
diff --git a/tests/python/relax/test_transform_annotate_tir_op_pattern.py b/tests/python/relax/test_transform_annotate_tir_op_pattern.py
index 8473b7bf08..27e6541055 100644
--- a/tests/python/relax/test_transform_annotate_tir_op_pattern.py
+++ b/tests/python/relax/test_transform_annotate_tir_op_pattern.py
@@ -416,11 +416,11 @@ def test_sum_sqsum():
assert new_mod["sum_sqsum"].attrs["op_pattern"] == OpPatternKind.kCommReduce
-def test_no_buffer_stores():
+def test_no_tensor_stores():
@tvm.script.ir_module
class Module:
@Ts.prim_func
- def no_buffer_stores(A: T.Tensor((32, 64), "float32"), vsum: T.Tensor((32,), "float32")):
+ def no_tensor_stores(A: T.Tensor((32, 64), "float32"), vsum: T.Tensor((32,), "float32")):
for ax0, k0 in T.grid(32, 64):
with Ts.sblock("block"):
v_ax0, v_k0 = Ts.axis.remap("SR", [ax0, k0])
@@ -432,7 +432,7 @@ def test_no_buffer_stores():
mod = Module
new_mod = relax.transform.AnnotateTIROpPattern()(mod)
- assert new_mod["no_buffer_stores"].attrs["op_pattern"] == OpPatternKind.kOpaque
+ assert new_mod["no_tensor_stores"].attrs["op_pattern"] == OpPatternKind.kOpaque
if __name__ == "__main__":
diff --git a/tests/python/runtime/test_runtime_module_load.py b/tests/python/runtime/test_runtime_module_load.py
index a84bb66200..e445393879 100644
--- a/tests/python/runtime/test_runtime_module_load.py
+++ b/tests/python/runtime/test_runtime_module_load.py
@@ -59,7 +59,7 @@ def test_dso_module_load():
0,
n - 1,
tvm.tirx.ForKind.SERIAL,
- tvm.tirx.BufferStore(Ab, tvm.tirx.TensorLoad(Ab, [i]) + 1, [i + 1]),
+ tvm.tirx.TensorStore(Ab, tvm.tirx.TensorLoad(Ab, [i]) + 1, [i + 1]),
)
mod = tvm.IRModule.from_expr(
tvm.tirx.PrimFunc([Ab], stmt).with_attr("global_symbol", "main")
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 4b6db504e7..6b2025c96b 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
@@ -69,7 +69,7 @@ def test_1d_compute():
for i in T.serial(1024):
B[i] = A[i] + 1.0
- expected = "Expected BufferStore's value to be TensorLoad"
+ expected = "Expected TensorStore's value to be TensorLoad"
_check_memcpy_results(func, expected)
@@ -82,7 +82,7 @@ def test_1d_conditional():
if i < 1024:
B[i] = A[i]
- expected = "Expected innermost loop to have BufferStore body"
+ expected = "Expected innermost loop to have TensorStore body"
_check_memcpy_results(func, expected)
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
index fe1a8861dc..99645370e3 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
@@ -114,7 +114,7 @@ def test_match_buffer_in_block_is_well_formed():
def test_error_undeclared_buffer_in_schedulable_tir():
"""In schedule-level TIR (with SBlock nodes), all buffers must be declared."""
- # Manually construct a BufferStore that uses a buffer without any declaration
+ # Manually construct a TensorStore that uses a buffer without any declaration
# inside a block context.
n = tvm.tirx.Var("n", "int32")
A = tvm.tirx.decl_tensor([n], "float32", name="A")
@@ -131,7 +131,7 @@ def test_error_undeclared_buffer_in_schedulable_tir():
iter_vars=[tvm.s_tir.IterVar(tvm.ir.Range(0, n), bi, 0)], # 0 = kDataPar
reads=[tvm.tirx.BufferRegion(A, [tvm.ir.Range(bi, bi + 1)])],
writes=[tvm.tirx.BufferRegion(B, [tvm.ir.Range(bi, bi + 1)])],
- body=tvm.tirx.BufferStore(B, tvm.tirx.TensorLoad(A, [bi]), [bi]),
+ body=tvm.tirx.TensorStore(B, tvm.tirx.TensorLoad(A, [bi]), [bi]),
name_hint="write_B",
)
block_realize = tvm.s_tir.SBlockRealize(
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py
index da010baa8e..3d76e5fcf1 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py
@@ -448,7 +448,7 @@ def test_cpu_fusion():
)
feature = feature.numpy()
assert feature.shape == (2, N_FEATURES)
- ## Features for BufferStore(B)
+ ## Features for TensorStore(B)
f = feature[0]
# Group 1.1: arith
assert_allclose(
@@ -580,7 +580,7 @@ def test_cpu_fusion():
rtol=1e-5,
atol=1e-5,
)
- ## Features for BufferStore(C)
+ ## Features for TensorStore(C)
f = feature[1]
# Group 1.1: arith
assert_allclose(
@@ -795,7 +795,7 @@ def test_gpu():
)
feature = feature.numpy()
assert feature.shape == (4, N_FEATURES)
- ### Check feature[0]: BufferStore(A_shared) <= A[...]
+ ### Check feature[0]: TensorStore(A_shared) <= A[...]
f = feature[0]
# Group 1.1: arith
assert_allclose(
@@ -1001,7 +1001,7 @@ def test_gpu():
rtol=1e-5,
atol=1e-5,
)
- ### Check feature[1]: BufferStore(B_shared) <= B[...]
+ ### Check feature[1]: TensorStore(B_shared) <= B[...]
f = feature[1]
# Group 1.1: arith
assert_allclose(
@@ -1207,7 +1207,7 @@ def test_gpu():
rtol=1e-5,
atol=1e-5,
)
- ### Check feature[2]: BufferStore(C_local) <= C_local[...] + A_shared[...] * B_shared[...]
+ ### Check feature[2]: TensorStore(C_local) <= C_local[...] + A_shared[...] * B_shared[...]
f = feature[2]
# Group 1.1: arith
assert_allclose(
@@ -1427,7 +1427,7 @@ def test_gpu():
rtol=1e-5,
atol=1e-5,
)
- ### Check feature[3]: BufferStore(C) <= C_local[...]
+ ### Check feature[3]: TensorStore(C) <= C_local[...]
f = feature[3]
# Group 1.1: arith
assert_allclose(
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_rfactor.py b/tests/python/s_tir/schedule/test_tir_schedule_rfactor.py
index f3a1d54be7..51de5783ef 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_rfactor.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_rfactor.py
@@ -422,7 +422,7 @@ def rowsum_wrong_reduce_pattern2(A: T.Tensor((128, 128)), B: T.Tensor((128,))) -
@Ts.prim_func
-def rowsum_init_not_bufferstore(A: T.Tensor((128, 128)), B: T.Tensor((128,))) -> None:
+def rowsum_init_not_tensorstore(A: T.Tensor((128, 128)), B: T.Tensor((128,))) -> None:
for i, k in T.grid(128, 128):
with Ts.sblock("B"):
vi, vk = Ts.axis.remap("SR", [i, k])
@@ -744,7 +744,7 @@ def argmax_split_different_indices(
@Ts.prim_func
-def argmax_split_init_not_bufferstore(
+def argmax_split_init_not_tensorstore(
idx: T.Tensor((128, 128), "int32"),
val: T.Tensor((128, 128), "float32"),
argmax_v0: T.Tensor((128,), "int32"),
@@ -845,7 +845,7 @@ def argmax_split_bind_more_than_init(
@Ts.prim_func
-def argmax_split_let_body_neither_seqstmt_nor_bufferstore(
+def argmax_split_let_body_neither_seqstmt_nor_tensorstore(
idx: T.Tensor((128, 128), "int32"),
val: T.Tensor((128, 128), "float32"),
argmax_v0: T.Tensor((128,), "int32"),
@@ -870,7 +870,7 @@ def argmax_split_let_body_neither_seqstmt_nor_bufferstore(
@Ts.prim_func
-def argmax_split_init_update_inconsistent_bufferstore_number(
+def argmax_split_init_update_inconsistent_tensorstore_number(
idx: T.Tensor((128, 128), "int32"),
val: T.Tensor((128, 128), "float32"),
argmax_v0: T.Tensor((128,), "int32"),
@@ -897,7 +897,7 @@ def argmax_split_init_update_inconsistent_bufferstore_number(
@Ts.prim_func
-def argmax_split_body_seq_not_bufferstore(
+def argmax_split_body_seq_not_tensorstore(
idx: T.Tensor((128, 128), "int32"),
val: T.Tensor((128, 128), "float32"),
argmax_v0: T.Tensor((128,), "int32"),
@@ -923,7 +923,7 @@ def argmax_split_body_seq_not_bufferstore(
@Ts.prim_func
-def argmax_split_body_bufferstore_value_not_var(
+def argmax_split_body_tensorstore_value_not_var(
idx: T.Tensor((128, 128), "int32"),
val: T.Tensor((128, 128), "float32"),
argmax_v0: T.Tensor((128,), "int32"),
@@ -953,7 +953,7 @@ v_unbound = T.dynamic("v_unbound", "int32")
@Ts.prim_func(check_well_formed=False)
-def argmax_split_body_bufferstore_value_unbound_var(
+def argmax_split_body_tensorstore_value_unbound_var(
idx: T.Tensor((128, 128), "int32"),
val: T.Tensor((128, 128), "float32"),
argmax_v0: T.Tensor((128,), "int32"),
@@ -1549,8 +1549,8 @@ def test_reduction_rfactor_wrong_reduce_pattern2():
s.rfactor(k, 0)
-def test_reduction_rfactor_init_not_bufferstore():
- s = tvm.s_tir.Schedule(rowsum_init_not_bufferstore, debug_mask="all")
+def test_reduction_rfactor_init_not_tensorstore():
+ s = tvm.s_tir.Schedule(rowsum_init_not_tensorstore, debug_mask="all")
_, k = s.get_loops(s.get_sblock("B"))
with pytest.raises(tvm.s_tir.ScheduleError):
s.rfactor(k, 0)
@@ -1680,8 +1680,8 @@ def test_reduction_rfactor_argmax_different_access_indices():
s.rfactor(ki, 1)
-def test_reduction_rfactor_argmax_init_not_bufferstore():
- s = tvm.s_tir.Schedule(argmax_split_init_not_bufferstore, debug_mask="all")
+def test_reduction_rfactor_argmax_init_not_tensorstore():
+ s = tvm.s_tir.Schedule(argmax_split_init_not_tensorstore, debug_mask="all")
argmax = s.get_sblock("argmax")
_, _, ki = s.get_loops(argmax)
with pytest.raises(tvm.s_tir.ScheduleError):
@@ -1712,17 +1712,17 @@ def test_reduction_rfactor_argmax_bind_more_than_init():
s.rfactor(ki, 1)
-def test_reduction_rfactor_argmax_let_body_neither_seqstmt_nor_bufferstore():
- s = tvm.s_tir.Schedule(argmax_split_let_body_neither_seqstmt_nor_bufferstore, debug_mask="all")
+def test_reduction_rfactor_argmax_let_body_neither_seqstmt_nor_tensorstore():
+ s = tvm.s_tir.Schedule(argmax_split_let_body_neither_seqstmt_nor_tensorstore, debug_mask="all")
argmax = s.get_sblock("argmax")
_, _, ki = s.get_loops(argmax)
with pytest.raises(tvm.s_tir.ScheduleError):
s.rfactor(ki, 1)
-def test_reduction_rfactor_argmax_init_update_inconsistent_bufferstore_number():
+def test_reduction_rfactor_argmax_init_update_inconsistent_tensorstore_number():
s = tvm.s_tir.Schedule(
- argmax_split_init_update_inconsistent_bufferstore_number, debug_mask="all"
+ argmax_split_init_update_inconsistent_tensorstore_number, debug_mask="all"
)
argmax = s.get_sblock("argmax")
_, _, ki = s.get_loops(argmax)
@@ -1730,16 +1730,16 @@ def test_reduction_rfactor_argmax_init_update_inconsistent_bufferstore_number():
s.rfactor(ki, 1)
-def test_reduction_rfactor_argmax_body_seq_not_bufferstore():
- s = tvm.s_tir.Schedule(argmax_split_body_seq_not_bufferstore, debug_mask="all")
+def test_reduction_rfactor_argmax_body_seq_not_tensorstore():
+ s = tvm.s_tir.Schedule(argmax_split_body_seq_not_tensorstore, debug_mask="all")
argmax = s.get_sblock("argmax")
_, _, ki = s.get_loops(argmax)
with pytest.raises(tvm.s_tir.ScheduleError):
s.rfactor(ki, 1)
-def test_reduction_rfactor_argmax_body_bufferstore_value_not_var():
- s = tvm.s_tir.Schedule(argmax_split_body_bufferstore_value_not_var, debug_mask="all")
+def test_reduction_rfactor_argmax_body_tensorstore_value_not_var():
+ s = tvm.s_tir.Schedule(argmax_split_body_tensorstore_value_not_var, debug_mask="all")
argmax = s.get_sblock("argmax")
_, _, ki = s.get_loops(argmax)
with pytest.raises(tvm.s_tir.ScheduleError):
@@ -1747,8 +1747,8 @@ def test_reduction_rfactor_argmax_body_bufferstore_value_not_var():
@pytest.mark.xfail(reason="The input IR is not well-formed")
-def test_reduction_rfactor_argmax_body_bufferstore_value_unbound_var():
- s = tvm.s_tir.Schedule(argmax_split_body_bufferstore_value_unbound_var, debug_mask="all")
+def test_reduction_rfactor_argmax_body_tensorstore_value_unbound_var():
+ s = tvm.s_tir.Schedule(argmax_split_body_tensorstore_value_unbound_var, debug_mask="all")
argmax = s.get_sblock("argmax")
_, _, ki = s.get_loops(argmax)
with pytest.raises(tvm.s_tir.ScheduleError):
diff --git a/tests/python/s_tir/script/test_s_tir_script_basic_usage.py b/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
index bcce69e259..1f5cb62942 100644
--- a/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
+++ b/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
@@ -889,7 +889,7 @@ def test_implicit_evaluate_call_extern():
def test_preserve_trivial_let_binding():
"""Trivial `T.let[...]` annotations survive the parser as LetStmt and are not inlined.
- In fork, bare `j = i` lowers to a local_scalar (AllocTensor + BufferStore); the
+ In fork, bare `j = i` lowers to a local_scalar (AllocTensor + TensorStore); the
LetStmt form is opt-in via `T.let[T.dtype]`. Both the explicit `T.bind(..., var=j)`
builder API and the `j: T.let[T.dtype]` annotation produce the same LetStmt IR.
"""
diff --git a/tests/python/s_tir/script/test_s_tir_script_blocks.py b/tests/python/s_tir/script/test_s_tir_script_blocks.py
index b700c7932a..e2e43145e9 100644
--- a/tests/python/s_tir/script/test_s_tir_script_blocks.py
+++ b/tests/python/s_tir/script/test_s_tir_script_blocks.py
@@ -721,8 +721,8 @@ def test_block_elements():
assert isinstance(rt_func.body.block.body, s_tir.SBlockRealize)
assert isinstance(rt_func.body.block.body.block, s_tir.SBlock)
block = rt_func.body.block.body.block
- assert isinstance(block.body, tirx.stmt.BufferStore)
- assert isinstance(block.init, tirx.stmt.BufferStore)
+ assert isinstance(block.body, tirx.stmt.TensorStore)
+ assert isinstance(block.init, tirx.stmt.TensorStore)
assert len(block.annotations) == 1
assert block.annotations["attr_key"] == "attr_value"
diff --git a/tests/python/s_tir/script/test_s_tir_script_error_handling.py b/tests/python/s_tir/script/test_s_tir_script_error_handling.py
index 7762703f3b..d23b4eb13b 100644
--- a/tests/python/s_tir/script/test_s_tir_script_error_handling.py
+++ b/tests/python/s_tir/script/test_s_tir_script_error_handling.py
@@ -239,7 +239,7 @@ def test_buffer_rebinding_preserves_distinct_allocations():
first, second = block.alloc_buffers
assert not first.same_as(second)
store = block.body
- assert isinstance(store, tirx.BufferStore)
+ assert isinstance(store, tirx.TensorStore)
assert store.buffer.same_as(second)
assert store.value.a.source.same_as(second)
diff --git a/tests/python/s_tir/script/test_s_tir_script_printer.py b/tests/python/s_tir/script/test_s_tir_script_printer.py
index 102232b06c..abd769b14f 100644
--- a/tests/python/s_tir/script/test_s_tir_script_printer.py
+++ b/tests/python/s_tir/script/test_s_tir_script_printer.py
@@ -2092,7 +2092,7 @@ def test_disable_concise_scoping_when_scope_annotated():
y: T.int32 = x + 1
T.evaluate(y - 1)
- # Explicit scalar declarations lower to AllocTensor + BufferStore (local_scalar).
+ # Explicit scalar declarations lower to AllocTensor + TensorStore (local_scalar).
# The printer fuses each pair into one line; annotate the allocation for y.
result = _func.with_attr("global_symbol", "main").script(
obj_to_annotate={
diff --git a/tests/python/s_tir/test_s_tir_renew_defs.py b/tests/python/s_tir/test_s_tir_renew_defs.py
index 9073f74e4e..d70daa64ef 100644
--- a/tests/python/s_tir/test_s_tir_renew_defs.py
+++ b/tests/python/s_tir/test_s_tir_renew_defs.py
@@ -151,11 +151,11 @@ def test_undefined_buffer():
# AllocTensor is now a flat statement in SeqStmt
assert f1.body.seq[0].var.data != f2.body.seq[0].var.data
- def _get_buffer_store_buffer(f):
- # SeqStmt: [AllocTensor, Evaluate, For]; For body has the BufferStore
+ def _get_tensor_store_buffer(f):
+ # SeqStmt: [AllocTensor, Evaluate, For]; For body has the TensorStore
return f.body.seq[2].body.buffer
- _check_buffer_decl(_get_buffer_store_buffer(f1), _get_buffer_store_buffer(f2))
+ _check_buffer_decl(_get_tensor_store_buffer(f1), _get_tensor_store_buffer(f2))
def test_symbolic_func():
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py
index 43256bb119..45273136db 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py
@@ -950,7 +950,7 @@ def lowered_zero_rank_buffer(
@Ts.prim_func
-def multiple_bufferstore(
+def multiple_tensorstore(
A: T.Tensor([128, 128], dtype="float32"), B: T.Tensor([128], dtype="float32")
) -> None:
C = Ts.sblock_alloc_buffer([], dtype="float32")
@@ -1803,8 +1803,8 @@ def test_zero_rank_buffer():
_check(zero_rank_buffer, lowered_zero_rank_buffer)
-def test_multiple_bufferstore():
- _check_fail(multiple_bufferstore)
+def test_multiple_tensorstore():
+ _check_fail(multiple_tensorstore)
def test_reduction_block_not_deepest():
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_mma_buffer_layout.py b/tests/python/s_tir/transform/test_s_tir_transform_mma_buffer_layout.py
index 8095f0de91..02906065d0 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_mma_buffer_layout.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_mma_buffer_layout.py
@@ -31,7 +31,7 @@ def test_explicit_matrix_ab_access_is_rejected(scope, shape, access_kind):
if access_kind == "load":
body = tirx.Evaluate(tirx.TensorLoad(buffer, [0, 0]))
else:
- body = tirx.BufferStore(buffer, 0.0, [0, 0])
+ body = tirx.TensorStore(buffer, 0.0, [0, 0])
block = s_tir.SBlock([], [], [], "root", body, alloc_buffers=[buffer])
func = tirx.PrimFunc([], s_tir.SBlockRealize([], True, block))
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_remove_undef.py b/tests/python/s_tir/transform/test_s_tir_transform_remove_undef.py
index 8162e0835f..6f826dffcb 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_remove_undef.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_remove_undef.py
@@ -114,7 +114,7 @@ def test_raise_error_for_undef_as_store_indices():
def test_raise_error_for_undef_as_load_indices():
"""Use of T.undef() as buffer indices is an error
- Even though this occurs as part of the BufferStore's value, the
+ Even though this occurs as part of the TensorStore's value, the
T.undef() may not appear in a buffer's indices.
"""
diff --git a/tests/python/tirx-base/test_tir_constructor.py b/tests/python/tirx-base/test_tir_constructor.py
index db1c3777b1..adf6d5bddd 100644
--- a/tests/python/tirx-base/test_tir_constructor.py
+++ b/tests/python/tirx-base/test_tir_constructor.py
@@ -310,8 +310,8 @@ def test_stmt_constructor():
buffer_var = tvm.tirx.Var("buf", tvm.ir.PointerType(tvm.ir.PrimType("bool")))
buffer = tvm.tirx.decl_tensor([16], "bool", data=buffer_var)
- x = tvm.tirx.BufferStore(buffer, tvm.tirx.IntImm("bool", 1), [10])
- assert isinstance(x, tvm.tirx.BufferStore)
+ x = tvm.tirx.TensorStore(buffer, tvm.tirx.IntImm("bool", 1), [10])
+ assert isinstance(x, tvm.tirx.TensorStore)
assert x.buffer == buffer
assert x.buffer.data.args[0].same_as(buffer)
assert x.buffer.data.ty == tvm.tirx.buffer_data_pointer_type(buffer)
diff --git a/tests/python/tirx-base/test_tir_nodes.py b/tests/python/tirx-base/test_tir_nodes.py
index 2f247b37d2..f83058f9ee 100644
--- a/tests/python/tirx-base/test_tir_nodes.py
+++ b/tests/python/tirx-base/test_tir_nodes.py
@@ -82,8 +82,8 @@ def test_ir2():
array = tvm.tirx.Var("array", handle_type)
buf = tvm.tirx.decl_tensor([buf_size], "int32", data=array)
- st = tvm.tirx.BufferStore(buf, x + 1, [1])
- assert isinstance(st, tvm.tirx.BufferStore)
+ st = tvm.tirx.TensorStore(buf, x + 1, [1])
+ assert isinstance(st, tvm.tirx.TensorStore)
assert st.buffer == buf
assert st.buffer.data.args[0].same_as(buf)
assert st.buffer.data.ty == array.ty
@@ -345,8 +345,8 @@ def test_buffer_load_store():
assert not hasattr(x, "buffer")
with pytest.raises(TypeError, match="cannot be constructed directly"):
tvm.ir.TensorLoad(b, [0])
- s = tvm.tirx.BufferStore(b, 0.1, [0])
- assert isinstance(s, tvm.tirx.BufferStore)
+ s = tvm.tirx.TensorStore(b, 0.1, [0])
+ assert isinstance(s, tvm.tirx.TensorStore)
def test_intimm_cond():
@@ -418,13 +418,13 @@ def test_buffer_load_scalable_vec():
assert load.ty.dtype == "float32xvscalex8"
-def test_buffer_store_scalable_vec():
+def test_tensor_store_scalable_vec():
b = tvm.tirx.decl_tensor((24,), "int32")
value = tvm.tirx.expr.Broadcast(1, 4 * tvm.tirx.vscale())
index = tvm.tirx.expr.Ramp(0, 1, 4 * tvm.tirx.vscale())
- store = tvm.tirx.BufferStore(b, value, [index])
+ store = tvm.tirx.TensorStore(b, value, [index])
- assert isinstance(store, tvm.tirx.BufferStore)
+ assert isinstance(store, tvm.tirx.TensorStore)
assert store.value.ty.dtype == "int32xvscalex4"
@@ -433,7 +433,7 @@ def test_scalable_vec_cast():
value = tvm.tirx.expr.Broadcast(1, 12 * tvm.tirx.vscale()).astype("float32xvscalex12")
index = tvm.tirx.expr.Ramp(0, 1, 12 * tvm.tirx.vscale())
- store = tvm.tirx.BufferStore(b, value, [index])
+ store = tvm.tirx.TensorStore(b, value, [index])
assert isinstance(store.value.value, tvm.tirx.expr.FloatImm)
diff --git a/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py b/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py
index 8cab7aeca8..f8f7e8b888 100644
--- a/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py
+++ b/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py
@@ -591,7 +591,7 @@ def test_let_body_no_extraction():
let_expr = tvm.tirx.Let(x, tvm.tirx.IntImm("int32", 1), (x + y) + (x + y))
buf = tvm.tirx.decl_tensor((10,), "int32", name="B")
i = tvm.tirx.Var("i", "int32")
- store = tvm.tirx.BufferStore(buf, let_expr, [i])
+ store = tvm.tirx.TensorStore(buf, let_expr, [i])
loop = tvm.tirx.For(
i,
tvm.tirx.const(0, "int32"),
@@ -621,7 +621,7 @@ def test_let_value_cse():
let_expr = tvm.tirx.Let(x, y + z, x + 1)
buf = tvm.tirx.decl_tensor((10,), "int32", name="B")
i = tvm.tirx.Var("i", "int32")
- store = tvm.tirx.BufferStore(buf, (y + z) + let_expr, [i])
+ store = tvm.tirx.TensorStore(buf, (y + z) + let_expr, [i])
loop = tvm.tirx.For(
i,
tvm.tirx.const(0, "int32"),
@@ -654,7 +654,7 @@ def test_nested_let_no_extraction():
)
buf = tvm.tirx.decl_tensor((10,), "int32", name="B")
i = tvm.tirx.Var("i", "int32")
- store = tvm.tirx.BufferStore(buf, nested_let, [i])
+ store = tvm.tirx.TensorStore(buf, nested_let, [i])
loop = tvm.tirx.For(
i,
tvm.tirx.const(0, "int32"),
@@ -697,7 +697,7 @@ def test_let_floordiv_pattern():
tvm.tirx.TensorLoad(buf_a, [i]),
tvm.tirx.Let(y, tvm.tirx.TensorLoad(buf_b, [i]), outer_let),
)
- store = tvm.tirx.BufferStore(buf_c, full_expr, [i])
+ store = tvm.tirx.TensorStore(buf_c, full_expr, [i])
loop = tvm.tirx.For(
i,
tvm.tirx.const(0, "int32"),
diff --git a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
index 9d8c700587..cf84dcef45 100644
--- a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
+++ b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
@@ -72,7 +72,7 @@ def check_value(expr, variables, data, fref):
result = expr
for j in range(num_vars - 1, -1, -1):
result = tvm.tirx.Let(variables[j], tvm.tirx.TensorLoad(input_bufs[j], [i_var]), result)
- return tvm.tirx.BufferStore(out_buf, result, [i_var])
+ return tvm.tirx.TensorStore(out_buf, result, [i_var])
loop = tvm.tirx.For(
loop_var,
diff --git a/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py b/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
index 82b7938d91..1d79272a25 100644
--- a/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
+++ b/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
@@ -162,13 +162,13 @@ def test_call_packed_return_non_i32():
# Build statements using direct TIR construction (no ir_builder)
# 1. Store packed_echo(const) result into Ab[0]
- store0 = tvm.tirx.BufferStore(
+ store0 = tvm.tirx.TensorStore(
Ab, packed_echo(tvm.tirx.const(expected_value[0], "float32")), [0]
)
# 2. Let binding: Aptr_dup = packed_echo(Ab.data), then store const into Ab[1]
Aptr_dup = tvm.tirx.Var("Aptr_dup", Ab.data.ty)
- store1 = tvm.tirx.BufferStore(Ab, tvm.tirx.const(expected_value[1], "float32"), [1])
+ store1 = tvm.tirx.TensorStore(Ab, tvm.tirx.const(expected_value[1], "float32"), [1])
bind_stmt = tvm.tirx.Bind(Aptr_dup, packed_echo(Ab.data))
# Combine into sequence
diff --git a/tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py b/tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py
index 38b25c9aea..b3d212dcd1 100644
--- a/tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py
+++ b/tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py
@@ -171,21 +171,21 @@ def test_decl_buffer_alias_chain_uses_flat_root_map():
assert func.params[0].ty.dtype == tvm.ir.PrimType("float32x4")
decl_buffers = []
- buffer_stores = []
+ tensor_stores = []
tvm_ffi.structural_walk(
func.body,
lambda node: (
decl_buffers.append(node)
if _is_buffer_binding(node, "tirx.decl_tensor")
- else buffer_stores.append(node)
- if isinstance(node, tvm.tirx.BufferStore)
+ else tensor_stores.append(node)
+ if isinstance(node, tvm.tirx.TensorStore)
else None
),
)
assert len(decl_buffers) == 2
assert all(decl.var.ty.dtype == tvm.ir.PrimType("float32x4") for decl in decl_buffers)
- assert len(buffer_stores) == 1
- assert buffer_stores[0].buffer.ty.dtype == tvm.ir.PrimType("float32x4")
+ assert len(tensor_stores) == 1
+ assert tensor_stores[0].buffer.ty.dtype == tvm.ir.PrimType("float32x4")
if __name__ == "__main__":
diff --git a/tests/python/tirx-transform/test_tir_transform_remove_no_op.py b/tests/python/tirx-transform/test_tir_transform_remove_no_op.py
index 5006c9573f..5b224887ad 100644
--- a/tests/python/tirx-transform/test_tir_transform_remove_no_op.py
+++ b/tests/python/tirx-transform/test_tir_transform_remove_no_op.py
@@ -60,7 +60,7 @@ def test_remove_no_op():
ret = tvm.tirx.transform.RemoveNoOp()(mod)["main"].body
assert isinstance(ret, tvm.tirx.Evaluate)
- store = tvm.tirx.BufferStore(Ab, tvm.tirx.TensorLoad(Ab, [i]) + 1, [i + 1])
+ store = tvm.tirx.TensorStore(Ab, tvm.tirx.TensorLoad(Ab, [i]) + 1, [i + 1])
stmt2 = tvm.tirx.SeqStmt([nop(), tvm.tirx.SeqStmt([store, nop()])])
mod = tvm.IRModule.from_expr(tvm.tirx.PrimFunc([Ab], stmt2))
diff --git a/tests/python/tirx-transform/test_tir_transform_simplify.py b/tests/python/tirx-transform/test_tir_transform_simplify.py
index 9dae622a53..0eadb918aa 100644
--- a/tests/python/tirx-transform/test_tir_transform_simplify.py
+++ b/tests/python/tirx-transform/test_tir_transform_simplify.py
@@ -42,7 +42,7 @@ def test_stmt_simplify():
# Find the For loop in the sequence
for_stmt = [s for s in stmts if isinstance(s, tvm.tirx.For)]
assert len(for_stmt) == 1, f"Expected one For loop, got {len(for_stmt)}"
- assert isinstance(for_stmt[0].body, tvm.tirx.BufferStore)
+ assert isinstance(for_stmt[0].body, tvm.tirx.TensorStore)
def test_thread_extent_simplify():
@@ -67,7 +67,7 @@ def test_thread_extent_simplify():
tx_loop = for_stmts[0]
assert isinstance(tx_loop, tvm.tirx.For) # tx loop
assert isinstance(tx_loop.body, tvm.tirx.For) # ty loop
- assert isinstance(tx_loop.body.body, tvm.tirx.BufferStore) # The if was eliminated
+ assert isinstance(tx_loop.body.body, tvm.tirx.TensorStore) # The if was eliminated
def test_if_likely():
@@ -1232,7 +1232,7 @@ def test_simplify_buffer_identity_well_formed():
The simplifier's Dispatch calls analyzer_->Simplify() directly, bypassing
normal ExprMutator dispatch. If VisitBufferDef remaps a buffer at a DeclTensor
- site (e.g. inlining n_val -> n in the shape), TensorLoad inside a BufferStore
+ site (e.g. inlining n_val -> n in the shape), TensorLoad inside a TensorStore
value would NOT pick up the remap because VisitBufferUse is never called.
This causes DeclTensor/TensorLoad buffer identity divergence.
"""
@@ -1318,9 +1318,9 @@ def test_mutable_branch_predicate_preserves_while_bound(else_branch, write_befor
count = tirx.decl_tensor((1,), "int32", name="count")
loop = tirx.While(
T.And(x[0] < 8, count[0] == 0),
- tirx.BufferStore(x, x[0] + 1, [0]),
+ tirx.TensorStore(x, x[0] + 1, [0]),
)
- body = tirx.SeqStmt([tirx.BufferStore(x, x[0] + 1, [0]), loop]) if write_before_loop else loop
+ body = tirx.SeqStmt([tirx.TensorStore(x, x[0] + 1, [0]), loop]) if write_before_loop else loop
branch = (
tirx.IfThenElse(T.int32(8) <= x[0], tirx.Evaluate(0), body)
if else_branch
diff --git a/tests/python/tirx-transform/test_tir_transform_unroll_loop.py b/tests/python/tirx-transform/test_tir_transform_unroll_loop.py
index 835198276c..4cf46f9bd6 100644
--- a/tests/python/tirx-transform/test_tir_transform_unroll_loop.py
+++ b/tests/python/tirx-transform/test_tir_transform_unroll_loop.py
@@ -87,7 +87,7 @@ def test_unroll_fake_loop():
}
):
ret = tvm.tirx.transform.UnrollLoop()(Module)["main"].body
- assert isinstance(ret[0], tvm.tirx.BufferStore)
+ assert isinstance(ret[0], tvm.tirx.TensorStore)
def test_unroll_allocations():
diff --git a/tests/python/tirx/script/test_tirx_script_error_handling.py b/tests/python/tirx/script/test_tirx_script_error_handling.py
index 7750742db8..9a0efc50a6 100644
--- a/tests/python/tirx/script/test_tirx_script_error_handling.py
+++ b/tests/python/tirx/script/test_tirx_script_error_handling.py
@@ -56,15 +56,15 @@ def test_optional_annotation_requires_jit_at_the_source_parameter():
def test_scalar_assign_error_not_swallowed():
- """Regression: genuine errors (non-TypeError) from buffer_store during
+ """Regression: genuine errors (non-TypeError) from tensor_store during
scalar-assignment sugar must propagate, not be silently swallowed.
- Before the fix, both eval_expr and buffer_store were wrapped in a single
- broad ``except Exception: pass``, so any error from buffer_store would be
+ Before the fix, both eval_expr and tensor_store were wrapped in a single
+ broad ``except Exception: pass``, so any error from tensor_store would be
swallowed and the assignment would silently fall through to eval_assign."""
from unittest.mock import patch
- original = tvm.tirx.script.ir_builder.parser_protocol.buffer_store
+ original = tvm.tirx.script.ir_builder.parser_protocol.tensor_store
def bomb(*args, **kwargs):
# Intercept only the scalar-assignment path (indices == [0])
@@ -74,7 +74,7 @@ def test_scalar_assign_error_not_swallowed():
# The ValueError propagates unchanged. A broad ``except Exception`` here
# previously swallowed it and fell through to eval_assign.
- with patch("tvm.tirx.script.ir_builder.parser_protocol.buffer_store", side_effect=bomb):
+ with patch("tvm.tirx.script.ir_builder.parser_protocol.tensor_store", side_effect=bomb):
try:
@T.prim_func
@@ -86,4 +86,4 @@ def test_scalar_assign_error_not_swallowed():
except ValueError as error:
assert str(error) == "boom"
else:
- pytest.fail("buffer_store error was swallowed")
+ pytest.fail("tensor_store error was swallowed")
diff --git a/tests/python/tirx/script/test_tirx_script_ir_builder.py b/tests/python/tirx/script/test_tirx_script_ir_builder.py
index 36d66cd04c..612e98e62e 100644
--- a/tests/python/tirx/script/test_tirx_script_ir_builder.py
+++ b/tests/python/tirx/script/test_tirx_script_ir_builder.py
@@ -191,35 +191,35 @@ def test_ir_builder_tir_if_then_else():
assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)
-def test_ir_builder_tir_buffer_store():
+def test_ir_builder_tir_tensor_store():
buffer_a = T.Var("buffer_a", T.Tensor((10, 10), "float32"))
i = T.int32()
with IRBuilder() as ib:
- T.buffer_store(buffer_a, 0.1, [0, i])
+ T.tensor_store(buffer_a, 0.1, [0, i])
# the buffer store generated by IRBuilder
ir_actual = ib.get()
# the expected buffer store
- ir_expected = tirx.BufferStore(buffer_a, 0.1, [0, i])
+ ir_expected = tirx.TensorStore(buffer_a, 0.1, [0, i])
# Check if the generated ir is expected
assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)
-def test_ir_builder_tir_buffer_store_scalable_vec():
+def test_ir_builder_tir_tensor_store_scalable_vec():
buffer_a = T.Var("buffer_a", T.Tensor((30,), "float32"))
value = T.broadcast(0.11, 4 * tvm.tirx.vscale())
index = T.ramp(0, 1, 4 * tvm.tirx.vscale())
with IRBuilder() as ib:
- T.buffer_store(buffer_a, value, [index])
+ T.tensor_store(buffer_a, value, [index])
# the buffer store generated by IRBuilder
ir_actual = ib.get()
# the expected buffer store
- ir_expected = tirx.BufferStore(buffer_a, value, [index])
+ ir_expected = tirx.TensorStore(buffer_a, value, [index])
# Check if the generated ir is expected
assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)
@@ -280,16 +280,16 @@ def test_symbolic_buffer_indices():
@pytest.mark.parametrize("layout", [None, TileLayout(S[(2, 3) : (1, 2)])])
-def test_flat_buffer_store_preserves_identity_and_emits_once(layout):
+def test_flat_tensor_store_preserves_identity_and_emits_once(layout):
buffer = tirx.decl_tensor((2, 3), "float32", strides=(5, 1), elem_offset=2, layout=layout)
indices = T.tensor_indices(buffer, 4)
load = buffer[indices]
assert load.source.same_as(buffer)
with IRBuilder() as ib:
- receipt = T.buffer_store(buffer, load + 1, indices)
+ receipt = T.tensor_store(buffer, load + 1, indices)
T.emit_(receipt)
store = ib.get()
- assert isinstance(store, tirx.BufferStore)
+ assert isinstance(store, tirx.TensorStore)
assert store.buffer.same_as(buffer)
assert list(store.indices) == [1, 1]
assert store.value.a.source.same_as(buffer)
@@ -302,7 +302,7 @@ def _check_stores(output):
for tensor in outputs:
tvm_ffi.structural_walk(
tensor.op.body,
- lambda node: stores.append(node) if isinstance(node, tirx.BufferStore) else None,
+ lambda node: stores.append(node) if isinstance(node, tirx.TensorStore) else None,
)
assert stores
assert all(len(store.indices) == len(store.buffer.shape) for store in stores)
@@ -463,7 +463,7 @@ def test_concrete_mutable_scalar(declare):
stores = []
tvm_ffi.structural_walk(
ib.get().body,
- lambda node: stores.append(node) if isinstance(node, tirx.BufferStore) else None,
+ lambda node: stores.append(node) if isinstance(node, tirx.TensorStore) else None,
)
assert len(stores) == 2
assert all(store.buffer.same_as(scalar.source) for store in stores)
diff --git a/tests/python/tirx/script/test_tirx_script_meta_programming.py b/tests/python/tirx/script/test_tirx_script_meta_programming.py
index c44cd2de3b..eac3ebbb3c 100644
--- a/tests/python/tirx/script/test_tirx_script_meta_programming.py
+++ b/tests/python/tirx/script/test_tirx_script_meta_programming.py
@@ -214,7 +214,7 @@ def test_scalar_assign_in_macro():
The parser narrowed ``except Exception: pass`` around the scalar-detection
path. This test verifies that Expr assignment to a scalar attribute in
- a macro still goes through buffer_store correctly.
+ a macro still goes through tensor_store correctly.
The full integration regression for the TypeError fallthrough path
(meta_var assigned to a scalar variable) is covered by
@@ -227,7 +227,7 @@ def test_scalar_assign_in_macro():
@T.inline
def add_one(self):
- # Expr assigned to scalar via self.attr → buffer_store succeeds
+ # Expr assigned to scalar via self.attr → tensor_store succeeds
self.counter = self.counter + T.int32(1)
@T.prim_func(private=True)
diff --git a/tests/python/tirx/script/test_tirx_script_printer.py b/tests/python/tirx/script/test_tirx_script_printer.py
index a808af2ffb..1319465e1e 100644
--- a/tests/python/tirx/script/test_tirx_script_printer.py
+++ b/tests/python/tirx/script/test_tirx_script_printer.py
@@ -173,10 +173,10 @@ A[128, 128]
)
-def test_buffer_store():
+def test_tensor_store():
a = tirx.decl_tensor((128, 128), "float16", name="A")
with IRBuilder() as ib:
- TB.buffer_store(a, a[128, 128] + 1, [128, 128])
+ TB.tensor_store(a, a[128, 128] + 1, [128, 128])
obj = ib.get()
_assert_print(
obj,
diff --git a/tests/python/tirx/script/test_tirx_script_source_locations.py b/tests/python/tirx/script/test_tirx_script_source_locations.py
index 51cf765934..b3107477b1 100644
--- a/tests/python/tirx/script/test_tirx_script_source_locations.py
+++ b/tests/python/tirx/script/test_tirx_script_source_locations.py
@@ -285,7 +285,7 @@ def test_native_view_keeps_producer_identity_name_and_span(monkeypatch):
)
nodes = list(main.body.seq)
assert len(nodes) == 3 and isinstance(nodes[0], tirx.Bind)
- assert isinstance(nodes[1], tirx.BufferStore) and isinstance(nodes[2], tirx.Bind)
+ assert isinstance(nodes[1], tirx.TensorStore) and isinstance(nodes[2], tirx.Bind)
assert nodes[0].var.same_as(value) and nodes[1].buffer.same_as(value)
assert nodes[2].var.same_as(produced[1][0])
ir.assert_structural_equal(nodes[0].value.args[0], captured.data)
diff --git a/tests/python/tirx/test_op.py b/tests/python/tirx/test_op.py
index db8068c568..ff0b8248e4 100644
--- a/tests/python/tirx/test_op.py
+++ b/tests/python/tirx/test_op.py
@@ -88,13 +88,13 @@ def test_buffer_replacer_no_shared_default():
def test_buffer_replacer_replaces_strides_and_elem_offset():
"""Vars in buffer strides/elem_offset must be replaced, not passed through."""
- from tvm.tirx import BufferStore, Var
+ from tvm.tirx import TensorStore, Var
from tvm.tirx.transform.common import BufferReplacer
n = Var("n", "int32")
m = Var("m", "int32")
A = decl_tensor((64,), "float32", strides=[n], elem_offset=n)
- store = BufferStore(A, 1.0, [0])
+ store = TensorStore(A, 1.0, [0])
new = BufferReplacer(var_map={n: m})(store)
assert new.buffer.strides[0].same_as(m)
diff --git a/tests/python/tirx/transform/test_transform_flatten_buffer.py b/tests/python/tirx/transform/test_transform_flatten_buffer.py
index b428a580fd..9282678eff 100644
--- a/tests/python/tirx/transform/test_transform_flatten_buffer.py
+++ b/tests/python/tirx/transform/test_transform_flatten_buffer.py
@@ -70,7 +70,7 @@ def _assert_loads_reference_defined_buffers(func):
tvm_ffi.structural_walk(expr, visit)
def visit(node):
- if isinstance(node, tvm.ir.TensorLoad | tvm.tirx.BufferStore):
+ if isinstance(node, tvm.ir.TensorLoad | tvm.tirx.TensorStore):
buffer = node.source if isinstance(node, tvm.ir.TensorLoad) else node.buffer
if not is_defined(buffer):
stale.append(f"access of {buffer.name}")
@@ -127,7 +127,7 @@ def test_flatten_remaps_loads_in_folded_elem_offset():
found = []
def visit(node):
- if isinstance(node, tvm.tirx.BufferStore) and node.buffer.name.startswith("mbar"):
+ if isinstance(node, tvm.tirx.TensorStore) and node.buffer.name.startswith("mbar"):
def inner(sub):
if isinstance(sub, tvm.ir.TensorLoad):
diff --git a/tests/python/tirx/transform/test_transform_lower_tirx.py b/tests/python/tirx/transform/test_transform_lower_tirx.py
index 5199840c01..a6da323dbb 100644
--- a/tests/python/tirx/transform/test_transform_lower_tirx.py
+++ b/tests/python/tirx/transform/test_transform_lower_tirx.py
@@ -1479,7 +1479,7 @@ def test_lower_alloc_decl_buffer_outside_of_parser():
def int_var1(val):
buf = T.local_scalar("int32")
if val is not None:
- T.buffer_store(buf.source, val, 0)
+ T.tensor_store(buf.source, val, 0)
return buf
TestMutableCells = SimpleNamespace(int_var1=int_var1)
@@ -1488,7 +1488,7 @@ def test_lower_alloc_decl_buffer_outside_of_parser():
def int_var2(val):
buf = T.alloc_local([1], "int32")
if val is not None:
- T.buffer_store(buf, val, 0)
+ T.tensor_store(buf, val, 0)
return buf
@T.prim_func(private=True)