From f03ec06a7afbaf727d9747d091fcf95e0160fc1b Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Mon, 27 Jul 2026 06:16:50 +0000 Subject: [PATCH 1/2] [REFACTOR][TIRX] Represent buffers as typed variables Make buffer identity an ordinary ir.Var carrying BufferType. The previous parallel Buffer object and pointer-typed data Var could diverge under generic IR substitution and rebuilding; BufferType is now the single source of truth, and buffer_data(buffer) is the explicit physical-pointer projection. This lets standard Var scoping, substitution, undefined-variable analysis, and identity maps handle buffers without a second object hierarchy. It also removes duplicate constructor state and representation-specific remapping from the common path. C++ migration: use BufferVar as a checked Var view, read logical metadata through buffer->..., use buffer.var() for identity, and use buffer.data() or builtin::buffer_data for the pointer projection. Metadata changes rebuild the BufferType and rebind the new Var through AllocBuffer or DeclBuffer. Python migration: constructors return tvm.ir.Var, value.ty carries BufferType, and is_buffer_var(value) is the runtime discriminator. Buffer remains a source-compatibility alias for imports and annotations, so isinstance(value, Buffer) must not be used to distinguish buffer variables. --- include/tvm/arith/bound.h | 2 +- .../tvm/relax/distributed/axis_group_graph.h | 12 +- include/tvm/relax/tir_pattern.h | 4 +- include/tvm/s_tir/analysis.h | 6 +- include/tvm/s_tir/sblock_scope.h | 2 +- include/tvm/s_tir/schedule/schedule.h | 4 +- include/tvm/s_tir/transform.h | 2 +- include/tvm/te/operation.h | 8 +- include/tvm/tirx/async_structs.h | 6 +- include/tvm/tirx/buffer.h | 244 ++++++++++---- include/tvm/tirx/builtin.h | 8 + include/tvm/tirx/expr.h | 4 +- include/tvm/tirx/function.h | 17 +- include/tvm/tirx/index_map.h | 4 +- include/tvm/tirx/script/builder/frame.h | 17 +- include/tvm/tirx/script/builder/ir.h | 40 +-- include/tvm/tirx/stmt.h | 44 +-- include/tvm/tirx/stmt_functor.h | 13 +- include/tvm/tirx/tirx_op.h | 8 +- include/tvm/tirx/tirx_stmt.h | 6 +- include/tvm/topi/contrib/cublas.h | 4 +- include/tvm/topi/detail/extern.h | 8 +- python/tvm/s_tir/schedule/schedule.py | 4 +- python/tvm/tirx/__init__.py | 2 +- python/tvm/tirx/buffer.py | 79 ++++- python/tvm/tirx/function.py | 4 +- python/tvm/tirx/op.py | 14 +- python/tvm/tirx/script/builder/ir.py | 14 +- python/tvm/tirx/script/builder/tirx.py | 56 ++-- python/tvm/tirx/script/parser/parser.py | 12 +- python/tvm/tirx/stmt.py | 6 +- python/tvm/tirx/transform/common.py | 14 +- src/arith/domain_touched.cc | 12 +- src/backend/cuda/codegen/codegen_cuda.cc | 14 +- .../cuda/codegen/llvm/codegen_nvptx.cc | 10 +- .../hexagon/codegen/llvm/codegen_hexagon.cc | 14 +- src/backend/metal/codegen/codegen_metal.cc | 12 +- src/backend/opencl/codegen/codegen_opencl.cc | 14 +- src/backend/opencl/codegen/codegen_opencl.h | 6 +- .../rocm/codegen/llvm/codegen_amdgpu.cc | 10 +- src/backend/trn/codegen/codegen_trn.cc | 16 +- src/backend/trn/codegen/codegen_trn.h | 4 +- .../trn/transform/lower_trainium_layout.cc | 72 +++-- src/backend/vulkan/codegen/codegen_spirv.cc | 19 +- src/backend/webgpu/codegen/codegen_webgpu.cc | 10 +- src/relax/analysis/layout_transformation.cc | 24 +- src/relax/analysis/tir_op_pattern_kind.cc | 26 +- .../backend/adreno/annotate_custom_storage.cc | 4 +- src/relax/backend/vm/vm_shape_lower.cc | 6 +- .../lower_global_view_to_local_view.cc | 43 +-- src/relax/ir/tir_pattern.cc | 2 +- src/relax/op/op.cc | 4 +- src/relax/op/tensor/inspect.cc | 14 +- src/relax/transform/dataflow_inplace.cc | 14 +- src/relax/transform/fuse_tir.cc | 112 +++---- .../transform/rewrite_dataflow_reshape.cc | 2 +- .../specialize_primfunc_based_on_callsite.cc | 14 +- .../transform/split_call_tir_by_pattern.cc | 37 ++- .../transform/split_layout_rewrite_preproc.cc | 20 +- .../analysis/calculate_allocated_memory.cc | 4 +- src/s_tir/analysis/is_pure_function.cc | 4 +- src/s_tir/analysis/oob_checker.cc | 4 +- .../analysis/sblock_access_region_detector.cc | 66 ++-- .../sblock_buffer_access_lca_detector.cc | 36 +-- .../backend/adreno/inject_texture_alloc.cc | 9 +- src/s_tir/backend/adreno/texture_flatten.cc | 20 +- src/s_tir/meta_schedule/arg_info.cc | 4 +- .../feature_extractor/per_store_feature.cc | 66 ++-- .../disallow_async_strided_mem_copy.cc | 4 +- .../meta_schedule/postproc/rewrite_layout.cc | 24 +- .../schedule_rule/multi_level_tiling.cc | 6 +- .../multi_level_tiling_tensor_core.cc | 14 +- src/s_tir/sblock_scope.cc | 4 +- src/s_tir/schedule/analysis.h | 12 +- src/s_tir/schedule/analysis/analysis.cc | 36 +-- src/s_tir/schedule/analysis/layout.cc | 6 +- src/s_tir/schedule/analysis/reducer.cc | 14 +- src/s_tir/schedule/concrete_schedule.cc | 8 +- src/s_tir/schedule/concrete_schedule.h | 2 +- src/s_tir/schedule/ir_comparator.cc | 22 +- src/s_tir/schedule/ir_comparator.h | 14 +- src/s_tir/schedule/primitive.h | 4 +- .../primitive/annotate_buffer_access.cc | 10 +- .../schedule/primitive/block_annotate.cc | 50 +-- .../schedule/primitive/blockize_tensorize.cc | 30 +- src/s_tir/schedule/primitive/cache_index.cc | 6 +- .../schedule/primitive/cache_read_write.cc | 112 ++++--- src/s_tir/schedule/primitive/compute_at.cc | 36 +-- .../schedule/primitive/compute_inline.cc | 132 ++++---- .../primitive/layout_transformation.cc | 81 ++--- .../schedule/primitive/loop_transformation.cc | 14 +- src/s_tir/schedule/primitive/pad_einsum.cc | 42 +-- src/s_tir/schedule/primitive/read_write_at.cc | 30 +- src/s_tir/schedule/primitive/reduction.cc | 44 ++- .../schedule/primitive/rolling_buffer.cc | 28 +- src/s_tir/schedule/schedule.cc | 2 +- src/s_tir/schedule/state.cc | 11 +- src/s_tir/schedule/traced_schedule.cc | 2 +- src/s_tir/schedule/traced_schedule.h | 2 +- src/s_tir/schedule/transform.cc | 62 ++-- src/s_tir/schedule/transform.h | 26 +- src/s_tir/transform/bound_checker.cc | 12 +- src/s_tir/transform/compact_buffer_region.cc | 128 ++++---- src/s_tir/transform/inject_double_buffer.cc | 16 +- src/s_tir/transform/inject_permuted_layout.cc | 26 +- src/s_tir/transform/inject_ptx_async_copy.cc | 16 +- src/s_tir/transform/inject_ptx_ldg32.cc | 6 +- .../transform/inject_software_pipeline.cc | 128 ++++---- src/s_tir/transform/inject_virtual_thread.cc | 28 +- src/s_tir/transform/lower_async_dma.cc | 4 +- .../transform/lower_cross_thread_reduction.cc | 50 +-- src/s_tir/transform/lower_match_buffer.cc | 39 ++- src/s_tir/transform/lower_opaque_block.cc | 4 +- src/s_tir/transform/lower_thread_allreduce.cc | 95 +++--- src/s_tir/transform/lower_vtcm_alloc.cc | 9 +- .../manifest_shared_memory_local_stage.cc | 32 +- .../transform/memhammer_intermediate_stage.cc | 15 +- .../transform/memhammer_lower_auto_copy.cc | 26 +- src/s_tir/transform/memhammer_rewrite_rule.h | 38 +-- .../transform/memhammer_tensorcore_rewrite.cc | 40 +-- .../merge_shared_memory_allocations.cc | 110 ++++--- .../plan_update_buffer_allocation_location.cc | 64 ++-- .../remove_weight_layout_rewrite_block.cc | 18 +- src/s_tir/transform/renew_defs.cc | 43 ++- src/s_tir/transform/storage_access.cc | 8 +- .../transform/tensorcore_infer_fragment.cc | 6 +- .../transform/transform_mma_buffer_layout.cc | 16 +- .../using_assume_to_reduce_branches.cc | 10 +- src/target/llvm/codegen_llvm.cc | 40 +-- src/target/llvm/codegen_llvm.h | 2 +- src/target/source/codegen_c.cc | 44 ++- src/target/source/codegen_c.h | 6 +- src/te/operation/create_primfunc.cc | 64 ++-- src/te/operation/extern_op.cc | 8 +- src/tirx/analysis/var_touch.cc | 4 +- src/tirx/analysis/var_use_def_analysis.cc | 53 ++-- src/tirx/analysis/var_use_def_analysis.h | 16 +- src/tirx/analysis/verify_memory.cc | 6 +- src/tirx/analysis/verify_ssa.cc | 6 +- src/tirx/analysis/verify_well_formed.cc | 28 +- src/tirx/ir/async_structs.cc | 8 +- src/tirx/ir/buffer.cc | 297 ++++++++++-------- src/tirx/ir/data_type_rewriter.cc | 26 +- src/tirx/ir/data_type_rewriter.h | 4 +- src/tirx/ir/expr.cc | 6 +- src/tirx/ir/function.cc | 4 +- src/tirx/ir/script/script_complete.cc | 34 +- src/tirx/ir/script/script_complete.h | 2 +- src/tirx/ir/specialize.cc | 87 ++--- src/tirx/ir/stmt.cc | 39 +-- src/tirx/ir/stmt_functor.cc | 127 +++++--- src/tirx/ir/tir_visitor_with_path.cc | 36 ++- src/tirx/ir/tir_visitor_with_path.h | 17 +- src/tirx/ir/tirx_stmt.cc | 4 +- src/tirx/op/builtin.cc | 4 + src/tirx/op/tirx.cc | 8 +- src/tirx/script/builder/frame.cc | 32 +- src/tirx/script/builder/ir.cc | 98 +++--- src/tirx/script/printer/block.cc | 2 +- src/tirx/script/printer/buffer.cc | 69 ++-- src/tirx/script/printer/expr.cc | 24 ++ src/tirx/script/printer/function.cc | 24 +- src/tirx/script/printer/stmt.cc | 39 +-- src/tirx/script/printer/utils.h | 27 +- src/tirx/transform/flatten_buffer.cc | 49 ++- .../transform/force_narrow_index_to_i32.cc | 2 +- .../transform/inline_private_functions.cc | 2 +- src/tirx/transform/ir_utils.cc | 156 +++++---- src/tirx/transform/ir_utils.h | 6 +- src/tirx/transform/lower_intrin.cc | 2 +- src/tirx/transform/lower_tirx_cleanup.cc | 71 +++-- src/tirx/transform/lower_tirx_opaque.cc | 20 +- src/tirx/transform/lower_tvm_builtin.cc | 38 ++- src/tirx/transform/lower_warp_memory.cc | 25 +- src/tirx/transform/make_packed_api.cc | 2 +- src/tirx/transform/remove_no_op.cc | 2 +- src/tirx/transform/split_host_device.cc | 49 ++- src/tirx/transform/stmt_simplify.cc | 4 +- src/tirx/transform/storage_rewrite.cc | 215 ++++++++----- src/tirx/transform/tile_primitive_dispatch.cc | 32 +- src/tirx/transform/tvm_ffi_binder.cc | 52 +-- src/tirx/transform/tvm_ffi_binder.h | 18 +- src/tirx/transform/unroll_loop.cc | 4 +- .../transform/unsupported_dtype_legalize.cc | 91 +++--- .../transform/update_pointer_storage_scope.cc | 22 +- .../transform/update_pointer_storage_scope.h | 4 +- src/tirx/transform/vectorize_loop.cc | 26 +- tests/cpp/ir_functor_test.cc | 26 +- .../analysis/test_sblock_access_region.py | 26 +- .../test_tir_analysis_verify_well_formed.py | 3 +- tests/python/tirx-base/test_tir_buffer.py | 30 +- tests/python/tirx-base/test_tir_op_types.py | 3 +- .../test_tir_transform_flatten_buffer.py | 4 +- .../test_tir_transform_storage_rewrite.py | 5 +- tests/python/tirx/test_parser_printer.py | 18 +- 195 files changed, 3029 insertions(+), 2566 deletions(-) diff --git a/include/tvm/arith/bound.h b/include/tvm/arith/bound.h index 7ae36fb289f0..13ec37b4116c 100644 --- a/include/tvm/arith/bound.h +++ b/include/tvm/arith/bound.h @@ -77,7 +77,7 @@ IntSet DeduceBound(PrimExpr v, PrimExpr cond, * \param consider_stores If stores are considered. * \return The domain that covers all the calls or provides within the given statement. */ -Region DomainTouched(const Stmt& body, const tirx::Buffer& buffer, bool consider_loads, +Region DomainTouched(const Stmt& body, const tirx::BufferVar& buffer, bool consider_loads, bool consider_stores); } // namespace arith diff --git a/include/tvm/relax/distributed/axis_group_graph.h b/include/tvm/relax/distributed/axis_group_graph.h index e7945a1b65d9..e7cd828b1f16 100644 --- a/include/tvm/relax/distributed/axis_group_graph.h +++ b/include/tvm/relax/distributed/axis_group_graph.h @@ -40,7 +40,7 @@ namespace tirx { // (var, axis) using TIRVarAxis = std::pair; // (buffer, axis) -using BufferAxis = std::pair; +using BufferAxis = std::pair; class BufferAxisHash { public: size_t operator()(const BufferAxis& buffer_axis) const { @@ -70,7 +70,7 @@ class BufferAxisGraphExtractor : public StmtExprVisitor { static std::vector> GetTIRVarAxisGraph(const PrimFunc& prim_func) { BufferAxisGraphExtractor extractor; extractor(prim_func->body); - ffi::Map inverse_buffer_map; + ffi::Map inverse_buffer_map; for (const auto& pr : prim_func->buffer_map) { inverse_buffer_map.Set(pr.second, pr.first); } @@ -78,7 +78,7 @@ class BufferAxisGraphExtractor : public StmtExprVisitor { std::unordered_set visited; for (const auto& pr : prim_func->buffer_map) { Var param = pr.first; - Buffer buffer = pr.second; + BufferVar buffer = pr.second; for (int i = 0; i < static_cast(buffer->shape.size()); i++) { if (extractor.buffer_axis_graph_.count({buffer, i})) { std::vector buffer_axis_group; @@ -163,14 +163,14 @@ class BufferAxisGraphExtractor : public StmtExprVisitor { } arith::Analyzer analyzer; for (const auto& access_pr : buffer_access_indices_) { - Buffer buffer = access_pr.first; + BufferVar buffer = access_pr.first; ffi::Array indices = access_pr.second; for (int i = 0; i < static_cast(indices.size()); i++) { for (const auto& another_access_pr : buffer_access_indices_) { if (another_access_pr.first.same_as(buffer)) { continue; } - Buffer another_buffer = another_access_pr.first; + BufferVar another_buffer = another_access_pr.first; ffi::Array another_indices = another_access_pr.second; for (int j = 0; j < static_cast(another_indices.size()); j++) { if (Match(indices[i], buffer->shape[i], another_indices[j], another_buffer->shape[j], @@ -194,7 +194,7 @@ class BufferAxisGraphExtractor : public StmtExprVisitor { buffer_axis_graph_[axis2].push_back(axis1); } - std::vector>> buffer_access_indices_; + std::vector>> buffer_access_indices_; std::unordered_map, BufferAxisHash> buffer_axis_graph_; ffi::Map iter_var_range_; std::string func_name; diff --git a/include/tvm/relax/tir_pattern.h b/include/tvm/relax/tir_pattern.h index 4adf7a505fc2..0dcc63841349 100644 --- a/include/tvm/relax/tir_pattern.h +++ b/include/tvm/relax/tir_pattern.h @@ -43,7 +43,7 @@ class MatchResultNode : public ffi::Object { /*! \brief The evaluated values of symbolic vars. */ ffi::Array symbol_values; /*! \brief The matched buffers of input and output. */ - ffi::Array matched_buffers; + ffi::Array matched_buffers; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; @@ -67,7 +67,7 @@ class MatchResult : public ffi::ObjectRef { * \param matched_buffers The matched buffers of input and output. */ TVM_DLL explicit MatchResult(TIRPattern pattern, ffi::Array symbol_values, - ffi::Array matched_buffers); + ffi::Array matched_buffers); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MatchResult, ffi::ObjectRef, MatchResultNode); }; diff --git a/include/tvm/s_tir/analysis.h b/include/tvm/s_tir/analysis.h index b0cf7b38b9d5..9d2ad25397e3 100644 --- a/include/tvm/s_tir/analysis.h +++ b/include/tvm/s_tir/analysis.h @@ -48,7 +48,7 @@ namespace tirx { * - third: opaque regions */ TVM_DLL ffi::Array> GetSBlockAccessRegion( - const SBlock& block, const ffi::Map& buffer_var_map); + const SBlock& block, const ffi::Map& buffer_var_map); /*! * \brief Auto detect the block read/write region according to its body stmt. An opaque access will @@ -59,7 +59,7 @@ TVM_DLL ffi::Array> GetSBlockAccessRegion( * \return An array only consisting of the read regions and write regions of the input block */ TVM_DLL ffi::Array> GetSBlockReadWriteRegion( - const SBlock& block, const ffi::Map& buffer_var_map); + const SBlock& block, const ffi::Map& buffer_var_map); /*! * \brief Detect the lowest common ancestor(LCA) of buffer access, including both high-level @@ -69,7 +69,7 @@ TVM_DLL ffi::Array> GetSBlockReadWriteRegion( * \return The Map from buffer to the LCA of all access to it. The lca is function root if the * return stmt is std::nullopt. */ -TVM_DLL ffi::Map> DetectBufferAccessLCA(const PrimFunc& func); +TVM_DLL ffi::Map> DetectBufferAccessLCA(const PrimFunc& func); /*! * \brief Find the "anchor block" of the given module. diff --git a/include/tvm/s_tir/sblock_scope.h b/include/tvm/s_tir/sblock_scope.h index 774b6d3f2ede..fc3e6b19c84c 100644 --- a/include/tvm/s_tir/sblock_scope.h +++ b/include/tvm/s_tir/sblock_scope.h @@ -264,7 +264,7 @@ class SBlockScopeNode : public ffi::Object { std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> dst2deps; /*! \brief The mapping from the buffer to the blocks who write it */ - std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> + std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> buffer_writers; static void RegisterReflection() { diff --git a/include/tvm/s_tir/schedule/schedule.h b/include/tvm/s_tir/schedule/schedule.h index dd3b91877777..f77d9b4edbd8 100644 --- a/include/tvm/s_tir/schedule/schedule.h +++ b/include/tvm/s_tir/schedule/schedule.h @@ -753,7 +753,7 @@ class ScheduleNode : public ffi::Object { /*! * \brief Apply a transformation represented by IndexMap to buffer * \details The indices and the access region to the target buffer is transformed by the given - * index_map. The index_map is used to infer the new shape of the buffer. Buffer must be either + * index_map. The index_map is used to infer the new shape of the buffer. BufferVar must be either * a function parameter, or allocated in a block (it cannot be a buffer subregion created via * 'match_buffer'). * \param block_rv The block that accesses the target buffer. @@ -825,7 +825,7 @@ class ScheduleNode : public ffi::Object { */ virtual void PadEinsum(const SBlockRV& block_rv, const ffi::Array& padding) = 0; - /******** Schedule: Buffer transformation ********/ + /******** Schedule: BufferVar transformation ********/ /*! * \brief Compute the target buffer via rolling buffering. * \details This primitive selects the outermost rollable axis with a positive bound overlap that diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h index 36424ab6cc19..c410ffbae512 100644 --- a/include/tvm/s_tir/transform.h +++ b/include/tvm/s_tir/transform.h @@ -35,7 +35,7 @@ namespace tvm { namespace s_tir { /*! - * \brief Renew the definition nodes for a TIR, including Var, Buffer and IterVar. + * \brief Renew the definition nodes for a TIR, including Var, BufferVar and IterVar. * This pass works as a simple DeepCopy to duplicate a function with different Vars and * Buffers but the same behavior * \param func The input PrimFunc. diff --git a/include/tvm/te/operation.h b/include/tvm/te/operation.h index 2428c264e3f7..cff916f94b5a 100644 --- a/include/tvm/te/operation.h +++ b/include/tvm/te/operation.h @@ -257,9 +257,9 @@ class ExternOpNode : public OperationNode { /*! \brief The input tensors */ ffi::Array inputs; /*! \brief Symbolic placeholder representation of inputs */ - ffi::Array input_placeholders; + ffi::Array input_placeholders; /*! \brief Symbolic placeholder representation of outputs */ - ffi::Array output_placeholders; + ffi::Array output_placeholders; /*! \brief the statement that generates the computation. */ Stmt body; @@ -289,8 +289,8 @@ class ExternOpNode : public OperationNode { class ExternOp : public Operation { public: TVM_DLL ExternOp(std::string name, std::string tag, ffi::Map attrs, - ffi::Array inputs, ffi::Array input_placeholders, - ffi::Array output_placeholders, Stmt body); + ffi::Array inputs, ffi::Array input_placeholders, + ffi::Array output_placeholders, Stmt body); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ExternOp, Operation, ExternOpNode); }; diff --git a/include/tvm/tirx/async_structs.h b/include/tvm/tirx/async_structs.h index eb140309cb17..704ed089fba4 100644 --- a/include/tvm/tirx/async_structs.h +++ b/include/tvm/tirx/async_structs.h @@ -45,7 +45,7 @@ class PipelineNode : public ffi::Object { ffi::String name_hint; /*! \brief The workspace of the pipeline. */ - ffi::Map workspace; + ffi::Map workspace; /*! \brief The schedule config of the pipeline. */ ffi::Map schedule_config; @@ -68,7 +68,7 @@ class Pipeline : public ffi::ObjectRef { public: TVM_DLL explicit Pipeline(ExecScope thread_scope, size_t depth = 0, bool separate_pc = false, ffi::String name_hint = "", - ffi::Map workspace = {}, + ffi::Map workspace = {}, ffi::Map schedule_config = {}); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Pipeline, ffi::ObjectRef, PipelineNode); @@ -90,7 +90,7 @@ class CopyPipeline : public Pipeline { public: TVM_DLL explicit CopyPipeline(ExecScope thread_scope, size_t depth = 0, bool separate_pc = false, ffi::String name_hint = "", - ffi::Map workspace = {}, + ffi::Map workspace = {}, ffi::Map schedule_config = {}); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(CopyPipeline, Pipeline, CopyPipelineNode); diff --git a/include/tvm/tirx/buffer.h b/include/tvm/tirx/buffer.h index e2a3e167d8db..66138d27fc4f 100644 --- a/include/tvm/tirx/buffer.h +++ b/include/tvm/tirx/buffer.h @@ -60,15 +60,18 @@ inline DLDataType DefaultIndexType() { // forward declare Stmt class Stmt; -/*! \brief Node to represent a buffer */ -class BufferNode : public ffi::Object { +/*! + * \brief Structural type of a TIRx buffer variable. + * + * A buffer value is an ordinary VarNode whose ExprNode::ty is BufferType. + * BufferType owns the immutable access contract. The physical pointer is + * deliberately not stored here; it is obtained with buffer_data(BufferVar) + * and is bound by the surrounding buffer definition. + */ +class BufferTypeNode : public TypeNode { public: - // Data fields. - /*! - * \brief The pointer to the head of the data - * \sa data_alignment The alignment of data in bytes. - */ - Var data; + /*! \brief Type of the physical pointer projected by buffer_data. */ + PointerType data_pointer_type = PointerType::VoidPointerTy(); /*! \brief dtype in the content of the tensor */ PrimType dtype = PrimType::Void(); /*! \brief The type of the buffer prior to flattening @@ -85,9 +88,6 @@ class BufferNode : public ffi::Object { ffi::Array strides; /*! \brief The offset in terms of number of dtype elements (including lanes) */ PrimExpr elem_offset; - // Meta data - /*! \brief optional name of the buffer */ - ffi::String name; /*! \brief Alignment requirement of data pointer in bytes. */ int data_alignment; /*! @@ -95,12 +95,6 @@ class BufferNode : public ffi::Object { * elem_offset is guaranteed to be multiple of offset_factor. */ int offset_factor; - /*! - * \brief Span that points to the original source code. - * Reserved debug information. - */ - mutable Span span; - /*! \brief The layout of the buffer */ ffi::Optional layout; @@ -111,27 +105,24 @@ class BufferNode : public ffi::Object { ffi::Array allocated_addr; /*! \brief constructor */ - BufferNode() {} + BufferTypeNode() {} static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - // TODO(tqchen): use SEqHashDefNonRecursive after the next pypi tvm-ffi release - .def_ro("data", &BufferNode::data, refl::AttachFieldFlag::SEqHashDefRecursive()) - .def_ro("dtype", &BufferNode::dtype) + refl::ObjectDef() + .def_ro("data_pointer_type", &BufferTypeNode::data_pointer_type) + .def_ro("dtype", &BufferTypeNode::dtype) // TODO(tqchen): use SEqHashDefNonRecursive after the next pypi tvm-ffi release - .def_ro("shape", &BufferNode::shape, refl::AttachFieldFlag::SEqHashDefRecursive()) + .def_ro("shape", &BufferTypeNode::shape, refl::AttachFieldFlag::SEqHashDefRecursive()) // TODO(tqchen): use SEqHashDefNonRecursive after the next pypi tvm-ffi release - .def_ro("strides", &BufferNode::strides, refl::AttachFieldFlag::SEqHashDefRecursive()) + .def_ro("strides", &BufferTypeNode::strides, refl::AttachFieldFlag::SEqHashDefRecursive()) // TODO(tqchen): use SEqHashDefNonRecursive after the next pypi tvm-ffi release - .def_ro("elem_offset", &BufferNode::elem_offset, + .def_ro("elem_offset", &BufferTypeNode::elem_offset, refl::AttachFieldFlag::SEqHashDefRecursive()) - .def_ro("name", &BufferNode::name, refl::AttachFieldFlag::SEqHashIgnore()) - .def_ro("data_alignment", &BufferNode::data_alignment) - .def_ro("offset_factor", &BufferNode::offset_factor) - .def_ro("span", &BufferNode::span, refl::AttachFieldFlag::SEqHashIgnore()) - .def_ro("layout", &BufferNode::layout) - .def_ro("allocated_addr", &BufferNode::allocated_addr); + .def_ro("data_alignment", &BufferTypeNode::data_alignment) + .def_ro("offset_factor", &BufferTypeNode::offset_factor) + .def_ro("layout", &BufferTypeNode::layout) + .def_ro("allocated_addr", &BufferTypeNode::allocated_addr); } /*! \return preferred index type for this buffer node */ @@ -153,31 +144,80 @@ class BufferNode : public ffi::Object { */ ffi::Array ElemOffset(ffi::Array index, bool inner = false) const; - static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferType", BufferTypeNode, TypeNode); +}; + +/*! + * \brief Managed reference to BufferTypeNode. + */ +class BufferType : public Type { + public: + TVM_DLL BufferType(PointerType data_pointer_type, PrimType dtype, ffi::Array shape, + ffi::Array strides, PrimExpr elem_offset, int data_alignment, + int offset_factor, ffi::Optional layout = std::nullopt, + ffi::Array allocated_addr = {}, Span span = Span()); + + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BufferType, Type, BufferTypeNode); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Buffer", BufferNode, ffi::Object); + explicit BufferType(ffi::ObjectPtr n) : Type(ffi::UnsafeInit{}) { + TVM_FFI_ICHECK(n != nullptr); + data_ = std::move(n); + } }; /*! - * \brief Buffer is a symbolic n-darray structure. - * It is a composition of primitive symbolic types, - * used to specify the memory layout of the Tensor used in program input. + * \brief Checked zero-state view over an ordinary VarNode with BufferType. + * + * BufferVar does not introduce a runtime object or a second identity. It + * safely widens to Var, and get() returns the underlying VarNode used by + * identity-sensitive maps. operator-> exposes the immutable BufferType + * access contract for concise compiler-side metadata access. */ -class Buffer : public ffi::ObjectRef { +class BufferVar : public Var { public: // User can specify data_alignment and offset_factor to be 0 // A default value will be picked. - TVM_DLL Buffer(Var data, PrimType dtype, ffi::Array shape, ffi::Array strides, - PrimExpr elem_offset, ffi::String name, int data_alignment, int offset_factor, - Span span = Span(), ffi::Optional layout = std::nullopt, - ffi::Array allocated_addr = {}); + TVM_DLL BufferVar(Var data, PrimType dtype, ffi::Array shape, + ffi::Array strides, PrimExpr elem_offset, ffi::String name, + int data_alignment, int offset_factor, Span span = Span(), + ffi::Optional layout = std::nullopt, + ffi::Array allocated_addr = {}); + + /*! \brief Construct a fresh buffer variable from an explicit BufferType. */ + TVM_DLL explicit BufferVar(ffi::String name, BufferType type, Span span = Span()); + + /*! \brief Create a checked buffer view over an existing ordinary Var. */ + explicit BufferVar(Var var) : Var(std::move(var)) { + TVM_FFI_ICHECK(get() == nullptr || get()->ty.as()) + << "Expected a Var with BufferType"; + } + + /*! \brief Return the ordinary variable view over the same identity. */ + Var var() const { return ffi::GetRef(get()); } + + /*! \brief Return the buffer type carried by the ordinary variable. */ + BufferType type() const { return get()->ty.as_or_throw(); } + + /*! \brief Return the buffer's diagnostic name. */ + const ffi::String& name() const { return get()->name; } + + /*! \brief Return the source span carried by the ordinary Var. */ + const Span& span() const { return get()->span; } + + /*! \brief Return a fresh buffer variable with the same type and a new name. */ + BufferVar CopyWithName(ffi::String name) const { + return BufferVar(std::move(name), type(), span()); + } + + /*! \brief Project the physical pointer established by the definition site. */ + TVM_DLL Expr data() const; /*! * \brief Return a new buffer that is equivalent with current one * but always add stride field. * \return The strided version of the buffer. */ - TVM_DLL Buffer MakeStrideView() const; + TVM_DLL BufferVar MakeStrideView() const; /*! * \brief Make a new symbolic buffer representing a slice of the buffer. * \param begins The beginning position of each dimension. @@ -186,7 +226,8 @@ class Buffer : public ffi::ObjectRef { * If stride is not needed in the slice, it won't be presented * \return the result buffer. */ - TVM_DLL Buffer MakeSlice(ffi::Array begins, ffi::Array extents) const; + TVM_DLL BufferVar MakeSlice(ffi::Array begins, + ffi::Array extents) const; /*! * \brief Get access ptr to the entire buffer. * \param access_mask The access mask @@ -220,7 +261,7 @@ class Buffer : public ffi::ObjectRef { /*! * \brief Get a flattened version of the buffer */ - Buffer GetFlattenedBuffer() const; + BufferVar GetFlattenedBuffer() const; /*! \brief Determine the offset in the buffer of the given index. * @@ -245,7 +286,7 @@ class Buffer : public ffi::ObjectRef { /*! * \brief Return a new buffer with the allocated address. */ - TVM_DLL Buffer with_allocated_addr(ffi::Array allocated_addr) const; + TVM_DLL BufferVar with_allocated_addr(ffi::Array allocated_addr) const; /*! * \brief Return true if the buffer is a scalar. @@ -257,20 +298,57 @@ class Buffer : public ffi::ObjectRef { /*! * \brief Return a new buffer with the dtype. */ - TVM_DLL Buffer with_dtype(PrimType dtype) const; + TVM_DLL BufferVar with_dtype(PrimType dtype) const; /*! \return primitive element type for compiler-side uses. */ PrimType ElementType() const { return (*this)->ElementType(); } - /*! - * \brief Return a new buffer with the data. - */ - TVM_DLL Buffer with_data(Var data) const; + BufferVar() = default; + explicit BufferVar(ffi::ObjectPtr n) : Var(std::move(n)) {} + explicit BufferVar(ffi::UnsafeInit tag) : Var(tag) {} + TVM_FFI_DEFINE_DEFAULT_COPY_MOVE_AND_ASSIGN(BufferVar); + + const BufferTypeNode* operator->() const { + const auto* var_node = static_cast(data_.get()); + TVM_FFI_ICHECK(var_node != nullptr); + const auto* type_node = var_node->ty.as(); + TVM_FFI_ICHECK(type_node != nullptr) + << "Expected a Var with BufferType, but " << var_node->name << " has type " + << var_node->ty; + return type_node; + } + + const VarNode* get() const { return static_cast(data_.get()); } - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Buffer, ffi::ObjectRef, BufferNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferNode); + [[maybe_unused]] static constexpr bool _type_is_nullable = true; + static constexpr bool _type_container_is_exact = false; + using ContainerType = VarNode; }; +inline bool operator==(const BufferVar& lhs, const BufferVar& rhs) { + return lhs.same_as(rhs); +} + +inline bool operator!=(const BufferVar& lhs, const BufferVar& rhs) { + return !lhs.same_as(rhs); +} + +/*! \brief Recover a checked buffer view from an ordinary VarNode pointer. */ +inline BufferVar GetBufferVar(const VarNode* var) { + return BufferVar(ffi::GetRef(var)); +} + +inline const BufferVar& GetBufferVar(const BufferVar& var) { return var; } + +inline ffi::ObjectPtr CopyBufferType(const BufferVar& var) { + return ffi::make_object(*var.operator->()); +} + +inline BufferVar RebuildBufferVar(const BufferVar& var, ffi::ObjectPtr type, + ffi::Optional name = std::nullopt) { + return BufferVar(name.value_or(var.name()), BufferType(std::move(type)), var.span()); +} + /*! * \brief Construct a new buffer given shape, and dtype. * \param shape The shape of the buffer, @@ -279,11 +357,12 @@ class Buffer : public ffi::ObjectRef { * \param storage_scope The storage scope associated with this buffer * \param span The location of this object in the source code. * \return The created buffer. - * \sa Buffer for complete constructor. + * \sa BufferVar for complete constructor. */ -TVM_DLL Buffer decl_buffer(ffi::Array shape, PrimType dtype = PrimType::Float(32), - ffi::String name = "buffer", ffi::String storage_scope = "", - Span span = Span()); +TVM_DLL BufferVar decl_buffer(ffi::Array shape, + PrimType dtype = PrimType::Float(32), + ffi::String name = "buffer", + ffi::String storage_scope = "", Span span = Span()); /*! * \brief Base node for data producers. @@ -329,7 +408,7 @@ class DataProducer : public PrimExprConvertible { }; /*! - * \brief Creates TIR Buffer for provided parameters + * \brief Creates TIR BufferVar for provided parameters * \param shape shape of the buffer * \param dtype data type * \param name buffer name @@ -340,9 +419,54 @@ class DataProducer : public PrimExprConvertible { * A default value will be picked. * \param memory_scope memory scope of the buffer */ -TVM_DLL tirx::Buffer BufferWithOffsetAlignment(ffi::Array shape, PrimType dtype, - std::string name, int data_alignment, - int offset_factor, std::string memory_scope = ""); +TVM_DLL tirx::BufferVar BufferWithOffsetAlignment(ffi::Array shape, PrimType dtype, + std::string name, int data_alignment, + int offset_factor, + std::string memory_scope = ""); } // namespace tirx } // namespace tvm + +namespace tvm::ffi { + +template <> +inline constexpr bool use_default_type_traits_v = false; + +template <> +struct TypeTraits : public ObjectRefTypeTraitsBase { + using Base = ObjectRefTypeTraitsBase; + using Base::CopyFromAnyViewAfterCheck; + using Base::CopyToAnyView; + using Base::GetMismatchTypeInfo; + using Base::MoveFromAnyAfterCheck; + using Base::MoveToAny; + using Base::TypeSchema; + using Base::TypeStr; + + TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { + if (src->type_index == TypeIndex::kTVMFFINone) { + return tirx::BufferVar::_type_is_nullable; + } + if (src->type_index != tirx::VarNode::RuntimeTypeIndex()) { + return false; + } + const auto* var = static_cast( + details::ObjectUnsafe::ObjectPtrFromUnowned(src->v_obj).get()); + return details::AnyUnsafe::CheckAnyStrict(var->ExprNode::ty); + } + + TVM_FFI_INLINE static std::optional TryCastFromAnyView( + const TVMFFIAny* src) { + if (CheckAnyStrict(src)) { + if (src->type_index == TypeIndex::kTVMFFINone) { + return details::ObjectUnsafe::ObjectRefFromObjectPtr(nullptr); + } + return details::ObjectUnsafe::ObjectRefFromObjectPtr( + details::ObjectUnsafe::ObjectPtrFromUnowned(src->v_obj)); + } + return std::nullopt; + } +}; + +} // namespace tvm::ffi + #endif // TVM_TIR_BUFFER_H_ diff --git a/include/tvm/tirx/builtin.h b/include/tvm/tirx/builtin.h index ee9d698ed2e9..879196823463 100644 --- a/include/tvm/tirx/builtin.h +++ b/include/tvm/tirx/builtin.h @@ -780,6 +780,14 @@ TVM_DLL const Op& ignore_loop_partition(); */ TVM_DLL const Op& buffer_offset(); +/*! + * \brief Project the physical pointer associated with a BufferVar definition. + * + * The result type is the BufferType::data_pointer_type of the sole BufferVar + * argument. This operation is consumed by TIRx lowering and code generation. + */ +TVM_DLL const Op& buffer_data(); + /*! \brief The kind of structure field info used in intrinsic */ enum TVMStructFieldKind : int { // DLTensor fields diff --git a/include/tvm/tirx/expr.h b/include/tvm/tirx/expr.h index 0ff426ca55a8..d4a63b62ee87 100644 --- a/include/tvm/tirx/expr.h +++ b/include/tvm/tirx/expr.h @@ -545,7 +545,7 @@ class Select : public PrimExpr { class BufferLoadNode : public ExprNode { public: /*! \brief The buffer variable. */ - Buffer buffer; + BufferVar buffer; /*! \brief The indices location to be loaded. */ ffi::Array indices; /*! \brief The predicate mask for loading values. */ @@ -582,7 +582,7 @@ class BufferLoadNode : public ExprNode { */ class BufferLoad : public PrimExpr { public: - TVM_DLL explicit BufferLoad(Buffer buffer, ffi::Array indices, + TVM_DLL explicit BufferLoad(BufferVar buffer, ffi::Array indices, ffi::Optional predicate = std::nullopt, Span span = Span()); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BufferLoad, PrimExpr, BufferLoadNode); static constexpr bool _type_container_is_exact = true; diff --git a/include/tvm/tirx/function.h b/include/tvm/tirx/function.h index 531355d2c50d..6b35c0bbb885 100644 --- a/include/tvm/tirx/function.h +++ b/include/tvm/tirx/function.h @@ -53,13 +53,13 @@ class PrimFuncNode : public BaseFuncNode { /*! \brief The return type of the function. */ Type ret_type = Type::Missing(); /*! - * \brief Maps some parameters to specific Buffer data structures. + * \brief Maps some parameters to specific BufferVar data structures. * * buffer_map provides a way to express data structure's field and shape * constraints. The provided information is used in the program analysis * and the code generation. * - * - It defines the vars in the Buffer (m, n) in the cases below when + * - It defines the vars in the BufferVar (m, n) in the cases below when * they appears in the buffer_map for the first time. * - When a var appears multiple times, they translate into runtime * assertion to check the field constraint. @@ -97,7 +97,7 @@ class PrimFuncNode : public BaseFuncNode { * all usage in the body of the function is done through a * flattened alias of the buffer. */ - ffi::Map buffer_map; + ffi::Map buffer_map; /*! \brief The body of the function */ tirx::Stmt body; @@ -106,7 +106,8 @@ class PrimFuncNode : public BaseFuncNode { refl::ObjectDef() .def_ro("params", &PrimFuncNode::params, refl::AttachFieldFlag::SEqHashDefRecursive()) .def_ro("ret_type", &PrimFuncNode::ret_type) - .def_ro("buffer_map", &PrimFuncNode::buffer_map) + .def_ro("buffer_map", &PrimFuncNode::buffer_map, + refl::AttachFieldFlag::SEqHashDefRecursive()) .def_ro("body", &PrimFuncNode::body); refl::TypeAttrDef() .def("__s_equal__", &PrimFuncNode::SEqual) @@ -121,7 +122,7 @@ class PrimFuncNode : public BaseFuncNode { return equal(attrs, other->attrs, false, "attrs") && equal(params, other->params, true, "params") && equal(ret_type, other->ret_type, false, "ret_type") && - equal(buffer_map, other->buffer_map, false, "buffer_map") && + equal(buffer_map, other->buffer_map, true, "buffer_map") && equal(body, other->body, false, "body"); } @@ -130,7 +131,7 @@ class PrimFuncNode : public BaseFuncNode { hash_value = hash(attrs, hash_value, false); hash_value = hash(params, hash_value, true); hash_value = hash(ret_type, hash_value, false); - hash_value = hash(buffer_map, hash_value, false); + hash_value = hash(buffer_map, hash_value, true); hash_value = hash(body, hash_value, false); return hash_value; } @@ -172,7 +173,7 @@ class PrimFunc : public BaseFunc { * \param span The location of this object in the source code. */ TVM_DLL PrimFunc(ffi::Array params, Stmt body, Type ret_type = VoidType(), - ffi::Map buffer_map = ffi::Map(), + ffi::Map buffer_map = ffi::Map(), DictAttrs attrs = DictAttrs(), Span span = Span()); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrimFunc, BaseFunc, PrimFuncNode); @@ -273,7 +274,7 @@ class TensorIntrin : public ffi::ObjectRef { * B[vi, vj] = A[vi, vj] * \endcode */ -PrimFunc Specialize(PrimFunc func, const ffi::Map>& param_map); +PrimFunc Specialize(PrimFunc func, const ffi::Map>& param_map); /*! * \brief PrimFunc specific attribute names. diff --git a/include/tvm/tirx/index_map.h b/include/tvm/tirx/index_map.h index cb6463e1fe58..776af5b5e885 100644 --- a/include/tvm/tirx/index_map.h +++ b/include/tvm/tirx/index_map.h @@ -21,7 +21,7 @@ * \file tvm/tirx/index_map.h * \brief Defines a remapping of buffer indices * - * For use with tvm::tirx::Buffer. + * For use with tvm::tirx::BufferVar. */ #ifndef TVM_TIR_INDEX_MAP_H_ #define TVM_TIR_INDEX_MAP_H_ @@ -46,7 +46,7 @@ namespace tirx { * \brief Defines a mapping between two representations of indices * into a buffer. * - * This is primarily used for layout transformations of Buffer + * This is primarily used for layout transformations of BufferVar * objects. */ class IndexMapNode : public ffi::Object { diff --git a/include/tvm/tirx/script/builder/frame.h b/include/tvm/tirx/script/builder/frame.h index 5b2d3953269b..1bd2d53981a8 100644 --- a/include/tvm/tirx/script/builder/frame.h +++ b/include/tvm/tirx/script/builder/frame.h @@ -77,14 +77,14 @@ class PrimFuncFrameNode : public TIRFrameNode { bool is_private; /*! \brief The return type of the function. */ ffi::Optional ret_type; - /*! \brief Maps some parameters to specific Buffer data structures. */ - ffi::Map buffer_map; + /*! \brief Maps some parameters to specific BufferVar data structures. */ + ffi::Map buffer_map; /*! \brief Additional attributes storing the meta-data */ ffi::Map attrs; /*! \brief The variable map bound to thread env. */ ffi::Map env_threads; /*! \brief The buffer allocated in root block. */ - ffi::Array root_alloc_buffers; + ffi::Array root_alloc_buffers; // TIR utils /*! \brief Whether this PrimFunc uses s_tir semantics (root SBlock wrap, @@ -150,7 +150,7 @@ class SBlockFrameNode : public TIRFrameNode { /*! \brief The init statement of the bolck. */ ffi::Optional init; /*! \brief The buffer allocated in the block. */ - ffi::Array alloc_buffers; + ffi::Array alloc_buffers; /*! \brief The match buffer regions. */ ffi::Array match_buffers; /*! \brief The annotation of the block. */ @@ -610,7 +610,9 @@ class ElseFrame : public TIRFrame { class DeclBufferFrameNode : public TIRFrameNode { public: /*! \brief The declared buffer. */ - tvm::tirx::Buffer buffer; + tvm::tirx::BufferVar buffer; + /*! \brief Optional physical pointer expression backing the declaration. */ + ffi::Optional data; /*! \brief The buffer allocated or not. */ bool allocated; @@ -618,6 +620,7 @@ class DeclBufferFrameNode : public TIRFrameNode { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def_ro("buffer", &DeclBufferFrameNode::buffer) + .def_ro("data", &DeclBufferFrameNode::data) .def_ro("allocated", &DeclBufferFrameNode::allocated); } TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.DeclBufferFrame", DeclBufferFrameNode, @@ -638,7 +641,7 @@ class DeclBufferFrame : public TIRFrame { class ComposeOpFrameNode : public TIRFrameNode { public: /*! \brief The workspace of the compose op. */ - ffi::Map workspace; + ffi::Map workspace; /*! \brief The config of the compose op. */ ffi::Map config; /*! \brief The optional dispatch variant name of the compose op. */ @@ -669,7 +672,7 @@ class ComposeOpFrame : public TIRFrame { class AllocBufferFrameNode : public TIRFrameNode { public: /*! \brief The allocated buffer. */ - tvm::tirx::Buffer buffer; + tvm::tirx::BufferVar buffer; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; diff --git a/include/tvm/tirx/script/builder/ir.h b/include/tvm/tirx/script/builder/ir.h index ff005566417e..80cd2aab8efb 100644 --- a/include/tvm/tirx/script/builder/ir.h +++ b/include/tvm/tirx/script/builder/ir.h @@ -37,7 +37,7 @@ namespace tirx { using tvm::ffi::Tuple; using tvm::ffi::Variant; using tvm::runtime::Tensor; -using tvm::tirx::Buffer; +using tvm::tirx::BufferVar; using tvm::tirx::ExecScope; using tvm::tirx::Layout; using tvm::tirx::Var; @@ -55,11 +55,12 @@ using tvm::tirx::Var; * \param offset_factor The factor of elem_offset field. * \return The declared buffer. */ -Buffer BufferDecl(ffi::Array shape, PrimType dtype, ffi::String buffer_name, - ffi::Optional data, ffi::Optional> strides, - ffi::Optional elem_offset, ffi::String storage_scope, int align, - int offset_factor, ffi::Optional layout = std::nullopt, - ffi::Array allocated_addr = {}); +BufferVar BufferDecl(ffi::Array shape, PrimType dtype, ffi::String buffer_name, + ffi::Optional data, + ffi::Optional> strides, + ffi::Optional elem_offset, ffi::String storage_scope, int align, + int offset_factor, ffi::Optional layout = std::nullopt, + ffi::Array allocated_addr = {}); /*! * \brief The primitive function statement. @@ -81,7 +82,7 @@ Var Arg(ffi::String name, Var var); * \param buffer The buffer argument. * \return The buffer. */ -Buffer Arg(ffi::String name, Buffer buffer); +BufferVar Arg(ffi::String name, BufferVar buffer); /*! * \brief The PrimFunc naming statement. @@ -115,11 +116,13 @@ Type FuncRet(Type ret_type); * \param offset_factor The factor of elem_offset field. * \return The matched buffer. */ -Buffer MatchBuffer(ffi::ObjectRef param, ffi::Array shape, - PrimType dtype = PrimType::Float(32), ffi::Optional data = std::nullopt, - ffi::Array strides = {}, PrimExpr elem_offset = PrimExpr(), - ffi::String storage_scope = "global", int align = -1, int offset_factor = 0, - ffi::Optional layout = std::nullopt); +BufferVar MatchBuffer(ffi::ObjectRef param, ffi::Array shape, + PrimType dtype = PrimType::Float(32), + ffi::Optional data = std::nullopt, + ffi::Array strides = {}, PrimExpr elem_offset = PrimExpr(), + ffi::String storage_scope = "global", int align = -1, + int offset_factor = 0, + ffi::Optional layout = std::nullopt); /*! * \brief The block declaration statement. @@ -186,9 +189,9 @@ void BlockAttrs(ffi::Map attrs); * \return The allocated buffer or the AllocBufferFrame if the function is called under * T.prim_func(tirx=True). */ -ffi::Variant SBlockAllocBuffer( +ffi::Variant SBlockAllocBuffer( ffi::Array shape, PrimType dtype = PrimType::Float(32), - ffi::Optional data = std::nullopt, ffi::Array strides = {}, + ffi::Optional data = std::nullopt, ffi::Array strides = {}, PrimExpr elem_offset = PrimExpr(), ffi::String storage_scope = "", int align = -1, int offset_factor = 0, ffi::Optional layout = std::nullopt, ffi::Array allocated_addr = {}); @@ -406,7 +409,8 @@ ElseFrame Else(); * \return The declaration frame. */ DeclBufferFrame DeclBuffer(ffi::Array shape, PrimType dtype, ffi::String buffer_name, - ffi::Optional data, ffi::Optional> strides, + ffi::Optional data, + ffi::Optional> strides, ffi::Optional elem_offset, ffi::String storage_scope, int align, int offset_factor, ffi::Optional layout = std::nullopt, @@ -420,7 +424,7 @@ DeclBufferFrame DeclBuffer(ffi::Array shape, PrimType dtype, ffi::Stri * \param annotations Optional annotations for the allocation. * \return The allocated buffer. */ -Buffer AllocBuffer(ffi::Array shape, PrimType dtype = PrimType::Float(32), +BufferVar AllocBuffer(ffi::Array shape, PrimType dtype = PrimType::Float(32), ffi::String storage_scope = "global", ffi::Optional> annotations = std::nullopt); @@ -447,7 +451,7 @@ LaunchThreadFrame LaunchThread(ffi::String thread_tag, PrimExpr extent); * \param dispatch The optional dispatch variant name. * \return The result ComposeOpFrame. */ -ComposeOpFrame ComposeOp(ffi::Map workspace, +ComposeOpFrame ComposeOp(ffi::Map workspace, ffi::Map config, ffi::Optional dispatch = std::nullopt); @@ -467,7 +471,7 @@ Var EnvThread(ffi::String thread_tag, PrimType dtype = PrimType::Int(32)); * \param predicate A vector mask of boolean values indicating which lanes of a vector are to be * stored. The number lanes of the mask must be equal to the number of lanes in value. */ -void BufferStore(Buffer buffer, PrimExpr value, ffi::Array indices, +void BufferStore(BufferVar buffer, PrimExpr value, ffi::Array indices, ffi::Optional predicate); /*! diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h index fa633b5a1125..cf762abf0d05 100644 --- a/include/tvm/tirx/stmt.h +++ b/include/tvm/tirx/stmt.h @@ -202,7 +202,7 @@ class AssertStmt : public Stmt { class BufferStoreNode : public StmtNode { public: /*! \brief The buffer variable. */ - Buffer buffer; + BufferVar buffer; /*! \brief The value to be stored. */ PrimExpr value; /*! \brief The indices location to be stored. */ @@ -227,7 +227,7 @@ class BufferStoreNode : public StmtNode { */ class BufferStore : public Stmt { public: - TVM_DLL explicit BufferStore(Buffer buffer, PrimExpr value, ffi::Array indices, + TVM_DLL explicit BufferStore(BufferVar buffer, PrimExpr value, ffi::Array indices, ffi::Optional predicate = std::nullopt, Span span = Span()); @@ -239,11 +239,16 @@ class BufferStore : public Stmt { class DeclBufferNode : public StmtNode { public: /*! \brief The buffer being declared */ - Buffer buffer; + BufferVar buffer; + /*! \brief Optional physical pointer expression backing the declaration. */ + ffi::Optional data; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("buffer", &DeclBufferNode::buffer); + refl::ObjectDef() + .def_ro("buffer", &DeclBufferNode::buffer, + refl::AttachFieldFlag::SEqHashDefRecursive()) + .def_ro("data", &DeclBufferNode::data); } TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.DeclBuffer", DeclBufferNode, StmtNode); }; @@ -251,7 +256,8 @@ class DeclBufferNode : public StmtNode { /*! \brief Managed reference to DeclBufferNode */ class DeclBuffer : public Stmt { public: - TVM_DLL DeclBuffer(Buffer buffer, Span span = Span()); + TVM_DLL DeclBuffer(BufferVar buffer, ffi::Optional data = std::nullopt, + Span span = Span()); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DeclBuffer, Stmt, DeclBufferNode); TVM_DEFINE_OBJECT_REF_COW_METHOD(DeclBufferNode); }; @@ -260,7 +266,7 @@ class DeclBuffer : public Stmt { class AllocBufferNode : public StmtNode { public: /*! \brief The buffer being allocated and declared */ - Buffer buffer; + BufferVar buffer; /*! * \brief Additional annotations about the allocation. * @@ -283,7 +289,7 @@ class AllocBufferNode : public StmtNode { class AllocBuffer : public Stmt { public: TVM_DLL AllocBuffer( - Buffer buffer, + BufferVar buffer, ffi::Map annotations = ffi::Map(), Span span = Span()); /*! @@ -772,7 +778,7 @@ class Continue : public Stmt { class BufferRegionNode : public PrimExprConvertibleNode { public: /*! \brief The buffer of the buffer region. */ - Buffer buffer; + BufferVar buffer; /*! \brief The region array of the buffer region. */ ffi::Array region; @@ -795,14 +801,14 @@ class BufferRegionNode : public PrimExprConvertibleNode { */ class BufferRegion : public PrimExprConvertible { public: - TVM_DLL explicit BufferRegion(Buffer buffer, ffi::Array region); + TVM_DLL explicit BufferRegion(BufferVar buffer, ffi::Array region); /*! * \brief Create a BufferRegion which is full region of the given buffer. * \param buffer The buffer to generate full BufferRegion. * \return The BufferRegion which covers all region of the given buffer */ - TVM_DLL static BufferRegion FullRegion(Buffer buffer); + TVM_DLL static BufferRegion FullRegion(BufferVar buffer); /*! * \brief Create a BufferRegion which is a single point of the given buffer. @@ -810,7 +816,7 @@ class BufferRegion : public PrimExprConvertible { * \param indices The access point indices of the buffer * \return The BufferRegion which is the single point of the given buffer. */ - TVM_DLL static BufferRegion FromPoint(Buffer buffer, ffi::Array indices); + TVM_DLL static BufferRegion FromPoint(BufferVar buffer, ffi::Array indices); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BufferRegion, PrimExprConvertible, BufferRegionNode); TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferRegionNode); @@ -828,14 +834,15 @@ class BufferRegion : public PrimExprConvertible { class MatchBufferRegionNode : public ffi::Object { public: /*! \brief The target buffer. */ - Buffer buffer; + BufferVar buffer; /*! \brief The source buffer region. */ BufferRegion source; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() - .def_ro("buffer", &MatchBufferRegionNode::buffer) + .def_ro("buffer", &MatchBufferRegionNode::buffer, + refl::AttachFieldFlag::SEqHashDefRecursive()) .def_ro("source", &MatchBufferRegionNode::source); } @@ -849,7 +856,7 @@ class MatchBufferRegionNode : public ffi::Object { */ class MatchBufferRegion : public ffi::ObjectRef { public: - TVM_DLL explicit MatchBufferRegion(Buffer buffer, BufferRegion source); + TVM_DLL explicit MatchBufferRegion(BufferVar buffer, BufferRegion source); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MatchBufferRegion, ffi::ObjectRef, MatchBufferRegionNode); @@ -888,7 +895,7 @@ class SBlockNode : public StmtNode { /*! \brief The name_hint of the block. */ ffi::String name_hint; /*! \brief The buffer allocated in the block. */ - ffi::Array alloc_buffers; + ffi::Array alloc_buffers; /*! \brief The match buffer regions. */ ffi::Array match_buffers; /*! \brief The annotation of the block. */ @@ -911,7 +918,8 @@ class SBlockNode : public StmtNode { .def_ro("reads", &SBlockNode::reads) .def_ro("writes", &SBlockNode::writes) .def_ro("name_hint", &SBlockNode::name_hint, refl::AttachFieldFlag::SEqHashIgnore()) - .def_ro("alloc_buffers", &SBlockNode::alloc_buffers) + .def_ro("alloc_buffers", &SBlockNode::alloc_buffers, + refl::AttachFieldFlag::SEqHashDefRecursive()) .def_ro("match_buffers", &SBlockNode::match_buffers) .def_ro("annotations", &SBlockNode::annotations) .def_ro("init", &SBlockNode::init) @@ -930,13 +938,13 @@ class SBlock : public Stmt { ffi::Array iter_vars, ffi::Array reads, ffi::Array writes, ffi::String name_hint, Stmt body, ffi::Optional init = std::nullopt, - ffi::Array alloc_buffers = ffi::Array(), + ffi::Array alloc_buffers = ffi::Array(), ffi::Array match_buffers = ffi::Array(), ffi::Map annotations = ffi::Map(), Span span = Span()); TVM_DLL explicit SBlock(ffi::String name_hint, Stmt body, - ffi::Array alloc_buffers = ffi::Array(), + ffi::Array alloc_buffers = ffi::Array(), Span span = Span()); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SBlock, Stmt, SBlockNode); diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h index 0df163d5f95a..f1aca2b94712 100644 --- a/include/tvm/tirx/stmt_functor.h +++ b/include/tvm/tirx/stmt_functor.h @@ -162,13 +162,13 @@ class TVM_DLL StmtVisitor : protected StmtFunctor { * \param alloc_data If true, the buffer's data pointer is a new allocation (AllocBuffer); * if false, data references an existing variable (DeclBuffer). */ - virtual void VisitBufferDef(const Buffer& buffer, bool alloc_data); + virtual void VisitBufferDef(const BufferVar& buffer, bool alloc_data); /*! * \brief Visit buffer at use site (BufferStore, BufferLoad, SBlock reads/writes). * By default, this is a no-op, as buffer fields (shape, strides, elem_offset) * are visited at their definition site. */ - virtual void VisitBufferUse(const Buffer& buffer); + virtual void VisitBufferUse(const BufferVar& buffer); // statement visitor void VisitStmt_(const BindNode* op) override; void VisitStmt_(const AttrStmtNode* op) override; @@ -210,7 +210,7 @@ class TVM_DLL StmtMutator : protected StmtFunctor { protected: /*! \brief Map from old buffer to new buffer, populated by VisitBufferDef. */ - ffi::Map buffer_remap_; + ffi::Map buffer_remap_; // We perform copy on write optimizations on the StmtMutator // so that an unique copy of parent can be mutated inplace // when some of its children changed. @@ -282,14 +282,14 @@ class TVM_DLL StmtMutator : protected StmtFunctor { * if false, data references an existing variable (DeclBuffer). * \return The (possibly new) buffer. */ - virtual Buffer VisitBufferDef(const Buffer& buffer, bool alloc_data); + virtual BufferVar VisitBufferDef(const BufferVar& buffer, bool alloc_data); /*! * \brief Visit buffer at use site (BufferStore, BufferLoad, SBlock reads/writes). * By default, returns the remapped buffer from buffer_remap_ if exists, otherwise - * returns the original buffer. Buffer fields are visited at their definition site. + * returns the original buffer. BufferVar fields are visited at their definition site. * \return The (possibly remapped) buffer. */ - virtual Buffer VisitBufferUse(const Buffer& buffer); + virtual BufferVar VisitBufferUse(const BufferVar& buffer); // statement visitor Stmt VisitStmt_(const BindNode* op) override; Stmt VisitStmt_(const AttrStmtNode* op) override; @@ -360,6 +360,7 @@ class TVM_DLL StmtExprMutator : public ExprMutator, public StmtMutator { using StmtMutator::VisitStmt; Expr VisitExpr(const Expr& e) override { return ExprMutator::VisitExpr(e); } + Expr VisitExpr_(const VarNode* op) override; Expr VisitExpr_(const BufferLoadNode* op) override; }; diff --git a/include/tvm/tirx/tirx_op.h b/include/tvm/tirx/tirx_op.h index ad2ec8e80fee..ef485470468e 100644 --- a/include/tvm/tirx/tirx_op.h +++ b/include/tvm/tirx/tirx_op.h @@ -51,7 +51,7 @@ constexpr const char* kDeviceInitStmt = "device_init_stmt"; */ constexpr const char* kHostInitStmt = "host_init_stmt"; /*! \brief Statements to be inserted after a specific buffer's definition (DeclBuffer/AllocBuffer). - * Stored as Map>. + * Stored as Map>. */ constexpr const char* kPostBufferDefStmt = "post_buffer_def_stmt"; } // namespace callback @@ -106,13 +106,13 @@ class DispatchContextNode : public ffi::Object { } /*! \brief Add a buffer to be allocated in the kernel. */ - void AddAllocBuffer(Buffer buffer); + void AddAllocBuffer(BufferVar buffer); /*! \brief Add an initialization statement to be inserted. */ void AddInitStmt(Stmt stmt, bool host = false); /*! \brief Add a statement to be inserted after a buffer's definition. */ - void AddPostBufferDefStmt(Buffer buffer, Stmt stmt); + void AddPostBufferDefStmt(BufferVar buffer, Stmt stmt); /*! \brief Set a value in the shared state cache. */ void SharedStateSet(ffi::String key, ffi::ObjectRef value); @@ -171,7 +171,7 @@ TVM_DLL const Op& fill(); /*! * \brief See pesudo code below: * - * Tx.gemm(Buffer A, Buffer B, Buffer C, Buffer D, PrimExpr alpha, PrimExpr beta) + * Tx.gemm(BufferVar A, BufferVar B, BufferVar C, BufferVar D, PrimExpr alpha, PrimExpr beta) */ TVM_DLL const Op& gemm(); diff --git a/include/tvm/tirx/tirx_stmt.h b/include/tvm/tirx/tirx_stmt.h index 41618eb1d96c..25dd48528428 100644 --- a/include/tvm/tirx/tirx_stmt.h +++ b/include/tvm/tirx/tirx_stmt.h @@ -35,7 +35,7 @@ namespace tirx { class TilePrimitiveCallNode : public StmtNode { public: TilePrimitiveCallNode(tvm::Op op, ffi::Array args, - ffi::Map workspace, + ffi::Map workspace, ffi::Map config, ffi::Optional dispatch, ExecScope scope) : op(std::move(op)), @@ -52,7 +52,7 @@ class TilePrimitiveCallNode : public StmtNode { ffi::Array args; // Workspace (pre-allocated buffers) for the operator. - ffi::Map workspace; + ffi::Map workspace; // Config for the operator/scheduler. ffi::Map config; @@ -84,7 +84,7 @@ class TilePrimitiveCallNode : public StmtNode { class TilePrimitiveCall : public Stmt { public: TVM_DLL TilePrimitiveCall(tvm::Op op, ffi::Array args, - ffi::Map workspace = {}, + ffi::Map workspace = {}, ffi::Map config = {}, ffi::Optional dispatch = std::nullopt, ExecScope scope = ExecScope(ScopeKind::kThread)); diff --git a/include/tvm/topi/contrib/cublas.h b/include/tvm/topi/contrib/cublas.h index 7d3f5f5b29d3..82d38eea8e36 100644 --- a/include/tvm/topi/contrib/cublas.h +++ b/include/tvm/topi/contrib/cublas.h @@ -49,7 +49,7 @@ inline Tensor cublas_matmul(const Tensor& lhs, const Tensor& rhs, bool transa, b return make_extern( {{n, m}}, {lhs->GetDataType()}, {lhs, rhs}, - [&](ffi::Array ins, ffi::Array outs) { + [&](ffi::Array ins, ffi::Array outs) { return call_packed({StringImm("tvm.contrib.cublas.matmul"), pack_buffer(ins[0]), pack_buffer(ins[1]), pack_buffer(outs[0]), IntImm::Int32(transa), IntImm::Int32(transb)}); @@ -75,7 +75,7 @@ inline Tensor cublas_batch_matmul(const Tensor& lhs, const Tensor& rhs, bool tra return make_extern( {{b, n, m}}, {lhs->GetDataType()}, {lhs, rhs}, - [&](ffi::Array ins, ffi::Array outs) { + [&](ffi::Array ins, ffi::Array outs) { return call_packed({StringImm("tvm.contrib.cublas.batch_matmul"), pack_buffer(ins[0]), pack_buffer(ins[1]), pack_buffer(outs[0]), IntImm::Int32(transa), IntImm::Int32(transb)}); diff --git a/include/tvm/topi/detail/extern.h b/include/tvm/topi/detail/extern.h index 00e9125ac686..1442e28e1e5f 100644 --- a/include/tvm/topi/detail/extern.h +++ b/include/tvm/topi/detail/extern.h @@ -42,7 +42,7 @@ using namespace tvm::te; * function. The function expects two arguments: an array of Buffers holding the input * tensor values, and a pre-allocated array of Buffers to be filled with the outputs. */ -using FExtern = std::function, ffi::Array)>; +using FExtern = std::function, ffi::Array)>; /*! * \brief Create tensors representing the result of invoking an external function. @@ -69,11 +69,11 @@ inline ffi::Array make_extern(const ffi::Array>& ou TVM_FFI_ICHECK_EQ(out_shapes.size(), out_types.size()) << "make_extern: out_shapes and out_types must have equal size"; - ffi::Array input_placeholders; + ffi::Array input_placeholders; for (auto t : inputs) { input_placeholders.push_back(tvm::tirx::decl_buffer(t->shape, t->dtype, t->op->name)); } - ffi::Array output_placeholders; + ffi::Array output_placeholders; for (size_t i = 0; i < out_shapes.size(); ++i) { output_placeholders.push_back(tvm::tirx::decl_buffer(out_shapes[i], out_types[i], name)); } @@ -98,7 +98,7 @@ inline ffi::Array make_extern(const ffi::Array>& ou * * \return An expression representing the pack operation */ -inline Expr pack_buffer(Buffer buf) { +inline Expr pack_buffer(BufferVar buf) { TVM_FFI_ICHECK_GT(buf->shape.size(), 0) << "buf shape must have at least one element"; Expr shape = Call(PointerType(PrimType::Int(64)), tvm::tirx::builtin::tvm_stack_make_shape(), buf->shape); diff --git a/python/tvm/s_tir/schedule/schedule.py b/python/tvm/s_tir/schedule/schedule.py index ecad5f1cb578..a53b9eb655b4 100644 --- a/python/tvm/s_tir/schedule/schedule.py +++ b/python/tvm/s_tir/schedule/schedule.py @@ -25,7 +25,7 @@ from tvm.error import register_error from tvm.ir import Expr, GlobalVar, IRModule, is_prim_expr from tvm.runtime import DataTypeCode, Object -from tvm.tirx import Buffer, FloatImm, For, IntImm, PrimFunc, SBlock +from tvm.tirx import Buffer, FloatImm, For, IntImm, PrimFunc, SBlock, is_buffer from tvm.tirx.function import IndexMap from . import _ffi_api @@ -3272,7 +3272,7 @@ def iter_buffers(): ) buffer_obj, (buffer_index_type, buffer_index) = next(iter(possible_buffers.items())) - elif isinstance(buffer, Buffer): + elif is_buffer(buffer): # Buffer lookup has unique id, can break out early found = False for buffer_index_type, buffer_index, buffer_obj in iter_buffers(): diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py index 8d530465e579..b7dcf734bd65 100644 --- a/python/tvm/tirx/__init__.py +++ b/python/tvm/tirx/__init__.py @@ -26,7 +26,7 @@ from tvm.ir import Expr from tvm.runtime import const -from .buffer import Buffer, decl_buffer, DataProducer +from .buffer import Buffer, BufferType, DataProducer, decl_buffer, is_buffer from .expr import convert from .expr import Var, Reduce, FloatImm, IntImm, StringImm, Cast from .expr import Add, Sub, Mul, Div, Mod, FloorDiv, FloorMod diff --git a/python/tvm/tirx/buffer.py b/python/tvm/tirx/buffer.py index b22688a09d06..a0e51bdef061 100644 --- a/python/tvm/tirx/buffer.py +++ b/python/tvm/tirx/buffer.py @@ -22,14 +22,34 @@ import tvm_ffi import tvm -from tvm.ir import PointerType, PrimType, Range +from tvm.ir import PointerType, PrimType, Range, Type from tvm.runtime import Object, Scriptable, convert from . import _ffi_api -@tvm_ffi.register_object("tirx.Buffer") -class Buffer(Object, Scriptable): +@tvm_ffi.register_object("tirx.BufferType") +class BufferType(Type): + """The structural type carried by an ordinary buffer variable.""" + + data_pointer_type: PointerType + dtype: PrimType + shape: list + strides: list + elem_offset: tvm.ir.Expr + data_alignment: int + offset_factor: int + layout: object | None + allocated_addr: list + + +def is_buffer(value) -> bool: + """Return whether ``value`` is an ordinary Var carrying BufferType.""" + + return isinstance(value, tvm.ir.Var) and isinstance(value.ty, BufferType) + + +class _BufferMethods: """Symbolic data buffer in TVM. Buffer provide a way to represent data layout @@ -193,10 +213,6 @@ def with_dtype(self, dtype): """Return a new buffer with the dtype.""" return _ffi_api.BufferWithDtype(self, dtype) # type: ignore - def with_data(self, data): - """Return a new buffer with the data.""" - return _ffi_api.BufferWithData(self, data) # type: ignore - def offset_of(self, indices): """Determine the offset of the provided indices in the flattened buffer. @@ -461,6 +477,9 @@ def permute(self, *dims) -> "Buffer": ) def __getitem__(self, indices): + if not is_buffer(self): + return _ORIGINAL_VAR_GETITEM(self, indices) + from ..arith import Analyzer # pylint: disable=import-outside-toplevel from .expr import BufferLoad, Ramp # pylint: disable=import-outside-toplevel from .stmt import BufferRegion # pylint: disable=import-outside-toplevel @@ -543,7 +562,7 @@ def decl_buffer( storage_type = dtype if isinstance(dtype, PrimType) else PrimType(dtype) storage_type = PrimType("int8") if storage_type.dtype == "bool" else storage_type data = Var(name, PointerType(storage_type, scope), span) - return _ffi_api.Buffer( # type: ignore + return _ffi_api.BufferVar( # type: ignore data, dtype, shape, @@ -557,6 +576,50 @@ def decl_buffer( ) +def _type_field(name): + def getter(self): + if not is_buffer(self): + raise AttributeError(f"{self.name} is not a Var with BufferType") + return getattr(self.ty, name) + + return property(getter) + + +# Buffer values intentionally retain runtime type key ``ir.Var``. Install the +# checked, type-directed convenience surface on that ordinary Python wrapper. +_ORIGINAL_VAR_GETITEM = tvm.ir.Var.__getitem__ +for _name, _value in _BufferMethods.__dict__.items(): + if _name.startswith("__") and _name != "__getitem__": + continue + if callable(_value) or isinstance(_value, property): + setattr(tvm.ir.Var, _name, _value) + +for _name in ( + "data_pointer_type", + "dtype", + "shape", + "strides", + "elem_offset", + "data_alignment", + "offset_factor", + "layout", + "allocated_addr", +): + setattr(tvm.ir.Var, _name, _type_field(_name)) + +tvm.ir.Var.data = property( + lambda self: _ffi_api.BufferData(self) + if is_buffer(self) + else (_ for _ in ()).throw(AttributeError(f"{self.name} is not a Var with BufferType")) +) +tvm.ir.Var.READ = _BufferMethods.READ +tvm.ir.Var.WRITE = _BufferMethods.WRITE + +# Source compatibility for annotations and imports only. There is no +# ``tirx.Buffer`` runtime object; constructors return ``tvm.ir.Var``. +Buffer = tvm.ir.Var + + @tvm_ffi.register_object("tirx.DataProducer") class DataProducer(Object): pass diff --git a/python/tvm/tirx/function.py b/python/tvm/tirx/function.py index 80b8fcaab0a6..5fdc2c3792dd 100644 --- a/python/tvm/tirx/function.py +++ b/python/tvm/tirx/function.py @@ -31,7 +31,7 @@ from ..runtime._tensor import Tensor from . import _ffi_api -from .buffer import Buffer +from .buffer import Buffer, is_buffer from .expr import Expr, Var @@ -72,7 +72,7 @@ def __init__(self, params, body, ret_type=None, buffer_map=None, attrs=None, spa buffer_map = {} if buffer_map is None else buffer_map for x in params: x = tvm.runtime.convert(x) if not isinstance(x, Object) else x - if isinstance(x, Buffer): + if is_buffer(x): var = Var(x.name, ty="handle") param_list.append(var) buffer_map[var] = x diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py index 782fbcfe58be..481b717c8f31 100644 --- a/python/tvm/tirx/op.py +++ b/python/tvm/tirx/op.py @@ -30,7 +30,7 @@ from tvm.runtime import const from . import _ffi_api -from .buffer import Buffer +from .buffer import Buffer, is_buffer from .expr import BufferLoad, CommReducer, ExprOp, ExprWithOp, IntImm, Var tir = tirx # alias for backward compat with upstream tir.convert() calls @@ -129,7 +129,7 @@ def call_packed_lowered(*args, span=None): -------- te.extern : Create tensor with extern function call. """ - call_args = [_pack_buffer(x) if isinstance(x, Buffer) else x for x in args] + call_args = [_pack_buffer(x) if is_buffer(x) else x for x in args] return Call(Op.get("tirx.tvm_call_packed_lowered"), call_args, span=span, ret_ty="int32") @@ -155,7 +155,7 @@ def call_cpacked_lowered(*args, span=None): -------- te.extern : Create tensor with extern function call. """ - call_args = [_pack_buffer(x) if isinstance(x, Buffer) else x for x in args] + call_args = [_pack_buffer(x) if is_buffer(x) else x for x in args] return Call(Op.get("tirx.tvm_call_cpacked_lowered"), call_args, span=span, ret_ty="int32") @@ -186,7 +186,7 @@ def call_packed(*args, span=None): -------- te.extern : Create tensor with extern function call. """ - call_args = [_pack_buffer(x) if isinstance(x, Buffer) else x for x in args] + call_args = [_pack_buffer(x) if is_buffer(x) else x for x in args] return Call(Op.get("tirx.tvm_call_packed"), call_args, span=span, ret_ty="int32") @@ -213,7 +213,7 @@ def call_cpacked(*args, span=None): -------- te.extern : Create tensor with extern function call. """ - call_args = [_pack_buffer(x) if isinstance(x, Buffer) else x for x in args] + call_args = [_pack_buffer(x) if is_buffer(x) else x for x in args] return Call(Op.get("tirx.tvm_call_cpacked"), call_args, span=span, ret_ty="int32") @@ -674,7 +674,7 @@ def address_of(obj: Buffer | BufferLoad | Var, span: Span | None = None) -> Expr call : Expr The call expression. """ - if isinstance(obj, Buffer): + if is_buffer(obj): n_dim = len(obj.shape) buffer_load = BufferLoad(obj, [0] * n_dim) return Call("tirx.address_of", [buffer_load], span=span, ret_ty=obj.data.ty) @@ -1267,7 +1267,7 @@ def trace(args, trace_action="tvm.default_trace_action"): """ if not isinstance(args, list): raise Exception("tvm.tirx.trace consumes the args as list type") - call_args = [_pack_buffer(x) if isinstance(x, Buffer) else x for x in args] + call_args = [_pack_buffer(x) if is_buffer(x) else x for x in args] call_args.insert(0, tvm.tirx.StringImm(trace_action)) tracing_value = args[-1] ret_ty = tracing_value.ty if isinstance(tracing_value, Expr) else tracing_value.dtype diff --git a/python/tvm/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py index 43a383a1d246..ee93a6404c15 100644 --- a/python/tvm/tirx/script/builder/ir.py +++ b/python/tvm/tirx/script/builder/ir.py @@ -44,7 +44,7 @@ # pylint: disable=unused-import from tvm.target.codegen import llvm_lookup_intrinsic_id -from tvm.tirx import Buffer, BufferRegion, Expr, IndexMap, type_annotation +from tvm.tirx import Buffer, BufferRegion, Expr, IndexMap, is_buffer, type_annotation from tvm.tirx import _ffi_api as _tirx_ffi_api from tvm.tirx import op as _tir_op from tvm.tirx.exec_scope import ExecScope, ScopeIdDef, Var @@ -1986,7 +1986,7 @@ def __invert__(self): def alloc_scalar(dtype: str = "float32", scope: str = "global") -> BufferLoad: """Allocate a zero-dimensional buffer (scalar).""" buf = alloc_buffer(shape=(1,), dtype=dtype, scope=scope, layout=TileLayout(S[1])) - assert isinstance(buf, Buffer) + assert is_buffer(buf) scalar = buf[0] if _current_meta_construction_scope() is not None: return scalar @@ -2006,7 +2006,7 @@ def decl_scalar(dtype, data, scope, elem_offset=None, byte_offset=None) -> Buffe offset_factor=0, layout=TileLayout(S[1]), ) - assert isinstance(buf, Buffer) + assert is_buffer(buf) scalar = buf[0] if _current_meta_construction_scope() is not None: return scalar @@ -2042,7 +2042,7 @@ def _meta_resource_for_value(value: Any) -> Any | None: return value.scalar.buffer if isinstance(value, BufferLoad): return value.buffer - if isinstance(value, Buffer): + if is_buffer(value): return value return None @@ -2913,19 +2913,19 @@ class WebGPUNamespace: @staticmethod def subgroup_shuffle(var, lane): - if isinstance(var, Buffer): + if is_buffer(var): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.webgpu.subgroup_shuffle", var, lane) @staticmethod def subgroup_shuffle_up(var, delta): - if isinstance(var, Buffer): + if is_buffer(var): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.webgpu.subgroup_shuffle_up", var, delta) @staticmethod def subgroup_shuffle_down(var, delta): - if isinstance(var, Buffer): + if is_buffer(var): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.webgpu.subgroup_shuffle_down", var, delta) diff --git a/python/tvm/tirx/script/builder/tirx.py b/python/tvm/tirx/script/builder/tirx.py index 6f2442ed82ef..9167918fca07 100644 --- a/python/tvm/tirx/script/builder/tirx.py +++ b/python/tvm/tirx/script/builder/tirx.py @@ -21,7 +21,7 @@ import tvm.tirx.operator as tirx_op from tvm.ir import Op -from tvm.tirx import Buffer, BufferRegion, Expr +from tvm.tirx import Buffer, BufferRegion, Expr, is_buffer from tvm.tirx.exec_scope import _SCOPE_KIND_TO_NAME, ExecScope from tvm.tirx.expr import FloatImm from tvm.tirx.lang.alloc_pool import SMEMPool, TMEMPool, TMEMStages @@ -123,11 +123,11 @@ def __getattr__(self, name: str): def _is_buffer_or_region(x): - return isinstance(x, Buffer | BufferRegion) + return is_buffer(x) or isinstance(x, BufferRegion) def _to_region(buffer: BufferRegion | Buffer): - if isinstance(buffer, Buffer): + if is_buffer(buffer): return buffer[[slice(None, None, None) for _ in range(len(buffer.shape))]] assert isinstance(buffer, BufferRegion) return buffer @@ -222,7 +222,7 @@ def sqrt( config = kwargs or {} dst = _to_region(dst) src = _to_region(src) - if bias is not None and isinstance(bias, Buffer): + if bias is not None and is_buffer(bias): bias = _to_region(bias) return f_insert( tirx_op.Sqrt( @@ -268,9 +268,9 @@ def add( workspace = {} config = kwargs or {} dst = _to_region(dst) - if isinstance(src1, Buffer): + if is_buffer(src1): src1 = _to_region(src1) - if isinstance(src2, Buffer): + if is_buffer(src2): src2 = _to_region(src2) return f_insert( tirx_op.Add( @@ -309,9 +309,9 @@ def sub( workspace = {} config = kwargs or {} dst = _to_region(dst) - if isinstance(src1, Buffer): + if is_buffer(src1): src1 = _to_region(src1) - if isinstance(src2, Buffer): + if is_buffer(src2): src2 = _to_region(src2) return f_insert( tirx_op.Sub( @@ -350,9 +350,9 @@ def mul( workspace = {} config = kwargs or {} dst = _to_region(dst) - if isinstance(src1, Buffer): + if is_buffer(src1): src1 = _to_region(src1) - if isinstance(src2, Buffer): + if is_buffer(src2): src2 = _to_region(src2) return f_insert( tirx_op.Mul( @@ -392,7 +392,7 @@ def fdiv( config = kwargs or {} dst = _to_region(dst) src1 = _to_region(src1) - if isinstance(src2, Buffer): + if is_buffer(src2): src2 = _to_region(src2) return f_insert( tirx_op.FDiv( @@ -436,9 +436,9 @@ def fma( config = kwargs or {} dst = _to_region(dst) src = _to_region(src) - if isinstance(scale, Buffer): + if is_buffer(scale): scale = _to_region(scale) - if isinstance(bias, Buffer): + if is_buffer(bias): bias = _to_region(bias) return f_insert( tirx_op.FMA( @@ -805,7 +805,7 @@ def max( """ from tvm import tirx as _tirx - if not isinstance(dst, BufferRegion | Buffer) or not isinstance(src, BufferRegion | Buffer): + if not _is_buffer_or_region(dst) or not _is_buffer_or_region(src): # Expression-level max return _tirx.max(dst, src) if workspace is None: @@ -846,7 +846,7 @@ def min( """ from tvm import tirx as _tirx - if not isinstance(dst, BufferRegion | Buffer) or not isinstance(src, BufferRegion | Buffer): + if not _is_buffer_or_region(dst) or not _is_buffer_or_region(src): return _tirx.min(dst, src) if workspace is None: workspace = {} @@ -1010,9 +1010,9 @@ def maximum( workspace = {} config = kwargs or {} dst = _to_region(dst) - if isinstance(src1, Buffer): + if is_buffer(src1): src1 = _to_region(src1) - if isinstance(src2, Buffer): + if is_buffer(src2): src2 = _to_region(src2) return f_insert( tirx_op.Maximum( @@ -1051,9 +1051,9 @@ def minimum( workspace = {} config = kwargs or {} dst = _to_region(dst) - if isinstance(src1, Buffer): + if is_buffer(src1): src1 = _to_region(src1) - if isinstance(src2, Buffer): + if is_buffer(src2): src2 = _to_region(src2) return f_insert( tirx_op.Minimum( @@ -1105,7 +1105,7 @@ def exp( config = kwargs or {} dst = _to_region(dst) src = _to_region(src) - if bias is not None and isinstance(bias, Buffer): + if bias is not None and is_buffer(bias): bias = _to_region(bias) return f_insert( tirx_op.Exp( @@ -1164,7 +1164,7 @@ def exp2( config = kwargs or {} dst = _to_region(dst) src = _to_region(src) - if bias is not None and isinstance(bias, Buffer): + if bias is not None and is_buffer(bias): bias = _to_region(bias) return f_insert( tirx_op.Exp2( @@ -1250,9 +1250,9 @@ def binary_reduce( workspace = {} binary_output = _to_region(binary_output) reduce_output = _to_region(reduce_output) - if isinstance(binary_input1, Buffer): + if is_buffer(binary_input1): binary_input1 = _to_region(binary_input1) - if isinstance(binary_input2, Buffer): + if is_buffer(binary_input2): binary_input2 = _to_region(binary_input2) reduce_axes = _wrap_elem_in_tuple(reduce_axes) @@ -1334,7 +1334,7 @@ def unary_reduce( reduce_output = _to_region(reduce_output) unary_input = _to_region(unary_input) - if bias is not None and isinstance(bias, Buffer): + if bias is not None and is_buffer(bias): bias = _to_region(bias) reduce_axes = _wrap_elem_in_tuple(reduce_axes) @@ -1418,9 +1418,9 @@ def binary_chain( output = _to_region(output) data = _to_region(data) - if isinstance(operand0, Buffer): + if is_buffer(operand0): operand0 = _to_region(operand0) - if isinstance(operand1, Buffer): + if is_buffer(operand1): operand1 = _to_region(operand1) if isinstance(op0, str): @@ -1533,9 +1533,9 @@ def select( The predicate to evaluate. The callable should take the same number of arguments as the dimensions of the destination buffer. """ # noqa: E501 dst = _to_region(dst) - if isinstance(true_value, Buffer): + if is_buffer(true_value): true_value = _to_region(true_value) - if isinstance(false_value, Buffer): + if is_buffer(false_value): false_value = _to_region(false_value) if not isinstance(pred, Predicate): pred = Predicate(pred) diff --git a/python/tvm/tirx/script/parser/parser.py b/python/tvm/tirx/script/parser/parser.py index ddbfac29ab18..8026f07b9efc 100644 --- a/python/tvm/tirx/script/parser/parser.py +++ b/python/tvm/tirx/script/parser/parser.py @@ -29,7 +29,7 @@ from tvm.script.ir_builder.base import IRBuilderFrame as Frame from tvm.script.parser._core import Parser, dispatch, doc from tvm.script.parser.core.doc import from_doc -from tvm.tirx import Buffer, IterVar, Layout +from tvm.tirx import Buffer, IterVar, Layout, is_buffer from tvm.tirx.script import builder as T from tvm.tirx.script.builder.ir import name_meta_class_value from tvm.tirx.stmt import BufferRegion @@ -119,7 +119,7 @@ def bind_with_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> for i, v in enumerate(value): bind_with_value(self, node, f"{var_name}_{i}", v) return value - elif isinstance(value, Buffer | tvm.ir.Var): + elif isinstance(value, tvm.ir.Var): IRBuilder.name(var_name, value) return value else: @@ -211,7 +211,7 @@ def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) - res = value.__enter__() IRBuilder.name(var_name, res) return res - elif isinstance(value, Buffer | IterVar | Layout) or ( + elif is_buffer(value) or isinstance(value, IterVar | Layout) or ( isinstance(value, tvm.ir.Var) and not self.var_table.exist(value) ): IRBuilder.name(var_name, value) @@ -407,7 +407,7 @@ def visit_assign(self: Parser, node: doc.Assign) -> None: # that genuine errors (e.g. wrong shape, bad store) are not swallowed. # Only TypeError from FFI type mismatch (e.g. rhs is a meta_var, not # a Expr or auto-convertible scalar) triggers fallthrough. - if isinstance(lhs_value, T.scalar_wrapper | T.BufferLoad | tvm.tirx.Buffer): + if isinstance(lhs_value, T.scalar_wrapper | T.BufferLoad) or is_buffer(lhs_value): if isinstance(lhs_value, T.scalar_wrapper): buffer = lhs_value.scalar.buffer else: @@ -487,7 +487,7 @@ def visit_aug_assign(self: Parser, node: doc.AugAssign) -> None: lhs_value = self.eval_expr(lhs_copy) except Exception: # pylint: disable=broad-except pass - if isinstance(lhs_value, T.scalar_wrapper | T.BufferLoad | tvm.tirx.Buffer): + if isinstance(lhs_value, T.scalar_wrapper | T.BufferLoad) or is_buffer(lhs_value): if isinstance(lhs_value, T.scalar_wrapper): buffer = lhs_value.scalar.buffer else: @@ -760,7 +760,7 @@ def visit_expr_stmt(self: Parser, node: doc.Expr) -> None: pass elif isinstance(res, tvm.tirx.stmt.BufferStore): T.buffer_store(res.buffer, res.value, res.indices, res.predicate) - elif isinstance(res, tvm.tirx.Buffer): + elif is_buffer(res): # ``T.match_buffer(...)`` used as a bare statement (no LHS) — the # buffer object is discarded; the underlying side effect (the # match_buffer node) has already been emitted into the frame. diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index 7270e04c8b3b..c723964d747f 100644 --- a/python/tvm/tirx/stmt.py +++ b/python/tvm/tirx/stmt.py @@ -66,7 +66,7 @@ def _normalize_legacy_stmt(stmt: Stmt | None) -> Stmt | None: cur = stmt while True: if isinstance(cur, DeclBuffer) and hasattr(cur, "body"): - prefix.append(DeclBuffer(cur.buffer, cur.span)) + prefix.append(DeclBuffer(cur.buffer, data=cur.data, span=cur.span)) cur = cur.body continue if isinstance(cur, AllocBuffer) and hasattr(cur, "body"): @@ -435,10 +435,12 @@ class DeclBuffer(Stmt): """ buffer: Buffer + data: Expr | None span: Span | None def __init__(self, buffer: Buffer, *args, **kwargs) -> None: body: Stmt | None = None + data: Expr | None = kwargs.pop("data", None) span: Span | None = None if len(args) == 1: @@ -479,7 +481,7 @@ def __init__(self, buffer: Buffer, *args, **kwargs) -> None: raise TypeError("DeclBuffer span specified by both args and kwargs") span = kw_span if kw_span is not None else span - self.__init_handle_by_constructor__(_ffi_api.DeclBuffer, buffer, span) + self.__init_handle_by_constructor__(_ffi_api.DeclBuffer, buffer, data, span) # Legacy compatibility. Body is carried on python side only. if body is not None: self.body = body diff --git a/python/tvm/tirx/transform/common.py b/python/tvm/tirx/transform/common.py index 4e0ccd848f26..513f1417f0b9 100644 --- a/python/tvm/tirx/transform/common.py +++ b/python/tvm/tirx/transform/common.py @@ -38,7 +38,8 @@ class BufferReplacer(StmtExprMutator): """ Replace buffer with another buffer. - Also replace the data of the buffer with another var. + Buffer values are ordinary Vars, so the same mapping also rewrites + ``buffer_data`` projections. """ def __init__( @@ -49,7 +50,7 @@ def __init__( self.var_map = var_map if var_map is not None else {} self.buffer_attr_var_mutated = False for old_buffer, new_buffer in self.buffer_map.items(): - self.var_map[old_buffer.data] = new_buffer.data + self.var_map[old_buffer] = new_buffer def mutate_buffer(self, buffer: Buffer): if buffer in self.buffer_map: @@ -59,7 +60,6 @@ def mutate_buffer(self, buffer: Buffer): # unrelated buffers can be spuriously cloned and introduce alias buffers. prev_mutated = self.buffer_attr_var_mutated self.buffer_attr_var_mutated = False - new_data = self.visit_expr(buffer.data) new_shape = [self.visit_expr(expr) for expr in buffer.shape] new_strides = [self.visit_expr(expr) for expr in buffer.strides] new_elem_offset = ( @@ -83,6 +83,7 @@ def mutate_buffer(self, buffer: Buffer): ) else: new_layout = buffer.layout + new_allocated_addr = [self.visit_expr(expr) for expr in buffer.allocated_addr] buffer_attr_mutated = self.buffer_attr_var_mutated self.buffer_attr_var_mutated = prev_mutated or buffer_attr_mutated if not buffer_attr_mutated: @@ -91,7 +92,7 @@ def mutate_buffer(self, buffer: Buffer): new_shape, buffer.dtype, buffer.name, - new_data, + Var(buffer.name, buffer.data_pointer_type), new_strides, new_elem_offset, buffer.scope(), @@ -99,7 +100,10 @@ def mutate_buffer(self, buffer: Buffer): buffer.offset_factor, layout=new_layout, ) + if new_allocated_addr: + new_buffer = new_buffer.with_allocated_addr(new_allocated_addr) self.buffer_map[buffer] = new_buffer + self.var_map[buffer] = new_buffer return new_buffer def visit_var_(self, op: Var): @@ -134,7 +138,7 @@ def visit_decl_buffer_(self, op: DeclBuffer): new_buffer = self.mutate_buffer(op.buffer) op = super().visit_decl_buffer_(op) if new_buffer is not None: - return DeclBuffer(new_buffer, op.span) + return DeclBuffer(new_buffer, data=op.data, span=op.span) return op def visit_array_prim_expr_(self, op: list[Expr]): diff --git a/src/arith/domain_touched.cc b/src/arith/domain_touched.cc index 6701beee3d69..94a90be1098e 100644 --- a/src/arith/domain_touched.cc +++ b/src/arith/domain_touched.cc @@ -64,11 +64,11 @@ class BufferTouchedDomain final : public IRVisitorWithAnalyzer { public: BufferTouchedDomain(const Stmt& stmt) { operator()(stmt); } - std::unordered_map& GetAccessedBufferRegions() { + std::unordered_map& GetAccessedBufferRegions() { return buffer_access_map_; } - Region FindUnion(const Buffer& buffer, bool consider_loads, bool consider_stores) { + Region FindUnion(const BufferVar& buffer, bool consider_loads, bool consider_stores) { Region ret; auto kv = buffer_access_map_.find(buffer.get()); if (kv == buffer_access_map_.end()) { @@ -130,17 +130,17 @@ class BufferTouchedDomain final : public IRVisitorWithAnalyzer { } } - std::unordered_map buffer_access_map_; + std::unordered_map buffer_access_map_; }; -Region DomainTouched(const Stmt& stmt, const Buffer& buffer, bool consider_loads, +Region DomainTouched(const Stmt& stmt, const BufferVar& buffer, bool consider_loads, bool consider_stores) { return BufferTouchedDomain(stmt).FindUnion(buffer, consider_loads, consider_stores); } -ffi::Map> DomainTouchedAccessMap(const PrimFunc& func) { +ffi::Map> DomainTouchedAccessMap(const PrimFunc& func) { auto buffer_access_map = BufferTouchedDomain(func->body).GetAccessedBufferRegions(); - ffi::Map> ret; + ffi::Map> ret; auto& buffer_map = func->buffer_map; for (auto& var : func->params) { auto& buffer = buffer_map[var]; diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc index 6f951f97896f..ba48f736ef03 100644 --- a/src/backend/cuda/codegen/codegen_cuda.cc +++ b/src/backend/cuda/codegen/codegen_cuda.cc @@ -1317,7 +1317,7 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op, std::ostream& os) { std::string guard = this->PrintExpr(op->args[1]); const BufferLoadNode* addr_buffer = op->args[2].as(); std::string global_addr = this->PrintExpr(addr_buffer->indices[0]); - std::string global_buffer = this->PrintExpr(addr_buffer->buffer->data); + std::string global_buffer = this->PrintExpr(addr_buffer->buffer.data()); std::string local_addr = this->PrintExpr(op->args[3]); this->stream << "asm volatile (\n"; this->stream << "\"{.reg .pred p;\\n\"\n"; @@ -1632,11 +1632,11 @@ void CodeGenCUDA::VisitStmt_(const AttrStmtNode* op) { void CodeGenCUDA::VisitStmt_(const AllocBufferNode* op) { TVM_FFI_ICHECK(op->buffer.defined()); - std::string vid = AllocVarID(op->buffer->data.get()); + std::string vid = AllocVarID(op->buffer.get()); this->PrintIndent(); - std::string scope = GetPtrStorageScope(op->buffer->data); - const VarNode* buffer = op->buffer->data.get(); + std::string scope = op->buffer.scope(); + const VarNode* buffer = op->buffer.get(); PrimType dtype = op->buffer->dtype; if (scope.find("wmma.") == 0) { @@ -1696,9 +1696,9 @@ void CodeGenCUDA::VisitStmt_(const AllocBufferNode* op) { stream << ' ' << vid << '[' << constant_size << "];\n"; } - RegisterHandleType(op->buffer->data.get(), dtype); + RegisterHandleType(op->buffer.get(), dtype); if (op->annotations.count(tirx::attr::kVolatile)) { - MarkVolatile(op->buffer->data.get()); + MarkVolatile(op->buffer.get()); } } @@ -2064,7 +2064,7 @@ void CodeGenCUDA::HandleVolatileLoads(const std::string& value, const BufferLoad PrimType op_ty = op->ty.as_or_throw(); if ((op_ty.MatchesElementType(DLDataTypeCode::kDLFloat, 16) || op_ty.MatchesElementType(DLDataTypeCode::kDLBfloat, 16)) && - IsVolatile(op->buffer->data.get())) { + IsVolatile(op->buffer.get())) { os << "("; PrintType(op_ty, os); os << ")(" << value << ")"; diff --git a/src/backend/cuda/codegen/llvm/codegen_nvptx.cc b/src/backend/cuda/codegen/llvm/codegen_nvptx.cc index 1649630405f7..36fa93f5bd2c 100644 --- a/src/backend/cuda/codegen/llvm/codegen_nvptx.cc +++ b/src/backend/cuda/codegen/llvm/codegen_nvptx.cc @@ -80,13 +80,13 @@ class CodeGenNVPTX : public CodeGenLLVM { void VisitStmt_(const AllocBufferNode* op) final { llvm::Value* buf = nullptr; - StorageInfo& info = alloc_storage_info_[op->buffer->data.get()]; + StorageInfo& info = alloc_storage_info_[op->buffer.get()]; // maximum necessary alignment in the NV devices if (info.alignment > 16) { info.alignment = 16; } - auto storage_scope = runtime::StorageScope::Create(GetPtrStorageScope(op->buffer->data)); + auto storage_scope = runtime::StorageScope::Create(GetPtrStorageScope(op->buffer.var())); PrimType dtype = op->buffer->dtype; if (storage_scope.rank == runtime::StorageRank::kShared && storage_scope.tag == ".dyn") { @@ -122,10 +122,10 @@ class CodeGenNVPTX : public CodeGenLLVM { buf = builder_->CreatePointerCast( buf, llvmGetPointerTo(DTypeToLLVMType(dtype), buf->getType()->getPointerAddressSpace())); - TVM_FFI_ICHECK(!var_map_.count(op->buffer->data.get())); - var_map_[op->buffer->data.get()] = buf; + TVM_FFI_ICHECK(!var_map_.count(op->buffer.get())); + var_map_[op->buffer.get()] = buf; if (op->annotations.count(tirx::attr::kVolatile)) { - volatile_buf_.insert(op->buffer->data.get()); + volatile_buf_.insert(op->buffer.get()); } } diff --git a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc index 665067ea3a53..2bd1a505c94b 100644 --- a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc +++ b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc @@ -107,7 +107,7 @@ class CodeGenHexagon final : public CodeGenCPU { bool IsQHLFunction(const std::string& func); - llvm::Value* VectorLookupLoad(Buffer buffer, PrimType buffer_type, ffi::Array indices); + llvm::Value* VectorLookupLoad(BufferVar buffer, PrimType buffer_type, ffi::Array indices); llvm::Value* Intrinsic(llvm::Intrinsic::ID, llvm::ArrayRef args); std::vector fqhl_list_ = { "tvm_vect_qhmath_hvx_cos_ahf", "tvm_vect_qhmath_hvx_tanh_ahf", @@ -199,12 +199,10 @@ llvm::Value* CodeGenHexagon::CreateCallExtern(Type ret_type, ffi::String global_ } llvm::Value* CodeGenHexagon::VisitExpr_(const BufferLoadNode* op) { - if (!op->buffer.same_as(op->buffer->data)) { - // Check if we can generate a vector lookup. - if (!op->indices[0].as()) { - if (auto* vlut = VectorLookupLoad(op->buffer, op->ty.as_or_throw(), op->indices)) { - return vlut; - } + // Check if we can generate a vector lookup. + if (!op->indices[0].as()) { + if (auto* vlut = VectorLookupLoad(op->buffer, op->ty.as_or_throw(), op->indices)) { + return vlut; } } return CodeGenCPU::VisitExpr_(op); @@ -327,7 +325,7 @@ llvm::Value* CodeGenHexagon::Intrinsic(llvm::Intrinsic::ID IntID, return builder_->CreateCall(intf_callee, conv_args); } -llvm::Value* CodeGenHexagon::VectorLookupLoad(Buffer buffer, PrimType buffer_type, +llvm::Value* CodeGenHexagon::VectorLookupLoad(BufferVar buffer, PrimType buffer_type, ffi::Array indices) { PrimExpr index = indices[0]; PrimType index_ty = index.ty(); diff --git a/src/backend/metal/codegen/codegen_metal.cc b/src/backend/metal/codegen/codegen_metal.cc index 1469423f0951..c4f180feee5a 100644 --- a/src/backend/metal/codegen/codegen_metal.cc +++ b/src/backend/metal/codegen/codegen_metal.cc @@ -310,7 +310,7 @@ void CodeGenMetal::PrintStorageScope(const std::string& scope, std::ostream& os) void CodeGenMetal::VisitStmt_(const AllocBufferNode* op) { TVM_FFI_ICHECK(op->buffer.defined()); - std::string vid = AllocVarID(op->buffer->data.get()); + std::string vid = AllocVarID(op->buffer.get()); this->PrintIndent(); // Compute constant_size from buffer shape @@ -322,8 +322,8 @@ void CodeGenMetal::VisitStmt_(const AllocBufferNode* op) { } TVM_FFI_ICHECK_GT(constant_size, 0) << "Can only handle constant size stack allocation for now"; - auto scope = GetPtrStorageScope(op->buffer->data); - alloc_storage_scope_[op->buffer->data.get()] = scope; + auto scope = op->buffer.scope(); + alloc_storage_scope_[op->buffer.get()] = scope; const PrimType& dtype = op->buffer->dtype; if (scope == "metal.simdgroup") { bool supported_simdgroup_dtype = dtype == PrimType::Float(16) || dtype == PrimType::Float(32) || @@ -337,7 +337,7 @@ void CodeGenMetal::VisitStmt_(const AllocBufferNode* op) { std::ostringstream dtype_os; PrintType(dtype, dtype_os); std::string dtype_str = dtype_os.str(); - simdgroup_dtype_[op->buffer->data.get()] = dtype_str; + simdgroup_dtype_[op->buffer.get()] = dtype_str; stream << "simdgroup_" << dtype_str << "8x8 " << vid << '[' << constant_size / 64 << "];\n"; } else { PrintStorageScope(scope, stream); @@ -345,9 +345,9 @@ void CodeGenMetal::VisitStmt_(const AllocBufferNode* op) { stream << ' ' << vid << '[' << constant_size << "];\n"; } - RegisterHandleType(op->buffer->data.get(), op->buffer->dtype); + RegisterHandleType(op->buffer.get(), op->buffer->dtype); if (op->annotations.count(tirx::attr::kVolatile)) { - MarkVolatile(op->buffer->data.get()); + MarkVolatile(op->buffer.get()); } } diff --git a/src/backend/opencl/codegen/codegen_opencl.cc b/src/backend/opencl/codegen/codegen_opencl.cc index d1d19a2fe413..8379c446172b 100644 --- a/src/backend/opencl/codegen/codegen_opencl.cc +++ b/src/backend/opencl/codegen/codegen_opencl.cc @@ -284,9 +284,9 @@ void CodeGenOpenCL::PrintType(const Type& type, std::ostream& os) { // NOLINT(* } } -void CodeGenOpenCL::PrintVecAddr(const BufferNode* buffer, const PrimType& t, PrimExpr base, +void CodeGenOpenCL::PrintVecAddr(const VarNode* buffer, const PrimType& t, PrimExpr base, std::ostream& os) { // NOLINT(*) - const VarNode* buffer_var = buffer->data.get(); + const VarNode* buffer_var = buffer; PrimType elem_type = t.WithLanes(1); if (!HandleTypeMatch(buffer_var, elem_type)) { os << '('; @@ -300,7 +300,7 @@ void CodeGenOpenCL::PrintVecAddr(const BufferNode* buffer, const PrimType& t, Pr os << GetVarID(buffer_var) << " + "; PrintExpr(base, os); } -std::string CodeGenOpenCL::GetVecLoad(const PrimType& t, const BufferNode* buffer, PrimExpr base) { +std::string CodeGenOpenCL::GetVecLoad(const PrimType& t, const VarNode* buffer, PrimExpr base) { std::ostringstream os; os << "vload" << t.lanes() << "(0, "; PrintVecAddr(buffer, t, base, os); @@ -308,7 +308,7 @@ std::string CodeGenOpenCL::GetVecLoad(const PrimType& t, const BufferNode* buffe return os.str(); } -void CodeGenOpenCL::PrintVecStore(const BufferNode* buffer, const PrimType& t, PrimExpr base, +void CodeGenOpenCL::PrintVecStore(const VarNode* buffer, const PrimType& t, PrimExpr base, const std::string& value) { this->PrintIndent(); stream << "vstore" << t.lanes() << "(" << value << ", 0, "; @@ -407,7 +407,7 @@ void CodeGenOpenCL::VisitStmt_(const AllocBufferNode* op) { TVM_FFI_ICHECK(dim_imm) << "Can only handle constant size stack allocation for now"; constant_size *= dim_imm->value; } - allocation_size_.insert({op->buffer->data.get(), constant_size * op->buffer->dtype.lanes()}); + allocation_size_.insert({op->buffer.get(), constant_size * op->buffer->dtype.lanes()}); CodeGenC::VisitStmt_(op); } @@ -419,12 +419,12 @@ void CodeGenOpenCL::VisitExpr_(const CallNode* op, std::ostream& os) { TVM_FFI_ICHECK_EQ(load->indices.size(), 1) << "CodeGenOpenCL only supports flat memory allocations."; os << "(("; - auto it = alloc_storage_scope_.find(load->buffer->data.get()); + auto it = alloc_storage_scope_.find(load->buffer.get()); if (it != alloc_storage_scope_.end()) { PrintStorageScope(it->second, os); } this->PrintType(load->ty.as_or_throw().WithLanes(1), os); - os << " *)" << this->GetVarID(load->buffer->data.get()) << " + "; + os << " *)" << this->GetVarID(load->buffer.get()) << " + "; this->PrintExpr(load->indices[0], os); os << ')'; } else if (op->op.same_as(builtin::texture2d_store())) { diff --git a/src/backend/opencl/codegen/codegen_opencl.h b/src/backend/opencl/codegen/codegen_opencl.h index a537a2ede698..38c33ecc2c1c 100644 --- a/src/backend/opencl/codegen/codegen_opencl.h +++ b/src/backend/opencl/codegen/codegen_opencl.h @@ -49,13 +49,13 @@ class CodeGenOpenCL final : public CodeGenC { using CodeGenC::PrintType; void PrintType(const PrimType& t, std::ostream& os) final; // NOLINT(*) void PrintType(const Type& type, std::ostream& os) final; // NOLINT(*) - std::string GetVecLoad(const PrimType& t, const BufferNode* buffer, PrimExpr base) final; - void PrintVecStore(const BufferNode* buffer, const PrimType& t, PrimExpr base, + std::string GetVecLoad(const PrimType& t, const VarNode* buffer, PrimExpr base) final; + void PrintVecStore(const VarNode* buffer, const PrimType& t, PrimExpr base, const std::string& value) final; // NOLINT(*) void PrintVecElemLoadExpr(const PrimType& t, int i, const std::string& value, std::ostream& os) final; // NOLINT(*) // the address of load/store - void PrintVecAddr(const BufferNode* buffer, const PrimType& t, PrimExpr base, + void PrintVecAddr(const VarNode* buffer, const PrimType& t, PrimExpr base, std::ostream& os); // NOLINT(*) void PrintRestrict(const Var& v, std::ostream& os) final; // NOLINT(*) std::string CastFromTo(std::string value, const PrimType& from, diff --git a/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc b/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc index 39a138e28f4b..277dec6a8d80 100644 --- a/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc +++ b/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc @@ -98,8 +98,8 @@ class CodeGenAMDGPU : public CodeGenLLVM { void VisitStmt_(const AllocBufferNode* op) final { llvm::Value* buf = nullptr; - StorageInfo& info = alloc_storage_info_[op->buffer->data.get()]; - auto storage_scope = runtime::StorageScope::Create(GetPtrStorageScope(op->buffer->data)); + StorageInfo& info = alloc_storage_info_[op->buffer.get()]; + auto storage_scope = runtime::StorageScope::Create(GetPtrStorageScope(op->buffer.var())); PrimType dtype = op->buffer->dtype; if (storage_scope.rank == runtime::StorageRank::kShared && storage_scope.tag == ".dyn") { @@ -140,10 +140,10 @@ class CodeGenAMDGPU : public CodeGenLLVM { buf = builder_->CreatePointerCast( buf, llvmGetPointerTo(DTypeToLLVMType(dtype), buf->getType()->getPointerAddressSpace())); - TVM_FFI_ICHECK(!var_map_.count(op->buffer->data.get())); - var_map_[op->buffer->data.get()] = buf; + TVM_FFI_ICHECK(!var_map_.count(op->buffer.get())); + var_map_[op->buffer.get()] = buf; if (op->annotations.count(tirx::attr::kVolatile)) { - volatile_buf_.insert(op->buffer->data.get()); + volatile_buf_.insert(op->buffer.get()); } } diff --git a/src/backend/trn/codegen/codegen_trn.cc b/src/backend/trn/codegen/codegen_trn.cc index f8d46370e9c7..466371ed9965 100644 --- a/src/backend/trn/codegen/codegen_trn.cc +++ b/src/backend/trn/codegen/codegen_trn.cc @@ -79,7 +79,7 @@ void CodeGenTrainium::AddFunction(const GlobalVar& gvar, const PrimFunc& func) { // We can switch to follow the flow with inter-function call process // after the Trainium function declaration is properly printed. // In Trainium, for PrimFuncs with signature - // def func(A: Buffer, B: Buffer, x: int, y: float) -> None + // def func(A: BufferVar, B: BufferVar, x: int, y: float) -> None // where there are trailing pod parameters, the codegen emits a struct // struct func_params{ x: int; y: float; } // for the function. In the flow of inter-function call process, @@ -103,7 +103,7 @@ void CodeGenTrainium::AddFunction(const GlobalVar& gvar, const PrimFunc& func) { // Function header. this->stream << "def " << static_cast(global_symbol.value()) << "("; - // Buffer arguments + // BufferVar arguments auto num_inputs = func->GetAttr(tvm::attr::kNumInputs); TVM_FFI_ICHECK(num_inputs.has_value()); std::vector output_vids; @@ -209,10 +209,10 @@ std::string CodeGenTrainium::GetStorageScopeStr(const std::string& scope) { // void CodeGenTrainium::VisitStmt_(const AllocBufferNode* op) { TVM_FFI_ICHECK(op->buffer.defined()); - std::string vid = AllocVarID(op->buffer->data.get()); + std::string vid = AllocVarID(op->buffer.get()); this->PrintIndent(); - auto scope = GetPtrStorageScope(op->buffer->data); + auto scope = op->buffer.scope(); std::ostringstream dtype_os; PrintType(op->buffer->dtype, dtype_os); std::string dtype_str = dtype_os.str(); @@ -229,7 +229,7 @@ void CodeGenTrainium::VisitStmt_(const AllocBufferNode* op) { if (auto allocated_addr = op->annotations.Get(tirx::attr::buffer_allocated_addr)) { addr = allocated_addr.value().as_or_throw>(); } else { - // AllocBuffer is a leaf stmt after rebase; in that path allocated_addr is carried by Buffer. + // AllocBuffer is a leaf stmt after rebase; in that path allocated_addr is carried by BufferVar. addr = op->buffer->allocated_addr; } if (addr.empty()) { @@ -342,7 +342,7 @@ void CodeGenTrainium::VisitExpr_(const BufferLoadNode* op, std::ostream& os) { if (buffer_idmap_.count(op->buffer)) { buffer_str = buffer_idmap_[op->buffer]; } else { - buffer_str = GetVarID(op->buffer->data.get()); + buffer_str = GetVarID(op->buffer.get()); } os << buffer_str << "["; os << PrintIndices(op->indices); @@ -607,10 +607,10 @@ void CodeGenTrainium::VisitStmt_(const DeclBufferNode* op) { if (op->buffer.scope() == "trn.psum" || op->buffer.scope() == "trn.sbuf") { return; } - const VarNode* data = op->buffer->data.get(); + const VarNode* data = op->buffer.get(); auto it = data_buffer_idmap_.find(data); if (it != data_buffer_idmap_.end()) { - const Buffer& prev_buffer = data_decl_buffer_map_.at(data); + const BufferVar& prev_buffer = data_decl_buffer_map_.at(data); if (ffi::StructuralEqual()(prev_buffer->shape, op->buffer->shape) && prev_buffer->dtype == op->buffer->dtype) { buffer_idmap_[op->buffer] = it->second; diff --git a/src/backend/trn/codegen/codegen_trn.h b/src/backend/trn/codegen/codegen_trn.h index bcfd9451b2cc..63d6712926d1 100644 --- a/src/backend/trn/codegen/codegen_trn.h +++ b/src/backend/trn/codegen/codegen_trn.h @@ -80,9 +80,9 @@ class CodeGenTrainium final : public CodeGenC { Target target_; NKIInstructionCtx ctx_; std::unordered_map opcode_map_; - std::unordered_map buffer_idmap_; + std::unordered_map buffer_idmap_; std::unordered_map data_buffer_idmap_; - std::unordered_map data_decl_buffer_map_; + std::unordered_map data_decl_buffer_map_; bool is_outermost_loop_ = true; }; } // namespace codegen diff --git a/src/backend/trn/transform/lower_trainium_layout.cc b/src/backend/trn/transform/lower_trainium_layout.cc index 5354ff5b34f7..977865762a6e 100644 --- a/src/backend/trn/transform/lower_trainium_layout.cc +++ b/src/backend/trn/transform/lower_trainium_layout.cc @@ -55,28 +55,29 @@ static bool IsTrainiumLayout(const TileLayoutNode* layout) { class TrainiumLayoutApplier : public arith::IRMutatorWithAnalyzer { public: - static std::pair> Lower( - const Stmt& stmt, const ffi::Map buffer_map) { + static std::pair> Lower( + const Stmt& stmt, const ffi::Map buffer_map) { arith::Analyzer ana; TrainiumLayoutApplier storage_lower(ana); - std::unordered_map new_buffer_map; - std::vector param_flattened_buffers; + std::unordered_map new_buffer_map; + std::vector> param_flattened_buffers; for (const auto& kv : buffer_map) { if (kv.second->layout.has_value()) { - param_flattened_buffers.push_back(storage_lower.GetFlattenedBuffer(kv.second)); - Buffer buffer = kv.second; - auto* writer = buffer.CopyOnWrite(); - writer->layout = std::nullopt; + BufferVar flattened = storage_lower.GetFlattenedBuffer(kv.second); + auto type = CopyBufferType(kv.second); + type->layout = std::nullopt; + BufferVar buffer = RebuildBufferVar(kv.second, std::move(type)); + param_flattened_buffers.emplace_back(flattened, buffer); new_buffer_map[kv.first] = buffer; } else { new_buffer_map[kv.first] = kv.second; } } auto new_stmt = storage_lower(stmt); - for (const auto& buf : param_flattened_buffers) { - new_stmt = SeqStmt::Flatten(DeclBuffer(buf), std::move(new_stmt)); + for (const auto& [buf, source] : param_flattened_buffers) { + new_stmt = SeqStmt::Flatten(DeclBuffer(buf, source.data()), std::move(new_stmt)); } - return std::make_pair(new_stmt, ffi::Map(new_buffer_map)); + return std::make_pair(new_stmt, ffi::Map(new_buffer_map)); } protected: @@ -90,7 +91,7 @@ class TrainiumLayoutApplier : public arith::IRMutatorWithAnalyzer { if (any == nullptr) { return any; } - if (auto buffer = any.as()) { + if (auto buffer = any.as()) { return GetFlattenedBuffer(buffer.value()); } else if (auto prim_expr = any.as()) { return VisitPrimExpr(prim_expr.value()); @@ -123,14 +124,14 @@ class TrainiumLayoutApplier : public arith::IRMutatorWithAnalyzer { return Stmt(n); } - Buffer GetFlattenedBuffer(Buffer buf, bool is_alloc = false) { + BufferVar GetFlattenedBuffer(BufferVar buf, bool is_alloc = false) { auto it = buffer_remap_.find(buf); if (it != buffer_remap_.end()) { return it->second; } auto trn_layout = buf->layout.as(); - Buffer flattened; - tirx::BufferNode* writer; + BufferVar flattened; + ffi::ObjectPtr type; if (IsTrainiumLayout(trn_layout)) { ffi::Array new_shape = buf.scope() == "trn.psum" ? ffi::Array{trn_layout->GetSpan(ffi::String("Bank")), @@ -139,9 +140,9 @@ class TrainiumLayoutApplier : public arith::IRMutatorWithAnalyzer { : ffi::Array{trn_layout->GetSize(ffi::String("P")), trn_layout->GetSpan(ffi::String("F"))}; flattened = buf; - writer = flattened.CopyOnWrite(); - writer->shape = new_shape; - writer->strides = {}; + type = CopyBufferType(flattened); + type->shape = new_shape; + type->strides = {}; } else if (is_alloc) { if (auto tile_layout = buf->layout.as(); tile_layout && tile_layout->HasThreadAxis()) { @@ -163,25 +164,27 @@ class TrainiumLayoutApplier : public arith::IRMutatorWithAnalyzer { } } flattened = buf; - writer = flattened.CopyOnWrite(); - writer->shape = {ana->Simplify(mem_span)}; - writer->strides = {}; + type = CopyBufferType(flattened); + type->shape = {ana->Simplify(mem_span)}; + type->strides = {}; } else { flattened = buf.GetFlattenedBuffer(); - writer = flattened.CopyOnWrite(); + type = CopyBufferType(flattened); } } else { flattened = buf.GetFlattenedBuffer(); - writer = flattened.CopyOnWrite(); + type = CopyBufferType(flattened); } if (flattened->dtype->dtype == DLDataType{kDLBool, 8, 1}) { - writer->dtype = PrimType::Int(8); + type->dtype = PrimType::Int(8); + type->data_pointer_type = PointerType(PrimType::Int(8), flattened.scope()); } for (size_t i = 0; i < flattened->shape.size(); ++i) { - writer->shape.Set(i, analyzer_->canonical_simplify(flattened->shape[i])); + type->shape.Set(i, analyzer_->canonical_simplify(flattened->shape[i])); } - writer->layout = std::nullopt; - writer->elem_offset = StmtExprMutator::VisitPrimExpr(buf->elem_offset); + type->layout = std::nullopt; + type->elem_offset = StmtExprMutator::VisitPrimExpr(buf->elem_offset); + flattened = RebuildBufferVar(flattened, std::move(type)); buffer_remap_[buf] = flattened; return flattened; @@ -230,7 +233,7 @@ class TrainiumLayoutApplier : public arith::IRMutatorWithAnalyzer { } } - ffi::Array GetSimplifiedElemOffset(const Buffer& buffer, + ffi::Array GetSimplifiedElemOffset(const BufferVar& buffer, const ffi::Array& indices) { if (buffer->layout.has_value()) { auto tile_layout = buffer->layout.value().as(); @@ -270,14 +273,14 @@ class TrainiumLayoutApplier : public arith::IRMutatorWithAnalyzer { return node; } auto flattened_indices = GetSimplifiedElemOffset(node->buffer, node->indices); - Buffer flattened_buffer = GetFlattenedBuffer(node->buffer); + BufferVar flattened_buffer = GetFlattenedBuffer(node->buffer); auto writer = node.CopyOnWrite(); writer->buffer = flattened_buffer; writer->indices = flattened_indices; return node; } - std::unordered_map buffer_remap_; + std::unordered_map buffer_remap_; }; class TrainiumBufferOffsetRemover : public StmtExprMutator { @@ -300,11 +303,12 @@ class TrainiumBufferOffsetRemover : public StmtExprMutator { if (elem_offset.same_as(buffer->elem_offset)) { return StmtExprMutator::VisitStmt_(op); } else { - auto n_buffer = buffer.CopyOnWrite(); - n_buffer->elem_offset = std::move(elem_offset); + auto type = CopyBufferType(buffer); + type->elem_offset = std::move(elem_offset); + buffer = RebuildBufferVar(buffer, std::move(type)); buffer_remap_[op->buffer] = buffer; auto n = CopyOnWrite(op); - n->buffer = ffi::GetRef(n_buffer); + n->buffer = buffer; return Stmt(n); } } @@ -336,7 +340,7 @@ class TrainiumBufferOffsetRemover : public StmtExprMutator { return node; } - std::unordered_map buffer_remap_; + std::unordered_map buffer_remap_; }; namespace transform { diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc b/src/backend/vulkan/codegen/codegen_spirv.cc index 7dbb23820bbd..b35137e74bf9 100644 --- a/src/backend/vulkan/codegen/codegen_spirv.cc +++ b/src/backend/vulkan/codegen/codegen_spirv.cc @@ -320,7 +320,10 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const LetNode* op) { } spirv::Value CodeGenSPIRV::VisitExpr_(const CallNode* op) { - if (op->op.same_as(builtin::call_spirv_pure_glsl450())) { + if (op->op.same_as(builtin::buffer_data())) { + TVM_FFI_ICHECK_EQ(op->args.size(), 1U); + return MakeValue(op->args[0]); + } else if (op->op.same_as(builtin::call_spirv_pure_glsl450())) { TVM_FFI_ICHECK_GE(op->args.size(), 2U); uint32_t inst_id = static_cast(op->args[0].as()->value); std::vector values; @@ -533,7 +536,7 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const CallNode* op) { return spirv::Value(); } else if (op->op.same_as(builtin::address_of())) { const BufferLoadNode* load = op->args[0].as(); - Var buffer_var = load->buffer->data; + Var buffer_var = load->buffer.var(); const VarNode* buffer_node = buffer_var.get(); PrimExpr index = load->indices[0]; PrimType ele_dtype = GetElementDataType(buffer_node); @@ -578,7 +581,7 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const BroadcastNode* op) { spirv::Value CodeGenSPIRV::VisitExpr_(const BufferLoadNode* op) { TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "SPIR-V codegen expects flat memory buffers"; TVM_FFI_ICHECK(!op->predicate.has_value()) << "Predicated buffer load is not supported."; - Var buffer_var = op->buffer->data; + Var buffer_var = op->buffer.var(); PrimExpr prim_index = op->indices[0]; PrimType desired_read_type = op->ty.as_or_throw(); @@ -665,7 +668,7 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const ShuffleNode* op) { void CodeGenSPIRV::VisitStmt_(const BufferStoreNode* op) { TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "SPIR-V codegen expects flat memory buffers"; TVM_FFI_ICHECK(!op->predicate.has_value()) << "Predicated buffer store is not supported."; - Var buffer_var = op->buffer->data; + Var buffer_var = op->buffer.var(); PrimExpr prim_index = op->indices[0]; auto it = storage_info_.find(buffer_var.get()); @@ -839,12 +842,12 @@ void CodeGenSPIRV::VisitStmt_(const AllocBufferNode* op) { TVM_FFI_ICHECK_GT(constant_size, 0) << "Can only handle constant size stack allocation in GPU"; spirv::Value buf; - const std::string scope = GetPtrStorageScope(op->buffer->data); + const std::string scope = GetPtrStorageScope(op->buffer.var()); auto storage_scope = runtime::StorageScope::Create(scope); spirv::SType etype = builder_->GetSType(op->buffer->dtype); runtime::StorageRank rank = storage_scope.rank; spv::StorageClass storage_class; - const VarNode* var_node = op->buffer->data.get(); + const VarNode* var_node = op->buffer.get(); switch (rank) { case runtime::StorageRank::kWMMAMatrixA: @@ -881,11 +884,11 @@ void CodeGenSPIRV::VisitStmt_(const AllocBufferNode* op) { TVM_FFI_THROW(InternalError) << "Can only allocate shared or local memory inside kernel"; } - builder_->SetName(buf, op->buffer->name); + builder_->SetName(buf, op->buffer.name()); StorageInfo& info = storage_info_[var_node]; TVM_FFI_ICHECK(!info.element_type_known); - info.SetContentType(op->buffer->dtype, op->buffer->name); + info.SetContentType(op->buffer->dtype, op->buffer.name()); TVM_FFI_ICHECK(!var_map_.count(var_node)); var_map_[var_node] = buf; diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc b/src/backend/webgpu/codegen/codegen_webgpu.cc index f24d6911d591..933a81f826fd 100644 --- a/src/backend/webgpu/codegen/codegen_webgpu.cc +++ b/src/backend/webgpu/codegen/codegen_webgpu.cc @@ -77,7 +77,7 @@ class WebGPUWorkgroupInfoCollector : public StmtExprVisitor { void VisitStmt_(const BufferStoreNode* op) final { StmtExprVisitor::VisitStmt_(op); - info_.write_access_set.insert(op->buffer->data); + info_.write_access_set.insert(op->buffer.var()); } void VisitStmt_(const AttrStmtNode* op) final { @@ -538,7 +538,7 @@ void CodeGenWebGPU::VisitExpr_(const BufferLoadNode* op, std::ostream& os) { // PrimType value_ty = op->ty.as_or_throw(); PrimExpr index = op->indices[0]; - Var buffer_var = op->buffer->data; + Var buffer_var = op->buffer.var(); const PrimType& element_ty = op->buffer->dtype; int lanes = value_ty.lanes(); @@ -611,7 +611,7 @@ void CodeGenWebGPU::VisitStmt_(const BufferStoreNode* op) { PrimType value_ty = op->value.ty(); const PrimType& element_ty = op->buffer->dtype; PrimExpr index = op->indices[0]; - Var buffer_var = op->buffer->data; + Var buffer_var = op->buffer.var(); std::string buffer_vid = GetVarID(buffer_var.get()); @@ -667,7 +667,7 @@ void CodeGenWebGPU::VisitStmt_(const BufferStoreNode* op) { void CodeGenWebGPU::VisitStmt_(const AllocBufferNode* op) { TVM_FFI_ICHECK(op->buffer.defined()); - std::string vid = AllocVarID(op->buffer->data.get()); + std::string vid = AllocVarID(op->buffer.get()); size_t constant_size = 1; for (const auto& dim : op->buffer->shape) { const IntImmNode* dim_imm = dim.as(); @@ -675,7 +675,7 @@ void CodeGenWebGPU::VisitStmt_(const AllocBufferNode* op) { constant_size *= dim_imm->value; } TVM_FFI_ICHECK_GT(constant_size, 0) << "Can only handle constant size stack allocation for now"; - auto storage_scope = runtime::StorageScope::Create(GetPtrStorageScope(op->buffer->data)); + auto storage_scope = runtime::StorageScope::Create(op->buffer.scope()); if (storage_scope.rank == runtime::StorageRank::kShared) { this->decl_stream << "var " << vid << " : array<"; diff --git a/src/relax/analysis/layout_transformation.cc b/src/relax/analysis/layout_transformation.cc index b0140a643cd8..6b48cd41eb37 100644 --- a/src/relax/analysis/layout_transformation.cc +++ b/src/relax/analysis/layout_transformation.cc @@ -273,7 +273,7 @@ static ffi::Optional InferLayoutTransformation(const SpatialLayout& sr }); if (depends_on_initial_indices) { LOG(WARNING) - << "[LayoutInference] Buffer access is dependent on both defined and undefined vars"; + << "[LayoutInference] BufferVar access is dependent on both defined and undefined vars"; return {}; } // It is ok to erase this final index expression as it only depends on undefined vars. @@ -328,7 +328,7 @@ static ffi::Optional InferLayoutTransformation(const SpatialLayout& sr class BlockAnalyzer : public StmtExprVisitor { public: explicit BlockAnalyzer(const SBlock& block, - const ffi::Map& transformation_cache, + const ffi::Map& transformation_cache, IndexMap write_transformation) : can_transform_block_(true), write_transformation_(write_transformation), @@ -361,7 +361,7 @@ class BlockAnalyzer : public StmtExprVisitor { } // Helper to get the spatial layout of buffer from buffer access map. - auto get_spatial_layout = [&](Buffer b) -> SpatialLayout { + auto get_spatial_layout = [&](BufferVar b) -> SpatialLayout { auto it = buffer_access_info_.find(b); if (it == buffer_access_info_.end()) { return {}; @@ -411,7 +411,7 @@ class BlockAnalyzer : public StmtExprVisitor { IndexMap read_transformation = maybe_read_transformation.value(); if (buffer_transformation_cache_.count(r->buffer) != 0) { if (!AreIdenticalTransforms(read_transformation, buffer_transformation_cache_[r->buffer])) - LOG(WARNING) << "[LayoutInference] Buffer: " << r->buffer + LOG(WARNING) << "[LayoutInference] BufferVar: " << r->buffer << " has conflicting transform proposals -- (preferred) " << buffer_transformation_cache_[r->buffer] << " vs. " << read_transformation; continue; @@ -511,7 +511,7 @@ class BlockAnalyzer : public StmtExprVisitor { } void VisitExpr_(const BufferLoadNode* op) final { - Buffer read_buffer = op->buffer; + BufferVar read_buffer = op->buffer; BufferAccessInfo& access_info = buffer_access_info_[op->buffer]; auto detected_spatial_layout = DetectBufferAccessIterMap(op->indices); @@ -526,7 +526,7 @@ class BlockAnalyzer : public StmtExprVisitor { public: bool CanBeTransformed() { return can_transform_block_; } IndexMap GetSBlockTransformation() { return block_transformation_; } - ffi::Map GetReadBufferTransformations() { return read_buffer_transformations_; } + ffi::Map GetReadBufferTransformations() { return read_buffer_transformations_; } private: bool can_transform_block_; @@ -537,9 +537,9 @@ class BlockAnalyzer : public StmtExprVisitor { SBlock block_; IndexMap block_transformation_; - ffi::Map read_buffer_transformations_; - const ffi::Map& buffer_transformation_cache_; - std::unordered_map + ffi::Map read_buffer_transformations_; + const ffi::Map& buffer_transformation_cache_; + std::unordered_map buffer_access_info_; }; @@ -560,7 +560,7 @@ class PrimFuncAnalyzer : public StmtExprVisitor { size_t first_write_index = func->params.size() - write_transformations.size(); for (size_t i = 0; i < write_transformations.size(); ++i) { auto param = func->params[first_write_index + i]; - ffi::Optional param_buf = func->buffer_map.Get(param); + ffi::Optional param_buf = func->buffer_map.Get(param); TVM_FFI_ICHECK(param_buf.has_value()); TVM_FFI_ICHECK_EQ(param_buf.value()->shape.size(), write_transformations[i]->initial_indices.size()) @@ -613,9 +613,9 @@ class PrimFuncAnalyzer : public StmtExprVisitor { } private: - ffi::Map buffer_transformation_cache_; + ffi::Map buffer_transformation_cache_; ffi::Map block_transformations_; - std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> + std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> block_to_buffer_; }; diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc index 912fb58305c8..5dc990bd1df9 100644 --- a/src/relax/analysis/tir_op_pattern_kind.cc +++ b/src/relax/analysis/tir_op_pattern_kind.cc @@ -39,7 +39,7 @@ class PatternKindAnalyzer : public StmtExprVisitor { public: explicit PatternKindAnalyzer(const tirx::PrimFunc& func) { for (const tirx::Var& param : func->params) { - ffi::Optional param_buf = func->buffer_map.Get(param); + ffi::Optional param_buf = func->buffer_map.Get(param); if (param_buf.has_value()) { param_buffers_.insert(param_buf.value()); } @@ -102,7 +102,7 @@ class PatternKindAnalyzer : public StmtExprVisitor { // while the order amount enums: kElemWise < kBroadcast < kInjective. // We can simply use `std::max` to detect these three patterns. // E.g Here is only one store node but two load nodes, like C[i, j] = A[i, j] + B[i] - // Buffer C and A are elemwise but C and B are broadcast. So the whole block follows + // BufferVar C and A are elemwise but C and B are broadcast. So the whole block follows // broadcast pattern. if (IsElemwisePattern(store, load)) { index_pair_pattern = std::max(index_pair_pattern, kElemWise); @@ -346,7 +346,7 @@ class PatternKindAnalyzer : public StmtExprVisitor { /*! \brief The result of op pattern. */ OpPatternKind kind_ = kElemWise; /*! \brief The buffers from function params. I.e. the input and output buffers. */ - std::unordered_set param_buffers_; + std::unordered_set param_buffers_; public: OpPatternKind GetResult() { return kind_; } @@ -361,14 +361,14 @@ OpPatternKind AnalyzeOpPatternKind(const PrimFunc& func) { bool HasReshapePattern(const PrimFunc& func) { class ReshapeDetector : public StmtVisitor { public: - static bool Detect(const Buffer& src_buffer, const Buffer& dst_buffer, Stmt stmt) { + static bool Detect(const BufferVar& src_buffer, const BufferVar& dst_buffer, Stmt stmt) { ReshapeDetector detector(src_buffer, dst_buffer); detector(stmt); return detector.is_reshape_; } private: - explicit ReshapeDetector(const Buffer& src_buffer, const Buffer& dst_buffer) + explicit ReshapeDetector(const BufferVar& src_buffer, const BufferVar& dst_buffer) : is_reshape_(false), src_buffer_(src_buffer), dst_buffer_(dst_buffer) {} void VisitStmt_(const ForNode* loop) final { @@ -436,7 +436,7 @@ bool HasReshapePattern(const PrimFunc& func) { // This check requires at least one of the src/dst side is a trivial buffer // access (e.g., buf[ax0, ax1, ax2]). - auto f_calc_flattened_idx = [&](const Buffer& buffer, const ffi::Array& indices) { + auto f_calc_flattened_idx = [&](const BufferVar& buffer, const ffi::Array& indices) { TVM_FFI_ICHECK_EQ(indices.size(), buffer->shape.size()); int ndim = indices.size(); PrimExpr idx = 0; @@ -453,7 +453,7 @@ bool HasReshapePattern(const PrimFunc& func) { /*simplify_trivial_iterators=*/true)[0]; }; - auto f_is_trivial_indices = [block, this](const Buffer& buffer, + auto f_is_trivial_indices = [block, this](const BufferVar& buffer, const ffi::Array& indices) { if (indices.size() != block->iter_vars.size()) { return false; @@ -470,7 +470,7 @@ bool HasReshapePattern(const PrimFunc& func) { }; ffi::Array nontrivial_indices{nullptr}; - Buffer nontrivial_buffer{nullptr}; + BufferVar nontrivial_buffer{nullptr}; if (f_is_trivial_indices(dst_buffer_, buffer_store->indices)) { nontrivial_indices = buffer_load->indices; nontrivial_buffer = src_buffer_; @@ -527,12 +527,12 @@ bool HasReshapePattern(const PrimFunc& func) { } bool is_reshape_; - const Buffer& src_buffer_; - const Buffer& dst_buffer_; + const BufferVar& src_buffer_; + const BufferVar& dst_buffer_; arith::Analyzer ana_; }; - ffi::Array buffer_args; + ffi::Array buffer_args; for (const auto& param : func->params) { if (auto buffer = func->buffer_map.Get(param)) { buffer_args.push_back(buffer.value()); @@ -542,8 +542,8 @@ bool HasReshapePattern(const PrimFunc& func) { if (buffer_args.size() < 2) { return false; } - Buffer src_buffer = buffer_args.front(); - Buffer dst_buffer = buffer_args.back(); + BufferVar src_buffer = buffer_args.front(); + BufferVar dst_buffer = buffer_args.back(); // To detect the reshape pattern, we require each For to have // either another For or a BlockRealize as body. diff --git a/src/relax/backend/adreno/annotate_custom_storage.cc b/src/relax/backend/adreno/annotate_custom_storage.cc index c3c47c65a786..19e02ed573fd 100644 --- a/src/relax/backend/adreno/annotate_custom_storage.cc +++ b/src/relax/backend/adreno/annotate_custom_storage.cc @@ -231,7 +231,7 @@ * - Fusion * - FoldVDeviceScopeChange: There existed some ToVDevice copies from texture to buffer * This pass removes the copes and updates producer scope to global. - * - SpecializePrimFuncBasedOnCallSite: Finally we updates the Buffer Var maps according to + * - SpecializePrimFuncBasedOnCallSite: Finally we updates the BufferVar Var maps according to * VDevice scopes. * */ @@ -256,7 +256,7 @@ namespace relax { namespace backend { namespace adreno { -using tvm::tirx::Buffer; +using tvm::tirx::BufferVar; static ffi::Array GetShapeFromTensorType(const TensorType& tensor_ty) { auto shape = tensor_ty->GetShape(); diff --git a/src/relax/backend/vm/vm_shape_lower.cc b/src/relax/backend/vm/vm_shape_lower.cc index b83a5ed72d3c..f671c1290440 100644 --- a/src/relax/backend/vm/vm_shape_lower.cc +++ b/src/relax/backend/vm/vm_shape_lower.cc @@ -263,7 +263,7 @@ class PrimExprSlotCollector : public ExprVisitor, public TypeVisitor { * \code * * @T.prim_func - * def shape_func(H: T.Buffer([3], "int64")): + * def shape_func(H: T.BufferVar([3], "int64")): * H[1] = H[2] + 1 * * \endcode @@ -708,9 +708,9 @@ class VMShapeLowerMutator TVM_FFI_ICHECK_GT(heap_size_->value, 0); // construct a PrimFunc that compute the shape. ffi::Array buffer_shape{heap_size_}; - tirx::Buffer buffer = tirx::decl_buffer(buffer_shape, PrimType(ShapeDType()), "H", "global"); + tirx::BufferVar buffer = tirx::decl_buffer(buffer_shape, PrimType(ShapeDType()), "H", "global"); tirx::Var heap("heap", PointerType::VoidPointerTy()); - ffi::Map buffer_map; + ffi::Map buffer_map; buffer_map.Set(heap, buffer); ffi::Map var_map; 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 aa1c54fd51c6..cac73cd97cc7 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 @@ -39,13 +39,13 @@ using s_tir::ReplaceBuffer; class DistBufferReplacer : public StmtExprMutator { public: - static Stmt BufferReplace(Stmt stmt, ffi::Map buffer_map) { + static Stmt BufferReplace(Stmt stmt, ffi::Map buffer_map) { DistBufferReplacer replacer(buffer_map); return replacer(stmt); } private: - explicit DistBufferReplacer(ffi::Map buffer_map) : buffer_map_(buffer_map) {} + explicit DistBufferReplacer(ffi::Map buffer_map) : buffer_map_(buffer_map) {} Stmt VisitStmt_(const BufferStoreNode* _store) final { BufferStore store = StmtExprMutator::VisitStmt_(_store).as_or_throw(); @@ -76,7 +76,7 @@ class DistBufferReplacer : public StmtExprMutator { return SBlock(new_block); } - ffi::Map buffer_map_; + ffi::Map buffer_map_; }; class DistSBlockInfoCollector : public StmtExprVisitor { @@ -136,10 +136,10 @@ class DistSBlockInfoCollector : public StmtExprVisitor { StmtExprVisitor::VisitExpr_(op); } - Buffer reduce_buffer_; + BufferVar reduce_buffer_; public: - std::unordered_map>, ffi::ObjectPtrHash, + std::unordered_map>, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> buffer_access_indices; std::string reduce_kind; @@ -155,10 +155,10 @@ class DistributedBufferCompactor : StmtExprMutator { const std::vector& sharding_specs, PrimFunc prim_func) { prim_func = s_tir::RenewDefs(prim_func); DistributedBufferCompactor compactor(sharding_specs, prim_func); - ffi::Map new_func_buffer_map; - ffi::Map replace_buffer_map; + ffi::Map new_func_buffer_map; + ffi::Map replace_buffer_map; for (const auto& pr : prim_func->buffer_map) { - Buffer shard_buffer = compactor.ShardBuffer(pr.second); + BufferVar shard_buffer = compactor.ShardBuffer(pr.second); new_func_buffer_map.Set(pr.first, shard_buffer); if (!shard_buffer.same_as(pr.second)) { replace_buffer_map.Set(pr.second, shard_buffer); @@ -186,7 +186,7 @@ class DistributedBufferCompactor : StmtExprMutator { if (!prim_func->buffer_map.count(param_var)) { continue; } - Buffer param_buffer = prim_func->buffer_map[param_var]; + BufferVar param_buffer = prim_func->buffer_map[param_var]; ShardingSpec spec = sharding_specs_[j++]; for (int mesh_dim = 0; mesh_dim < static_cast(spec.first->shape.size()); mesh_dim++) { @@ -205,9 +205,9 @@ class DistributedBufferCompactor : StmtExprMutator { ffi::Array ShardIterVar( SBlock block, - const std::unordered_map>, ffi::ObjectPtrHash, + const std::unordered_map>, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>& buffer_access_indices) { - std::vector buffers; + std::vector buffers; for (const auto& read : block->reads) { buffers.push_back(read->buffer); } @@ -257,7 +257,7 @@ class DistributedBufferCompactor : StmtExprMutator { return new_iter_vars; } - Buffer ShardBuffer(Buffer buffer) { + BufferVar ShardBuffer(BufferVar buffer) { if (buffer_shards_.count(buffer) == 0) { return buffer; } @@ -270,9 +270,10 @@ class DistributedBufferCompactor : StmtExprMutator { shape.push_back(buffer->shape[i]); } } - ffi::ObjectPtr new_buffer = ffi::make_object(*buffer.get()); - new_buffer->shape = shape; - return Buffer(new_buffer); + BufferType new_type(buffer->data_pointer_type, buffer->dtype, std::move(shape), + buffer->strides, buffer->elem_offset, buffer->data_alignment, + buffer->offset_factor, buffer->layout, buffer->allocated_addr); + return BufferVar(buffer.name(), std::move(new_type), buffer.span()); } Stmt VisitStmt_(const SBlockNode* op) final { @@ -280,10 +281,10 @@ class DistributedBufferCompactor : StmtExprMutator { DistSBlockInfoCollector collector; collector(block); ffi::Array new_iter_vars = ShardIterVar(block, collector.buffer_access_indices); - ffi::Array new_alloc_buffers; - ffi::Map buffer_map; - for (const Buffer& buffer : block->alloc_buffers) { - Buffer sharded_buffer = ShardBuffer(buffer); + ffi::Array new_alloc_buffers; + ffi::Map buffer_map; + for (const BufferVar& buffer : block->alloc_buffers) { + BufferVar sharded_buffer = ShardBuffer(buffer); if (!sharded_buffer.same_as(buffer)) { buffer_map.Set(buffer, sharded_buffer); } @@ -344,10 +345,10 @@ class DistributedBufferCompactor : StmtExprMutator { std::unordered_map iter_var_shards_; std::unordered_map loop_var_shards_; - ffi::Array allocated_buffer_under_root; + ffi::Array allocated_buffer_under_root; BufferAxisGraphExtractor extractor_; std::vector sharding_specs_; - std::unordered_map buffer_shards_; + std::unordered_map buffer_shards_; std::string add_allreduce_kind_; }; diff --git a/src/relax/ir/tir_pattern.cc b/src/relax/ir/tir_pattern.cc index 4d00d17976dd..cfd332a94f4e 100644 --- a/src/relax/ir/tir_pattern.cc +++ b/src/relax/ir/tir_pattern.cc @@ -25,7 +25,7 @@ namespace relax { TVM_FFI_STATIC_INIT_BLOCK() { MatchResultNode::RegisterReflection(); } MatchResult::MatchResult(TIRPattern pattern, ffi::Array symbol_values, - ffi::Array matched_buffers) { + ffi::Array matched_buffers) { auto n = ffi::make_object(); n->pattern = std::move(pattern); n->symbol_values = std::move(symbol_values); diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc index 80de3c8f8eda..e6f31492a4bf 100644 --- a/src/relax/op/op.cc +++ b/src/relax/op/op.cc @@ -263,8 +263,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { * * For dynamic shapes, it is not always possible to infer the output * of a TIR PrimFunc from its inputs. For example, a PrimFunc that - * accepts input buffer `T.Buffer([16], "float32")` and output buffer - * `T.Buffer([M, N], "float32")` infers the values of `M` and `N` from + * accepts input buffer `T.BufferVar([16], "float32")` and output buffer + * `T.BufferVar([M, N], "float32")` infers the values of `M` and `N` from * the shape of the provided output buffer. * * If the arguments provided are not compatible with the PrimFunc's diff --git a/src/relax/op/tensor/inspect.cc b/src/relax/op/tensor/inspect.cc index b6abbfa54b24..5a15dab73518 100644 --- a/src/relax/op/tensor/inspect.cc +++ b/src/relax/op/tensor/inspect.cc @@ -257,7 +257,7 @@ Expr LegalizeTensorShape(const BlockBuilder& bb, const Call& call) { tirx::Var ndim("ndim", PrimType::Int(32)); - tirx::Buffer shape_buffer = + tirx::BufferVar shape_buffer = tirx::decl_buffer({ndim.as_or_throw()}, field_ty, "shape"); tirx::Var extent("extent", field_ty); @@ -276,11 +276,11 @@ Expr LegalizeTensorShape(const BlockBuilder& bb, const Call& call) { tirx::StringImm("RuntimeError"), {tirx::StringImm( "Specified axis may not be larger than the tensor's dimensionality")}), - tirx::Bind(shape_buffer->data, - tvm::Call(shape_buffer->data->ty, tirx::builtin::tvm_struct_get(), - {dlpack_handle, IntImm::Int32(0), - IntImm::Int32(tirx::builtin::TVMStructFieldKind::kDLTensorShape)})), - tirx::DeclBuffer(shape_buffer), + tirx::DeclBuffer( + shape_buffer, + tvm::Call(shape_buffer->data_pointer_type, tirx::builtin::tvm_struct_get(), + {dlpack_handle, IntImm::Int32(0), + IntImm::Int32(tirx::builtin::TVMStructFieldKind::kDLTensorShape)})), tirx::Bind(extent, tirx::BufferLoad(shape_buffer, {axis.as_or_throw()})), tirx::Return(extent)}); @@ -328,7 +328,7 @@ Type InferTypeTensorStride(const Call& call, const BlockBuilder&) { // `FLegalize` function for most operators is implemented in terms // of `topi`, and is then converted from TE to `tirx::PrimFunc` // using `tvm::tirx::CreatePrimFunc`. The `te::Tensor` is - // converted to a `tirx::Buffer` in `RewriteStageToBlock`, and uses + // converted to a `tirx::BufferVar` in `RewriteStageToBlock`, and uses // the default empty list for the strides. The empty strides // represent a compact data array. // diff --git a/src/relax/transform/dataflow_inplace.cc b/src/relax/transform/dataflow_inplace.cc index 9ea3be24dc46..d6a129722697 100644 --- a/src/relax/transform/dataflow_inplace.cc +++ b/src/relax/transform/dataflow_inplace.cc @@ -763,10 +763,10 @@ FindInplaceOpportunities(const DataflowBlock& block, const ffi::Array& inpu // Replace buffers in a PrimFunc according to the mapping. tirx::Stmt RemapBuffers(const tirx::Stmt& stmt, - const ffi::Map& buffer_map) { + const ffi::Map& buffer_map) { class BufferMapper : public tirx::StmtExprMutator { public: - explicit BufferMapper(const ffi::Map& buffer_map) + explicit BufferMapper(const ffi::Map& buffer_map) : buffer_map_(buffer_map) {} tirx::Stmt Remap(const tirx::Stmt& stmt) { return VisitStmt(stmt); } @@ -804,7 +804,7 @@ tirx::Stmt RemapBuffers(const tirx::Stmt& stmt, auto* node_cow = node.CopyOnWrite(); // need the lambdas because class methods are not first-class (how ironic) node_cow->alloc_buffers = - node->alloc_buffers.Map([this](const tirx::Buffer& b) { return AttemptRemap(b); }); + node->alloc_buffers.Map([this](const tirx::BufferVar& b) { return AttemptRemap(b); }); node_cow->reads = node->reads.Map([this](const tirx::BufferRegion& br) { return VisitBufferRegion(br); }); node_cow->writes = @@ -815,7 +815,7 @@ tirx::Stmt RemapBuffers(const tirx::Stmt& stmt, } private: - tirx::Buffer AttemptRemap(const tirx::Buffer& buffer) { + tirx::BufferVar AttemptRemap(const tirx::BufferVar& buffer) { if (buffer_map_.count(buffer)) { return buffer_map_.at(buffer); } @@ -834,7 +834,7 @@ tirx::Stmt RemapBuffers(const tirx::Stmt& stmt, return region; } - const ffi::Map& buffer_map_; + const ffi::Map& buffer_map_; }; BufferMapper mapper(buffer_map); @@ -957,7 +957,7 @@ class ModuleInplaceTransformer : public ExprMutator { // 2. For each output var, replace its instances with the corresponding inplace index var // 3. Do the same for the *buffer vars* corresponding to the output vars // 4. Remove the output vars from the param list and buffer map - ffi::Map buffer_subst_map; + ffi::Map buffer_subst_map; ffi::Map var_subst_map; for (size_t i = 0; i < num_outs; i++) { // we will substitute output i with the corresponding param indicated by inplace indices @@ -968,7 +968,7 @@ class ModuleInplaceTransformer : public ExprMutator { // also do the same with the buffer vars auto output_buffer = old_primfunc->buffer_map.at(output_var); auto inplace_buffer = old_primfunc->buffer_map.at(inplace_var); - var_subst_map.Set(output_buffer->data, inplace_buffer->data); + var_subst_map.Set(output_buffer.var(), inplace_buffer.var()); buffer_subst_map.Set(output_buffer, inplace_buffer); } diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc index 5db564c9bf8a..3a03c1c186bc 100644 --- a/src/relax/transform/fuse_tir.cc +++ b/src/relax/transform/fuse_tir.cc @@ -166,20 +166,20 @@ class SymbolicMatcher : ExprFunctor */ class FuseTIRBufferSubstitutor : private StmtExprMutator { public: - explicit FuseTIRBufferSubstitutor(const ffi::Map& buffer_map, + explicit FuseTIRBufferSubstitutor(const ffi::Map& buffer_map, const ffi::Map& var_map) { buffer_remap_ = buffer_map; for (const auto& [var, value] : var_map) { var_remap_.Set(var, value); } for (const auto& [src, tgt] : buffer_map) { - var_remap_.Set(src->data, tgt->data); + var_remap_.Set(src.var(), tgt.var()); } } Stmt Substitute(Stmt stmt) { return this->VisitStmt(std::move(stmt)); } - Buffer SubstituteAllocatedBuffer(Buffer buffer) { + BufferVar SubstituteAllocatedBuffer(BufferVar buffer) { TVM_FFI_ICHECK(buffer_remap_.find(buffer) == buffer_remap_.end()); ffi::Array shape = MutateArray( buffer->shape, [this](const PrimExpr& expr) { return this->VisitPrimExpr(expr); }); @@ -190,11 +190,11 @@ class FuseTIRBufferSubstitutor : private StmtExprMutator { elem_offset.same_as(buffer->elem_offset)) { return buffer; } else { - auto n = ffi::make_object(*buffer.get()); - n->shape = std::move(shape); - n->strides = std::move(strides); - n->elem_offset = std::move(elem_offset); - Buffer new_buffer(n); + BufferType new_type( + buffer->data_pointer_type, buffer->dtype, std::move(shape), + std::move(strides), std::move(elem_offset), buffer->data_alignment, + buffer->offset_factor, buffer->layout, buffer->allocated_addr); + BufferVar new_buffer(buffer.name(), std::move(new_type), buffer.span()); this->buffer_remap_.Set(buffer, new_buffer); return new_buffer; } @@ -211,7 +211,7 @@ class FuseTIRBufferSubstitutor : private StmtExprMutator { Expr VisitExpr_(const BufferLoadNode* _op) final { BufferLoad load = StmtExprMutator::VisitExpr_(_op).as_or_throw(); - const Buffer& buffer = SubstituteBuffer(load->buffer); + const BufferVar& buffer = SubstituteBuffer(load->buffer); if (buffer.same_as(load->buffer)) { return load; @@ -224,7 +224,7 @@ class FuseTIRBufferSubstitutor : private StmtExprMutator { Stmt VisitStmt_(const BufferStoreNode* _op) final { BufferStore store = StmtExprMutator::VisitStmt_(_op).as_or_throw(); - const Buffer& buffer = SubstituteBuffer(store->buffer); + const BufferVar& buffer = SubstituteBuffer(store->buffer); if (buffer.same_as(store->buffer)) { return store; @@ -241,8 +241,8 @@ class FuseTIRBufferSubstitutor : private StmtExprMutator { // Define the mutation functions. auto f_mutate_match_buffers = [this](const MatchBufferRegion& match_buffer) { - const Buffer& src_buffer = SubstituteBuffer(match_buffer->source->buffer); - const Buffer& tgt_buffer = SubstituteAllocatedBuffer(match_buffer->buffer); + const BufferVar& src_buffer = SubstituteBuffer(match_buffer->source->buffer); + const BufferVar& tgt_buffer = SubstituteAllocatedBuffer(match_buffer->buffer); ffi::Array region = MutateRegion(match_buffer->source->region); if (src_buffer.same_as(match_buffer->source->buffer) && tgt_buffer.same_as(match_buffer->buffer) && @@ -257,7 +257,7 @@ class FuseTIRBufferSubstitutor : private StmtExprMutator { }; auto f_mutate_read_write_region = [this](const BufferRegion& buffer_region) { - const Buffer& buffer = SubstituteBuffer(buffer_region->buffer); + const BufferVar& buffer = SubstituteBuffer(buffer_region->buffer); const ffi::Array& region = MutateRegion(buffer_region->region); if (buffer.same_as(buffer_region->buffer) && region.same_as(buffer_region->region)) { return buffer_region; @@ -273,9 +273,9 @@ class FuseTIRBufferSubstitutor : private StmtExprMutator { ffi::Array reads = MutateArray(block->reads, f_mutate_read_write_region); ffi::Array writes = MutateArray(block->writes, f_mutate_read_write_region); // Step 3. Mutate the Allocate Buffers. - ffi::Array alloc_buffers = + ffi::Array alloc_buffers = MutateArray(block->alloc_buffers, - [this](const Buffer& buffer) { return SubstituteAllocatedBuffer(buffer); }); + [this](const BufferVar& buffer) { return SubstituteAllocatedBuffer(buffer); }); reads = UnionAccessRegion(reads); writes = UnionAccessRegion(writes); @@ -298,17 +298,17 @@ class FuseTIRBufferSubstitutor : private StmtExprMutator { private: /*! \brief Mapping from src buffer to tgt buffer. */ - ffi::Map buffer_remap_; + ffi::Map buffer_remap_; /*! \brief Mapping from src tirx var to tgt var. */ ffi::Map var_remap_; ffi::Array UnionAccessRegion(const ffi::Array& regions) const { - // For now we only allow Buffer access the same elements. + // For now we only allow BufferVar access the same elements. // e.g. `[A[vi, vj], A[vi, vj]]` is a legal pattern but need to union to `A[vi, vj]` // However, `A[vi, vj], A[vi, vj + 1]` is not allow for now. // Note: the order of return region should remain the same as the first occurrence of the region ffi::Array ret; - std::unordered_map> buffer_region_set; + std::unordered_map> buffer_region_set; for (const BufferRegion& region : regions) { auto it = buffer_region_set.find(region->buffer.get()); @@ -325,7 +325,7 @@ class FuseTIRBufferSubstitutor : private StmtExprMutator { } } - inline Buffer SubstituteBuffer(const Buffer& buffer) const { + inline BufferVar SubstituteBuffer(const BufferVar& buffer) const { auto it = buffer_remap_.find(buffer); if (it != buffer_remap_.end()) { return (*it).second; @@ -441,7 +441,7 @@ static ffi::Array GetInplaceOutputIndices(const ffi::Array& in class RelaxToTIRVarMapCollector : public ExprVisitor { public: explicit RelaxToTIRVarMapCollector(const IRModule& mod) : mod_(mod) {} - static ffi::Map Collect(const IRModule& mod, const Function& func) { + static ffi::Map Collect(const IRModule& mod, const Function& func) { RelaxToTIRVarMapCollector visitor(mod); visitor(func->body); return visitor.relax_to_tir_var_map_; @@ -496,7 +496,7 @@ class RelaxToTIRVarMapCollector : public ExprVisitor { // If the `expr` is already seen (present in the map), validate whether the mapped buffer is // structurally equal to the `new_buf` passed - auto ValidateBufferCompatibility = [this](tirx::Buffer new_buf, Expr expr) { + auto ValidateBufferCompatibility = [this](tirx::BufferVar new_buf, Expr expr) { if (auto it = relax_to_tir_var_map_.find(expr); it != relax_to_tir_var_map_.end()) { TVM_FFI_ICHECK(ffi::StructuralEqual()((*it).second, new_buf)) << "Inconsistent buffers " << (*it).second << " and " << new_buf @@ -525,7 +525,7 @@ class RelaxToTIRVarMapCollector : public ExprVisitor { private: /*! \brief The IRModule */ const IRModule& mod_; - ffi::Map relax_to_tir_var_map_; + ffi::Map relax_to_tir_var_map_; Var current_var_; }; @@ -560,15 +560,15 @@ class FusedTIRConstructor : public ExprVisitor { void VisitExpr_(const FunctionNode* func) final { auto relax_to_tir_var_map = RelaxToTIRVarMapCollector::Collect(mod_, ffi::GetRef(func)); - std::vector> prim_func_params; + std::vector> prim_func_params; for (const Var& relax_param : func->params) { size_t size_before = prim_func_params.size(); CollectPrimFuncParams(relax_param, &prim_func_params, relax_to_tir_var_map.Get(relax_param)); - auto param_buffers = [&]() -> ffi::Array { - ffi::Array out; + auto param_buffers = [&]() -> ffi::Array { + ffi::Array out; for (size_t i = size_before; i < prim_func_params.size(); i++) { - if (auto buf = prim_func_params[i].as()) { + if (auto buf = prim_func_params[i].as()) { out.push_back(buf.value()); } } @@ -582,15 +582,15 @@ class FusedTIRConstructor : public ExprVisitor { // parameters are both explicit call_tir arguments, while output buffers // are appended after the complete explicit argument prefix. for (const auto& param : prim_func_params) { - if (auto opt = param.as()) { + if (auto opt = param.as()) { auto buffer = opt.value(); // Differentiate buffer name and param name by adding prefix // `p_` to the buffer name. Every symbol should be unique in // TVMScript, and while they can be de-deplicated when // printed, it's more readable when done explicitly. Since - // Buffer is used more than param it gets the name with better + // BufferVar is used more than param it gets the name with better // readability. - tirx::Var param = tirx::Var("p_" + buffer->name, PointerType::VoidPointerTy()); + tirx::Var param = tirx::Var("p_" + buffer.name(), PointerType::VoidPointerTy()); func_info_.params.push_back(param); func_info_.buffer_map.Set(param, buffer); } else if (auto var = param.as()) { @@ -607,10 +607,10 @@ class FusedTIRConstructor : public ExprVisitor { TVM_FFI_ICHECK(it != func_info_.expr2buffers.end()) << "Fail to detect output buffers for function body"; - const ffi::Array& buffers = (*it).second; + const ffi::Array& buffers = (*it).second; // map of input buffers to indices (helpful for detecting in-place inputs) - std::unordered_map buffer_to_idx; + std::unordered_map buffer_to_idx; std::unordered_map input_to_idx; for (size_t i = 0; i < func_info_.params.size(); i++) { input_to_idx[func_info_.params[i]] = i; @@ -722,7 +722,7 @@ class FusedTIRConstructor : public ExprVisitor { void VisitExpr_(const TupleNode* tuple) final { ExprVisitor::VisitExpr_(tuple); - ffi::Array buffers; + ffi::Array buffers; for (const Expr& expr : tuple->fields) { auto it = func_info_.expr2buffers.find(expr); if (it != func_info_.expr2buffers.end()) { @@ -775,16 +775,16 @@ class FusedTIRConstructor : public ExprVisitor { } /*! \brief Map old TIR func param buffer to new buffer, and then update `buffer_subst_map` */ - void MapArgsToBuffer(const ffi::Array args, const ffi::Array& buffers) { + void MapArgsToBuffer(const ffi::Array args, const ffi::Array& buffers) { size_t buffer_idx = 0; for (const Expr& arg : args) { if (const auto* v = arg.as()) { auto it = func_info_.expr2buffers.find(ffi::GetRef(v)); // Substitute the buffer with the already allocated one if it is an intermediate var if (it != func_info_.expr2buffers.end()) { - for (const tirx::Buffer& target_buffer : (*it).second) { + for (const tirx::BufferVar& target_buffer : (*it).second) { TVM_FFI_ICHECK_LT(buffer_idx, buffers.size()); - const tirx::Buffer& buffer = buffers[buffer_idx]; + const tirx::BufferVar& buffer = buffers[buffer_idx]; func_info_.symbolic_var_matcher.Match(buffer->shape, target_buffer->shape); func_info_.buffer_subst_map.Set(buffer, target_buffer); buffer_idx++; @@ -803,7 +803,7 @@ class FusedTIRConstructor : public ExprVisitor { */ void MapInputBuffer(const tirx::PrimFunc& func, const relax::Expr& args) { ffi::Array arg_list; - ffi::Array buffer_list; + ffi::Array buffer_list; ffi::Array call_args = args.as_or_throw()->fields; TVM_FFI_ICHECK_GE(func->params.size(), call_args.size()); @@ -859,7 +859,7 @@ class FusedTIRConstructor : public ExprVisitor { int num_inputs = call->args[1].as_or_throw()->fields.size(); size_t output_size = output_shapes.size(); TVM_FFI_ICHECK_GE(n, output_size); - ffi::Array output_buffers; + ffi::Array output_buffers; ffi::Array output_idxs; if (is_inplace) { const auto* attrs = call->attrs.as(); @@ -874,7 +874,7 @@ class FusedTIRConstructor : public ExprVisitor { ffi::Array output_params = GetPrimFuncOutputParams(func, output_idxs); for (size_t i = 0; i < output_size; ++i) { const tirx::Var& param = output_params[i]; - const tirx::Buffer& buffer = func->buffer_map.at(param); + const tirx::BufferVar& buffer = func->buffer_map.at(param); // if this is an inplace output, do not do an intermediate allocation if (output_idxs[i] < num_inputs) { @@ -886,13 +886,13 @@ class FusedTIRConstructor : public ExprVisitor { } auto unify_name_hints = [this, &buffer]() { - ffi::String base_name = buffer->name; + ffi::String base_name = buffer.name(); ffi::String unique_name = base_name + "_intermediate"; size_t unique_id = 0; std::unordered_set names; for (auto& _buffer : func_info_.alloc_buffers) { - names.insert(_buffer->name); + names.insert(_buffer.name()); } while (names.find(unique_name) != names.end()) { @@ -901,15 +901,17 @@ class FusedTIRConstructor : public ExprVisitor { return unique_name; }; // Update buffer with new symbolic shape according to the ty - auto n = ffi::make_object(*buffer.get()); - n->shape = output_shapes[i]; - n->name = unify_name_hints(); - tirx::Buffer new_buffer(n); + tirx::BufferType new_type( + buffer->data_pointer_type, buffer->dtype, output_shapes[i], + buffer->strides, buffer->elem_offset, buffer->data_alignment, + buffer->offset_factor, buffer->layout, buffer->allocated_addr); + tirx::BufferVar new_buffer(unify_name_hints(), std::move(new_type), + buffer.span()); func_info_.alloc_buffers.push_back(new_buffer); output_buffers.push_back(new_buffer); // Match the shape of the output buffer with the shape - func_info_.symbolic_var_matcher.Match(buffer->shape, n->shape); + func_info_.symbolic_var_matcher.Match(buffer->shape, new_buffer->shape); func_info_.buffer_subst_map.Set(buffer, new_buffer); } // Update expr2buffers @@ -923,8 +925,8 @@ class FusedTIRConstructor : public ExprVisitor { * \param out The vector into which to collect the params/buffers */ static void CollectPrimFuncParams(const Var& relax_param, - std::vector>* out, - const ffi::Optional& tir_buffer_param) { + std::vector>* out, + const ffi::Optional& tir_buffer_param) { auto ty = GetType(relax_param); TVM_FFI_CHECK(!ty.as(), InternalError) @@ -938,7 +940,7 @@ class FusedTIRConstructor : public ExprVisitor { const auto* shape_expr = tensor->shape.as(); TVM_FFI_ICHECK(shape_expr) << "FuseTIR expects all Tensor parameters have a known shape."; PrimType dtype = tensor->dtype.value(); - tirx::Buffer buffer; + tirx::BufferVar buffer; if (tir_buffer_param.has_value()) { buffer = tirx::decl_buffer(shape_expr->values, dtype, name_hint, tir_buffer_param.value().scope()); @@ -977,8 +979,8 @@ class FusedTIRConstructor : public ExprVisitor { func_info_.symbolic_var_remap); TVM_FFI_ICHECK(func_info_.global_name != "fused"); // Remove output buffers from func_info_.alloc_buffers - ffi::Array alloc_buffers; - for (const tirx::Buffer& buf : func_info_.alloc_buffers) { + ffi::Array alloc_buffers; + for (const tirx::BufferVar& buf : func_info_.alloc_buffers) { if (func_info_.output_buffers.count(buf.get()) == 0) { alloc_buffers.push_back(subst.SubstituteAllocatedBuffer(buf)); } @@ -1020,9 +1022,9 @@ class FusedTIRConstructor : public ExprVisitor { * \brief The map from each dataflow var (intermediate var) to the corresponding buffers * allocated in the fused func */ - ffi::Map> expr2buffers; + ffi::Map> expr2buffers; /*! \brief The buffers to allocate in the fused func*/ - ffi::Array alloc_buffers; + ffi::Array alloc_buffers; /*! \brief The bodies of the original funcs, which is also the body of the fused func. */ ffi::Array bodies; /*! \brief The params of the fused function*/ @@ -1031,11 +1033,11 @@ class FusedTIRConstructor : public ExprVisitor { * \brief The map from buffer in original functions to corresponding buffer in the fused * function */ - ffi::Map buffer_subst_map; + ffi::Map buffer_subst_map; /*! \brief The `buffer_map` in the fused function*/ - ffi::Map buffer_map; + ffi::Map buffer_map; /*! \brief The output buffers in the function buffer_map*/ - std::unordered_set output_buffers; + std::unordered_set output_buffers; /*! \brief The name of the fused function */ std::string global_name = "fused"; diff --git a/src/relax/transform/rewrite_dataflow_reshape.cc b/src/relax/transform/rewrite_dataflow_reshape.cc index 98898f62b5c7..97a3658f5347 100644 --- a/src/relax/transform/rewrite_dataflow_reshape.cc +++ b/src/relax/transform/rewrite_dataflow_reshape.cc @@ -40,7 +40,7 @@ std::vector GetUsedTensorArgIndices(const tirx::PrimFunc& fn, size_t num std::vector indices; for (size_t i = 0; i < num_args; ++i) { if (auto buffer = fn->buffer_map.Get(fn->params[i])) { - auto buffer_var = buffer.value()->data; + auto buffer_var = buffer.value().var(); if (tirx::UsesVar(fn->body, [=](const tirx::VarNode* var) { return var == buffer_var.get(); })) { indices.push_back(i); diff --git a/src/relax/transform/specialize_primfunc_based_on_callsite.cc b/src/relax/transform/specialize_primfunc_based_on_callsite.cc index 942569410284..f6a99c9efd9a 100644 --- a/src/relax/transform/specialize_primfunc_based_on_callsite.cc +++ b/src/relax/transform/specialize_primfunc_based_on_callsite.cc @@ -37,7 +37,7 @@ namespace tvm { namespace relax { -using tvm::tirx::Buffer; +using tvm::tirx::BufferVar; static ffi::Array GetShapeFromTensorType(const TensorType& tensor_ty) { auto shape = tensor_ty->GetShape(); @@ -80,7 +80,7 @@ class SpecializeTIRCallArgs : ExprMutator { auto gv = call->args[0].as_or_throw(); auto pfunc = mod_->Lookup(gv).as_or_throw(); auto args = call->args[1].as_or_throw()->fields; - ffi::Map> param_map; + ffi::Map> param_map; for (size_t i = 0; i < args.size(); ++i) { auto ty = GetType(args[i]); @@ -99,7 +99,7 @@ class SpecializeTIRCallArgs : ExprMutator { name = std::string({static_cast('A' + i)}); } - const Buffer& buffer = tirx::decl_buffer(GetShapeFromTensorType(tensor_ty), + const BufferVar& buffer = tirx::decl_buffer(GetShapeFromTensorType(tensor_ty), tensor_ty->dtype.value(), name, scope); param_map.Set(pfunc->params[i], buffer); } @@ -110,7 +110,7 @@ class SpecializeTIRCallArgs : ExprMutator { if (ty->vdevice.has_value()) { scope = ty->vdevice.value()->memory_scope; } - const Buffer& buffer = + const BufferVar& buffer = tirx::decl_buffer(GetShapeFromTensorType(ty), ty->dtype.value(), "ret_val", scope); param_map.Set(pfunc->params[pfunc->params.size() - 1], buffer); } else { @@ -132,7 +132,7 @@ class SpecializeTIRCallArgs : ExprMutator { scope = ty->vdevice.value()->memory_scope; } - const Buffer& buffer = tirx::decl_buffer(GetShapeFromTensorType(ty), ty->dtype.value(), + const BufferVar& buffer = tirx::decl_buffer(GetShapeFromTensorType(ty), ty->dtype.value(), "ret_val_" + std::to_string(index), scope); param_map.Set(pfunc->params[args.size() + index], buffer); index++; @@ -141,8 +141,8 @@ class SpecializeTIRCallArgs : ExprMutator { auto new_pfunc = Specialize(pfunc, param_map); for (const auto& [var, buffer] : new_pfunc->buffer_map) { - auto* ptr = buffer->data->ty.as(); - TVM_FFI_ICHECK(ptr) << "Buffer Var's type annotation must be of PointerType"; + auto* ptr = buffer->data_pointer_type.as(); + TVM_FFI_ICHECK(ptr) << "BufferVar Var's type annotation must be of PointerType"; } auto new_prim_func = WithAttr(new_pfunc, "scoped", static_cast(1)); updates_->Add(gv, new_prim_func); diff --git a/src/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc index 14373e043353..a385f0a7e03e 100644 --- a/src/relax/transform/split_call_tir_by_pattern.cc +++ b/src/relax/transform/split_call_tir_by_pattern.cc @@ -81,7 +81,7 @@ class ForMatcher : public TensorizeComparator { } std::vector evaluated_symbols; - std::vector evaluated_buffers; + std::vector evaluated_buffers; private: using ExprComparator::VisitExpr_; @@ -324,7 +324,7 @@ class ForMatcher : public TensorizeComparator { return CompareBufferAccess(op, rhs); } - bool CompareBuffer(const Buffer& lhs, const Buffer& rhs) { + bool CompareBuffer(const BufferVar& lhs, const BufferVar& rhs) { if (lhs.same_as(rhs)) return true; auto it = rhs_buffer_map_.find(rhs); bool equal; @@ -336,9 +336,8 @@ class ForMatcher : public TensorizeComparator { for (size_t i = 0; i < lhs->shape.size(); ++i) { if (!VisitExpr(lhs->shape[i], rhs->shape[i])) return false; } - // Remap both buffer itself and buffer data - equal = - DefEqual(lhs->data, rhs->data) && lhs->dtype == rhs->dtype && lhs.scope() == rhs.scope(); + equal = DefEqual(lhs.var(), rhs.var()) && lhs->dtype == rhs->dtype && + lhs.scope() == rhs.scope(); if (equal) { rhs_buffer_map_[rhs] = lhs; } @@ -449,18 +448,18 @@ class FunctionPartitioner : public StmtExprVisitor { public: explicit FunctionPartitioner(int num_matched_ops) : num_matched_ops_(num_matched_ops) {} /*! \brief alloc_buffers for the first function */ - std::unordered_set allocs1; + std::unordered_set allocs1; /*! \brief alloc_buffers for the second function */ - std::unordered_set allocs2; + std::unordered_set allocs2; /*! \brief whether the current block is in the first function */ ffi::Map block_partition; /*! \brief input buffers for the first function */ - std::unordered_set input1; + std::unordered_set input1; /*! \brief input buffers for the second function */ - std::unordered_set input2; + std::unordered_set input2; /*! \brief The output buffer for the first function, which is also the input buffer for the second function */ - Buffer intermediate_buffer; + BufferVar intermediate_buffer; /*! \brief Indicate whether we have failed. If failed, we will not do any further analysis and directly return the original one. */ bool fail = false; @@ -506,7 +505,7 @@ class BlockRemover : public StmtExprMutator { public: static Stmt RemoveBlockByPartition( Stmt stmt, const ffi::Map& block_partition, - const std::unordered_set& allocs, + const std::unordered_set& allocs, bool is_library_part) { BlockRemover remover(block_partition, allocs, is_library_part); return remover(stmt); @@ -514,7 +513,7 @@ class BlockRemover : public StmtExprMutator { private: BlockRemover(const ffi::Map& block_partition, - const std::unordered_set& allocs, + const std::unordered_set& allocs, bool is_library_part) : block_partition(block_partition), allocs_(allocs), is_library_part_(is_library_part) {} @@ -530,8 +529,8 @@ class BlockRemover : public StmtExprMutator { erased_ = true; } } - ffi::Array alloc_buffers; - for (const Buffer& b : block->alloc_buffers) { + ffi::Array alloc_buffers; + for (const BufferVar& b : block->alloc_buffers) { if (allocs_.count(b)) { alloc_buffers.push_back(b); } @@ -555,7 +554,7 @@ class BlockRemover : public StmtExprMutator { bool erased_ = false; ffi::Map block_partition; - std::unordered_set allocs_; + std::unordered_set allocs_; bool is_library_part_ = false; }; @@ -583,7 +582,7 @@ std::pair> SplitFunctions( TVM_FFI_ICHECK(codegen_result.size() == 3); ffi::String library_code = codegen_result[0].as_or_throw(); int num_matched_ops = codegen_result[1].as_or_throw()->value; - ffi::Array func1_args = codegen_result[2].as_or_throw>(); + ffi::Array func1_args = codegen_result[2].as_or_throw>(); if (num_matched_ops == 0) { return {func, std::nullopt}; } @@ -624,7 +623,7 @@ std::pair> SplitFunctions( } arg_partition->push_back(arg_partition1); new_params1.push_back(Var("output", PointerType::VoidPointerTy())); - ffi::Map new_buffer_map1; + ffi::Map new_buffer_map1; for (const auto& kv : func->buffer_map) { if (partitioner.input1.count(kv.second)) { new_buffer_map1.Set(kv.first, kv.second); @@ -647,7 +646,7 @@ std::pair> SplitFunctions( } } arg_partition->push_back(arg_partition2); - ffi::Map new_buffer_map2; + ffi::Map new_buffer_map2; new_buffer_map2.Set(new_params2[0], partitioner.intermediate_buffer); for (const auto& kv : func->buffer_map) { if (partitioner.input2.count(kv.second)) { @@ -752,7 +751,7 @@ class SplitMutator : public ExprMutator { if (lib_func->IsInstance()) return ffi::GetRef(op); TVM_FFI_ICHECK(lib_func->IsInstance()); builder_->UpdateFunction(gv, lib_func); - tirx::Buffer intermediate_buffer = func1->buffer_map.at(func1->params.back()); + tirx::BufferVar intermediate_buffer = func1->buffer_map.at(func1->params.back()); PrimType dtype = intermediate_buffer->dtype; Call call1(Type::Missing(), call_dps_packed_, {lib_func, Tuple(args1)}, call->attrs, {TensorType(ShapeExpr(intermediate_buffer->shape), dtype)}); diff --git a/src/relax/transform/split_layout_rewrite_preproc.cc b/src/relax/transform/split_layout_rewrite_preproc.cc index be17d8dbf872..ff340d153fec 100644 --- a/src/relax/transform/split_layout_rewrite_preproc.cc +++ b/src/relax/transform/split_layout_rewrite_preproc.cc @@ -62,14 +62,14 @@ class SplitPrimFuncLayoutRewrite : public StmtMutator { // Step 2: Create the params for the new PrimFunc ffi::Array params; - ffi::Map buffer_map; + ffi::Map buffer_map; for (const auto& info : rewrite_infos_) { - params.push_back(Var(info.pre_rewrite_buffer->name, PointerType::VoidPointerTy())); + params.push_back(Var(info.pre_rewrite_buffer.name(), PointerType::VoidPointerTy())); buffer_map.Set(params.back(), info.pre_rewrite_buffer); } for (const auto& info : rewrite_infos_) { - params.push_back(Var(info.post_rewrite_buffer->name, PointerType::VoidPointerTy())); + params.push_back(Var(info.post_rewrite_buffer.name(), PointerType::VoidPointerTy())); buffer_map.Set(params.back(), info.post_rewrite_buffer); } @@ -102,7 +102,7 @@ class SplitPrimFuncLayoutRewrite : public StmtMutator { PrimFunc create_compute_func() const { // Step 1: Create the params for the new PrimFunc ffi::Array params = original_func_->params; - ffi::Map buffer_map = original_func_->buffer_map; + ffi::Map buffer_map = original_func_->buffer_map; for (const auto& info : rewrite_infos_) { const Var& param = params[info.buffer_index]; TVM_FFI_ICHECK(buffer_map[param] == info.pre_rewrite_buffer); @@ -112,7 +112,7 @@ class SplitPrimFuncLayoutRewrite : public StmtMutator { // Step 2: Create the body for the new PrimFunc Stmt body = compute_stmts_.size() == 1 ? compute_stmts_[0] : SeqStmt(compute_stmts_); SBlock original_block = original_func_->body.as()->block; - ffi::Array alloc_buffers; + ffi::Array alloc_buffers; for (const auto& buffer : original_block->alloc_buffers) { auto it = std::find_if(rewrite_infos_.begin(), rewrite_infos_.end(), @@ -190,10 +190,10 @@ class SplitPrimFuncLayoutRewrite : public StmtMutator { << "There should be no alloc buffer in the layout rewrite"; TVM_FFI_ICHECK(op->match_buffers.empty()) << "There should be no match buffer in the layout rewrite"; - const Buffer& preproc_buffer = op->reads[0]->buffer; + const BufferVar& preproc_buffer = op->reads[0]->buffer; int buffer_index = -1; for (size_t i = 0; i < original_func_->params.size(); ++i) { - const Buffer& buffer = original_func_->buffer_map[original_func_->params[i]]; + const BufferVar& buffer = original_func_->buffer_map[original_func_->params[i]]; if (buffer == preproc_buffer) { buffer_index = i; break; @@ -216,8 +216,8 @@ class SplitPrimFuncLayoutRewrite : public StmtMutator { public: struct RewriteInfo { int buffer_index; - Buffer pre_rewrite_buffer; - Buffer post_rewrite_buffer; + BufferVar pre_rewrite_buffer; + BufferVar post_rewrite_buffer; }; std::vector rewrite_infos_; @@ -303,7 +303,7 @@ class SplitLayoutRewritePreproc : public ExprMutator { ffi::Array preproc_ty_list; for (const auto& info : rewrite_infos) { preproc_args.push_back(call_tir_args[info.buffer_index]); - tirx::Buffer rewritten_buffer = info.post_rewrite_buffer; + tirx::BufferVar rewritten_buffer = info.post_rewrite_buffer; for (const auto& shape_expr : rewritten_buffer->shape) { TVM_FFI_ICHECK(shape_expr.as()) << "Currently does not support rewrite buffer with " diff --git a/src/s_tir/analysis/calculate_allocated_memory.cc b/src/s_tir/analysis/calculate_allocated_memory.cc index 3759e7ca421d..09fa73a5e1c2 100644 --- a/src/s_tir/analysis/calculate_allocated_memory.cc +++ b/src/s_tir/analysis/calculate_allocated_memory.cc @@ -41,7 +41,7 @@ using namespace tvm::tirx; std::string GetStorageScope(const Var& var) { auto* ptr = var->ty.as(); - TVM_FFI_ICHECK(ptr) << "Buffer Var's type annotation must be of PointerType"; + TVM_FFI_ICHECK(ptr) << "BufferVar Var's type annotation must be of PointerType"; return ptr->storage_scope; } @@ -61,7 +61,7 @@ class AllocBufferCalculator : public StmtExprVisitor { private: void VisitStmt_(const AllocBufferNode* op) override { - std::string storage_scope = GetStorageScope(op->buffer->data); + std::string storage_scope = op->buffer.scope(); auto search = _current_size.find(storage_scope); if (search == _current_size.end()) { _current_size[storage_scope] = 0; diff --git a/src/s_tir/analysis/is_pure_function.cc b/src/s_tir/analysis/is_pure_function.cc index d7ae4c856ef9..d806df1760b4 100644 --- a/src/s_tir/analysis/is_pure_function.cc +++ b/src/s_tir/analysis/is_pure_function.cc @@ -46,14 +46,14 @@ class PurityChecker : TIRVisitorWithPath { explicit PurityChecker(bool assert_on_error) : assert_on_error_(assert_on_error) {} void VisitStmt_(const AllocBufferNode* op, ffi::reflection::AccessPath path) override { - internal_allocations_.insert(op->buffer->data); + internal_allocations_.insert(op->buffer.var()); TIRVisitorWithPath::VisitStmt_(op, path); } void VisitStmt_(const BufferStoreNode* op, ffi::reflection::AccessPath path) override { TIRVisitorWithPath::VisitStmt_(op, path); - if (!internal_allocations_.count(op->buffer->data)) { + if (!internal_allocations_.count(op->buffer.var())) { is_pure_ = false; if (assert_on_error_) { TVM_FFI_THROW(AssertionError) << "Pure functions must not write to buffers, " diff --git a/src/s_tir/analysis/oob_checker.cc b/src/s_tir/analysis/oob_checker.cc index a37c8387731e..6dc7c0719fe1 100644 --- a/src/s_tir/analysis/oob_checker.cc +++ b/src/s_tir/analysis/oob_checker.cc @@ -32,7 +32,7 @@ namespace s_tir { using namespace tvm::tirx; namespace transform { struct OOBLocation { - Buffer buf; + BufferVar buf; size_t dimension; ffi::ObjectRef index; arith::IntSet index_bounds; @@ -47,7 +47,7 @@ class OOBError : public s_tir::ScheduleError { ffi::String DetailRenderTemplate() const final { std::stringstream s; for (const auto& oob : locations_) { - s << "Out of bounds memory access on buffer " << oob.buf->name << " dimension " + s << "Out of bounds memory access on buffer " << oob.buf.name() << " dimension " << oob.dimension << "."; s << " index " << oob.index << " with bounds [" << oob.index_bounds.min() << ", " << oob.index_bounds.max() << "] is outside the range [0, " << oob.shape_bounds.min() diff --git a/src/s_tir/analysis/sblock_access_region_detector.cc b/src/s_tir/analysis/sblock_access_region_detector.cc index c1656a1e4645..516078f569ef 100644 --- a/src/s_tir/analysis/sblock_access_region_detector.cc +++ b/src/s_tir/analysis/sblock_access_region_detector.cc @@ -42,15 +42,19 @@ namespace tirx { */ class BlockReadWriteDetector : public StmtExprVisitor { public: - explicit BlockReadWriteDetector(const ffi::Map& buffer_var_map) - : buffer_var_map_(buffer_var_map) {} + explicit BlockReadWriteDetector(const ffi::Map& buffer_var_map) + : buffer_var_map_(buffer_var_map) { + for (const auto& [raw_pointer, buffer] : buffer_var_map) { + buffer_var_map_.Set(buffer.var(), buffer); + } + } /*! \brief Return read regions of the block */ ffi::Array CollectReads( - const std::unordered_set* excluded_buffers = nullptr); + const std::unordered_set* excluded_buffers = nullptr); /*! \brief Return write regions of the block */ ffi::Array CollectWrites( - const std::unordered_set* excluded_buffers = nullptr); + const std::unordered_set* excluded_buffers = nullptr); /*! * \brief Return opaque buffer regions of the block * \note The buffer accessed by load/store or call with buffer.data will @@ -68,11 +72,11 @@ class BlockReadWriteDetector : public StmtExprVisitor { /*! \brief Unresolved conditions within current scope. */ std::vector pending_conditions_; /*! \brief The buffers that the current block reads */ - std::vector read_buffers_; + std::vector read_buffers_; /*! \brief The buffers that the current block writes */ - std::vector writes_buffers_; + std::vector writes_buffers_; /*! \brief The opaque buffer which is access by buffer.data */ - std::vector opaque_buffers_; + std::vector opaque_buffers_; /*! \brief The read regions of the current block */ std::vector> read_regions_; /*! \brief The write regions of the current block */ @@ -80,7 +84,7 @@ class BlockReadWriteDetector : public StmtExprVisitor { /*! \brief The opaque regions of the current block */ std::vector> opaque_regions_; /*! \brief The outside buffer data mapping to its buffer */ - ffi::Map buffer_var_map_; + ffi::Map buffer_var_map_; /*! \brief The target buffer var mapping to its matching */ std::unordered_map match_buffers_; /*! \brief let bindings inside the block */ @@ -95,14 +99,14 @@ class BlockReadWriteDetector : public StmtExprVisitor { * \param buffer The provided buffer * \param region The provided region */ - void Update(std::vector* buffers, std::vector>* regions, - Buffer buffer, std::vector region); + void Update(std::vector* buffers, std::vector>* regions, + BufferVar buffer, std::vector region); /*! \brief Helper function to collect access regions. */ ffi::Array CollectRegions( - const std::vector& buffers, + const std::vector& buffers, const std::vector>& regions, - const std::unordered_set* excluded_buffers = nullptr); + const std::unordered_set* excluded_buffers = nullptr); /*! \brief Helper function to convert matched access region to source region. */ std::vector ConvertMatchedRegion(const MatchBufferRegion& match_buffer, @@ -129,8 +133,8 @@ void BlockReadWriteDetector::operator()(const Stmt& stmt) { TVM_FFI_ICHECK(block != nullptr) << "Only visiting Blocks is allowed, but got " << stmt->GetTypeKey(); for (const MatchBufferRegion& match_buffer : block->match_buffers) { - const Var& target_var = match_buffer->buffer->data; - const Var& source_var = match_buffer->source->buffer->data; + const Var target_var = match_buffer->buffer.var(); + const Var source_var = match_buffer->source->buffer.var(); if (buffer_var_map_.find(source_var) != buffer_var_map_.end()) { match_buffers_[target_var.get()] = match_buffer; buffer_var_map_.Set(target_var, match_buffer->buffer); @@ -140,12 +144,12 @@ void BlockReadWriteDetector::operator()(const Stmt& stmt) { } ffi::Array BlockReadWriteDetector::CollectReads( - const std::unordered_set* excluded_buffers) { + const std::unordered_set* excluded_buffers) { return CollectRegions(read_buffers_, read_regions_, excluded_buffers); } ffi::Array BlockReadWriteDetector::CollectWrites( - const std::unordered_set* excluded_buffers) { + const std::unordered_set* excluded_buffers) { return CollectRegions(writes_buffers_, write_regions_, excluded_buffers); } @@ -200,11 +204,15 @@ void BlockReadWriteDetector::VisitStmt_(const BindNode* op) { void BlockReadWriteDetector::VisitExpr_(const CallNode* op) { if (op->op.same_as(builtin::tvm_access_ptr())) { const VarNode* buffer_var = op->args[1].as(); + if (const auto* data = op->args[1].as(); + data && data->op.same_as(builtin::buffer_data())) { + buffer_var = data->args[0].as(); + } const IntImmNode* access_mask = op->args[4].as(); if (buffer_var && access_mask) { auto it = buffer_var_map_.find(ffi::GetRef(buffer_var)); if (it != buffer_var_map_.end()) { - const Buffer& buffer = (*it).second; + const BufferVar& buffer = (*it).second; const BufferRegion buffer_region = BufferRegion::FullRegion(buffer); const ffi::Array& region = buffer_region->region; std::vector int_set; @@ -288,7 +296,7 @@ void BlockReadWriteDetector::VisitStmt_(const SBlockRealizeNode* op) { std::vector BlockReadWriteDetector::ConvertMatchedRegion( const MatchBufferRegion& match_buffer, const std::vector& int_sets) const { - const Buffer& buffer = match_buffer->buffer; + const BufferVar& buffer = match_buffer->buffer; ffi::Array region; region.reserve(int_sets.size()); @@ -308,12 +316,12 @@ std::vector BlockReadWriteDetector::ConvertMatchedRegion( return result; } -void BlockReadWriteDetector::Update(std::vector* buffers, - std::vector>* regions, Buffer buffer, +void BlockReadWriteDetector::Update(std::vector* buffers, + std::vector>* regions, BufferVar buffer, std::vector region) { - if (buffer_var_map_.find(buffer->data) == buffer_var_map_.end()) return; + if (buffer_var_map_.find(buffer.var()) == buffer_var_map_.end()) return; // Handle match_buffer remap - auto it = match_buffers_.find(buffer->data.get()); + auto it = match_buffers_.find(buffer.get()); if (it != match_buffers_.end()) { const MatchBufferRegion& match_buffer = it->second; buffer = match_buffer->source->buffer; @@ -335,8 +343,8 @@ void BlockReadWriteDetector::Update(std::vector* buffers, } ffi::Array BlockReadWriteDetector::CollectRegions( - const std::vector& buffers, const std::vector>& regions, - const std::unordered_set* excluded_buffers) { + const std::vector& buffers, const std::vector>& regions, + const std::unordered_set* excluded_buffers) { TVM_FFI_ICHECK_EQ(buffers.size(), regions.size()); ffi::Array res; res.reserve(buffers.size()); @@ -364,7 +372,7 @@ ffi::Array BlockReadWriteDetector::CollectRegions( void BlockReadWriteDetector::UpdateOpaque(const Var& buffer_var) { auto it = buffer_var_map_.find(buffer_var); if (it != buffer_var_map_.end()) { - const Buffer& buffer = (*it).second; + const BufferVar& buffer = (*it).second; const BufferRegion buffer_region = BufferRegion::FullRegion(buffer); const ffi::Array& region = buffer_region->region; std::vector int_set; @@ -377,11 +385,11 @@ void BlockReadWriteDetector::UpdateOpaque(const Var& buffer_var) { } ffi::Array> GetSBlockAccessRegion( - const SBlock& block, const ffi::Map& buffer_var_map) { + const SBlock& block, const ffi::Map& buffer_var_map) { BlockReadWriteDetector detector(buffer_var_map); detector(block); ffi::Array writes = detector.CollectWrites(); - std::unordered_set excluded_buffers; + std::unordered_set excluded_buffers; // exclude write buffers from read regions for reductions if init block is defined. if (block->init.has_value()) { for (const BufferRegion& write_access : writes) { @@ -394,11 +402,11 @@ ffi::Array> GetSBlockAccessRegion( } ffi::Array> GetSBlockReadWriteRegion( - const SBlock& block, const ffi::Map& buffer_var_map) { + const SBlock& block, const ffi::Map& buffer_var_map) { BlockReadWriteDetector detector(buffer_var_map); detector(block); ffi::Array opaques = detector.CollectOpaques(); - std::unordered_set excluded_buffers; + std::unordered_set excluded_buffers; for (const BufferRegion& opaque_access : opaques) { excluded_buffers.insert(opaque_access->buffer.get()); } 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 6b07bc2469c6..9a0002dd6a3c 100644 --- a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc +++ b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc @@ -34,7 +34,7 @@ namespace tvm { namespace tirx { /*! - * \brief Detect the lowest common ancestor(LCA) position of Buffer access. + * \brief Detect the lowest common ancestor(LCA) position of BufferVar access. * \note * - Only consider BlockNode and ForNode to be the LCA nodes. * - In the LCA locator, we are aware of the buffer scope and CUDA hierarchy so that any buffer in @@ -43,11 +43,11 @@ namespace tirx { */ class LCADetector : public StmtExprVisitor { public: - static ffi::Map> Detect(const PrimFunc& func) { + static ffi::Map> Detect(const PrimFunc& func) { LCADetector detector; for (const auto& kv : func->buffer_map) { - const Buffer& buffer = kv.second; - detector.buffer_var_map_.emplace(buffer->data.get(), buffer.get()); + const BufferVar& buffer = kv.second; + detector.buffer_var_map_.emplace(buffer.get(), buffer.get()); } // The root node must be explicitly present in the list of @@ -61,9 +61,9 @@ class LCADetector : public StmtExprVisitor { detector.UpdateWithBlockidx(); // Prepare the return - ffi::Map> buffer_lca; + ffi::Map> buffer_lca; for (const auto& kv : detector.buffer_lca_) { - const Buffer& buffer = ffi::GetRef(kv.first); + BufferVar buffer(ffi::GetRef(kv.first)); const ffi::Optional stmt = kv.second ? ffi::Optional(ffi::GetRef(kv.second->stmt)) : std::nullopt; buffer_lca.Set(buffer, stmt); @@ -111,8 +111,8 @@ class LCADetector : public StmtExprVisitor { void VisitStmt_(const SBlockRealizeNode* op) final { const SBlockNode* block = op->block.get(); int n = ancestor_scopes_.size(); - for (const Buffer& buf : block->alloc_buffers) { - buffer_var_map_.emplace(buf->data.get(), buf.get()); + for (const BufferVar& buf : block->alloc_buffers) { + buffer_var_map_.emplace(buf.get(), buf.get()); } const ScopeInfo* parent_scope = ancestor_scopes_.back(); @@ -196,7 +196,7 @@ class LCADetector : public StmtExprVisitor { // relate to, which is record in `itervar_to_dom_scope`. auto do_update = [this, &opaque_var_scope, highest_reduce_scope](const BufferRegion& region, bool is_reduce_write = false) { - const Buffer& buffer = region->buffer; + const BufferVar& buffer = region->buffer; const ScopeInfo* scope = ancestor_scopes_.back(); auto handle_itervar = [&opaque_var_scope, &scope](const ffi::ObjectRef& obj) { @@ -274,8 +274,8 @@ class LCADetector : public StmtExprVisitor { } } - void UpdateBufferLCA(const BufferNode* buffer, const ScopeInfo* scope) { - buffer_var_map_.emplace(buffer->data.get(), buffer); + void UpdateBufferLCA(const VarNode* buffer, const ScopeInfo* scope) { + buffer_var_map_.emplace(buffer, buffer); if (match_buffers_.find(buffer) == match_buffers_.end()) { // Ingore buffer created by block match_buffer const ScopeInfo*& lca = buffer_lca_[buffer]; @@ -286,7 +286,7 @@ class LCADetector : public StmtExprVisitor { void UpdateWithBlockidx() { for (const auto& it : buffer_lca_) { const runtime::StorageScope& scope = - runtime::StorageScope::Create(ffi::GetRef(it.first).scope()); + runtime::StorageScope::Create(BufferVar(ffi::GetRef(it.first)).scope()); if (scope.rank == runtime::StorageRank::kGlobal) { const ScopeInfo*& lca = buffer_lca_[it.first]; for (const ScopeInfo* blockidx_scope : blockidx_scopes_) { @@ -326,12 +326,12 @@ class LCADetector : public StmtExprVisitor { * the root scope. */ std::vector ancestor_scopes_ = {}; - /*! \brief The map from Buffer to its LCA ForNode/BlockNode. */ - std::unordered_map buffer_lca_ = {}; - /*! \brief The map from Buffer data to the Buffer. */ - std::unordered_map buffer_var_map_ = {}; + /*! \brief The map from BufferVar to its LCA ForNode/BlockNode. */ + std::unordered_map buffer_lca_ = {}; + /*! \brief The map from BufferVar data to the BufferVar. */ + std::unordered_map buffer_var_map_ = {}; /*! \brief The match buffers inside blocks. */ - std::unordered_set match_buffers_ = {}; + std::unordered_set match_buffers_ = {}; /*! \brief The ForNodes/BlockNodes which contain immediate `blockIdx` launch. */ std::vector blockidx_scopes_ = {}; /*! \brief The map from loop var to the corresponding scope. */ @@ -340,7 +340,7 @@ class LCADetector : public StmtExprVisitor { support::Arena arena_; }; -ffi::Map> DetectBufferAccessLCA(const PrimFunc& func) { +ffi::Map> DetectBufferAccessLCA(const PrimFunc& func) { return LCADetector::Detect(func); } diff --git a/src/s_tir/backend/adreno/inject_texture_alloc.cc b/src/s_tir/backend/adreno/inject_texture_alloc.cc index 1fee1ecf62a9..2876df779915 100644 --- a/src/s_tir/backend/adreno/inject_texture_alloc.cc +++ b/src/s_tir/backend/adreno/inject_texture_alloc.cc @@ -63,7 +63,7 @@ class TextureAllocInjector : public arith::IRMutatorWithAnalyzer { Stmt VisitStmt_(const AllocBufferNode* op) final { Stmt stmt = StmtExprMutator::VisitStmt_(op); - std::string storage_scope = GetStorageScope(op->buffer->data); + std::string storage_scope = op->buffer.scope(); if (IsTextureStorage(storage_scope)) { op = stmt.as(); const auto& extents = op->buffer->shape; @@ -82,8 +82,9 @@ class TextureAllocInjector : public arith::IRMutatorWithAnalyzer { args.push_back(Call(PointerType(PrimType::Int(64)), builtin::tvm_stack_make_shape(), {texture.width, texture.height, texture.depth})); args.push_back(IntImm::Int64(channel_size)); - stmt = Bind(op->buffer->data, - Call(op->buffer->data->ty, builtin::nd_mem_alloc_with_scope(), args)); + stmt = DeclBuffer( + op->buffer, + Call(op->buffer->data_pointer_type, builtin::nd_mem_alloc_with_scope(), args)); } return stmt; } @@ -91,7 +92,7 @@ class TextureAllocInjector : public arith::IRMutatorWithAnalyzer { protected: std::string GetStorageScope(const Var& buffer_var) { auto* ptr = buffer_var->ty.as(); - TVM_FFI_ICHECK(ptr) << "Buffer Var's type annotation must be of PointerType"; + TVM_FFI_ICHECK(ptr) << "BufferVar Var's type annotation must be of PointerType"; return ptr->storage_scope; } }; diff --git a/src/s_tir/backend/adreno/texture_flatten.cc b/src/s_tir/backend/adreno/texture_flatten.cc index 389edf1e8050..1d222d5b6dd4 100644 --- a/src/s_tir/backend/adreno/texture_flatten.cc +++ b/src/s_tir/backend/adreno/texture_flatten.cc @@ -48,7 +48,7 @@ using runtime::IsTextureStorage; class TextureLoweringBase : public StmtExprMutator { public: - explicit TextureLoweringBase(const ffi::Map& extern_buffer_map, + explicit TextureLoweringBase(const ffi::Map& extern_buffer_map, IRVisitorWithAnalyzer* bound_analyzer) : bound_analyzer_{bound_analyzer} { for (auto kv : extern_buffer_map) { @@ -71,14 +71,14 @@ class TextureLoweringBase : public StmtExprMutator { } protected: - std::string GetStorageScope(const Buffer& buffer) { - auto* ptr = buffer->data->ty.as(); - TVM_FFI_ICHECK(ptr) << "Buffer Var's type annotation must be of PointerType"; + std::string GetStorageScope(const BufferVar& buffer) { + auto* ptr = buffer->data_pointer_type.as(); + TVM_FFI_ICHECK(ptr) << "BufferVar Var's type annotation must be of PointerType"; return ptr->storage_scope; } // Set of all external input and output buffers - std::unordered_set extern_buf_; + std::unordered_set extern_buf_; // Bound analzer IRVisitorWithAnalyzer* bound_analyzer_; }; @@ -88,7 +88,7 @@ class TextureLoweringBase : public StmtExprMutator { class TextureFlattener : public TextureLoweringBase { public: using StmtExprMutator::VisitStmt_; - explicit TextureFlattener(const ffi::Map& extern_buffer_map, + explicit TextureFlattener(const ffi::Map& extern_buffer_map, IRVisitorWithAnalyzer* bound_analyzer) : TextureLoweringBase(extern_buffer_map, bound_analyzer) {} @@ -122,12 +122,12 @@ class TextureFlattener : public TextureLoweringBase { protected: template - ffi::Array GetTextureAccessArgs(const T* op, const Buffer& buffer) { + ffi::Array GetTextureAccessArgs(const T* op, const BufferVar& buffer) { ffi::Array args; - if (let_binding_.count(op->buffer->data)) { - args.push_back(let_binding_[op->buffer->data]); + if (let_binding_.count(op->buffer.var())) { + args.push_back(let_binding_[op->buffer.var()]); } else { - args.push_back(buffer->data); + args.push_back(buffer.data()); } ffi::Array row_dims, row_indices, col_dims, col_indices, depth_dims, depth_indices; size_t axis = DefaultTextureLayoutSeparator(op->buffer->shape.size(), GetStorageScope(buffer)); diff --git a/src/s_tir/meta_schedule/arg_info.cc b/src/s_tir/meta_schedule/arg_info.cc index 73fa41773883..94fec3bdf988 100644 --- a/src/s_tir/meta_schedule/arg_info.cc +++ b/src/s_tir/meta_schedule/arg_info.cc @@ -96,8 +96,8 @@ ffi::Array ArgInfo::FromPrimFunc(const tirx::PrimFunc& func) { ffi::Array result; result.reserve(func->params.size()); for (const tirx::Var& arg : func->params) { - if (ffi::Optional _buffer = func->buffer_map.Get(arg)) { - tirx::Buffer buffer = _buffer.value(); + if (ffi::Optional _buffer = func->buffer_map.Get(arg)) { + tirx::BufferVar buffer = _buffer.value(); result.push_back(TensorInfo(/*dtype=*/buffer->dtype->dtype, /*shape=*/AsVector(buffer->shape))); } else { 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 530da5c15fe5..73bbc7eac617 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 @@ -49,7 +49,7 @@ using ForVec = std::vector; * \tparam V The value type */ template -using ForBufferMap = std::unordered_map>; +using ForBufferMap = std::unordered_map>; /*! \brief Given x, compute log2(|x| + 1) */ inline double slog(double x) { return x >= 0 ? std::log2(x + 1) : std::log2(-x + 1); } @@ -62,11 +62,11 @@ namespace utils { * \param analyzer The analyzer * \return The shape of the buffer */ -std::vector GetBufferShape(const Buffer& buffer, arith::AnalyzerObj* analyzer) { - int ndim = buffer->shape.size(); +std::vector GetBufferShape(const BufferVar& buffer, arith::AnalyzerObj* analyzer) { + int ndim = GetBufferVar(buffer)->shape.size(); std::vector result; result.reserve(ndim); - for (const PrimExpr& i : buffer->shape) { + for (const PrimExpr& i : GetBufferVar(buffer)->shape) { if (const IntImmNode* int_imm = i.as()) { result.push_back(int_imm->value); continue; @@ -655,9 +655,9 @@ struct Feature { kUnknownRW = 3, }; enum class ReuseType : int { - /*! Buffer reuse because accessed on each iteration of a loop */ + /*! BufferVar reuse because accessed on each iteration of a loop */ kLoopMultipleRead = 0, - /*! Buffer reuse because it is serially accessed */ + /*! BufferVar reuse because it is serially accessed */ kSerialMultipleReadWrite = 1, /*! No buffer reuse */ kNoReuse = 2, @@ -665,14 +665,14 @@ struct Feature { struct SubFeature { /*! \brief The buffer this feature is for */ - const BufferNode* buffer = nullptr; + const VarNode* buffer = nullptr; /*! \brief The access type of the buffer */ AccessType access_type = AccessType::kUnknownRW; /*! \brief A list of multi-dimensonal indices used to access the buffer */ std::vector multi_indices = {}; // Access information /*! \brief loop_accessed_numel[i][...] means the number of elements accessed by loops[i] */ - std::vector> loop_accessed_numel = {}; + std::vector> loop_accessed_numel = {}; /*! \brief The shape of the data access */ IntVec access_shape; /*! \brief The bytes that are continuously accessed */ @@ -750,7 +750,7 @@ struct Feature { void SetFeature(const LoopNest& loop_nest, int64_t cache_line_bytes); - explicit SubFeature(const BufferNode* buffer, AccessType access_type, + explicit SubFeature(const VarNode* buffer, AccessType access_type, std::vector multi_indices, int n_loops) : buffer(buffer), access_type(access_type), @@ -788,7 +788,7 @@ void Feature::Init(const BufferStoreNode* store, int n_loops) { AccessType access_type = AccessType::kUnknownRW; std::vector multi_indices; }; - std::unordered_map buffer_info; + std::unordered_map buffer_info; { Info& info = buffer_info[store->buffer.get()]; info.access_type = AccessType::kWrite; @@ -796,7 +796,7 @@ void Feature::Init(const BufferStoreNode* store, int n_loops) { } PostOrderVisit(store->value, [&buffer_info](const ffi::ObjectRef& obj) -> void { if (const BufferLoadNode* load = obj.as()) { - const BufferNode* buffer = load->buffer.get(); + const VarNode* buffer = load->buffer.get(); Info& info = buffer_info[buffer]; switch (info.access_type) { case AccessType::kRead: @@ -838,7 +838,7 @@ void Feature::SetRegion(const LoopNest& loop_nest, IntVec* for_touched_bytes, if (n_loops == 0) { // In this case, the `access_shape` is not calculated for (SubFeature& feature : sub_features) { - feature.access_shape = IntVec(feature.buffer->shape.size(), 1); + feature.access_shape = IntVec(GetBufferVar(feature.buffer)->shape.size(), 1); } return; } @@ -850,14 +850,14 @@ void Feature::SetRegion(const LoopNest& loop_nest, IntVec* for_touched_bytes, /*allow_override=*/true); int64_t& touched_bytes = (*for_touched_bytes)[i] = 0; for (SubFeature& feature : sub_features) { - const BufferNode* buffer = feature.buffer; + const VarNode* buffer = feature.buffer; // Note: `feature.access_shape` for `i == 0` is the only one preserved, // while others are discarded int64_t numel; feature.access_shape = utils::RelaxAndUnion(feature.multi_indices, &numel, analyzer); numel = std::max(0, numel); feature.loop_accessed_numel[i][buffer] = numel; - touched_bytes += numel * static_cast(buffer->dtype.StorageBytes()); + touched_bytes += numel * static_cast(GetBufferVar(buffer)->dtype.StorageBytes()); (*buffer_touched_under_loop)[loop][buffer].push_back(numel); } } @@ -867,9 +867,9 @@ void Feature::SubFeature::SetStride(const LoopNest& loop_nest, arith::AnalyzerOb int n_loops = loop_nest.loops.size(); const std::vector& loops = loop_nest.loops; // For each buffer, we find the loop stride on it - const BufferNode* buffer = this->buffer; - int ndim = this->buffer->shape.size(); - IntVec buffer_shape = utils::GetBufferShape(ffi::GetRef(buffer), analyzer); + const VarNode* buffer = this->buffer; + int ndim = GetBufferVar(this->buffer)->shape.size(); + IntVec buffer_shape = utils::GetBufferShape(BufferVar(ffi::GetRef(buffer)), analyzer); // Calculate the buffer's stride from its shape IntVec buffer_stride(ndim); if (ndim >= 1) { @@ -885,7 +885,7 @@ void Feature::SubFeature::SetStride(const LoopNest& loop_nest, arith::AnalyzerOb TVM_FFI_ICHECK_EQ(access_shape.size(), buffer_shape.size()); for (int i = ndim - 1; i >= 0; --i) { if (access_shape[i] == buffer_shape[i]) { - num_continuous_bytes = buffer_shape[i] * static_cast(buffer->dtype.StorageBytes()); + num_continuous_bytes = buffer_shape[i] * static_cast(GetBufferVar(buffer)->dtype.StorageBytes()); break; } } @@ -913,7 +913,7 @@ void Feature::SubFeature::SetStride(const LoopNest& loop_nest, arith::AnalyzerOb void Feature::SubFeature::SetReuse(const LoopNest& loop_nest, int64_t top_loop_touch_bytes, const ForBufferMap& buffer_touched_under_loop) { - const BufferNode* buffer = this->buffer; + const VarNode* buffer = this->buffer; // Step 3.1. Collect all `Var`s that appears in the buffer region std::unordered_set region_vars; for (const MultiIndex& multi_index : this->multi_indices) { @@ -955,10 +955,10 @@ void Feature::SubFeature::SetReuse(const LoopNest& loop_nest, int64_t top_loop_t reuse_dis_bytes = top_loop_touch_bytes; } else { for (const auto& iter : buffer_touched_under_loop.at(loops[i + 1])) { - const BufferNode* buffer = iter.first; + const VarNode* buffer = iter.first; const IntVec& numels = iter.second; int64_t numel = std::accumulate(numels.begin(), numels.end(), int64_t(0)); - reuse_dis_bytes += numel * static_cast(buffer->dtype.StorageBytes()); + reuse_dis_bytes += numel * static_cast(GetBufferVar(buffer)->dtype.StorageBytes()); } } break; @@ -975,10 +975,10 @@ void Feature::SubFeature::SetReuse(const LoopNest& loop_nest, int64_t top_loop_t reuse_dis_iter = *std::min_element(touched.begin(), touched.end()); reuse_dis_bytes = 0.0; for (const auto& iter : buffer_touched_under_loop.at(loop)) { - const BufferNode* buffer = iter.first; + const VarNode* buffer = iter.first; const IntVec& numels = iter.second; int64_t numel = std::accumulate(numels.begin(), numels.end(), int64_t(0)); - reuse_dis_bytes += numel * static_cast(buffer->dtype.StorageBytes()); + reuse_dis_bytes += numel * static_cast(GetBufferVar(buffer)->dtype.StorageBytes()); } reuse_dis_iter /= extent; reuse_dis_bytes /= extent; @@ -988,7 +988,7 @@ void Feature::SubFeature::SetReuse(const LoopNest& loop_nest, int64_t top_loop_t } void Feature::SubFeature::SetFeature(const LoopNest& loop_nest, int64_t cache_line_bytes) { - int64_t dtype_bytes = static_cast(this->buffer->dtype.StorageBytes()); + int64_t dtype_bytes = static_cast(GetBufferVar(this->buffer)->dtype.StorageBytes()); this->stride = this->innermost_stride; this->bytes = dtype_bytes * loop_nest.prod; if (loop_nest.loops.empty()) { @@ -1028,7 +1028,7 @@ Feature::Feature(const BufferStoreNode* store, const LoopNest& loop_nest, int64_ int64_t top_loop_touch_bytes = 0.0; if (n_loops > 0) { for (const SubFeature& feature : sub_features) { - int64_t bytes = static_cast(feature.buffer->dtype.StorageBytes()); + int64_t bytes = static_cast(GetBufferVar(feature.buffer)->dtype.StorageBytes()); int64_t n_buffer = feature.loop_accessed_numel[0].size(); top_loop_touch_bytes += bytes * n_buffer; } @@ -1048,7 +1048,7 @@ Feature::Feature(const BufferStoreNode* store, const LoopNest& loop_nest, int64_ if (a.bytes != b.bytes) { return a.bytes > b.bytes; } - return a.buffer->name < b.buffer->name; + return GetBufferVar(a.buffer).name() < GetBufferVar(b.buffer).name(); }); } @@ -1160,13 +1160,13 @@ struct Feature { Feature() = default; - explicit Feature(const LoopNest& loop_nest, const Buffer& buffer, arith::AnalyzerObj* analyzer) { + explicit Feature(const LoopNest& loop_nest, const BufferVar& buffer, arith::AnalyzerObj* analyzer) { std::vector shape = utils::GetBufferShape(buffer, analyzer); int64_t numel = 1; for (int64_t x : shape) { numel *= x; } - alloc_size = numel * static_cast(buffer->dtype.StorageBytes()); + alloc_size = numel * static_cast(GetBufferVar(buffer)->dtype.StorageBytes()); alloc_prod = numel * loop_nest.prod; alloc_outer_prod = loop_nest.prod; } @@ -1262,7 +1262,7 @@ struct Feature { /*! \brief The feature extracted */ struct Feature { - const BufferNode* buffer = nullptr; + const VarNode* buffer = nullptr; int buffer_order = -1; std::unique_ptr group1 = nullptr; std::unique_ptr group2 = nullptr; @@ -1320,7 +1320,7 @@ class PerStoreFeatureCollector : private StmtVisitor { if (store->value->IsInstance() || store->value->IsInstance()) { return; } - const BufferNode* buffer = store->buffer.get(); + const VarNode* buffer = store->buffer.get(); Feature& feature = buffer_features_[buffer]; if (feature.buffer == nullptr) { feature.buffer = buffer; @@ -1338,12 +1338,12 @@ class PerStoreFeatureCollector : private StmtVisitor { void VisitStmt_(const SBlockNode* block) final { StmtVisitor::VisitStmt_(block); - for (const Buffer& buffer : block->alloc_buffers) { + for (const BufferVar& buffer : block->alloc_buffers) { HandleBufferAlloc(buffer); } } - void HandleBufferAlloc(const Buffer& buffer) { + void HandleBufferAlloc(const BufferVar& buffer) { Feature& feature = buffer_features_[buffer.get()]; feature.group4 = std::make_unique(loop_nest_, buffer, analyzer_.get()); } @@ -1361,7 +1361,7 @@ class PerStoreFeatureCollector : private StmtVisitor { LoopNest loop_nest_ = {}; IntVec for_touched_bytes_ = {}; ForBufferMap buffer_touched_under_loop_ = {}; - std::unordered_map buffer_features_ = {}; + std::unordered_map buffer_features_ = {}; }; } // namespace s_tir 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 be9389437523..352fb33df906 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 @@ -77,10 +77,10 @@ struct AsyncStridedMemCopyFinder : private StmtExprVisitor { } // get store buffer; assert it exists and is contiguous given it uses a single index - auto bufferstore = bufferstorenode->buffer.as(); + auto bufferstore = bufferstorenode->buffer.as(); // get load buffer; assert it exists and is contiguous given it uses a single index - auto bufferload = bufferloadnode->buffer.as(); + auto bufferload = bufferloadnode->buffer.as(); if (!bufferstore || !bufferload) { StmtExprVisitor::VisitStmt_(attrStmt); diff --git a/src/s_tir/meta_schedule/postproc/rewrite_layout.cc b/src/s_tir/meta_schedule/postproc/rewrite_layout.cc index d1f38db1608d..f6a439b650d1 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_layout.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_layout.cc @@ -35,7 +35,7 @@ using namespace tvm::tirx; */ class BufferReadPosCollector : public StmtExprVisitor { public: - explicit BufferReadPosCollector(const Buffer& buffer) : buffer_(buffer.get()) {} + explicit BufferReadPosCollector(const BufferVar& buffer) : buffer_(buffer.get()) {} const std::pair& GetBufferLocation() const { return buffer_loc_; } @@ -58,7 +58,7 @@ class BufferReadPosCollector : public StmtExprVisitor { void VisitExpr_(const BufferLoadNode* op) final { TVM_FFI_ICHECK(cur_realize_.defined()) << "BufferLoad occurred outside of any block"; - const Buffer& buffer = op->buffer; + const BufferVar& buffer = op->buffer; if (buffer_ == buffer.get()) { ffi::Map subst_map; for (size_t i = 0; i < cur_realize_->iter_values.size(); i++) { @@ -81,7 +81,7 @@ class BufferReadPosCollector : public StmtExprVisitor { } } - static int GetReadBufferIndex(const SBlock& block, const Buffer& buffer) { + static int GetReadBufferIndex(const SBlock& block, const BufferVar& buffer) { for (size_t i = 0; i < block->reads.size(); i++) { if (block->reads[i]->buffer.same_as(buffer)) { return i; @@ -92,7 +92,7 @@ class BufferReadPosCollector : public StmtExprVisitor { private: /*! \brief The buffer of interest. */ - const BufferNode* buffer_; + const VarNode* buffer_; /*! \brief The block that consumes the buffer and the corresponding read index. */ std::pair buffer_loc_; /*! \brief The proposed IndexMap. */ @@ -111,21 +111,21 @@ class LayoutFreeBufferCollector : public StmtVisitor { void VisitStmt_(const SBlockNode* block) final { StmtVisitor::VisitStmt_(block); if (auto ann = block->annotations.Get("layout_free_placeholders")) { - for (Buffer buffer : ann.value().as_or_throw>()) { + for (BufferVar buffer : ann.value().as_or_throw>()) { buffers.insert(buffer); } } } - std::unordered_set buffers; + std::unordered_set buffers; }; -ffi::Array CollectLayoutFreeBuffers(const PrimFuncNode* func) { +ffi::Array CollectLayoutFreeBuffers(const PrimFuncNode* func) { // Only rewrite PrimFuncs with attr "layout_free_buffers" ffi::Array layout_free_buffer_index = func->GetAttr(s_tir::attr::layout_free_buffers, ffi::Array()).value(); - ffi::Array layout_free_buffers; + ffi::Array layout_free_buffers; for (int64_t index : layout_free_buffer_index) { TVM_FFI_ICHECK(static_cast(index) < func->params.size()); const Var& param = func->params[index]; @@ -142,7 +142,7 @@ ffi::Array CollectLayoutFreeBuffers(const PrimFuncNode* func) { } std::optional> GetSuggestedIndexMap( - Buffer buffer, const PrimFuncNode* prim_func) { + BufferVar buffer, const PrimFuncNode* prim_func) { BufferReadPosCollector collector(buffer); collector(prim_func->body); @@ -158,10 +158,10 @@ std::optional> GetSuggestedIndexMap( } /*! \brief Get a chain of cache-read blocks, starting from the one consuming buf. */ -std::vector GetCacheReadChain(const Buffer& buf, const PrimFuncNode* prim_func) { +std::vector GetCacheReadChain(const BufferVar& buf, const PrimFuncNode* prim_func) { class BufferReadChainCollector : public StmtVisitor { public: - explicit BufferReadChainCollector(const Buffer& buffer) : cur_buffer_(buffer.get()) {} + explicit BufferReadChainCollector(const BufferVar& buffer) : cur_buffer_(buffer.get()) {} void VisitStmt_(const SBlockNode* op) final { // Check if this block is doing cache_read or a similar operation that consumes cur_buffer_. @@ -176,7 +176,7 @@ std::vector GetCacheReadChain(const Buffer& buf, const PrimFuncNode std::vector cache_read_chain; private: - const BufferNode* cur_buffer_; + const VarNode* cur_buffer_; }; BufferReadChainCollector collector(buf); diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc index 1cd504dfee68..50cffabb86ff 100644 --- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc +++ b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc @@ -35,13 +35,13 @@ using namespace tvm::tirx; std::vector GetReadBufferNDims(const StmtSRef& block_sref) { const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); - const BufferNode* write_buffer = block->writes[0]->buffer.get(); + const VarNode* write_buffer = block->writes[0]->buffer.get(); int n = block->reads.size(); std::vector results(n, -1); for (int i = 0; i < n; ++i) { - const BufferNode* read_buffer = block->reads[i]->buffer.get(); + const VarNode* read_buffer = block->reads[i]->buffer.get(); if (read_buffer != write_buffer) { - results[i] = read_buffer->shape.size(); + results[i] = GetBufferVar(read_buffer)->shape.size(); } } return results; diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc index c966358d5ce1..1d5fc2c45e6b 100644 --- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc +++ b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc @@ -623,7 +623,7 @@ std::vector MultiLevelTilingTensorCoreNode::AddReadReuseTensorCore( // Inline the reindex / padding block sch->ComputeInline(sch->GetProducers(cache_read)[0]); const tirx::SBlockNode* cache_read_block = sch->GetSRef(cache_read)->StmtAs(); - tirx::Buffer cache_read_buffer = + tirx::BufferVar cache_read_buffer = s_tir::GetNthAccessBuffer(sch->state(), ffi::GetRef(cache_read_block), 0, s_tir::BufferIndexType::kWrite); const DLDataType dtype = cache_read_buffer->dtype->dtype; @@ -824,11 +824,11 @@ ffi::Optional MultiLevelTilingTensorCoreNode::TransformWithTensorIntrin( rhs_to_index_map_tgt[mapping_info->rhs_iters[i - offset]->var] = index_map->final_indices[i]; } - auto f_get_sub_index_map = [&](const tirx::Buffer& lhs_buffer, + auto f_get_sub_index_map = [&](const tirx::BufferVar& lhs_buffer, const ffi::Array& lhs_region) { std::vector sub_index_map_src; std::vector sub_index_map_tgt; - const tirx::Buffer& rhs_buffer = mapping_info->lhs_buffer_map[lhs_buffer]; + const tirx::BufferVar& rhs_buffer = mapping_info->lhs_buffer_map[lhs_buffer]; for (const Range& range : lhs_region) { TVM_FFI_ICHECK(tirx::is_one(range->extent)); auto var = range->min.as(); @@ -847,13 +847,13 @@ ffi::Optional MultiLevelTilingTensorCoreNode::TransformWithTensorIntrin( return tirx::IndexMap(sub_index_map_src, sub_index_map_tgt); }; - std::unordered_set visited_buffers; + std::unordered_set visited_buffers; - ffi::Map buffer_sub_index_map; // cache of the sub index map + ffi::Map buffer_sub_index_map; // cache of the sub index map // associated with each buffer auto f_transform_buffer_layout = [&](s_tir::BufferIndexType index_type, int buffer_index) { - const tirx::Buffer& lhs_buffer = s_tir::GetNthAccessBuffer( + const tirx::BufferVar& lhs_buffer = s_tir::GetNthAccessBuffer( state->sch->state(), block_before_reindex, buffer_index, index_type); if (visited_buffers.count(lhs_buffer)) { return; @@ -879,7 +879,7 @@ ffi::Optional MultiLevelTilingTensorCoreNode::TransformWithTensorIntrin( // Transform the layout of current block and reindex blocks auto f_transform_reindex_block_layout = [&](const SBlockRV& block_rv, s_tir::BufferIndexType buffer_type) { - tirx::Buffer buffer = + tirx::BufferVar buffer = s_tir::GetNthAccessBuffer(state->sch->state(), state->sch->Get(block_rv), 0, buffer_type); const auto& sub_index_map = buffer_sub_index_map.at(buffer); state->sch->TransformBlockLayout(block_rv, sub_index_map); diff --git a/src/s_tir/sblock_scope.cc b/src/s_tir/sblock_scope.cc index 606ab2034d20..5f61cf1daa4b 100644 --- a/src/s_tir/sblock_scope.cc +++ b/src/s_tir/sblock_scope.cc @@ -82,8 +82,8 @@ SBlockScope::SBlockScope() { data_ = ffi::make_object(); } SBlockScope::SBlockScope(const ffi::Array& child_block_srefs) { ffi::ObjectPtr n = ffi::make_object(); - SMap> buffer_readers; - SMap>& buffer_writers = n->buffer_writers; + SMap> buffer_readers; + SMap>& buffer_writers = n->buffer_writers; for (const StmtSRef& child_block_sref : child_block_srefs) { const SBlockNode* child_block = TVM_SREF_TO_SBLOCK(child_block_sref); // Step 1. Update `buffer_readers` and `buffer_writers` for each buffer diff --git a/src/s_tir/schedule/analysis.h b/src/s_tir/schedule/analysis.h index 27454e5e6434..4f1065bfe585 100644 --- a/src/s_tir/schedule/analysis.h +++ b/src/s_tir/schedule/analysis.h @@ -450,7 +450,7 @@ struct ProducerConsumerSplit { * \return The buffer of the n-th read/write region of the block. * \throw ScheduleError If the buffer index is out of bound. */ -Buffer GetNthAccessBuffer(const ScheduleState& self, const SBlock& block, int n, +BufferVar GetNthAccessBuffer(const ScheduleState& self, const SBlock& block, int n, BufferIndexType index_type); /*! @@ -473,7 +473,7 @@ BufferRegion GetNthAccessBufferRegion(const ScheduleState& self, const SBlock& b * buffer is from match_buffer). */ std::pair, bool> GetBufferDefiningSite(const StmtSRef& block_sref, - const Buffer& buffer); + const BufferVar& buffer); /******** Reduction SBlock Related ********/ @@ -600,7 +600,7 @@ bool CanReverseComputeAt(const ScheduleState& self, const StmtSRef& block_sref, * \param predicate The predicate of the access * \param analyzer Arithmetic analyzer */ -ffi::Optional SuggestIndexMap(const Buffer& buffer, const ffi::Array& indices, +ffi::Optional SuggestIndexMap(const BufferVar& buffer, const ffi::Array& indices, const ffi::Array& loops, const PrimExpr& predicate, arith::AnalyzerObj* analyzer); @@ -788,9 +788,9 @@ class AutoTensorizeMappingInfoNode : public ffi::Object { /* Additional information from AutoTensorizeComparator */ /*! \brief Mapping from LHS buffer to RHS buffer */ - ffi::Map lhs_buffer_map; - /*! \brief Buffer indices on RHS */ - ffi::Map> rhs_buffer_indices; + ffi::Map lhs_buffer_map; + /*! \brief BufferVar indices on RHS */ + ffi::Map> rhs_buffer_indices; /*! \brief SBlock iters on LHS */ ffi::Array lhs_iters; /*! \brief SBlock iters on RHS */ diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc index d1cbb409a257..83f7c31c3c09 100644 --- a/src/s_tir/schedule/analysis/analysis.cc +++ b/src/s_tir/schedule/analysis/analysis.cc @@ -181,7 +181,7 @@ void CheckSRefHigherOrEqual(const StmtSRef& sref_a, const StmtSRef& sref_b) { */ bool IsDominantBlock(const ScheduleState& self, const StmtSRef& scope_root_sref, const StmtSRef& block_sref) { - std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> + std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> buffer_writers; CheckSRefHigherOrEqual(scope_root_sref, block_sref); const SBlockNode* maybe_root_block = scope_root_sref->StmtAs(); @@ -234,7 +234,7 @@ int CheckCompleteBlockErrorCode(const ScheduleState& self, const StmtSRef& block return 2; } // Cond 3. No overlap between the buffers the block reads and writes - std::unordered_set written_buffers; + std::unordered_set written_buffers; written_buffers.reserve(block->writes.size()); for (const BufferRegion& write : block->writes) { written_buffers.insert(write->buffer.get()); @@ -484,9 +484,9 @@ bool IsOutputBlock(const ScheduleState& self, const StmtSRef& block_sref, const StmtSRef& scope_root_sref) { const SBlockNode* scope_root = TVM_SREF_TO_SBLOCK(scope_root_sref); const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); - std::unordered_set scope_allocated; + std::unordered_set scope_allocated; scope_allocated.reserve(scope_root->alloc_buffers.size()); - for (const Buffer& buffer : scope_root->alloc_buffers) { + for (const BufferVar& buffer : scope_root->alloc_buffers) { scope_allocated.insert(buffer.get()); } for (const BufferRegion& buffer_region : block->writes) { @@ -1263,13 +1263,13 @@ BufferRegion GetNthAccessBufferRegion(const ScheduleState& self, const SBlock& b return access_region[n]; } -Buffer GetNthAccessBuffer(const ScheduleState& self, const SBlock& block, int n, +BufferVar GetNthAccessBuffer(const ScheduleState& self, const SBlock& block, int n, BufferIndexType index_type) { return GetNthAccessBufferRegion(self, block, n, index_type)->buffer; } std::pair, bool> GetBufferDefiningSite(const StmtSRef& block_sref, - const Buffer& buffer) { + const BufferVar& buffer) { // Climb up along the sref tree, and find the block where `buffer` is in alloc_buffers or // match_buffers. const StmtSRefNode* defining_site_sref = block_sref.get(); @@ -1281,7 +1281,7 @@ std::pair, bool> GetBufferDefiningSite(const StmtSRef& b continue; } // Try to find the buffer in `allloc_buffers` - for (const Buffer& alloc_buffer : block->alloc_buffers) { + for (const BufferVar& alloc_buffer : block->alloc_buffers) { if (buffer.same_as(alloc_buffer)) { return {ffi::GetRef(defining_site_sref), true}; } @@ -1315,7 +1315,7 @@ void AddShapeVarBounds(const ScheduleState& state, const StmtSRefNode* sref, } const PrimFuncNode* f = GetRootPrimFunc(state->mod, sref->stmt, nullptr); for (const auto& kv : f->buffer_map) { - const Buffer& buffer = kv.second; + const BufferVar& buffer = kv.second; for (const PrimExpr& e : buffer->shape) { analyzer->MarkGlobalNonNegValue(e); } @@ -1528,7 +1528,7 @@ bool NeedsMultiLevelTiling(const ScheduleState& self, const StmtSRef& block_sref !IsTrivialBinding(self, block_sref)) { return false; } - const BufferNode* write_buffer = block->writes[0]->buffer.get(); + const VarNode* write_buffer = block->writes[0]->buffer.get(); // Step 1. Sort out spatial block variables. Skip the block iters of domain [0, 1), since such // block iters distracts the following check of the unused block iters. std::vector spatial_block_vars; @@ -1545,10 +1545,10 @@ bool NeedsMultiLevelTiling(const ScheduleState& self, const StmtSRef& block_sref // Step 2. Enumerate each read region, check the number of block vars that are not used // to index the read region int total_unused_block_vars = 0; - std::unordered_set read_buffers; + std::unordered_set read_buffers; read_buffers.reserve(block->reads.size()); for (const BufferRegion& buffer_region : block->reads) { - const BufferNode* buffer = buffer_region->buffer.get(); + const VarNode* buffer = buffer_region->buffer.get(); const ffi::Array& regions = buffer_region->region; // Step 2.1. Duplication of read buffers are not allowed if (read_buffers.insert(buffer).second == false) { @@ -1973,13 +1973,13 @@ class AutoTensorizeMappingProposer { using BufferMask = std::vector; // Step 1: Assign an index to each buffer in LHS and RHS - std::unordered_map rhs_buffer_index; - std::unordered_map lhs_buffer_index; + std::unordered_map rhs_buffer_index; + std::unordered_map lhs_buffer_index; { int i = 0; for (const auto& kv : extractor_->rhs_buffer_map_) { - const Buffer& rhs_buffer = kv.first; - const Buffer& lhs_buffer = kv.second; + const BufferVar& rhs_buffer = kv.first; + const BufferVar& lhs_buffer = kv.second; rhs_buffer_index[rhs_buffer] = i; lhs_buffer_index[lhs_buffer] = i; ++i; @@ -2000,20 +2000,20 @@ class AutoTensorizeMappingProposer { }; for (const auto& it : extractor_->rhs_buffer_indices_map_) { - const Buffer& rhs_buffer = it.first; + const BufferVar& rhs_buffer = it.first; for (const PrimExpr& rhs_index : it.second) { if (auto var = rhs_index.as()) { update_mask(var.value().get(), &rhs_buffer_masks, rhs_buffer_index.at(rhs_buffer)); } else { TVM_FFI_THROW(ValueError) - << "Buffer index " << rhs_index + << "BufferVar index " << rhs_index << " other that variables in tensor intrinsics is not supported."; } } auto lhs_buffer_it = extractor_->rhs_buffer_map_.find(rhs_buffer); TVM_FFI_ICHECK(lhs_buffer_it != extractor_->rhs_buffer_map_.end()); - const Buffer& lhs_buffer = lhs_buffer_it->second; + const BufferVar& lhs_buffer = lhs_buffer_it->second; for (const PrimExpr& index : extractor_->lhs_buffer_indices_map_.at(lhs_buffer)) { PreOrderVisit(index, [&](const ffi::ObjectRef& obj) -> bool { if (auto var = obj.as()) { diff --git a/src/s_tir/schedule/analysis/layout.cc b/src/s_tir/schedule/analysis/layout.cc index 5619422ff26d..12efe4122431 100644 --- a/src/s_tir/schedule/analysis/layout.cc +++ b/src/s_tir/schedule/analysis/layout.cc @@ -30,7 +30,7 @@ using namespace tvm::tirx; * \param buffer The buffer * \return The strides */ -ffi::Array GetStrides(const Buffer& buffer) { +ffi::Array GetStrides(const BufferVar& buffer) { if (!buffer->strides.empty()) { TVM_FFI_ICHECK_EQ(buffer->strides.size(), buffer->shape.size()); return buffer->strides; @@ -129,7 +129,7 @@ class SplitExprCollector { std::vector exprs_; }; -ffi::Optional SuggestIndexMap(const Buffer& buffer, const ffi::Array& indices, +ffi::Optional SuggestIndexMap(const BufferVar& buffer, const ffi::Array& indices, const ffi::Array& loops, const PrimExpr& predicate, arith::AnalyzerObj* analyzer) { int ndim = buffer->shape.size(); @@ -249,7 +249,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def( "s_tir.schedule.SuggestIndexMap", - [](Buffer buffer, ffi::Array indices, ffi::Array loops, PrimExpr predicate) { + [](BufferVar buffer, ffi::Array indices, ffi::Array loops, PrimExpr predicate) { arith::Analyzer analyzer; return SuggestIndexMap(buffer, indices, loops, predicate, analyzer.get()); }); diff --git a/src/s_tir/schedule/analysis/reducer.cc b/src/s_tir/schedule/analysis/reducer.cc index 4eb5affe7063..f5ba1b300469 100644 --- a/src/s_tir/schedule/analysis/reducer.cc +++ b/src/s_tir/schedule/analysis/reducer.cc @@ -360,7 +360,7 @@ void ErrorRFactorCrossThreadReductionNotApplicable(const ffi::Optional& self, SBlock block, const ffi::Array& stmts, int n_buffers, ffi::Array* updates, - std::unordered_map* buf2index) { + std::unordered_map* buf2index) { std::unordered_map var2index; ffi::Array let_values; let_values.reserve(n_buffers); @@ -441,7 +441,7 @@ std::pair, ffi::Array> GetInitValuesAndUpdates if (auto init = block->init.as()) { inits.push_back(init.value()); } else if (const auto* seq_init = block->init.as()) { - std::unordered_set init_buffers; + std::unordered_set init_buffers; for (const Stmt& stmt : seq_init->seq) { auto init = stmt.as(); if (!init) { @@ -459,7 +459,7 @@ std::pair, ffi::Array> GetInitValuesAndUpdates // Step 2. Extract the block updates, in the form of BufferStores. int n_buffers = inits.size(); - std::unordered_map buf2index; + std::unordered_map buf2index; if (const auto* update = block->body.as()) { updates.push_back(ffi::GetRef(update)); buf2index[update->buffer.get()] = 0; @@ -537,15 +537,15 @@ bool ReductionIterNotIndexOutputBuffer(const SBlock& block) { } } // Step 2. Check if the reduction block iters are used to index the output buffer. - std::unordered_set buffer_written; + std::unordered_set buffer_written; buffer_written.reserve(block->writes.size()); for (const BufferRegion& write_region : block->writes) { buffer_written.insert(write_region->buffer.get()); } - std::unordered_set buffer_allocated; + std::unordered_set buffer_allocated; buffer_allocated.reserve(block->alloc_buffers.size()); - for (const Buffer& buffer : block->alloc_buffers) { + for (const BufferVar& buffer : block->alloc_buffers) { buffer_allocated.insert(buffer.get()); } @@ -555,7 +555,7 @@ bool ReductionIterNotIndexOutputBuffer(const SBlock& block) { }); }; - std::unordered_map match_buffer_sources; + std::unordered_map match_buffer_sources; for (const MatchBufferRegion& region : block->match_buffers) { match_buffer_sources[region->buffer.get()] = region->source->buffer.get(); } diff --git a/src/s_tir/schedule/concrete_schedule.cc b/src/s_tir/schedule/concrete_schedule.cc index 2443d7966e7a..b22b41d25f85 100644 --- a/src/s_tir/schedule/concrete_schedule.cc +++ b/src/s_tir/schedule/concrete_schedule.cc @@ -131,9 +131,9 @@ class ScheduleCopier { return result; } - /*! \brief Copy SMap> */ - SMap> Copy(const SMap>& map) { - SMap> result; + /*! \brief Copy SMap> */ + SMap> Copy(const SMap>& map) { + SMap> result; result.reserve(map.size()); for (const auto& kv : map) { result[kv.first] = Copy(kv.second); @@ -1067,7 +1067,7 @@ void ConcreteScheduleNode::PadEinsum(const SBlockRV& block_rv, const ffi::Array< this->state_->DebugVerify(); } -/******** Schedule: Buffer Transformation ********/ +/******** Schedule: BufferVar Transformation ********/ void ConcreteScheduleNode::RollingBuffer(const SBlockRV& block_rv, int write_buffer_index) { TVM_TIR_SCHEDULE_BEGIN(); diff --git a/src/s_tir/schedule/concrete_schedule.h b/src/s_tir/schedule/concrete_schedule.h index ab78f8af923d..899c15b3a24b 100644 --- a/src/s_tir/schedule/concrete_schedule.h +++ b/src/s_tir/schedule/concrete_schedule.h @@ -183,7 +183,7 @@ class ConcreteScheduleNode : public ScheduleNode { void TransformBlockLayout(const SBlockRV& block_rv, const IndexMap& index_map) override; /******** Schedule: Padding decomposition ********/ SBlockRV DecomposePadding(const SBlockRV& block_rv, const LoopRV& loop_rv) override; - /******** Schedule: Buffer transformation ********/ + /******** Schedule: BufferVar transformation ********/ void RollingBuffer(const SBlockRV& block_rv, int write_buffer_index) override; /******** Schedule: Misc ********/ void EnterPostproc() override {} diff --git a/src/s_tir/schedule/ir_comparator.cc b/src/s_tir/schedule/ir_comparator.cc index 3917601f518a..269952e0cd0c 100644 --- a/src/s_tir/schedule/ir_comparator.cc +++ b/src/s_tir/schedule/ir_comparator.cc @@ -480,22 +480,22 @@ bool TensorizeComparator::CompareAnnotationMap(const ffi::Mapdata, rhs->data) && lhs->dtype == rhs->dtype && lhs.scope() == rhs.scope(); + DefEqual(lhs.var(), rhs.var()) && lhs->dtype == rhs->dtype && lhs.scope() == rhs.scope(); if (equal) { rhs_buffer_map_[rhs] = lhs; } else { if (assert_mode_) { std::ostringstream os; - os << "CompareBuffer buffer mismatch. data: " << lhs->data << " vs " << rhs->data + os << "CompareBuffer buffer mismatch: " << lhs << " vs " << rhs << ", dtypes: " << lhs->dtype << " vs " << rhs->dtype << ", scope(): " << lhs.scope() << " vs " << rhs.scope(); EmitError(os.str()); @@ -581,7 +581,7 @@ bool TensorizeComparator::CompareBufferRegion(const BufferRegion& lhs, const Buf if (!lhs_analyzer_->CanProveEqual(indices_base[i], lhs->region[i]->min)) { if (assert_mode_) { std::ostringstream os; - os << "Buffer base index consistency check failed due to unequal index base: " + os << "BufferVar base index consistency check failed due to unequal index base: " "indices_base[i]=" << indices_base[i] << " vs lhs->region[i]->min=" << lhs->region[i]->min; EmitError(os.str()); @@ -734,7 +734,7 @@ bool AutoTensorizeComparator::VisitStmt_(const SBlockNode* op, const Stmt& other return VisitStmt(op->body, rhs->body); } -bool AutoTensorizeComparator::CompareBuffer(const Buffer& lhs, const Buffer& rhs) { +bool AutoTensorizeComparator::CompareBuffer(const BufferVar& lhs, const BufferVar& rhs) { if (lhs.same_as(rhs)) return true; auto it = rhs_buffer_map_.find(rhs); bool equal; @@ -744,16 +744,16 @@ bool AutoTensorizeComparator::CompareBuffer(const Buffer& lhs, const Buffer& rhs // Remap both buffer itself and buffer data, skipping buffer shape and storage scope. Auto // tensorization inserts the cache stages that move workload buffers into an intrinsic's // required scope, while the pointer element type must still agree. - auto data_it = equal_map_.find(lhs->data); + auto data_it = equal_map_.find(lhs.var()); if (data_it != equal_map_.end()) { - equal = data_it->second.same_as(rhs->data); + equal = data_it->second.same_as(rhs.var()); } else { - const auto* lhs_ptr = lhs->data->ty.as(); - const auto* rhs_ptr = rhs->data->ty.as(); + const auto* lhs_ptr = lhs->data_pointer_type.as(); + const auto* rhs_ptr = rhs->data_pointer_type.as(); equal = lhs_ptr && rhs_ptr && ffi::StructuralEqual()(lhs_ptr->element_type, rhs_ptr->element_type); if (equal) { - equal_map_[lhs->data] = rhs->data; + equal_map_[lhs.var()] = rhs.var(); } } equal = equal && lhs->dtype == rhs->dtype; diff --git a/src/s_tir/schedule/ir_comparator.h b/src/s_tir/schedule/ir_comparator.h index aabe4239f1dd..4aeab6c0028f 100644 --- a/src/s_tir/schedule/ir_comparator.h +++ b/src/s_tir/schedule/ir_comparator.h @@ -79,15 +79,15 @@ class TensorizeComparator : public ExprComparator, public StmtComparator { bool VisitExpr_(const SelectNode* op, const PrimExpr& other) override; /*! \brief Map from RHS buffer to LHS buffer */ - std::unordered_map rhs_buffer_map_; + std::unordered_map rhs_buffer_map_; /*! \brief Base indices of the LHS buffer. */ - std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> + std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> buffer_indices_; protected: bool DefEqual(const Var& lhs, const Var& rhs); bool CompareExpr(const Expr& lhs, const Expr& rhs); - virtual bool CompareBuffer(const Buffer& lhs, const Buffer& rhs); + virtual bool CompareBuffer(const BufferVar& lhs, const BufferVar& rhs); bool CompareBufferRegion(const BufferRegion& lhs, const BufferRegion& rhs); bool CompareAnnotation(const std::pair& lhs, const std::pair& rhs); @@ -144,7 +144,7 @@ class AutoTensorizeComparator : public TensorizeComparator { bool VisitExpr_(const BufferLoadNode* op, const PrimExpr& other) override; - bool CompareBuffer(const Buffer& lhs, const Buffer& rhs) override; + bool CompareBuffer(const BufferVar& lhs, const BufferVar& rhs) override; template bool CompareBufferAccess(const T* lhs, const T* rhs); @@ -156,13 +156,13 @@ class AutoTensorizeComparator : public TensorizeComparator { /*! \brief SBlock iters in the RHS stmt. */ std::vector rhs_iters_; /*! \brief The buffer and its access indices in the LHS stmt. */ - std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> + std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> lhs_buffer_indices_map_; /*! \brief The buffer and its access indices in the RHS stmt. */ - std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> + std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> rhs_buffer_indices_map_; /*! \brief Map from LHS buffer to RHS buffer */ - std::unordered_map lhs_buffer_map_; + std::unordered_map lhs_buffer_map_; private: /*! \brief The domain of the inner block iters. */ diff --git a/src/s_tir/schedule/primitive.h b/src/s_tir/schedule/primitive.h index 0c2c02af86e3..1d1f9cf4ce60 100644 --- a/src/s_tir/schedule/primitive.h +++ b/src/s_tir/schedule/primitive.h @@ -640,7 +640,7 @@ TVM_DLL void Unannotate(ScheduleState self, const StmtSRef& sref, const ffi::Str /*! * \brief Apply a transformation represented by IndexMap to buffer * \details The indices and the access region to the target buffer is transformed by the given - * index_map. The index_map is also used to infer the new shape of the buffer. Buffer must be + * index_map. The index_map is also used to infer the new shape of the buffer. BufferVar must be * one of the parameter of the function, or allocated in some blocks (it cannot be a buffer * subregion created via match_buffer). * \param self The state of the schedule @@ -691,7 +691,7 @@ TVM_DLL StmtSRef DecomposePadding(ScheduleState self, const StmtSRef& block_sref */ TVM_DLL void PadEinsum(ScheduleState self, const StmtSRef& block_sref, const ffi::Array& padding); -/******** Schedule: Buffer transformation ********/ +/******** Schedule: BufferVar transformation ********/ /*! * \brief Compute the target buffer via rolling buffering. * \details This primitive selects the outermost rollable axis with a positive bound overlap that diff --git a/src/s_tir/schedule/primitive/annotate_buffer_access.cc b/src/s_tir/schedule/primitive/annotate_buffer_access.cc index 0e2448f87cec..cbe626cac8fe 100644 --- a/src/s_tir/schedule/primitive/annotate_buffer_access.cc +++ b/src/s_tir/schedule/primitive/annotate_buffer_access.cc @@ -27,7 +27,7 @@ using namespace tvm::tirx; class AnnotateRegionRewriter : public StmtExprMutator { public: - AnnotateRegionRewriter(Buffer buffer, int buffer_index, BufferRegion new_region, + AnnotateRegionRewriter(BufferVar buffer, int buffer_index, BufferRegion new_region, BufferIndexType buffer_index_type) : buffer_(buffer), buffer_index_(buffer_index), @@ -39,9 +39,9 @@ class AnnotateRegionRewriter : public StmtExprMutator { ffi::Array regions = buffer_index_type_ == BufferIndexType::kWrite ? block->writes : block->reads; - TVM_FFI_ICHECK_GE(buffer_index_, 0) << "Buffer index must be non-negative"; + TVM_FFI_ICHECK_GE(buffer_index_, 0) << "BufferVar index must be non-negative"; TVM_FFI_ICHECK_LT(buffer_index_, static_cast(regions.size())) - << "Buffer index out of range"; + << "BufferVar index out of range"; regions.Set(buffer_index_, new_region_); ffi::ObjectPtr n = CopyOnWrite(block.get()); @@ -79,7 +79,7 @@ class AnnotateRegionRewriter : public StmtExprMutator { } private: - Buffer buffer_; + BufferVar buffer_; int buffer_index_; BufferRegion new_region_; BufferIndexType buffer_index_type_; @@ -88,7 +88,7 @@ class AnnotateRegionRewriter : public StmtExprMutator { void AnnotateBufferAccess(ScheduleState self, const StmtSRef& block_sref, int buffer_index, BufferIndexType buffer_index_type, const IndexMap& index_map) { const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); - Buffer buffer = + BufferVar buffer = GetNthAccessBuffer(self, ffi::GetRef(block), buffer_index, buffer_index_type); arith::Analyzer analyzer; diff --git a/src/s_tir/schedule/primitive/block_annotate.cc b/src/s_tir/schedule/primitive/block_annotate.cc index af59b02ba3fb..5ec65c97b2f2 100644 --- a/src/s_tir/schedule/primitive/block_annotate.cc +++ b/src/s_tir/schedule/primitive/block_annotate.cc @@ -30,7 +30,7 @@ using namespace tvm::tirx; class StorageAlignAxisOutOfRangeError : public ScheduleError { public: - explicit StorageAlignAxisOutOfRangeError(IRModule mod, Buffer buffer, int axis) + explicit StorageAlignAxisOutOfRangeError(IRModule mod, BufferVar buffer, int axis) : mod_(std::move(mod)), buffer_(std::move(buffer)), axis_(axis) {} ffi::String FastErrorString() const final { @@ -42,7 +42,7 @@ class StorageAlignAxisOutOfRangeError : public ScheduleError { ffi::String DetailRenderTemplate() const final { std::ostringstream os; int ndim = static_cast(buffer_->shape.size()); - os << "The buffer to set storage alignment of, " << buffer_->name << ", has " << ndim + os << "The buffer to set storage alignment of, " << buffer_.name() << ", has " << ndim << " dimension(s), so `axis` is required to be in [" << -(ndim) << ", " << ndim << ") for storage_align. However, the input `axis` is " << axis_ << ", which is out of the expected range."; @@ -52,7 +52,7 @@ class StorageAlignAxisOutOfRangeError : public ScheduleError { IRModule mod() const final { return mod_; } ffi::Array LocationsOfInterest() const final { return {}; } - static int CheckAndUpdate(const IRModule& mod, const Buffer& buffer, int axis) { + static int CheckAndUpdate(const IRModule& mod, const BufferVar& buffer, int axis) { int ndim = static_cast(buffer->shape.size()); if (axis < -ndim || axis >= ndim) { throw StorageAlignAxisOutOfRangeError(mod, buffer, axis); @@ -66,13 +66,13 @@ class StorageAlignAxisOutOfRangeError : public ScheduleError { private: IRModule mod_; - Buffer buffer_; + BufferVar buffer_; int axis_; }; class NonAllocatedBufferError : public ScheduleError { public: - explicit NonAllocatedBufferError(IRModule mod, Buffer buffer) : mod_(mod), buffer_(buffer) {} + explicit NonAllocatedBufferError(IRModule mod, BufferVar buffer) : mod_(mod), buffer_(buffer) {} ffi::String FastErrorString() const final { return "ScheduleError: The input buffer is not allocated by a block. This means the buffer is " @@ -81,14 +81,14 @@ class NonAllocatedBufferError : public ScheduleError { ffi::String DetailRenderTemplate() const final { std::ostringstream os; - os << "The input buffer " << buffer_->name + os << "The input buffer " << buffer_.name() << " is not allocated by a block. This means the buffer is either a function parameter or " "defined in `match_buffer` of a block."; return os.str(); } static StmtSRef CheckAndGetBufferAllocationSite(const IRModule& mod, const StmtSRef& block_sref, - const Buffer& buffer) { + const BufferVar& buffer) { auto [defining_site_sref, is_alloc] = GetBufferDefiningSite(block_sref, buffer); if (!defining_site_sref.has_value() || !is_alloc) { throw NonAllocatedBufferError(mod, buffer); @@ -102,7 +102,7 @@ class NonAllocatedBufferError : public ScheduleError { private: IRModule mod_; - Buffer buffer_; + BufferVar buffer_; }; class StorageAlignInvalidFactorError : public ScheduleError { @@ -196,25 +196,25 @@ class StorageScopeMutator : private ReplaceBufferMutator { * \param block_sref_reuse The block sref reuse map to be updated * \return The new block after the mutation */ - static SBlock Mutate(const SBlock& allocate_site, const Buffer& old_buffer, + static SBlock Mutate(const SBlock& allocate_site, const BufferVar& old_buffer, const ffi::String& storage_scope, ffi::Map* block_sref_reuse) { - Buffer new_buffer = WithScope(old_buffer, storage_scope); + BufferVar new_buffer = WithScope(old_buffer, storage_scope); StorageScopeMutator mutator(old_buffer, new_buffer, storage_scope, block_sref_reuse); Stmt new_block = mutator.VisitStmt(allocate_site); return new_block.as_or_throw(); } private: - StorageScopeMutator(const Buffer& old_buffer, Buffer new_buffer, ffi::String storage_scope, + StorageScopeMutator(const BufferVar& old_buffer, BufferVar new_buffer, ffi::String storage_scope, ffi::Map* block_sref_reuse) : ReplaceBufferMutator(old_buffer, std::move(new_buffer), block_sref_reuse) {} MatchBufferRegion VisitMatchBufferRegion(const MatchBufferRegion& match_buffer) final { - auto it = buffer_var_map_.find(match_buffer->source->buffer->data.get()); + auto it = buffer_var_map_.find(match_buffer->source->buffer.get()); if (it != buffer_var_map_.end()) { - Buffer new_target_buffer = WithScope(match_buffer->buffer, it->second.scope()); - buffer_var_map_[match_buffer->buffer->data.get()] = new_target_buffer; + BufferVar new_target_buffer = WithScope(match_buffer->buffer, it->second.scope()); + buffer_var_map_[match_buffer->buffer.get()] = new_target_buffer; return MatchBufferRegion(new_target_buffer, BufferRegion(it->second, match_buffer->source->region)); } else { @@ -226,7 +226,7 @@ class StorageScopeMutator : private ReplaceBufferMutator { void StorageAlign(ScheduleState self, const StmtSRef& block_sref, int buffer_index, int axis, int factor, int offset) { const SBlockNode* block_ptr = TVM_SREF_TO_SBLOCK(block_sref); - Buffer buffer = GetNthAccessBuffer(self, ffi::GetRef(block_ptr), buffer_index, + BufferVar buffer = GetNthAccessBuffer(self, ffi::GetRef(block_ptr), buffer_index, BufferIndexType::kWrite); StorageAlignInvalidFactorError::Check(self->mod, factor); axis = StorageAlignAxisOutOfRangeError::CheckAndUpdate(self->mod, buffer, axis); @@ -261,7 +261,7 @@ void StorageAlign(ScheduleState self, const StmtSRef& block_sref, int buffer_ind void SetScope(ScheduleState self, const StmtSRef& block_sref, int buffer_index, const ffi::String& storage_scope) { const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); - Buffer buffer = + BufferVar buffer = GetNthAccessBuffer(self, ffi::GetRef(block), buffer_index, BufferIndexType::kWrite); // Step 1. If `storage_scope` equals the original storage scope of the buffer, just return. @@ -298,26 +298,26 @@ class DTypeMutator : private ReplaceBufferMutator { * \param block_sref_reuse The block sref reuse map to be updated * \return The new block after the mutation */ - static SBlock Mutate(const SBlock& allocate_site, const Buffer& old_buffer, PrimType dtype, + static SBlock Mutate(const SBlock& allocate_site, const BufferVar& old_buffer, PrimType dtype, ffi::Map* block_sref_reuse) { - Buffer new_buffer = WithDType(old_buffer, dtype); + BufferVar new_buffer = WithDType(old_buffer, dtype); DTypeMutator mutator(old_buffer, new_buffer, dtype, block_sref_reuse); Stmt new_block = mutator.VisitStmt(allocate_site); return new_block.as_or_throw(); } private: - DTypeMutator(const Buffer& old_buffer, Buffer new_buffer, PrimType dtype, + DTypeMutator(const BufferVar& old_buffer, BufferVar new_buffer, PrimType dtype, ffi::Map* block_sref_reuse) : ReplaceBufferMutator(old_buffer, std::move(new_buffer), block_sref_reuse), src_dtype_(old_buffer->dtype), tgt_dtype_(dtype) {} MatchBufferRegion VisitMatchBufferRegion(const MatchBufferRegion& match_buffer) final { - auto it = buffer_var_map_.find(match_buffer->source->buffer->data.get()); + auto it = buffer_var_map_.find(match_buffer->source->buffer.get()); if (it != buffer_var_map_.end()) { - Buffer new_target_buffer = WithDType(match_buffer->buffer, it->second->dtype); - buffer_var_map_[match_buffer->buffer->data.get()] = new_target_buffer; + BufferVar new_target_buffer = WithDType(match_buffer->buffer, it->second->dtype); + buffer_var_map_[match_buffer->buffer.get()] = new_target_buffer; return MatchBufferRegion(new_target_buffer, BufferRegion(it->second, match_buffer->source->region)); } else { @@ -327,7 +327,7 @@ class DTypeMutator : private ReplaceBufferMutator { Stmt VisitStmt_(const BufferStoreNode* op) final { BufferStore node = StmtExprMutator::VisitStmt_(op).as_or_throw(); - auto it = buffer_var_map_.find(node->buffer->data.get()); + auto it = buffer_var_map_.find(node->buffer.get()); if (it != buffer_var_map_.end()) { node.CopyOnWrite()->buffer = it->second; node.CopyOnWrite()->value = Cast(tgt_dtype_, node->value); @@ -337,7 +337,7 @@ class DTypeMutator : private ReplaceBufferMutator { Expr VisitExpr_(const BufferLoadNode* op) final { BufferLoad node = StmtExprMutator::VisitExpr_(op).as_or_throw(); - auto it = buffer_var_map_.find(node->buffer->data.get()); + auto it = buffer_var_map_.find(node->buffer.get()); if (it != buffer_var_map_.end()) { return Cast(src_dtype_, BufferLoad(it->second, node->indices)); } @@ -350,7 +350,7 @@ class DTypeMutator : private ReplaceBufferMutator { void UnsafeSetDType(ScheduleState self, const StmtSRef& block_sref, int buffer_index, const ffi::String& dtype) { const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); - Buffer buffer = + BufferVar buffer = GetNthAccessBuffer(self, ffi::GetRef(block), buffer_index, BufferIndexType::kWrite); PrimType target_dtype(ffi::StringToDLDataType(dtype)); diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc b/src/s_tir/schedule/primitive/blockize_tensorize.cc index 4054edb9e901..4f2f15a230c0 100644 --- a/src/s_tir/schedule/primitive/blockize_tensorize.cc +++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc @@ -432,7 +432,7 @@ ffi::Array EvalSetRegions(const ffi::Array& regions, ffi::Array results; results.reserve(regions.size()); for (const BufferRegion& buffer_region : regions) { - const Buffer& buffer = buffer_region->buffer; + const BufferVar& buffer = buffer_region->buffer; ffi::Array relaxed = arith::EvalSet(buffer_region->region, dom_map); TVM_FFI_ICHECK_EQ(relaxed.size(), buffer->shape.size()); int ndim = buffer->shape.size(); @@ -453,9 +453,9 @@ ffi::Array EvalSetRegions(const ffi::Array& regions, */ ffi::Array UnionRegions(const ffi::Array& regions) { typedef std::vector> ranges_t; - std::unordered_map intset_map; + std::unordered_map intset_map; for (const BufferRegion& buffer_region : regions) { - const Buffer& buffer = buffer_region->buffer; + const BufferVar& buffer = buffer_region->buffer; if (intset_map.find(buffer) == intset_map.end()) { intset_map[buffer] = {buffer->shape.size(), ffi::Array()}; } @@ -467,7 +467,7 @@ ffi::Array UnionRegions(const ffi::Array& regions) { } ffi::Array results; for (const auto& it : intset_map) { - const Buffer& buffer = it.first; + const BufferVar& buffer = it.first; ffi::Array regions; for (size_t dim = 0; dim < buffer->shape.size(); ++dim) { const arith::IntSet intset = arith::Union(it.second[dim]); @@ -790,24 +790,24 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int TensorizeComparator comparator(self->mod, /*assert_mode=*/true); comparator.VisitStmt(block_realize, intrin_desc->body); // Step 3: Prepare necessary mapping - // 1) Buffer mapping from intrin impl buffers to intrin desc buffers. - // 2) Buffer mapping from intrin impl buffers to buffers in the current AST. + // 1) BufferVar mapping from intrin impl buffers to intrin desc buffers. + // 2) BufferVar mapping from intrin impl buffers to buffers in the current AST. // 3) Mapping impl buffers to their accessed regions. - std::unordered_map impl2desc; + std::unordered_map impl2desc; TVM_FFI_ICHECK_EQ(intrin_desc->params.size(), intrin_impl->params.size()); for (int i = 0, n = intrin_desc->params.size(); i < n; ++i) { - const Buffer& desc = intrin_desc->buffer_map[intrin_desc->params[i]]; - const Buffer& impl = intrin_impl->buffer_map[intrin_impl->params[i]]; + const BufferVar& desc = intrin_desc->buffer_map[intrin_desc->params[i]]; + const BufferVar& impl = intrin_impl->buffer_map[intrin_impl->params[i]]; impl2desc[impl] = desc; } - std::unordered_map impl2cur; + std::unordered_map impl2cur; for (const auto& pair : impl2desc) { - const Buffer& impl = pair.first; - const Buffer& desc = pair.second; + const BufferVar& impl = pair.first; + const BufferVar& desc = pair.second; TVM_FFI_ICHECK(comparator.rhs_buffer_map_.count(desc)); impl2cur[impl] = comparator.rhs_buffer_map_[desc]; } - std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> + std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> impl2region; SBlock impl_block = intrin_impl->body.as_or_throw()->block; for (const BufferRegion& read : impl_block->reads) { @@ -821,8 +821,8 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int ffi::Array match_buffer_regions; match_buffer_regions.reserve(intrin_impl->params.size()); for (int i = 0, n = intrin_impl->params.size(); i < n; ++i) { - const Buffer& impl = intrin_impl->buffer_map.at(intrin_impl->params[i]); - const Buffer& cur = impl2cur.at(impl); + const BufferVar& impl = intrin_impl->buffer_map.at(intrin_impl->params[i]); + const BufferVar& cur = impl2cur.at(impl); const ffi::Array& old_region = impl2region.at(impl); const std::vector& indices_base = comparator.buffer_indices_.at(cur); int offset = static_cast(indices_base.size()) - static_cast(old_region.size()); diff --git a/src/s_tir/schedule/primitive/cache_index.cc b/src/s_tir/schedule/primitive/cache_index.cc index d3775e0ca330..e089a58135ba 100644 --- a/src/s_tir/schedule/primitive/cache_index.cc +++ b/src/s_tir/schedule/primitive/cache_index.cc @@ -36,7 +36,7 @@ struct IndexInfo { /*! \brief Record the common subexpr extract threshold */ size_t cse_thresh; /*! \brief The cache buffer to store the precomputed index */ - std::vector cache_buffer; + std::vector cache_buffer; /*! \brief The expr to be precomputed */ std::vector index_exprs; /*! \brief The range of the loop vars relating to index computation */ @@ -270,7 +270,7 @@ ffi::Array MakeIndexCacheStage(IndexInfo* info, const ffi::String& stora arith::EvalSet(info->var_binding.at(it), arith::AsIntSet(info->range_map)).max() + 1); } info->cache_buffer.push_back( - Buffer(index_buffer_var, data_ty, buffer_shape, {1}, {0}, index_buffer_var->name, 0, 0)); + BufferVar(index_buffer_var, data_ty, buffer_shape, {1}, {0}, index_buffer_var->name, 0, 0)); // Create loop vars and block vars' binding_value std::vector loop_vars; @@ -400,7 +400,7 @@ class CacheIndexRewriter : public StmtExprMutator { // If so, put buffer allocation and insert cache stages on the parent scope ffi::ObjectPtr n = ffi::make_object(*stmt.as()); n->body = InsertIndexStage(n->body, info_->loc_pos, info_->cache_stage); - for (const Buffer& it : info_->cache_buffer) { + for (const BufferVar& it : info_->cache_buffer) { n->alloc_buffers.push_back(it); } stmt = SBlock(n); diff --git a/src/s_tir/schedule/primitive/cache_read_write.cc b/src/s_tir/schedule/primitive/cache_read_write.cc index c1d2025e9b4e..f5483119e6c7 100644 --- a/src/s_tir/schedule/primitive/cache_read_write.cc +++ b/src/s_tir/schedule/primitive/cache_read_write.cc @@ -33,7 +33,7 @@ using namespace tvm::tirx; class NotSingleWriteBlock : public ScheduleError { public: - explicit NotSingleWriteBlock(IRModule mod, Buffer buffer, ffi::Array write_blocks) + explicit NotSingleWriteBlock(IRModule mod, BufferVar buffer, ffi::Array write_blocks) : mod_(std::move(mod)), buffer_(std::move(buffer)) { TVM_FFI_ICHECK_GT(write_blocks.size(), 1); write_blocks_.reserve(write_blocks.size()); @@ -49,7 +49,7 @@ class NotSingleWriteBlock : public ScheduleError { ffi::String DetailRenderTemplate() const final { size_t k = write_blocks_.size(); - return "The buffer " + buffer_->name + " is expected to be written by single block, but got " + + return "The buffer " + buffer_.name() + " is expected to be written by single block, but got " + std::to_string(k) + " blocks who write it."; } @@ -60,7 +60,7 @@ class NotSingleWriteBlock : public ScheduleError { private: IRModule mod_; - Buffer buffer_; + BufferVar buffer_; ffi::Array write_blocks_; }; @@ -69,11 +69,11 @@ class NotSingleWriteBlock : public ScheduleError { /*! \brief The auxiliary info used for the insertion point and content of the cache stage. */ struct CacheStageInfo { /*! \brief The buffer to be read. */ - Buffer read_buffer; + BufferVar read_buffer; /*! \brief The buffer to be written. */ - Buffer write_buffer; + BufferVar write_buffer; /*! \brief The buffer allocation to be inserted into the block signature. */ - ffi::Optional alloc; + ffi::Optional alloc; /*! \brief The AST node whose body is where the cache stage should be inserted. */ StmtSRef loc_sref; /*! \brief The index to insert the cache_read/cache_write stage. */ @@ -90,7 +90,7 @@ struct CacheStageInfo { /*! \brief Return the buffer region related with the buffer */ ffi::Optional GetBufferRegionFromBuffer( - const ffi::Array& buffer_regions, const Buffer& buffer) { + const ffi::Array& buffer_regions, const BufferVar& buffer) { ffi::Optional res = std::nullopt; for (const auto& region : buffer_regions) { if (region->buffer.same_as(buffer)) { @@ -203,7 +203,7 @@ SBlock MakeReindexCacheStage(const BufferRegion& cache_region, ReindexCacheStage /*iter_vars*/ std::move(block_vars), /*reads=*/{BufferRegion(info->read_buffer, read_access_region)}, /*writes=*/{BufferRegion(info->write_buffer, write_access_region)}, - /*name_hint*/ cache_region->buffer->name + "_" + storage_scope, + /*name_hint*/ cache_region->buffer.name() + "_" + storage_scope, /*body=*/ BufferStore(info->write_buffer, BufferLoad(info->read_buffer, read_access_indices), write_access_indices), @@ -303,7 +303,7 @@ SBlock MakeCacheStage(const BufferRegion& cache_region, CacheStageInfo* info, /*iter_vars=*/std::move(block_vars), /*reads=*/{BufferRegion(info->read_buffer, read_access_region)}, /*writes=*/{BufferRegion(info->write_buffer, write_access_region)}, - /*name_hint=*/cache_region->buffer->name + "_" + storage_scope, + /*name_hint=*/cache_region->buffer.name() + "_" + storage_scope, /*body=*/ BufferStore(info->write_buffer, BufferLoad(info->read_buffer, read_access_indices), write_access_indices), @@ -402,7 +402,7 @@ SBlock MakeReIndexStage(const SBlock& block, CacheStageInfo* info, /*iter_vars=*/new_block_iters, /*reads=*/{BufferRegion::FromPoint(info->read_buffer, src_indices)}, /*writes=*/{BufferRegion::FromPoint(info->write_buffer, dst_indices)}, - /*name_hint=*/info->write_buffer->name + "_reindex", + /*name_hint=*/info->write_buffer.name() + "_reindex", /*body=*/ BufferStore(info->write_buffer, BufferLoad(info->read_buffer, src_indices), dst_indices)); @@ -494,7 +494,7 @@ Stmt InsertCacheStage(const Stmt& stmt, int pos, const Stmt& stage) { * \throw NotSingleWriteBlock if there are more than one interested block. */ ffi::Optional GetOnlyWriteBlock(ScheduleState self, const StmtSRef& scope_sref, - const Buffer& buffer) { + const BufferVar& buffer) { SBlockScope scope = self->GetSBlockScope(scope_sref); auto it = scope->buffer_writers.find(buffer); if (it == scope->buffer_writers.end()) { @@ -519,7 +519,7 @@ ffi::Optional GetOnlyWriteBlock(ScheduleState self, const StmtSRef& sc * \return A boolean indicating if all the consumer blocks of the input buffer * meet the requirement. */ -bool AllConsumersUnderStmt(ScheduleState self, Buffer buffer, StmtSRef scope_sref, +bool AllConsumersUnderStmt(ScheduleState self, BufferVar buffer, StmtSRef scope_sref, StmtSRef stmt_sref) { // Collect all children blocks of the target stmt. std::unordered_set blocks_under_target; @@ -558,10 +558,10 @@ bool AllConsumersUnderStmt(ScheduleState self, Buffer buffer, StmtSRef scope_sre * \param index_type Whether to look for reads (kRead) or writes (kWrite). * \return The OR-combination of all nested block predicates found. */ -static PrimExpr CollectNestedBlockPredicates(const Stmt& body, const Buffer& buffer, +static PrimExpr CollectNestedBlockPredicates(const Stmt& body, const BufferVar& buffer, BufferIndexType index_type) { struct Collector : public StmtVisitor { - Collector(const Buffer& buf, BufferIndexType idx_type) + Collector(const BufferVar& buf, BufferIndexType idx_type) : buffer_(buf), index_type_(idx_type), result_(IntImm::Bool(false)), found_(false) {} void VisitStmt_(const SBlockRealizeNode* realize) final { @@ -594,7 +594,7 @@ static PrimExpr CollectNestedBlockPredicates(const Stmt& body, const Buffer& buf StmtVisitor::VisitStmt_(realize); } - const Buffer& buffer_; + const BufferVar& buffer_; BufferIndexType index_type_; PrimExpr result_; bool found_; @@ -624,7 +624,7 @@ BufferRegion RelaxBufferRegion(ScheduleState self, const BufferRegion& buffer_re PrimExpr extra_predicate = IntImm::Bool(true)) { SBlockRealize realize = GetSBlockRealize(self, block_sref); ffi::Map binding = GetBindings(realize); - const Buffer& buffer = buffer_region->buffer; + const BufferVar& buffer = buffer_region->buffer; arith::Analyzer analyzer; BufferRegion subst_region = BufferRegion(buffer, Substitute(buffer_region->region, binding)); ffi::Array int_sets = AnalyzeRegionUpperBound( @@ -1037,8 +1037,8 @@ class CacheReadRewriter : public StmtExprMutator { } Expr VisitExpr_(const VarNode* op) final { - if (op == info_->read_buffer->data.get()) { - return info_->write_buffer->data; + if (op == info_->read_buffer.get()) { + return info_->write_buffer.var(); } return ffi::GetRef(op); } @@ -1321,8 +1321,8 @@ class CacheWriteRewriter : public StmtExprMutator { } Expr VisitExpr_(const VarNode* op) final { - if (op == info_->write_buffer->data.get()) { - return info_->read_buffer->data; + if (op == info_->write_buffer.get()) { + return info_->read_buffer.var(); } return ffi::GetRef(op); } @@ -1439,10 +1439,9 @@ class ReindexCacheWriteRewriter : public CacheWriteRewriter { * \param covered Set of block iter vars covered by the buffer access indices * \return The new buffer with target shape. */ -Buffer CreateReindexBuffer(const Buffer& buffer, const ffi::Array& block_iters, +BufferVar CreateReindexBuffer(const BufferVar& buffer, const ffi::Array& block_iters, const std::unordered_set& covered) { - ffi::ObjectPtr new_buffer = ffi::make_object(*buffer.get()); - ffi::ObjectPtr new_var = ffi::make_object(*buffer->data.get()); + ffi::ObjectPtr new_buffer = CopyBufferType(buffer); std::vector new_shape; std::vector new_strides; for (const auto& iter : block_iters) { @@ -1453,9 +1452,7 @@ Buffer CreateReindexBuffer(const Buffer& buffer, const ffi::Array& bloc new_strides.clear(); new_buffer->shape = new_shape; new_buffer->strides = new_strides; - new_buffer->data = buffer->data.CopyWithSuffix("_reindex"); - new_buffer->name = buffer->name + "_reindex"; - return Buffer(new_buffer); + return RebuildBufferVar(buffer, std::move(new_buffer), buffer.name() + "_reindex"); } /*! @@ -1487,7 +1484,7 @@ class InvalidBufferAccessError : public ScheduleError { kOpaqueAccess, // opaque access to the buffer }; - InvalidBufferAccessError(IRModule mod, Buffer buffer, SBlock block, ErrorKind kind) + InvalidBufferAccessError(IRModule mod, BufferVar 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 BufferLoad or BufferStore. The " @@ -1496,7 +1493,7 @@ class InvalidBufferAccessError : public ScheduleError { ffi::String DetailRenderTemplate() const final { std::ostringstream os; - os << "The target buffer " << buffer_->name + os << "The target buffer " << buffer_.name() << " should be accessed in the leaf block {0} via BufferLoad or BufferStore. The indices " "should be the same if there are multiple accesses to the target buffer. "; if (kind_ == ErrorKind::kNoAccess) { @@ -1513,7 +1510,7 @@ class InvalidBufferAccessError : public ScheduleError { private: IRModule mod_; - Buffer buffer_; + BufferVar buffer_; SBlock block_; ErrorKind kind_; }; @@ -1521,7 +1518,7 @@ class InvalidBufferAccessError : public ScheduleError { /*! \brief Collect the related Load/Store to reindex */ class ReIndexCollector : public StmtExprVisitor { public: - static ffi::Array Collect(const IRModule& mod, const Buffer& buffer, + static ffi::Array Collect(const IRModule& mod, const BufferVar& buffer, const SBlock& block) { ReIndexCollector collector(mod, buffer, block); collector(block->body); @@ -1533,7 +1530,7 @@ class ReIndexCollector : public StmtExprVisitor { } private: - explicit ReIndexCollector(const IRModule& mod, const Buffer& buffer, const SBlock& block) + explicit ReIndexCollector(const IRModule& mod, const BufferVar& buffer, const SBlock& block) : mod_(mod), buffer_(buffer), block_(block) {} void VisitExpr_(const BufferLoadNode* load) final { @@ -1568,7 +1565,7 @@ class ReIndexCollector : public StmtExprVisitor { } void VisitExpr_(const VarNode* var) final { - if (var == buffer_->data.get()) { + if (var == buffer_.get()) { throw InvalidBufferAccessError(mod_, buffer_, block_, InvalidBufferAccessError::ErrorKind::kOpaqueAccess); } @@ -1576,7 +1573,7 @@ class ReIndexCollector : public StmtExprVisitor { /*! \brief The IR module */ IRModule mod_; /*! \brief The buffer to rewrite */ - Buffer buffer_; + BufferVar buffer_; /*! \brief The block to visit */ SBlock block_; /*! \brief The indices of buffer acess to rewrite */ @@ -1674,16 +1671,16 @@ class ReIndexRewriter : public StmtExprMutator { /*! \brief Whether the current block is scope block */ bool is_scope_{true}; /*! \brief The buffer to be replaced */ - Buffer old_buffer_; + BufferVar old_buffer_; /*! \brief The reindex buffer */ - Buffer new_buffer_; + BufferVar new_buffer_; /*! \brief The new indices */ ffi::Array indices_; /*! \brief The new region */ Region region_; }; -void CheckRegionCover(const ScheduleState& self, StmtSRef scope_root, Buffer read_buffer) { +void CheckRegionCover(const ScheduleState& self, StmtSRef scope_root, BufferVar read_buffer) { class NotRegionCoverError : public ScheduleError { public: explicit NotRegionCoverError(IRModule mod, SBlock block) : mod_(mod), block_(block) {} @@ -1734,7 +1731,7 @@ StmtSRef CacheRead(ScheduleState self, const StmtSRef& block_sref, int read_buff // Step 1. Check index, getting the target buffer and the parent scope const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); - Buffer read_buffer = GetNthAccessBuffer(self, ffi::GetRef(block), read_buffer_index, + BufferVar read_buffer = GetNthAccessBuffer(self, ffi::GetRef(block), read_buffer_index, BufferIndexType::kRead); StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false); // Check required region cover for cache_read @@ -1800,12 +1797,13 @@ StmtSRef CacheRead(ScheduleState self, const StmtSRef& block_sref, int read_buff info.cache_region = cache_region; info.write_buffer = WithScope(read_buffer, storage_scope); if (!cache_full_region) { - auto* write_buffer = info.write_buffer.CopyOnWrite(); + auto write_buffer = CopyBufferType(info.write_buffer); std::vector shape; for (auto cache_range : info.cache_region->region) { shape.push_back(cache_range->extent); } write_buffer->shape = std::move(shape); + info.write_buffer = RebuildBufferVar(info.write_buffer, std::move(write_buffer)); } info.alloc = info.write_buffer; @@ -1843,7 +1841,7 @@ StmtSRef CacheWrite(ScheduleState self, const StmtSRef& block_sref, int write_bu // Step 1. Checking index, getting the target buffer and the parent scope const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); - Buffer write_buffer = GetNthAccessBuffer(self, ffi::GetRef(block), write_buffer_index, + BufferVar write_buffer = GetNthAccessBuffer(self, ffi::GetRef(block), write_buffer_index, BufferIndexType::kWrite); StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false); @@ -1889,12 +1887,13 @@ StmtSRef CacheWrite(ScheduleState self, const StmtSRef& block_sref, int write_bu info.cache_region = cache_region; info.read_buffer = WithScope(write_buffer, storage_scope); if (!cache_full_region) { - auto* read_buffer = info.read_buffer.CopyOnWrite(); + auto read_buffer_type = CopyBufferType(info.read_buffer); std::vector shape; for (auto cache_range : info.cache_region->region) { shape.push_back(cache_range->extent); } - read_buffer->shape = std::move(shape); + read_buffer_type->shape = std::move(shape); + info.read_buffer = RebuildBufferVar(info.read_buffer, std::move(read_buffer_type)); } info.alloc = info.read_buffer; @@ -1985,7 +1984,7 @@ template void CollectReindexCacheStageInfoAndCreateBuffer( ReindexCacheStageInfo* info, const IRModule& mod, const StmtSRef& block_sref, const ffi::String& storage_scope, const IndexMap& index_map, const SBlock& block, - const SBlockRealize& realize, const Buffer& old_buffer, const BufferRegion& cache_region) { + const SBlockRealize& realize, const BufferVar& old_buffer, const BufferRegion& cache_region) { arith::Analyzer analyzer; ffi::Array block_iter_vars, block_shape; for (const IterVar& iter_var : block->iter_vars) { @@ -2035,19 +2034,18 @@ void CollectReindexCacheStageInfoAndCreateBuffer( } // Create new buffer - ffi::ObjectPtr new_buffer = ffi::make_object(*old_buffer.get()); - ffi::ObjectPtr new_var = ffi::make_object(*old_buffer->data.get()); - const auto* ptr_type = TVM_TYPE_AS(old_buffer->data->ty, PointerTypeNode); - new_var->ty = PointerType(ptr_type->element_type, storage_scope); - new_buffer->data = Var(new_var->name + "_" + storage_scope, new_var->ty); - new_buffer->name = old_buffer->name + "_" + storage_scope; + ffi::ObjectPtr new_buffer = CopyBufferType(old_buffer); + const auto* ptr_type = TVM_TYPE_AS(old_buffer->data_pointer_type, PointerTypeNode); + new_buffer->data_pointer_type = PointerType(ptr_type->element_type, storage_scope); new_buffer->shape = new_shape; + BufferVar rebuilt = + RebuildBufferVar(old_buffer, std::move(new_buffer), old_buffer.name() + "_" + storage_scope); if (is_cache_read) { - info->write_buffer = Buffer(new_buffer); + info->write_buffer = rebuilt; info->alloc = info->write_buffer; } else { - info->read_buffer = Buffer(new_buffer); + info->read_buffer = rebuilt; info->alloc = info->read_buffer; } } @@ -2086,7 +2084,7 @@ StmtSRef ReindexCacheRead(ScheduleState self, const StmtSRef& block_sref, int re // Step 1. Check index, getting the target buffer and the parent scope SBlock block = ffi::GetRef(TVM_SREF_TO_SBLOCK(block_sref)); SBlockRealize realize = GetSBlockRealize(self, block_sref); - Buffer read_buffer = GetNthAccessBuffer(self, block, read_buffer_index, BufferIndexType::kRead); + BufferVar read_buffer = GetNthAccessBuffer(self, block, read_buffer_index, BufferIndexType::kRead); StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/true); // Step 2. Create CacheStageInfo @@ -2157,7 +2155,7 @@ StmtSRef ReindexCacheWrite(ScheduleState self, const StmtSRef& block_sref, int w // Step 1. Checking index, getting the target buffer and the parent scope SBlock block = ffi::GetRef(TVM_SREF_TO_SBLOCK(block_sref)); SBlockRealize realize = GetSBlockRealize(self, block_sref); - Buffer write_buffer = + BufferVar write_buffer = GetNthAccessBuffer(self, block, write_buffer_index, BufferIndexType::kWrite); StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/true); @@ -2208,7 +2206,7 @@ StmtSRef ReindexCacheWrite(ScheduleState self, const StmtSRef& block_sref, int w /*! \brief The schedule error that the target block doesn't both read&write target buffer. */ class NotReadWriteError : public ScheduleError { public: - NotReadWriteError(IRModule mod, SBlock block, Buffer buffer) + NotReadWriteError(IRModule mod, SBlock block, BufferVar buffer) : mod_(std::move(mod)), block_(std::move(block)), buffer_(std::move(buffer)) {} ffi::String FastErrorString() const final { return "ScheduleError: The target block does not both read & write target buffer."; @@ -2222,7 +2220,7 @@ class NotReadWriteError : public ScheduleError { ffi::Array LocationsOfInterest() const final { return {block_, buffer_}; } IRModule mod_; SBlock block_; - Buffer buffer_; + BufferVar buffer_; }; ffi::Array CacheInplace(ScheduleState self, const StmtSRef& block_sref, @@ -2236,7 +2234,7 @@ ffi::Array CacheInplace(ScheduleState self, const StmtSRef& block_sref // Check 1. Check index, get the target buffer and the parent scope const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); - Buffer buffer = GetNthAccessBuffer(self, ffi::GetRef(block), read_buffer_index, + BufferVar buffer = GetNthAccessBuffer(self, ffi::GetRef(block), read_buffer_index, BufferIndexType::kRead); StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false); @@ -2252,7 +2250,7 @@ ffi::Array CacheInplace(ScheduleState self, const StmtSRef& block_sref } ffi::Array results_block_sref; - Buffer new_buffer = WithScope(buffer, storage_scope); + BufferVar new_buffer = WithScope(buffer, storage_scope); // Do cache read // Cache read step 0. Create CacheStageInfo @@ -2318,7 +2316,7 @@ StmtSRef ReIndex(ScheduleState self, const StmtSRef& block_sref, int buffer_inde BufferIndexType buffer_index_type) { const SBlockNode* block_ptr = TVM_SREF_TO_SBLOCK(block_sref); SBlock block = ffi::GetRef(block_ptr); - Buffer buffer = GetNthAccessBuffer(self, block, buffer_index, buffer_index_type); + BufferVar buffer = GetNthAccessBuffer(self, block, buffer_index, buffer_index_type); StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/true); arith::Analyzer analyzer; diff --git a/src/s_tir/schedule/primitive/compute_at.cc b/src/s_tir/schedule/primitive/compute_at.cc index e622602a5f5f..535541ddbd3d 100644 --- a/src/s_tir/schedule/primitive/compute_at.cc +++ b/src/s_tir/schedule/primitive/compute_at.cc @@ -366,14 +366,14 @@ void RelaxBufferRegions(const ffi::Map& binding, const ffi::Array& buffer_regions, const StmtSRef& relax_path_low_inclusive, const StmtSRef& relax_path_high_exclusive, - std::unordered_map>* relaxed) { + std::unordered_map>* relaxed) { runtime::StorageScope global_scope{runtime::StorageRank::kGlobal, ""}; // We cache the variable domains runtime::StorageRank previous_rank = runtime::StorageRank::kGlobal; ffi::Optional> var_dom = std::nullopt; // Enumerate every buffer region for (const BufferRegion& buffer_region : buffer_regions) { - const Buffer& buffer = buffer_region->buffer; + const BufferVar& buffer = buffer_region->buffer; const ffi::Array& region = buffer_region->region; // Skip the buffer regions we are not interested in auto it = relaxed->find(buffer.get()); @@ -485,14 +485,14 @@ std::pair SolveBlockVarDomain(const arith::IntSet& prov * \param iter_doms The result iteration domains to be updated */ void UpdateBlockVarDomainDimwise( - const BufferNode* buffer, const NDIntSet& provided_region, const NDIntSet& required_region, + const VarNode* buffer, const NDIntSet& provided_region, const NDIntSet& required_region, arith::AnalyzerObj* analyzer, std::unordered_map* iter_doms) { - size_t ndim = buffer->shape.size(); + size_t ndim = GetBufferVar(buffer)->shape.size(); for (size_t i = 0; i < ndim; ++i) { arith::IntSet provided = provided_region[i]; arith::IntSet required = required_region[i]; - PrimExpr dim_max = max(buffer->shape[i] - 1, 0); + PrimExpr dim_max = max(GetBufferVar(buffer)->shape[i] - 1, 0); arith::Analyzer analyzer_ref = ffi::GetRef(analyzer); if (provided.CanProveSinglePoint(analyzer_ref) && is_const_int(provided.min())) { @@ -551,7 +551,7 @@ ffi::Map InverseAffineIterMap(const ffi::Array& iter_vars, +bool UpdateBlockVarDomainAffine(const VarNode* buffer, const ffi::Array& iter_vars, const NDIntSet& provided_region, const NDIntSet& required_region, arith::AnalyzerObj* analyzer, std::unordered_map* iter_doms) { @@ -565,7 +565,7 @@ bool UpdateBlockVarDomainAffine(const BufferNode* buffer, const ffi::Arrayvar, iter_var->dom); } - size_t ndim = buffer->shape.size(); + size_t ndim = GetBufferVar(buffer)->shape.size(); ffi::Array provide_indices; provide_indices.reserve(ndim); for (size_t i = 0; i < ndim; ++i) { @@ -580,7 +580,7 @@ bool UpdateBlockVarDomainAffine(const BufferNode* buffer, const ffi::Arrayshape[i].ty(), 0), max(buffer->shape[i] - 1, 0))); + arith::IntSet::Interval(IntImm(GetBufferVar(buffer)->shape[i].ty(), 0), max(GetBufferVar(buffer)->shape[i] - 1, 0))); } ffi::Map var_dom = InverseAffineIterMap(res->indices, required_region, analyzer); @@ -605,8 +605,8 @@ bool UpdateBlockVarDomainAffine(const BufferNode* buffer, const ffi::Array CalculateBlockVarDomain( const ffi::Array& iter_vars, - std::unordered_map> provided_regions, - std::unordered_map> required_regions, + std::unordered_map> provided_regions, + std::unordered_map> required_regions, arith::AnalyzerObj* analyzer) { int n_iters = iter_vars.size(); // Step 1. Construct the mapping from block var to their iteration domain (initialized to empty) @@ -617,7 +617,7 @@ std::vector CalculateBlockVarDomain( } // Step 2. For each buffer, update the domain according to the provided and required regions for (const auto& kv : provided_regions) { - const BufferNode* buffer = kv.first; + const VarNode* buffer = kv.first; const std::vector& many_provided_regions = kv.second; // Calculate `provided_region` and `required_region` auto it = required_regions.find(buffer); @@ -626,8 +626,8 @@ std::vector CalculateBlockVarDomain( } NDIntSet required_region = support::NDIntSetUnion(it->second); NDIntSet provided_region = support::NDIntSetUnion(many_provided_regions); - TVM_FFI_ICHECK_EQ(provided_region.size(), buffer->shape.size()); - TVM_FFI_ICHECK_EQ(required_region.size(), buffer->shape.size()); + TVM_FFI_ICHECK_EQ(provided_region.size(), GetBufferVar(buffer)->shape.size()); + TVM_FFI_ICHECK_EQ(required_region.size(), GetBufferVar(buffer)->shape.size()); // Try update iter var domains with current required and provided region pair. if (!UpdateBlockVarDomainAffine(buffer, iter_vars, provided_region, required_region, analyzer, &iter_doms)) { @@ -669,14 +669,14 @@ void CalculateProvidedRequiredRegions( const SBlockNode* block, const StmtSRef& loop_sref, std::unordered_map block2realize, ffi::Array producer_srefs, ffi::Array consumer_srefs, - std::unordered_map>* provided_regions, - std::unordered_map>* required_regions) { + std::unordered_map>* provided_regions, + std::unordered_map>* required_regions) { // Step 1. Calculate the region provided by a single execution instance of `block` const ffi::Array& provided_buffers = is_compute_at ? block->writes : block->reads; provided_regions->reserve(provided_buffers.size()); required_regions->reserve(provided_buffers.size()); for (const BufferRegion& provided_buffer_region : provided_buffers) { - const BufferNode* buffer = provided_buffer_region->buffer.get(); + const VarNode* buffer = provided_buffer_region->buffer.get(); const ffi::Array& region = provided_buffer_region->region; (*provided_regions)[buffer].push_back(support::NDIntSetFromRegion(region)); (*required_regions)[buffer].clear(); @@ -739,8 +739,8 @@ void ComputeAtOrReverseComputeAtImpl(ScheduleState self, const StmtSRef& block_s // Here is the definition of `provide` and `require`: // - In compute-at, `provide` means `produce`, and `require` means `consume` // - In reverse-compute-at, `provide` means `consume`, and `require` means `produce` - std::unordered_map> provided_regions; - std::unordered_map> required_regions; + std::unordered_map> provided_regions; + std::unordered_map> required_regions; CalculateProvidedRequiredRegions( /*block=*/block, /*loop_sref=*/loop_sref, /*block2realize=*/std::move(block2realize), /*producer_srefs=*/std::move(producer_srefs), diff --git a/src/s_tir/schedule/primitive/compute_inline.cc b/src/s_tir/schedule/primitive/compute_inline.cc index 1d90e5e27450..c6997a19bbfb 100644 --- a/src/s_tir/schedule/primitive/compute_inline.cc +++ b/src/s_tir/schedule/primitive/compute_inline.cc @@ -91,17 +91,17 @@ class NotSingleReadWriteBuffer : public ScheduleError { bool is_read_; SBlock block_; - static Buffer GetSingleRead(const ScheduleState& self, const SBlock& block, + static BufferVar GetSingleRead(const ScheduleState& self, const SBlock& block, const StmtSRef& scope_root_sref) { - const std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>& + const std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>& buffer_writers = self->block_info.at(scope_root_sref).scope->buffer_writers; - const BufferNode* read_buffer = nullptr; + const VarNode* read_buffer = nullptr; for (const BufferRegion& read_region : block->reads) { - const BufferNode* buffer = read_region->buffer.get(); + const VarNode* buffer = read_region->buffer.get(); if (buffer == read_buffer) { continue; } - if (buffer_writers.count(ffi::GetRef(buffer)) > 0) { + if (buffer_writers.count(BufferVar(ffi::GetRef(buffer))) > 0) { if (read_buffer != nullptr) { throw NotSingleReadWriteBuffer(self->mod, true, block); } @@ -111,10 +111,10 @@ class NotSingleReadWriteBuffer : public ScheduleError { if (read_buffer == nullptr) { throw NotSingleReadWriteBuffer(self->mod, true, block); } - return ffi::GetRef(read_buffer); + return BufferVar(ffi::GetRef(read_buffer)); } - static Buffer GetSingleWrite(const ScheduleState& self, const SBlock& block) { + static BufferVar GetSingleWrite(const ScheduleState& self, const SBlock& block) { if (block->writes.size() != 1) { throw NotSingleReadWriteBuffer(self->mod, false, block); } @@ -179,12 +179,12 @@ class NonSingleProducerError : public ScheduleError { const StmtSRef& scope_root_sref) { const SBlockNode* scope_block = TVM_SREF_TO_SBLOCK(scope_root_sref); const SBlockNode* consumer_block = TVM_SREF_TO_SBLOCK(consumer_block_sref); - Buffer consumer_buffer = NotSingleReadWriteBuffer::GetSingleRead( + BufferVar consumer_buffer = NotSingleReadWriteBuffer::GetSingleRead( self, ffi::GetRef(consumer_block), scope_root_sref); class ProducerFinder : public StmtVisitor { public: static std::vector GetProducer(const ScheduleState& self, - const StmtSRef& scope_root_sref, const Buffer& buffer, + const StmtSRef& scope_root_sref, const BufferVar& buffer, const SBlock& scope_block) { ProducerFinder finder(self, scope_root_sref, buffer); finder(scope_block); @@ -193,7 +193,7 @@ class NonSingleProducerError : public ScheduleError { private: explicit ProducerFinder(const ScheduleState& self, const StmtSRef& scope_root_sref, - const Buffer& buffer) + const BufferVar& buffer) : self_(self), scope_root_sref_(scope_root_sref), buffer_(buffer) { producer_across_scope_.push_back({}); } @@ -226,7 +226,7 @@ class NonSingleProducerError : public ScheduleError { } ScheduleState self_; StmtSRef scope_root_sref_; - Buffer buffer_; + BufferVar buffer_; std::vector> producer_across_scope_; }; std::vector producer_across_scope = ProducerFinder::GetProducer( @@ -297,7 +297,7 @@ class ProducerHasNonTrivialPredicateError : public ScheduleError { */ class BaseInliner : public StmtExprMutator { protected: - explicit BaseInliner(const Buffer& inlined_buffer, const SBlock& inlined_block, + explicit BaseInliner(const BufferVar& inlined_buffer, const SBlock& inlined_block, const StmtSRef& scope_root_sref) : inlined_buffer_(inlined_buffer), inlined_store_(inlined_block->body.as()), @@ -341,15 +341,15 @@ class BaseInliner : public StmtExprMutator { */ void AddBuffersInBlockSignature(const SBlockNode* block) { for (const BufferRegion& buffer_region : block->reads) { - const Buffer& buffer = buffer_region->buffer; - buffer_var_map_.Set(buffer->data, buffer); + const BufferVar& buffer = buffer_region->buffer; + buffer_var_map_.Set(buffer.var(), buffer); } for (const BufferRegion& buffer_region : block->writes) { - const Buffer& buffer = buffer_region->buffer; - buffer_var_map_.Set(buffer->data, buffer); + const BufferVar& buffer = buffer_region->buffer; + buffer_var_map_.Set(buffer.var(), buffer); } - for (const Buffer& buffer : block->alloc_buffers) { - buffer_var_map_.Set(buffer->data, buffer); + for (const BufferVar& buffer : block->alloc_buffers) { + buffer_var_map_.Set(buffer.var(), buffer); } } @@ -364,10 +364,10 @@ class BaseInliner : public StmtExprMutator { */ SBlock UpdateBuffersInBlockSignature(SBlock block, bool is_scope_root) { // Step 1. Update `BlockNode::alloc_buffers` - ffi::Array alloc_buffers; + ffi::Array alloc_buffers; if (is_scope_root) { alloc_buffers.reserve(block->alloc_buffers.size()); - for (const Buffer& alloc_buffer : block->alloc_buffers) { + for (const BufferVar& alloc_buffer : block->alloc_buffers) { if (!alloc_buffer.same_as(inlined_buffer_)) { alloc_buffers.push_back(alloc_buffer); } @@ -402,7 +402,7 @@ class BaseInliner : public StmtExprMutator { * \param buffer_var The buffer var to be checked */ void CheckOpaqueAccess(const VarNode* buffer_var) { - if (inlined_buffer_->data.get() == buffer_var) { + if (inlined_buffer_.get() == buffer_var) { this->has_opaque_access = true; } } @@ -414,7 +414,7 @@ class BaseInliner : public StmtExprMutator { */ void CheckMatchBufferRegion(const SBlockNode* block) { for (const MatchBufferRegion& match_buffer_region : block->match_buffers) { - const Buffer& matched = match_buffer_region->source->buffer; + const BufferVar& matched = match_buffer_region->source->buffer; if (matched.same_as(inlined_buffer_)) { this->has_opaque_access = true; } @@ -423,13 +423,13 @@ class BaseInliner : public StmtExprMutator { protected: /*! \brief The buffer to be inlined */ - Buffer inlined_buffer_{nullptr}; + BufferVar inlined_buffer_{nullptr}; /*! \brief The body of the block to be inlined */ const BufferStoreNode* inlined_store_{nullptr}; /*! \brief The scope root */ StmtSRef scope_root_sref_{nullptr}; /*! \brief Maps a buffer's data field to itself */ - ffi::Map buffer_var_map_; + ffi::Map buffer_var_map_; /*! \brief The indices used for indexing the buffer to be inlined */ std::vector idx_vars_; /*! \brief The mapping to substitute index variables to PrimExprs */ @@ -459,7 +459,7 @@ class BaseInliner : public StmtExprMutator { */ class ComputeInliner : public BaseInliner { public: - explicit ComputeInliner(const Buffer& inlined_buffer, const SBlock& producer_block, + explicit ComputeInliner(const BufferVar& inlined_buffer, const SBlock& producer_block, const StmtSRef& scope_root_sref) : BaseInliner(inlined_buffer, producer_block, scope_root_sref) {} @@ -620,7 +620,7 @@ class ReverseComputeInliner : public BaseInliner { }; public: - explicit ReverseComputeInliner(const Buffer& inlined_buffer, const SBlockNode* producer_block, + explicit ReverseComputeInliner(const BufferVar& inlined_buffer, const SBlockNode* producer_block, const SBlockRealize& consumer_block_realize, const StmtSRef& scope_root_sref, const IRModule& mod) : BaseInliner(inlined_buffer, consumer_block_realize->block, scope_root_sref), @@ -829,7 +829,7 @@ class ReverseComputeInliner : public BaseInliner { * \param from The BufferStore statement to be extracted from * \return A list of `BufferLoad` expressions */ - static std::vector ExtractBufferLoad(const Buffer& buffer, + static std::vector ExtractBufferLoad(const BufferVar& buffer, const BufferStoreNode* from) { struct Extractor : public ExprVisitor { void VisitExpr_(const BufferLoadNode* load) final { @@ -838,7 +838,7 @@ class ReverseComputeInliner : public BaseInliner { } ExprVisitor::VisitExpr_(load); } - const BufferNode* buffer; + const VarNode* buffer; std::vector result; } extractor; extractor.buffer = buffer.get(); @@ -890,7 +890,7 @@ void ComputeInlineImpl(ScheduleState self, const StmtSRef& producer_block_sref, const SBlockNode* _producer_block = TVM_SREF_TO_SBLOCK(producer_block_sref); SBlock producer_block = ffi::GetRef(_producer_block); HasInitBlock::Check(self->mod, producer_block); - Buffer inlined_buffer = NotSingleReadWriteBuffer::GetSingleWrite(self, producer_block); + BufferVar inlined_buffer = NotSingleReadWriteBuffer::GetSingleWrite(self, producer_block); // Step 1. Get the scope block StmtSRef scope_root_sref = GetScopeRoot(self, producer_block_sref, /*require_stage_pipeline=*/true); @@ -939,7 +939,7 @@ void ReverseComputeInlineImpl(ScheduleState self, const StmtSRef& consumer_block // Step 1. Get the scope block StmtSRef scope_root_sref = GetScopeRoot(self, consumer_block_sref, // /*require_stage_pipeline=*/true); - Buffer inlined_buffer = + BufferVar inlined_buffer = NotSingleReadWriteBuffer::GetSingleRead(self, consumer_block, scope_root_sref); // Step 2. Check completeness CheckCompleteBlock(self, consumer_block_sref, scope_root_sref); @@ -996,7 +996,7 @@ void ReverseComputeInline(ScheduleState self, const StmtSRef& consumer_block_sre */ class ReductionEpilogueFuser : public BaseInliner { public: - explicit ReductionEpilogueFuser(const Buffer& reduction_buffer, const SBlockNode* reduction_block, + explicit ReductionEpilogueFuser(const BufferVar& reduction_buffer, const SBlockNode* reduction_block, const SBlockRealize& epilogue_block_realize, const StmtSRef& scope_root_sref) : BaseInliner(reduction_buffer, epilogue_block_realize->block, scope_root_sref), @@ -1024,7 +1024,7 @@ class ReductionEpilogueFuser : public BaseInliner { bool IsReductionBlock(const SBlockNode* block); void ExtractEpilogueInfo(); // Helper function to extract BufferLoad nodes from BufferStore - static std::vector ExtractBufferLoad(const Buffer& buffer, + static std::vector ExtractBufferLoad(const BufferVar& buffer, const BufferStoreNode* from) { struct Extractor : public ExprVisitor { void VisitExpr_(const BufferLoadNode* load) final { @@ -1034,7 +1034,7 @@ class ReductionEpilogueFuser : public BaseInliner { // Continue visiting child nodes (indices) ExprVisitor::VisitExpr_(load); } - Buffer buffer; + BufferVar buffer; std::vector result; } extractor; extractor.buffer = buffer; @@ -1054,10 +1054,10 @@ class ReductionEpilogueFuser : public BaseInliner { nullptr}; // The entire epilogue expression (e.g., temp + C, max(temp + C, 0)) const BufferLoadNode* reduction_buffer_load_{ nullptr}; // The reduction buffer load in epilogue expression - Buffer epilogue_output_buffer_{nullptr}; // Output buffer D + BufferVar epilogue_output_buffer_{nullptr}; // Output buffer D ffi::Array epilogue_output_indices_{nullptr}; // Indices of D[vi, vj] BufferRegion epilogue_output_region_{nullptr}; // Write region of D - Buffer epilogue_addend_buffer_{nullptr}; // Additional buffer (e.g., bias buffer C) + BufferVar epilogue_addend_buffer_{nullptr}; // Additional buffer (e.g., bias buffer C) BufferRegion epilogue_addend_region_{nullptr}; // Read region of additional buffer }; @@ -1098,7 +1098,7 @@ bool ReductionEpilogueFuser::BodyPatternAllowFusion(const SBlockRealize& epilogu // We only allow the reduction result to be combined via Add/Min/Max shells. class ScalingDetector : public ExprVisitor { public: - explicit ScalingDetector(const Buffer& buffer) : buffer_(buffer) {} + explicit ScalingDetector(const BufferVar& buffer) : buffer_(buffer) {} bool HasScaling(const PrimExpr& expr) { has_scaling_ = false; @@ -1111,7 +1111,7 @@ bool ReductionEpilogueFuser::BodyPatternAllowFusion(const SBlockRealize& epilogu bool ContainsTarget(const PrimExpr& expr) { class TargetFinder : public ExprVisitor { public: - explicit TargetFinder(const Buffer& buffer) : buffer_(buffer) {} + explicit TargetFinder(const BufferVar& buffer) : buffer_(buffer) {} bool Find(const PrimExpr& e) { found_ = false; @@ -1128,7 +1128,7 @@ bool ReductionEpilogueFuser::BodyPatternAllowFusion(const SBlockRealize& epilogu ExprVisitor::VisitExpr_(op); } - Buffer buffer_; + BufferVar buffer_; bool found_{false}; }; @@ -1165,7 +1165,7 @@ bool ReductionEpilogueFuser::BodyPatternAllowFusion(const SBlockRealize& epilogu ExprVisitor::VisitExpr_(op); } - Buffer buffer_; + BufferVar buffer_; bool has_scaling_{false}; }; @@ -1221,8 +1221,8 @@ void ReductionEpilogueFuser::ExtractEpilogueInfo() { } ExprVisitor::VisitExpr_(load); } - Buffer reduction_buffer; - std::unordered_set other_buffers; + BufferVar reduction_buffer; + std::unordered_set other_buffers; } extractor; extractor.reduction_buffer = inlined_buffer_; extractor(epilogue_expression_); @@ -1230,8 +1230,8 @@ void ReductionEpilogueFuser::ExtractEpilogueInfo() { // Extract the first non-reduction buffer and its region // In most cases, there's one additional buffer (e.g., bias buffer) if (!extractor.other_buffers.empty()) { - const BufferNode* first_buffer = *extractor.other_buffers.begin(); - epilogue_addend_buffer_ = ffi::GetRef(first_buffer); + const VarNode* first_buffer = *extractor.other_buffers.begin(); + epilogue_addend_buffer_ = BufferVar(ffi::GetRef(first_buffer)); // Find the read region from epilogue block reads for (const BufferRegion& read : epilogue_block_->reads) { if (read->buffer.get() == first_buffer) { @@ -1273,7 +1273,7 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( // Create a substituter to replace reduction_buffer_load_ with identity element class InitSubstituter : public ExprMutator { public: - InitSubstituter(const Buffer& target_buffer, PrimExpr identity_elem) + InitSubstituter(const BufferVar& target_buffer, PrimExpr identity_elem) : target_buffer_(target_buffer), identity_elem_(identity_elem) {} Expr VisitExpr_(const BufferLoadNode* op) final { @@ -1285,7 +1285,7 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( } private: - Buffer target_buffer_; + BufferVar target_buffer_; PrimExpr identity_elem_; }; @@ -1312,7 +1312,7 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( // remove that operand (bias addend) from update expression class UpdateSubstituter : public StmtExprMutator { public: - UpdateSubstituter(const Buffer& old_buf, const Buffer& new_buf, const Buffer& reduction_buf, + UpdateSubstituter(const BufferVar& old_buf, const BufferVar& new_buf, const BufferVar& reduction_buf, const PrimExpr& epilogue_expr, const std::unordered_map& var_map) : old_buffer_(old_buf), new_buffer_(new_buf), @@ -1327,7 +1327,7 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( // expression This ensures store->value references new_buffer_ instead of old_buffer_ class ReductionUpdateReplacer : public ExprMutator { public: - ReductionUpdateReplacer(const Buffer& old_buf, const Buffer& new_buf) + ReductionUpdateReplacer(const BufferVar& old_buf, const BufferVar& new_buf) : old_buffer_(old_buf), new_buffer_(new_buf) {} Expr VisitExpr_(const BufferLoadNode* op) final { @@ -1339,8 +1339,8 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( } private: - Buffer old_buffer_; - Buffer new_buffer_; + BufferVar old_buffer_; + BufferVar new_buffer_; }; ReductionUpdateReplacer reduction_replacer(old_buffer_, new_buffer_); @@ -1351,7 +1351,7 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( // buffer, remove that operand (bias addend) from the update expression class GeneralizedEpilogueApplier : public ExprMutator { public: - GeneralizedEpilogueApplier(const Buffer& target_buf, const Buffer& reduction_buf, + GeneralizedEpilogueApplier(const BufferVar& target_buf, const BufferVar& reduction_buf, const PrimExpr& replacement) : target_buffer_(target_buf), reduction_buffer_(reduction_buf), @@ -1414,8 +1414,8 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( } private: - const Buffer& target_buffer_; - const Buffer& reduction_buffer_; + const BufferVar& target_buffer_; + const BufferVar& reduction_buffer_; const PrimExpr& replacement_; bool found_target_load_; }; @@ -1440,9 +1440,9 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( } private: - Buffer old_buffer_; - Buffer new_buffer_; - Buffer reduction_buffer_; + BufferVar old_buffer_; + BufferVar new_buffer_; + BufferVar reduction_buffer_; PrimExpr epilogue_expression_; std::unordered_map var_map_; }; @@ -1468,7 +1468,7 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( // 5. Update read regions: add all buffers from epilogue expression (except reduction buffer) ffi::Array new_reads; - std::unordered_set read_bufs; + std::unordered_set read_bufs; // Add all non-reduction buffers from epilogue expression for (const BufferRegion& read : epilogue_block_->reads) { @@ -1497,10 +1497,10 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( /*! * \brief Check if a buffer is still referenced by other blocks in the scope */ -static bool CheckBufferStillUsed(const SBlock& scope_root, const Buffer& buffer) { +static bool CheckBufferStillUsed(const SBlock& scope_root, const BufferVar& buffer) { class BufferUsageChecker : public StmtVisitor { public: - explicit BufferUsageChecker(const Buffer& buffer) : buffer_(buffer) {} + explicit BufferUsageChecker(const BufferVar& buffer) : buffer_(buffer) {} bool CheckStmt(const Stmt& stmt) { found_usage_ = false; @@ -1548,7 +1548,7 @@ static bool CheckBufferStillUsed(const SBlock& scope_root, const Buffer& buffer) if (!op) return; // Check alloc_buffers - for (const Buffer& buf : op->alloc_buffers) { + for (const BufferVar& buf : op->alloc_buffers) { if (buf.same_as(buffer_)) { found_usage_ = true; return; @@ -1558,7 +1558,7 @@ static bool CheckBufferStillUsed(const SBlock& scope_root, const Buffer& buffer) StmtVisitor::VisitStmt_(op); } - const Buffer& buffer_; + const BufferVar& buffer_; bool found_usage_{false}; }; @@ -1576,7 +1576,7 @@ static bool CheckBufferStillUsed(const SBlock& scope_root, const Buffer& buffer) class SingleBlockFusionReplacer : public StmtMutator { public: static SBlock Replace(SBlock old_scope_root, SBlock new_fused_block, SBlock old_reduction_block, - SBlock old_epilogue_block, Buffer reduction_buffer) { + SBlock old_epilogue_block, BufferVar reduction_buffer) { SingleBlockFusionReplacer replacer(std::move(new_fused_block), std::move(old_reduction_block), std::move(old_epilogue_block), std::move(reduction_buffer)); SBlock result = replacer(std::move(old_scope_root)).as_or_throw(); @@ -1587,8 +1587,8 @@ class SingleBlockFusionReplacer : public StmtMutator { // Remove intermediate temp buffer only if it's not used by other blocks if (!buffer_still_used) { SBlockNode* p = result.CopyOnWrite(); - ffi::Array new_alloc_buffers; - for (const Buffer& buf : p->alloc_buffers) { + ffi::Array new_alloc_buffers; + for (const BufferVar& buf : p->alloc_buffers) { if (!buf.same_as(reduction_buffer)) { new_alloc_buffers.push_back(buf); } @@ -1601,7 +1601,7 @@ class SingleBlockFusionReplacer : public StmtMutator { private: explicit SingleBlockFusionReplacer(SBlock new_fused_block, SBlock old_reduction_block, - SBlock old_epilogue_block, Buffer reduction_buffer) + SBlock old_epilogue_block, BufferVar reduction_buffer) : new_fused_block_(std::move(new_fused_block)), old_reduction_block_(std::move(old_reduction_block)), old_epilogue_block_(std::move(old_epilogue_block)), @@ -1647,7 +1647,7 @@ class SingleBlockFusionReplacer : public StmtMutator { SBlock new_fused_block_; SBlock old_reduction_block_; SBlock old_epilogue_block_; - Buffer reduction_buffer_; + BufferVar reduction_buffer_; }; void FuseReductionEpilogueImpl(ScheduleState self, const StmtSRef& reduction_block_sref, @@ -1664,7 +1664,7 @@ void FuseReductionEpilogueImpl(ScheduleState self, const StmtSRef& reduction_blo GetScopeRoot(self, epilogue_block_sref, /*require_stage_pipeline=*/true); // Step 2. Get the reduction buffer (intermediate buffer) - Buffer reduction_buffer = NotSingleReadWriteBuffer::GetSingleWrite(self, reduction_block); + BufferVar reduction_buffer = NotSingleReadWriteBuffer::GetSingleWrite(self, reduction_block); // Step 3. Check completeness and reduction block properties CheckReductionBlock(self, reduction_block_sref, scope_root_sref); diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc b/src/s_tir/schedule/primitive/layout_transformation.cc index 97c0f00a759e..c9de8f6fbe76 100644 --- a/src/s_tir/schedule/primitive/layout_transformation.cc +++ b/src/s_tir/schedule/primitive/layout_transformation.cc @@ -95,7 +95,7 @@ class TransformLayoutPlanner : private StmtExprVisitor { using TransformPlan = std::variant; - static TransformPlan Plan(SBlock block, Buffer old_buffer, Buffer new_buffer, IndexMap index_map, + static TransformPlan Plan(SBlock block, BufferVar old_buffer, BufferVar new_buffer, IndexMap index_map, IndexMap inverse, PrimExpr padding_predicate, ffi::Optional pad_value, arith::AnalyzerObj* analyzer) { TVM_FFI_ICHECK(!pad_value.has_value() || pad_value.value()->final_indices.size() == 1) @@ -125,7 +125,7 @@ class TransformLayoutPlanner : private StmtExprVisitor { bool contains_row_major_traversal{false}; }; - explicit TransformLayoutPlanner(Buffer old_buffer) : old_buffer_(old_buffer) {} + explicit TransformLayoutPlanner(BufferVar old_buffer) : old_buffer_(old_buffer) {} void VisitStmt_(const ForNode* op) override { BindLoopVar context(this, ffi::GetRef(op)); @@ -222,7 +222,7 @@ class TransformLayoutPlanner : private StmtExprVisitor { class BufferStoreReplacer : public StmtExprMutator { public: - BufferStoreReplacer(const WriteInfo& info, const Buffer& new_buffer, PrimExpr padding_predicate, + BufferStoreReplacer(const WriteInfo& info, const BufferVar& new_buffer, PrimExpr padding_predicate, const IndexMap& inverse, const ffi::Optional& pad_value, ffi::Map* new_block_to_old, arith::AnalyzerObj* analyzer) : info(info), @@ -428,7 +428,7 @@ class TransformLayoutPlanner : private StmtExprVisitor { } const WriteInfo& info; - const Buffer& new_buffer; + const BufferVar& new_buffer; ffi::Array new_indices; ffi::Array new_iter_vars; ffi::Array new_iter_values; @@ -442,7 +442,7 @@ class TransformLayoutPlanner : private StmtExprVisitor { ffi::Map var_remap; }; - TransformPlan Finalize(Buffer new_buffer, IndexMap index_map, IndexMap inverse, + TransformPlan Finalize(BufferVar new_buffer, IndexMap index_map, IndexMap inverse, PrimExpr padding_predicate, ffi::Optional pad_value, arith::AnalyzerObj* analyzer) const { if (auto prologue_plan = FinalizeProloguePlan(new_buffer, index_map, inverse, padding_predicate, @@ -462,7 +462,7 @@ class TransformLayoutPlanner : private StmtExprVisitor { } } - std::optional FinalizeProloguePlan(Buffer new_buffer, IndexMap index_map, + std::optional FinalizeProloguePlan(BufferVar new_buffer, IndexMap index_map, IndexMap inverse, PrimExpr padding_predicate, ffi::Optional pad_value, arith::AnalyzerObj* analyzer) const { @@ -493,7 +493,7 @@ class TransformLayoutPlanner : private StmtExprVisitor { Stmt stmt = Evaluate(Call(PrimType::Bool(), builtin::assume(), {expr}).as_or_throw()); std::stringstream block_name; - block_name << "buffer_" << new_buffer->name << "_assumptions"; + block_name << "buffer_" << new_buffer.name() << "_assumptions"; auto read_region = BufferRegion::FromPoint(new_buffer, indices); stmt = SBlockRealize(iter_values, IntImm::Bool(true), SBlock(iter_vars, {read_region}, {}, block_name.str(), stmt)); @@ -507,7 +507,7 @@ class TransformLayoutPlanner : private StmtExprVisitor { return ProloguePlan{stmt}; } - std::optional FinalizeReplacementPlan(Buffer new_buffer, IndexMap index_map, + std::optional FinalizeReplacementPlan(BufferVar new_buffer, IndexMap index_map, IndexMap inverse, PrimExpr padding_predicate, ffi::Optional pad_value, @@ -558,7 +558,7 @@ class TransformLayoutPlanner : private StmtExprVisitor { } } - std::optional FinalizeEpiloguePlan(Buffer new_buffer, IndexMap index_map, + std::optional FinalizeEpiloguePlan(BufferVar new_buffer, IndexMap index_map, IndexMap inverse, PrimExpr padding_predicate, ffi::Optional pad_value, arith::AnalyzerObj* analyzer) const { @@ -585,7 +585,7 @@ class TransformLayoutPlanner : private StmtExprVisitor { Stmt stmt = BufferStore(new_buffer, pad_value_at_index, indices); std::stringstream block_name; - block_name << "buffer_" << new_buffer->name << "_padding"; + block_name << "buffer_" << new_buffer.name() << "_padding"; auto write_region = BufferRegion::FromPoint(new_buffer, indices); stmt = SBlockRealize(iter_values, padding_predicate, SBlock(iter_vars, {}, {write_region}, block_name.str(), stmt)); @@ -707,7 +707,7 @@ class TransformLayoutPlanner : private StmtExprVisitor { ffi::Optional innermost_block_realize_{std::nullopt}; /*! \brief The buffer to be replaced */ - Buffer old_buffer_; + BufferVar old_buffer_; }; /*! @@ -759,7 +759,7 @@ class TransformLayoutRewriter : private arith::IRMutatorWithAnalyzer { * new block */ static std::pair> Rewrite( - const SBlock& scope_stmt, const Buffer& old_buffer, const Buffer& new_buffer, + const SBlock& scope_stmt, const BufferVar& old_buffer, const BufferVar& new_buffer, const IndexMap& index_map, const ffi::Optional& opt_inverse, const PrimExpr& padding_predicate, const ffi::Optional& pad_value) { arith::Analyzer analyzer; @@ -783,7 +783,7 @@ class TransformLayoutRewriter : private arith::IRMutatorWithAnalyzer { } private: - TransformLayoutRewriter(const Buffer& old_buffer, const Buffer& new_buffer, + TransformLayoutRewriter(const BufferVar& old_buffer, const BufferVar& new_buffer, const IndexMap& index_map, const TransformLayoutPlanner::TransformPlan& plan, const arith::Analyzer& analyzer) @@ -792,13 +792,13 @@ class TransformLayoutRewriter : private arith::IRMutatorWithAnalyzer { new_buffer_(new_buffer), index_map_(index_map), plan_(plan), - buffer_data_to_buffer_{{new_buffer->data, new_buffer}} { + buffer_data_to_buffer_{{new_buffer.var(), new_buffer}} { if (auto plan_ptr = std::get_if(&plan_)) { new_block_to_old_ = plan_ptr->new_block_to_old; } } - void RewriteBufferAccess(Buffer* buffer, ffi::Array* indices) { + void RewriteBufferAccess(BufferVar* buffer, ffi::Array* indices) { *buffer = new_buffer_; *indices = index_map_->MapIndices(*indices, index_simplifier_); *indices = this->IterMapSimplifyWithContext(*indices, true); @@ -901,7 +901,7 @@ class TransformLayoutRewriter : private arith::IRMutatorWithAnalyzer { } return match_buf; }); - n->alloc_buffers.MutateByApply([this](const Buffer& buffer) { + n->alloc_buffers.MutateByApply([this](const BufferVar& buffer) { if (buffer.same_as(old_buffer_)) { return new_buffer_; } else { @@ -931,18 +931,18 @@ class TransformLayoutRewriter : private arith::IRMutatorWithAnalyzer { new_block_to_old_.Set(after, before); } - const Buffer& old_buffer_; - const Buffer& new_buffer_; + const BufferVar& old_buffer_; + const BufferVar& new_buffer_; const IndexMap& index_map_; const TransformLayoutPlanner::TransformPlan& plan_; - ffi::Map buffer_data_to_buffer_; + ffi::Map buffer_data_to_buffer_; ffi::Map new_block_to_old_; arith::Analyzer index_simplifier_; }; class BufferIsSubregionError : public ScheduleError { public: - explicit BufferIsSubregionError(IRModule mod, Buffer buffer) : mod_(mod), buffer_(buffer) {} + explicit BufferIsSubregionError(IRModule mod, BufferVar buffer) : mod_(mod), buffer_(buffer) {} ffi::String FastErrorString() const final { return "ScheduleError: The input buffer is defined in `match_buffer` of a block, it is expected" @@ -951,7 +951,7 @@ class BufferIsSubregionError : public ScheduleError { ffi::String DetailRenderTemplate() const final { std::ostringstream os; - os << "ScheduleError: The input buffer " << buffer_->name << " is defined in `match_buffer` of " + os << "ScheduleError: The input buffer " << buffer_.name() << " is defined in `match_buffer` of " << "a block, it is expected to be a function parameter or allocated by a block."; return os.str(); } @@ -961,7 +961,7 @@ class BufferIsSubregionError : public ScheduleError { private: IRModule mod_; - Buffer buffer_; + BufferVar buffer_; }; class TransformationPaddingIndexMapError : public ScheduleError { @@ -993,7 +993,7 @@ class TransformationPaddingIndexMapError : public ScheduleError { class TransformationPaddingTypeError : public ScheduleError { public: - TransformationPaddingTypeError(IRModule mod, Buffer buffer, IndexMap pad_value) + TransformationPaddingTypeError(IRModule mod, BufferVar buffer, IndexMap pad_value) : mod_(mod), buffer_(buffer), pad_value_(pad_value) { TVM_FFI_ICHECK_EQ(pad_value_->final_indices.size(), 1); pad_value_dtype_ = pad_value_->final_indices[0].ty()->dtype; @@ -1007,7 +1007,7 @@ class TransformationPaddingTypeError : public ScheduleError { ffi::String DetailRenderTemplate() const final { std::ostringstream ss; - ss << "ScheduleError: Buffer " << buffer_->name << " has elements of type " << buffer_->dtype + ss << "ScheduleError: BufferVar " << buffer_.name() << " has elements of type " << buffer_->dtype << ", but the transformation fills padding with " << pad_value_ << ", which is of type " << pad_value_dtype_; return ss.str(); @@ -1018,14 +1018,14 @@ class TransformationPaddingTypeError : public ScheduleError { private: IRModule mod_; - Buffer buffer_; + BufferVar buffer_; IndexMap pad_value_; DLDataType pad_value_dtype_; }; class TransformationPaddingExpressionError : public ScheduleError { public: - static void Check(IRModule mod, Buffer buffer, IndexMap pad_value) { + static void Check(IRModule mod, BufferVar buffer, IndexMap pad_value) { Visitor visitor(buffer); TVM_FFI_ICHECK_EQ(pad_value->final_indices.size(), 1) << "Internal error: Should be caught by ScheduleError checks prior to this point"; @@ -1038,7 +1038,7 @@ class TransformationPaddingExpressionError : public ScheduleError { private: struct Visitor : ExprVisitor { - explicit Visitor(const Buffer& buffer) : buffer_(buffer) {} + explicit Visitor(const BufferVar& buffer) : buffer_(buffer) {} void VisitExpr_(const BufferLoadNode* op) final { if (!op->buffer.same_as(buffer_)) { @@ -1047,24 +1047,24 @@ class TransformationPaddingExpressionError : public ScheduleError { ExprVisitor::VisitExpr_(op); } - const Buffer& buffer_; + const BufferVar& buffer_; ffi::Optional illegal_load; }; - TransformationPaddingExpressionError(IRModule mod, Buffer buffer, IndexMap pad_value, + TransformationPaddingExpressionError(IRModule mod, BufferVar buffer, IndexMap pad_value, BufferLoad illegal_load) : mod_(mod), buffer_(buffer), pad_value_(pad_value), illegal_load_(illegal_load) {} ffi::String FastErrorString() const final { std::ostringstream ss; - ss << "ScheduleError: Pad value may not contain load from " << illegal_load_->buffer->name; + ss << "ScheduleError: Pad value may not contain load from " << illegal_load_->buffer.name(); return ss.str(); } ffi::String DetailRenderTemplate() const final { std::ostringstream ss; ss << "ScheduleError: Pad value may only contain BufferLoad from the transformed buffer " - << buffer_->name << ", but pad_value " << pad_value_ << " contains expression " + << buffer_.name() << ", but pad_value " << pad_value_ << " contains expression " << illegal_load_; return ss.str(); } @@ -1073,14 +1073,14 @@ class TransformationPaddingExpressionError : public ScheduleError { ffi::Array LocationsOfInterest() const final { return {}; } IRModule mod_; - Buffer buffer_; + BufferVar buffer_; IndexMap pad_value_; BufferLoad illegal_load_; }; class TransformationIntroducesPaddingError : public ScheduleError { public: - TransformationIntroducesPaddingError(IRModule mod, Buffer buffer, IndexMap index_map, + TransformationIntroducesPaddingError(IRModule mod, BufferVar buffer, IndexMap index_map, PrimExpr padding_predicate) : mod_(std::move(mod)), buffer_(std::move(buffer)), @@ -1097,7 +1097,7 @@ class TransformationIntroducesPaddingError : public ScheduleError { arith::Analyzer analyzer; auto new_shape = index_map_->MapShape(buffer_->shape, analyzer); std::ostringstream os; - os << "The transformation " << index_map_ << " applied on buffer " << buffer_->name + os << "The transformation " << index_map_ << " applied on buffer " << buffer_.name() << " of shape " << buffer_->shape << " would result in shape " << new_shape << ". However, this would introduce padding wherever " << padding_predicate_ << " is true."; return os.str(); @@ -1108,7 +1108,7 @@ class TransformationIntroducesPaddingError : public ScheduleError { private: IRModule mod_; - Buffer buffer_; + BufferVar buffer_; IndexMap index_map_; PrimExpr padding_predicate_; }; @@ -1127,7 +1127,7 @@ IndexMap LegalizeIndexMapDType(const IndexMap& index_map, const ffi::Arraydtype; if (index_dtype.has_value()) { TVM_FFI_ICHECK_EQ(*index_dtype, arg_dtype) - << "Buffer index " << args[i] << " has dtype " << arg_dtype + << "BufferVar index " << args[i] << " has dtype " << arg_dtype << ", but previous index for the same buffer access used index type " << *index_dtype; } else { index_dtype = arg_dtype; @@ -1173,7 +1173,7 @@ void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_ AddShapeVarBounds(self, block_sref.get(), analyzer.get()); // Step 1: Input handling and error checking const SBlockNode* block_ptr = TVM_SREF_TO_SBLOCK(block_sref); - Buffer old_buffer = + BufferVar old_buffer = GetNthAccessBuffer(self, ffi::GetRef(block_ptr), buffer_index, buffer_index_type); auto index_map = LegalizeIndexMapDType(index_map_orig, old_buffer->shape); @@ -1216,8 +1216,9 @@ void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_ } // Step 2: Infer the shape of the new buffer - Buffer new_buffer = old_buffer; - new_buffer.CopyOnWrite()->shape = index_map->MapShape(old_buffer->shape, analyzer); + auto new_buffer_type = CopyBufferType(old_buffer); + new_buffer_type->shape = index_map->MapShape(old_buffer->shape, analyzer); + BufferVar new_buffer = RebuildBufferVar(old_buffer, std::move(new_buffer_type)); // Step 3: Rewrite BufferLoad/BufferStore access indices, block read/write regions, and block // alloc_buffers. @@ -1233,7 +1234,7 @@ void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_ IRModuleNode* new_mod = self->mod.CopyOnWrite(); ffi::MapObj* new_map = new_mod->functions.CopyOnWrite(); - ffi::Map new_buffer_map; + ffi::Map new_buffer_map; for (auto [var, buffer] : old_func->buffer_map) { if (buffer.same_as(old_buffer)) { buffer = new_buffer; diff --git a/src/s_tir/schedule/primitive/loop_transformation.cc b/src/s_tir/schedule/primitive/loop_transformation.cc index 845e88fe640b..f02a31defed7 100644 --- a/src/s_tir/schedule/primitive/loop_transformation.cc +++ b/src/s_tir/schedule/primitive/loop_transformation.cc @@ -487,8 +487,8 @@ class BufferIndicesMapExtractor : public StmtExprVisitor { } indices.push_back(var.value()->name); } - if (buffer_indices_map.find(store->buffer->name) == buffer_indices_map.end() && !check_) - buffer_indices_map.Set(store->buffer->name, indices); + if (buffer_indices_map.find(store->buffer.name()) == buffer_indices_map.end() && !check_) + buffer_indices_map.Set(store->buffer.name(), indices); StmtExprVisitor::VisitStmt_(store); } @@ -503,8 +503,8 @@ class BufferIndicesMapExtractor : public StmtExprVisitor { } indices.push_back(var.value()->name); } - if (buffer_indices_map.find(load->buffer->name) == buffer_indices_map.end() && !check_) - buffer_indices_map.Set(load->buffer->name, indices); + if (buffer_indices_map.find(load->buffer.name()) == buffer_indices_map.end() && !check_) + buffer_indices_map.Set(load->buffer.name(), indices); StmtExprVisitor::VisitExpr_(load); } @@ -521,10 +521,10 @@ ffi::Array MutateBufferRegion( ffi::Array new_region_arr = MutateArray(region_arr, [&buffer_indices_map, &index_range_map](const BufferRegion& region) { BufferRegion new_region = region; - auto it = buffer_indices_map.find(new_region->buffer->name); + auto it = buffer_indices_map.find(new_region->buffer.name()); if (it == buffer_indices_map.end()) return new_region; - ffi::Array old_indices = buffer_indices_map[new_region->buffer->name]; + ffi::Array old_indices = buffer_indices_map[new_region->buffer.name()]; ffi::Array new_ranges; for (size_t i = 0; i < old_indices.size(); i++) { new_ranges.push_back(index_range_map[old_indices[i]]); @@ -582,7 +582,7 @@ class BlockMutator : public StmtExprMutator { index_range_map.Set(iter->var->name, iter->dom); } - // Get the (Buffer, indices) map + // Get the (BufferVar, indices) map ffi::Map> buffer_indices_map = BufferIndicesMapExtractor::Extract(new_loop_var_, new_block); ffi::Array new_writes = diff --git a/src/s_tir/schedule/primitive/pad_einsum.cc b/src/s_tir/schedule/primitive/pad_einsum.cc index e10f8ff7d78e..f7c2d97c4f7d 100644 --- a/src/s_tir/schedule/primitive/pad_einsum.cc +++ b/src/s_tir/schedule/primitive/pad_einsum.cc @@ -123,18 +123,18 @@ class NonEinsumError : public ScheduleError { /*! \brief Data structure that represents a Einsum computation. */ struct Einsum { // The output buffer - ffi::Array output_buffers; + ffi::Array output_buffers; // The indices of the output buffer - ffi::Map> output_indices; + ffi::Map> output_indices; // The input buffers - ffi::Array input_buffers; + ffi::Array input_buffers; // The indices of the input buffers - ffi::Map> input_indices; + ffi::Map> input_indices; }; struct BufferPadding { - Buffer buffer; - Buffer padded_buffer; + BufferVar buffer; + BufferVar padded_buffer; static BufferPadding FromBufferRegion(const BufferRegion& buffer_region, const ffi::Map& iter_extents) { @@ -154,7 +154,7 @@ struct BufferPadding { shape.push_back(buffer_region->buffer->shape[i]); } } - result.padded_buffer = decl_buffer(shape, result.buffer->dtype, result.buffer->name + "_pad", + result.padded_buffer = decl_buffer(shape, result.buffer->dtype, result.buffer.name() + "_pad", result.buffer.scope()); return result; } @@ -200,7 +200,7 @@ struct BufferPadding { if (!is_read) { std::swap(read_region, write_region); } - SBlock new_block(iter_vars, {read_region}, {write_region}, padded_buffer->name, + SBlock new_block(iter_vars, {read_region}, {write_region}, padded_buffer.name(), std::move(body)); blocks->push_back(new_block); ffi::Array prim_loop_vars; @@ -217,10 +217,10 @@ struct BufferPadding { Einsum ExtractEinsum(const ScheduleState& self, const SBlock& block) { Einsum result; - std::unordered_set buffer_used; + std::unordered_set buffer_used; int n_reads = block->reads.size(); for (int i = 0; i < n_reads; ++i) { - const Buffer& buffer = block->reads[i]->buffer; + const BufferVar& buffer = block->reads[i]->buffer; if (buffer_used.count(buffer.get()) != 0) { throw NonEinsumError(self->mod, block); } @@ -234,7 +234,7 @@ Einsum ExtractEinsum(const ScheduleState& self, const SBlock& block) { } int n_writes = block->writes.size(); for (int i = 0; i < n_writes; ++i) { - const Buffer& buffer = block->writes[i]->buffer; + const BufferVar& buffer = block->writes[i]->buffer; if (buffer_used.count(buffer.get()) != 0) { throw NonEinsumError(self->mod, block); } @@ -251,7 +251,7 @@ Einsum ExtractEinsum(const ScheduleState& self, const SBlock& block) { class BufferNotAllocatedInScopeError : public ScheduleError { public: - explicit BufferNotAllocatedInScopeError(IRModule mod, Buffer buffer) + explicit BufferNotAllocatedInScopeError(IRModule mod, BufferVar buffer) : mod_(std::move(mod)), buffer_(std::move(buffer)) {} ffi::String FastErrorString() const final { @@ -261,7 +261,7 @@ class BufferNotAllocatedInScopeError : public ScheduleError { ffi::String DetailRenderTemplate() const final { std::ostringstream os; - os << "The buffer " << buffer_->name + os << "The buffer " << buffer_.name() << " is not allocated as an intermediate buffer in current PrimFunc."; return os.str(); } @@ -271,7 +271,7 @@ class BufferNotAllocatedInScopeError : public ScheduleError { private: IRModule mod_; - Buffer buffer_; + BufferVar buffer_; }; /*! \brief The schedule error class when the producer block cannot be padded. */ @@ -296,7 +296,7 @@ class InvalidProducerError : public ScheduleError { private: IRModule mod_; - Buffer buffer_; + BufferVar buffer_; SBlock producer_; }; @@ -319,7 +319,7 @@ class PadEinsumBufferReplacer : public StmtExprMutator { ffi::Array reads; reads.reserve(block->reads.size()); for (const BufferRegion& read : block->reads) { - if (ffi::Optional buffer = buffer_map_.Get(read->buffer)) { + if (ffi::Optional buffer = buffer_map_.Get(read->buffer)) { reads.push_back(BufferRegion(buffer.value(), read->region)); } else { reads.push_back(read); @@ -328,7 +328,7 @@ class PadEinsumBufferReplacer : public StmtExprMutator { ffi::Array writes; writes.reserve(block->writes.size()); for (const BufferRegion& write : block->writes) { - if (ffi::Optional buffer = buffer_map_.Get(write->buffer)) { + if (ffi::Optional buffer = buffer_map_.Get(write->buffer)) { writes.push_back(BufferRegion(buffer.value(), write->region)); } else { writes.push_back(write); @@ -354,7 +354,7 @@ class PadEinsumBufferReplacer : public StmtExprMutator { Stmt VisitStmt_(const BufferStoreNode* old_store_ptr) final { BufferStore store = StmtMutator::VisitStmt_(old_store_ptr).as_or_throw(); - if (ffi::Optional buffer = buffer_map_.Get(store->buffer)) { + if (ffi::Optional buffer = buffer_map_.Get(store->buffer)) { return BufferStore(buffer.value(), store->value, store->indices); } else { return store; @@ -363,7 +363,7 @@ class PadEinsumBufferReplacer : public StmtExprMutator { Expr VisitExpr_(const BufferLoadNode* old_load_ptr) final { BufferLoad load = ExprMutator::VisitExpr_(old_load_ptr).as_or_throw(); - if (ffi::Optional buffer = buffer_map_.Get(load->buffer)) { + if (ffi::Optional buffer = buffer_map_.Get(load->buffer)) { return BufferLoad(buffer.value(), load->indices); } else { return load; @@ -372,7 +372,7 @@ class PadEinsumBufferReplacer : public StmtExprMutator { ffi::Map iter2padded_extents; ffi::Map loop_var2padded_extent; - ffi::Map buffer_map_; + ffi::Map buffer_map_; ffi::Map block_sref_reuse_; }; @@ -437,7 +437,7 @@ void PadEinsum(ScheduleState self, const StmtSRef& block_sref, const ffi::Array< ffi::Array read_blocks; ffi::Array write_blocks; ffi::Array new_copy_blocks; - ffi::Array alloc_buffers; + ffi::Array alloc_buffers; for (const BufferRegion& buffer_region : block->reads) { if (f_needs_padding(buffer_region->region)) { BufferPadding bp = diff --git a/src/s_tir/schedule/primitive/read_write_at.cc b/src/s_tir/schedule/primitive/read_write_at.cc index 8151894e393d..cd0e3656517a 100644 --- a/src/s_tir/schedule/primitive/read_write_at.cc +++ b/src/s_tir/schedule/primitive/read_write_at.cc @@ -30,7 +30,7 @@ using namespace tvm::tirx; using support::NDIntSet; -bool HasBuffer(const ffi::Array& buffer_regions, const Buffer& buffer) { +bool HasBuffer(const ffi::Array& buffer_regions, const BufferVar& buffer) { for (const BufferRegion& buffer_region : buffer_regions) { if (buffer_region->buffer.same_as(buffer)) { return true; @@ -40,7 +40,7 @@ bool HasBuffer(const ffi::Array& buffer_regions, const Buffer& buf } void RelaxBufferRegions(const ffi::Array& buffer_regions, - const Buffer& buffer, // + const BufferVar& buffer, // const ffi::Map& var_dom, // const ffi::Map& bindings, // std::vector* relaxed_regions) { @@ -55,7 +55,7 @@ void RelaxBufferRegions(const ffi::Array& buffer_regions, class ScopeReplacer : public StmtMutator { public: - static SBlock Replace(const SBlockNode* scope_block, const Buffer& dst, const ForNode* old_loop, + static SBlock Replace(const SBlockNode* scope_block, const BufferVar& dst, const ForNode* old_loop, const ForNode* new_loop) { ffi::ObjectPtr new_scope_block = ffi::make_object(*scope_block); new_scope_block->body = ScopeReplacer(old_loop, new_loop)(std::move(new_scope_block->body)); @@ -84,7 +84,7 @@ class ScopeReplacer : public StmtMutator { class ReadWriteAtBufferReplacer : public StmtExprMutator { public: - explicit ReadWriteAtBufferReplacer(const Buffer& src, const Buffer& dst, + explicit ReadWriteAtBufferReplacer(const BufferVar& src, const BufferVar& dst, ffi::Map* block_sref_reuse) : src_(src), dst_(dst), block_sref_reuse_(block_sref_reuse) {} @@ -119,8 +119,8 @@ class ReadWriteAtBufferReplacer : public StmtExprMutator { return SBlock(new_block); } - const Buffer& src_; - const Buffer& dst_; + const BufferVar& src_; + const BufferVar& dst_; ffi::Map* block_sref_reuse_; }; @@ -130,12 +130,12 @@ struct ReadWriteAtImpl { int buffer_index, const ffi::String& storage_scope, ffi::Map annotations) { const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); - Buffer src = GetNthAccessBuffer(self, ffi::GetRef(block), buffer_index, + BufferVar src = GetNthAccessBuffer(self, ffi::GetRef(block), buffer_index, is_read ? BufferIndexType::kRead : BufferIndexType::kWrite); - Buffer dst = WithScope(src, storage_scope); + BufferVar dst = WithScope(src, storage_scope); ReadWriteAtImpl impl(self, loop_sref, src, dst, annotations); std::pair new_loop_block = - impl.MakeLoopAndBlock(src->name + "_" + storage_scope); + impl.MakeLoopAndBlock(src.name() + "_" + storage_scope); StmtSRef result_block_sref = impl.ReplaceScopeBlock(new_loop_block.first.get(), new_loop_block.second->block.get()); impl.UpdateSBlockInfo(result_block_sref, !new_loop_block.second->iter_values.empty()); @@ -225,7 +225,7 @@ struct ReadWriteAtImpl { TVM_FFI_ICHECK(w_pos.empty() || w_pos.back() < r_pos.front()); // Can be inserted at [0, r_pos.front()], i.e. before the first read insert_pos = r_pos.front(); - // Buffer reads in [insert_pos, +oo) is rewritten + // BufferVar reads in [insert_pos, +oo) is rewritten st = insert_pos; ed = n_subtrees; } else { @@ -265,7 +265,7 @@ struct ReadWriteAtImpl { return {For(new_loop), realize}; } - SBlockRealize MakeSBlock(const Buffer& copy_from, const Buffer& copy_to, + SBlockRealize MakeSBlock(const BufferVar& copy_from, const BufferVar& copy_to, const ffi::String& name_hint, const ffi::Map& loop_domain, ffi::Array domain) const { int n = domain.size(); @@ -322,8 +322,8 @@ struct ReadWriteAtImpl { /*annotations=*/annotations_)); } - explicit ReadWriteAtImpl(ScheduleState self, const StmtSRef& loop_sref, const Buffer& src, - const Buffer& dst, ffi::Map annotations) + explicit ReadWriteAtImpl(ScheduleState self, const StmtSRef& loop_sref, const BufferVar& src, + const BufferVar& dst, ffi::Map annotations) : self_(self), loop_sref_(loop_sref), loop_(nullptr), @@ -338,8 +338,8 @@ struct ReadWriteAtImpl { ScheduleState self_; const StmtSRef& loop_sref_; const ForNode* loop_; - const Buffer& src_; - const Buffer& dst_; + const BufferVar& src_; + const BufferVar& dst_; ffi::Map annotations_; ffi::Map block_sref_reuse_; arith::Analyzer analyzer_; diff --git a/src/s_tir/schedule/primitive/reduction.cc b/src/s_tir/schedule/primitive/reduction.cc index cf5e74a51efe..b2f6fcfecd00 100644 --- a/src/s_tir/schedule/primitive/reduction.cc +++ b/src/s_tir/schedule/primitive/reduction.cc @@ -70,7 +70,7 @@ class DecomposeReductionBlockReplacer : public StmtMutator { p_new_block->init = std::nullopt; // Add write regions back to read regions in update block. ffi::Array new_reads; - std::unordered_set read_bufs; + std::unordered_set read_bufs; for (const BufferRegion& read_access : block->reads) { read_bufs.insert(read_access->buffer.get()); } @@ -500,7 +500,7 @@ class NotSerialLoopKindError : public ScheduleError { class FactorAxisOutOfRangeError : public ScheduleError { public: - explicit FactorAxisOutOfRangeError(IRModule mod, Buffer buffer, int factor_axis) + explicit FactorAxisOutOfRangeError(IRModule mod, BufferVar buffer, int factor_axis) : mod_(std::move(mod)), buffer_(std::move(buffer)), factor_axis_(factor_axis) {} ffi::String FastErrorString() const final { @@ -511,7 +511,7 @@ class FactorAxisOutOfRangeError : public ScheduleError { ffi::String DetailRenderTemplate() const final { std::ostringstream os; int ndim = static_cast(buffer_->shape.size()); - os << "The write buffer " << buffer_->name << " has " << ndim + os << "The write buffer " << buffer_.name() << " has " << ndim << " dimension(s), so `factor_axis` is required to be in [" << -(ndim + 1) << ", " << ndim << "] for rfactor. However, the input `factor_axis` is " << factor_axis_ << ", which is out of the expected range"; @@ -521,7 +521,7 @@ class FactorAxisOutOfRangeError : public ScheduleError { IRModule mod() const final { return mod_; } ffi::Array LocationsOfInterest() const final { return {}; } - static int CheckAndUpdate(const IRModule& mod, const Buffer& buffer, int factor_axis) { + static int CheckAndUpdate(const IRModule& mod, const BufferVar& buffer, int factor_axis) { int ndim = static_cast(buffer->shape.size()); if (factor_axis < -(ndim + 1) || factor_axis > ndim) { throw FactorAxisOutOfRangeError(mod, buffer, factor_axis); @@ -534,7 +534,7 @@ class FactorAxisOutOfRangeError : public ScheduleError { } IRModule mod_; - Buffer buffer_; + BufferVar buffer_; int factor_axis_; }; @@ -654,20 +654,18 @@ std::unordered_map GetLoopVar2LoopMap(const ffi::Array * \param rf_loop The rfactor loop * \return The new created intermediate rfactor buffer */ -ffi::Array CreateRFactorBuffers(const ffi::Array& buf_stores, int factor_axis, +ffi::Array CreateRFactorBuffers(const ffi::Array& buf_stores, int factor_axis, const ForNode* rf_loop) { - ffi::Array rf_buffers; + ffi::Array rf_buffers; rf_buffers.reserve(buf_stores.size()); for (const BufferStore& buf_store : buf_stores) { - Buffer buffer = buf_store->buffer; + BufferVar buffer = buf_store->buffer; ffi::Array rf_shape = buffer->shape; rf_shape.insert(rf_shape.begin() + factor_axis, rf_loop->extent); - ffi::ObjectPtr n = ffi::make_object(*buffer.get()); + ffi::ObjectPtr n = CopyBufferType(buffer); n->shape = rf_shape; - n->name = buffer->name + ".rf"; - n->data = buffer->data.CopyWithSuffix(".rf"); - rf_buffers.push_back(Buffer(n)); + rf_buffers.push_back(RebuildBufferVar(buffer, std::move(n), buffer.name() + ".rf")); } return rf_buffers; } @@ -684,7 +682,7 @@ class BaseBlockCreator { public: explicit BaseBlockCreator(SBlockRealize old_block_realize, For rf_loop, ffi::Array old_reduction_updates, CommReducer reducer, - ffi::Array rf_buffers, bool is_rf_block) + ffi::Array rf_buffers, bool is_rf_block) : old_block_realize_(std::move(old_block_realize)), rf_loop_(std::move(rf_loop)), old_reduction_updates_(std::move(old_reduction_updates)), @@ -771,7 +769,7 @@ class BaseBlockCreator { ffi::Array let_vars; let_vars.reserve(n_buffers_); for (int i = 0; i < n_buffers_; ++i) { - Var var("v_" + update_buffers_[i]->name, stored_values[i].ty()); + 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(), update_indices_[i])); @@ -820,7 +818,7 @@ class BaseBlockCreator { /*! \brief The matched commutative reducer */ CommReducer reducer_; /*! \brief The intermediate rfactor buffers */ - ffi::Array rf_buffers_; + ffi::Array rf_buffers_; /*! \brief The number of rfactor buffers. */ const int n_buffers_; /*! @@ -836,7 +834,7 @@ class BaseBlockCreator { /*! \brief The new block iter bindings of the new created block-realize */ std::vector iter_values_; /*! \brief The buffers updated in this block */ - ffi::Array update_buffers_; + ffi::Array update_buffers_; /*! \brief The indices of the buffers updated in this block, respectively */ ffi::Array> update_indices_; /*! \brief The LHS values of the reduction in this block */ @@ -875,7 +873,7 @@ class RFactorBlockCreator : public BaseBlockCreator { public: explicit RFactorBlockCreator(SBlockRealize old_block_realize, For rf_loop, ffi::Array old_reduction_updates, CommReducer reducer, - ffi::Array rf_buffers, + ffi::Array rf_buffers, std::unordered_map loop_vars2loop, int factor_axis, ffi::Array combiner_rhs) : BaseBlockCreator(std::move(old_block_realize), std::move(rf_loop), @@ -948,7 +946,7 @@ class RFactorBlockCreator : public BaseBlockCreator { } void CreateReadWriteRegions() final { - ffi::Map buffer_map; + ffi::Map buffer_map; for (int i = 0; i < n_buffers_; ++i) { buffer_map.Set(old_reduction_updates_[i]->buffer, rf_buffers_[i]); } @@ -964,7 +962,7 @@ class RFactorBlockCreator : public BaseBlockCreator { region.insert( region.begin() + factor_axis_, Range::FromMinExtent(additional_iter_->var, IntImm(additional_iter_->var.ty(), 1))); - ffi::Optional rf_buffer = buffer_map.Get(write_region->buffer); + ffi::Optional rf_buffer = buffer_map.Get(write_region->buffer); TVM_FFI_ICHECK(rf_buffer.has_value()); write_regions_.push_back(BufferRegion(rf_buffer.value(), Substitute(region, var_map_))); } @@ -1000,7 +998,7 @@ class WriteBackBlockCreator : public BaseBlockCreator { public: explicit WriteBackBlockCreator(SBlockRealize old_block_realize, For rf_loop, ffi::Array old_reduction_updates, CommReducer reducer, - ffi::Array rf_buffers, IterVar rf_additional_iter, + ffi::Array rf_buffers, IterVar rf_additional_iter, ffi::Array combiner_lhs, ffi::Array rf_buf_access_indices) : BaseBlockCreator(std::move(old_block_realize), std::move(rf_loop), @@ -1141,14 +1139,14 @@ class BlockReplacer : public StmtMutator { SBlockRealize wb_block_realize, SBlockRealize old_block_realize, For rf_loop, std::unordered_set reduce_loop_vars, std::unordered_map loop_vars2loop, - const ffi::Array& rf_buffers) { + const ffi::Array& rf_buffers) { BlockReplacer replacer(std::move(rf_body), std::move(outermost_loop), std::move(wb_block_realize), std::move(old_block_realize), std::move(rf_loop), std::move(reduce_loop_vars), std::move(loop_vars2loop)); SBlock new_scope_root = replacer(std::move(scope_root_block)).as_or_throw(); SBlockNode* p = new_scope_root.CopyOnWrite(); - for (const Buffer& rf_buffer : rf_buffers) { + for (const BufferVar& rf_buffer : rf_buffers) { p->alloc_buffers.push_back(rf_buffer); } return new_scope_root; @@ -1285,7 +1283,7 @@ StmtSRef RFactor(ScheduleState self, const StmtSRef& rf_loop_sref, int factor_ax // Step 1. Create the intermediate buffer (a.k.a. rfactor buffer), which has an additional // dimension that specified by `factor_axis` and `rf_loop`. - ffi::Array rf_buffers = CreateRFactorBuffers(updates, factor_axis, rf_loop); + ffi::Array rf_buffers = CreateRFactorBuffers(updates, factor_axis, rf_loop); // Step 2. Create the rfactor block. RFactorBlockCreator rf_block_creator(block_realize, ffi::GetRef(rf_loop), updates, reducer, diff --git a/src/s_tir/schedule/primitive/rolling_buffer.cc b/src/s_tir/schedule/primitive/rolling_buffer.cc index 86a7a97e092c..abfd27f69e49 100644 --- a/src/s_tir/schedule/primitive/rolling_buffer.cc +++ b/src/s_tir/schedule/primitive/rolling_buffer.cc @@ -30,8 +30,8 @@ using namespace tvm::tirx; namespace { struct RollingBufferInfo { - Buffer old_buffer; - Buffer new_buffer; + BufferVar old_buffer; + BufferVar new_buffer; int rolling_axis; PrimExpr rolling_extent; std::vector axis_overlaps; @@ -108,7 +108,7 @@ class RollingBufferMatchError : public ScheduleError { } ffi::String DetailRenderTemplate() const final { std::ostringstream os; - os << "The target buffer " << buffer_region_->buffer->name << " with region " + os << "The target buffer " << buffer_region_->buffer.name() << " with region " << buffer_region_->region << " should have at least one dimension range that matches a rolling pattern " "such as hh.outer * stride + hh.inner. "; @@ -126,7 +126,7 @@ class RollingBufferMatchError : public ScheduleError { class RollingBufferInsertionError : public ScheduleError { public: - RollingBufferInsertionError(IRModule mod, Buffer buffer, SBlock block) + RollingBufferInsertionError(IRModule mod, BufferVar buffer, SBlock block) : mod_(mod), buffer_(std::move(buffer)), block_(block) {} ffi::String FastErrorString() const final { return "ScheduleError: rolling_buffer injection is invalid, the lca of the access " @@ -136,7 +136,7 @@ class RollingBufferInsertionError : public ScheduleError { ffi::String DetailRenderTemplate() const final { std::ostringstream os; os << "rolling_buffer injection is invalid. The block {0} should be tiled so that " - << "the lca of the access location of the target buffer " << buffer_->name + << "the lca of the access location of the target buffer " << buffer_.name() << " is a for loop. "; return os.str(); } @@ -145,7 +145,7 @@ class RollingBufferInsertionError : public ScheduleError { private: IRModule mod_; - Buffer buffer_; + BufferVar buffer_; SBlock block_; }; @@ -164,7 +164,7 @@ class RollingBufferInfoCollector { private: bool MatchRollingBuffer(const StmtSRef& block_sref, const BufferRegion& buffer_region) { - const Buffer& buffer = buffer_region->buffer; + const BufferVar& buffer = buffer_region->buffer; const Region& region = buffer_region->region; std::vector> bound_iter_vars; @@ -238,8 +238,9 @@ class RollingBufferInfoCollector { } ffi::Array new_shape = buffer->shape; new_shape.Set(roll_axis, region[roll_axis]->extent); - Buffer new_buffer = buffer; - new_buffer.CopyOnWrite()->shape = new_shape; + auto new_buffer_type = CopyBufferType(buffer); + new_buffer_type->shape = new_shape; + BufferVar new_buffer = RebuildBufferVar(buffer, std::move(new_buffer_type)); info_.old_buffer = buffer; info_.new_buffer = new_buffer; @@ -277,7 +278,7 @@ class RollingBufferRewriter : public StmtExprMutator { (*old_access_regions).MutateByApply(fmutate); } - void RewriteBufferAccess(Buffer* buffer, ffi::Array* indices) const { + void RewriteBufferAccess(BufferVar* buffer, ffi::Array* indices) const { ffi::Array new_indices; new_indices.reserve(indices->size()); // First modify the access indices to use modulo arithmetic @@ -299,8 +300,8 @@ class RollingBufferRewriter : public StmtExprMutator { SBlock stmt = StmtExprMutator::VisitStmt_(block).as_or_throw(); SBlockNode* n = stmt.CopyOnWrite(); if (block == scope_sref_->stmt) { - ffi::Array new_alloc_buffers; - for (const Buffer& buffer : stmt->alloc_buffers) { + ffi::Array new_alloc_buffers; + for (const BufferVar& buffer : stmt->alloc_buffers) { if (buffer != info_->old_buffer) { new_alloc_buffers.push_back(buffer); } else { @@ -326,7 +327,8 @@ class RollingBufferRewriter : public StmtExprMutator { new_iter_vars.push_back(old_iter_var); } } - ffi::Map buffer_data_to_buffer = {{info_->new_buffer->data, info_->new_buffer}}; + ffi::Map buffer_data_to_buffer = { + {info_->new_buffer.var(), info_->new_buffer}}; auto infered_access_regions = GetSBlockReadWriteRegion(stmt, buffer_data_to_buffer); n->iter_vars = std::move(new_iter_vars); diff --git a/src/s_tir/schedule/schedule.cc b/src/s_tir/schedule/schedule.cc index e8e666672f08..6275b8463332 100644 --- a/src/s_tir/schedule/schedule.cc +++ b/src/s_tir/schedule/schedule.cc @@ -333,7 +333,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def_method("s_tir.schedule.ScheduleDecomposePadding", &ScheduleNode::DecomposePadding) .def_method("s_tir.schedule.SchedulePadEinsum", &ScheduleNode::PadEinsum); } -/******** (FFI) Buffer transformation ********/ +/******** (FFI) BufferVar transformation ********/ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def_method("s_tir.schedule.ScheduleRollingBuffer", diff --git a/src/s_tir/schedule/state.cc b/src/s_tir/schedule/state.cc index 0c3deac1b87f..9f6528fc7b86 100644 --- a/src/s_tir/schedule/state.cc +++ b/src/s_tir/schedule/state.cc @@ -272,11 +272,11 @@ class SBlockInfoCollector : private StmtVisitor { continue; } // For each buffer, record the regions generated under this loop - std::unordered_map>> + std::unordered_map>> touched_regions; // Step 2.3.1. Find all the regions read by the consumer that we care about for (const BufferRegion& region : block_reads_unbound.at(consumer_block_sref.get())) { - const BufferNode* buffer = region->buffer.get(); + const VarNode* buffer = region->buffer.get(); touched_regions[buffer] = {}; } // Step 2.3.2. Find all the regions written by each producer @@ -284,7 +284,7 @@ class SBlockInfoCollector : private StmtVisitor { const SBlockRealize& producer_realize = block2realize_.at(producer_block_sref->stmt); StmtSRef parent_sref = ffi::GetRef(producer_block_sref->parent); for (const BufferRegion& region : block_writes_unbound.at(producer_block_sref)) { - const BufferNode* buffer = region->buffer.get(); + const VarNode* buffer = region->buffer.get(); auto it = touched_regions.find(buffer); // Skip the regions that is not read by the consumer if (it != touched_regions.end()) { @@ -306,7 +306,7 @@ class SBlockInfoCollector : private StmtVisitor { { StmtSRef parent_sref = ffi::GetRef(consumer_block_sref->parent); for (const BufferRegion& region : block_reads_unbound.at(consumer_block_sref.get())) { - const BufferNode* buffer = region->buffer.get(); + const VarNode* buffer = region->buffer.get(); const std::vector>& touched_region = touched_regions.at(buffer); if (!touched_region.empty()) { @@ -318,7 +318,8 @@ class SBlockInfoCollector : private StmtVisitor { /*dom_low_inclusive=*/parent_sref, /*dom_high_exclusive=*/lca, /*analyzer=*/analyzer_.get()); - if (!ProducerCoversConsumer(buffer->shape, produced_region, consumed_region, + if (!ProducerCoversConsumer(GetBufferVar(buffer)->shape, produced_region, + consumed_region, analyzer_.get())) { region_cover = false; self_->block_info.at(consumer_block_sref).region_cover = region_cover; diff --git a/src/s_tir/schedule/traced_schedule.cc b/src/s_tir/schedule/traced_schedule.cc index 646269abfb7e..2849183bc43a 100644 --- a/src/s_tir/schedule/traced_schedule.cc +++ b/src/s_tir/schedule/traced_schedule.cc @@ -759,7 +759,7 @@ void TracedScheduleNode::PadEinsum(const SBlockRV& block_rv, const ffi::Array& padding) final; - /******** Schedule: Buffer transformation ********/ + /******** Schedule: BufferVar transformation ********/ void RollingBuffer(const SBlockRV& block_rv, int write_buffer_index) final; /******** Schedule: Misc ********/ void EnterPostproc() final; diff --git a/src/s_tir/schedule/transform.cc b/src/s_tir/schedule/transform.cc index 4c7ed8a074cb..976c39f55d1a 100644 --- a/src/s_tir/schedule/transform.cc +++ b/src/s_tir/schedule/transform.cc @@ -38,27 +38,25 @@ SBlock WithAnnotation(const SBlockNode* block, const ffi::String& attr_key, return SBlock(new_block); } -/******** Buffer Related ********/ -Buffer WithScope(const Buffer& buffer, const ffi::String& scope) { - ffi::ObjectPtr new_buffer = ffi::make_object(*buffer.get()); - const auto* ptr_type = TVM_TYPE_AS(buffer->data->ty, PointerTypeNode); - Type new_type = PointerType(ptr_type->element_type, scope); - new_buffer->data = tirx::Var(buffer->data->name + "_" + scope, new_type); - new_buffer->name = buffer->name + "_" + scope; - return Buffer(new_buffer); +/******** BufferVar Related ********/ +BufferVar WithScope(const BufferVar& buffer, const ffi::String& scope) { + const auto* ptr_type = TVM_TYPE_AS(buffer->data_pointer_type, PointerTypeNode); + BufferType new_type(PointerType(ptr_type->element_type, scope), buffer->dtype, buffer->shape, + buffer->strides, buffer->elem_offset, buffer->data_alignment, + buffer->offset_factor, buffer->layout, buffer->allocated_addr); + return BufferVar(buffer.name() + "_" + scope, new_type, buffer.span()); } -Buffer WithDType(const Buffer& buffer, PrimType dtype) { - ffi::ObjectPtr new_buffer = ffi::make_object(*buffer.get()); - new_buffer->dtype = dtype; - const auto* ptr_type = TVM_TYPE_AS(buffer->data->ty, PointerTypeNode); - new_buffer->data = tirx::Var(buffer->data->name, PointerType(dtype, ptr_type->storage_scope)); - new_buffer->name = buffer->name; - return Buffer(new_buffer); +BufferVar WithDType(const BufferVar& buffer, PrimType dtype) { + const auto* ptr_type = TVM_TYPE_AS(buffer->data_pointer_type, PointerTypeNode); + BufferType new_type(PointerType(dtype, ptr_type->storage_scope), dtype, buffer->shape, + buffer->strides, buffer->elem_offset, buffer->data_alignment, + buffer->offset_factor, buffer->layout, buffer->allocated_addr); + return BufferVar(buffer.name(), new_type, buffer.span()); } -ffi::Array ReplaceBuffer(ffi::Array regions, const Buffer& source, - const Buffer& target) { +ffi::Array ReplaceBuffer(ffi::Array regions, const BufferVar& source, + const BufferVar& target) { regions.MutateByApply([&source, &target](BufferRegion region) -> BufferRegion { if (region->buffer.same_as(source)) { ffi::ObjectPtr n = ffi::make_object(*region.get()); @@ -71,7 +69,7 @@ ffi::Array ReplaceBuffer(ffi::Array regions, const B } ffi::Array ReplaceBuffer(ffi::Array regions, - const ffi::Map& buffer_map) { + const ffi::Map& buffer_map) { regions.MutateByApply([&buffer_map](BufferRegion region) -> BufferRegion { if (buffer_map.count(region->buffer)) { ffi::ObjectPtr n = ffi::make_object(*region.get()); @@ -84,7 +82,7 @@ ffi::Array ReplaceBuffer(ffi::Array regions, } ffi::Array ReplaceBuffer(ffi::Array match_buffers, - const Buffer& source, const Buffer& target) { + const BufferVar& source, const BufferVar& target) { match_buffers.MutateByApply( [&source, &target](MatchBufferRegion match_buffer) -> MatchBufferRegion { if (match_buffer->source->buffer.same_as(source)) { @@ -99,7 +97,7 @@ ffi::Array ReplaceBuffer(ffi::Array match_ } ffi::Array ReplaceBufferRegion(ffi::Array regions, - const Buffer& source_buffer, + const BufferVar& source_buffer, const BufferRegion& target) { regions.MutateByApply([&source_buffer, &target](const BufferRegion& region) -> BufferRegion { if (region->buffer.same_as(source_buffer)) { @@ -111,7 +109,7 @@ ffi::Array ReplaceBufferRegion(ffi::Array regions, } ffi::Array ReplaceBufferRegion(ffi::Array match_buffers, - const Buffer& source_buffer, + const BufferVar& source_buffer, const BufferRegion& target) { match_buffers.MutateByApply( [&source_buffer, &target](const MatchBufferRegion& match_buffer) -> MatchBufferRegion { @@ -127,23 +125,23 @@ ffi::Array ReplaceBufferRegion(ffi::Array } /******** ReplaceBufferMutator ********/ -ReplaceBufferMutator::ReplaceBufferMutator(const Buffer& old_buffer, Buffer new_buffer, +ReplaceBufferMutator::ReplaceBufferMutator(const BufferVar& old_buffer, BufferVar new_buffer, ffi::Map* block_sref_reuse) : block_sref_reuse_(block_sref_reuse) { - buffer_var_map_[old_buffer->data.get()] = std::move(new_buffer); + buffer_var_map_[old_buffer.get()] = std::move(new_buffer); } -ReplaceBufferMutator::ReplaceBufferMutator(const ffi::Map& buffer_map, +ReplaceBufferMutator::ReplaceBufferMutator(const ffi::Map& buffer_map, ffi::Map* block_sref_reuse) : block_sref_reuse_(block_sref_reuse) { for (const auto& [old_buffer, new_buffer] : buffer_map) { - buffer_var_map_[old_buffer->data.get()] = new_buffer; + buffer_var_map_[old_buffer.get()] = new_buffer; } } Expr ReplaceBufferMutator::VisitExpr_(const VarNode* var) { auto it = buffer_var_map_.find(var); - return it != buffer_var_map_.end() ? it->second->data : ffi::GetRef(var); + return it != buffer_var_map_.end() ? it->second.var() : ffi::GetRef(var); } Stmt ReplaceBufferMutator::VisitStmt_(const BufferStoreNode* op) { @@ -158,7 +156,7 @@ Expr ReplaceBufferMutator::VisitExpr_(const BufferLoadNode* op) { MatchBufferRegion ReplaceBufferMutator::VisitMatchBufferRegion( const MatchBufferRegion& match_buffer) { - auto it = buffer_var_map_.find(match_buffer->source->buffer->data.get()); + auto it = buffer_var_map_.find(match_buffer->source->buffer.get()); if (it != buffer_var_map_.end()) { return MatchBufferRegion(match_buffer->buffer, BufferRegion(it->second, match_buffer->source->region)); @@ -186,8 +184,8 @@ Stmt ReplaceBufferMutator::VisitStmt_(const SBlockNode* block) { } }); - Buffer buf = [&]() { - auto it = buffer_var_map_.find(buffer_region->buffer->data.get()); + BufferVar buf = [&]() { + auto it = buffer_var_map_.find(buffer_region->buffer.get()); if (it == buffer_var_map_.end()) { return buffer_region->buffer; } else { @@ -201,8 +199,8 @@ Stmt ReplaceBufferMutator::VisitStmt_(const SBlockNode* block) { return BufferRegion(buf, region); } }; - auto f_mutate_alloc_buffers = [this](const Buffer& buffer) { - auto it = buffer_var_map_.find(buffer->data.get()); + auto f_mutate_alloc_buffers = [this](const BufferVar& buffer) { + auto it = buffer_var_map_.find(buffer.get()); return it == buffer_var_map_.end() ? buffer : it->second; }; @@ -212,7 +210,7 @@ Stmt ReplaceBufferMutator::VisitStmt_(const SBlockNode* block) { ffi::Array reads = block->reads.Map(f_mutate_read_write_region); ffi::Array writes = block->writes.Map(f_mutate_read_write_region); // Step 3. Mutate `alloc_buffers` for the old buffer allocated in this block. - ffi::Array alloc_buffers = block->alloc_buffers.Map(f_mutate_alloc_buffers); + ffi::Array alloc_buffers = block->alloc_buffers.Map(f_mutate_alloc_buffers); // Step 4. Recursively mutate the block. SBlock mutated_block = StmtMutator::VisitStmt_(block).as_or_throw(); diff --git a/src/s_tir/schedule/transform.h b/src/s_tir/schedule/transform.h index af00f9f8cc66..57e422b1d947 100644 --- a/src/s_tir/schedule/transform.h +++ b/src/s_tir/schedule/transform.h @@ -45,7 +45,7 @@ using namespace tvm::tirx; SBlock WithAnnotation(const SBlockNode* block, const ffi::String& attr_key, const ffi::ObjectRef& attr_value); -/******** Buffer Related ********/ +/******** BufferVar Related ********/ /*! * \brief Create a new buffer by changing the storage scope. @@ -53,7 +53,7 @@ SBlock WithAnnotation(const SBlockNode* block, const ffi::String& attr_key, * \param scope The target storage scope. * \return The new buffer with target storage scope. */ -Buffer WithScope(const Buffer& buffer, const ffi::String& scope); +BufferVar WithScope(const BufferVar& buffer, const ffi::String& scope); /*! * \brief Create a new buffer by changint the data type. @@ -61,7 +61,7 @@ Buffer WithScope(const Buffer& buffer, const ffi::String& scope); * \param scope The target data type. * \return The new buffer with target data type. */ -Buffer WithDType(const Buffer& buffer, PrimType dtype); +BufferVar WithDType(const BufferVar& buffer, PrimType dtype); /*! * \brief Replaces the buffer within the specific sequence of regions @@ -70,8 +70,8 @@ Buffer WithDType(const Buffer& buffer, PrimType dtype); * \param target The buffer to be replaced to * \return The new sequence of regions after replacement */ -ffi::Array ReplaceBuffer(ffi::Array regions, const Buffer& source, - const Buffer& target); +ffi::Array ReplaceBuffer(ffi::Array regions, const BufferVar& source, + const BufferVar& target); /*! * \brief Replaces the buffer within the specific sequence of regions @@ -80,7 +80,7 @@ ffi::Array ReplaceBuffer(ffi::Array regions, const B * \return The new sequence of regions after replacement */ ffi::Array ReplaceBuffer(ffi::Array regions, - const ffi::Map& buffer_map); + const ffi::Map& buffer_map); /*! * \brief Replaces the buffer within the specific sequence of match_buffers @@ -90,7 +90,7 @@ ffi::Array ReplaceBuffer(ffi::Array regions, * \return The new sequence of match_buffers after replacement */ ffi::Array ReplaceBuffer(ffi::Array match_buffers, - const Buffer& source, const Buffer& target); + const BufferVar& source, const BufferVar& target); /*! * \brief Replaces the buffer region within the specific sequence of regions @@ -100,7 +100,7 @@ ffi::Array ReplaceBuffer(ffi::Array match_ * \return The new sequence of regions after replacement */ ffi::Array ReplaceBufferRegion(ffi::Array regions, - const Buffer& source_buffer, + const BufferVar& source_buffer, const BufferRegion& target); /*! @@ -111,7 +111,7 @@ ffi::Array ReplaceBufferRegion(ffi::Array regions, * \return The new sequence of match_buffers after replacement */ ffi::Array ReplaceBufferRegion(ffi::Array match_buffers, - const Buffer& source_buffer, + const BufferVar& source_buffer, const BufferRegion& target); /*! @@ -131,10 +131,10 @@ class ReplaceBufferMutator : public StmtExprMutator { * \param block_sref_reuse Optional map to record mapping between old and new blocks that reuse * sref. */ - ReplaceBufferMutator(const Buffer& old_buffer, Buffer new_buffer, + ReplaceBufferMutator(const BufferVar& old_buffer, BufferVar new_buffer, ffi::Map* block_sref_reuse); - ReplaceBufferMutator(const ffi::Map& buffer_map, + ReplaceBufferMutator(const ffi::Map& buffer_map, ffi::Map* block_sref_reuse); protected: @@ -145,7 +145,7 @@ class ReplaceBufferMutator : public StmtExprMutator { template Node VisitBufferAccess(Node node) { - auto it = buffer_var_map_.find(node->buffer->data.get()); + auto it = buffer_var_map_.find(node->buffer.get()); if (it != buffer_var_map_.end()) { node.CopyOnWrite()->buffer = it->second; } @@ -164,7 +164,7 @@ class ReplaceBufferMutator : public StmtExprMutator { * \brief A mapping which maps old buffer vars to new buffers, including the buffers defined in * MatchBufferRegion. */ - std::unordered_map buffer_var_map_; + std::unordered_map buffer_var_map_; /*! \brief The block sref reuse map for the following replacement */ ffi::Map* block_sref_reuse_; }; diff --git a/src/s_tir/transform/bound_checker.cc b/src/s_tir/transform/bound_checker.cc index 73436123d06d..9dfb216ca095 100644 --- a/src/s_tir/transform/bound_checker.cc +++ b/src/s_tir/transform/bound_checker.cc @@ -72,8 +72,8 @@ class BoundChecker : public StmtExprMutator { : mem_to_shape_(mem_to_shape) {} Stmt VisitStmt_(const AllocBufferNode* op) final { - if (UpdateIsNeeded(op->buffer->data)) { - Update(op->buffer->data, op->buffer->shape, op->buffer->dtype); + if (UpdateIsNeeded(op->buffer.var())) { + Update(op->buffer.var(), op->buffer->shape, op->buffer->dtype); } return StmtExprMutator::VisitStmt_(op); } @@ -91,8 +91,8 @@ class BoundChecker : public StmtExprMutator { unsafe_rewritten_ = false; StmtExprMutator::VisitStmt_(op); process_store_ = false; - if (CanInstrument(op->indices, op->buffer->data)) { - Collect(op->indices, op->buffer->data); + if (CanInstrument(op->indices, op->buffer.var())) { + Collect(op->indices, op->buffer.var()); } // The collector should has at least one item. if (store_scope_bound_collector_.size()) { @@ -109,8 +109,8 @@ class BoundChecker : public StmtExprMutator { } Expr VisitExpr_(const BufferLoadNode* op) final { - if (CanInstrument(op->indices, op->buffer->data)) { - Collect(op->indices, op->buffer->data); + if (CanInstrument(op->indices, op->buffer.var())) { + Collect(op->indices, op->buffer.var()); } return StmtExprMutator::VisitExpr_(op); } diff --git a/src/s_tir/transform/compact_buffer_region.cc b/src/s_tir/transform/compact_buffer_region.cc index b963af8b4069..c119956ca296 100644 --- a/src/s_tir/transform/compact_buffer_region.cc +++ b/src/s_tir/transform/compact_buffer_region.cc @@ -70,38 +70,38 @@ NDIntSet NDIntSetEval(Region region, PrimExpr predicate, class Var2BufferCollector : public StmtExprVisitor { public: /*! \brief Map the buffer var to all aliased buffers. */ - std::unordered_map> + std::unordered_map> var2buffer_; private: void VisitStmt_(const BufferStoreNode* op) final { - var2buffer_[op->buffer->data].insert(op->buffer); + var2buffer_[op->buffer.var()].insert(op->buffer); StmtExprVisitor::VisitStmt_(op); } void VisitExpr_(const BufferLoadNode* op) final { - var2buffer_[op->buffer->data].insert(op->buffer); + var2buffer_[op->buffer.var()].insert(op->buffer); StmtExprVisitor::VisitExpr_(op); } void VisitStmt_(const SBlockNode* op) final { - for (const Buffer& buffer : op->alloc_buffers) { - var2buffer_[buffer->data].insert(buffer); + for (const BufferVar& buffer : op->alloc_buffers) { + var2buffer_[buffer.var()].insert(buffer); } for (const MatchBufferRegion& region : op->match_buffers) { - var2buffer_[region->buffer->data].insert(region->buffer); - var2buffer_[region->source->buffer->data].insert(region->source->buffer); + var2buffer_[region->buffer.var()].insert(region->buffer); + var2buffer_[region->source->buffer.var()].insert(region->source->buffer); } StmtExprVisitor::VisitStmt_(op); } void VisitStmt_(const DeclBufferNode* op) final { - var2buffer_[op->buffer->data].insert(op->buffer); + var2buffer_[op->buffer.var()].insert(op->buffer); StmtExprVisitor::VisitStmt_(op); } void VisitStmt_(const AllocBufferNode* op) final { - var2buffer_[op->buffer->data].insert(op->buffer); + var2buffer_[op->buffer.var()].insert(op->buffer); StmtExprVisitor::VisitStmt_(op); } }; @@ -112,7 +112,7 @@ class Var2BufferCollector : public StmtExprVisitor { */ class BufferAccessRegionCollector : public StmtExprVisitor { public: - static std::unordered_map Collect( + static std::unordered_map Collect( const PrimFunc& f, bool collect_inbound) { BufferAccessRegionCollector region_collector(collect_inbound); @@ -133,11 +133,11 @@ class BufferAccessRegionCollector : public StmtExprVisitor { struct BufferAccessInfo { /*! \brief The buffer. */ - Buffer buffer; + BufferVar buffer; /*! \brief The buffer access region, which can be updated during visiting. */ NDIntSet accessed_region; - explicit BufferAccessInfo(const Buffer& buffer, const NDIntSet& region) + explicit BufferAccessInfo(const BufferVar& buffer, const NDIntSet& region) : buffer(buffer), accessed_region(region) {} }; @@ -243,7 +243,7 @@ class BufferAccessRegionCollector : public StmtExprVisitor { TVM_FFI_ICHECK(!op->init.has_value()); TVM_FFI_ICHECK_EQ(op->iter_vars.size(), 0) << "CompactBufferRegion only works on opaque blocks"; // Step 1. Record and update current read/write region annotations - std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> + std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> cur_access_annotations; for (const BufferRegion& region : op->reads) { cur_access_annotations[region->buffer].push_back(region); @@ -277,8 +277,8 @@ class BufferAccessRegionCollector : public StmtExprVisitor { record_explicit_region(s_tir::attr::explicit_write_region, BufferIndexType::kWrite); // Step 3. Record relax position of ancestor_loops_ - for (const Buffer& buffer : op->alloc_buffers) { - VisitBufferDef(buffer->data); + for (const BufferVar& buffer : op->alloc_buffers) { + VisitBufferDef(buffer.var()); } // Step 4. Visit match buffers for (const MatchBufferRegion& region : op->match_buffers) { @@ -298,8 +298,8 @@ class BufferAccessRegionCollector : public StmtExprVisitor { // Step 7. Clear explicit access annotations explicit_access_annotations_.clear(); // Step 8. Update buffer_access_region_ from relaxed_accesses_ for inner buffers. - for (const Buffer& buffer : op->alloc_buffers) { - TVM_FFI_ICHECK_EQ(var2buffer_[buffer->data].size(), 1) + for (const BufferVar& buffer : op->alloc_buffers) { + TVM_FFI_ICHECK_EQ(var2buffer_[buffer.var()].size(), 1) << "Block allocation buffer shoud not be alised"; SimplifyAndNarrowBufferRegionFromNDIntSet(buffer); } @@ -312,7 +312,7 @@ class BufferAccessRegionCollector : public StmtExprVisitor { void VisitStmt_(const AllocBufferNode* op) final { // AllocBuffer is flat: register the buffer def and track for post-scope compaction. - VisitBufferDef(op->buffer->data); + VisitBufferDef(op->buffer.var()); pending_flat_alloc_buffers_.push_back(op->buffer); return StmtExprVisitor::VisitStmt_(op); } @@ -347,8 +347,8 @@ class BufferAccessRegionCollector : public StmtExprVisitor { } void VisitBufferAccess(const BufferRegion& buffer_region) { - const Buffer& buffer = buffer_region->buffer; - auto it = buffer_scope_depth_.find(buffer->data); + const BufferVar& buffer = buffer_region->buffer; + auto it = buffer_scope_depth_.find(buffer.var()); if (it != buffer_scope_depth_.end()) { size_t n_ancestor_loops = it->second; // Step 1. Stop ancestor loop vars out of the allocation block from @@ -400,7 +400,7 @@ class BufferAccessRegionCollector : public StmtExprVisitor { if (it == var2buffer_.end()) { return; } - for (const Buffer& buffer : it->second) { + for (const BufferVar& buffer : it->second) { auto annotation_it = access_annotations_.find(buffer); if (annotation_it != access_annotations_.end()) { // opaque buffer has explicit accessed region annotations @@ -428,7 +428,7 @@ class BufferAccessRegionCollector : public StmtExprVisitor { * Update the `relaxed_accesses_` dict. If `collect_inbound_` is true, * the result region would never exceed the original buffer shape. */ - void SimplifyAndNarrowBufferRegionFromNDIntSet(const Buffer& buffer) { + void SimplifyAndNarrowBufferRegionFromNDIntSet(const BufferVar& buffer) { auto it = relaxed_accesses_.find(buffer); TVM_FFI_ICHECK(it != relaxed_accesses_.end()) << buffer << " is allocated but not accessed within block scope"; @@ -488,7 +488,7 @@ class BufferAccessRegionCollector : public StmtExprVisitor { */ void CompactPendingFlatAllocBuffers(size_t n_before = 0) { for (size_t i = n_before; i < pending_flat_alloc_buffers_.size(); ++i) { - const Buffer& buf = pending_flat_alloc_buffers_[i]; + const BufferVar& buf = pending_flat_alloc_buffers_[i]; auto it = relaxed_accesses_.find(buf); if (it != relaxed_accesses_.end()) { SimplifyAndNarrowBufferRegionFromNDIntSet(buf); @@ -501,7 +501,7 @@ class BufferAccessRegionCollector : public StmtExprVisitor { /*! \brief Only collect accessed region within original buffer shape bound. */ bool collect_inbound_{true}; /*! \brief Pending flat AllocBuffer nodes to compact when leaving scope. */ - std::vector pending_flat_alloc_buffers_; + std::vector pending_flat_alloc_buffers_; /*! \brief The iteration scopes from the current node up to the root. */ std::vector ancestor_iters_; @@ -514,7 +514,7 @@ class BufferAccessRegionCollector : public StmtExprVisitor { std::unordered_map buffer_scope_depth_; /*! \brief Map the buffer var to all aliased buffers. */ - std::unordered_map> + std::unordered_map> var2buffer_; /*! \brief The map from loop vars to their iter range. */ @@ -525,21 +525,21 @@ class BufferAccessRegionCollector : public StmtExprVisitor { std::vector pending_conditions_; /*! \brief The analyzer aware of loop domains. */ arith::Analyzer dom_analyzer_; - /*! \brief The map from Buffer to it's relaxed access set. */ - std::unordered_map relaxed_accesses_; + /*! \brief The map from BufferVar to it's relaxed access set. */ + std::unordered_map relaxed_accesses_; /*! - * \brief The map from Buffer to it entire access region, used for returning. + * \brief The map from BufferVar to it entire access region, used for returning. * The entire access region should get updated on the buffer's define point * and we sanity check that every buffer is defined only once. */ - std::unordered_map buffer_access_region_; + std::unordered_map buffer_access_region_; - /*! \brief The map from Buffer to it's access regions annotated by current block. */ - std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> + /*! \brief The map from BufferVar to it's access regions annotated by current block. */ + std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> access_annotations_; - /*! \brief The map from Buffer to its explicit access region annotated by the block. */ - std::unordered_map + /*! \brief The map from BufferVar to its explicit access region annotated by the block. */ + std::unordered_map explicit_access_annotations_; }; @@ -560,14 +560,24 @@ struct BufferAllocInfo { * \brief The reallocated buffer with minimal size. * \note The value if std::nullopt if the buffer do not need reallocate (e.g parameter buffer). */ - Buffer new_buffer; + BufferVar new_buffer; }; /*! \brief Reallocate the buffers with minimal region. */ class BufferCompactor : public StmtExprMutator { public: explicit BufferCompactor(std::unordered_map buffer_info) - : buffer_info_(std::move(buffer_info)) {} + : buffer_info_(std::move(buffer_info)) { + std::vector> remapped_entries; + for (const auto& [old_buffer, info] : buffer_info_) { + if (!old_buffer.same_as(info.new_buffer.var())) { + remapped_entries.emplace_back(info.new_buffer.var(), info); + } + } + for (auto& [new_buffer, info] : remapped_entries) { + buffer_info_.emplace(std::move(new_buffer), std::move(info)); + } + } Stmt VisitStmt_(const BufferStoreNode* _op) final { BufferStore store = StmtExprMutator::VisitStmt_(_op).as_or_throw(); @@ -587,8 +597,8 @@ class BufferCompactor : public StmtExprMutator { // Step 0. Check there is no Init part. TVM_FFI_ICHECK(!op->init.has_value()); // Step 1. Reallocate and rewrite alloc_buffers, also update BufferAllocInfo. - ffi::Array alloc_buffers = - op->alloc_buffers.Map([this](const Buffer& buf) { return RewriteAllocBuffer(buf); }); + ffi::Array alloc_buffers = + op->alloc_buffers.Map([this](const BufferVar& buf) { return RewriteAllocBuffer(buf); }); // Step 2. Recursively rewrite BufferLoad/BufferStore. SBlock block = StmtExprMutator::VisitStmt_(op).as_or_throw(); // Step 3. Update block signature. @@ -601,22 +611,18 @@ class BufferCompactor : public StmtExprMutator { } Stmt VisitStmt_(const DeclBufferNode* op) final { - Buffer new_buffer = RewriteAllocBuffer(op->buffer); - if (new_buffer.same_as(op->buffer)) { - return ffi::GetRef(op); - } - auto n = CopyOnWrite(op); - n->buffer = std::move(new_buffer); - return DeclBuffer(n); + RewriteAllocBuffer(op->buffer); + return StmtExprMutator::VisitStmt_(op); } Stmt VisitStmt_(const AllocBufferNode* op) final { + RewriteAllocBuffer(op->buffer); AllocBuffer alloc_buf = StmtExprMutator::VisitStmt_(op).as_or_throw(); - auto it = buffer_info_.find(alloc_buf->buffer->data); + auto it = buffer_info_.find(alloc_buf->buffer.var()); if (it == buffer_info_.end()) { return alloc_buf; } - const Buffer& new_buffer = it->second.new_buffer; + const BufferVar& new_buffer = it->second.new_buffer; if (op->buffer->dtype != new_buffer->dtype) { return alloc_buf; } @@ -624,16 +630,20 @@ class BufferCompactor : public StmtExprMutator { return alloc_buf; } - Buffer RewriteAllocBuffer(const Buffer& buffer) { - auto it = buffer_info_.find(buffer->data); + BufferVar RewriteAllocBuffer(const BufferVar& buffer) { + auto it = buffer_info_.find(buffer.var()); if (it != buffer_info_.end()) { - return it->second.new_buffer; + const BufferVar& new_buffer = it->second.new_buffer; + if (!new_buffer.same_as(buffer)) { + buffer_remap_.Set(buffer, new_buffer); + } + return new_buffer; } return buffer; } - void RewriteBufferAccess(Buffer* buffer, ffi::Array* indices) const { - auto it = buffer_info_.find((*buffer)->data); + void RewriteBufferAccess(BufferVar* buffer, ffi::Array* indices) const { + auto it = buffer_info_.find((*buffer).var()); if (it == buffer_info_.end()) { return; } @@ -649,8 +659,8 @@ class BufferCompactor : public StmtExprMutator { *indices = std::move(new_indices); } - void RewriteBufferRegion(Buffer* buffer, Region* region) const { - auto it = buffer_info_.find((*buffer)->data); + void RewriteBufferRegion(BufferVar* buffer, Region* region) const { + auto it = buffer_info_.find((*buffer).var()); if (it == buffer_info_.end()) { // Skip if the buffer is parameter return; @@ -721,16 +731,16 @@ ffi::Array CalcStrides(const BufferAllocInfo& alloc_info, Stmt BufferCompactorCompact( const PrimFunc& f, - const std::unordered_map& regions, + const std::unordered_map& regions, const std::unordered_map& storage_align) { // collect buffer allocation info for no-alias buffers std::unordered_map buffer_info; for (const auto& kv : regions) { - const Buffer& buffer = kv.first; + const BufferVar& buffer = kv.first; // set dim alignment info Region region = kv.second; BufferAllocInfo alloc_info; - auto it = storage_align.find(buffer->data); + auto it = storage_align.find(buffer.var()); if (it != storage_align.end()) { std::vector dim_aligns(buffer->shape.size()); for (const StorageAlignTuple& dim_align : (*it).second) { @@ -745,12 +755,12 @@ Stmt BufferCompactorCompact( // prepare new buffer ffi::Array shape = region.Map([](const Range& range) { return range->extent; }); ffi::Array strides = CalcStrides(alloc_info, shape); - ffi::ObjectPtr n = ffi::make_object(*buffer.get()); + ffi::ObjectPtr n = CopyBufferType(buffer); n->shape = std::move(shape); n->strides = std::move(strides); - alloc_info.new_buffer = Buffer(std::move(n)); + alloc_info.new_buffer = RebuildBufferVar(buffer, std::move(n)); alloc_info.region = region; - buffer_info.emplace(buffer->data, std::move(alloc_info)); + buffer_info.emplace(buffer.var(), std::move(alloc_info)); } BufferCompactor compactor(std::move(buffer_info)); Stmt stmt = compactor(f->body); diff --git a/src/s_tir/transform/inject_double_buffer.cc b/src/s_tir/transform/inject_double_buffer.cc index d01ce9e6b49e..63cf54062659 100644 --- a/src/s_tir/transform/inject_double_buffer.cc +++ b/src/s_tir/transform/inject_double_buffer.cc @@ -114,7 +114,7 @@ class DoubleBufferInjector : public StmtExprMutator { } Stmt VisitStmt_(const AllocBufferNode* op) final { - const VarNode* buf = op->buffer->data.get(); + const VarNode* buf = op->buffer.get(); auto it = dbuffer_info_.find(buf); if (it != dbuffer_info_.end()) { StorageEntry& entry = it->second; @@ -201,7 +201,7 @@ class DoubleBufferInjector : public StmtExprMutator { Stmt VisitStmt_(const BufferStoreNode* op) final { auto node = StmtExprMutator::VisitStmt_(op).as_or_throw(); - auto it = dbuffer_info_.find(node->buffer->data.get()); + auto it = dbuffer_info_.find(node->buffer.get()); if (it != dbuffer_info_.end()) { const StorageEntry& e = it->second; TVM_FFI_ICHECK(in_double_buffer_scope_); @@ -221,7 +221,7 @@ class DoubleBufferInjector : public StmtExprMutator { Expr VisitExpr_(const BufferLoadNode* op) final { auto node = StmtExprMutator::VisitExpr_(op).as_or_throw(); - auto it = dbuffer_info_.find(node->buffer->data.get()); + auto it = dbuffer_info_.find(node->buffer.get()); if (it != dbuffer_info_.end()) { const StorageEntry& e = it->second; TVM_FFI_ICHECK(e.switch_read_var.defined()); @@ -237,7 +237,7 @@ class DoubleBufferInjector : public StmtExprMutator { return node; } - Buffer GetRemappedBuffer(Buffer buf, PrimExpr stride) { + BufferVar GetRemappedBuffer(BufferVar buf, PrimExpr stride) { auto key = buf.get(); auto it = buf_remap_.find(key); if (it != buf_remap_.end()) { @@ -254,7 +254,9 @@ class DoubleBufferInjector : public StmtExprMutator { // Stride gives the distance between the two halves of the // double-buffer, not the stride of the buffer's index. - buf.CopyOnWrite()->shape = {buf->shape[0] + stride}; + auto type = CopyBufferType(buf); + type->shape = {buf->shape[0] + stride}; + buf = RebuildBufferVar(buf, std::move(type)); buf_remap_[key] = buf; return buf; @@ -321,8 +323,8 @@ class DoubleBufferInjector : public StmtExprMutator { std::unordered_map> loop_pre_; // The allocation size of the buffer std::unordered_map dbuffer_info_; - // The updated Buffer objects - std::unordered_map buf_remap_; + // The updated BufferVar objects + std::unordered_map buf_remap_; // Pending double-buffer AllocBuffer nodes (deferred from flat AllocBuffer visit) std::unordered_map pending_dbuffer_allocs_; }; diff --git a/src/s_tir/transform/inject_permuted_layout.cc b/src/s_tir/transform/inject_permuted_layout.cc index 3b335472abb1..b5686e77b148 100644 --- a/src/s_tir/transform/inject_permuted_layout.cc +++ b/src/s_tir/transform/inject_permuted_layout.cc @@ -119,12 +119,12 @@ class PermutedLayoutInjector : private IRMutatorWithAnalyzer { } Stmt VisitStmt_(const SBlockNode* op) final { - // Record the mapping from buffer data var to buffer for later lookup + // Record the mapping from buffer identity to buffer for later lookup. for (auto buffer : op->alloc_buffers) { - buffer_map_.insert({buffer->data, buffer}); + buffer_map_.insert({buffer.var(), buffer}); } for (auto match_buffer : op->match_buffers) { - buffer_map_.insert({match_buffer->buffer->data, match_buffer->buffer}); + buffer_map_.insert({match_buffer->buffer.var(), match_buffer->buffer}); } if (op->annotations.count("permuted_layout") == 0 || @@ -145,9 +145,9 @@ class PermutedLayoutInjector : private IRMutatorWithAnalyzer { return block; } - int CheckAndGetBufferRowSize(Buffer buffer) { + int CheckAndGetBufferRowSize(BufferVar buffer) { TVM_FFI_ICHECK(buffer->shape.size() >= 2) - << "The dimension of Buffer \"" << buffer->name << "\" with shape " << buffer->shape + << "The dimension of BufferVar \"" << buffer.name() << "\" with shape " << buffer->shape << " should be at least 2"; auto dim = buffer->shape.size(); @@ -156,10 +156,10 @@ class PermutedLayoutInjector : private IRMutatorWithAnalyzer { if (buffer_row_size % 64 != 0) { TVM_FFI_ICHECK(buffer_row_size % 32 == 0) - << "Permuted SLayout for Buffer \"" << buffer->name << "\" with shape " << buffer->shape + << "Permuted SLayout for BufferVar \"" << buffer.name() << "\" with shape " << buffer->shape << " is not supported since its second dimension is not divisible by 32"; TVM_FFI_ICHECK(buffer_col_size % 2 == 0) - << "Permuted SLayout for Buffer \"" << buffer->name << "\" with shape " << buffer->shape + << "Permuted SLayout for BufferVar \"" << buffer.name() << "\" with shape " << buffer->shape << " is not supported since its first dimension is not divisible by 2 and second " "dimension is not divisible by 64"; } @@ -167,7 +167,7 @@ class PermutedLayoutInjector : private IRMutatorWithAnalyzer { return buffer_row_size; } - ffi::Array HandleBufferIndices(Buffer buffer, ffi::Array indices) { + ffi::Array HandleBufferIndices(BufferVar buffer, ffi::Array indices) { auto buffer_row_size = CheckAndGetBufferRowSize(buffer); // Mutate the last two indices @@ -190,7 +190,7 @@ class PermutedLayoutInjector : private IRMutatorWithAnalyzer { return store; } - auto scope = StorageScope::Create(GetPtrStorageScope(store->buffer->data)); + auto scope = StorageScope::Create(store->buffer.scope()); if (scope.rank != StorageRank::kShared) { return store; } @@ -208,7 +208,7 @@ class PermutedLayoutInjector : private IRMutatorWithAnalyzer { return load; } - auto scope = StorageScope::Create(GetPtrStorageScope(load->buffer->data)); + auto scope = StorageScope::Create(load->buffer.scope()); if (scope.rank != StorageRank::kShared) { return load; } @@ -273,7 +273,7 @@ class PermutedLayoutInjector : private IRMutatorWithAnalyzer { return call; } else if (call->op.same_as(mma_store_op)) { // TODO(yixin): mma_store is not fully tested yet - // because we will directly store result to Buffer instead of calling mma_store now + // because we will directly store result to BufferVar instead of calling mma_store now Expr access_ptr = call->args[2]; auto new_access_ptr = HandleAccessPtrAndOffset(access_ptr); auto new_call = call.CopyOnWrite(); @@ -287,8 +287,8 @@ class PermutedLayoutInjector : private IRMutatorWithAnalyzer { static constexpr size_t VECTORIZE_FACTOR = 8; static constexpr size_t BANK_SIZE_BYTES = 128; - // Mapping from data Var of a Buffer to Buffer, for lookup - std::unordered_map buffer_map_; + // Mapping from data Var of a BufferVar to BufferVar, for lookup + std::unordered_map buffer_map_; bool permute_ = false; }; diff --git a/src/s_tir/transform/inject_ptx_async_copy.cc b/src/s_tir/transform/inject_ptx_async_copy.cc index 5e4683ac951c..93c9261c532b 100644 --- a/src/s_tir/transform/inject_ptx_async_copy.cc +++ b/src/s_tir/transform/inject_ptx_async_copy.cc @@ -62,8 +62,8 @@ class PTXAsyncCopyInjector : public StmtMutator { const int bytes = indices_lanes * static_cast(load->buffer->dtype.StorageBytes()); if (bytes == 4 || bytes == 8 || bytes == 16) { - auto dst_elem_type = GetPointerType(store->buffer->data->ty); - auto src_elem_type = GetPointerType(load->buffer->data->ty); + auto dst_elem_type = GetPointerType(store->buffer->data_pointer_type); + auto src_elem_type = GetPointerType(load->buffer->data_pointer_type); TVM_FFI_ICHECK(dst_elem_type.has_value() && src_elem_type.has_value()) << "Both store and load buffer should have a pointer type annotation."; @@ -84,8 +84,8 @@ class PTXAsyncCopyInjector : public StmtMutator { if (indices_lanes == 1) { auto src_offset = load->indices[0]; auto dst_offset = store->indices[0]; - ffi::Array args = {store->buffer->data, mul(dst_offset, PrimExpr(index_factor)), - load->buffer->data, src_offset, PrimExpr(bytes)}; + ffi::Array args = {store->buffer.data(), mul(dst_offset, PrimExpr(index_factor)), + load->buffer.data(), src_offset, PrimExpr(bytes)}; // use arguments size to indicate whether or not to use predicated cp.async if (predicated) { args.push_back(predicate_value); @@ -122,8 +122,8 @@ class PTXAsyncCopyInjector : public StmtMutator { if (src_offset.defined() && dst_offset.defined()) { static const Op& ptx_cp_async_op = Op::Get("tirx.ptx.cp_async_raw"); return Evaluate(Call(store->buffer->dtype, ptx_cp_async_op, - {store->buffer->data, mul(dst_offset, PrimExpr(index_factor)), - load->buffer->data, src_offset, PrimExpr(bytes)}) + {store->buffer.data(), mul(dst_offset, PrimExpr(index_factor)), + load->buffer.data(), src_offset, PrimExpr(bytes)}) .as_or_throw()); } } else { @@ -153,8 +153,8 @@ class PTXAsyncCopyInjector : public StmtMutator { if (src_offset.defined() && dst_offset.defined()) { static const Op& ptx_cp_async_op = Op::Get("tirx.ptx.cp_async_raw"); return Evaluate(Call(store->buffer->dtype, ptx_cp_async_op, - {store->buffer->data, mul(dst_offset, PrimExpr(index_factor)), - load->buffer->data, src_offset, PrimExpr(bytes), predicate_value}) + {store->buffer.data(), mul(dst_offset, PrimExpr(index_factor)), + load->buffer.data(), src_offset, PrimExpr(bytes), predicate_value}) .as_or_throw()); } } diff --git a/src/s_tir/transform/inject_ptx_ldg32.cc b/src/s_tir/transform/inject_ptx_ldg32.cc index 54baa9de094f..0ee797a8313a 100644 --- a/src/s_tir/transform/inject_ptx_ldg32.cc +++ b/src/s_tir/transform/inject_ptx_ldg32.cc @@ -59,7 +59,7 @@ class PTXRewriter : public StmtMutator { Stmt VisitStmt_(const BufferStoreNode* store) final { Stmt result = StmtMutator::VisitStmt_(store); - Buffer load_buffer = store->buffer; + BufferVar load_buffer = store->buffer; PrimExpr load_value = store->value; // const BufferLoadNode* gload = load_value.as(); // take // the place of instance of @@ -99,7 +99,7 @@ class PTXRewriter : public StmtMutator { BufferStore value_store(store->buffer, imm_value, {new_indice}); static const Op& ptx_ldg32_op = Op::Get("tirx.ptx.ldg32"); Evaluate ptx_load(Call(store->buffer->dtype, ptx_ldg32_op, - {store->buffer->data, new_predicate, new_lhs, new_indice}) + {store->buffer.data(), new_predicate, new_lhs, new_indice}) .as_or_throw()); ffi::Array tmp_seq = {addr_store, local_addr_store, predicate_store, value_store, ptx_load}; @@ -122,7 +122,7 @@ class PTXRewriter : public StmtMutator { bool has_buffer_1 = false, has_buffer_2 = false; bool needs_buffer = false; - Buffer addr_buffer, predicate_buffer; + BufferVar addr_buffer, predicate_buffer; }; namespace transform { diff --git a/src/s_tir/transform/inject_software_pipeline.cc b/src/s_tir/transform/inject_software_pipeline.cc index 2c837e9ed322..598fb876e7fa 100644 --- a/src/s_tir/transform/inject_software_pipeline.cc +++ b/src/s_tir/transform/inject_software_pipeline.cc @@ -54,7 +54,7 @@ namespace software_pipeline { * \param buffer_data_to_buffer The map from buffer data to buffer. * \return The result block. */ -SBlock MakeSBlock(const Stmt& body, const ffi::Map& buffer_data_to_buffer) { +SBlock MakeSBlock(const Stmt& body, const ffi::Map& buffer_data_to_buffer) { if (const SBlockRealizeNode* block_realize = body.as()) { if (is_one(block_realize->predicate)) { // no need to create a new block @@ -96,8 +96,8 @@ class PipelineOpaqueAccessRewriter { * \param fragment_info Information about tensor core fragment */ PipelineOpaqueAccessRewriter( - const ffi::Map& buffer_data_to_buffer, - const ffi::Map& buffer_remap, const For& pipeline_loop, + const ffi::Map& buffer_data_to_buffer, + const ffi::Map& buffer_remap, const For& pipeline_loop, const std::unordered_map& fragment_info) : buffer_data_to_buffer_(buffer_data_to_buffer), buffer_remap_(buffer_remap), @@ -114,11 +114,11 @@ class PipelineOpaqueAccessRewriter { static const Op& ptx_ldmatrix_legacy = Op::Get("tirx.ptx.ldmatrix_legacy"); static const Op& ptx_mma_legacy = Op::Get("tirx.ptx.mma_legacy"); if (call->op.same_as(load_matrix_sync) || call->op.same_as(store_matrix_sync)) { - const Buffer& buffer = buffer_data_to_buffer_.at(call->args[0].as_or_throw()); + const BufferVar& buffer = buffer_data_to_buffer_.at(call->args[0].as_or_throw()); auto it = buffer_remap_.find(buffer); if (it != buffer_remap_.end()) { ffi::Array new_args = call->args; - const Buffer& new_buffer = (*it).second; + const BufferVar& new_buffer = (*it).second; new_args.Set( 4, RewriteWmmaFragmentIndex(buffer, new_buffer, call->args[4].as_or_throw())); return Call(call->ty, call->op, new_args, call->attrs, {}, call->span); @@ -128,7 +128,7 @@ class PipelineOpaqueAccessRewriter { for (int i = 0; i < 4; i++) { const Var& buffer_var = call->args[i * 2].as_or_throw(); PrimExpr index = call->args[i * 2 + 1].as_or_throw(); - const Buffer& buffer = buffer_data_to_buffer_.at(buffer_var); + const BufferVar& buffer = buffer_data_to_buffer_.at(buffer_var); auto it = buffer_remap_.find(buffer); if (it != buffer_remap_.end()) { PrimExpr new_index = RewriteWmmaFragmentIndex(buffer, (*it).second, index); @@ -147,14 +147,14 @@ class PipelineOpaqueAccessRewriter { } private: - int GetWmmaFragmentSize(const Buffer& buffer) { - auto it = fragment_info_.find(buffer->data.get()); + int GetWmmaFragmentSize(const BufferVar& buffer) { + auto it = fragment_info_.find(buffer.get()); TVM_FFI_ICHECK(it != fragment_info_.end()); const FragmentInfo& info = (*it).second; return info.GetSize(); } - PrimExpr RewriteWmmaFragmentIndex(const Buffer& old_buffer, const Buffer& new_buffer, + PrimExpr RewriteWmmaFragmentIndex(const BufferVar& old_buffer, const BufferVar& new_buffer, const PrimExpr& old_index) { PrimExpr new_buffer_offset = old_index; @@ -175,10 +175,10 @@ class PipelineOpaqueAccessRewriter { }; ffi::Array new_args = call->args; for (int i : arg_indices) { - const Buffer& buffer = buffer_data_to_buffer_.at(call->args[i].as_or_throw()); + const BufferVar& buffer = buffer_data_to_buffer_.at(call->args[i].as_or_throw()); auto it = buffer_remap_.find(buffer); if (it != buffer_remap_.end()) { - const Buffer& new_buffer = (*it).second; + const BufferVar& new_buffer = (*it).second; PrimExpr old_index = call->args[i + 1].as_or_throw(); PrimExpr offset; if (new_buffer->strides.empty()) { @@ -201,8 +201,8 @@ class PipelineOpaqueAccessRewriter { return Call(call->ty, call->op, new_args, call->attrs, {}, call->span); } - const ffi::Map& buffer_data_to_buffer_; - const ffi::Map& buffer_remap_; + const ffi::Map& buffer_data_to_buffer_; + const ffi::Map& buffer_remap_; const For& pipeline_loop_; const std::unordered_map& fragment_info_; }; @@ -224,8 +224,8 @@ class PipelineBodyRewriter : public StmtExprMutator { * of a two-stage software pipeline, only one version of these buffers are accessed. * \param fragment_info Information about tensor core fragment */ - PipelineBodyRewriter(const ffi::Map& buffer_data_to_buffer, - const ffi::Map& buffer_remap, For pipeline_loop, + PipelineBodyRewriter(const ffi::Map& buffer_data_to_buffer, + const ffi::Map& buffer_remap, For pipeline_loop, bool access_all_versions, const std::unordered_map& fragment_info) : buffer_data_to_buffer_(buffer_data_to_buffer), @@ -240,7 +240,7 @@ class PipelineBodyRewriter : public StmtExprMutator { auto it = buffer_remap_.find(buffer_region->buffer); if (it != buffer_remap_.end()) { Region new_region = buffer_region->region; - const Buffer& new_buffer = (*it).second; + const BufferVar& new_buffer = (*it).second; // For pipeline buffers, relax the access region of the first dimension to full extent // if access_all_versions == true Range accessed_version = @@ -256,8 +256,8 @@ class PipelineBodyRewriter : public StmtExprMutator { } Stmt VisitStmt_(const SBlockNode* op) final { - for (const Buffer& alloc_buffer : op->alloc_buffers) { - buffer_data_to_buffer_.Set(alloc_buffer->data, alloc_buffer); + for (const BufferVar& alloc_buffer : op->alloc_buffers) { + buffer_data_to_buffer_.Set(alloc_buffer.var(), alloc_buffer); } SBlock block = StmtExprMutator::VisitStmt_(op).as_or_throw(); SBlockNode* n = block.CopyOnWrite(); @@ -267,8 +267,8 @@ class PipelineBodyRewriter : public StmtExprMutator { n->writes.MutateByApply([this](const BufferRegion& buffer_region) { return RewritePipelineBufferRegion(buffer_region); }); - for (const Buffer& alloc_buffer : op->alloc_buffers) { - buffer_data_to_buffer_.erase(alloc_buffer->data); + for (const BufferVar& alloc_buffer : op->alloc_buffers) { + buffer_data_to_buffer_.erase(alloc_buffer.var()); } return block; } @@ -279,7 +279,7 @@ class PipelineBodyRewriter : public StmtExprMutator { if (it == buffer_remap_.end()) { return store; } - const Buffer& new_buffer = (*it).second; + const BufferVar& new_buffer = (*it).second; auto* n = store.CopyOnWrite(); n->buffer = new_buffer; PrimExpr version = @@ -294,7 +294,7 @@ class PipelineBodyRewriter : public StmtExprMutator { if (it == buffer_remap_.end()) { return load; } - const Buffer& new_buffer = (*it).second; + const BufferVar& new_buffer = (*it).second; auto* n = load.CopyOnWrite(); n->buffer = new_buffer; PrimExpr version = @@ -308,8 +308,8 @@ class PipelineBodyRewriter : public StmtExprMutator { return opaque_access_rewriter_.Rewrite(call); } - ffi::Map buffer_data_to_buffer_; - ffi::Map buffer_remap_; + ffi::Map buffer_data_to_buffer_; + ffi::Map buffer_remap_; For pipeline_loop_; bool access_all_versions_; PipelineOpaqueAccessRewriter opaque_access_rewriter_; @@ -321,9 +321,9 @@ class PipelineBodyRewriter : public StmtExprMutator { class PipelineRewriter : public StmtExprMutator { public: static Stmt Rewrite( - ffi::Map buffer_data_to_buffer, - const std::unordered_set& double_buffers, - const ffi::Array pipeline_allocs, const For& pipeline_loop, + ffi::Map buffer_data_to_buffer, + const std::unordered_set& double_buffers, + const ffi::Array pipeline_allocs, const For& pipeline_loop, const PipelineInfo& pipeline_info, const std::unordered_map& fragment_info, const ffi::Map preserved_annotations) { @@ -334,9 +334,9 @@ class PipelineRewriter : public StmtExprMutator { private: PipelineRewriter( - ffi::Map buffer_data_to_buffer, - const std::unordered_set& double_buffers, - const ffi::Array& pipeline_allocs, const For& pipeline_loop, + ffi::Map buffer_data_to_buffer, + const std::unordered_set& double_buffers, + const ffi::Array& pipeline_allocs, const For& pipeline_loop, const PipelineInfo& pipeline_info, const std::unordered_map& fragment_info, const ffi::Map preserved_annotations) @@ -352,9 +352,9 @@ class PipelineRewriter : public StmtExprMutator { Stmt BuildPipeline() { // Step 1: Analyze accesses to the buffers in the pipeline and compute the number of versions // need to maintain for each buffer. - std::unordered_map infos = + std::unordered_map infos = GetBufferAccessInfo(); - for (const Buffer& buffer : pipeline_allocs_) { + for (const BufferVar& buffer : pipeline_allocs_) { int num_versions = ComputeBufferVersions(buffer, infos.at(buffer)); if (num_versions > 1) { buffer_remap_.Set(buffer, RewriteAllocBuffer(buffer, num_versions)); @@ -392,10 +392,10 @@ class PipelineRewriter : public StmtExprMutator { SeqStmt stmt = SeqStmt({prologue, body, epilogue}); // Step 3: Make a new block that contains new buffer allocations after pipeline rewriting. - ffi::Array alloc_buffers; + ffi::Array alloc_buffers; for (const auto& alloc : pipeline_allocs_) { alloc_buffers.push_back(buffer_remap_.Get(alloc).value_or(alloc)); - buffer_data_to_buffer_.erase(alloc->data); + buffer_data_to_buffer_.erase(alloc.var()); } SBlock block = MakeSBlock(stmt, buffer_data_to_buffer_); block.CopyOnWrite()->alloc_buffers = std::move(alloc_buffers); @@ -409,9 +409,9 @@ class PipelineRewriter : public StmtExprMutator { * This method check the 'define' and 'use' stage of the buffers in the software pipeline, which * can be used to compute the number of versions needed to maintain after rewriting. */ - std::unordered_map + std::unordered_map GetBufferAccessInfo() { - std::unordered_map infos; + std::unordered_map infos; for (const auto& pair : pipeline_info_) { const SBlock& block = pair.first; int stage = pair.second.stage; @@ -474,7 +474,7 @@ class PipelineRewriter : public StmtExprMutator { * \param buffer_info The access information of the target buffer. * \return The number of versions required for the target buffer. */ - int ComputeBufferVersions(const Buffer& buffer, const BufferAccessInfo& buffer_info) { + int ComputeBufferVersions(const BufferVar& buffer, const BufferAccessInfo& buffer_info) { if (buffer_info.def == -1) { // Keep the original number of versions as buffers defined outside the software pipeline // should not be mutated. @@ -536,21 +536,21 @@ class PipelineRewriter : public StmtExprMutator { * \param num_versions The number of versions to keep. * \return The resized buffer. */ - Buffer RewriteAllocBuffer(const Buffer& buffer, int num_versions) { - ffi::ObjectPtr new_buffer = ffi::make_object(*(buffer.get())); + BufferVar RewriteAllocBuffer(const BufferVar& buffer, int num_versions) { + ffi::ObjectPtr new_buffer = CopyBufferType(buffer); new_buffer->shape.insert(new_buffer->shape.begin(), PrimExpr(num_versions)); if (new_buffer->strides.size()) { TVM_FFI_ICHECK(new_buffer->strides.size() + 1 == new_buffer->shape.size()); PrimExpr stride_0 = new_buffer->strides[0] * new_buffer->shape[1]; new_buffer->strides.insert(new_buffer->strides.begin(), stride_0); } - return Buffer(new_buffer); + return RebuildBufferVar(buffer, std::move(new_buffer)); } // Per-stage states that need to be tracked across pipeline prologue, body, and epilogue. struct AsyncStateGlobal { // Buffers that this stage asynchronously writes. - std::unordered_set dst_buffers; + std::unordered_set dst_buffers; // An imaginary index that the latest async operation associated with this stage has written // into. Only valid if all associated predicates are true, so that we can count the number of // async invocations exactly. When it is valid, it is the "sum of extents of loops that have @@ -558,7 +558,7 @@ class PipelineRewriter : public StmtExprMutator { // is only needed to compute wait count for epilogue without async producers. ffi::Optional producer_head{PrimExpr(-1)}; - bool writes(Buffer buf) const { return dst_buffers.count(buf.get()) > 0; } + bool writes(BufferVar buf) const { return dst_buffers.count(buf.get()) > 0; } }; // Per-stage states that are local to each of pipeline prologue, body, and epilogue. @@ -584,7 +584,7 @@ class PipelineRewriter : public StmtExprMutator { // until a point where we encounter a consumer of async result buffers. This is used to decide // if the producer_head of each buffer points to a copy written in the current or previous // iteration. - std::unordered_set seen; + std::unordered_set seen; // A symbolic expression representing the index the latest async operation associated with this // stage has written into, at the "current" iteration. @@ -612,7 +612,7 @@ class PipelineRewriter : public StmtExprMutator { // Determine where to insert async_wait and the corresponding wait count. void PopulateWaitCounts(const std::vector& new_blocks, arith::AnalyzerObj* ana_normalized, - const std::unordered_map& buffer_to_commit_group, + const std::unordered_map& buffer_to_commit_group, std::map* async_states_local) { for (size_t i = 0; i < new_blocks.size(); ++i) { if (new_blocks[i].is_async) { @@ -853,7 +853,7 @@ class PipelineRewriter : public StmtExprMutator { // Async related std::map async_states_local; - std::unordered_map buffer_to_commit_group; + std::unordered_map buffer_to_commit_group; for (const SBlock& block : ordered_stmts_) { int stage = pipeline_info_.at(block).stage; @@ -984,14 +984,14 @@ class PipelineRewriter : public StmtExprMutator { } arith::Analyzer analyzer_; - ffi::Map buffer_data_to_buffer_; - const std::unordered_set& double_buffers_; - ffi::Array pipeline_allocs_; + ffi::Map buffer_data_to_buffer_; + const std::unordered_set& double_buffers_; + ffi::Array pipeline_allocs_; For pipeline_loop_; PipelineInfo pipeline_info_; const std::unordered_map& fragment_info_; int max_stage_ = -1; - ffi::Map buffer_remap_; + ffi::Map buffer_remap_; ffi::Array ordered_stmts_; std::map async_states; ffi::Map preserved_annotations_; @@ -1014,7 +1014,7 @@ void BuildDependencyGraph(const ffi::Array& blocks, for (const SBlock& block : blocks) { for (const BufferRegion& read : block->reads) { - auto it = buffer_writers.find(read->buffer->data); + auto it = buffer_writers.find(read->buffer.var()); if (it != buffer_writers.end()) { for (const SBlock& writer : it->second) { if (dep_src2dst != nullptr) { @@ -1027,7 +1027,7 @@ void BuildDependencyGraph(const ffi::Array& blocks, } } for (const BufferRegion& write : block->writes) { - buffer_writers[write->buffer->data].push_back(block); + buffer_writers[write->buffer.var()].push_back(block); } } } @@ -1038,8 +1038,8 @@ class PipelineInjector : private StmtExprMutator { auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); PipelineInjector injector(global_symbol); for (const auto& kv : func->buffer_map) { - const Buffer& buffer = kv.second; - injector.buffer_data_to_buffer_.Set(buffer->data, buffer); + const BufferVar& buffer = kv.second; + injector.buffer_data_to_buffer_.Set(buffer.var(), buffer); } injector.fragment_info_ = GetTensorCoreFragmentInfo(func->body); return injector(func->body); @@ -1104,12 +1104,12 @@ class PipelineInjector : private StmtExprMutator { // the for-loop. If the for-loop has BlockRealize as its child, the pipeline body will be the // child of the block. Stmt pipeline_body{nullptr}; - ffi::Array pipeline_allocs; + ffi::Array pipeline_allocs; if (const auto* realize = for_node->body.as()) { const auto& block = realize->block; for (const auto& buffer : block->alloc_buffers) { - TVM_FFI_ICHECK(buffer->IsInstance()); - buffer_data_to_buffer_.Set(buffer->data, buffer); + TVM_FFI_ICHECK(buffer->IsInstance()); + buffer_data_to_buffer_.Set(buffer.var(), buffer); } pipeline_body = block->body; pipeline_allocs = block->alloc_buffers; @@ -1139,7 +1139,7 @@ class PipelineInjector : private StmtExprMutator { nested_pipeline_block->match_buffers.empty()); // match_buffer should have been lowered for (const auto& buffer : nested_pipeline_block->alloc_buffers) { pipeline_allocs.push_back(buffer); - buffer_data_to_buffer_.Set(buffer->data, buffer); + buffer_data_to_buffer_.Set(buffer.var(), buffer); } const auto* nested_seq = nested_pipeline_block->body.as(); for (size_t j = 0; j < nested_seq->seq.size(); j++) { @@ -1198,7 +1198,7 @@ class PipelineInjector : private StmtExprMutator { if (const auto* realize = op->body.as()) { const auto& block = realize->block; for (const auto& buffer : block->alloc_buffers) { - buffer_data_to_buffer_.erase(buffer->data); + buffer_data_to_buffer_.erase(buffer.var()); } } return pipeline; @@ -1209,8 +1209,8 @@ class PipelineInjector : private StmtExprMutator { * \param n The block pointer to which the buffer allocations are added. * \param alloc_buffers The buffer allocations to be added. */ - void AddAllocBuffers(SBlockNode* n, const ffi::Array alloc_buffers) { - for (const Buffer& alloc_buffer : alloc_buffers) { + void AddAllocBuffers(SBlockNode* n, const ffi::Array alloc_buffers) { + for (const BufferVar& alloc_buffer : alloc_buffers) { n->alloc_buffers.push_back(alloc_buffer); Region region; region.reserve(alloc_buffer->shape.size()); @@ -1223,7 +1223,7 @@ class PipelineInjector : private StmtExprMutator { Stmt VisitStmt_(const SBlockNode* op) final { for (const auto& buffer : op->alloc_buffers) { - buffer_data_to_buffer_.Set(buffer->data, buffer); + buffer_data_to_buffer_.Set(buffer.var(), buffer); } auto it = op->annotations.find(s_tir::attr::double_buffer_scope); @@ -1238,7 +1238,7 @@ class PipelineInjector : private StmtExprMutator { SBlock block = StmtExprMutator::VisitStmt_(op).as_or_throw(); for (const auto& buffer : op->alloc_buffers) { - buffer_data_to_buffer_.erase(buffer->data); + buffer_data_to_buffer_.erase(buffer.var()); } return block; } @@ -1260,9 +1260,9 @@ class PipelineInjector : private StmtExprMutator { return false; } - ffi::Map buffer_data_to_buffer_; + ffi::Map buffer_data_to_buffer_; std::unordered_map fragment_info_; - std::unordered_set double_buffers; + std::unordered_set double_buffers; ffi::Optional global_symbol_; }; diff --git a/src/s_tir/transform/inject_virtual_thread.cc b/src/s_tir/transform/inject_virtual_thread.cc index e54270702809..f6e4c1a7add9 100644 --- a/src/s_tir/transform/inject_virtual_thread.cc +++ b/src/s_tir/transform/inject_virtual_thread.cc @@ -55,7 +55,7 @@ class ExprTouched final : public StmtExprVisitor { StmtExprVisitor::VisitStmt(n); } void VisitExpr_(const BufferLoadNode* op) final { - HandleUseVar(op->buffer->data.get()); + HandleUseVar(op->buffer.get()); StmtExprVisitor::VisitExpr_(op); } void VisitExpr_(const VarNode* op) final { HandleUseVar(op); } @@ -119,7 +119,7 @@ class VarTouchedAnalysis : public StmtVisitor { for (const auto& index : op->indices) { tc(index); } - Record(op->buffer->data.get(), tc); + Record(op->buffer.get(), tc); } void VisitStmt_(const ForNode* op) final { ExprTouched tc(touched_var_, false); @@ -141,7 +141,7 @@ class VarTouchedAnalysis : public StmtVisitor { for (size_t i = 0; i < op->buffer->shape.size(); ++i) { tc(op->buffer->shape[i]); } - Record(op->buffer->data.get(), tc); + Record(op->buffer.get(), tc); StmtVisitor::VisitStmt_(op); } void Record(const VarNode* var, const ExprTouched& tc) { @@ -212,7 +212,7 @@ class VTInjector : public arith::IRMutatorWithAnalyzer { } // Variable Expr VisitExpr_(const VarNode* op) final { - TVM_FFI_ICHECK(!alloc_remap_.count(op)) << "Buffer address may get rewritten in virtual thread"; + TVM_FFI_ICHECK(!alloc_remap_.count(op)) << "BufferVar address may get rewritten in virtual thread"; if (touched_var_.count(op)) { visit_touched_var_ = true; } @@ -263,11 +263,11 @@ class VTInjector : public arith::IRMutatorWithAnalyzer { template Node VisitBufferAccess(Node node) { - if (touched_var_.count(node->buffer->data.get())) { + if (touched_var_.count(node->buffer.get())) { visit_touched_var_ = true; } - auto it = alloc_remap_.find(node->buffer->data.get()); + auto it = alloc_remap_.find(node->buffer.get()); if (it != alloc_remap_.end()) { TVM_FFI_ICHECK_EQ(node->indices.size(), 1) << "InjectVirtualThread expects rewritten allocations to be flat memory."; @@ -279,7 +279,7 @@ class VTInjector : public arith::IRMutatorWithAnalyzer { return node; } - Buffer GetRemappedBuffer(Buffer buf, PrimExpr alloc_extent) { + BufferVar GetRemappedBuffer(BufferVar buf, PrimExpr alloc_extent) { auto key = buf.get(); auto it = buf_remap_.find(key); if (it != buf_remap_.end()) { @@ -288,8 +288,9 @@ class VTInjector : public arith::IRMutatorWithAnalyzer { TVM_FFI_ICHECK_EQ(buf->shape.size(), 1) << "Expected buffers being rewritten to already be flattened."; - auto writer = buf.CopyOnWrite(); + auto writer = CopyBufferType(buf); writer->shape = {buf->shape[0] * alloc_extent}; + buf = RebuildBufferVar(buf, std::move(writer)); buf_remap_[key] = buf; return buf; @@ -440,19 +441,20 @@ class VTInjector : public arith::IRMutatorWithAnalyzer { visit_touched_var_ = false; - if (touched_var_.count(op->buffer->data.get()) || !allow_share_) { + if (touched_var_.count(op->buffer.get()) || !allow_share_) { TVM_FFI_ICHECK_EQ(shape.size(), 1) << "InjectVirtualThread expects rewritten allocations to be flat memory."; PrimExpr stride = shape[0]; shape = {stride * num_threads_}; - alloc_remap_[op->buffer->data.get()] = stride; + alloc_remap_[op->buffer.get()] = stride; } if (shape.same_as(op->buffer->shape)) { return ffi::GetRef(op); } else { - Buffer new_buffer = op->buffer; - new_buffer.CopyOnWrite()->shape = shape; + auto type = CopyBufferType(op->buffer); + type->shape = shape; + BufferVar new_buffer = RebuildBufferVar(op->buffer, std::move(type)); return AllocBuffer(new_buffer, op->annotations); } } @@ -520,7 +522,7 @@ class VTInjector : public arith::IRMutatorWithAnalyzer { * the allocated buffer size, then modifying the indices at which * each virtual thread accesses the buffer. */ - std::unordered_map buf_remap_; + std::unordered_map buf_remap_; }; class VirtualThreadInjector : public arith::IRMutatorWithAnalyzer { diff --git a/src/s_tir/transform/lower_async_dma.cc b/src/s_tir/transform/lower_async_dma.cc index ccde9f2e8dce..b4ecd26b3ac1 100644 --- a/src/s_tir/transform/lower_async_dma.cc +++ b/src/s_tir/transform/lower_async_dma.cc @@ -80,9 +80,9 @@ class AsyncDMALowerer : public arith::IRMutatorWithAnalyzer { return Evaluate( Call(PrimType::Int(32), builtin::dma_copy(), ffi::Array{PrimExpr(async_queue_id_.value()), - Call(mem_copy->dest->buffer->data->ty, builtin::address_of(), + Call(mem_copy->dest->buffer->data_pointer_type, builtin::address_of(), ffi::Array{dst}, Attrs(), {}, Span()), - Call(mem_copy->source->buffer->data->ty, builtin::address_of(), + Call(mem_copy->source->buffer->data_pointer_type, builtin::address_of(), ffi::Array{src}, Attrs(), {}, Span()), dst_nbytes, PrimExpr(dma_bypass_cache_)}, Attrs(), {}, Span()) diff --git a/src/s_tir/transform/lower_cross_thread_reduction.cc b/src/s_tir/transform/lower_cross_thread_reduction.cc index 994f461c2007..5564a1eb2132 100644 --- a/src/s_tir/transform/lower_cross_thread_reduction.cc +++ b/src/s_tir/transform/lower_cross_thread_reduction.cc @@ -77,7 +77,7 @@ bool IsBoundToThreadIdx(const ForNode* loop) { */ bool IsDominantBlock(const SBlock& scope_block, const SBlock& block) { // Step 1. Count the number of writers for each buffer written by the scope block. - std::unordered_map buffer_writer_cnt; + std::unordered_map buffer_writer_cnt; PreOrderVisit(scope_block->body, [&buffer_writer_cnt](const ffi::ObjectRef& obj) { if (const auto* block = obj.as()) { for (const BufferRegion& buffer_region : block->writes) { @@ -140,14 +140,14 @@ bool IsReductionBlock(const SBlockRealize& realize, const ffi::Map& * computation results or not, which is used for determine the buffer name prefix * \return The created buffers */ -ffi::Array MakeScratchpads(const ffi::Array& reduction_buffers, +ffi::Array MakeScratchpads(const ffi::Array& reduction_buffers, bool is_cross_thread_buffer) { - ffi::Array new_buffers; + ffi::Array new_buffers; new_buffers.reserve(reduction_buffers.size()); - for (const Buffer& buffer : reduction_buffers) { + for (const BufferVar& buffer : reduction_buffers) { ffi::String name = is_cross_thread_buffer ? "cross" : "in"; - name = name + "_thread_" + buffer->name; - new_buffers.push_back(Buffer(/*ptr=*/Var(name, PointerType(buffer->dtype, "local")), + name = name + "_thread_" + buffer.name(); + new_buffers.push_back(BufferVar(/*ptr=*/Var(name, PointerType(buffer->dtype, "local")), /*dtype=*/buffer->dtype, /*shape=*/{IntImm::Int32(1)}, /*strides=*/{IntImm::Int32(1)}, @@ -165,8 +165,8 @@ ffi::Array MakeScratchpads(const ffi::Array& reduction_buffers, */ class BufferReplacer : private StmtExprMutator { public: - static Stmt Run(ffi::Array src_buffers, ffi::Array tgt_buffers, Stmt stmt) { - ffi::Map buffer_map; + static Stmt Run(ffi::Array src_buffers, ffi::Array tgt_buffers, Stmt stmt) { + ffi::Map buffer_map; TVM_FFI_ICHECK_EQ(src_buffers.size(), tgt_buffers.size()); int n_buffers = src_buffers.size(); for (int i = 0; i < n_buffers; ++i) { @@ -176,7 +176,7 @@ class BufferReplacer : private StmtExprMutator { } private: - explicit BufferReplacer(ffi::Map buffer_map) + explicit BufferReplacer(ffi::Map buffer_map) : buffer_map_(std::move(buffer_map)) {} Expr VisitExpr_(const BufferLoadNode* load) final { @@ -194,7 +194,7 @@ class BufferReplacer : private StmtExprMutator { } } - ffi::Map buffer_map_; + ffi::Map buffer_map_; }; /*! @@ -298,9 +298,9 @@ class InThreadReducerMaker : private StmtMutator { * \param reduction_loops The reduction loops */ Stmt TransformReductionBlock(const SBlockRealizeNode* realize, // - const ffi::Optional>& it_buffers, // - const ffi::Array& ct_buffers, // - const ffi::Array& wb_buffers, // + const ffi::Optional>& it_buffers, // + const ffi::Array& ct_buffers, // + const ffi::Array& wb_buffers, // const ffi::Array& old_wb_indices, // const CommReducer& reducer, // const ffi::Array& combiner_rhs, // @@ -308,10 +308,10 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize, int n_buffers = wb_buffers.size(); const SBlockNode* block = realize->block.get(); - auto f_create_buffer_regions = [](ffi::Array buffers) { + auto f_create_buffer_regions = [](ffi::Array buffers) { ffi::Array regions; regions.reserve(buffers.size()); - for (const Buffer& buffer : buffers) { + for (const BufferVar& buffer : buffers) { regions.push_back(BufferRegion(buffer, {Range::FromMinExtent(0, 1)})); } return regions; @@ -635,7 +635,7 @@ class CrossThreadReductionTransformer : public StmtMutator { * - the RHS values of the reduction updates, * - the indices which is used to access the reduction buffers when storing the reduction results */ - std::tuple, ffi::Array, ffi::Array> + std::tuple, ffi::Array, ffi::Array> CheckCanApplyCrossThreadReduction(const SBlockNode* block, const std::vector& reduction_loops) const { // Condition 1. All the reduction-related loops should be the deepest among all statements @@ -688,7 +688,7 @@ class CrossThreadReductionTransformer : public StmtMutator { // Condition 4. All reduction buffers should be all local or all non-local. int is_local_buf = -1; - ffi::Array reduction_buffers; + ffi::Array reduction_buffers; reduction_buffers.reserve(updates.size()); for (const BufferStore& buf_store : updates) { reduction_buffers.push_back(buf_store->buffer); @@ -792,7 +792,7 @@ class CrossThreadReductionTransformer : public StmtMutator { auto it = block2new_buffers_.find(block); if (it != block2new_buffers_.end()) { SBlockNode* p_new_block = new_block.CopyOnWrite(); - for (const Buffer& new_buffer : it->second) { + for (const BufferVar& new_buffer : it->second) { if (new_buffer.defined()) { p_new_block->alloc_buffers.push_back(new_buffer); } @@ -809,7 +809,7 @@ class CrossThreadReductionTransformer : public StmtMutator { // which condition the block violates. int n_bound_reduction_loops = 0; CommReducer reducer{nullptr}; - ffi::Array reduction_buffers{nullptr}; + ffi::Array reduction_buffers{nullptr}; ffi::Array combiner_rhs{nullptr}; ffi::Array wb_indices{nullptr}; std::tie(n_bound_reduction_loops, reducer, reduction_buffers, combiner_rhs, wb_indices) = @@ -822,11 +822,11 @@ class CrossThreadReductionTransformer : public StmtMutator { !is_one(realize->predicate); // Step 3. Create intermediate buffers, storing them in `ct_buffers` and // `it_buffers`. Let the scope block allocate these new buffers. - ffi::Array& new_buffers = block2new_buffers_[block_stack_.back()]; - ffi::Array ct_buffers = + ffi::Array& new_buffers = block2new_buffers_[block_stack_.back()]; + ffi::Array ct_buffers = MakeScratchpads(reduction_buffers, /*is_cross_thread_buffer=*/true); new_buffers.insert(new_buffers.end(), ct_buffers.begin(), ct_buffers.end()); - ffi::Optional> it_buffers = std::nullopt; + ffi::Optional> it_buffers = std::nullopt; if (need_in_thread_reduction) { it_buffers = MakeScratchpads(reduction_buffers, /*is_cross_thread_buffer=*/false); new_buffers.insert(new_buffers.end(), it_buffers.value().begin(), it_buffers.value().end()); @@ -847,7 +847,7 @@ class CrossThreadReductionTransformer : public StmtMutator { Range::FromMinExtent(loop->min, loop->extent)); } } - for (const Buffer& reduction_buf : reduction_buffers) { + for (const BufferVar& reduction_buf : reduction_buffers) { crt_buf2threads_[reduction_buf.get()] = reduction_threads; } } @@ -921,14 +921,14 @@ class CrossThreadReductionTransformer : public StmtMutator { std::vector statement_stack_; std::vector loop_stack_; std::vector block_stack_; - std::unordered_map> block2new_buffers_; + std::unordered_map> block2new_buffers_; std::unordered_map loop2new_stmt_; ffi::Map loop_range_map_; arith::Analyzer analyzer_; int block_idx_depth = 0; int thread_idx_depth = 0; - std::unordered_map>> + std::unordered_map>> crt_buf2threads_; }; diff --git a/src/s_tir/transform/lower_match_buffer.cc b/src/s_tir/transform/lower_match_buffer.cc index 2fccadb13ffc..7163a5f7f2bb 100644 --- a/src/s_tir/transform/lower_match_buffer.cc +++ b/src/s_tir/transform/lower_match_buffer.cc @@ -59,15 +59,15 @@ class MatchBufferLower : public StmtExprMutator { // (e.g., remapping elem_offset). We need match_buffers_ to also // contain the remapped buffer keys for lookups in the body. // Snapshot current match_buffers_ keys before base visit. - std::vector orig_buffers; + std::vector orig_buffers; for (const auto& kv : match_buffers_) { orig_buffers.push_back(kv.first); } Stmt stmt = StmtExprMutator ::VisitStmt_(op); // Add remapped buffer keys to match_buffers_ - for (const Buffer& orig_buf : orig_buffers) { + for (const BufferVar& orig_buf : orig_buffers) { if (auto remap_it = buffer_remap_.find(orig_buf); remap_it != buffer_remap_.end()) { - Buffer remapped = (*remap_it).second; + BufferVar remapped = (*remap_it).second; if (!match_buffers_.count(remapped)) { match_buffers_.Set(remapped, match_buffers_[orig_buf]); } @@ -108,7 +108,7 @@ class MatchBufferLower : public StmtExprMutator { Stmt VisitStmt_(const BufferStoreNode* op) final { // Save the original buffer before base class mutation may remap it - Buffer orig_buffer = op->buffer; + BufferVar orig_buffer = op->buffer; Stmt stmt = StmtExprMutator::VisitStmt_(op); op = stmt.as(); TVM_FFI_ICHECK(op != nullptr); @@ -118,7 +118,7 @@ class MatchBufferLower : public StmtExprMutator { if (it == match_buffers_.end()) { return stmt; } else { - const Buffer& buffer = (*it).first; + const BufferVar& buffer = (*it).first; const BufferRegion& source = (*it).second; auto n = CopyOnWrite(op); @@ -132,7 +132,7 @@ class MatchBufferLower : public StmtExprMutator { Expr VisitExpr_(const BufferLoadNode* op) final { // Save the original buffer before base class mutation may remap it - Buffer orig_buffer = op->buffer; + BufferVar orig_buffer = op->buffer; PrimExpr expr = StmtExprMutator::VisitExpr_(op).as_or_throw(); op = expr.as(); TVM_FFI_ICHECK(op != nullptr); @@ -141,7 +141,7 @@ class MatchBufferLower : public StmtExprMutator { if (it == match_buffers_.end()) { return expr; } else { - const Buffer& buffer = (*it).first; + const BufferVar& buffer = (*it).first; const BufferRegion& source = (*it).second; ffi::Array indices = ConvertIndices(MatchBufferRegion(buffer, source), op->indices); TVM_FFI_ICHECK(!op->predicate.has_value()) @@ -151,7 +151,7 @@ class MatchBufferLower : public StmtExprMutator { } BufferRegion VisitBufferRegion(const BufferRegion& buffer_region) { - const Buffer& buffer = buffer_region->buffer; + const BufferVar& buffer = buffer_region->buffer; auto it = match_buffers_.find(buffer); if (it == match_buffers_.end()) { return buffer_region; @@ -165,9 +165,9 @@ class MatchBufferLower : public StmtExprMutator { private: void CheckAndUpdateVarMap(const MatchBufferRegion& match_buffer) { // Step.1. Check - const Buffer& buffer = match_buffer->buffer; + const BufferVar& buffer = match_buffer->buffer; const BufferRegion& source = VisitBufferRegion(match_buffer->source); - const Buffer& source_buffer = source->buffer; + const BufferVar& source_buffer = source->buffer; // Step.1.1. Check scope & dtype TVM_FFI_ICHECK_EQ(buffer.scope(), source_buffer.scope()) @@ -185,17 +185,14 @@ class MatchBufferLower : public StmtExprMutator { } if (is_zero(buffer->elem_offset)) { TVM_FFI_ICHECK(is_zero(source_buffer->elem_offset)) - << "Trying to bind a Buffer with offset into one without offset " + << "Trying to bind a BufferVar with offset into one without offset " << " required elem_offset=" << buffer->elem_offset << ", provided elem_offset=" << source_buffer->elem_offset; } // Step.2. Update match_buffers_.Set(buffer, source); - // Step.2.1. Update buffer data - Bind(buffer->data, source_buffer->data, buffer->name + ".data"); - - // Step.2.2. Update element offset + // Step.2.1. Update element offset // We use the ElemOffset method to avoid duplicating the index calculation. { ffi::Array indices; @@ -206,7 +203,7 @@ class MatchBufferLower : public StmtExprMutator { ffi::Array buffer_start_indices = source_buffer->ElemOffset(indices); if (buffer_start_indices.size() == 1) { - Bind(buffer->elem_offset, buffer_start_indices[0], buffer->name + ".elem_offset"); + Bind(buffer->elem_offset, buffer_start_indices[0], buffer.name() + ".elem_offset"); TVM_FFI_ICHECK( analyzer_->CanProve(truncmod(buffer->elem_offset, buffer->offset_factor) == 0)) << "The source elem_offset " << buffer_start_indices[0] @@ -229,14 +226,14 @@ class MatchBufferLower : public StmtExprMutator { PrimExpr stride = MakeConst(buffer->strides.back().ty(), 1); for (size_t i = buffer->shape.size(); i > 0; --i) { const PrimExpr& shape = source_buffer->shape[i - 1 + offset]; - Bind(buffer->strides[i - 1], stride, buffer->name + ".strides_" + std::to_string(i - 1)); + Bind(buffer->strides[i - 1], stride, buffer.name() + ".strides_" + std::to_string(i - 1)); stride *= shape; } } else { TVM_FFI_ICHECK_EQ(buffer->shape.size() + offset, source_buffer->strides.size()); for (size_t i = buffer->shape.size(); i > 0; --i) { const PrimExpr& stride = source_buffer->strides[i - 1 + offset]; - Bind(buffer->strides[i - 1], stride, buffer->name + ".strides_" + std::to_string(i - 1)); + Bind(buffer->strides[i - 1], stride, buffer.name() + ".strides_" + std::to_string(i - 1)); } } } @@ -244,7 +241,7 @@ class MatchBufferLower : public StmtExprMutator { // Step 2.4. Check and update shape for (size_t i = 0; i < buffer->shape.size(); ++i) { const Range& range = source->region[i + offset]; - Bind(buffer->shape[i], range->extent, buffer->name + ".shape_" + std::to_string(i)); + Bind(buffer->shape[i], range->extent, buffer.name() + ".shape_" + std::to_string(i)); } } @@ -295,8 +292,8 @@ class MatchBufferLower : public StmtExprMutator { } private: - /*! \brief Buffer region mapping. */ - ffi::Map match_buffers_; + /*! \brief BufferVar region mapping. */ + ffi::Map match_buffers_; /*! \brief Var mapping for buffer signature (data, strides, element_offset, etc.) */ ffi::Map var_map_; /*! \brief The analyzer */ diff --git a/src/s_tir/transform/lower_opaque_block.cc b/src/s_tir/transform/lower_opaque_block.cc index 88b366c9efb6..a9d5cf34f78d 100644 --- a/src/s_tir/transform/lower_opaque_block.cc +++ b/src/s_tir/transform/lower_opaque_block.cc @@ -60,9 +60,9 @@ class OpaqueBlockLower : public StmtExprMutator { } // Step 3. Handle allocations in reverse order for (size_t i = new_block->alloc_buffers.size(); i > 0; --i) { - const Buffer& buffer = new_block->alloc_buffers[i - 1]; + const BufferVar& buffer = new_block->alloc_buffers[i - 1]; ffi::Map allocate_annotations; - auto it = storage_align_.find(buffer->data); + auto it = storage_align_.find(buffer.var()); if (it != storage_align_.end()) { StorageAlignAnnotation allocate_aligns; for (auto tuple : it->second) { diff --git a/src/s_tir/transform/lower_thread_allreduce.cc b/src/s_tir/transform/lower_thread_allreduce.cc index 898b9cf57396..06dde28c34af 100644 --- a/src/s_tir/transform/lower_thread_allreduce.cc +++ b/src/s_tir/transform/lower_thread_allreduce.cc @@ -80,7 +80,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { // (the remap is set up by MakeAllreduce which runs during AttrStmt/Evaluate visit // that appears later in the sequence). We record the original data pointer and // attempt the remap; if it's not ready, the post-processing pass will handle it. - const VarNode* orig_data_ptr = op->buffer->data.get(); + const VarNode* orig_data_ptr = op->buffer.get(); auto node = StmtExprMutator::VisitStmt_(op).as_or_throw(); if (auto it = alloc_remap_.find(orig_data_ptr); it != alloc_remap_.end()) { @@ -97,7 +97,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { * \param replacement The replacement buffer. * \return The remapped statement(s). */ - Stmt RemapAllocBuffer(AllocBuffer node, const Buffer& replacement) { + Stmt RemapAllocBuffer(AllocBuffer node, const BufferVar& replacement) { auto* cow = node.CopyOnWrite(); cow->buffer = replacement; if (replacement.scope() == "shared") { @@ -108,14 +108,13 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { return node; } - ffi::Optional GetRemappedBuffer(const Buffer& buf) { + ffi::Optional GetRemappedBuffer(const BufferVar& buf) { if (auto it = buf_remap_.find(buf.get()); it != buf_remap_.end()) { return it->second; } - if (auto it = var_remap_.find(buf->data.get()); it != var_remap_.end()) { - Buffer new_buf = buf; - new_buf.CopyOnWrite()->data = it->second; + if (auto it = var_remap_.find(buf.get()); it != var_remap_.end()) { + BufferVar new_buf(it->second); buf_remap_[buf.get()] = new_buf; return new_buf; } @@ -132,7 +131,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { } Expr VisitExpr_(const BufferLoadNode* op) final { - if (auto it = load_remap_.find(op->buffer->data.get()); it != load_remap_.end()) { + if (auto it = load_remap_.find(op->buffer.get()); it != load_remap_.end()) { for (const auto& index : op->indices) { TVM_FFI_ICHECK(is_zero(index)); } @@ -190,7 +189,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { } dtypes.push_back(values[idx].ty()); } - std::vector buffers(size); + std::vector buffers(size); for (size_t idx = 0; idx < size; ++idx) { PrimExpr arg = call->args[2 + size + idx].as_or_throw(); // Loads from boolean buffers may have cast nodes inserted by @@ -277,7 +276,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { } std::vector seq; - std::vector new_alloc_bufs; + std::vector new_alloc_bufs; // // This is an optimization. For small reduction sizes, it may be beneficial // for a single warp to performance the entire reduction. No trips to shared @@ -322,7 +321,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { // This avoids to emit predicated stores, as all threads are // uniformly writing the same result. for (size_t i = 0; i < size; ++i) { - Buffer buf = reduce_results[i].as_or_throw()->buffer; + BufferVar buf = reduce_results[i].as_or_throw()->buffer; PrimExpr val = BufferLoad(buf, {zero_index}); TVM_FFI_ICHECK_EQ(val.ty(), dtypes[i]); PrimExpr splat = WarpShuffle(builtin::tvm_warp_shuffle(), new_alloc_bufs.back(), val, @@ -331,13 +330,13 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { } } else { int n_warps = reduce_extent / warp_size_; - std::vector local_bufs; + std::vector local_bufs; // 1. Create the staging buffer in shared memory. - std::vector staging_shared_bufs; + std::vector staging_shared_bufs; staging_shared_bufs.reserve(size); for (size_t i = 0; i < size; ++i) { - Buffer staging_shared_buf = decl_buffer( + BufferVar staging_shared_buf = decl_buffer( /*shape=*/{IntImm(reduce_index.ty(), n_warps * group_extent)}, /*dtype=*/buffers[i]->dtype, /*name=*/"red_buf_staging", /*storage_scope=*/"shared"); staging_shared_bufs.push_back(staging_shared_buf); @@ -381,7 +380,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { write_result.reserve(size); for (size_t i = 0; i < size; ++i) { new_alloc_bufs.push_back(reduce_results[i].as_or_throw()->buffer); - Buffer broadcast_shared_buf = decl_buffer( + BufferVar broadcast_shared_buf = decl_buffer( /*shape=*/{IntImm(reduce_index.ty(), group_extent)}, /*dtype=*/buffers[i]->dtype, /*name=*/"red_result", /*storage_scope=*/"shared"); write_result.push_back( @@ -395,19 +394,19 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { // Write back allreduce results and update existing allocations. for (size_t i = 0; i < size; ++i) { - TVM_FFI_ICHECK(!load_remap_.count(buffers[i]->data.get())); - Buffer buf = reduce_results[i].as_or_throw()->buffer; + TVM_FFI_ICHECK(!load_remap_.count(buffers[i].get())); + BufferVar buf = reduce_results[i].as_or_throw()->buffer; TVM_FFI_ICHECK_EQ(reduce_results[i].ty(), dtypes[i]); - load_remap_[buffers[i]->data.get()] = reduce_results[i]; + load_remap_[buffers[i].get()] = reduce_results[i]; // The AllocBuffer doesn't need to be emitted here since alloc_remap_ // will cause the existing allocation to be rewritten in VisitStmt_(AllocBufferNode*). - alloc_remap_[buffers[i]->data.get()] = buf; - var_remap_[buffers[i]->data.get()] = buf->data; + alloc_remap_[buffers[i].get()] = buf; + var_remap_[buffers[i].get()] = buf.var(); buf_remap_[buffers[i].get()] = buf; } } else { - std::vector shared_bufs(size); + std::vector shared_bufs(size); if (reduce_extent == 1) { // special case, no reduction is needed. std::vector stores; @@ -429,21 +428,21 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { seq.emplace_back(MakeBufAllreduce(combiner, dtypes, shared_bufs, reduce_index, group_index, reduce_extent, group_extent, contiguous_reduce_extent)); for (size_t idx = 0; idx < size; ++idx) { - TVM_FFI_ICHECK(!load_remap_.count(buffers[idx]->data.get())); + TVM_FFI_ICHECK(!load_remap_.count(buffers[idx].get())); PrimExpr pred = MakeConst(PrimType::Bool(static_cast(dtypes[idx].lanes())), true); BufferLoad load(shared_bufs[idx], {BufIndex(IntImm(reduce_index.ty(), 0), group_index, reduce_extent)}); TVM_FFI_ICHECK_EQ(load.ty(), dtypes[idx]); - load_remap_[buffers[idx]->data.get()] = load; - alloc_remap_[buffers[idx]->data.get()] = shared_bufs[idx]; - var_remap_[buffers[idx]->data.get()] = shared_bufs[idx]->data; + load_remap_[buffers[idx].get()] = load; + alloc_remap_[buffers[idx].get()] = shared_bufs[idx]; + var_remap_[buffers[idx].get()] = shared_bufs[idx].var(); buf_remap_[buffers[idx].get()] = shared_bufs[idx]; } } // Fix all local allocations as all statements are built. ffi::Array alloc_stmts; - for (Buffer buf : new_alloc_bufs) { + for (BufferVar buf : new_alloc_bufs) { alloc_stmts.push_back(AllocBuffer(buf)); } // Prepend allocations before the sequence @@ -455,7 +454,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { return body; } - std::pair, std::vector> MakeWarpAllreduce( + std::pair, std::vector> MakeWarpAllreduce( std::vector src_values, // std::vector dtypes, // const CommReducerNode* combiner, // @@ -465,8 +464,8 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { std::vector* seq) { int n_buffers = src_values.size(); - std::vector shared_bufs; - std::vector local_bufs; + std::vector shared_bufs; + std::vector local_bufs; shared_bufs.reserve(n_buffers); // This is the index to the reduction variable, one reduction @@ -496,7 +495,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { // The mask for this reducer, as this reducer may sit inside // a divergent control flow. Here it uses a variable to cache the current // active channels. - ffi::Optional mask_buffer; + ffi::Optional mask_buffer; if (need_warp_shuffle_mask_) { mask_buffer = decl_buffer(shape, mask.ty(), "mask", "local"); seq->emplace_back(BufferStore(mask_buffer.value(), mask, zero_indices)); @@ -514,7 +513,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { // Load reduction values, no synchronization needed. ffi::Array a, b; for (int i = 0; i < n_buffers; ++i) { - Buffer shared_buf = shared_bufs[i]; + BufferVar shared_buf = shared_bufs[i]; BufferLoad val(shared_buf, zero_indices); TVM_FFI_ICHECK_EQ(val.ty(), dtypes[i]); a.push_back(val); @@ -532,7 +531,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { // The former may cause dead lock as there is a divergent // branch with a warp sync call inside. PrimExpr other = WarpShuffle(builtin::tvm_warp_shuffle_down(), mask_buffer, val, offset); - Buffer local_buf = local_bufs[i]; + BufferVar local_buf = local_bufs[i]; Stmt s = BufferStore(local_buf, other, zero_indices); seq->push_back(s); @@ -548,7 +547,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { std::vector stores; stores.reserve(n_buffers); for (int i = 0; i < n_buffers; ++i) { - Buffer buf = shared_bufs[i]; + BufferVar buf = shared_bufs[i]; stores.push_back(BufferStore(buf, ret[i], zero_indices)); } @@ -577,7 +576,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { // make allreduce. Stmt MakeBufAllreduce(const CommReducerNode* combiner, const std::vector& dtypes, - const ffi::Array& shared_bufs, PrimExpr reduce_index, + const ffi::Array& shared_bufs, PrimExpr reduce_index, PrimExpr group_index, int reduce_extent, int group_extent, int contiguous_reduce_extent) { // Get next power of two @@ -725,7 +724,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { } // Emit warp shuffle calls. - PrimExpr WarpShuffle(const Op& op, ffi::Optional mask_buffer, PrimExpr val, + PrimExpr WarpShuffle(const Op& op, ffi::Optional mask_buffer, PrimExpr val, PrimExpr delta_or_lane) { ffi::Array indices = {0}; PrimExpr mask; @@ -823,11 +822,11 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { public: // These members are public for post-processing by DeferredRemapper. // Allocate remap - std::unordered_map alloc_remap_; + std::unordered_map alloc_remap_; // BufferVar remap std::unordered_map var_remap_; - // Buffer remap - std::unordered_map buf_remap_; + // BufferVar remap + std::unordered_map buf_remap_; // Pending AllocBuffer original data pointers (for flat IR deferred remapping) std::vector pending_alloc_buffers_; }; @@ -844,9 +843,9 @@ namespace transform { */ class DeferredRemapper : public StmtExprMutator { public: - DeferredRemapper(const std::unordered_map& alloc_remap, + DeferredRemapper(const std::unordered_map& alloc_remap, const std::unordered_map& var_remap, - const std::unordered_map& buf_remap, + const std::unordered_map& buf_remap, const std::vector& pending) : alloc_remap_(alloc_remap), var_remap_(var_remap), buf_remap_(buf_remap) { for (const VarNode* ptr : pending) { @@ -863,10 +862,10 @@ class DeferredRemapper : public StmtExprMutator { Stmt VisitStmt_(const AllocBufferNode* op) final { auto node = StmtExprMutator::VisitStmt_(op).as_or_throw(); - const VarNode* data_ptr = op->buffer->data.get(); + const VarNode* data_ptr = op->buffer.get(); if (pending_set_.count(data_ptr)) { if (auto it = alloc_remap_.find(data_ptr); it != alloc_remap_.end()) { - const Buffer& replacement = it->second; + const BufferVar& replacement = it->second; auto* cow = node.CopyOnWrite(); cow->buffer = replacement; if (replacement.scope() == "shared") { @@ -883,7 +882,7 @@ class DeferredRemapper : public StmtExprMutator { // If the DeclBuffer's original data var was remapped by alloc_remap_, // the corresponding AllocBuffer was also remapped, making this DeclBuffer // redundant. Remove it by replacing with a no-op. - const VarNode* orig_data = op->buffer->data.get(); + const VarNode* orig_data = op->buffer.get(); if (pending_set_.count(orig_data) && alloc_remap_.count(orig_data)) { return Evaluate(0); } @@ -895,21 +894,19 @@ class DeferredRemapper : public StmtExprMutator { } private: - ffi::Optional GetRemappedBuffer(const Buffer& buf) { + ffi::Optional GetRemappedBuffer(const BufferVar& buf) { if (auto it = buf_remap_.find(buf.get()); it != buf_remap_.end()) { return it->second; } - if (auto it = var_remap_.find(buf->data.get()); it != var_remap_.end()) { - Buffer new_buf = buf; - new_buf.CopyOnWrite()->data = it->second; - return new_buf; + if (auto it = var_remap_.find(buf.get()); it != var_remap_.end()) { + return BufferVar(it->second); } return std::nullopt; } - const std::unordered_map& alloc_remap_; + const std::unordered_map& alloc_remap_; const std::unordered_map& var_remap_; - const std::unordered_map& buf_remap_; + const std::unordered_map& buf_remap_; std::unordered_set pending_set_; }; diff --git a/src/s_tir/transform/lower_vtcm_alloc.cc b/src/s_tir/transform/lower_vtcm_alloc.cc index 9e4a31255362..002ef491dc8f 100644 --- a/src/s_tir/transform/lower_vtcm_alloc.cc +++ b/src/s_tir/transform/lower_vtcm_alloc.cc @@ -38,15 +38,16 @@ class VtcmAllocator : public StmtExprMutator { VtcmAllocator() {} Stmt VisitStmt_(const AllocBufferNode* op) final { - std::string storage_scope = GetStorageScope(op->buffer->data); + std::string storage_scope = op->buffer.scope(); if (IsVtcmStorage(storage_scope)) { ffi::Array args; args.push_back(StringImm(storage_scope)); args.push_back(IntImm::Int64(op->buffer->shape.size())); args.push_back( Call(PointerType(PrimType::Int(64)), builtin::tvm_stack_make_shape(), op->buffer->shape)); - return Bind(op->buffer->data, - Call(op->buffer->data->ty, builtin::nd_mem_alloc_with_scope(), args)); + return DeclBuffer( + op->buffer, + Call(op->buffer->data_pointer_type, builtin::nd_mem_alloc_with_scope(), args)); } return StmtExprMutator::VisitStmt_(op); } @@ -54,7 +55,7 @@ class VtcmAllocator : public StmtExprMutator { protected: std::string GetStorageScope(const Var& var) { auto* ptr = var->ty.as(); - TVM_FFI_ICHECK(ptr) << "Buffer Var's type annotation must be of PointerType"; + TVM_FFI_ICHECK(ptr) << "BufferVar Var's type annotation must be of PointerType"; return ptr->storage_scope; } }; 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 45a218d85017..a515a6583013 100644 --- a/src/s_tir/transform/manifest_shared_memory_local_stage.cc +++ b/src/s_tir/transform/manifest_shared_memory_local_stage.cc @@ -53,14 +53,14 @@ class IntermediateStageRewriter { explicit IntermediateStageRewriter(const std::vector& ancestor_loop_or_blocks) : ancestor_loop_or_blocks_(ancestor_loop_or_blocks) {} - std::tuple Rewrite(const SBlockNode* block) { + std::tuple Rewrite(const SBlockNode* block) { const BufferStoreNode* store = block->body.as(); 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."; - const Buffer& target_buffer = store->buffer; + const BufferVar& target_buffer = store->buffer; // Step 0: Collect relaxed loops std::vector relaxed_loops = CollectRelaxedOuterLoops(block, target_buffer); @@ -87,7 +87,7 @@ class IntermediateStageRewriter { private: /*! \brief Collect relaxed outer loops from innermost to outermost */ std::vector CollectRelaxedOuterLoops(const SBlockNode* block, - const Buffer& target_buffer) { + const BufferVar& target_buffer) { std::vector relaxed_loops; for (int n = static_cast(ancestor_loop_or_blocks_.size()) - 1, i = n - 1; i >= 0; --i) { const Stmt& ancestor = ancestor_loop_or_blocks_[i]; @@ -113,7 +113,7 @@ class IntermediateStageRewriter { const SBlockNode* ancestor_block = ancestor_block_realize->block.get(); auto it = std::find_if( ancestor_block->alloc_buffers.begin(), ancestor_block->alloc_buffers.end(), - [&target_buffer](const Buffer& buffer) { return buffer.same_as(target_buffer); }); + [&target_buffer](const BufferVar& buffer) { return buffer.same_as(target_buffer); }); TVM_FFI_CHECK(it != ancestor_block->alloc_buffers.end(), ValueError) << "Expect the shared memory allocation to be in the parent block."; break; @@ -123,7 +123,7 @@ class IntermediateStageRewriter { } /*! \brief Create the intermediate stage. */ - Stmt MakeLocalStage(const SBlockNode* block, const Buffer& new_buffer, + Stmt MakeLocalStage(const SBlockNode* block, const BufferVar& new_buffer, ffi::Array local_stage_indices, std::vector relaxed_loops, const BufferStoreNode* store) { // Step 0: Create the body of the local stage, which is BufferStore to the intermediate buffer. @@ -153,8 +153,8 @@ class IntermediateStageRewriter { } /*! \brief Create the intermediate buffer with the extents of the relaxed outer loops. */ - std::pair> CreateIntermediateBuffer( - const std::vector relaxed_loops, const Buffer& buffer) const { + std::pair> CreateIntermediateBuffer( + const std::vector relaxed_loops, const BufferVar& buffer) const { ffi::Array buffer_indices; ffi::Array new_buffer_shape; @@ -166,8 +166,10 @@ class IntermediateStageRewriter { buffer_indices.push_back(relaxed_loop->min + relaxed_loop->loop_var); new_buffer_shape.push_back(relaxed_loop->extent); } - Buffer new_buffer = WithScope(buffer, "local"); - new_buffer.CopyOnWrite()->shape = new_buffer_shape; + BufferVar new_buffer = WithScope(buffer, "local"); + ffi::ObjectPtr type = CopyBufferType(new_buffer); + type->shape = new_buffer_shape; + new_buffer = RebuildBufferVar(new_buffer, std::move(type)); return {new_buffer, buffer_indices}; } @@ -207,11 +209,11 @@ class SharedMemoryLocalStageInserter : public StmtMutator { return new_block; } - std::unordered_set allocated_buffers( + std::unordered_set allocated_buffers( op->alloc_buffers.begin(), op->alloc_buffers.end()); // Visit children and insert local stages (if any) to the proper location. - ffi::Array new_alloc_buffers; + ffi::Array new_alloc_buffers; ffi::Array new_seq; // Helper function to check if the subtree (body of the block) contains any target buffers. @@ -219,7 +221,7 @@ class SharedMemoryLocalStageInserter : public StmtMutator { // block. auto f_check_subtree = [&](int start, int end) { for (int i = start; i < end; ++i) { - const Buffer& buffer = target_buffers_[i]; + const BufferVar& buffer = target_buffers_[i]; if (allocated_buffers.count(buffer)) { new_seq.push_back(buffer_local_stage_.at(buffer)); new_alloc_buffers.push_back(buffer_remap_.at(buffer)); @@ -265,10 +267,10 @@ class SharedMemoryLocalStageInserter : public StmtMutator { } std::vector ancestor_loop_or_blocks_; // ancestor loops or block realize - ffi::Map + ffi::Map buffer_remap_; // mapping from the target buffer to the intermediate buffer - ffi::Map buffer_local_stage_; // mapping from the target buffer to the local stage - ffi::Array target_buffers_; // the target buffers for rewriting + ffi::Map buffer_local_stage_; // mapping from the target buffer to the local stage + ffi::Array target_buffers_; // the target buffers for rewriting }; namespace transform { diff --git a/src/s_tir/transform/memhammer_intermediate_stage.cc b/src/s_tir/transform/memhammer_intermediate_stage.cc index b9e13d0222dc..2a95ed00be16 100644 --- a/src/s_tir/transform/memhammer_intermediate_stage.cc +++ b/src/s_tir/transform/memhammer_intermediate_stage.cc @@ -202,7 +202,7 @@ class IndexPatternFinder : public ExprVisitor { class BufferLoadReplacer : public StmtExprMutator { public: - BufferLoadReplacer(const Buffer& tgt_buffer, const BufferLoad& new_buffer_load) + BufferLoadReplacer(const BufferVar& tgt_buffer, const BufferLoad& new_buffer_load) : tgt_buffer_(tgt_buffer), new_buffer_load_(new_buffer_load) {} Expr VisitExpr_(const BufferLoadNode* op) { @@ -213,7 +213,7 @@ class BufferLoadReplacer : public StmtExprMutator { } private: - Buffer tgt_buffer_; + BufferVar tgt_buffer_; BufferLoad new_buffer_load_; }; @@ -231,7 +231,7 @@ class BufferLoadReplacer : public StmtExprMutator { std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::String storage_scope, ffi::Optional compute_location, const ffi::Array& outer_loops, - Buffer* alloc_buffer) { + BufferVar* alloc_buffer) { Stmt body = stmt; std::vector loops; std::vector loops_under_compute_location; @@ -358,7 +358,7 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S subst_cache_indices.push_back(Substitute(e, subst_map)); } - Buffer new_buffer; + BufferVar new_buffer; if (is_write_cache) { // this is needed for global <- cast(load(wmma)) // shared stage should have the same dtype as wmma @@ -366,8 +366,9 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S } else { new_buffer = WithScope(buf_store->buffer, storage_scope); } - BufferNode* buffer_ptr = new_buffer.CopyOnWrite(); - buffer_ptr->shape = new_shape; + ffi::ObjectPtr buffer_type = CopyBufferType(new_buffer); + buffer_type->shape = new_shape; + new_buffer = RebuildBufferVar(new_buffer, std::move(buffer_type)); *alloc_buffer = new_buffer; Stmt generate_body; @@ -440,7 +441,7 @@ Stmt CreateLocalStage::Rewrite(const Stmt& stmt, const ConstraintSet& constraint Stmt body; For compute_location; std::tie(body, compute_location) = LiftThreadBindingLoops(std::move(stmt)); - Buffer cache_buffer; + BufferVar cache_buffer; Stmt after_caching = InsertCacheStage(body, false, "local", compute_location, constraints.outer_loops, &cache_buffer) .first; diff --git a/src/s_tir/transform/memhammer_lower_auto_copy.cc b/src/s_tir/transform/memhammer_lower_auto_copy.cc index 95123836a011..4caa23e1c075 100644 --- a/src/s_tir/transform/memhammer_lower_auto_copy.cc +++ b/src/s_tir/transform/memhammer_lower_auto_copy.cc @@ -94,10 +94,10 @@ class AutoPadder { * \param buffers the given buffers * \return the list of new padded buffers */ - ffi::Array PadSharedMemory(const ffi::Array& buffers) { - ffi::Array result; + ffi::Array PadSharedMemory(const ffi::Array& buffers) { + ffi::Array result; - for (const Buffer& buffer : buffers) { + for (const BufferVar& buffer : buffers) { runtime::StorageScope scope = runtime::StorageScope::Create(buffer.scope()); if (scope.rank == runtime::StorageRank::kShared) { auto iter_spaces = iter_spaces_[buffer.get()]; @@ -169,14 +169,14 @@ class AutoPadder { reverse_strides.push_back(stride); } // Step 3. create the new padded buffer - ffi::ObjectPtr b = ffi::make_object(*buffer.get()); + ffi::ObjectPtr b = CopyBufferType(buffer); ffi::Array strides; for (int i = static_cast(reverse_strides.size()) - 1; i >= 0; i--) { strides.push_back(reverse_strides[i]); } strides.push_back(1); b->strides = strides; - Buffer new_buffer(b); + BufferVar new_buffer = RebuildBufferVar(buffer, std::move(b)); result.push_back(new_buffer); padded_buffer_map_.Set(buffer, new_buffer); } else { @@ -194,7 +194,7 @@ class AutoPadder { Stmt RewriteBufferAccess(const Stmt& stmt) { class Rewriter : public StmtExprMutator { public: - explicit Rewriter(const ffi::Map& buffer_map) : buffer_map_(buffer_map) {} + explicit Rewriter(const ffi::Map& buffer_map) : buffer_map_(buffer_map) {} private: Expr VisitExpr_(const BufferLoadNode* _op) final { @@ -246,7 +246,7 @@ class AutoPadder { for (const MatchBufferRegion& match_buffer : op->match_buffers) { if (buffer_map_.count(match_buffer->source->buffer)) { changed = true; - Buffer new_buffer = buffer_map_[match_buffer->source->buffer]; + BufferVar new_buffer = buffer_map_[match_buffer->source->buffer]; match_buffers.push_back(MatchBufferRegion( match_buffer->buffer, BufferRegion(new_buffer, match_buffer->source->region))); } else { @@ -269,7 +269,7 @@ class AutoPadder { return ffi::GetRef(op); } } - const ffi::Map& buffer_map_; + const ffi::Map& buffer_map_; }; Rewriter rewriter(padded_buffer_map_); return rewriter(stmt); @@ -573,7 +573,7 @@ class AutoPadder { if (call->op.same_as(tvm_load_matrix_sync_op) || call->op.same_as(tvm_store_matrix_sync_op)) { for (const MatchBufferRegion& r : op->match_buffers) { - Buffer src_buffer = r->source->buffer; + BufferVar src_buffer = r->source->buffer; runtime::StorageScope scope = runtime::StorageScope::Create(src_buffer.scope()); if (scope.rank == runtime::StorageRank::kShared) { Region region = r->source->region; @@ -644,11 +644,11 @@ class AutoPadder { private: /*! \brief A map from the old buffers to the new padded buffers */ - ffi::Map padded_buffer_map_; + ffi::Map padded_buffer_map_; /*! \brief A map from each buffer to the iteration spaces of the accesses*/ - std::unordered_map>>> iter_spaces_; + std::unordered_map>>> iter_spaces_; /*! \brief A map from each buffer to their minimal padding size */ - ffi::Map padding_min_; + ffi::Map padding_min_; /*! \brief max padding size in relative to the original shape*/ const double max_pad_factor_ = 0.25; @@ -702,7 +702,7 @@ class AutoCopyMutator : public StmtExprMutator { for (RewriteRule* rule : rules) { n->body = rule->Apply(std::move(n->body), constraints, &outputs); } - for (const Buffer& buffer : outputs.alloc_buffer) { + for (const BufferVar& buffer : outputs.alloc_buffer) { n->alloc_buffers.push_back(buffer); } for (const auto& p : outputs.padding_min) { diff --git a/src/s_tir/transform/memhammer_rewrite_rule.h b/src/s_tir/transform/memhammer_rewrite_rule.h index af2866836ce9..124f2ec58c6c 100644 --- a/src/s_tir/transform/memhammer_rewrite_rule.h +++ b/src/s_tir/transform/memhammer_rewrite_rule.h @@ -75,9 +75,9 @@ struct ConstraintSet { /*! \brief The set containing all possible outputs of a rewrite rule */ struct OutputSet { /*! \brief New buffers allocated after rewrite */ - ffi::Array alloc_buffer; + ffi::Array alloc_buffer; /*! \brief The minimal padding size of a buffer in base 2 logarithm */ - ffi::Map padding_min; + ffi::Map padding_min; }; /*! @@ -114,14 +114,14 @@ class RewriteRule { } }; -inline bool IsCopyBetweenScope(const Buffer& src_buffer, const Buffer& tgt_buffer, +inline bool IsCopyBetweenScope(const BufferVar& src_buffer, const BufferVar& tgt_buffer, runtime::StorageRank src_rank, runtime::StorageRank tgt_rank) { runtime::StorageScope src_scope = runtime::StorageScope::Create(src_buffer.scope()); runtime::StorageScope tgt_scope = runtime::StorageScope::Create(tgt_buffer.scope()); return src_scope.rank == src_rank && tgt_scope.rank == tgt_rank; } -inline bool IsScope(const Buffer& src_buffer, runtime::StorageRank src_rank) { +inline bool IsScope(const BufferVar& src_buffer, runtime::StorageRank src_rank) { runtime::StorageScope src_scope = runtime::StorageScope::Create(src_buffer.scope()); return src_scope.rank == src_rank; } @@ -134,8 +134,8 @@ class CoalescedAccess : public RewriteRule { CoalescedAccess() = default; Stmt Rewrite(const Stmt& stmt, const ConstraintSet& constraints, OutputSet* output) const final; bool CanApply(const Stmt& stmt, const ConstraintSet& constraints) const final { - Buffer src_buffer = constraints.read_region->buffer; - Buffer tgt_buffer = constraints.write_region->buffer; + BufferVar src_buffer = constraints.read_region->buffer; + BufferVar tgt_buffer = constraints.write_region->buffer; return IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kGlobal, runtime::StorageRank::kShared) || IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kShared, @@ -151,8 +151,8 @@ class InverseMapping : public RewriteRule { InverseMapping() = default; Stmt Rewrite(const Stmt& stmt, const ConstraintSet& constraints, OutputSet* output) const final; bool CanApply(const Stmt& stmt, const ConstraintSet& constraints) const final { - Buffer src_buffer = constraints.read_region->buffer; - Buffer tgt_buffer = constraints.write_region->buffer; + BufferVar src_buffer = constraints.read_region->buffer; + BufferVar tgt_buffer = constraints.write_region->buffer; return IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kShared, runtime::StorageRank::kGlobal); } @@ -166,8 +166,8 @@ class CreateLocalStage : public RewriteRule { CreateLocalStage() = default; Stmt Rewrite(const Stmt& stmt, const ConstraintSet& constraints, OutputSet* output) const final; bool CanApply(const Stmt& stmt, const ConstraintSet& constraints) const final { - Buffer src_buffer = constraints.read_region->buffer; - Buffer tgt_buffer = constraints.write_region->buffer; + BufferVar src_buffer = constraints.read_region->buffer; + BufferVar tgt_buffer = constraints.write_region->buffer; return IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kGlobal, runtime::StorageRank::kShared) && is_one(constraints.add_local_stage); @@ -183,8 +183,8 @@ class WmmaToGlobal : public RewriteRule { WmmaToGlobal() = default; Stmt Rewrite(const Stmt& stmt, const ConstraintSet& constraints, OutputSet* output) const final; bool CanApply(const Stmt& stmt, const ConstraintSet& constraints) const final { - Buffer src_buffer = constraints.read_region->buffer; - Buffer tgt_buffer = constraints.write_region->buffer; + BufferVar src_buffer = constraints.read_region->buffer; + BufferVar tgt_buffer = constraints.write_region->buffer; return IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kWMMAAccumulator, runtime::StorageRank::kGlobal); } @@ -199,8 +199,8 @@ class MmaToGlobal : public RewriteRule { MmaToGlobal() = default; Stmt Rewrite(const Stmt& stmt, const ConstraintSet& constraints, OutputSet* output) const final; bool CanApply(const Stmt& stmt, const ConstraintSet& constraints) const final { - Buffer src_buffer = constraints.read_region->buffer; - Buffer tgt_buffer = constraints.write_region->buffer; + BufferVar src_buffer = constraints.read_region->buffer; + BufferVar tgt_buffer = constraints.write_region->buffer; return IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kMMAMatrixC, runtime::StorageRank::kGlobal); } @@ -214,8 +214,8 @@ class SharedToWmma : public RewriteRule { SharedToWmma() = default; Stmt Rewrite(const Stmt& stmt, const ConstraintSet& constraints, OutputSet* output) const final; bool CanApply(const Stmt& stmt, const ConstraintSet& constraints) const final { - Buffer src_buffer = constraints.read_region->buffer; - Buffer tgt_buffer = constraints.write_region->buffer; + BufferVar src_buffer = constraints.read_region->buffer; + BufferVar tgt_buffer = constraints.write_region->buffer; return IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kShared, runtime::StorageRank::kWMMAMatrixA) || IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kShared, @@ -231,8 +231,8 @@ class WmmaToShared : public RewriteRule { WmmaToShared() = default; Stmt Rewrite(const Stmt& stmt, const ConstraintSet& constraints, OutputSet* output) const final; bool CanApply(const Stmt& stmt, const ConstraintSet& constraints) const final { - Buffer src_buffer = constraints.read_region->buffer; - Buffer tgt_buffer = constraints.write_region->buffer; + BufferVar src_buffer = constraints.read_region->buffer; + BufferVar tgt_buffer = constraints.write_region->buffer; return IsCopyBetweenScope(src_buffer, tgt_buffer, runtime::StorageRank::kWMMAAccumulator, runtime::StorageRank::kShared); } @@ -251,7 +251,7 @@ class WmmaToShared : public RewriteRule { */ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::String storage_scope, ffi::Optional compute_location, - const ffi::Array& outer_loops, Buffer* alloc_buffer); + const ffi::Array& outer_loops, BufferVar* alloc_buffer); } // namespace s_tir } // namespace tvm diff --git a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc index 07bddb3f1295..ab24b0483a22 100644 --- a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc +++ b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc @@ -131,10 +131,10 @@ Stmt RewriteWmmaLoad(Stmt stmt) { const BufferStoreNode* buf_store = TVM_TYPE_AS(body, BufferStoreNode); const BufferLoadNode* buf_load = TVM_TYPE_AS(buf_store->value, BufferLoadNode); - Buffer src_buffer = buf_load->buffer; - Buffer tgt_buffer = buf_store->buffer; + BufferVar src_buffer = buf_load->buffer; + BufferVar tgt_buffer = buf_store->buffer; std::string layout = tgt_buffer.scope() == "wmma.matrix_a" ? "row_major" : "col_major"; - Buffer new_src_buffer( + BufferVar new_src_buffer( /*data=*/Var("src", PointerType(dtype_ty, src_buffer.scope())), /*dtype=*/dtype, /*shape=*/{IntImm::Int32(16), IntImm::Int32(16)}, @@ -143,7 +143,7 @@ Stmt RewriteWmmaLoad(Stmt stmt) { /*name=*/"src", /*data_alignment=*/64, /*offset_factor=*/16); - Buffer new_tgt_buffer( + BufferVar new_tgt_buffer( /*data=*/Var("tgt", PointerType(dtype_ty, tgt_buffer.scope())), /*dtype=*/dtype, /*shape=*/{IntImm::Int32(16), IntImm::Int32(16)}, @@ -168,7 +168,7 @@ Stmt RewriteWmmaLoad(Stmt stmt) { /*data=*/PrimType::Void(), /*op=*/tvm_load_matrix_sync_op, ffi::Array{ - /*0:*/ new_tgt_buffer->data, + /*0:*/ new_tgt_buffer.data(), /*1:*/ PrimExpr(16), /*2:*/ PrimExpr(16), /*3:*/ PrimExpr(16), @@ -176,12 +176,12 @@ Stmt RewriteWmmaLoad(Stmt stmt) { floordiv(floormod(new_tgt_buffer->elem_offset, 256), 16), /*5:*/ Call( - /*dtype=*/new_src_buffer->data->ty, + /*dtype=*/new_src_buffer.data()->ty, /*op=*/builtin::tvm_access_ptr(), /*args=*/ ffi::Array{ /*0:*/ TypeAnnotation(new_src_buffer->dtype), - /*1:*/ new_src_buffer->data, + /*1:*/ new_src_buffer.data(), /*2:*/ new_src_buffer->elem_offset, /*3:*/ new_src_buffer->strides[new_src_buffer->strides.size() - 2] * 16, /*4:*/ PrimExpr(1), @@ -238,13 +238,13 @@ Stmt RewriteWmmaStore(Stmt stmt) { } return true; }); - Buffer src_buffer = buf_load->buffer; - Buffer tgt_buffer = buf_store->buffer; + BufferVar src_buffer = buf_load->buffer; + BufferVar tgt_buffer = buf_store->buffer; PrimType dtype_ty = src_buffer->dtype; const PrimType& dtype = dtype_ty; - Buffer new_src_buffer(/*data=*/Var("src", PointerType(dtype_ty, src_buffer.scope())), + BufferVar new_src_buffer(/*data=*/Var("src", PointerType(dtype_ty, src_buffer.scope())), /*dtype=*/dtype, /*shape=*/{IntImm::Int32(16), IntImm::Int32(16)}, /*strides=*/{}, @@ -252,7 +252,7 @@ Stmt RewriteWmmaStore(Stmt stmt) { /*name=*/"src", /*data_alignment=*/64, /*offset_factor=*/16); - Buffer new_tgt_buffer(/*data=*/Var("tgt", PointerType(dtype_ty, tgt_buffer.scope())), + BufferVar new_tgt_buffer(/*data=*/Var("tgt", PointerType(dtype_ty, tgt_buffer.scope())), /*dtype=*/dtype, /*shape=*/{IntImm::Int32(16), IntImm::Int32(16)}, /*strides=*/{PrimVar("s1", int32_ty), PrimVar("s0", int32_ty)}, @@ -274,7 +274,7 @@ Stmt RewriteWmmaStore(Stmt stmt) { Evaluate(Call( /*data=*/PrimType::Void(), /*op=*/tvm_store_matrix_sync_op, - ffi::Array{/*0:*/ new_src_buffer->data, + ffi::Array{/*0:*/ new_src_buffer.data(), /*1:*/ PrimExpr(16), /*2:*/ PrimExpr(16), /*3:*/ PrimExpr(16), @@ -282,11 +282,11 @@ Stmt RewriteWmmaStore(Stmt stmt) { floordiv(floormod(new_src_buffer->elem_offset, 256), 16), /*5:*/ Call( - /*data=*/new_tgt_buffer->data->ty, + /*data=*/new_tgt_buffer.data()->ty, /*op=*/builtin::tvm_access_ptr(), ffi::Array{ /*0:*/ TypeAnnotation(new_tgt_buffer->dtype), - /*1:*/ new_tgt_buffer->data, + /*1:*/ new_tgt_buffer.data(), /*2:*/ new_tgt_buffer->elem_offset, /*3:*/ new_tgt_buffer->strides[0] * 16, /*4:*/ PrimExpr(2), @@ -350,7 +350,7 @@ Stmt WmmaToGlobal::Rewrite(const Stmt& stmt, const ConstraintSet& constraints, ffi::Optional compute_location; std::tie(body, compute_location) = TileWmmaBlock(stmt); SeqStmt seq{nullptr}; - Buffer cache_buffer; + BufferVar cache_buffer; // Step 1. add a shared memory cache std::tie(body, seq) = InsertCacheStage(std::move(body), true, "shared.dyn", compute_location, constraints.outer_loops, &cache_buffer); @@ -466,11 +466,11 @@ Stmt RewriteMmaStore(Stmt stmt) { // https://docs.nvidia.com/cuda/archive/11.1.0/pdf/ptx_isa_7.1.pdf // Step 3.1. Generate new buffer - Buffer src_buffer = buf_load->buffer; - Buffer tgt_buffer = buf_store->buffer; + BufferVar src_buffer = buf_load->buffer; + BufferVar tgt_buffer = buf_store->buffer; PrimType dtype_ty = src_buffer->dtype; const PrimType& dtype = dtype_ty; - Buffer new_src_buffer(/*data=*/Var("src", PointerType(dtype_ty, src_buffer.scope())), + BufferVar new_src_buffer(/*data=*/Var("src", PointerType(dtype_ty, src_buffer.scope())), /*dtype=*/dtype, /*shape=*/{IntImm::Int32(8), IntImm::Int32(8)}, /*strides=*/{}, @@ -478,7 +478,7 @@ Stmt RewriteMmaStore(Stmt stmt) { /*name=*/"src", /*data_alignment=*/64, /*offset_factor=*/8); - Buffer new_tgt_buffer(/*data=*/Var("tgt", PointerType(dtype_ty, tgt_buffer.scope())), + BufferVar new_tgt_buffer(/*data=*/Var("tgt", PointerType(dtype_ty, tgt_buffer.scope())), /*dtype=*/dtype, /*shape=*/{IntImm::Int32(8), IntImm::Int32(8)}, /*strides=*/{PrimVar("s1", int32_ty), PrimVar("s0", int32_ty)}, @@ -566,7 +566,7 @@ Stmt MmaToGlobal::Rewrite(const Stmt& stmt, const ConstraintSet& constraints, ffi::Optional compute_location; std::tie(body, compute_location) = TileMmaToGlobalBlock(stmt); SeqStmt seq{nullptr}; - Buffer cache_buffer; + BufferVar cache_buffer; // Step 1. add a shared memory cache std::tie(body, seq) = InsertCacheStage(std::move(body), true, "shared.dyn", compute_location, constraints.outer_loops, &cache_buffer); diff --git a/src/s_tir/transform/merge_shared_memory_allocations.cc b/src/s_tir/transform/merge_shared_memory_allocations.cc index 6c6956d22b39..93ccb33ee647 100644 --- a/src/s_tir/transform/merge_shared_memory_allocations.cc +++ b/src/s_tir/transform/merge_shared_memory_allocations.cc @@ -60,6 +60,16 @@ bool IsStaticSharedMemory(Var buffer_var) { return storage_scope.rank == runtime::StorageRank::kShared && storage_scope.tag == ""; } +bool IsDynamicSharedMemory(const BufferVar& buffer) { + StorageScope storage_scope = runtime::StorageScope::Create(buffer.scope()); + return storage_scope.rank == runtime::StorageRank::kShared && storage_scope.tag == ".dyn"; +} + +bool IsStaticSharedMemory(const BufferVar& buffer) { + StorageScope storage_scope = runtime::StorageScope::Create(buffer.scope()); + return storage_scope.rank == runtime::StorageRank::kShared && storage_scope.tag == ""; +} + /*! * \brief Compute constant allocation size from buffer's allocation shape. * \return Product of extents if all constant, 0 otherwise. @@ -78,23 +88,23 @@ static int64_t ConstantAllocationSize(const ffi::Array& extents) { } /*! - * \brief collect the mapping from the buffer var to its Buffer within a subtree + * \brief collect the mapping from the buffer var to its BufferVar within a subtree */ class AllocateCollector : public StmtExprVisitor { public: explicit AllocateCollector(bool is_dynamic) : is_dynamic_(is_dynamic) {} void VisitStmt_(const AllocBufferNode* op) final { - if (is_dynamic_ && IsDynamicSharedMemory(op->buffer->data)) { - shmem_allocs_[op->buffer->data.get()] = op->buffer; - } else if (!is_dynamic_ && IsStaticSharedMemory(op->buffer->data)) { - shmem_allocs_[op->buffer->data.get()] = op->buffer; + if (is_dynamic_ && IsDynamicSharedMemory(op->buffer)) { + shmem_allocs_[op->buffer.get()] = op->buffer; + } else if (!is_dynamic_ && IsStaticSharedMemory(op->buffer)) { + shmem_allocs_[op->buffer.get()] = op->buffer; } StmtExprVisitor::VisitStmt_(op); } - // The mapping from the original buffer var to its Buffer - std::unordered_map shmem_allocs_; + // The mapping from the original buffer var to its BufferVar + std::unordered_map shmem_allocs_; private: bool is_dynamic_; @@ -133,12 +143,12 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { // the level in the scope stack size_t level{0}; // The buffer object - Buffer buffer; + BufferVar buffer; }; void VisitStmt_(const AllocBufferNode* op) final { size_t level = scope_.size(); - const VarNode* buf = op->buffer->data.get(); + const VarNode* buf = op->buffer.get(); alloc_info_[buf].buffer = op->buffer; alloc_info_[buf].level = level; StmtExprVisitor::VisitStmt_(op); @@ -149,7 +159,7 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { // visit subexpr StmtExprVisitor::VisitStmt_(op); // Add write access. - const VarNode* buf = op->buffer->data.get(); + const VarNode* buf = op->buffer.get(); auto it = alloc_info_.find(buf); if (it != alloc_info_.end() && it->second.buffer.defined()) { TVM_FFI_ICHECK_LT(it->second.level, scope_.size()); @@ -180,7 +190,7 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { void VisitExpr_(const BufferLoadNode* op) final { // Add read access. StmtExprVisitor::VisitExpr_(op); - const VarNode* buf = op->buffer->data.get(); + const VarNode* buf = op->buffer.get(); auto it = alloc_info_.find(buf); if (it != alloc_info_.end() && it->second.buffer.defined()) { TVM_FFI_ICHECK_LT(it->second.level, scope_.size()) @@ -317,15 +327,15 @@ class SharedMemoryRewriter : public StmtExprMutator { */ struct KernelScope { // The merged buffer var for THIS kernel launch. - Var merged_buf_var; + BufferVar merged_buffer; // Total byte size of THIS kernel's merged buffer. PrimExpr merged_alloc_size{0}; // Allocations from THIS kernel's subtree. - std::unordered_map shmem_allocs; + std::unordered_map shmem_allocs; // Per-buffer byte offset into merged_buf_var. std::unordered_map buffer_byte_offsets; - // Buffer-object remap: original Buffer -> merged-data-var Buffer. - std::unordered_map buffer_remap; + // BufferVar-object remap: original BufferVar -> merged-data-var BufferVar. + std::unordered_map buffer_remap; // Has any original alloc in this scope been marked volatile? bool has_volatile_alloc{false}; // Liveness data (event_map, alloc_map, const_free_map, sym_free_list) — all per-scope. @@ -339,12 +349,10 @@ class SharedMemoryRewriter : public StmtExprMutator { * \brief Create a fresh merged buffer Var for a new kernel scope. * Same name string is fine — Var identity is by pointer, not name. */ - Var MakeMergedBufferVar() { - if (is_dynamic_) { - return Var("buf_dyn_shmem", PointerType(PrimType::UInt(8), "shared.dyn")); - } else { - return Var("buf_shmem", PointerType(PrimType::UInt(8), "shared")); - } + BufferVar MakeMergedBuffer(PrimExpr size) { + return decl_buffer({std::move(size)}, PrimType::UInt(8), + is_dynamic_ ? "buf_dyn_shmem" : "buf_shmem", + is_dynamic_ ? "shared.dyn" : "shared"); } Stmt VisitStmt_(const AttrStmtNode* op) final { @@ -354,8 +362,6 @@ class SharedMemoryRewriter : public StmtExprMutator { // 1. Push a fresh scope. scope_stack_.emplace_back(); KernelScope& scope = scope_stack_.back(); - scope.merged_buf_var = MakeMergedBufferVar(); - // 2. Collect shmem allocs that belong to THIS subtree. AllocateCollector collector(is_dynamic_); collector(op->body); @@ -380,9 +386,14 @@ class SharedMemoryRewriter : public StmtExprMutator { // 4. Compute byte offsets / merged_alloc_size. this->ComputeOffsets(scope); + scope.merged_buffer = MakeMergedBuffer(scope.merged_alloc_size); // 5. Recursively mutate the body — reads scope_stack_.back() for all rewrites. Stmt visited_body = StmtExprMutator::VisitStmt(op->body); + for (const auto& [_, remapped] : scope.buffer_remap) { + visited_body = + SeqStmt::Flatten(DeclBuffer(remapped, scope.merged_buffer.data()), visited_body); + } in_thread_env_ = false; @@ -393,13 +404,11 @@ class SharedMemoryRewriter : public StmtExprMutator { } // 7. Wrap with the merged-buffer AllocBuffer. - Buffer merged_buf(scope.merged_buf_var, PrimType::UInt(8), {scope.merged_alloc_size}, {}, - PrimExpr(), scope.merged_buf_var->name, 0, 0); ffi::Map annotations; if (scope.has_volatile_alloc) { annotations.Set(tirx::attr::kVolatile, true); } - Stmt alloc_stmt = AllocBuffer(merged_buf, annotations); + Stmt alloc_stmt = AllocBuffer(scope.merged_buffer, annotations); Stmt new_body = SeqStmt::Flatten(alloc_stmt, visited_body); // 8. Pop the scope. @@ -411,10 +420,10 @@ class SharedMemoryRewriter : public StmtExprMutator { } Stmt VisitStmt_(const AllocBufferNode* op) final { - if (IsAppropriateSharedMemory(op->buffer->data)) { + if (IsAppropriateSharedMemory(op->buffer)) { if (!scope_stack_.empty()) { KernelScope& scope = scope_stack_.back(); - if (scope.shmem_allocs.count(op->buffer->data.get())) { + if (scope.shmem_allocs.count(op->buffer.get())) { if (op->annotations.count(tirx::attr::kVolatile)) { scope.has_volatile_alloc = true; } @@ -447,14 +456,15 @@ class SharedMemoryRewriter : public StmtExprMutator { template Node VisitBufferAccess(Node node) { - if (IsAppropriateSharedMemory(node->buffer->data) && !scope_stack_.empty() && - scope_stack_.back().shmem_allocs.count(node->buffer->data.get())) { + if (IsAppropriateSharedMemory(node->buffer) && !scope_stack_.empty() && + scope_stack_.back().shmem_allocs.count(node->buffer.get())) { TVM_FFI_ICHECK_EQ(node->indices.size(), 1) << "MergeSharedMemoryAllocations expects flat memory buffers, " << "and is to be run after " << "FlattenBuffer"; ffi::Array indices = { - node->indices[0] + this->GetBufferOffset(node->buffer->data, node->buffer->dtype->dtype)}; + node->indices[0] + + this->GetBufferOffset(node->buffer.var(), node->buffer->dtype->dtype)}; auto writer = node.CopyOnWrite(); writer->buffer = GetUpdatedBuffer(node->buffer); @@ -464,10 +474,10 @@ class SharedMemoryRewriter : public StmtExprMutator { return node; } - Buffer GetUpdatedBuffer(Buffer buffer) { + BufferVar GetUpdatedBuffer(BufferVar buffer) { if (scope_stack_.empty()) return buffer; KernelScope& scope = scope_stack_.back(); - if (!scope.shmem_allocs.count(buffer->data.get())) return buffer; + if (!scope.shmem_allocs.count(buffer.get())) return buffer; auto key = buffer.get(); auto it = scope.buffer_remap.find(key); @@ -475,14 +485,13 @@ class SharedMemoryRewriter : public StmtExprMutator { return it->second; } - if (IsAppropriateSharedMemory(buffer->data)) { + if (IsAppropriateSharedMemory(buffer)) { TVM_FFI_ICHECK_EQ(buffer->shape.size(), 1) - << "Buffer " << buffer << " has shape " << buffer->shape << ". " + << "BufferVar " << buffer << " has shape " << buffer->shape << ". " << "MergeSharedMemoryAllocations expects flat memory buffers, " << "and is to be run after " << "FlattenBuffer"; - auto writer = buffer.CopyOnWrite(); - writer->data = scope.merged_buf_var; + buffer = RebuildBufferVar(buffer, CopyBufferType(buffer)); } scope.buffer_remap[key] = buffer; @@ -508,7 +517,8 @@ class SharedMemoryRewriter : public StmtExprMutator { PrimExpr offset = this->VisitPrimExpr(op->args[2].as_or_throw()); PrimExpr extent = this->VisitPrimExpr(op->args[3].as_or_throw()); return Call(op->ty, op->op, - {op->args[0], scope_stack_.back().merged_buf_var, extra_offset + offset, extent, + {op->args[0], scope_stack_.back().merged_buffer.data(), + extra_offset + offset, extent, op->args[4]}); } else if (op->op.same_as(ptx_cp_async_op)) { TVM_FFI_ICHECK((op->args.size() == 5U) || (op->args.size() == 6U)); @@ -530,13 +540,13 @@ class SharedMemoryRewriter : public StmtExprMutator { int index_factor = (static_cast(dtype.bits) * static_cast(dtype.lanes) + 7) / 8; if (op->args.size() == 5) return Call(op->ty.as_or_throw(), op->op, - {scope_stack_.back().merged_buf_var, + {scope_stack_.back().merged_buffer.data(), mul(extra_offset + offset, PrimExpr(index_factor)), op->args[2], op->args[3].as_or_throw(), op->args[4].as_or_throw()}) .as_or_throw(); else return Call(op->ty.as_or_throw(), op->op, - {scope_stack_.back().merged_buf_var, + {scope_stack_.back().merged_buffer.data(), mul(extra_offset + offset, PrimExpr(index_factor)), op->args[2], op->args[3].as_or_throw(), op->args[4].as_or_throw(), op->args[5].as_or_throw()}) @@ -560,6 +570,10 @@ class SharedMemoryRewriter : public StmtExprMutator { return is_dynamic_ ? IsDynamicSharedMemory(var) : IsStaticSharedMemory(var); } + bool IsAppropriateSharedMemory(const BufferVar& buffer) { + return is_dynamic_ ? IsDynamicSharedMemory(buffer) : IsStaticSharedMemory(buffer); + } + /*! * \brief Liveness analysis to find gen and kill point of each variable. * \param seq the linear pattern of storage access @@ -622,7 +636,7 @@ class SharedMemoryRewriter : public StmtExprMutator { if (it != scope.event_map.end() && seq[i].scope_pair_offset >= 0) { for (const VarNode* var : it->second.gen) { TVM_FFI_ICHECK(scope.shmem_allocs.count(var)); - const Buffer& buf = scope.shmem_allocs.at(var); + const BufferVar& buf = scope.shmem_allocs.at(var); StorageEntry* dst_entry = FindAlloc(buf, scope); scope.alloc_map[var] = dst_entry; } @@ -656,7 +670,7 @@ class SharedMemoryRewriter : public StmtExprMutator { for (const StorageEntry* e : all_entry) { for (int i = 0; i < static_cast(e->allocs.size()); i++) { for (const VarNode* buffer : e->allocs[i]) { - const Buffer& buf = scope.shmem_allocs.at(buffer); + const BufferVar& buf = scope.shmem_allocs.at(buffer); int elem_bytes = static_cast(buf->dtype.StorageBytes()); align[i] = std::max(align[i], elem_bytes); } @@ -668,7 +682,7 @@ class SharedMemoryRewriter : public StmtExprMutator { for (int i = 0; i < static_cast(e->allocs.size()); i++) { PrimExpr inner_offset = 0; for (const VarNode* buffer : e->allocs[i]) { - const Buffer& buf = scope.shmem_allocs.at(buffer); + const BufferVar& buf = scope.shmem_allocs.at(buffer); ffi::Array alloc_shape = GetBufferAllocationShape(buf); int elem_bytes = static_cast(buf->dtype.StorageBytes()); int align_bytes = std::max(align[i], elem_bytes); @@ -696,10 +710,10 @@ class SharedMemoryRewriter : public StmtExprMutator { * \param const_nbits the size of the allocation in bits * \return the new storage entry */ - StorageEntry* NewAlloc(const Buffer& buf, size_t const_nbits) { + StorageEntry* NewAlloc(const BufferVar& buf, size_t const_nbits) { // Re-use not successful, allocate a new buffer. StorageEntry* entry = arena_.make(); - entry->allocs.push_back({buf->data.get()}); + entry->allocs.push_back({buf.get()}); entry->const_nbits = const_nbits; return entry; } @@ -710,7 +724,7 @@ class SharedMemoryRewriter : public StmtExprMutator { * \param scope the kernel scope whose free lists to search * \return the storage entry */ - StorageEntry* FindAlloc(const Buffer& buf, KernelScope& scope) { + StorageEntry* FindAlloc(const BufferVar& buf, KernelScope& scope) { // skip plan for local variable, // compiler can do a better job with register allocation. const uint64_t match_range = 16; @@ -737,7 +751,7 @@ class SharedMemoryRewriter : public StmtExprMutator { StorageEntry* e = it->second; e->const_nbits = std::max(const_nbits, e->const_nbits); scope.const_free_map.erase(it); - e->allocs.push_back({buf->data.get()}); + e->allocs.push_back({buf.get()}); return e; } // Then start looking at smaller buffers. @@ -764,7 +778,7 @@ class SharedMemoryRewriter : public StmtExprMutator { break; } } - reuse_allocs.push_back({buf->data.get()}); + reuse_allocs.push_back({buf.get()}); if (mem_ct != 0) { StorageEntry* e = arena_.make(); e->const_nbits = std::max(const_nbits, mem_ct); 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 ecb66770ff4e..ccf5400cae19 100644 --- a/src/s_tir/transform/plan_update_buffer_allocation_location.cc +++ b/src/s_tir/transform/plan_update_buffer_allocation_location.cc @@ -38,10 +38,10 @@ class CollectManagedAllocations : public StmtExprVisitor { public: void VisitStmt_(const SBlockNode* op) final { for (const auto& buf : op->alloc_buffers) { - managed_allocations.insert(buf->data.get()); + managed_allocations.insert(buf.get()); } for (const auto& buf : op->match_buffers) { - managed_allocations.insert(buf->buffer->data.get()); + managed_allocations.insert(buf->buffer.get()); } StmtExprVisitor::VisitStmt_(op); } @@ -54,7 +54,7 @@ class CollectManagedAllocations : public StmtExprVisitor { /*! \brief Collect the allocate buffer order. */ class BufferAllocateOrderCollector : public StmtExprVisitor { public: - static ffi::Array Collect(const PrimFunc& func) { + static ffi::Array Collect(const PrimFunc& func) { BufferAllocateOrderCollector collector; for (const auto& kv : func->buffer_map) { collector.buffer_alloc_recorder_.push_back(kv.second); @@ -64,13 +64,13 @@ class BufferAllocateOrderCollector : public StmtExprVisitor { } private: - bool find(const Buffer& buf) { + bool find(const BufferVar& buf) { return std::find(buffer_alloc_recorder_.begin(), buffer_alloc_recorder_.end(), buf) != buffer_alloc_recorder_.end(); } void VisitStmt_(const SBlockNode* op) final { - for (const Buffer& buffer : op->alloc_buffers) { + for (const BufferVar& buffer : op->alloc_buffers) { buffer_alloc_recorder_.push_back(buffer); } // Also visit match_buffers to collect buffers that only appear in read and match_buffer @@ -99,38 +99,38 @@ class BufferAllocateOrderCollector : public StmtExprVisitor { } /*! \brief The buffer allocated order recorder. */ - ffi::Array buffer_alloc_recorder_; + ffi::Array buffer_alloc_recorder_; }; class BufferAllocationLocator : public StmtExprMutator { public: explicit BufferAllocationLocator(const PrimFunc& func) { - ffi::Map> buffer_lca = DetectBufferAccessLCA(func); + ffi::Map> buffer_lca = DetectBufferAccessLCA(func); // The buffer_alloc_recorder Array is used to keep the buffer allocation order // since the buffer_lca Map is unordered. - ffi::Array buffer_alloc_recorder = BufferAllocateOrderCollector::Collect(func); + ffi::Array buffer_alloc_recorder = BufferAllocateOrderCollector::Collect(func); std::unordered_set arg_buffer_vars; CollectManagedAllocations collector; collector(func->body); managed_allocations_ = collector.managed_allocations; for (const auto& kv : func->buffer_map) { - const Buffer& buffer = kv.second; - arg_buffer_vars.emplace(buffer->data.get()); - buffer_data_to_buffer_.Set(buffer->data, buffer); + const BufferVar& buffer = kv.second; + arg_buffer_vars.emplace(buffer.get()); + buffer_data_to_buffer_.Set(buffer.var(), buffer); } // create buffers to be allocated at each stmts for (const auto& buffer : buffer_alloc_recorder) { auto it = buffer_lca.find(buffer); if (it != buffer_lca.end()) { const StmtNode* stmt = (*it).second.has_value() ? (*it).second.value().get() : nullptr; - if (arg_buffer_vars.count(buffer->data.get())) { + if (arg_buffer_vars.count(buffer.get())) { continue; } - if (managed_allocations_.count(buffer->data.get())) { + if (managed_allocations_.count(buffer.get())) { alloc_buffers_[stmt].push_back(buffer); } - buffer_data_to_buffer_.Set(buffer->data, buffer); + buffer_data_to_buffer_.Set(buffer.var(), buffer); } } } @@ -141,15 +141,15 @@ class BufferAllocationLocator : public StmtExprMutator { if (it == alloc_buffers_.end()) { return StmtMutator::VisitStmt_(op); } - for (const Buffer& buf : it->second) { - buffer_data_to_buffer_.Set(buf->data, buf); + for (const BufferVar& buf : it->second) { + buffer_data_to_buffer_.Set(buf.var(), buf); } auto node = StmtMutator::VisitStmt_(op).as_or_throw(); - ffi::Array new_block_alloc_bufs; - for (const Buffer& buf : it->second) { - if (managed_allocations_.count(buf->data.get())) { - buffer_data_to_buffer_.erase(buf->data); + ffi::Array new_block_alloc_bufs; + for (const BufferVar& buf : it->second) { + if (managed_allocations_.count(buf.get())) { + buffer_data_to_buffer_.erase(buf.var()); new_block_alloc_bufs.push_back(buf); } } @@ -163,17 +163,17 @@ class BufferAllocationLocator : public StmtExprMutator { Stmt VisitStmt_(const SBlockNode* op) final { TVM_FFI_ICHECK(!op->init.has_value()); - ffi::Array alloc_buffers; + ffi::Array alloc_buffers; auto it = alloc_buffers_.find(op); if (it != alloc_buffers_.end()) { alloc_buffers = it->second; - for (const Buffer& buf : it->second) { - buffer_data_to_buffer_.Set(buf->data, buf); + for (const BufferVar& buf : it->second) { + buffer_data_to_buffer_.Set(buf.var(), buf); } } for (const MatchBufferRegion match_buffer : op->match_buffers) { - const Var& target_var = match_buffer->buffer->data; - const Var& source_var = match_buffer->source->buffer->data; + const Var target_var = match_buffer->buffer.var(); + const Var source_var = match_buffer->source->buffer.var(); TVM_FFI_ICHECK(buffer_data_to_buffer_.count(source_var)); buffer_data_to_buffer_.Set(target_var, match_buffer->buffer); } @@ -184,13 +184,13 @@ class BufferAllocationLocator : public StmtExprMutator { // No longer consider buffers created by match_buffer inside the block when updating access // region. for (const MatchBufferRegion match_buffer : op->match_buffers) { - const Var& target_var = match_buffer->buffer->data; + const Var target_var = match_buffer->buffer.var(); buffer_data_to_buffer_.erase(target_var); } // No longer consider buffers allocated inside the block when updating access region. if (it != alloc_buffers_.end()) { - for (const Buffer& buf : it->second) { - buffer_data_to_buffer_.erase(buf->data); + for (const BufferVar& buf : it->second) { + buffer_data_to_buffer_.erase(buf.var()); } } @@ -202,7 +202,7 @@ class BufferAllocationLocator : public StmtExprMutator { return Stmt(n); } - Stmt InjectOpaqueBlock(Stmt body, const ffi::Array& alloc_buffers) { + Stmt InjectOpaqueBlock(Stmt body, const ffi::Array& alloc_buffers) { TVM_FFI_ICHECK(!alloc_buffers.empty()); SBlock opaque_block(/*iter_vars=*/{}, /*reads=*/{}, @@ -224,7 +224,7 @@ class BufferAllocationLocator : public StmtExprMutator { const ffi::Array& region) const { ffi::Array result; for (const BufferRegion& buffer_region : region) { - if (buffer_data_to_buffer_.count(buffer_region->buffer->data)) { + if (buffer_data_to_buffer_.count(buffer_region->buffer.var())) { result.push_back(buffer_region); } } @@ -232,9 +232,9 @@ class BufferAllocationLocator : public StmtExprMutator { } /*! \brief The map from stmt to the buffers to be allocated under it. */ - std::unordered_map> alloc_buffers_; + std::unordered_map> alloc_buffers_; /*! \brief The buffer already allocated during recursive visiting. */ - ffi::Map buffer_data_to_buffer_; + ffi::Map buffer_data_to_buffer_; /*! \brief Buffers that are allocated within a BlockNode, and may be moved. */ std::unordered_set managed_allocations_; }; 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 cfab35ad343f..8cb514954040 100644 --- a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc +++ b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc @@ -37,7 +37,7 @@ using namespace tvm::tirx; class RemoveLayoutRewriteBlock : public StmtMutator { public: - static std::tuple, + static std::tuple, std::unordered_map, std::unordered_map>> Rewrite(PrimFunc f) { @@ -57,8 +57,8 @@ class RemoveLayoutRewriteBlock : public StmtMutator { if (it == block->annotations.end() || !is_one((*it).second.cast())) { // The block is not a weight layout block // Remove allocates if needed - ffi::Array alloc_buffers; - for (const Buffer& buffer : block->alloc_buffers) { + ffi::Array alloc_buffers; + for (const BufferVar& buffer : block->alloc_buffers) { if (!rewritten_buffers_.count(buffer)) { alloc_buffers.push_back(buffer); } @@ -84,7 +84,7 @@ class RemoveLayoutRewriteBlock : public StmtMutator { const auto* load = store->value.as(); TVM_FFI_ICHECK(load); - // Step 3. Update Buffer + // Step 3. Update BufferVar buf_map_.Set(load->buffer, store->buffer); rewritten_buffers_.insert(store->buffer); @@ -99,18 +99,18 @@ class RemoveLayoutRewriteBlock : public StmtMutator { TVM_FFI_ICHECK(ind.as()); load_indices.push_back(ind.as_or_throw()); } - buffer_var_to_index_map_[load->buffer->data.get()] = IndexMap(load_indices, store->indices); + buffer_var_to_index_map_[load->buffer.get()] = IndexMap(load_indices, store->indices); - buffer_var_to_rewritten_shape_[load->buffer->data.get()] = store->buffer->shape; + buffer_var_to_rewritten_shape_[load->buffer.get()] = store->buffer->shape; return Stmt(n); } private: /*! \brief The buffer map from original layout buffer to rewritten buffer */ - ffi::Map buf_map_; + ffi::Map buf_map_; /*! \brief The buffer map from original layout buffer to rewritten buffer */ - std::unordered_set rewritten_buffers_; + std::unordered_set rewritten_buffers_; /*! \brief Maps a buffer load to an index map associated with the load / store in a layout rewrite block. */ std::unordered_map buffer_var_to_index_map_; @@ -126,7 +126,7 @@ class WeightLayoutRewriteBlockRemover : public StmtMutator { PrimFuncNode* n = f_.CopyOnWrite(); - ffi::Map buffer_map; + ffi::Map buffer_map; for (const auto& [param, buffer] : f_->buffer_map) { auto it = buf_map.find(buffer); if (it != buf_map.end()) { diff --git a/src/s_tir/transform/renew_defs.cc b/src/s_tir/transform/renew_defs.cc index 23410ae42280..0a88920b4e5a 100644 --- a/src/s_tir/transform/renew_defs.cc +++ b/src/s_tir/transform/renew_defs.cc @@ -19,7 +19,7 @@ /*! * \file renew_defs.cc - * \brief Renew the definition nodes for a TIR, including Var, Buffer and IterVar. + * \brief Renew the definition nodes for a TIR, including Var, BufferVar and IterVar. */ #include @@ -56,7 +56,7 @@ class RenewDefMutator : public StmtExprMutator { for (const auto& param : func->params) { auto it = func->buffer_map.find(param); if (it != func->buffer_map.end()) { - const Buffer& buffer = (*it).second; + const BufferVar& buffer = (*it).second; for (const PrimExpr& e : buffer->shape) { if (auto var = e.as()) { if (generator.remap_.count(var.value()) == 0) { @@ -68,13 +68,13 @@ class RenewDefMutator : public StmtExprMutator { } // Redefine buffers in order // TODO(Siyuan Feng): checking var is used after define - ffi::Map buffer_map; + ffi::Map buffer_map; for (const auto& param : func->params) { auto it = func->buffer_map.find(param); if (it != func->buffer_map.end()) { - const Buffer& buffer = (*it).second; + const BufferVar& buffer = (*it).second; Var new_param = generator.VisitExpr(param).as_or_throw(); - Buffer new_buffer = generator.DefineBuffer(buffer); + BufferVar new_buffer = generator.DefineBuffer(buffer); buffer_map.Set(new_param, new_buffer); } } @@ -107,13 +107,13 @@ class RenewDefMutator : public StmtExprMutator { // Override VisitBufferDef to create fresh buffer copies at definition sites // (AllocBuffer, DeclBuffer, SBlock alloc_buffers, match_buffers) - Buffer VisitBufferDef(const Buffer& buffer, bool alloc_data) final { + BufferVar VisitBufferDef(const BufferVar& buffer, bool alloc_data) final { return DefineBuffer(buffer); } // Override VisitBufferUse to remap buffers at use sites // (BufferStore, BufferLoad, SBlock reads/writes) - Buffer VisitBufferUse(const Buffer& buffer) final { return UseOrRemapBuffer(buffer); } + BufferVar VisitBufferUse(const BufferVar& buffer) final { return UseOrRemapBuffer(buffer); } Stmt VisitStmt_(const SBlockNode* op) final { // Step 0. Re-define Itervars @@ -121,8 +121,8 @@ class RenewDefMutator : public StmtExprMutator { op->iter_vars.Map(std::bind(&RenewDefMutator::VisitIterVar, this, std::placeholders::_1)); // Step 1. Re-define buffers allocated under the block - ffi::Array alloc_buffers = - op->alloc_buffers.Map([this](const Buffer& buf) { return this->DefineBuffer(buf); }); + ffi::Array alloc_buffers = + op->alloc_buffers.Map([this](const BufferVar& buf) { return this->DefineBuffer(buf); }); // Step 2. Re-define match_buffers ffi::Array match_buffers = op->match_buffers.Map( @@ -167,10 +167,10 @@ class RenewDefMutator : public StmtExprMutator { remap_.Set(source, target); } - Buffer DefineBuffer(const Buffer& buffer) { + BufferVar DefineBuffer(const BufferVar& buffer) { auto it = remap_.find(buffer); if (it != remap_.end()) { - return (*it).second.as_or_throw(); + return (*it).second.as_or_throw(); } auto redefine_if_is_var = [this](const Expr& expr) -> Expr { @@ -184,8 +184,6 @@ class RenewDefMutator : public StmtExprMutator { } }; - // data is DEFINED by this buffer — needs a fresh copy - Var data = redefine_if_is_var(buffer->data).as_or_throw(); // shape is USED (references existing definitions like buffer_map shape vars), // remap via VisitExpr to avoid creating spurious new var definitions auto visit_expr = [this](const PrimExpr& e) -> PrimExpr { return this->VisitPrimExpr(e); }; @@ -196,35 +194,32 @@ class RenewDefMutator : public StmtExprMutator { [&](const PrimExpr& expr) { return redefine_if_is_var(expr).as_or_throw(); }); PrimExpr elem_offset = redefine_if_is_var(buffer->elem_offset).as_or_throw(); - auto n = ffi::make_object(*buffer.get()); - n->data = std::move(data); + auto n = CopyBufferType(buffer); n->shape = std::move(shape); n->strides = std::move(strides); n->elem_offset = std::move(elem_offset); - Buffer new_buffer(n); + BufferVar new_buffer = RebuildBufferVar(buffer, std::move(n)); this->AddDefRemap(buffer, new_buffer); return new_buffer; } - Buffer UseOrRemapBuffer(const Buffer& buffer) { + BufferVar UseOrRemapBuffer(const BufferVar& buffer) { // If the buffer has been remapped, return the remapped buffer, otherwise, // remap it without creating new var definitions. auto it = remap_.find(buffer); if (it != remap_.end()) { - return (*it).second.as_or_throw(); + return (*it).second.as_or_throw(); } - Var data = VisitExpr(buffer->data).as_or_throw(); auto visit_expr = [this](const PrimExpr& e) -> PrimExpr { return this->VisitPrimExpr(e); }; ffi::Array shape = buffer->shape.Map(visit_expr); ffi::Array strides = buffer->strides.Map(visit_expr); PrimExpr elem_offset = VisitPrimExpr(buffer->elem_offset); - auto n = ffi::make_object(*buffer.get()); - n->data = std::move(data); + auto n = CopyBufferType(buffer); n->shape = std::move(shape); n->strides = std::move(strides); n->elem_offset = std::move(elem_offset); - Buffer new_buffer(n); + BufferVar new_buffer = RebuildBufferVar(buffer, std::move(n)); this->AddDefRemap(buffer, new_buffer); return new_buffer; } @@ -243,7 +238,7 @@ class RenewDefMutator : public StmtExprMutator { } MatchBufferRegion VisitMatchBuffer(const MatchBufferRegion& match_buffer) { - Buffer buffer = DefineBuffer(match_buffer->buffer); + BufferVar buffer = DefineBuffer(match_buffer->buffer); BufferRegion region = VisitBufferRegion(match_buffer->source); return MatchBufferRegion(std::move(buffer), std::move(region)); } @@ -259,7 +254,7 @@ class RenewDefMutator : public StmtExprMutator { } BufferRegion VisitBufferRegion(const BufferRegion& buffer_region) { - Buffer buffer = UseOrRemapBuffer(buffer_region->buffer); + BufferVar buffer = UseOrRemapBuffer(buffer_region->buffer); ffi::Array region = buffer_region->region.Map( std::bind(&RenewDefMutator::VisitRange, this, std::placeholders::_1)); if (buffer.same_as(buffer_region->buffer) && region.same_as(buffer_region->region)) { diff --git a/src/s_tir/transform/storage_access.cc b/src/s_tir/transform/storage_access.cc index ce6bfa62efad..57b3ce33cbdd 100644 --- a/src/s_tir/transform/storage_access.cc +++ b/src/s_tir/transform/storage_access.cc @@ -36,8 +36,8 @@ namespace s_tir { using namespace tvm::tirx; void StorageAccessVisitor::VisitExpr_(const BufferLoadNode* op) { - Var buf = op->buffer->data; - StorageScope scope = GetScope(buf); + Var buf = op->buffer.var(); + StorageScope scope = StorageScope::Create(op->buffer.scope()); if (Enabled(buf.get(), scope)) { TVM_FFI_ICHECK(allow_append_) << op << " " << scope.to_string(); AccessEntry e; @@ -60,8 +60,8 @@ void StorageAccessVisitor::VisitStmt_(const BufferStoreNode* op) { TVM_FFI_ICHECK_EQ(curr_stmt_.access.size(), 0U); curr_stmt_.stmt = op; - Var buf = op->buffer->data; - StorageScope scope = GetScope(buf); + Var buf = op->buffer.var(); + StorageScope scope = StorageScope::Create(op->buffer.scope()); if (Enabled(buf.get(), scope)) { AccessEntry e; e.threads = env_threads(); diff --git a/src/s_tir/transform/tensorcore_infer_fragment.cc b/src/s_tir/transform/tensorcore_infer_fragment.cc index 28b92a6e6285..a97d5a166882 100644 --- a/src/s_tir/transform/tensorcore_infer_fragment.cc +++ b/src/s_tir/transform/tensorcore_infer_fragment.cc @@ -185,16 +185,16 @@ class InferFragmenter : public StmtMutator { Stmt VisitStmt_(const AllocBufferNode* op) final { Stmt stmt = StmtMutator::VisitStmt_(op); - const VarNode* buffer = op->buffer->data.get(); + const VarNode* buffer = op->buffer.get(); if (fragment_getter.fragments.count(buffer)) { FragmentInfo info = fragment_getter.fragments.at(buffer); std::string shape = std::to_string(info.m) + ", " + std::to_string(info.n) + ", " + std::to_string(info.k); PrimExpr shape_expr = StringImm(shape); - Stmt shape_attr = AttrStmt(op->buffer->data, s_tir::attr::fragment_shape, shape_expr, stmt); + Stmt shape_attr = AttrStmt(op->buffer.var(), s_tir::attr::fragment_shape, shape_expr, stmt); if (info.layout != "") { - Stmt layout_attr = AttrStmt(op->buffer->data, s_tir::attr::fragment_layout, + Stmt layout_attr = AttrStmt(op->buffer.var(), s_tir::attr::fragment_layout, StringImm(info.layout), shape_attr); return layout_attr; } else { diff --git a/src/s_tir/transform/transform_mma_buffer_layout.cc b/src/s_tir/transform/transform_mma_buffer_layout.cc index ee60179657ea..6ce119f40572 100644 --- a/src/s_tir/transform/transform_mma_buffer_layout.cc +++ b/src/s_tir/transform/transform_mma_buffer_layout.cc @@ -48,7 +48,7 @@ class MmaBufferLayoutTransformer : public StmtExprMutator { Stmt VisitStmt_(const SBlockNode* op) { SBlock block = ffi::GetRef(op); auto* n = block.CopyOnWrite(); - auto fmutate = [this](const Buffer& buffer) { + auto fmutate = [this](const BufferVar& buffer) { // m16n8k8.matrix[A/B/C] buffers are composed ofseveral small blocks. Assume the block's // shape is [bi, bj]. Inside each small block, we have 8 threads in stride dimension and 4 // threads in contiguous dimension, so we change the buffer's shape from [i, j] @@ -70,9 +70,9 @@ class MmaBufferLayoutTransformer : public StmtExprMutator { new_shape.insert(new_shape.end(), {IntImm::Int32(dim0->value / 16), IntImm::Int32(dim1->value / 8), 2, 2}); - Buffer new_buffer = decl_buffer(std::move(new_shape), buffer->dtype, buffer->name, "local"); + BufferVar new_buffer = decl_buffer(std::move(new_shape), buffer->dtype, buffer.name(), "local"); this->buffer_map_.insert({buffer, new_buffer}); - this->buffer_var_map_.insert({buffer->data, new_buffer->data}); + this->buffer_var_map_.insert({buffer.var(), new_buffer.var()}); return new_buffer; } else if (buffer.scope() == "m16n8k8.matrixA") { @@ -91,9 +91,9 @@ class MmaBufferLayoutTransformer : public StmtExprMutator { new_shape.insert(new_shape.end(), {IntImm::Int32(dim0->value / 32), IntImm::Int32(dim1->value / 8), 4, 2}); - Buffer new_buffer = decl_buffer(std::move(new_shape), buffer->dtype, buffer->name, "local"); + BufferVar new_buffer = decl_buffer(std::move(new_shape), buffer->dtype, buffer.name(), "local"); this->buffer_map_.insert({buffer, new_buffer}); - this->buffer_var_map_.insert({buffer->data, new_buffer->data}); + this->buffer_var_map_.insert({buffer.var(), new_buffer.var()}); return new_buffer; } else if (buffer.scope() == "m16n8k8.matrixB") { @@ -112,9 +112,9 @@ class MmaBufferLayoutTransformer : public StmtExprMutator { new_shape.insert(new_shape.end(), {IntImm::Int32(dim0->value / 8), IntImm::Int32(dim1->value / 32), 1, 8}); - Buffer new_buffer = decl_buffer(std::move(new_shape), buffer->dtype, buffer->name, "local"); + BufferVar new_buffer = decl_buffer(std::move(new_shape), buffer->dtype, buffer.name(), "local"); this->buffer_map_.insert({buffer, new_buffer}); - this->buffer_var_map_.insert({buffer->data, new_buffer->data}); + this->buffer_var_map_.insert({buffer.var(), new_buffer.var()}); return new_buffer; } return buffer; @@ -170,7 +170,7 @@ class MmaBufferLayoutTransformer : public StmtExprMutator { } private: - std::unordered_map buffer_map_; + std::unordered_map buffer_map_; std::unordered_map buffer_var_map_; arith::Analyzer analyzer; }; 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 2a1e6b443c59..7e0a2f4edfc5 100644 --- a/src/s_tir/transform/using_assume_to_reduce_branches.cc +++ b/src/s_tir/transform/using_assume_to_reduce_branches.cc @@ -45,7 +45,7 @@ #include #include -#include +#include #include "../../arith/constraint_extract.h" #include "../../arith/ir_mutator_with_analyzer.h" @@ -135,8 +135,10 @@ class ParseAssumeAndOvercompute : public IRMutatorWithAnalyzer { std::vector conditions_; // Storing all the buffer assumptions data in map - std::map map_buffer_assumption; - tirx::Buffer current_bufferstorenode_name; + std::unordered_map + map_buffer_assumption; + tirx::BufferVar current_bufferstorenode_name; struct InternalConstraintContext { /* This stuct appends the constraint passed to it in the conditions list. @@ -350,7 +352,7 @@ class ParseAssumeAndOvercompute : public IRMutatorWithAnalyzer { auto has_side_effect = tirx::SideEffect(value) > tirx::CallEffectKind::kPure; TVM_FFI_ICHECK(!has_side_effect) - << "Buffer value in constraint must be pure expression, but was " << value; + << "BufferVar value in constraint must be pure expression, but was " << value; if (has_side_effect) { return; } diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc index 97bd1b0f2644..6da359bcbed3 100644 --- a/src/target/llvm/codegen_llvm.cc +++ b/src/target/llvm/codegen_llvm.cc @@ -109,7 +109,7 @@ PrimType WithScalableVScaleFactor(const PrimType& dtype, int vscale_factor) { return PrimType::ScalableVector(dtype.code(), dtype.bits(), vscale_factor); } -// Underlying access type for a Buffer: bool is backed by int8 so vectorized +// Underlying access type for a BufferVar: bool is backed by int8 so vectorized // accesses lower to real loads/stores instead of i1 predicate registers. PrimType BufferAccessType(const PrimType& dtype) { if (!dtype.MatchesCode(DLDataTypeCode::kDLBool)) return dtype; @@ -1007,7 +1007,7 @@ CodeGenLLVM::TypedPointer CodeGenLLVM::CreateBufferPtr(llvm::Value* buffer_ptr, "no padding for alignment."; } else { TVM_FFI_ICHECK(buffer_element_type.as()) - << "Buffer elements must have primitive or pointer type, but got " << buffer_element_type; + << "BufferVar elements must have primitive or pointer type, but got " << buffer_element_type; } llvm::Value* value_ptr = builder_->CreateInBoundsGEP(llvm_element_type, buffer_ptr, index); @@ -1415,7 +1415,7 @@ llvm::Value* CodeGenLLVM::CreateIntrinsic(const CallNode* op) { } TypedPointer buffer_ptr = - CreateBufferPtr(MakeValue(load->buffer->data), load->buffer->dtype, indices_val, + CreateBufferPtr(MakeValue(load->buffer.var()), load->buffer->dtype, indices_val, PrimType(load->ty.as_or_throw()->dtype)); return buffer_ptr.addr; } else if (op->op.same_as(builtin::reinterpret()) && args[0].as() && @@ -1724,7 +1724,7 @@ bool CodeGenLLVM::HasAlignmentPadding(PrimType dtype) { } void CodeGenLLVM::BufferAccessHelper( - Buffer buffer, ffi::Array indices, ffi::Optional predicate, + BufferVar buffer, ffi::Array indices, ffi::Optional predicate, PrimType value_dtype, std::function @@ -1732,7 +1732,7 @@ void CodeGenLLVM::BufferAccessHelper( PrimType buffer_element_dtype = BufferAccessType(buffer->dtype); TVM_FFI_ICHECK_GE(indices.size(), 1) - << "Buffer " << buffer->name << " is accessed with no indices. " + << "BufferVar " << buffer.name() << " is accessed with no indices. " << "0-d scalar buffers are expected to be flattened to 1-d buffers prior to codegen."; // Only the last index is allowed to be multi-lane. All earlier @@ -1742,7 +1742,7 @@ void CodeGenLLVM::BufferAccessHelper( std::vector earlier_index_values; for (size_t i = 0; i < indices.size() - 1; i++) { TVM_FFI_ICHECK_EQ(PrimType(indices[i].ty()->dtype).lanes(), 1) - << "Buffer " << buffer->name << " is accessed with a multi-lane index at position " << i + << "BufferVar " << buffer.name() << " is accessed with a multi-lane index at position " << i << ". Multi-lane indices are only supported as the last index."; earlier_index_values.push_back(MakeValue(indices[i])); } @@ -1756,7 +1756,7 @@ void CodeGenLLVM::BufferAccessHelper( PrimExpr last_index_origin = last_index; PrimType buffer_element_dtype_origin = buffer_element_dtype; - bool is_volatile = volatile_buf_.count(buffer->data.get()); + bool is_volatile = volatile_buf_.count(buffer.get()); // If the buffer index is a contiguous ramp node, we only need to // access the first element, then cast to the value type. @@ -1784,7 +1784,7 @@ void CodeGenLLVM::BufferAccessHelper( // element being accessed may require more alignment than the // underlying data type. int native_bits; - GetAlignment(value_dtype, buffer->data.get(), last_index, &alignment, &native_bits); + GetAlignment(value_dtype, buffer.get(), last_index, &alignment, &native_bits); } else { // Otherwise, alignment is based on the return value's scalar // type. @@ -1821,14 +1821,14 @@ void CodeGenLLVM::BufferAccessHelper( TypedPointer buffer_ptr = value_dtype.IsScalableVector() - ? CreateBufferPtr(MakeValue(buffer->data), buffer_element_dtype, all_index_values, + ? CreateBufferPtr(MakeValue(buffer.var()), buffer_element_dtype, all_index_values, WithScalableVScaleFactor( value_dtype, value_dtype.VScaleFactor() / last_index_lanes)) - : CreateBufferPtr(MakeValue(buffer->data), buffer_element_dtype, all_index_values, + : CreateBufferPtr(MakeValue(buffer.var()), buffer_element_dtype, all_index_values, value_dtype.WithLanes(value_dtype.lanes() / last_index_lanes)); auto instruction = make_instruction(buffer_ptr, subelement_i, predicate_value, alignment, is_volatile); - AddAliasInfo(instruction, buffer->data.get(), last_index_origin, buffer_element_dtype_origin); + AddAliasInfo(instruction, buffer.get(), last_index_origin, buffer_element_dtype_origin); } } @@ -1877,6 +1877,10 @@ llvm::Value* CodeGenLLVM::VisitExpr_(const BufferLoadNode* op) { llvm::Value* CodeGenLLVM::VisitExpr_(const CallNode* op) { const ffi::Array& args = op->args; + if (op->op.same_as(builtin::buffer_data())) { + TVM_FFI_ICHECK_EQ(args.size(), 1U); + return MakeValue(args[0]); + } if (auto opt_call_op = op->op.as()) { auto call_op = opt_call_op.value(); if (op->op.same_as(builtin_call_extern_) || op->op.same_as(builtin_call_pure_extern_)) { @@ -1986,7 +1990,7 @@ llvm::Value* CodeGenLLVM::VisitExpr_(const BroadcastNode* op) { void CodeGenLLVM::VisitStmt_(const BufferStoreNode* op) { EmitDebugLocation(op); PrimType value_dtype = PrimType(op->value.ty()->dtype); - Var buffer_var = op->buffer->data; + Var buffer_var = op->buffer.var(); llvm::Value* value = MakeValue(op->value); @@ -2099,7 +2103,7 @@ void CodeGenLLVM::VisitStmt_(const AllocBufferNode* op) { EmitDebugLocation(op); TVM_FFI_ICHECK_EQ(op->buffer->shape.size(), 1) << "LLVM codegen only supports flat 1-d buffer allocation, but allocation of " - << op->buffer->name << " is " << op->buffer->shape << "-d"; + << op->buffer.name() << " is " << op->buffer->shape << "-d"; llvm::Value* buf = nullptr; @@ -2108,7 +2112,7 @@ void CodeGenLLVM::VisitStmt_(const AllocBufferNode* op) { int32_t constant_size = static_cast(dim_imm->value); TVM_FFI_ICHECK_GT(constant_size, 0) << "Can only handle constant size stack allocation"; - StorageInfo& info = alloc_storage_info_[op->buffer->data.get()]; + StorageInfo& info = alloc_storage_info_[op->buffer.get()]; // Use buffer's data_alignment if specified, otherwise compute from shape. if (op->buffer->data_alignment > 0) { info.alignment = op->buffer->data_alignment; @@ -2133,12 +2137,12 @@ void CodeGenLLVM::VisitStmt_(const AllocBufferNode* op) { buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(DTypeToLLVMType(op->buffer->dtype), buf->getType()->getPointerAddressSpace())); - AddDebugInformation(buf, op->buffer->data); + AddDebugInformation(buf, op->buffer.var()); - TVM_FFI_ICHECK(!var_map_.count(op->buffer->data.get())); - var_map_[op->buffer->data.get()] = buf; + TVM_FFI_ICHECK(!var_map_.count(op->buffer.get())); + var_map_[op->buffer.get()] = buf; if (op->annotations.count(tirx::attr::kVolatile)) { - volatile_buf_.insert(op->buffer->data.get()); + volatile_buf_.insert(op->buffer.get()); } } diff --git a/src/target/llvm/codegen_llvm.h b/src/target/llvm/codegen_llvm.h index 698e2571e9b4..41b237ad4326 100644 --- a/src/target/llvm/codegen_llvm.h +++ b/src/target/llvm/codegen_llvm.h @@ -358,7 +358,7 @@ class CodeGenLLVM : public ExprFunctor, * - Should return the generated expression. */ void BufferAccessHelper( - Buffer buffer, ffi::Array indices, ffi::Optional predicate, + BufferVar buffer, ffi::Array indices, ffi::Optional predicate, PrimType value_dtype, std::function diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc index 6f9c25c592f2..a813d4fb5e34 100644 --- a/src/target/source/codegen_c.cc +++ b/src/target/source/codegen_c.cc @@ -237,8 +237,8 @@ void CodeGenC::PrintSSAAssign(const std::string& target, const std::string& src, } // Print a reference expression to a buffer. -std::string CodeGenC::GetBufferRef(const PrimType& t, const BufferNode* buffer, PrimExpr index) { - const VarNode* buffer_var = buffer->data.get(); +std::string CodeGenC::GetBufferRef(const PrimType& t, const VarNode* buffer, PrimExpr index) { + const VarNode* buffer_var = buffer; std::ostringstream os; std::string vid = GetVarID(buffer_var); std::string scope; @@ -261,7 +261,7 @@ std::string CodeGenC::GetBufferRef(const PrimType& t, const BufferNode* buffer, return ptr_os.str(); }; - const PrimType& buffer_element_dtype = buffer->dtype; + const PrimType& buffer_element_dtype = GetBufferVar(buffer)->dtype; std::string buffer_str = vid; if (!HandleTypeMatch(buffer_var, buffer_element_dtype) || is_vol) { @@ -437,11 +437,11 @@ void CodeGenC::PrintVecElemStore(const std::string& vec, const PrimType& t, int stream << vec << ".s" << std::hex << i << " = " << value << ";\n" << std::dec; } -std::string CodeGenC::GetVecLoad(const PrimType& t, const BufferNode* buffer, PrimExpr base) { +std::string CodeGenC::GetVecLoad(const PrimType& t, const VarNode* buffer, PrimExpr base) { return GetBufferRef(t, buffer, base); } -void CodeGenC::PrintVecStore(const BufferNode* buffer, const PrimType& t, PrimExpr base, +void CodeGenC::PrintVecStore(const VarNode* buffer, const PrimType& t, PrimExpr base, const std::string& value) { std::string ref = GetBufferRef(t, buffer, base); this->PrintIndent(); @@ -673,7 +673,13 @@ void CodeGenC::VisitExpr_(const CallNode* op, std::ostream& os) { // NOLINT(*) if (auto opt_call_op = op->op.as()) { auto call_op = opt_call_op.value(); - if (op->op.same_as(builtin::continue_loop())) { + if (op->op.same_as(builtin::buffer_data())) { + TVM_FFI_ICHECK_EQ(op->args.size(), 1U); + const auto* buffer = op->args[0].as(); + TVM_FFI_ICHECK(buffer && buffer->ty.as()) + << "buffer_data expects a Var with BufferType"; + os << GetVarID(buffer); + } else if (op->op.same_as(builtin::continue_loop())) { os << "continue;"; } else if (op->op.same_as(builtin::break_loop())) { os << "break;"; @@ -758,7 +764,7 @@ void CodeGenC::VisitExpr_(const CallNode* op, std::ostream& os) { // NOLINT(*) if (load) { TVM_FFI_ICHECK_EQ(load->indices.size(), 1) << "CodeGenC only supports flat memory allocations."; - const VarNode* data = load->buffer->data.get(); + const VarNode* data = load->buffer.get(); if (pointer_offset_vars_.count(data) && HandleTypeMatch(data, load->buffer->dtype) && !IsVolatile(data)) { os << "(" << GetVarID(data) << " + "; @@ -899,7 +905,15 @@ void CodeGenC::PrintVecBinaryOp(const std::string& op, const PrimType& t, PrimEx } void CodeGenC::VisitStmt_(const DeclBufferNode* op) { - // DeclBuffer is a flat statement with no body — nothing to emit. + if (!op->data.has_value()) { + return; + } + this->PrintIndent(); + PrintType(op->buffer->data_pointer_type, stream); + stream << ' ' << AllocVarID(op->buffer.get()) << " = "; + PrintExpr(op->data.value(), stream); + stream << ";\n"; + RegisterHandleType(op->buffer.get(), op->buffer->dtype); } void CodeGenC::VisitExpr_(const BufferLoadNode* op, std::ostream& os) { // NOLINT(*) @@ -908,7 +922,7 @@ void CodeGenC::VisitExpr_(const BufferLoadNode* op, std::ostream& os) { // NOLI PrimType value_ty = op->ty.as_or_throw(); PrimExpr index = op->indices[0]; - Var buffer_var = op->buffer->data; + Var buffer_var = op->buffer.var(); const PrimType& element_ty = op->buffer->dtype; int lanes = value_ty.lanes(); @@ -984,7 +998,7 @@ void CodeGenC::VisitStmt_(const BufferStoreNode* op) { PrimType value_ty = op->value.ty(); const PrimType& element_ty = op->buffer->dtype; PrimExpr index_expr = op->indices[0]; - Var buffer_var = op->buffer->data; + Var buffer_var = op->buffer.var(); if (value_ty.lanes() == element_ty.lanes()) { std::string value = this->PrintExpr(op->value); @@ -1190,7 +1204,7 @@ void CodeGenC::VisitStmt_(const BindNode* op) { void CodeGenC::VisitStmt_(const AllocBufferNode* op) { TVM_FFI_ICHECK(op->buffer.defined()); - std::string vid = AllocVarID(op->buffer->data.get()); + std::string vid = AllocVarID(op->buffer.get()); this->PrintIndent(); const auto& shape = op->buffer->shape; @@ -1202,16 +1216,16 @@ void CodeGenC::VisitStmt_(const AllocBufferNode* op) { } TVM_FFI_ICHECK_GT(constant_size, 0) << "Can only handle constant size stack allocation for now"; - auto scope = GetPtrStorageScope(op->buffer->data); - alloc_storage_scope_[op->buffer->data.get()] = scope; + auto scope = op->buffer.scope(); + alloc_storage_scope_[op->buffer.get()] = scope; PrintStorageScope(scope, stream); PrintType(op->buffer->dtype, stream); stream << ' ' << vid << '[' << constant_size << "];\n"; - RegisterHandleType(op->buffer->data.get(), op->buffer->dtype); + RegisterHandleType(op->buffer.get(), op->buffer->dtype); if (op->annotations.count(tirx::attr::kVolatile)) { - MarkVolatile(op->buffer->data.get()); + MarkVolatile(op->buffer.get()); } } diff --git a/src/target/source/codegen_c.h b/src/target/source/codegen_c.h index 02fdc563919a..354800756162 100644 --- a/src/target/source/codegen_c.h +++ b/src/target/source/codegen_c.h @@ -220,9 +220,9 @@ class CodeGenC : public ExprFunctor, PrimExpr rhs, std::ostream& os); // NOLINT(*) // print vector load - virtual std::string GetVecLoad(const PrimType& t, const BufferNode* buffer, PrimExpr base); + virtual std::string GetVecLoad(const PrimType& t, const VarNode* buffer, PrimExpr base); // print vector store - virtual void PrintVecStore(const BufferNode* buffer, const PrimType& t, PrimExpr base, + virtual void PrintVecStore(const VarNode* buffer, const PrimType& t, PrimExpr base, const std::string& value); // NOLINT(*) // print load of single element virtual void PrintVecElemLoad(const std::string& vec, const PrimType& t, int i, @@ -250,7 +250,7 @@ class CodeGenC : public ExprFunctor, // Print reference to struct location std::string GetStructRef(const Type& t, const Expr& buffer, const PrimExpr& index, int kind); // Print reference to a buffer as type t in index. - virtual std::string GetBufferRef(const PrimType& t, const BufferNode* buffer, PrimExpr index); + virtual std::string GetBufferRef(const PrimType& t, const VarNode* buffer, PrimExpr index); /*! * \brief Handle volatile loads. diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc index ec485da36345..41b28c116d4f 100644 --- a/src/te/operation/create_primfunc.cc +++ b/src/te/operation/create_primfunc.cc @@ -48,7 +48,7 @@ namespace tirx { /*! \brief The helper mutator that transforms ProducerLoad to BufferLoad */ class ProducerToBufferTransformer : public StmtExprMutator { public: - explicit ProducerToBufferTransformer(const std::unordered_map& tensor2buffers) + explicit ProducerToBufferTransformer(const std::unordered_map& tensor2buffers) : tensor2buffers_(tensor2buffers) {} Expr VisitExpr_(const ProducerLoadNode* op) final { @@ -56,20 +56,20 @@ class ProducerToBufferTransformer : public StmtExprMutator { te::Tensor tensor = visited_op->producer.as_or_throw(); auto it = tensor2buffers_.find(tensor); TVM_FFI_ICHECK(it != tensor2buffers_.end()) << "IndexError: Cannot find the tensor " << tensor; - const Buffer& buffer = it->second; + const BufferVar& buffer = it->second; return BufferLoad(buffer, visited_op->indices); } private: /*! \brief The Map from Operations to buffers */ - const std::unordered_map& tensor2buffers_; + const std::unordered_map& tensor2buffers_; }; /*! \brief The helper mutator to rewrite buffer and buffer var accessed by block body */ class BufferSubstituter : public StmtExprMutator { public: explicit BufferSubstituter(const std::unordered_map& var_map, - const std::unordered_map& buffer_map) + const std::unordered_map& buffer_map) : var_map_(var_map), buffer_map_(buffer_map) {} Expr VisitExpr_(const VarNode* op) final { @@ -100,7 +100,7 @@ class BufferSubstituter : public StmtExprMutator { private: const std::unordered_map& var_map_; - const std::unordered_map& buffer_map_; + const std::unordered_map& buffer_map_; }; /*! \brief Helper data structure to store information. */ @@ -108,11 +108,11 @@ struct CreateFuncInfo { /*! \brief The Tensor arg_list. */ ffi::Array arg_list; /*! \brief The map from each Tensor to its corresponding buffer. */ - std::unordered_map tensor2buffers; + std::unordered_map tensor2buffers; /*! \brief The transformer from ProducerLoad to BufferLoad. */ ProducerToBufferTransformer transformer; /*! \brief The buffers should be allocated at function root. */ - ffi::Array root_alloc; + ffi::Array root_alloc; /*! \brief The unique name supply to make block name unique. */ UniqueNameSupply name_supply; @@ -132,7 +132,7 @@ class LayoutFreePlaceholdersNormalizer : public StmtMutator { PrimFunc Process(PrimFunc func) { for (int i = 0, n = func->params.size(); i < n; ++i) { if (auto v = func->params[i].as()) { - if (ffi::Optional buffer = func->buffer_map.Get(v.value())) { + if (ffi::Optional buffer = func->buffer_map.Get(v.value())) { buffer2index_[buffer.value()] = i; } } @@ -154,8 +154,8 @@ class LayoutFreePlaceholdersNormalizer : public StmtMutator { SBlock block = StmtMutator::VisitStmt_(_block).as_or_throw(); SBlockNode* n = block.CopyOnWrite(); if (auto opt_ann = n->annotations.Get(topi_attr)) { - ffi::Array new_buffers; - for (Buffer buffer : opt_ann.value().as_or_throw>()) { + ffi::Array new_buffers; + for (BufferVar buffer : opt_ann.value().as_or_throw>()) { auto it = buffer2index_.find(buffer); if (it != buffer2index_.end()) { layout_free_buffer_indices_.insert(it->second); @@ -178,7 +178,7 @@ class LayoutFreePlaceholdersNormalizer : public StmtMutator { return block; } - std::unordered_map buffer2index_; + std::unordered_map buffer2index_; std::set layout_free_buffer_indices_; ffi::String topi_attr = "layout_free_placeholders"; std::vector blocklist = {"const_matrix", @@ -246,7 +246,7 @@ NestedIterLevels GenerateNestedIterLevels(const ffi::Array& axes, * \param info Generation context info. * \returns The output buffer objects, ordered by compute op's outputs. **/ -ffi::Array GenerateOutputBuffers(const te::ComputeOp& compute_op, CreateFuncInfo* info) { +ffi::Array GenerateOutputBuffers(const te::ComputeOp& compute_op, CreateFuncInfo* info) { // Step 1. Collect output tensors in TE operation. ffi::Array tensors; if (compute_op->body[0]->IsInstance()) { @@ -280,9 +280,9 @@ ffi::Array GenerateOutputBuffers(const te::ComputeOp& compute_op, Create // - Declare buffers // - Update `op2buffers` // - Add the non-argument tensors to `alloc_buffer` of the root block - ffi::Array buffers; + ffi::Array buffers; for (const te::Tensor& tensor : tensors) { - Buffer buffer = decl_buffer(tensor->shape, tensor->dtype, tensor->GetNameHint(), "global"); + BufferVar buffer = decl_buffer(tensor->shape, tensor->dtype, tensor->GetNameHint(), "global"); info->tensor2buffers[tensor] = buffer; buffers.push_back(buffer); if (!info->IsArg(tensor)) { @@ -333,7 +333,7 @@ ffi::Map GenerateBlockAnnotations(const te::ComputeOp& co * \param info Generation context info. * \returns Init stmt. **/ -Stmt GenerateInitStmt(const ffi::Array& indices, const ffi::Array& buffers, +Stmt GenerateInitStmt(const ffi::Array& indices, const ffi::Array& buffers, const ReduceNode* reduce, const ffi::Map& var_map, CreateFuncInfo* info) { // helper to transform the expr and remap iters to the block domain @@ -346,7 +346,7 @@ Stmt GenerateInitStmt(const ffi::Array& indices, const ffi::Array init_stmts; init_stmts.reserve(n_buffers); for (int i = 0; i < n_buffers; ++i) { - const Buffer& buffer = buffers[i]; + const BufferVar& buffer = buffers[i]; PrimExpr identity = f_transform_and_remap(reduce->combiner->identity_element[i]); init_stmts.push_back(BufferStore(buffer, identity, indices)); } @@ -363,7 +363,7 @@ Stmt GenerateInitStmt(const ffi::Array& indices, const ffi::Array& indices, const ffi::Array& buffers, +Stmt GenerateBodyStmt(const ffi::Array& indices, const ffi::Array& buffers, const ffi::Map& var_map, PrimExpr expr_body, CreateFuncInfo* info, arith::AnalyzerObj* analyzer) { // helper to transform the expr and remap iters to the block domain @@ -402,10 +402,10 @@ Stmt GenerateBodyStmt(const ffi::Array& indices, const ffi::Array 1) { - temp_vars.push_back(Var("v_" + buffer->name, lhs[i].ty())); + temp_vars.push_back(Var("v_" + buffer.name(), lhs[i].ty())); value = temp_vars.back().as_or_throw(); } else { PrimExpr combined = reduce->combiner.get()->operator()(lhs, rhs)[i]; @@ -535,7 +535,7 @@ Stmt GenerateStmtFromCompute(const te::ComputeOp& compute_op, CreateFuncInfo* in } // Step 3. Generate output buffers for each output tensor - ffi::Array buffers = GenerateOutputBuffers(compute_op, info); + ffi::Array buffers = GenerateOutputBuffers(compute_op, info); // Step 4. Generate leaf block stmts. ffi::Array seq_stmt; @@ -575,7 +575,7 @@ Stmt GenerateStmtFromCompute(const te::ComputeOp& compute_op, CreateFuncInfo* in SBlock(/*iter_vars=*/leaf.block_iters, /*reads=*/{}, /*writes=*/{}, - /*name_hint=*/info->FreshName(buffers[i]->name), + /*name_hint=*/info->FreshName(buffers[i].name()), /*body=*/body, /*init=*/std::nullopt, /*alloc_buffers=*/{}, @@ -626,29 +626,31 @@ Stmt GenerateStmtFromCompute(const te::ComputeOp& compute_op, CreateFuncInfo* in Stmt GenerateStmtFromExternOp(const te::ExternOp& extern_op, CreateFuncInfo* info) { // Step 1. Check all inputs are visited before and update var_map. std::unordered_map var_map; - std::unordered_map input_buffer_map; + std::unordered_map input_buffer_map; TVM_FFI_ICHECK_EQ(extern_op->inputs.size(), extern_op->input_placeholders.size()); for (size_t i = 0; i < extern_op->inputs.size(); ++i) { - const Buffer& placeholder = extern_op->input_placeholders[i]; + const BufferVar& placeholder = extern_op->input_placeholders[i]; const te::Tensor& input_tensor = extern_op->inputs[i]; auto it = info->tensor2buffers.find(input_tensor); TVM_FFI_ICHECK(it != info->tensor2buffers.end()); - var_map[placeholder->data.get()] = it->second->data; + var_map[placeholder.get()] = it->second.var(); input_buffer_map[placeholder.get()] = it->second; } // Step 2. Update info with its output tensor and placeholder buffer. TVM_FFI_ICHECK_EQ(extern_op->num_outputs(), extern_op->output_placeholders.size()); for (int i = 0; i < extern_op->num_outputs(); ++i) { - const Buffer& placeholder = extern_op->output_placeholders[i]; + const BufferVar& placeholder = extern_op->output_placeholders[i]; const te::Tensor& output_tensor = extern_op.output(i); - Buffer output_buffer = placeholder; + BufferVar output_buffer = placeholder; if (!info->IsArg(output_tensor)) { PrimExpr zero_offset = IntImm(placeholder->elem_offset.ty(), 0); if (auto offset_var = placeholder->elem_offset.as()) { var_map[offset_var.value().get()] = zero_offset; } - output_buffer.CopyOnWrite()->elem_offset = zero_offset; + ffi::ObjectPtr type = CopyBufferType(output_buffer); + type->elem_offset = zero_offset; + output_buffer = RebuildBufferVar(output_buffer, std::move(type)); input_buffer_map[placeholder.get()] = output_buffer; info->root_alloc.push_back(output_buffer); } @@ -709,7 +711,7 @@ void InitializeBufferBinds(const ffi::Array& ordered_ops, CreateF TVM_FFI_ICHECK_EQ(extern_op->inputs.size(), extern_op->input_placeholders.size()); for (size_t i = 0; i < extern_op->inputs.size(); ++i) { const te::Tensor& input = extern_op->inputs[i]; - const Buffer& buffer = extern_op->input_placeholders[i]; + const BufferVar& buffer = extern_op->input_placeholders[i]; info->tensor2buffers[input] = buffer; } } @@ -730,7 +732,7 @@ void RewriteStageToBlock(const te::Operation& op, CreateFuncInfo* info, // Declare a buffer for any argument tensors without a pre-existing // buffer declaration recorded in the tensor2buffer binds map if (info->tensor2buffers.count(tensor) == 0) { - const Buffer& buffer = + const BufferVar& buffer = decl_buffer(placeholder->shape, placeholder->dtype, placeholder->name, "global"); info->tensor2buffers[tensor] = buffer; } @@ -749,7 +751,7 @@ void RewriteStageToBlock(const te::Operation& op, CreateFuncInfo* info, PrimFunc GenerateAndCompletePrimFunc(const ffi::Array& arg_list, const ffi::Array& root_stmts, CreateFuncInfo* info) { ffi::Array parameters; - ffi::Map buffer_map; + ffi::Map buffer_map; for (const te::Tensor& tensor : arg_list) { auto it = info->tensor2buffers.find(tensor); TVM_FFI_ICHECK(it != info->tensor2buffers.end()); @@ -815,7 +817,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { PrimFunc GenerateAndCompletePrimFunc(const ffi::Array& arg_tir_var_list, const ffi::Array& root_stmts, CreateFuncInfo* info) { ffi::Array parameters; - ffi::Map buffer_map; + ffi::Map buffer_map; for (const ffi::ObjectRef& arg : arg_tir_var_list) { if (auto opt_tensor = arg.as()) { te::Tensor tensor = opt_tensor.value(); diff --git a/src/te/operation/extern_op.cc b/src/te/operation/extern_op.cc index 6fbaf4482b5c..7008724969b6 100644 --- a/src/te/operation/extern_op.cc +++ b/src/te/operation/extern_op.cc @@ -46,8 +46,8 @@ ffi::Array ExternOpNode::output_shape(size_t i) const { } ExternOp::ExternOp(std::string name, std::string tag, ffi::Map attrs, - ffi::Array inputs, ffi::Array input_placeholders, - ffi::Array output_placeholders, Stmt body) { + ffi::Array inputs, ffi::Array input_placeholders, + ffi::Array output_placeholders, Stmt body) { if (!attrs.defined()) { attrs = ffi::Map(); } @@ -76,8 +76,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def( "te.ExternOp", [](std::string name, std::string tag, ffi::Optional> attrs, - ffi::Array inputs, ffi::Array input_placeholders, - ffi::Array output_placeholders, Stmt body) { + ffi::Array inputs, ffi::Array input_placeholders, + ffi::Array output_placeholders, Stmt body) { return ExternOp(name, tag, attrs.value_or({}), inputs, input_placeholders, output_placeholders, body); }); diff --git a/src/tirx/analysis/var_touch.cc b/src/tirx/analysis/var_touch.cc index 2e918d8f14ca..c2e6a2850ec4 100644 --- a/src/tirx/analysis/var_touch.cc +++ b/src/tirx/analysis/var_touch.cc @@ -45,12 +45,12 @@ class VarTouchVisitor : public StmtExprVisitor { void VisitExpr_(const VarNode* op) final { Handle(op); } void VisitStmt_(const BufferStoreNode* op) final { - Handle(op->buffer->data.get()); + Handle(op->buffer.get()); StmtVisitor::VisitStmt_(op); } void VisitExpr_(const BufferLoadNode* op) final { - Handle(op->buffer->data.get()); + Handle(op->buffer.get()); ExprVisitor::VisitExpr_(op); } diff --git a/src/tirx/analysis/var_use_def_analysis.cc b/src/tirx/analysis/var_use_def_analysis.cc index 12e5e451a6e2..6d5c860f9fd3 100644 --- a/src/tirx/analysis/var_use_def_analysis.cc +++ b/src/tirx/analysis/var_use_def_analysis.cc @@ -66,7 +66,8 @@ void VarUseDefAnalyzer::VisitStmt_(const ForNode* op) { } void VarUseDefAnalyzer::VisitStmt_(const AllocBufferNode* op) { - // VisitBufferDef (called by base) defines buffer->data and the buffer itself. + // VisitBufferDef (called by base) defines the typed buffer Var and visits + // its dependent BufferType expressions. StmtExprVisitor::VisitStmt_(op); } @@ -101,42 +102,30 @@ void VarUseDefAnalyzer::VisitExpr_(const ReduceNode* op) { StmtExprVisitor::VisitExpr_(op); } -void VarUseDefAnalyzer::VisitBufferDef(const Buffer& buffer, bool alloc_data) { - if (alloc_data) { - // AllocBuffer / SBlock: data is a new allocation — define it. - if (!use_count_.count(buffer->data.get())) { - HandleDef(buffer->data); - } - } else { - // DeclBuffer: data references an existing variable — use it. - // TMEM DeclBuffer data vars are internal lowering symbols and should - // not become external free vars in host packed-api generation. - if (buffer.scope() != "tmem") { - HandleUse(buffer->data); +void VarUseDefAnalyzer::VisitBufferDef(const BufferVar& buffer, bool alloc_data) { + bool is_first_buffer_definition = !buffer_def_count_.count(buffer.get()); + HandleDef(buffer); + if (is_first_buffer_definition) { + auto it = use_count_.find(buffer.get()); + if (it == use_count_.end()) { + HandleDef(buffer.var()); + } else { + TVM_FFI_ICHECK_GE(it->second, 0) + << "variable " << buffer.name() << " has been used before definition!"; } } - HandleDef(buffer); // Visit shape/strides/elem_offset as uses of vars from the enclosing scope. for (const auto& e : buffer->shape) this->VisitExpr(e); for (const auto& e : buffer->strides) this->VisitExpr(e); this->VisitExpr(buffer->elem_offset); } -void VarUseDefAnalyzer::VisitBufferUse(const Buffer& buffer) { +void VarUseDefAnalyzer::VisitBufferUse(const BufferVar& buffer) { HandleUse(buffer); - // Buffer data pointer must be tracked as a use — the use site - // reads/writes through this pointer. Without this, UndefinedVars - // misses data vars for buffers whose DeclBuffer is outside the scope. - HandleUse(buffer->data); + HandleUse(buffer.var()); } -void VarUseDefAnalyzer::VisitBuffer(const Buffer& buffer) { - // TMEM buffers can carry symbolic data vars that are internal to lowering - // and should not become external free vars during host/device splitting. - if (buffer.scope() != "tmem") { - this->HandleUse(buffer->data); - } - +void VarUseDefAnalyzer::VisitBuffer(const BufferVar& buffer) { auto visit_arr = [&](ffi::Array arr) { for (const auto& element : arr) { this->VisitExpr(element); @@ -170,10 +159,10 @@ void VarUseDefAnalyzer::HandleUse(const Var& var) { } } -void VarUseDefAnalyzer::HandleDef(const Buffer& buf) { +void VarUseDefAnalyzer::HandleDef(const BufferVar& buf) { auto ptr = buf.get(); // Some lowering pipelines may duplicate identical DeclBuffer nodes that - // reference the same Buffer object. Treat repeated definition of the same + // reference the same BufferVar object. Treat repeated definition of the same // buffer object as idempotent. if (buffer_def_count_.count(ptr)) { VisitBuffer(buf); @@ -183,11 +172,11 @@ void VarUseDefAnalyzer::HandleDef(const Buffer& buf) { << "buffer " << ptr->name << " has been used before definition!"; buffer_use_count_[ptr] = 0; buffer_def_count_[ptr] = 1; - // Buffer fields (data, shape, strides) are visited by the caller + // BufferVar fields (data, shape, strides) are visited by the caller // (VisitBufferDef) via the base class, not here. } -void VarUseDefAnalyzer::HandleUse(const Buffer& buf) { +void VarUseDefAnalyzer::HandleUse(const BufferVar& buf) { auto ptr = buf.get(); auto it = buffer_use_count_.find(ptr); if (it != buffer_use_count_.end()) { @@ -195,10 +184,10 @@ void VarUseDefAnalyzer::HandleUse(const Buffer& buf) { ++it->second; } } else { - undefined_buffers_.push_back(ffi::GetRef(ptr)); + undefined_buffers_.push_back(BufferVar(ffi::GetRef(ptr))); buffer_use_count_[ptr] = -1; } - // Buffer fields (shape, strides, data) are visited at the definition + // BufferVar fields (shape, strides, data) are visited at the definition // site via VisitBufferDef. Do not re-visit them at use sites, as the // buffer's shape variables may not be in scope at the point of use. } diff --git a/src/tirx/analysis/var_use_def_analysis.h b/src/tirx/analysis/var_use_def_analysis.h index b96887eb306f..293f5574a0d5 100644 --- a/src/tirx/analysis/var_use_def_analysis.h +++ b/src/tirx/analysis/var_use_def_analysis.h @@ -45,12 +45,12 @@ class VarUseDefAnalyzer : public StmtExprVisitor { // be accessible to the users. bool visit_thread_extent_{true}; ffi::Array undefined_; - ffi::Array undefined_buffers_; + ffi::Array undefined_buffers_; std::unordered_map use_count_; std::unordered_map def_count_; - std::unordered_map buffer_use_count_; - std::unordered_map buffer_def_count_; + std::unordered_map buffer_use_count_; + std::unordered_map buffer_def_count_; private: ExprDeepEqual deep_equal_; @@ -72,16 +72,16 @@ class VarUseDefAnalyzer : public StmtExprVisitor { // Piggyback on base class VisitBufferDef/VisitBufferUse to handle buffer // def/use tracking. Base class calls these from AllocBuffer, DeclBuffer, // BufferStore, BufferLoad, and SBlock visitors. - void VisitBufferDef(const Buffer& buffer, bool alloc_data) final; - void VisitBufferUse(const Buffer& buffer) final; + void VisitBufferDef(const BufferVar& buffer, bool alloc_data) final; + void VisitBufferUse(const BufferVar& buffer) final; void HandleDef(const Var& v); void HandleUse(const Var& v); - void HandleDef(const Buffer& buf); - void HandleUse(const Buffer& buf); + void HandleDef(const BufferVar& buf); + void HandleUse(const BufferVar& buf); - void VisitBuffer(const Buffer& buffer); + void VisitBuffer(const BufferVar& buffer); }; } // namespace tirx diff --git a/src/tirx/analysis/verify_memory.cc b/src/tirx/analysis/verify_memory.cc index fffccd1f15c1..2656bc521bd1 100644 --- a/src/tirx/analysis/verify_memory.cc +++ b/src/tirx/analysis/verify_memory.cc @@ -90,12 +90,12 @@ class MemoryAccessVerifier final : protected StmtExprVisitor { } void VisitExpr_(const BufferLoadNode* op) final { - HandleLoadStoreToVariable(op->buffer->data); + HandleLoadStoreToVariable(op->buffer.var()); return StmtExprVisitor::VisitExpr_(op); } void VisitStmt_(const BufferStoreNode* op) final { - HandleLoadStoreToVariable(op->buffer->data); + HandleLoadStoreToVariable(op->buffer.var()); return StmtExprVisitor::VisitStmt_(op); } //@} @@ -104,7 +104,7 @@ class MemoryAccessVerifier final : protected StmtExprVisitor { bool IsFromFunctionArgs(const VarNode* var) const { const VarNode* V = var; for (auto kv : func_->buffer_map) { - if (V == kv.second->data.get()) return true; + if (V == kv.second.get()) return true; } while (true) { diff --git a/src/tirx/analysis/verify_ssa.cc b/src/tirx/analysis/verify_ssa.cc index 844d6a2a1dcd..b14dd6062e3d 100644 --- a/src/tirx/analysis/verify_ssa.cc +++ b/src/tirx/analysis/verify_ssa.cc @@ -77,7 +77,7 @@ class SSAVerifier final : public StmtExprVisitor { StmtExprVisitor::VisitStmt_(op); } void VisitStmt_(const AllocBufferNode* op) final { - MarkDef(op->buffer->data, op->buffer->data); + MarkDef(op->buffer.var(), op->buffer.var()); StmtExprVisitor::VisitStmt_(op); } @@ -99,9 +99,9 @@ class SSAVerifier final : public StmtExprVisitor { this->VisitStmt(func->body); } - void DefineBuffer(const Buffer& buffer) { + void DefineBuffer(const BufferVar& buffer) { match_scope_ = true; - this->VisitExpr(buffer->data); + this->VisitExpr(buffer.var()); for (size_t i = 0; i < buffer->shape.size(); ++i) { this->VisitExpr(buffer->shape[i]); } diff --git a/src/tirx/analysis/verify_well_formed.cc b/src/tirx/analysis/verify_well_formed.cc index 59e79d6f4b63..a2bb35aae6d5 100644 --- a/src/tirx/analysis/verify_well_formed.cc +++ b/src/tirx/analysis/verify_well_formed.cc @@ -167,13 +167,7 @@ class UndefinedVarVerifier : public Verifier { } } - void EnterDef(const Buffer& buffer, AccessPath path) override { - // A buffer definition implicitly defines its data Var when that Var has no - // prior definition (e.g., tmem buffers where DeclBuffer auto-creates data). - if (currently_defined_.find(buffer->data) == currently_defined_.end() && - previously_defined_.find(buffer->data) == previously_defined_.end()) { - currently_defined_.insert({buffer->data, path->Attr("data")}); - } + void EnterDef(const BufferVar& buffer, AccessPath path) override { Verifier::EnterDef(buffer, path); } @@ -269,13 +263,13 @@ class UndefinedBufferVerifier : public Verifier { previously_defined_.clear(); } - void EnterDef(const Buffer& buffer, AccessPath path) override { + void EnterDef(const BufferVar& buffer, AccessPath path) override { // Call the base class to visit buffer's internal vars (shape, strides, etc.) Verifier::EnterDef(buffer, path); currently_defined_.insert({buffer, path}); } - void ExitDef(const Buffer& buffer, AccessPath path) override { + void ExitDef(const BufferVar& buffer, AccessPath path) override { auto active_def = currently_defined_.find(buffer); if (active_def != currently_defined_.end()) { currently_defined_.erase(active_def); @@ -283,30 +277,30 @@ class UndefinedBufferVerifier : public Verifier { previously_defined_.insert({buffer, path}); } - void VisitBufferUse(const Buffer& buffer, AccessPath path) override { + void VisitBufferUse(const BufferVar& buffer, AccessPath path) override { bool is_declared = currently_defined_.count(buffer); bool was_declared = previously_defined_.count(buffer); if (was_declared && !is_declared) { - // Buffer was previously declared but is now out of scope — always an error. + // BufferVar was previously declared but is now out of scope — always an error. auto prev_def = previously_defined_.find(buffer); - Verify(false) << "TIR is ill-formed: buffer " << buffer->name << " is used at " << path + Verify(false) << "TIR is ill-formed: buffer " << buffer.name() << " is used at " << path << " but its declaration is no longer in-scope. " << "It was declared at " << prev_def->second << "."; } else if (!is_declared && !was_declared) { - // Buffer was never declared — error. - Verify(false) << "TIR is ill-formed: buffer " << buffer->name << " is used at " << path + // BufferVar was never declared — error. + Verify(false) << "TIR is ill-formed: buffer " << buffer.name() << " is used at " << path << " without a prior DeclBuffer or other declaration."; } - // Buffer fields are visited at definition site (EnterDef), not here. + // BufferVar fields are visited at definition site (EnterDef), not here. Verifier::VisitBufferUse(buffer, path); } // Buffers defined in the currently-visited scope. - std::unordered_map + std::unordered_map currently_defined_; // Buffers that were previously defined and are now out of scope. - std::unordered_map + std::unordered_map previously_defined_; }; diff --git a/src/tirx/ir/async_structs.cc b/src/tirx/ir/async_structs.cc index 95f821be698b..0259af5e211f 100644 --- a/src/tirx/ir/async_structs.cc +++ b/src/tirx/ir/async_structs.cc @@ -36,7 +36,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { /*************************** Pipeline ***************************/ Pipeline::Pipeline(ExecScope thread_scope, size_t depth, bool separate_pc, ffi::String name_hint, - ffi::Map workspace, + ffi::Map workspace, ffi::Map schedule_config) { auto n = ffi::make_object(); n->thread_scope = std::move(thread_scope); @@ -53,7 +53,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def( "tirx.Pipeline", [](ExecScope thread_scope, size_t depth, bool separate_pc, ffi::String name_hint, - ffi::Map workspace, ffi::Map schedule_config) { + ffi::Map workspace, ffi::Map schedule_config) { return Pipeline(thread_scope, depth, separate_pc, name_hint, workspace, schedule_config); }); } @@ -61,7 +61,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { /*************************** CopyPipeline ***************************/ CopyPipeline::CopyPipeline(ExecScope thread_scope, size_t depth, bool separate_pc, - ffi::String name_hint, ffi::Map workspace, + ffi::String name_hint, ffi::Map workspace, ffi::Map schedule_config) { auto n = ffi::make_object(); n->thread_scope = std::move(thread_scope); @@ -77,7 +77,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def("tirx.CopyPipeline", [](ExecScope thread_scope, size_t depth, bool separate_pc, ffi::String name_hint, - ffi::Map workspace, + ffi::Map workspace, ffi::Map schedule_config) { return CopyPipeline(thread_scope, depth, separate_pc, name_hint, workspace, schedule_config); }); diff --git a/src/tirx/ir/buffer.cc b/src/tirx/ir/buffer.cc index 6deea08ac617..c0fa6e065ca4 100644 --- a/src/tirx/ir/buffer.cc +++ b/src/tirx/ir/buffer.cc @@ -39,11 +39,61 @@ namespace tvm { namespace tirx { -TVM_FFI_STATIC_INIT_BLOCK() { BufferNode::RegisterReflection(); } +TVM_FFI_STATIC_INIT_BLOCK() { BufferTypeNode::RegisterReflection(); } using IndexMod = tirx::FloorModNode; using IndexDiv = tirx::FloorDivNode; +BufferType::BufferType(PointerType data_pointer_type, PrimType dtype, + ffi::Array shape, ffi::Array strides, + PrimExpr elem_offset, int data_alignment, int offset_factor, + ffi::Optional layout, ffi::Array allocated_addr, + Span span) + : Type(ffi::UnsafeInit{}) { + auto n = ffi::make_object(); + n->data_pointer_type = std::move(data_pointer_type); + n->dtype = std::move(dtype); + n->shape = std::move(shape); + n->strides = std::move(strides); + if (!elem_offset.defined()) { + elem_offset = IntImm(PrimType(n->DefaultIndexType()), 0); + } + n->elem_offset = std::move(elem_offset); + n->data_alignment = + data_alignment <= 0 ? static_cast(runtime::kAllocAlignment) : data_alignment; + n->offset_factor = offset_factor == 0 ? 1 : offset_factor; + n->layout = std::move(layout); + n->allocated_addr = std::move(allocated_addr); + n->span = std::move(span); + data_ = std::move(n); +} + +namespace { + +BufferType MakeBufferType(Var data, PrimType dtype, ffi::Array shape, + ffi::Array strides, PrimExpr elem_offset, + int data_alignment, int offset_factor, + ffi::Optional layout, + ffi::Array allocated_addr) { + TVM_FFI_ICHECK(!data->ty.IsMissing()) + << "Variable " << data->name << " is missing a type annotation."; + const auto* pointer_type = data->ty.as(); + TVM_FFI_ICHECK(pointer_type) << "Variable " << data->name << " is not a pointer."; + TVM_FFI_ICHECK(pointer_type->element_type.as()) + << "Variable " << data->name << " does not point to a primitive."; + return BufferType(ffi::GetRef(pointer_type), std::move(dtype), + std::move(shape), std::move(strides), std::move(elem_offset), + data_alignment, offset_factor, std::move(layout), + std::move(allocated_addr)); +} + +BufferVar RebuildBufferVar(const BufferVar& buffer, BufferType type, + ffi::String name_suffix = "") { + return BufferVar(buffer.name() + name_suffix, std::move(type), buffer.span()); +} + +} // namespace + ffi::Array SimplifyArray(arith::AnalyzerObj* ana, ffi::Array array) { for (size_t i = 0; i < array.size(); ++i) { array.Set(i, ana->Simplify(array[i])); @@ -51,12 +101,12 @@ ffi::Array SimplifyArray(arith::AnalyzerObj* ana, ffi::Array return array; } -Buffer decl_buffer(ffi::Array shape, PrimType dtype, ffi::String name, - ffi::String storage_scope, Span span) { +BufferVar decl_buffer(ffi::Array shape, PrimType dtype, ffi::String name, + ffi::String storage_scope, Span span) { PrimType storage_type = (dtype == PrimType::Bool() ? PrimType::Int(8) : dtype); - return Buffer(Var(name, PointerType(storage_type, storage_scope), span), dtype, shape, - ffi::Array(), PrimExpr(), name, 0, 0, span, std::nullopt, - ffi::Array()); + return BufferVar(Var(name, PointerType(storage_type, storage_scope), span), dtype, shape, + ffi::Array(), PrimExpr(), name, 0, 0, span, std::nullopt, + ffi::Array()); } // Split the given expression w.r.t the add operator @@ -251,17 +301,17 @@ inline PrimExpr MergeMulMod(arith::AnalyzerObj* analyzer, const PrimExpr& base) return no_opt_sum; } -ffi::Array Buffer::OffsetOf(ffi::Array input_indices) const { +ffi::Array BufferVar::OffsetOf(ffi::Array input_indices) const { return (*this)->ElemOffset(std::move(input_indices)); } // The buffer offset in convention of number of elements of // original data ignoring number of lanes. // We also perform optimization to simplify the indexing expression. -ffi::Array BufferNode::ElemOffset(ffi::Array input_indices, bool inner) const { +ffi::Array BufferTypeNode::ElemOffset(ffi::Array input_indices, bool inner) const { TVM_FFI_ICHECK_EQ(shape.size(), input_indices.size()) - << "Buffer " << this->name << " is " << shape.size() - << "-dimensional, cannot be indexed with the " << input_indices.size() + << "BufferType is " << shape.size() << "-dimensional, cannot be indexed with the " + << input_indices.size() << "-dimensional indices provided."; if (strides.size()) { @@ -292,10 +342,10 @@ ffi::Array BufferNode::ElemOffset(ffi::Array input_indices, return SimplifyArray(ana.get(), {output_index}); } -inline ffi::Array BufferOffset(const BufferNode* n, ffi::Array index, +inline ffi::Array BufferOffset(const BufferTypeNode* n, ffi::Array index, PrimType dtype) { ffi::Array offsets = n->ElemOffset(index); - // If the Buffer has element type with more than one lane, scale to + // If the BufferVar has element type with more than one lane, scale to // get the offset in number of scalars. if (PrimType(n->dtype).lanes() != 1) { PrimExpr last_offset = offsets[offsets.size() - 1]; @@ -313,7 +363,7 @@ inline ffi::Array BufferOffset(const BufferNode* n, ffi::Array(); ffi::Array output_shape{1}; @@ -333,24 +383,23 @@ Buffer Buffer::GetFlattenedBuffer() const { if (output_shape.size() == self->shape.size() && self->strides.empty()) { return *this; } else { - Buffer output = *this; - auto writer = output.CopyOnWrite(); - writer->shape = output_shape; - writer->strides = {}; // Keep `layout` in sync with `shape`. The old layout describes the // pre-flatten N-D shape (e.g. `S[(16,16):(16,1)]`); after collapsing // shape to 1-D, that layout no longer matches the buffer's rank and // structural compares against a freshly-decl'd 1-D buffer would diff // (see test_tir_transform_flatten_buffer). Reset to the default layout // for the new shape so the buffer stays internally consistent. - writer->layout = TileLayoutNode::DefaultLayout(output_shape); - return output; + return RebuildBufferVar( + *this, + BufferType(self->data_pointer_type, self->dtype, output_shape, {}, + self->elem_offset, self->data_alignment, self->offset_factor, + TileLayoutNode::DefaultLayout(output_shape), self->allocated_addr)); } } -PrimExpr Buffer::vload(ffi::Array begin, PrimType value_dtype, +PrimExpr BufferVar::vload(ffi::Array begin, PrimType value_dtype, ffi::Optional predicate) const { - const BufferNode* n = operator->(); + const BufferTypeNode* n = operator->(); TVM_FFI_ICHECK(n != nullptr); PrimType buffer_dtype(n->dtype); int value_lanes = @@ -373,9 +422,9 @@ PrimExpr Buffer::vload(ffi::Array begin, PrimType value_dtype, return BufferLoad(*this, indices, predicate); } -Stmt Buffer::vstore(ffi::Array begin, PrimExpr value, +Stmt BufferVar::vstore(ffi::Array begin, PrimExpr value, ffi::Optional predicate) const { - const BufferNode* n = operator->(); + const BufferTypeNode* n = operator->(); TVM_FFI_ICHECK(n != nullptr); PrimType value_dtype = value.ty(); PrimType buffer_dtype(n->dtype); @@ -399,35 +448,36 @@ Stmt Buffer::vstore(ffi::Array begin, PrimExpr value, return BufferStore(*this, value, indices, predicate); } -ffi::String Buffer::scope() const { - const auto* ptr_type = (*this)->data->ty.as(); - TVM_FFI_ICHECK(ptr_type) << "Buffer variable is not of pointer type"; - if (ptr_type->storage_scope.empty()) { +ffi::String BufferVar::scope() const { + if ((*this)->data_pointer_type->storage_scope.empty()) { return "global"; } - return ptr_type->storage_scope; + return (*this)->data_pointer_type->storage_scope; } -Buffer Buffer::MakeStrideView() const { +BufferVar BufferVar::MakeStrideView() const { if ((*this)->strides.size() != 0) return *this; if ((*this)->shape.size() == 0) return *this; - std::vector temp; - const BufferNode* self = operator->(); + const BufferTypeNode* self = operator->(); TVM_FFI_ICHECK(self != nullptr); - auto n = ffi::make_object(*self); - PrimExpr acc = IntImm(PrimType(n->DefaultIndexType()), 1); - for (size_t i = n->shape.size(); i != 0; --i) { + PrimExpr acc = IntImm(PrimType(self->DefaultIndexType()), 1); + std::vector temp; + for (size_t i = self->shape.size(); i != 0; --i) { temp.push_back(acc); - acc = acc * n->shape[i - 1]; + acc = acc * self->shape[i - 1]; } + ffi::Array strides; for (size_t i = temp.size(); i != 0; --i) { - n->strides.push_back(temp[i - 1]); + strides.push_back(temp[i - 1]); } - return Buffer(n); + return RebuildBufferVar( + *this, BufferType(self->data_pointer_type, self->dtype, self->shape, + std::move(strides), self->elem_offset, self->data_alignment, + self->offset_factor, self->layout, self->allocated_addr)); } -Buffer Buffer::MakeSlice(ffi::Array begins, ffi::Array extents) const { - const BufferNode* n = operator->(); +BufferVar BufferVar::MakeSlice(ffi::Array begins, ffi::Array extents) const { + const BufferTypeNode* n = operator->(); TVM_FFI_ICHECK(n != nullptr); arith::Analyzer ana; begins = SimplifyArray(ana.get(), begins); @@ -452,22 +502,22 @@ Buffer Buffer::MakeSlice(ffi::Array begins, ffi::Array exten return MakeStrideView().MakeSlice(begins, extents); } } - Buffer slice(n->data, n->dtype, extents, strides, elem_offset[0], n->name + "_slice", - n->data_alignment, 0); - return slice; + return RebuildBufferVar( + *this, + BufferType(n->data_pointer_type, n->dtype, extents, strides, elem_offset[0], + n->data_alignment, 0, TileLayoutNode::DefaultLayout(extents)), + "_slice"); } -Expr Buffer::access_ptr(int access_mask, PointerType ptr_type, int content_lanes, PrimExpr offset, +Expr BufferVar::access_ptr(int access_mask, PointerType ptr_type, int content_lanes, PrimExpr offset, ffi::Optional input_extent) const { - const BufferNode* self = operator->(); + const BufferTypeNode* self = operator->(); TVM_FFI_ICHECK(self != nullptr); - const auto* data_pointer_type = self->data->ty.as(); - TVM_FFI_ICHECK(data_pointer_type) - << "Buffer data must have PointerType, but got " << self->data->ty; // An access pointer addresses the same allocation as the buffer data. The // requested type controls its pointee, while the buffer controls its address // space (for example, shared or local memory). - ptr_type = PointerType(ptr_type->element_type, data_pointer_type->storage_scope); + ptr_type = + PointerType(ptr_type->element_type, self->data_pointer_type->storage_scope); PrimExpr e_dtype; PrimExpr extent; if (self->shape.size() == 0) { @@ -492,64 +542,31 @@ Expr Buffer::access_ptr(int access_mask, PointerType ptr_type, int content_lanes if (input_extent.has_value()) { extent = input_extent.value(); } - ffi::Array acc_args{e_dtype, self->data, elem_offset, extent, IntImm::Int32(access_mask)}; + ffi::Array acc_args{e_dtype, data(), elem_offset, extent, + IntImm::Int32(access_mask)}; return Call(ptr_type, tirx::builtin::tvm_access_ptr(), acc_args); } -Buffer::Buffer(Var data, PrimType dtype, ffi::Array shape, ffi::Array strides, - PrimExpr elem_offset, ffi::String name, int data_alignment, int offset_factor, - Span span, ffi::Optional layout, ffi::Array allocated_addr) { - DLDataType storage_dtype = dtype->dtype; - // specially handle bool - if (storage_dtype == DLDataType{kDLBool, 8, 1}) { - storage_dtype = DLDataType{kDLInt, 8, 1}; - } - // The buffer dtype may differ from the dtype of the underlying - // allocation, such as a single allocation that backs multiple - // tensors without a common datatype. Therefore, we check that the - // data pointer is a pointer, but not the exact type of the - // pointed-to values. - - // TODO(Lunderberg): Use an explicit pointer cast for the data - // pointer. Should be done alongside extensions to StmtExprMutator - // to more easily handle buffer/buffer_var updates. - TVM_FFI_ICHECK(!data->ty.IsMissing()) - << "Variable " << data->name << " is missing a type annotation."; - TVM_FFI_ICHECK(data->ty.as()) - << "Variable " << data->name << " is not a pointer."; - TVM_FFI_ICHECK(data->ty.as()->element_type.as()) - << "Variable " << data->name << " does not point to a primitive."; - - auto n = ffi::make_object(); - n->data = std::move(data); - n->dtype = dtype; - - n->shape = std::move(shape); - n->strides = std::move(strides); - n->name = std::move(name); - if (!elem_offset.defined()) { - elem_offset = IntImm(PrimType(n->DefaultIndexType()), 0); - } - if (data_alignment <= 0) { - data_alignment = runtime::kAllocAlignment; - } - if (offset_factor == 0) { - offset_factor = 1; - } - n->elem_offset = std::move(elem_offset); - n->data_alignment = data_alignment; - n->offset_factor = offset_factor; - n->span = std::move(span); - // `layout=nullopt` is a meaningful sentinel: it tells the printer that the - // user opted out of layout sugar (e.g., the `local_scalar` shorthand keys - // off `layout` being defined). Don't default-fill here — callers that want - // the default `TileLayout::DefaultLayout(shape)` must pass it explicitly. - n->layout = std::move(layout); - n->allocated_addr = std::move(allocated_addr); - data_ = std::move(n); +BufferVar::BufferVar(ffi::String name, BufferType type, Span span) + : Var(Var(std::move(name), std::move(type), std::move(span))) {} + +BufferVar::BufferVar(Var data, PrimType dtype, ffi::Array shape, + ffi::Array strides, PrimExpr elem_offset, + ffi::String name, int data_alignment, int offset_factor, + Span span, ffi::Optional layout, + ffi::Array allocated_addr) + : BufferVar( + std::move(name), + MakeBufferType(std::move(data), std::move(dtype), std::move(shape), + std::move(strides), std::move(elem_offset), data_alignment, + offset_factor, std::move(layout), std::move(allocated_addr)), + std::move(span)) {} + +Expr BufferVar::data() const { + return Call((*this)->data_pointer_type, builtin::buffer_data(), {var()}); } -tirx::Buffer BufferWithOffsetAlignment(ffi::Array shape, PrimType dtype, std::string name, +tirx::BufferVar BufferWithOffsetAlignment(ffi::Array shape, PrimType dtype, std::string name, int data_alignment, int offset_factor, std::string memory_scope) { PrimType storage_type = (dtype == PrimType::Bool() ? PrimType::Int(8) : dtype); @@ -562,37 +579,33 @@ tirx::Buffer BufferWithOffsetAlignment(ffi::Array shape, PrimType dtyp elem_offset = PrimExpr(); } - return tirx::Buffer(data, dtype, shape, ffi::Array(), elem_offset, name, data_alignment, + return tirx::BufferVar(data, dtype, shape, ffi::Array(), elem_offset, name, data_alignment, offset_factor); } -Buffer Buffer::with_allocated_addr(ffi::Array allocated_addr) const { - Buffer output = *this; - auto writer = output.CopyOnWrite(); - writer->allocated_addr = std::move(allocated_addr); - return output; -} - -Buffer Buffer::with_dtype(PrimType dtype) const { - Buffer output = *this; - auto writer = output.CopyOnWrite(); - writer->dtype = dtype; - return output; +BufferVar BufferVar::with_allocated_addr(ffi::Array allocated_addr) const { + const auto* self = operator->(); + return RebuildBufferVar( + *this, BufferType(self->data_pointer_type, self->dtype, self->shape, + self->strides, self->elem_offset, self->data_alignment, + self->offset_factor, self->layout, + std::move(allocated_addr))); } -Buffer Buffer::with_data(Var data) const { - Buffer output = *this; - auto writer = output.CopyOnWrite(); - writer->data = data; - return output; +BufferVar BufferVar::with_dtype(PrimType dtype) const { + const auto* self = operator->(); + return RebuildBufferVar( + *this, BufferType(self->data_pointer_type, std::move(dtype), self->shape, + self->strides, self->elem_offset, self->data_alignment, + self->offset_factor, self->layout, self->allocated_addr)); } -PrimExpr Buffer::OffsetOf_p(const Array& indices) const { +PrimExpr BufferVar::OffsetOf_p(const Array& indices) const { return Call(PrimType::Int(32), tirx::builtin::buffer_offset(), {BufferLoad(*this, indices)}) .as_or_throw(); } -bool Buffer::IsScalar(bool alloc_or_decl) const { +bool BufferVar::IsScalar(bool alloc_or_decl) const { // TODO(@bohan): logical scope is not considered return (*this)->shape.size() == 1 && is_one((*this)->shape[0]) && (*this)->strides.size() == 0 && (!alloc_or_decl || tirx::is_zero((*this)->elem_offset)) && (*this)->data_alignment == 64 && @@ -604,7 +617,7 @@ bool Buffer::IsScalar(bool alloc_or_decl) const { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() - .def_packed("tirx.Buffer", + .def_packed("tirx.BufferVar", [](ffi::PackedArgs args, ffi::Any* ret) { TVM_FFI_ICHECK_EQ(args.size(), 10); auto data = args[0].cast(); @@ -617,24 +630,36 @@ TVM_FFI_STATIC_INIT_BLOCK() { auto offset_factor = args[7].cast(); auto span = args[8].cast(); auto layout = args[9].cast(); - *ret = Buffer(data, dtype, shape, strides, elem_offset, name, data_alignment, + *ret = BufferVar(data, dtype, shape, strides, elem_offset, name, data_alignment, offset_factor, span, layout); }) .def_method("tirx.BufferAccessPtr", - static_cast) const>(&Buffer::access_ptr)) - .def_method("tirx.BufferGetFlattenedBuffer", &Buffer::GetFlattenedBuffer) - .def_method("tirx.BufferOffsetOf", &Buffer::OffsetOf) - .def_method("tirx.BufferOffsetOfp", &Buffer::OffsetOf_p) + static_cast) const>(&BufferVar::access_ptr)) + .def_method("tirx.BufferGetFlattenedBuffer", &BufferVar::GetFlattenedBuffer) + .def_method("tirx.BufferOffsetOf", &BufferVar::OffsetOf) + .def_method("tirx.BufferOffsetOfp", &BufferVar::OffsetOf_p) .def_method("tirx.BufferVLoad", - static_cast, PrimType, - ffi::Optional) const>(&Buffer::vload)) - .def_method("tirx.BufferVStore", &Buffer::vstore) - .def_method("tirx.BufferStorageScope", &Buffer::scope) - .def_method("tirx.BufferWithAllocatedAddr", &Buffer::with_allocated_addr) - .def_method("tirx.BufferWithDtype", &Buffer::with_dtype) - .def_method("tirx.BufferWithData", &Buffer::with_data) - .def_method("tirx.BufferIsScalar", &Buffer::IsScalar); + static_cast, PrimType, + ffi::Optional) const>(&BufferVar::vload)) + .def_method("tirx.BufferVStore", &BufferVar::vstore) + .def_method("tirx.BufferStorageScope", &BufferVar::scope) + .def_method("tirx.BufferWithAllocatedAddr", &BufferVar::with_allocated_addr) + .def_method("tirx.BufferWithDtype", &BufferVar::with_dtype) + .def_method("tirx.BufferIsScalar", &BufferVar::IsScalar) + .def_method("tirx.BufferData", &BufferVar::data) + .def("tirx.BufferType", + [](PointerType data_pointer_type, PrimType dtype, + ffi::Array shape, ffi::Array strides, + PrimExpr elem_offset, int data_alignment, int offset_factor, + ffi::Optional layout, + ffi::Array allocated_addr, Span span) { + return BufferType(std::move(data_pointer_type), std::move(dtype), + std::move(shape), std::move(strides), + std::move(elem_offset), data_alignment, offset_factor, + std::move(layout), std::move(allocated_addr), + std::move(span)); + }); } } // namespace tirx diff --git a/src/tirx/ir/data_type_rewriter.cc b/src/tirx/ir/data_type_rewriter.cc index d20346d90567..c8bd53f100ef 100644 --- a/src/tirx/ir/data_type_rewriter.cc +++ b/src/tirx/ir/data_type_rewriter.cc @@ -297,15 +297,15 @@ Stmt IndexDataTypeRewriter::VisitStmt_(const AttrStmtNode* op) { return DataTypeLegalizer::VisitStmt_(op); } -Buffer IndexDataTypeRewriter::VisitBufferDef(const Buffer& buffer, bool alloc_data) { +BufferVar IndexDataTypeRewriter::VisitBufferDef(const BufferVar& buffer, bool alloc_data) { bool is_enabled = is_enabled_; is_enabled_ = true; - Buffer new_buf = StmtMutator::VisitBufferDef(buffer, alloc_data); + BufferVar new_buf = StmtMutator::VisitBufferDef(buffer, alloc_data); is_enabled_ = is_enabled; return new_buf; } -Buffer IndexDataTypeRewriter::VisitBufferUse(const Buffer& buffer) { +BufferVar IndexDataTypeRewriter::VisitBufferUse(const BufferVar& buffer) { return StmtMutator::VisitBufferUse(buffer); } @@ -336,11 +336,11 @@ Stmt IndexDataTypeRewriter::VisitStmt_(const SBlockRealizeNode* op) { } Stmt IndexDataTypeRewriter::VisitStmt_(const SBlockNode* op) { - ffi::Array new_alloc_buffers = op->alloc_buffers.Map( - [this](const Buffer& buffer) { return this->VisitBufferDef(buffer, /*alloc_data=*/true); }); + ffi::Array new_alloc_buffers = op->alloc_buffers.Map( + [this](const BufferVar& buffer) { return this->VisitBufferDef(buffer, /*alloc_data=*/true); }); ffi::Array new_match_buffers = op->match_buffers.Map([this](const MatchBufferRegion& match_buffer_region) { - Buffer new_buffer = this->VisitBufferDef(match_buffer_region->buffer, /*alloc_data=*/true); + BufferVar new_buffer = this->VisitBufferDef(match_buffer_region->buffer, /*alloc_data=*/true); BufferRegion new_buffer_region = this->VisitBufferRegion(match_buffer_region->source); if (!new_buffer.same_as(match_buffer_region->buffer) || !new_buffer_region.same_as(match_buffer_region->source)) { @@ -390,9 +390,9 @@ ffi::Map IndexDataTypeRewriter::VisitBlockAnnotations( if (obj == nullptr) { return obj; } - if (obj.as()) { - Buffer buffer = obj.as_or_throw(); - if (Buffer new_buffer = VisitBufferUse(buffer); !new_buffer.same_as(buffer)) { + if (obj.as()) { + BufferVar buffer = obj.as_or_throw(); + if (BufferVar new_buffer = VisitBufferUse(buffer); !new_buffer.same_as(buffer)) { return new_buffer; } } else if (obj.as()) { @@ -430,7 +430,7 @@ IterVar IndexDataTypeRewriter::VisitIterVar(const IterVar& iter_var) { } BufferRegion IndexDataTypeRewriter::VisitBufferRegion(const BufferRegion& buffer_region) { - Buffer remapped_buffer = VisitBufferUse(buffer_region->buffer); + BufferVar remapped_buffer = VisitBufferUse(buffer_region->buffer); bool is_enabled = is_enabled_; is_enabled_ = true; @@ -451,7 +451,7 @@ BufferRegion IndexDataTypeRewriter::VisitBufferRegion(const BufferRegion& buffer Stmt IndexDataTypeRewriter::VisitStmt_(const BufferStoreNode* op) { BufferStore store = ffi::GetRef(op); - Buffer new_buffer = VisitBufferUse(op->buffer); + BufferVar new_buffer = VisitBufferUse(op->buffer); auto value = this->VisitPrimExpr(op->value); PrimType value_dtype = value.ty(); if (new_buffer->dtype != value_dtype && value_dtype.IsScalar()) { @@ -473,7 +473,7 @@ Stmt IndexDataTypeRewriter::VisitStmt_(const BufferStoreNode* op) { Expr IndexDataTypeRewriter::VisitExpr_(const BufferLoadNode* op) { BufferLoad load = ffi::GetRef(op); - Buffer new_buffer = VisitBufferUse(op->buffer); + BufferVar new_buffer = VisitBufferUse(op->buffer); auto indices = VisitIndices(op->indices); if (!new_buffer.same_as(op->buffer) || !indices.same_as(op->indices)) { @@ -630,7 +630,7 @@ PrimFunc IndexDataTypeNormalizer::Rewrite(PrimFunc func) { buffer_remap_.clear(); ivmap_.clear(); // start rewrite - ffi::Map new_buffer_map = func->buffer_map; + ffi::Map new_buffer_map = func->buffer_map; for (const auto& [var, buffer] : func->buffer_map) { new_buffer_map.Set(var, VisitBufferDef(buffer, /*alloc_data=*/true)); } diff --git a/src/tirx/ir/data_type_rewriter.h b/src/tirx/ir/data_type_rewriter.h index 7a385ac37628..883bde94b096 100644 --- a/src/tirx/ir/data_type_rewriter.h +++ b/src/tirx/ir/data_type_rewriter.h @@ -101,8 +101,8 @@ class IndexDataTypeRewriter : public DataTypeLegalizer { using Parent::VisitExpr_; using Parent::VisitStmt_; - Buffer VisitBufferDef(const Buffer& buffer, bool alloc_data) override; - Buffer VisitBufferUse(const Buffer& buffer) override; + BufferVar VisitBufferDef(const BufferVar& buffer, bool alloc_data) override; + BufferVar VisitBufferUse(const BufferVar& buffer) override; Stmt VisitStmt_(const SBlockRealizeNode* op) override; Stmt VisitStmt_(const SBlockNode* op) override; Stmt VisitStmt_(const BufferStoreNode* op) override; diff --git a/src/tirx/ir/expr.cc b/src/tirx/ir/expr.cc index c7a044a5aa41..a975b31153d9 100644 --- a/src/tirx/ir/expr.cc +++ b/src/tirx/ir/expr.cc @@ -745,10 +745,10 @@ void BufferLoadNode::LegalizeDType() { } } -BufferLoad::BufferLoad(Buffer buffer, ffi::Array indices, +BufferLoad::BufferLoad(BufferVar buffer, ffi::Array indices, ffi::Optional predicate, Span span) { TVM_FFI_ICHECK_EQ(buffer->shape.size(), indices.size()) - << "Buffer " << buffer->name << " is " << buffer->shape.size() + << "BufferVar " << buffer.name() << " is " << buffer->shape.size() << "-dimensional, cannot be indexed with the " << indices.size() << "-dimensional indices provided."; @@ -784,7 +784,7 @@ BufferLoad::BufferLoad(Buffer buffer, ffi::Array indices, TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("tirx.BufferLoad", [](Buffer buffer, ffi::Array indices, + refl::GlobalDef().def("tirx.BufferLoad", [](BufferVar buffer, ffi::Array indices, ffi::Optional predicate, Span span) { return BufferLoad(buffer, indices, predicate, span); }); diff --git a/src/tirx/ir/function.cc b/src/tirx/ir/function.cc index 7ae0e153f2be..c991e572dc28 100644 --- a/src/tirx/ir/function.cc +++ b/src/tirx/ir/function.cc @@ -80,7 +80,7 @@ tvm::Type InferType(const PrimFunc& prim_func) { // Get the function type of a PrimFunc PrimFunc::PrimFunc(ffi::Array params, Stmt body, Type ret_type, - ffi::Map buffer_map, DictAttrs attrs, Span span) { + ffi::Map buffer_map, DictAttrs attrs, Span span) { if (ret_type.IsMissing()) { ret_type = VoidType(); } @@ -165,7 +165,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef() .def("tirx.PrimFunc", [](ffi::Array params, Stmt body, Type ret_type, - ffi::Map buffer_map, DictAttrs attrs, + ffi::Map buffer_map, DictAttrs attrs, Span span) { return PrimFunc(params, body, ret_type, buffer_map, attrs, span); }) .def("tirx.TensorIntrin", [](PrimFunc desc_func, PrimFunc intrin_func) { diff --git a/src/tirx/ir/script/script_complete.cc b/src/tirx/ir/script/script_complete.cc index bb915e96acf8..ae128f20ad08 100644 --- a/src/tirx/ir/script/script_complete.cc +++ b/src/tirx/ir/script/script_complete.cc @@ -38,11 +38,11 @@ namespace tirx { /*! \brief Generate surrounding loops automatically */ class ScriptCompleter : public StmtMutator { public: - explicit ScriptCompleter(ffi::Map* buffer_var_map, bool s_tir = false) + explicit ScriptCompleter(ffi::Map* buffer_var_map, bool s_tir = false) : buffer_var_map_(buffer_var_map), s_tir_(s_tir) {} private: - ffi::Map* buffer_var_map_; + ffi::Map* buffer_var_map_; Stmt VisitStmt_(const SBlockRealizeNode* op) final { for (const PrimExpr& value : op->iter_values) { PrimType value_ty = value.ty(); @@ -55,11 +55,11 @@ class ScriptCompleter : public StmtMutator { Stmt VisitStmt_(const SBlockNode* op) final { // Buffers allocated in the block can be accessed by its body. for (const auto& alloc_buffer : op->alloc_buffers) { - buffer_var_map_->Set(alloc_buffer->data, alloc_buffer); + buffer_var_map_->Set(alloc_buffer.var(), alloc_buffer); } for (const auto& match_buffer : op->match_buffers) { - const Buffer& target_buffer = match_buffer->buffer; - buffer_var_map_->Set(target_buffer->data, target_buffer); + const BufferVar& target_buffer = match_buffer->buffer; + buffer_var_map_->Set(target_buffer.var(), target_buffer); } bool is_root_block = this->is_root_block_; @@ -69,11 +69,11 @@ class ScriptCompleter : public StmtMutator { // Remove buffers allocated inside block to detect its access region for (const auto& alloc_buffer : op->alloc_buffers) { - buffer_var_map_->erase(alloc_buffer->data); + buffer_var_map_->erase(alloc_buffer.var()); } for (const auto& match_buffer : op->match_buffers) { - const Buffer& target_buffer = match_buffer->buffer; - buffer_var_map_->erase(target_buffer->data); + const BufferVar& target_buffer = match_buffer->buffer; + buffer_var_map_->erase(target_buffer.var()); } // Get access detection mask // 0 for provided region, 1 and 3 for need detect read, 2 and 3 for need detect write @@ -106,16 +106,16 @@ class ScriptCompleter : public StmtMutator { Stmt VisitStmt_(const AllocBufferNode* op) final { // AllocBuffer is flat: register buffer for subsequent siblings - if (!buffer_var_map_->count(op->buffer->data)) { - buffer_var_map_->Set(op->buffer->data, op->buffer); + if (!buffer_var_map_->count(op->buffer.var())) { + buffer_var_map_->Set(op->buffer.var(), op->buffer); } return StmtMutator::VisitStmt_(op); } Stmt VisitStmt_(const DeclBufferNode* op) final { // DeclBuffer is flat: register buffer for subsequent siblings - if (!buffer_var_map_->count(op->buffer->data)) { - buffer_var_map_->Set(op->buffer->data, op->buffer); + if (!buffer_var_map_->count(op->buffer.var())) { + buffer_var_map_->Set(op->buffer.var(), op->buffer); } return StmtMutator::VisitStmt_(op); } @@ -124,14 +124,14 @@ class ScriptCompleter : public StmtMutator { bool s_tir_ = false; }; -PrimFunc ScriptComplete(PrimFunc func, const ffi::Array& root_allocates, bool s_tir) { - ffi::Map buffer_var_map; +PrimFunc ScriptComplete(PrimFunc func, const ffi::Array& root_allocates, bool s_tir) { + ffi::Map buffer_var_map; for (const auto& pair : func->buffer_map) { - const Buffer& buffer = pair.second; - buffer_var_map.Set(buffer->data, buffer); + const BufferVar& buffer = pair.second; + buffer_var_map.Set(buffer.var(), buffer); } for (const auto& alloc : root_allocates) { - buffer_var_map.Set(alloc->data, alloc); + buffer_var_map.Set(alloc.var(), alloc); } Stmt res = func->body; diff --git a/src/tirx/ir/script/script_complete.h b/src/tirx/ir/script/script_complete.h index 775a00aab0c3..5f93f7955b0f 100644 --- a/src/tirx/ir/script/script_complete.h +++ b/src/tirx/ir/script/script_complete.h @@ -30,7 +30,7 @@ namespace tvm { namespace tirx { -PrimFunc ScriptComplete(PrimFunc func, const ffi::Array& root_allocates, +PrimFunc ScriptComplete(PrimFunc func, const ffi::Array& root_allocates, bool s_tir = false); } // namespace tirx diff --git a/src/tirx/ir/specialize.cc b/src/tirx/ir/specialize.cc index 9cfeb2d455cb..fcea0f7bf349 100644 --- a/src/tirx/ir/specialize.cc +++ b/src/tirx/ir/specialize.cc @@ -78,13 +78,13 @@ class PrimFuncSpecializer : public StmtExprMutator { static PrimFunc Specialize(PrimFunc f, const VarMap& var_map) { PrimFuncSpecializer specializer(var_map); - // Updating Buffer map - ffi::Map buffer_map; + // Updating BufferVar map + ffi::Map buffer_map; bool buffer_map_updated = false; for (const auto& it : f->buffer_map) { const Var& var = it.first; - const Buffer& buffer = it.second; - Buffer new_buffer = specializer.MutateBuffer(buffer); + const BufferVar& buffer = it.second; + BufferVar new_buffer = specializer.MutateBuffer(buffer); buffer_map.Set(var, new_buffer); if (!new_buffer.same_as(buffer)) { buffer_map_updated = true; @@ -117,7 +117,7 @@ class PrimFuncSpecializer : public StmtExprMutator { private: Stmt VisitStmt_(const SBlockNode* op) final { // Step.0. Define buffer mappings which is allocated inside the block - ffi::Array alloc_buffers = + ffi::Array alloc_buffers = op->alloc_buffers.Map([this](const auto& buf) { return MutateAllocBuffer(buf); }); // Step.1. Recursively visit block body @@ -145,8 +145,7 @@ class PrimFuncSpecializer : public StmtExprMutator { Stmt VisitStmt_(const DeclBufferNode* op) final { // Visit the buffer before delegating to StmtExprMutator, so the // buffer's replacement will be defined before the point of use. - Var old_buffer_var = op->buffer->data; - Buffer new_buf = MutateAllocBuffer(op->buffer); + BufferVar new_buf = MutateAllocBuffer(op->buffer); auto node = StmtExprMutator::VisitStmt_(op).as_or_throw(); @@ -154,33 +153,11 @@ class PrimFuncSpecializer : public StmtExprMutator { node.CopyOnWrite()->buffer = new_buf; } - // If the buffer variable is being remapped to an expression, we - // still need a tirx::Var to be used as a the buffer variable. - // Therefore, generate a Bind that will provide a tirx::Var for - // the buffer to use. - // - // This step is only required when a buffer definition is using a - // previously-defined buffer variable, which is therefore eligible - // for specialization. An allocation in the - // `BlockNode::alloc_buffers` defines both the buffer variable and - // the buffer, this check is unnecessary there. In addition, if - // the buffer var has been remapped to another variable, it has already - // been handled as part of the buffer mutation. - Var new_buffer_var = node->buffer->data; - Stmt stmt = std::move(node); - - if (new_buffer_var.same_as(old_buffer_var)) { - auto remapped_data = VisitExpr(old_buffer_var); - if (!remapped_data.same_as(old_buffer_var)) { - stmt = SeqStmt({Bind(old_buffer_var, remapped_data), stmt}); - } - } - - return stmt; + return node; } // Override VisitBufferUse to use our own buffer_map_ instead of base class field visiting. - Buffer VisitBufferUse(const Buffer& buffer) final { return GetNewBuffer(buffer); } + BufferVar VisitBufferUse(const BufferVar& buffer) final { return GetNewBuffer(buffer); } Expr VisitExpr_(const VarNode* op) final { auto it = var_map_.find(ffi::GetRef(op)); @@ -211,15 +188,7 @@ class PrimFuncSpecializer : public StmtExprMutator { DEFINE_SPECIALIZER_UNARY_OP_MUTATE(NotNode, logical_not); private: - Buffer MutateBuffer(const Buffer& buffer) { - // For the data variable, only Var-to-Var remapping can be handled - // in MutateBuffer. See the DeclBuffer visitor for the handling - // of Var-to-PrimExpr remapping. - Var data = buffer->data; - if (auto new_data = VisitExpr(buffer->data).as()) { - data = new_data.value(); - } - + BufferVar MutateBuffer(const BufferVar& buffer) { ffi::Array shape = buffer->shape.Map([this](const PrimExpr& e) { return VisitPrimExpr(e); }); ffi::Array strides = @@ -251,19 +220,18 @@ class PrimFuncSpecializer : public StmtExprMutator { } } - if (buffer->data.same_as(data) && buffer->elem_offset.same_as(elem_offset) && - buffer->shape.same_as(shape) && buffer->strides.same_as(strides) && !layout_changed) { + if (buffer->elem_offset.same_as(elem_offset) && buffer->shape.same_as(shape) && + buffer->strides.same_as(strides) && !layout_changed) { return buffer; } else { - auto n = ffi::make_object(*buffer.get()); - n->data = std::move(data); + auto n = CopyBufferType(buffer); n->elem_offset = std::move(elem_offset); n->shape = std::move(shape); n->strides = std::move(strides); if (layout_changed) { n->layout = std::move(layout); } - return Buffer(n); + return RebuildBufferVar(buffer, std::move(n)); } } @@ -277,23 +245,23 @@ class PrimFuncSpecializer : public StmtExprMutator { } } - Buffer MutateAllocBuffer(const Buffer& alloc_buf) { + BufferVar MutateAllocBuffer(const BufferVar& alloc_buf) { TVM_FFI_ICHECK(!buffer_map_.count(alloc_buf)) << "Multiple points of definition found for buffer " << alloc_buf; - Buffer buf = MutateBuffer(alloc_buf); + BufferVar buf = MutateBuffer(alloc_buf); buffer_map_[alloc_buf] = buf; return buf; } - Buffer GetNewBuffer(const Buffer& old_buffer) { + BufferVar GetNewBuffer(const BufferVar& old_buffer) { if (auto it = buffer_map_.find(old_buffer); it != buffer_map_.end()) { return it->second; } auto mutated = MutateBuffer(old_buffer); TVM_FFI_ICHECK(mutated.same_as(old_buffer)) - << "Buffer " << old_buffer << " (shape = " << old_buffer->shape << ")" + << "BufferVar " << old_buffer << " (shape = " << old_buffer->shape << ")" << " was used without a declaration, " << "and would be specialized into " << mutated << " (shape = " << mutated->shape << "). " << "While usage of an undeclared buffer is currently allowed in TIR, " @@ -309,7 +277,7 @@ class PrimFuncSpecializer : public StmtExprMutator { BufferRegion MutateBufferRegion(const BufferRegion& buffer_region) { auto it = buffer_map_.find(buffer_region->buffer); - const Buffer& buffer = it != buffer_map_.end() ? it->second : buffer_region->buffer; + const BufferVar& buffer = it != buffer_map_.end() ? it->second : buffer_region->buffer; ffi::Array region = buffer_region->region.Map( std::bind(&PrimFuncSpecializer::MutateRange, this, std::placeholders::_1)); if (it == buffer_map_.end() && region.same_as(buffer_region->region)) { @@ -323,7 +291,7 @@ class PrimFuncSpecializer : public StmtExprMutator { /*! \brief The vars to be substitute and their values */ const VarMap& var_map_; /*! \brief map from old buffer to mutated buffer */ - std::unordered_map buffer_map_; + std::unordered_map buffer_map_; }; /*! @@ -343,7 +311,7 @@ class PrimFuncSpecializer : public StmtExprMutator { * If the buffer signature is not a Var, the mapping will fail. * e.g. A = T.match_buffer(a, [m * 2, n + 1]) */ -void UpdateSpecializeVarMap(const PrimFunc& func, const Var& param, const Buffer& specific_buf, +void UpdateSpecializeVarMap(const PrimFunc& func, const Var& param, const BufferVar& specific_buf, VarMap* var_map) { // preliminaries tirx::ExprDeepEqual equal; @@ -351,7 +319,7 @@ void UpdateSpecializeVarMap(const PrimFunc& func, const Var& param, const Buffer auto it = func->buffer_map.find(param); TVM_FFI_CHECK(it != func->buffer_map.end(), ValueError) << "specialize expects param to be in PrimFunc's buffer_map"; - const Buffer& buf_to_specialize = (*it).second; + const BufferVar& buf_to_specialize = (*it).second; // build var mapping using specific_buf's parameters auto build_var_mapping = [&](const Expr& new_expr, const Expr& old_expr) { @@ -389,7 +357,6 @@ void UpdateSpecializeVarMap(const PrimFunc& func, const Var& param, const Buffer << specific_buf->strides.size() << "."; // Updating var mapping using specific_expr - build_var_mapping(specific_buf->data, buf_to_specialize->data); for (size_t i = 0; i < specific_buf->shape.size(); ++i) { build_var_mapping(specific_buf->shape[i], buf_to_specialize->shape[i]); } @@ -430,17 +397,17 @@ void UpdateSpecializeVarMap(const PrimFunc& func, const Var& param, const Expr& /**************** Implementation ****************/ -PrimFunc Specialize(PrimFunc func, const ffi::Map>& param_map) { +PrimFunc Specialize(PrimFunc func, const ffi::Map>& param_map) { VarMap var_map; for (const auto& kv : param_map) { const Var& param = kv.first; - const ffi::Variant& instance = kv.second; - if (auto opt_buffer = instance.as()) { + const ffi::Variant& instance = kv.second; + if (auto opt_buffer = instance.as()) { UpdateSpecializeVarMap(func, param, opt_buffer.value(), &var_map); } else if (auto opt_expr = instance.as()) { UpdateSpecializeVarMap(func, param, opt_expr.value(), &var_map); } else { - TVM_FFI_THROW(TypeError) << "specialize expected instance to be Buffer or Expr"; + TVM_FFI_THROW(TypeError) << "specialize expected instance to be BufferVar or Expr"; } } return PrimFuncSpecializer::Specialize(func, std::move(var_map)); @@ -454,14 +421,14 @@ TVM_FFI_STATIC_INIT_BLOCK() { const ffi::Map& param_map) { VarMap var_map; for (const auto& [param, instance] : param_map) { - if (auto buffer = instance.as()) { + if (auto buffer = instance.as()) { UpdateSpecializeVarMap(func, param, buffer.value(), &var_map); } else if (const ExprNode* expr = instance.as()) { UpdateSpecializeVarMap(func, param, ffi::GetRef(expr), &var_map); } else if (instance.type_index() < ffi::TypeIndex::kTVMFFISmallStr) { UpdateSpecializeVarMap(func, param, instance.cast(), &var_map); } else { - TVM_FFI_THROW(TypeError) << "specialize expected instance to be Buffer or Expr, but got " + TVM_FFI_THROW(TypeError) << "specialize expected instance to be BufferVar or Expr, but got " << instance.GetTypeKey(); } } diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index c768e450a6cb..c7b494c4b2c0 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -277,7 +277,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } // DeclBuffer -DeclBuffer::DeclBuffer(Buffer buffer, Span span) { +DeclBuffer::DeclBuffer(BufferVar buffer, ffi::Optional data, Span span) { // Enforce storage scope rules for DeclBuffer. std::string scope = static_cast(buffer.scope()); if (scope.empty()) { @@ -293,18 +293,21 @@ DeclBuffer::DeclBuffer(Buffer buffer, Span span) { } ffi::ObjectPtr node = ffi::make_object(); node->buffer = std::move(buffer); + node->data = std::move(data); node->span = std::move(span); data_ = std::move(node); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("tirx.DeclBuffer", - [](Buffer buffer, Span span) { return DeclBuffer(buffer, span); }); + refl::GlobalDef().def( + "tirx.DeclBuffer", [](BufferVar buffer, ffi::Optional data, Span span) { + return DeclBuffer(buffer, data, span); + }); } // AllocBuffer -AllocBuffer::AllocBuffer(Buffer buffer, ffi::Map annotations, Span span) { +AllocBuffer::AllocBuffer(BufferVar buffer, ffi::Map annotations, Span span) { ffi::ObjectPtr node = ffi::make_object(); node->buffer = std::move(buffer); node->annotations = std::move(annotations); @@ -316,7 +319,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def( "tirx.AllocBuffer", - [](Buffer buffer, ffi::Optional> annotations, Span span) { + [](BufferVar buffer, ffi::Optional> annotations, Span span) { return AllocBuffer(buffer, annotations.value_or(ffi::Map()), span); }); } @@ -398,10 +401,10 @@ TVM_FFI_INLINE int GetLanesOrVScaleFactor(const PrimType& ty) { return ty.IsScalableVector() ? ty.VScaleFactor() : ty.lanes(); } -BufferStore::BufferStore(Buffer buffer, PrimExpr value, ffi::Array indices, +BufferStore::BufferStore(BufferVar buffer, PrimExpr value, ffi::Array indices, ffi::Optional predicate, Span span) { TVM_FFI_ICHECK_EQ(buffer->shape.size(), indices.size()) - << "Buffer " << buffer->name << " is " << buffer->shape.size() + << "BufferVar " << buffer.name() << " is " << buffer->shape.size() << "-dimensional, cannot be indexed with the " << indices.size() << "-dimensional indices provided."; @@ -480,7 +483,7 @@ BufferStore::BufferStore(Buffer buffer, PrimExpr value, ffi::Array ind TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def("tirx.BufferStore", - [](Buffer buffer, PrimExpr value, ffi::Array indices, + [](BufferVar buffer, PrimExpr value, ffi::Array indices, ffi::Optional predicate, Span span) { return BufferStore(buffer, value, indices, predicate, span); }); @@ -504,7 +507,7 @@ PrimExpr BufferRegionNode::ToPrimExpr() const { return tirx::BufferLoad(this->buffer, indices); } -BufferRegion::BufferRegion(Buffer buffer, ffi::Array region) { +BufferRegion::BufferRegion(BufferVar buffer, ffi::Array region) { TVM_FFI_ICHECK_EQ(buffer->shape.size(), region.size()) << "The dimension between " << buffer << " and region " << region << " mismatched, the buffer is " << buffer; @@ -514,7 +517,7 @@ BufferRegion::BufferRegion(Buffer buffer, ffi::Array region) { data_ = std::move(node); } -BufferRegion BufferRegion::FullRegion(Buffer buffer) { +BufferRegion BufferRegion::FullRegion(BufferVar buffer) { ffi::Array region; for (PrimExpr extent : buffer->shape) { region.push_back(Range::FromMinExtent(0, extent)); @@ -522,7 +525,7 @@ BufferRegion BufferRegion::FullRegion(Buffer buffer) { return BufferRegion(buffer, region); } -BufferRegion BufferRegion::FromPoint(Buffer buffer, ffi::Array indices) { +BufferRegion BufferRegion::FromPoint(BufferVar buffer, ffi::Array indices) { ffi::Array region; for (const PrimExpr& index : indices) { if (const RampNode* ramp_index = index.as()) { @@ -537,14 +540,14 @@ BufferRegion BufferRegion::FromPoint(Buffer buffer, ffi::Array indices TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("tirx.BufferRegion", [](Buffer buffer, ffi::Array region) { + refl::GlobalDef().def("tirx.BufferRegion", [](BufferVar buffer, ffi::Array region) { return BufferRegion(buffer, region); }); } // MatchBufferRegion -MatchBufferRegion::MatchBufferRegion(Buffer buffer, BufferRegion source) { - const Buffer& source_buffer = source->buffer; +MatchBufferRegion::MatchBufferRegion(BufferVar buffer, BufferRegion source) { + const BufferVar& source_buffer = source->buffer; arith::Analyzer analyzer; // Check scope and dtype TVM_FFI_ICHECK_EQ(buffer.scope(), source_buffer.scope()) @@ -591,7 +594,7 @@ MatchBufferRegion::MatchBufferRegion(Buffer buffer, BufferRegion source) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("tirx.MatchBufferRegion", [](Buffer buffer, BufferRegion source) { + refl::GlobalDef().def("tirx.MatchBufferRegion", [](BufferVar buffer, BufferRegion source) { return MatchBufferRegion(buffer, source); }); } @@ -599,7 +602,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { // Block SBlock::SBlock(ffi::Array iter_vars, ffi::Array reads, ffi::Array writes, ffi::String name_hint, Stmt body, - ffi::Optional init, ffi::Array alloc_buffers, + ffi::Optional init, ffi::Array alloc_buffers, ffi::Array match_buffers, ffi::Map annotations, Span span) { ffi::ObjectPtr node = ffi::make_object(); @@ -616,7 +619,7 @@ SBlock::SBlock(ffi::Array iter_vars, ffi::Array reads, data_ = std::move(node); } -SBlock::SBlock(ffi::String name_hint, Stmt body, ffi::Array alloc_buffers, Span span) { +SBlock::SBlock(ffi::String name_hint, Stmt body, ffi::Array alloc_buffers, Span span) { ffi::ObjectPtr node = ffi::make_object(); node->iter_vars = {}; node->reads = {}; @@ -636,7 +639,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("tirx.SBlock", [](ffi::Array iter_vars, ffi::Array reads, ffi::Array writes, ffi::String name_hint, Stmt body, - ffi::Optional init, ffi::Array alloc_buffers, + ffi::Optional init, ffi::Array alloc_buffers, ffi::Array match_buffers, ffi::Map annotations, Span span) { return SBlock(iter_vars, reads, writes, name_hint, body, init, diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc index 7ae935601f89..904d11f0b8a7 100644 --- a/src/tirx/ir/stmt_functor.cc +++ b/src/tirx/ir/stmt_functor.cc @@ -65,17 +65,30 @@ void StmtVisitor::VisitStmt_(const BreakNode* op) {} void StmtVisitor::VisitStmt_(const ContinueNode* op) {} -void StmtVisitor::VisitBufferDef(const Buffer& buffer, bool alloc_data) { +void StmtVisitor::VisitBufferDef(const BufferVar& buffer, bool alloc_data) { for (const auto& e : buffer->shape) this->VisitExpr(e); for (const auto& e : buffer->strides) this->VisitExpr(e); this->VisitExpr(buffer->elem_offset); + for (const auto& e : buffer->allocated_addr) this->VisitExpr(e); + if (buffer->layout.has_value()) { + const auto* layout = buffer->layout.value().as(); + if (layout == nullptr) return; + for (const Iter& iter : layout->shard) { + this->VisitExpr(iter->extent); + this->VisitExpr(iter->stride); + } + for (const Iter& iter : layout->replica) { + this->VisitExpr(iter->extent); + this->VisitExpr(iter->stride); + } + } } // Default VisitBufferUse is empty: buffer fields (shape, strides, elem_offset) // are visited at the definition site (VisitBufferDef) and should not be // re-visited at each use site, as the use site may be in a different scope // where the buffer's shape variables are not defined. -void StmtVisitor::VisitBufferUse(const Buffer& buffer) {} +void StmtVisitor::VisitBufferUse(const BufferVar& buffer) {} void StmtExprVisitor::VisitExpr_(const BufferLoadNode* op) { this->VisitBufferUse(op->buffer); @@ -87,6 +100,9 @@ void StmtVisitor::VisitStmt_(const AllocBufferNode* op) { } void StmtVisitor::VisitStmt_(const DeclBufferNode* op) { + if (op->data.has_value()) { + this->VisitExpr(op->data.value()); + } this->VisitBufferDef(op->buffer, /*alloc_data=*/false); } @@ -129,7 +145,7 @@ void StmtVisitor::VisitStmt_(const SBlockNode* op) { this->VisitExpr(iter_var->dom->extent); }); VisitArray(op->alloc_buffers, - [this](const Buffer& buf) { this->VisitBufferDef(buf, /*alloc_data=*/true); }); + [this](const BufferVar& buf) { this->VisitBufferDef(buf, /*alloc_data=*/true); }); VisitArray(op->reads, fvisit_buffer_region); VisitArray(op->writes, fvisit_buffer_region); VisitArray(op->match_buffers, @@ -257,7 +273,7 @@ class StmtMutator::Internal { static ffi::Array Mutate(StmtMutator* self, const ffi::Array& arr) { auto fmutate = [self](const BufferRegion& buffer_region) { - Buffer new_buf = self->VisitBufferUse(buffer_region->buffer); + BufferVar new_buf = self->VisitBufferUse(buffer_region->buffer); ffi::Array region = Mutate(self, buffer_region->region); if (new_buf.same_as(buffer_region->buffer) && region.same_as(buffer_region->region)) { return buffer_region; @@ -271,8 +287,8 @@ class StmtMutator::Internal { static ffi::Array Mutate(StmtMutator* self, const ffi::Array& arr) { auto fmutate = [self](const MatchBufferRegion& match_buffer_region) { - Buffer new_buf = self->VisitBufferDef(match_buffer_region->buffer, /*alloc_data=*/true); - Buffer new_source_buf = self->VisitBufferUse(match_buffer_region->source->buffer); + BufferVar new_buf = self->VisitBufferDef(match_buffer_region->buffer, /*alloc_data=*/true); + BufferVar new_source_buf = self->VisitBufferUse(match_buffer_region->source->buffer); ffi::Array region = Mutate(self, match_buffer_region->source->region); if (new_buf.same_as(match_buffer_region->buffer) && new_source_buf.same_as(match_buffer_region->source->buffer) && @@ -361,7 +377,7 @@ Stmt StmtMutator::VisitStmt_(const BreakNode* op) { return ffi::GetRef(op) Stmt StmtMutator::VisitStmt_(const ContinueNode* op) { return ffi::GetRef(op); } -Buffer StmtMutator::VisitBufferDef(const Buffer& buffer, bool alloc_data) { +BufferVar StmtMutator::VisitBufferDef(const BufferVar& buffer, bool alloc_data) { if (auto it = buffer_remap_.find(buffer); it != buffer_remap_.end()) { return (*it).second; } @@ -372,6 +388,8 @@ Buffer StmtMutator::VisitBufferDef(const Buffer& buffer, bool alloc_data) { auto shape = buffer->shape.Map([this](const PrimExpr& e) { return this->VisitPrimExpr(e); }); auto strides = buffer->strides.Map([this](const PrimExpr& e) { return this->VisitPrimExpr(e); }); PrimExpr elem_offset = this->VisitPrimExpr(buffer->elem_offset); + auto allocated_addr = + buffer->allocated_addr.Map([this](const PrimExpr& e) { return this->VisitPrimExpr(e); }); // Visit the layout's per-iter extent/stride PrimExprs too: they share dtype // semantics with the shape, e.g. ``IndexDataTypeRewriter`` (int32 -> int64) @@ -399,30 +417,37 @@ Buffer StmtMutator::VisitBufferDef(const Buffer& buffer, bool alloc_data) { } if (shape.same_as(buffer->shape) && strides.same_as(buffer->strides) && - elem_offset.same_as(buffer->elem_offset) && !layout_changed) { + elem_offset.same_as(buffer->elem_offset) && + allocated_addr.same_as(buffer->allocated_addr) && !layout_changed) { return buffer; } - Buffer new_buf = buffer; - auto* n = new_buf.CopyOnWrite(); - n->shape = std::move(shape); - n->strides = std::move(strides); - n->elem_offset = std::move(elem_offset); - if (layout_changed) { - n->layout = std::move(new_layout); - } + BufferType new_type( + buffer->data_pointer_type, buffer->dtype, std::move(shape), + std::move(strides), std::move(elem_offset), buffer->data_alignment, + buffer->offset_factor, std::move(new_layout), std::move(allocated_addr), + buffer->span); + BufferVar new_buf(buffer.name(), std::move(new_type), buffer.span()); buffer_remap_.Set(buffer, new_buf); return new_buf; } -Buffer StmtMutator::VisitBufferUse(const Buffer& buffer) { +BufferVar StmtMutator::VisitBufferUse(const BufferVar& buffer) { if (auto it = buffer_remap_.find(buffer); it != buffer_remap_.end()) { return (*it).second; } return buffer; } +Expr StmtExprMutator::VisitExpr_(const VarNode* op) { + Var var = ffi::GetRef(op); + if (var->ty.as()) { + return VisitBufferUse(BufferVar(var)).var(); + } + return var; +} + Expr StmtExprMutator::VisitExpr_(const BufferLoadNode* op) { - Buffer new_buf = this->VisitBufferUse(op->buffer); + BufferVar new_buf = this->VisitBufferUse(op->buffer); PrimExpr expr = ExprMutator::VisitExpr_(op).as_or_throw(); op = expr.as(); TVM_FFI_ICHECK(op != nullptr); @@ -435,7 +460,7 @@ Expr StmtExprMutator::VisitExpr_(const BufferLoadNode* op) { } Stmt StmtMutator::VisitStmt_(const AllocBufferNode* op) { - Buffer new_buf = this->VisitBufferDef(op->buffer, /*alloc_data=*/true); + BufferVar new_buf = this->VisitBufferDef(op->buffer, /*alloc_data=*/true); if (new_buf.same_as(op->buffer)) { return ffi::GetRef(op); @@ -447,12 +472,17 @@ Stmt StmtMutator::VisitStmt_(const AllocBufferNode* op) { } Stmt StmtMutator::VisitStmt_(const DeclBufferNode* op) { - Buffer new_buf = this->VisitBufferDef(op->buffer, /*alloc_data=*/false); + ffi::Optional data = std::nullopt; + if (op->data.has_value()) { + data = this->VisitExpr(op->data.value()); + } + BufferVar new_buf = this->VisitBufferDef(op->buffer, /*alloc_data=*/false); - if (new_buf.same_as(op->buffer)) { + if (new_buf.same_as(op->buffer) && data.same_as(op->data)) { return ffi::GetRef(op); } else { auto n = CopyOnWrite(op); + n->data = std::move(data); n->buffer = std::move(new_buf); return Stmt(n); } @@ -478,7 +508,7 @@ Stmt StmtMutator::VisitStmt_(const IfThenElseNode* op) { } Stmt StmtMutator::VisitStmt_(const BufferStoreNode* op) { - Buffer new_buf = this->VisitBufferUse(op->buffer); + BufferVar new_buf = this->VisitBufferUse(op->buffer); PrimExpr value = this->VisitPrimExpr(op->value); ffi::Array indices = Internal::Mutate(this, op->indices); @@ -579,9 +609,9 @@ Stmt StmtMutator::VisitStmt_(const EvaluateNode* op) { Stmt StmtMutator::VisitStmt_(const SBlockNode* op) { ffi::Array iter_vars = Internal::Mutate(this, op->iter_vars); - ffi::Array alloc_buffers = Internal::MutateArray( + ffi::Array alloc_buffers = Internal::MutateArray( this, op->alloc_buffers, - [this](const Buffer& buf) { return this->VisitBufferDef(buf, /*alloc_data=*/true); }); + [this](const BufferVar& buf) { return this->VisitBufferDef(buf, /*alloc_data=*/true); }); ffi::Array reads = Internal::Mutate(this, op->reads); ffi::Array writes = Internal::Mutate(this, op->writes); ffi::Array match_buffers = Internal::Mutate(this, op->match_buffers); @@ -705,8 +735,8 @@ class IRApplyVisit : public StmtExprVisitor { f_(node); } - void VisitBufferDef(const Buffer& buffer, bool alloc_data) override {} - void VisitBufferUse(const Buffer& buffer) override {} + void VisitBufferDef(const BufferVar& buffer, bool alloc_data) override {} + void VisitBufferUse(const BufferVar& buffer) override {} private: std::function f_; @@ -795,28 +825,37 @@ class IRSubstitute : public StmtExprMutator { } return ret.value(); } - return var; - } - - // Override VisitBufferDef to also remap buffer->data (the backing allocation var). - // The base class only visits shape/strides/elem_offset. - Buffer VisitBufferDef(const Buffer& buffer, bool alloc_data) final { - Buffer new_buf = StmtExprMutator::VisitBufferDef(buffer, alloc_data); - // Additionally handle data var substitution (base does not visit data). - Expr new_data_expr = VisitExpr(new_buf->data); - auto new_data = new_data_expr.as(); - TVM_FFI_ICHECK(new_data) << "Buffer " << new_buf << " uses backing allocation " << new_buf->data - << ", which was substituted into the expression " << new_data_expr - << " and the backing allocation must be a tirx::Var"; - Var data = new_data.value(); - if (!data.same_as(new_buf->data)) { - auto* n = new_buf.CopyOnWrite(); - n->data = std::move(data); + return StmtExprMutator::VisitExpr_(op); + } + + // Buffer variables share the ordinary Var identity model. A caller-provided + // substitution may therefore replace the definition with another checked + // BufferVar; metadata-only rewrites are handled by the base implementation. + BufferVar VisitBufferDef(const BufferVar& buffer, bool alloc_data) final { + BufferVar new_buf = StmtExprMutator::VisitBufferDef(buffer, alloc_data); + if (auto mapped = vmap_(new_buf.var())) { + auto mapped_var = mapped.value().as(); + TVM_FFI_ICHECK(mapped_var && mapped_var.value()->ty.as()) + << "BufferVar " << new_buf << " was substituted into " + << mapped.value() << ", which is not a Var with BufferType"; + new_buf = BufferVar(mapped_var.value()); buffer_remap_.Set(buffer, new_buf); } return new_buf; } + BufferVar VisitBufferUse(const BufferVar& buffer) final { + BufferVar new_buf = StmtExprMutator::VisitBufferUse(buffer); + if (auto mapped = vmap_(new_buf.var())) { + auto mapped_var = mapped.value().as(); + TVM_FFI_ICHECK(mapped_var && mapped_var.value()->ty.as()) + << "BufferVar " << new_buf << " was substituted into " + << mapped.value() << ", which is not a Var with BufferType"; + new_buf = BufferVar(mapped_var.value()); + } + return new_buf; + } + Stmt VisitStmt_(const AttrStmtNode* op) final { Stmt ret = StmtExprMutator::VisitStmt_(op); op = ret.as(); @@ -898,7 +937,7 @@ class IRSubstituteWithDataTypeLegalization : public DataTypeLegalizer { if (ret.has_value()) { return ret.value(); } - return var; + return StmtExprMutator::VisitExpr_(op); } Stmt VisitStmt_(const AttrStmtNode* op) final { diff --git a/src/tirx/ir/tir_visitor_with_path.cc b/src/tirx/ir/tir_visitor_with_path.cc index 0d0df45f2b39..dd9bcc88d15f 100644 --- a/src/tirx/ir/tir_visitor_with_path.cc +++ b/src/tirx/ir/tir_visitor_with_path.cc @@ -82,11 +82,16 @@ void TIRVisitorWithPath::Visit(const PrimFunc& func, AccessPath path) { // variable has occurred. Therefore, to ensure that we only avoid // duplicate calls to VisitVarDef, these semantics need to be // checked. - std::vector, DefContext>> context; + std::vector, DefContext>> context; auto ppath = path->Attr("params"); for (size_t i = 0; i < func->params.size(); i++) { - context.push_back(WithDef(func->params[i], ppath->ArrayItem(i))); + const Var& param = func->params[i]; + if (param->ty.as()) { + context.push_back(WithDef(BufferVar(param), ppath->ArrayItem(i))); + } else { + context.push_back(WithDef(param, ppath->ArrayItem(i))); + } } auto buffer_map_path = path->Attr("buffer_map"); @@ -126,26 +131,31 @@ void TIRVisitorWithPath::ExitDef(const IterVar& iter_var, AccessPath path) { ExitDef(iter_var->var, path->Attr("var")); } -void TIRVisitorWithPath::EnterDef(const Buffer& buffer, AccessPath path) { +void TIRVisitorWithPath::EnterDef(const BufferVar& buffer, AccessPath path) { + // BufferVar is a checked view over an ordinary Var. Its definition + // therefore introduces both the variable identity and the buffer metadata. + EnterDef(buffer.var(), path); // Defining a buffer counts as using all parameters in the buffer // (e.g. shape/strides). VisitBufferDef(buffer, path); } -void TIRVisitorWithPath::ExitDef(const Buffer& buffer, AccessPath path) {} +void TIRVisitorWithPath::ExitDef(const BufferVar& buffer, AccessPath path) { + ExitDef(buffer.var(), path); +} -void TIRVisitorWithPath::VisitBufferDef(const Buffer& buffer, AccessPath path) { - Visit(buffer->data, path->Attr("data")); +void TIRVisitorWithPath::VisitBufferDef(const BufferVar& buffer, AccessPath path) { Visit(buffer->shape, path->Attr("shape")); Visit(buffer->strides, path->Attr("strides")); Visit(buffer->elem_offset, path->Attr("elem_offset")); + Visit(buffer->allocated_addr, path->Attr("allocated_addr")); } -// Default: buffer use sites do not re-visit buffer fields. Buffer fields +// Default: buffer use sites do not re-visit buffer fields. BufferVar fields // (shape, strides, elem_offset) are visited at the definition site via // VisitBufferDef/EnterDef. Re-visiting at use sites would require those // variables to be in scope at every use, which may not hold when buffers // are allocated in a different scope than where they are used. -void TIRVisitorWithPath::VisitBufferUse(const Buffer& buffer, AccessPath path) {} +void TIRVisitorWithPath::VisitBufferUse(const BufferVar& buffer, AccessPath path) {} void TIRVisitorWithPath::Visit(const BufferRegion& region, AccessPath path) { VisitBufferUse(region->buffer, path->Attr("buffer")); @@ -182,7 +192,7 @@ void TIRVisitorWithPath::VisitStmt_(const BindNode* op, AccessPath path) { void TIRVisitorWithPath::VisitStmt_(const AttrStmtNode* op, AccessPath path) { Visit(op->value, path->Attr("value")); - std::vector, DefContext, DefContext>> context; + std::vector, DefContext, DefContext>> context; if (auto iter_var = op->node.as(); iter_var && (op->attr_key == attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread)) { @@ -221,14 +231,13 @@ void TIRVisitorWithPath::VisitStmt_(const BreakNode* op, AccessPath path) {} void TIRVisitorWithPath::VisitStmt_(const ContinueNode* op, AccessPath path) {} void TIRVisitorWithPath::VisitStmt_(const AllocBufferNode* op, AccessPath path) { - // AllocBuffer both allocates the data variable and declares the buffer. // Push definitions into the current scope so they are visible to subsequent siblings. auto buf_path = path->Attr("buffer"); - bind_scope_.Current().push_back(WithDef(op->buffer->data, buf_path->Attr("data"))); bind_scope_.Current().push_back(WithDef(op->buffer, buf_path)); } void TIRVisitorWithPath::VisitStmt_(const DeclBufferNode* op, AccessPath path) { + Visit(op->data, path->Attr("data")); // Push buffer definition into the current scope so it is visible to subsequent siblings. bind_scope_.Current().push_back(WithDef(op->buffer, path->Attr("buffer"))); } @@ -263,7 +272,7 @@ void TIRVisitorWithPath::VisitStmt_(const EvaluateNode* op, AccessPath path) { } void TIRVisitorWithPath::VisitStmt_(const SBlockNode* op, AccessPath path) { - std::vector, DefContext, DefContext>> context; + std::vector, DefContext, DefContext>> context; { auto iter_path = path->Attr("iter_vars"); @@ -279,7 +288,6 @@ void TIRVisitorWithPath::VisitStmt_(const SBlockNode* op, AccessPath path) { for (size_t i = 0; i < op->alloc_buffers.size(); i++) { auto buffer_path = alloc_path->ArrayItem(i); auto buf = op->alloc_buffers[i]; - context.push_back(WithDef(buf->data, buffer_path->Attr("data"))); context.push_back(WithDef(buf, buffer_path)); } } @@ -325,7 +333,7 @@ void TIRVisitorWithPath::VisitStmt_(const tirx::TilePrimitiveCallNode* op, Acces Visit(expr.value(), path->Attr("args")->ArrayItem(i)); } else if (auto stmt = op->args[i].as()) { Visit(stmt.value(), path->Attr("args")->ArrayItem(i)); - } else if (auto buf = op->args[i].as()) { + } else if (auto buf = op->args[i].as()) { VisitBufferUse(buf.value(), path->Attr("args")->ArrayItem(i)); } } diff --git a/src/tirx/ir/tir_visitor_with_path.h b/src/tirx/ir/tir_visitor_with_path.h index eb3974e4ee26..860759dc4b5b 100644 --- a/src/tirx/ir/tir_visitor_with_path.h +++ b/src/tirx/ir/tir_visitor_with_path.h @@ -74,9 +74,9 @@ class TIRVisitorWithPath : protected ExprFunctor @@ -242,7 +242,7 @@ class TIRVisitorWithPath : protected ExprFunctor> WithMatchBufferDefs(Buffer buf, ffi::reflection::AccessPath path) { + std::vector> WithMatchBufferDefs(BufferVar buf, ffi::reflection::AccessPath path) { std::vector> context; auto try_visit_implicit_var_def = [this, &context](const Expr& expr, @@ -262,7 +262,6 @@ class TIRVisitorWithPath : protected ExprFunctordata, path->Attr("data")); try_visit_implicit_var_def_array(buf->shape, path->Attr("shape")); try_visit_implicit_var_def_array(buf->strides, path->Attr("strides")); try_visit_implicit_var_def(buf->elem_offset, path->Attr("elem_offset")); @@ -278,7 +277,7 @@ class TIRVisitorWithPath : protected ExprFunctor, DefContext>; + using BindScopeEntry = std::variant, DefContext>; ScopeStack> bind_scope_; }; diff --git a/src/tirx/ir/tirx_stmt.cc b/src/tirx/ir/tirx_stmt.cc index e2838fe2efa1..d8be8942479d 100644 --- a/src/tirx/ir/tirx_stmt.cc +++ b/src/tirx/ir/tirx_stmt.cc @@ -33,7 +33,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { TilePrimitiveCallNode::RegisterReflection(); } // TilePrimitiveCall TilePrimitiveCall::TilePrimitiveCall(tvm::Op op, ffi::Array args, - ffi::Map workspace, + ffi::Map workspace, ffi::Map config, ffi::Optional dispatch, ExecScope scope) { TVM_FFI_CHECK(op.defined(), ValueError) << "TilePrimitiveCall expects a defined operator"; @@ -50,7 +50,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def( "tirx.TilePrimitiveCall", - [](tvm::Op op, ffi::Array args, ffi::Map workspace, + [](tvm::Op op, ffi::Array args, ffi::Map workspace, ffi::Map config, ffi::Optional dispatch, ExecScope scope) { return TilePrimitiveCall(op, args, workspace, config, dispatch, scope); diff --git a/src/tirx/op/builtin.cc b/src/tirx/op/builtin.cc index a5d7b342cee8..737d857e09f5 100644 --- a/src/tirx/op/builtin.cc +++ b/src/tirx/op/builtin.cc @@ -388,6 +388,10 @@ TIR_DEFINE_BUILTIN_FUNC(buffer_offset) .set_num_inputs(2) .set_attr("TCallEffectKind", static_cast(CallEffectKind::kPure)); +TIR_DEFINE_BUILTIN_FUNC(buffer_data) + .set_num_inputs(1) + .set_attr("TCallEffectKind", static_cast(CallEffectKind::kPure)); + TIR_DEFINE_BUILTIN_FUNC(print_buffer) .set_attr("TCallEffectKind", static_cast(CallEffectKind::kOpaque)); diff --git a/src/tirx/op/tirx.cc b/src/tirx/op/tirx.cc index 2bfba2fc7cfa..1f96ad3e0008 100644 --- a/src/tirx/op/tirx.cc +++ b/src/tirx/op/tirx.cc @@ -60,8 +60,8 @@ Value getOrSetDefault(ffi::Map& m, const Key& key, /********************* DispatchContext **********************/ -void DispatchContextNode::AddAllocBuffer(Buffer buffer) { - auto buffers = getOrSetDefault(callbacks, callback::kPrivateAlloc, ffi::Array()); +void DispatchContextNode::AddAllocBuffer(BufferVar buffer) { + auto buffers = getOrSetDefault(callbacks, callback::kPrivateAlloc, ffi::Array()); buffers.push_back(buffer); callbacks.Set(callback::kPrivateAlloc, buffers); } @@ -73,9 +73,9 @@ void DispatchContextNode::AddInitStmt(Stmt stmt, bool host) { callbacks.Set(tag, stmts); } -void DispatchContextNode::AddPostBufferDefStmt(Buffer buffer, Stmt stmt) { +void DispatchContextNode::AddPostBufferDefStmt(BufferVar buffer, Stmt stmt) { auto mapping = getOrSetDefault(callbacks, callback::kPostBufferDefStmt, - ffi::Map>()); + ffi::Map>()); auto it = mapping.find(buffer); ffi::Array stmts; if (it != mapping.end()) { diff --git a/src/tirx/script/builder/frame.cc b/src/tirx/script/builder/frame.cc index e99bdc29a582..63647ee1ada4 100644 --- a/src/tirx/script/builder/frame.cc +++ b/src/tirx/script/builder/frame.cc @@ -37,11 +37,11 @@ namespace tirx { namespace { // In s_tir functions, buffer-typed parameters must not carry a layout (the -// s_tir IR doesn't track per-buffer layouts on params). When `T.Buffer(...)` is +// s_tir IR doesn't track per-buffer layouts on params). When `T.BufferVar(...)` is // used as a parameter annotation, the parser evaluates the annotation outside // the PrimFunc frame; if the annotation captures an outer-scope variable (e.g. // `dtype` in a closure-based generator), the evaluation happens *before* -// `_current_s_tir()` becomes true, so the resulting Buffer is built with the +// `_current_s_tir()` becomes true, so the resulting BufferVar is built with the // default tile layout instead of None. Direct annotations using only literals // are re-evaluated inside the frame and correctly get layout=None. // @@ -51,11 +51,11 @@ namespace { // `buffer_remap_` machinery, so the body remains well-formed. class STirBufferLayoutNormalizer : public tvm::tirx::StmtExprMutator { public: - void Register(const tvm::tirx::Buffer& old_buf, const tvm::tirx::Buffer& new_buf) { + void Register(const tvm::tirx::BufferVar& old_buf, const tvm::tirx::BufferVar& new_buf) { this->buffer_remap_.Set(old_buf, new_buf); } bool Empty() const { return this->buffer_remap_.empty(); } - tvm::tirx::Buffer Lookup(const tvm::tirx::Buffer& buf) const { + tvm::tirx::BufferVar Lookup(const tvm::tirx::BufferVar& buf) const { auto it = this->buffer_remap_.find(buf); if (it != this->buffer_remap_.end()) { return (*it).second; @@ -113,17 +113,19 @@ void PrimFuncFrameNode::ExitWithScope() { } // s_tir-mode normalization: drop stale default layouts (see comment on // STirBufferLayoutNormalizer above) and rewrite body references coherently. - ffi::Map effective_buffer_map = buffer_map; - ffi::Array effective_root_alloc_buffers = root_alloc_buffers; + ffi::Map effective_buffer_map = buffer_map; + ffi::Array effective_root_alloc_buffers = root_alloc_buffers; tvm::tirx::Stmt body = AsStmt(stmts); if (s_tir) { STirBufferLayoutNormalizer normalizer; - ffi::Map new_buffer_map; + ffi::Map new_buffer_map; for (const auto& kv : buffer_map) { - tvm::tirx::Buffer buf = kv.second; + tvm::tirx::BufferVar buf = kv.second; if (buf->layout.has_value()) { - tvm::tirx::Buffer new_buf = buf; - new_buf.CopyOnWrite()->layout = std::nullopt; + ffi::ObjectPtr type = tvm::tirx::CopyBufferType(buf); + type->layout = std::nullopt; + tvm::tirx::BufferVar new_buf = + tvm::tirx::RebuildBufferVar(buf, std::move(type)); normalizer.Register(buf, new_buf); new_buffer_map.Set(kv.first, new_buf); } else { @@ -132,8 +134,8 @@ void PrimFuncFrameNode::ExitWithScope() { } if (!normalizer.Empty()) { body = normalizer(std::move(body)); - ffi::Array new_root_alloc_buffers; - for (const tvm::tirx::Buffer& buf : root_alloc_buffers) { + ffi::Array new_root_alloc_buffers; + for (const tvm::tirx::BufferVar& buf : root_alloc_buffers) { new_root_alloc_buffers.push_back(normalizer.Lookup(buf)); } effective_buffer_map = std::move(new_buffer_map); @@ -177,8 +179,8 @@ void SBlockFrameNode::ExitWithScope() { // Allow SBlock construction in raw IRBuilder context (no enclosing PrimFuncFrame) // so test fixtures can construct blocks/block-realizes directly. - ffi::Array tir_alloc_buffers; - for (const tvm::tirx::Buffer& buffer : alloc_buffers) { + ffi::Array tir_alloc_buffers; + for (const tvm::tirx::BufferVar& buffer : alloc_buffers) { tir_alloc_buffers.push_back(buffer); } ffi::Map attrs = annotations.value_or({}); @@ -297,7 +299,7 @@ void ElseFrameNode::ExitWithScope() { void DeclBufferFrameNode::ExitWithScope() { TIRFrameNode::ExitWithScope(); if (allocated) { - AddToParent(tvm::tirx::SeqStmt::Flatten(tvm::tirx::DeclBuffer(buffer), AsStmt(stmts))); + AddToParent(tvm::tirx::SeqStmt::Flatten(tvm::tirx::DeclBuffer(buffer, data), AsStmt(stmts))); } else { // data is undefined in `decl_buffer(...)`, lower to `alloc_buffer(...)`. AddToParent(tvm::tirx::SeqStmt::Flatten(tvm::tirx::AllocBuffer(buffer), AsStmt(stmts))); diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc index 0e67fd151458..fa973fced3ba 100644 --- a/src/tirx/script/builder/ir.cc +++ b/src/tirx/script/builder/ir.cc @@ -41,33 +41,37 @@ namespace tirx { using tvm::tirx::IterVar; using tvm::tirx::Layout; -Buffer BufferDecl(ffi::Array shape, PrimType dtype, ffi::String buffer_name, - ffi::Optional data, ffi::Optional> strides, - ffi::Optional elem_offset, ffi::String storage_scope, int align, - int offset_factor, ffi::Optional layout, - ffi::Array allocated_addr) { +BufferVar BufferDecl(ffi::Array shape, PrimType dtype, ffi::String buffer_name, + ffi::Optional data, + ffi::Optional> strides, + ffi::Optional elem_offset, ffi::String storage_scope, int align, + int offset_factor, ffi::Optional layout, + ffi::Array allocated_addr) { if (!allocated_addr.empty()) { TVM_FFI_ICHECK(!data.has_value() && !elem_offset.has_value() && !offset_factor) << "ValueError: `allocated_addr` can only be used with `data`, `elem_offset`, and " "`offset_factor` undefined"; } - Var buffer_data; + PointerType data_pointer_type = PointerType::VoidPointerTy(); if (!data.has_value()) { DLDataType storage_dtype = dtype->dtype; if (storage_dtype == DLDataType{kDLBool, 8, 1}) { storage_dtype = DLDataType{kDLInt, 8, 1}; } - buffer_data = tvm::tirx::Var(buffer_name, PointerType(PrimType(storage_dtype), storage_scope)); + data_pointer_type = PointerType(PrimType(storage_dtype), storage_scope); } else { - buffer_data = data.value(); + data_pointer_type = data.value()->ty.as_or_throw(); } if (!elem_offset.has_value() && offset_factor) { PrimType shape_dtype = shape.empty() ? PrimType::Int(32) : shape[0].ty(); elem_offset = tvm::tirx::PrimVar("elem_offset", shape_dtype); } - return Buffer(buffer_data, dtype, shape, strides.value_or(ffi::Array()), - elem_offset.value_or(PrimExpr()), buffer_name, align, offset_factor, Span(), layout, - allocated_addr); + return BufferVar( + buffer_name, + tvm::tirx::BufferType(data_pointer_type, dtype, shape, + strides.value_or(ffi::Array()), + elem_offset.value_or(PrimExpr()), align, offset_factor, layout, + allocated_addr)); } PrimFuncFrame PrimFunc(bool is_private, bool s_tir, bool persistent) { @@ -92,12 +96,12 @@ Var Arg(ffi::String name, Var var) { return var; } -Buffer Arg(ffi::String name, Buffer buffer) { +BufferVar Arg(ffi::String name, BufferVar buffer) { PrimFuncFrame frame = FindPrimFuncFrame("T.Arg"); details::Namer::Name(buffer, name); - // A Buffer parameter is an opaque ABI handle. The Buffer's data Var + // A BufferVar parameter is an opaque ABI handle. The BufferVar's data Var // carries the exact pointee type used within the function body. - Var handle(buffer->name + "_handle", PointerType::VoidPointerTy()); + Var handle(buffer.name() + "_handle", PointerType::VoidPointerTy()); frame->args.push_back(handle); frame->buffer_map.Set(handle, buffer); return buffer; @@ -145,11 +149,11 @@ tvm::Type FuncRet(tvm::Type ret_type) { return ret_type; } -Buffer MatchBuffer(ffi::ObjectRef param, ffi::Array shape, PrimType dtype, - ffi::Optional data, ffi::Array strides, PrimExpr elem_offset, - ffi::String storage_scope, int align, int offset_factor, - ffi::Optional layout) { - Buffer buffer = BufferDecl(shape, dtype, "", data, strides, elem_offset, storage_scope, align, +BufferVar MatchBuffer(ffi::ObjectRef param, ffi::Array shape, PrimType dtype, + ffi::Optional data, ffi::Array strides, + PrimExpr elem_offset, ffi::String storage_scope, int align, + int offset_factor, ffi::Optional layout) { + BufferVar buffer = BufferDecl(shape, dtype, "", data, strides, elem_offset, storage_scope, align, offset_factor, layout, {}); if (auto var = param.as()) { PrimFuncFrame frame = FindPrimFuncFrame("T.match_buffer"); @@ -366,8 +370,8 @@ void BlockAttrs(ffi::Map attrs) { << "frame, but T.sblock_attr occurred outside of any such frame"; } -ffi::Variant SBlockAllocBuffer( - ffi::Array shape, PrimType dtype, ffi::Optional data, +ffi::Variant SBlockAllocBuffer( + ffi::Array shape, PrimType dtype, ffi::Optional data, ffi::Array strides, PrimExpr elem_offset, ffi::String storage_scope, int align, int offset_factor, ffi::Optional layout, ffi::Array allocated_addr) { std::string scope = static_cast(storage_scope); @@ -381,7 +385,7 @@ ffi::Variant SBlockAllocBuffer( } ffi::Optional opt_elem_offset = elem_offset.defined() ? ffi::Optional(elem_offset) : std::nullopt; - Buffer buffer = BufferDecl(shape, dtype, "", std::nullopt, strides, opt_elem_offset, + BufferVar buffer = BufferDecl(shape, dtype, "", std::nullopt, strides, opt_elem_offset, storage_scope, align, offset_factor, layout, allocated_addr); IRBuilder builder = IRBuilder::Current(); auto opt_func_frame = builder->FindFrame(); @@ -715,7 +719,7 @@ HintFrame Hint(ffi::String message, ffi::Map attrs) { return HintFrame(n); } -ComposeOpFrame ComposeOp(ffi::Map workspace, +ComposeOpFrame ComposeOp(ffi::Map workspace, ffi::Map config, ffi::Optional dispatch) { ffi::ObjectPtr n = ffi::make_object(); @@ -737,7 +741,7 @@ Var EnvThread(ffi::String thread_tag, PrimType dtype) { return var; } -void BufferStore(Buffer buffer, PrimExpr value, ffi::Array indices, +void BufferStore(BufferVar buffer, PrimExpr value, ffi::Array indices, ffi::Optional predicate = std::nullopt) { PrimType buffer_dtype = buffer->dtype; PrimType index_ty = indices.empty() ? PrimType::Int(32) : indices.back().ty(); @@ -808,7 +812,8 @@ void BufferStore(Buffer buffer, PrimExpr value, ffi::Array indices, } DeclBufferFrame DeclBuffer(ffi::Array shape, PrimType dtype, ffi::String buffer_name, - ffi::Optional data, ffi::Optional> strides, + ffi::Optional data, + ffi::Optional> strides, ffi::Optional elem_offset, ffi::String storage_scope, int align, int offset_factor, ffi::Optional layout, ffi::Optional allocated_addr) { @@ -842,14 +847,17 @@ DeclBufferFrame DeclBuffer(ffi::Array shape, PrimType dtype, ffi::Stri ffi::ObjectPtr n = ffi::make_object(); n->buffer = BufferDecl(shape, dtype, buffer_name, data, strides, elem_offset, storage_scope, align, offset_factor, layout, allocated_addr_arr); + if (data.has_value()) { + n->data = data.value(); + } // For tmem, even without `data`, we should not emit an Allocate node. n->allocated = (scope == "tmem") || data.has_value(); return DeclBufferFrame(n); } -Buffer AllocBuffer(ffi::Array shape, PrimType dtype, ffi::String storage_scope, +BufferVar AllocBuffer(ffi::Array shape, PrimType dtype, ffi::String storage_scope, ffi::Optional> annotations) { - Buffer buffer = BufferDecl(shape, dtype, "", std::nullopt, std::nullopt, std::nullopt, + BufferVar buffer = BufferDecl(shape, dtype, "", std::nullopt, std::nullopt, std::nullopt, storage_scope, 0, 0, std::nullopt, {}); AddToParent( tvm::tirx::AllocBuffer(buffer, annotations.value_or(ffi::Map()))); @@ -865,28 +873,6 @@ Var Ptr(PrimType dtype, ffi::String storage_scope = "global") { using tvm::script::ir_builder::details::Namer; -TVM_STATIC_IR_FUNCTOR(Namer, vtable) - .set_dispatch([](const ffi::ObjectRef& node, ffi::String name) -> void { - tvm::tirx::BufferNode* buffer = - const_cast(node.as()); - if (!buffer->name.empty() && buffer->name != std::string(name)) { - TVM_FFI_THROW(InternalError) - << "Buffer name conflict: buffer was created with name \"" << buffer->name - << "\", but the parser is trying to rename it to \"" << name - << "\". Remove the explicit `name=` argument and let the parser " - << "auto-name the buffer from the LHS variable."; - } - buffer->name = name; - Namer::Name(buffer->data, name + "_ptr"); - int n = buffer->strides.size(); - for (int i = 0; i < n; ++i) { - PrimExpr e = buffer->strides[i]; - if (auto v = e.as()) { - Namer::Name(v.value(), name + "_s" + std::to_string(i)); - } - } - }); - TVM_STATIC_IR_FUNCTOR(Namer, vtable) .set_dispatch([](const ffi::ObjectRef& node, ffi::String name) -> void { @@ -912,20 +898,20 @@ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() .def("script.ir_builder.tirx.Buffer", - static_cast, PrimType, ffi::String, ffi::Optional, - ffi::Optional>, ffi::Optional, - ffi::String, int, int, ffi::Optional, - ffi::Array)>(BufferDecl)) + static_cast, PrimType, ffi::String, ffi::Optional, + ffi::Optional>, ffi::Optional, ffi::String, int, + int, ffi::Optional, ffi::Array)>(BufferDecl)) .def("script.ir_builder.tirx.PrimFunc", PrimFunc) .def("script.ir_builder.tirx.Arg", [](ffi::String name, ffi::ObjectRef obj) -> ffi::ObjectRef { using namespace tvm::tirx; + if (auto buffer = obj.as()) { + return Arg(name, buffer.value()); + } if (auto var = obj.as()) { return Arg(name, var.value()); } - if (auto buffer = obj.as()) { - return Arg(name, buffer.value()); - } TVM_FFI_THROW(InternalError) << "ValueError: Unexpected type for TIR Arg: " << obj->GetTypeKey(); throw; diff --git a/src/tirx/script/printer/block.cc b/src/tirx/script/printer/block.cc index 3c0f64a9eafd..79271f10629e 100644 --- a/src/tirx/script/printer/block.cc +++ b/src/tirx/script/printer/block.cc @@ -179,7 +179,7 @@ Doc PrintBlock(IRDocsifier d, tirx::SBlock block, AccessPath block_p, // } // Step 5. Handle `alloc_buffer` for (int i = 0, n = block->alloc_buffers.size(); i < n; ++i) { - tirx::Buffer buffer = block->alloc_buffers[i]; + tirx::BufferVar buffer = block->alloc_buffers[i]; AccessPath buffer_p = block_p->Attr("alloc_buffers")->ArrayItem(i); IdDoc lhs = DefineBuffer(buffer, *frame, d); ExprDoc rhs = BufferDecl(buffer, "sblock_alloc_buffer", {}, buffer_p, *frame, d, diff --git a/src/tirx/script/printer/buffer.cc b/src/tirx/script/printer/buffer.cc index 732de462a285..b251bc4eb448 100644 --- a/src/tirx/script/printer/buffer.cc +++ b/src/tirx/script/printer/buffer.cc @@ -26,9 +26,10 @@ namespace tvm { namespace script { namespace printer { -ffi::Map BufferAttrs(tirx::Buffer buffer, const AccessPath& buffer_p, +ffi::Map BufferAttrs(tirx::BufferVar buffer, const AccessPath& buffer_p, const Frame& frame, const IRDocsifier& d, - BufferVarDefinition var_definitions) { + BufferVarDefinition var_definitions, + ffi::Optional data = std::nullopt) { using tvm::tirx::Var; using tvm::tirx::VarNode; ffi::Map kwargs; @@ -45,7 +46,9 @@ ffi::Map BufferAttrs(tirx::Buffer buffer, const AccessPath }); }; update_use_count(buffer->elem_offset); - update_use_count(buffer->data); + if (data.has_value()) { + update_use_count(data.value()); + } for (const PrimExpr& e : buffer->strides) { update_use_count(e); } @@ -98,23 +101,21 @@ ffi::Map BufferAttrs(tirx::Buffer buffer, const AccessPath } // Step 3. Handle `buffer.data` // For tmem scope, DeclBuffer does not accept `data` (it auto-creates the data var). - bool is_tmem_scope = false; - if (auto* ptr_type = buffer->data->ty.as()) { - is_tmem_scope = (ptr_type->storage_scope == "tmem"); - } + bool is_tmem_scope = buffer.scope() == "tmem"; bool is_inline_data = false; - if (!is_tmem_scope) { - if (is_new_var(buffer->data)) { + if (!is_tmem_scope && data.has_value()) { + Expr source = data.value(); + if (is_new_var(source)) { if (var_definitions >= BufferVarDefinition::DataPointer) { - is_inline_data = try_inline_def(buffer->data, buffer_p->Attr("data"), [=]() { + is_inline_data = try_inline_def(source, buffer_p->Attr("data"), [=]() { return d->AsDoc(buffer, buffer_p)->Attr("data"); }); } else { - add_out_of_line_var_def(buffer->data, buffer_p->Attr("data")); + add_out_of_line_var_def(source.as_or_throw(), buffer_p->Attr("data")); } } if (!is_inline_data) { - kwargs.Set("data", d->AsDoc(buffer->data, buffer_p->Attr("data"))); + kwargs.Set("data", d->AsDoc(source, buffer_p->Attr("data"))); } } // Step 4. Handle `buffer.strides` @@ -206,16 +207,16 @@ ffi::Map BufferAttrs(tirx::Buffer buffer, const AccessPath // Unwrap single-element array: DeclBuffer expects Optional, not Array. // For BufferLoad from scalar buffers, we must explicitly print buf[idx] because // the scalar shorthand (which drops the index) produces just the variable name, - // and the parser resolves that to a Buffer object rather than a PrimExpr value. + // and the parser resolves that to a BufferVar object rather than a PrimExpr value. PrimExpr addr = buffer->allocated_addr[0]; AccessPath addr_p = buffer_p->Attr("allocated_addr")->ArrayItem(0); if (const auto* bl = addr.as()) { - // Ensure the buffer variable is defined (may emit a T.Buffer(...) statement). + // Ensure the buffer variable is defined (may emit a T.BufferVar(...) statement). d->AsDoc(bl->buffer, addr_p->Attr("buffer")); // Get the variable name bound to this buffer. ffi::Optional buf_var = d->GetVarDoc(bl->buffer); TVM_FFI_ICHECK(buf_var.has_value()) - << "Buffer in allocated_addr is not defined: " << bl->buffer; + << "BufferVar in allocated_addr is not defined: " << bl->buffer; // Build var[indices] explicitly instead of going through the default BufferLoad // printer, which would use the scalar shorthand and drop the index. int n_idx = bl->indices.size(); @@ -262,11 +263,12 @@ ExprDoc BufferCall(const ExprDoc& prefix, const ffi::Map& return prefix->Call(args, kwargs_keys, kwargs_values); } -ExprDoc BufferDecl(const tirx::Buffer& buffer, const ffi::String& method, +ExprDoc BufferDecl(const tirx::BufferVar& buffer, const ffi::String& method, const ffi::Array& args, const AccessPath& p, const Frame& frame, - const IRDocsifier& d, BufferVarDefinition var_definitions) { + const IRDocsifier& d, BufferVarDefinition var_definitions, + ffi::Optional data) { auto prefix = TIR(d, method); - auto attrs = BufferAttrs(buffer, p, frame, d, var_definitions); + auto attrs = BufferAttrs(buffer, p, frame, d, var_definitions, data); if (method == "alloc_buffer") { if (buffer.IsScalar()) { // The buffer can be allocated by the alloc_scalar function @@ -303,15 +305,17 @@ ExprDoc BufferDecl(const tirx::Buffer& buffer, const ffi::String& method, auto dtype = d->AsDoc(buffer->dtype, p->Attr("dtype")); auto scope = d->AsDoc(buffer.scope(), p->Attr("scope")); auto elem_offset = d->AsDoc(buffer->elem_offset, p->Attr("elem_offset")); - auto data = d->AsDoc(buffer->data, p->Attr("data")); attrs = ffi::Map( - {{"dtype", dtype}, {"scope", scope}, {"elem_offset", elem_offset}, {"data", data}}); + {{"dtype", dtype}, {"scope", scope}, {"elem_offset", elem_offset}}); + if (data.has_value()) { + attrs.Set("data", d->AsDoc(data.value(), p->Attr("data"))); + } } } return BufferCall(prefix, attrs, args); } -ExprDoc BufferAttn(const tirx::Buffer& buffer, const AccessPath& p, const Frame& frame, +ExprDoc BufferAttn(const tirx::BufferVar& buffer, const AccessPath& p, const Frame& frame, const IRDocsifier& d) { ffi::Map attrs = BufferAttrs(buffer, p, frame, d, BufferVarDefinition::DataPointer); @@ -431,27 +435,6 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) return buffer[BufferIndices(load->indices, p->Attr("indices"), d)]; }); -TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) // - .set_dispatch("", [](tirx::Buffer buffer, AccessPath p, IRDocsifier d) -> Doc { - if (!d->IsVarDefined(buffer)) { - if (ffi::Optional opt_f = FindLowestVarDef(buffer, d)) { - ExprDoc lhs = DefineBuffer(buffer, opt_f.value(), d); - ExprDoc rhs = BufferDecl(buffer, "Buffer", {}, p, opt_f.value(), d, - BufferVarDefinition::DataPointer); - opt_f.value()->stmts.push_back(AssignDoc(lhs, rhs, std::nullopt)); - } - } - if (ffi::Optional doc = d->GetVarDoc(buffer)) { - // special case for scalar buffer - if (buffer.IsScalar()) { - return doc.value()->Attr("buffer"); - } - return doc.value(); - } - TVM_FFI_THROW(IndexError) << "Buffer is not defined in the environment: " << buffer; - TVM_FFI_UNREACHABLE(); - }); - TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) .set_dispatch("", [](tirx::Axis axis, AccessPath p, IRDocsifier d) -> Doc { return LiteralDoc::Str(axis->name, p->Attr("name")); @@ -599,7 +582,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) TVM_SCRIPT_REPR(tirx::BufferRegionNode, ReprPrintTIR); TVM_SCRIPT_REPR(tirx::BufferLoadNode, ReprPrintTIR); TVM_SCRIPT_REPR(tirx::BufferStoreNode, ReprPrintTIR); -TVM_SCRIPT_REPR(tirx::BufferNode, ReprPrintTIR); +TVM_SCRIPT_REPR(tirx::BufferTypeNode, ReprPrintTIR); TVM_SCRIPT_REPR(tirx::IterNode, ReprPrintTIR); TVM_SCRIPT_REPR(tirx::TileLayoutNode, ReprPrintTIR); TVM_SCRIPT_REPR(tirx::ComposeLayoutNode, ReprPrintTIR); diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc index 9beb59799e21..23c2ce4f4dd5 100644 --- a/src/tirx/script/printer/expr.cc +++ b/src/tirx/script/printer/expr.cc @@ -87,6 +87,25 @@ Doc PrintVar(const tirx::Var& var, const AccessPath& var_p, const IRDocsifier& d TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) // .set_dispatch("", [](tirx::Var var, AccessPath p, IRDocsifier d) -> Doc { + if (var->ty.as()) { + tirx::BufferVar buffer(var); + if (!d->IsVarDefined(buffer)) { + if (ffi::Optional opt_f = FindLowestVarDef(buffer, d)) { + ExprDoc lhs = DefineBuffer(buffer, opt_f.value(), d); + ExprDoc rhs = BufferDecl(buffer, "Buffer", {}, p, opt_f.value(), d, + BufferVarDefinition::DataPointer); + opt_f.value()->stmts.push_back(AssignDoc(lhs, rhs, std::nullopt)); + } + } + if (ffi::Optional doc = d->GetVarDoc(buffer)) { + // special case for scalar buffer + if (buffer.IsScalar()) { + return doc.value()->Attr("buffer"); + } + return doc.value(); + } + TVM_FFI_THROW(IndexError) << "BufferVar is not defined in the environment: " << buffer; + } if (var->ty.as() || var->ty.as()) { return PrintVar(var, p, d); } @@ -266,6 +285,11 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) }); Doc PrintTIRCall(Call call, AccessPath call_p, IRDocsifier d) { + if (call->op.same_as(tirx::builtin::buffer_data())) { + TVM_FFI_ICHECK_EQ(call->args.size(), 1); + return d->AsDoc(call->args[0], call_p->Attr("args")->ArrayItem(0)) + ->Attr("data"); + } ffi::Optional call_prim_type = call->ty.as(); auto get_call_type_doc = [&](AccessPath type_p) -> ExprDoc { if (call_prim_type.has_value()) { diff --git a/src/tirx/script/printer/function.cc b/src/tirx/script/printer/function.cc index b3a553c17905..a1eddde6f2d7 100644 --- a/src/tirx/script/printer/function.cc +++ b/src/tirx/script/printer/function.cc @@ -25,7 +25,7 @@ namespace tvm { namespace script { namespace printer { -bool IsSimpleBuffer(const tirx::Buffer& buf, bool s_tir) { +bool IsSimpleBuffer(const tirx::BufferVar& buf, bool s_tir) { if (!buf->strides.empty()) { return false; } @@ -73,7 +73,7 @@ int CountVarOccurrence(const tirx::PrimFunc& f, const tirx::Var& v) { } for (const auto& pair : f->buffer_map) { counter.VisitVar(pair.first); - counter.VisitBuffer(pair.second.get()); + counter.VisitBuffer(pair.second); } return counter.count; } @@ -84,29 +84,29 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) (*f)->AddDispatchToken(d, "tirx"); IdDoc func_name = IdDoc(FindFunctionName(d, func).value_or("main")); d->SetCommonPrefix(func, [](const ffi::ObjectRef& obj) { - return obj->IsInstance() || obj->IsInstance(); + return obj->IsInstance() || obj->IsInstance(); }); int n_args = func->params.size(); std::unordered_map buffer_data_counter; for (const auto& pair : func->buffer_map) { - const tirx::VarNode* data_var = pair.second->data.get(); - if (!buffer_data_counter.count(data_var)) { - buffer_data_counter.insert({data_var, 0}); + const tirx::VarNode* buffer_var = pair.second.get(); + if (!buffer_data_counter.count(buffer_var)) { + buffer_data_counter.insert({buffer_var, 0}); } - ++buffer_data_counter.at(data_var); + ++buffer_data_counter.at(buffer_var); } // Step 1. Handle `func->params` ffi::Array args; args.reserve(n_args); - std::unordered_set buffer_inlined; + std::unordered_set buffer_inlined; for (int i = 0; i < n_args; ++i) { tirx::Var var = func->params[i]; AccessPath var_p = p->Attr("params")->ArrayItem(i); if (d->cfg->syntax_sugar && CountVarOccurrence(func, var) == 2 && func->buffer_map.count(var)) { - tirx::Buffer buffer = func->buffer_map[var]; + tirx::BufferVar buffer = func->buffer_map[var]; bool s_tir = func->attrs->dict.count(tvm::attr::kSTir); - if (IsSimpleBuffer(buffer, s_tir) && buffer_data_counter.at(buffer->data.get()) == 1) { + if (IsSimpleBuffer(buffer, s_tir) && buffer_data_counter.at(buffer.get()) == 1) { AccessPath buffer_p = p->Attr("buffer_map")->MapItem(var); IdDoc lhs = DefineBuffer(buffer, *f, d); ExprDoc annotation = BufferAttn(buffer, buffer_p, *f, d); @@ -151,7 +151,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) for (int i = 0; i < n_args; ++i) { tirx::Var param = func->params[i]; if (func->buffer_map.count(param)) { - tirx::Buffer buffer = func->buffer_map[param]; + tirx::BufferVar buffer = func->buffer_map[param]; if (buffer_inlined.count(buffer.get())) { continue; } @@ -190,7 +190,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) (*f)->stmts.push_back(CommentDoc("with T.sblock(\"root\"):")); // Handle root block `alloc_buffer` for (int i = 0, n = root_block->alloc_buffers.size(); i < n; ++i) { - tirx::Buffer buffer = root_block->alloc_buffers[i]; + tirx::BufferVar buffer = root_block->alloc_buffers[i]; AccessPath buffer_p = root_block_p->Attr("alloc_buffers")->ArrayItem(i); IdDoc lhs = DefineBuffer(buffer, *f, d); ExprDoc rhs = BufferDecl(buffer, "sblock_alloc_buffer", {}, buffer_p, *f, d, diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc index 8601ecb62a5b..661685a3b281 100644 --- a/src/tirx/script/printer/stmt.cc +++ b/src/tirx/script/printer/stmt.cc @@ -256,18 +256,8 @@ namespace { * \param d The IRDocsifier. * \return A list of candidate parent buffers. */ -std::vector FindParentBuffers(const tirx::Buffer& child, const IRDocsifier& d) { - std::vector results; - for (const auto& [obj, info] : d->obj2info) { - if (const auto* buf = obj.as()) { - tirx::Buffer parent = ffi::GetRef(buf); - if (parent.same_as(child)) continue; - if (parent->data.same_as(child->data)) { - results.push_back(parent); - } - } - } - return results; +std::vector FindParentBuffers(const tirx::BufferVar& child, const IRDocsifier& d) { + return {}; } /*! @@ -284,9 +274,9 @@ bool IsDefaultLayout(const ffi::Optional& layout, const ffi::Array * * Returns std::nullopt if no sugar pattern matches. */ -ffi::Optional TryDeclBufferSugarWithParent(const tirx::Buffer& child, const AccessPath& p, +ffi::Optional TryDeclBufferSugarWithParent(const tirx::BufferVar& child, const AccessPath& p, const IRDocsifier& d, - const tirx::Buffer& parent) { + const tirx::BufferVar& parent) { ffi::Optional parent_doc = d->GetVarDoc(parent); if (!parent_doc.has_value()) return std::nullopt; ExprDoc pdoc = parent_doc.value(); @@ -615,7 +605,7 @@ ffi::Optional TryDeclBufferSugarWithParent(const tirx::Buffer& child, c // --- (f) View(*shape, layout=L): different shape/layout, same dtype and elem_offset --- if (same_elem_offset && same_dtype && !same_shape) { - // Buffer.view(...) copies the parent's strides onto the child (see + // BufferVar.view(...) copies the parent's strides onto the child (see // python/tvm/tirx/buffer.py:view). If parent has strides but child // doesn't (or vice versa), the sugar can't faithfully round-trip // through view — fall back to T.decl_buffer where strides is an @@ -659,7 +649,7 @@ ffi::Optional TryDeclBufferSugarWithParent(const tirx::Buffer& child, c /*! * \brief Try to produce a DeclBuffer sugar expression, trying all parent buffer candidates. */ -ffi::Optional TryDeclBufferSugar(const tirx::Buffer& child, const AccessPath& p, +ffi::Optional TryDeclBufferSugar(const tirx::BufferVar& child, const AccessPath& p, const IRDocsifier& d) { auto parents = FindParentBuffers(child, d); for (const auto& parent : parents) { @@ -676,18 +666,12 @@ Doc DeclBufferDoc(tirx::DeclBuffer stmt, AccessPath p, IRDocsifier d, if (d->cfg->syntax_sugar) { if (auto sugar = TryDeclBufferSugar(stmt->buffer, p, d)) { ExprDoc lhs = DefineBuffer(stmt->buffer, d->frames.back(), d); - // Define data pointer inline if needed - if (!d->IsVarDefined(stmt->buffer->data)) { - tirx::Buffer buf = stmt->buffer; - d->Define(stmt->buffer->data, d->frames.back(), [d, buf, p]() { - return d->AsDoc(buf, p->Attr("buffer"))->Attr("data"); - }); - } return AssignDoc(lhs, sugar.value(), std::nullopt); } } - ExprDoc rhs = BufferDecl(stmt->buffer, "decl_buffer", {}, p->Attr("buffer"), d->frames.back(), d, - var_definitions); + ExprDoc rhs = + BufferDecl(stmt->buffer, "decl_buffer", {}, p->Attr("buffer"), d->frames.back(), d, + var_definitions, stmt->data); ExprDoc lhs = DefineBuffer(stmt->buffer, d->frames.back(), d); return AssignDoc(lhs, rhs, std::nullopt); } @@ -703,11 +687,6 @@ namespace { Doc AllocBufferDoc(tirx::AllocBuffer stmt, AccessPath p, IRDocsifier d) { if (d->cfg->syntax_sugar && stmt->buffer.IsScalar(true)) { ExprDoc lhs = DefineBuffer(stmt->buffer, d->frames.back(), d); - if (!d->IsVarDefined(stmt->buffer->data)) { - tirx::Buffer buf = stmt->buffer; - d->Define(stmt->buffer->data, d->frames.back(), - [d, buf, p]() { return d->AsDoc(buf, p->Attr("buffer"))->Attr("data"); }); - } ExprDoc type_ann = TIR(d, DType2Str(stmt->buffer->dtype->dtype)); return AssignDoc(lhs, std::nullopt, type_ann); } diff --git a/src/tirx/script/printer/utils.h b/src/tirx/script/printer/utils.h index b6e07ffd67c1..f1f607ea8dff 100644 --- a/src/tirx/script/printer/utils.h +++ b/src/tirx/script/printer/utils.h @@ -102,8 +102,8 @@ inline ExprDoc DefineVar(const tirx::Var& var, const Frame& frame, const IRDocsi * \param d The IRDocsifier * \return The IdDoc corresponding to the buffer */ -inline IdDoc DefineBuffer(const tirx::Buffer& buffer, const Frame& frame, const IRDocsifier& d) { - return d->Define(buffer, frame, buffer->name.empty() ? "buffer" : buffer->name); +inline IdDoc DefineBuffer(const tirx::BufferVar& buffer, const Frame& frame, const IRDocsifier& d) { + return d->Define(buffer, frame, buffer.name().empty() ? "buffer" : buffer.name()); } /*! @@ -116,7 +116,7 @@ inline IdDoc DefineBuffer(const tirx::Buffer& buffer, const Frame& frame, const inline void AsDocBody(const tirx::Stmt& stmt, AccessPath p, TIRFrameNode* f, const IRDocsifier& d) { if (const auto* seq_stmt = stmt.as()) { ffi::Array body = seq_stmt->seq; - auto value_refs_buffer = [](const PrimExpr& value, const tirx::Buffer& buffer) { + auto value_refs_buffer = [](const PrimExpr& value, const tirx::BufferVar& buffer) { bool found = false; tirx::PostOrderVisit(value, [&](const ffi::ObjectRef& node) { if (const auto* load = node.as()) { @@ -259,7 +259,7 @@ inline ffi::Optional FindLowestVarDef(const ffi::ObjectRef& var, const IR inline std::string ReprPrintTIR(const ffi::ObjectRef& obj, const PrinterConfig& cfg) { IRDocsifier d(cfg); d->SetCommonPrefix(obj, [](const ffi::ObjectRef& obj) { - return obj->IsInstance() || obj->IsInstance(); + return obj->IsInstance() || obj->IsInstance(); }); With f(d, ffi::ObjectRef{nullptr}); (*f)->AddDispatchToken(d, "tirx"); @@ -303,9 +303,10 @@ enum class BufferVarDefinition { * the buffer. * \return The ExprDoc corresponding to the buffer declaration */ -ExprDoc BufferDecl(const tirx::Buffer& buffer, const ffi::String& method, +ExprDoc BufferDecl(const tirx::BufferVar& buffer, const ffi::String& method, const ffi::Array& args, const AccessPath& p, const Frame& frame, - const IRDocsifier& d, BufferVarDefinition var_definitions); + const IRDocsifier& d, BufferVarDefinition var_definitions, + ffi::Optional data = std::nullopt); /*! * \brief Declare and define a buffer as annotation @@ -315,7 +316,7 @@ ExprDoc BufferDecl(const tirx::Buffer& buffer, const ffi::String& method, * \param d The IRDocsifier * \return The ExprDoc corresponding to the buffer declaration */ -ExprDoc BufferAttn(const tirx::Buffer& buffer, const AccessPath& p, const Frame& frame, +ExprDoc BufferAttn(const tirx::BufferVar& buffer, const AccessPath& p, const Frame& frame, const IRDocsifier& d); /*! @@ -345,27 +346,27 @@ class OccurrenceCounter : public tirx::StmtExprVisitor { } void VisitStmt_(const tirx::BufferStoreNode* op) final { - VisitBuffer(op->buffer.get()); + VisitBuffer(op->buffer); tirx::StmtExprVisitor::VisitStmt_(op); } void VisitExpr_(const tirx::BufferLoadNode* op) final { - VisitBuffer(op->buffer.get()); + VisitBuffer(op->buffer); tirx::StmtExprVisitor::VisitExpr_(op); } void VisitStmt_(const tirx::AllocBufferNode* op) final { - VisitBuffer(op->buffer.get()); + VisitBuffer(op->buffer); tirx::StmtExprVisitor::VisitStmt_(op); } void VisitStmt_(const tirx::DeclBufferNode* op) final { - VisitBuffer(op->buffer.get()); + VisitBuffer(op->buffer); tirx::StmtExprVisitor::VisitStmt_(op); } - void VisitBuffer(const tirx::BufferNode* buffer) { - VisitExpr(buffer->data); + void VisitBuffer(const tirx::BufferVar& buffer) { + VisitExpr(buffer.var()); for (const PrimExpr& shape_i : buffer->shape) { VisitExpr(shape_i); } diff --git a/src/tirx/transform/flatten_buffer.cc b/src/tirx/transform/flatten_buffer.cc index 5c27c90e6ab8..e525c2e3feff 100644 --- a/src/tirx/transform/flatten_buffer.cc +++ b/src/tirx/transform/flatten_buffer.cc @@ -61,7 +61,7 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { if (pass.buffers_used_.count(old_buf)) { auto new_buf = pass.GetFlattenedBuffer(old_buf); if (!old_buf.same_as(new_buf)) { - body = SeqStmt::Flatten(DeclBuffer(new_buf), std::move(body)); + body = SeqStmt::Flatten(DeclBuffer(new_buf, old_buf.data()), std::move(body)); } } } @@ -88,8 +88,8 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { SBlock block = ffi::GetRef(op); - ffi::Array alloc_buffers = op->alloc_buffers; - alloc_buffers.MutateByApply([this](Buffer buf) { return GetFlattenedBuffer(buf); }); + ffi::Array alloc_buffers = op->alloc_buffers; + alloc_buffers.MutateByApply([this](BufferVar buf) { return GetFlattenedBuffer(buf); }); if (!alloc_buffers.same_as(op->alloc_buffers)) { block.CopyOnWrite()->alloc_buffers = alloc_buffers; } @@ -131,48 +131,50 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { return std::move(node); } - Buffer GetFlattenedBuffer(Buffer buf) { - auto it = buffer_remap_.find(buf); - if (it != buffer_remap_.end()) { - return it->second; + BufferVar GetFlattenedBuffer(BufferVar buf) { + if (auto remapped = buffer_remap_.Get(buf)) { + return remapped.value(); } auto flattened = buf.GetFlattenedBuffer(); - auto writer = flattened.CopyOnWrite(); + ffi::ObjectPtr type = CopyBufferType(flattened); // canonicalize shape for (size_t i = 0; i < flattened->shape.size(); ++i) { - writer->shape.Set(i, analyzer_->canonical_simplify(flattened->shape[i])); + type->shape.Set(i, analyzer_->canonical_simplify(flattened->shape[i])); } - writer->layout = std::nullopt; + type->layout = std::nullopt; + flattened = RebuildBufferVar(flattened, std::move(type)); - buffer_remap_[buf] = flattened; + buffer_remap_.Set(buf, flattened); return flattened; } Stmt VisitStmt_(const BufferStoreNode* op) final { + BufferVar original_buffer = op->buffer; BufferStore store = StmtExprMutator::VisitStmt_(op).as_or_throw(); - store = VisitBufferAccess(store); + store = VisitBufferAccess(store, original_buffer); return store; } Expr VisitExpr_(const BufferLoadNode* op) final { + BufferVar original_buffer = op->buffer; BufferLoad load = StmtExprMutator::VisitExpr_(op).as_or_throw(); - load = VisitBufferAccess(load); + load = VisitBufferAccess(load, original_buffer); return load; } - ffi::Array GetSimplifiedElemOffset(const Buffer& buffer, + ffi::Array GetSimplifiedElemOffset(const BufferVar& buffer, const ffi::Array& indices) { auto flattened_indices = buffer->ElemOffset(indices); return this->IterMapSimplifyWithContext(flattened_indices, false); } template - Node VisitBufferAccess(Node node) { + Node VisitBufferAccess(Node node, const BufferVar& original_buffer) { TVM_FFI_ICHECK(node->buffer.defined()); - buffers_used_.insert(node->buffer); - auto flattened_indices = GetSimplifiedElemOffset(node->buffer, node->indices); - Buffer flattened_buffer = GetFlattenedBuffer(node->buffer); + buffers_used_.insert(original_buffer); + auto flattened_indices = GetSimplifiedElemOffset(original_buffer, node->indices); + BufferVar flattened_buffer = GetFlattenedBuffer(original_buffer); auto writer = node.CopyOnWrite(); writer->buffer = flattened_buffer; @@ -181,8 +183,8 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { } BufferRegion MutateBufferRegion(BufferRegion region) { - Buffer orig_buf = region->buffer; - Buffer flattened_buf = GetFlattenedBuffer(orig_buf); + BufferVar orig_buf = region->buffer; + BufferVar flattened_buf = GetFlattenedBuffer(orig_buf); if (flattened_buf.same_as(orig_buf)) { return region; } @@ -206,15 +208,12 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { return BufferRegion(flattened_buf, flattened_ranges); } - /*! \brief Map of buffers being remapped. */ - std::unordered_map buffer_remap_; - /*! \brief Set of buffers accessed during visitation (used to emit DeclBuffer for param buffers). */ - std::unordered_set buffers_used_; + std::unordered_set buffers_used_; /*! \brief The updated external buffer map. */ - ffi::Map updated_extern_buffer_map_; + ffi::Map updated_extern_buffer_map_; }; PrimFunc FlattenBuffer(PrimFunc f) { return BufferFlattener::Flatten(f); } diff --git a/src/tirx/transform/force_narrow_index_to_i32.cc b/src/tirx/transform/force_narrow_index_to_i32.cc index 8d2a770f132b..5a4eb54d9e17 100644 --- a/src/tirx/transform/force_narrow_index_to_i32.cc +++ b/src/tirx/transform/force_narrow_index_to_i32.cc @@ -65,7 +65,7 @@ class Int32DTypeNarrower : public IndexDataTypeNormalizer { Stmt VisitStmt_(const SBlockNode* block) final { SBlock block_ = IndexDataTypeNormalizer::VisitStmt_(block).as_or_throw(); // Check if the allocated integer buffers have dtype other than int32. - for (const Buffer& buf : block_->alloc_buffers) { + for (const BufferVar& buf : block_->alloc_buffers) { if (buf->dtype.MatchesCode(DLDataTypeCode::kDLInt) && buf->dtype.bits() > 32) { TVM_FFI_THROW(InternalError) << "The buffer " << buf << " allocated in the function has dtype " << buf->dtype diff --git a/src/tirx/transform/inline_private_functions.cc b/src/tirx/transform/inline_private_functions.cc index 46552bb8b1dc..abe5d2850a54 100644 --- a/src/tirx/transform/inline_private_functions.cc +++ b/src/tirx/transform/inline_private_functions.cc @@ -232,7 +232,7 @@ class PrimFuncInliner : StmtExprMutator { << "Inlining of PrimFuncs with buffer arguments is not yet supported, " << "but callee " << gvar << " has non-empty buffer map " << callee->buffer_map; - ffi::Map> param_map; + ffi::Map> param_map; for (size_t i = 0; i < callee->params.size(); i++) { param_map.Set(callee->params[i], args[i]); } diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc index 5ced08b335f8..dfb359c2225d 100644 --- a/src/tirx/transform/ir_utils.cc +++ b/src/tirx/transform/ir_utils.cc @@ -135,15 +135,15 @@ class IRConvertSSA final : public StmtExprMutator { // Update the buffer map, based on the redefined parameters auto buffer_map = [&]() { - ffi::Map buffer_map; + ffi::Map buffer_map; bool made_change = false; for (const auto& [var, buffer] : func->buffer_map) { auto new_var = GetRemappedVar(var); - if (defined_.count(buffer->data.get())) { - Var new_data = MakeNewVar(buffer->data); - PushVarRemap(buffer->data, new_data); + if (defined_.count(buffer.get())) { + Var new_buffer_var = MakeNewVar(buffer.var()); + PushVarRemap(buffer.var(), new_buffer_var); } else { - defined_.insert(buffer->data.get()); + defined_.insert(buffer.get()); } auto new_buf = GetRemappedBuffer(buffer); @@ -202,7 +202,7 @@ class IRConvertSSA final : public StmtExprMutator { // would create a conflicting second remap (into base buffer_remap_) when called // from the default DeclBuffer/AllocBuffer handlers, producing buffers with // undefined SSA-renamed variables. - Buffer VisitBufferDef(const Buffer& buffer, bool alloc_data) override { return buffer; } + BufferVar VisitBufferDef(const BufferVar& buffer, bool alloc_data) override { return buffer; } Expr VisitExpr_(const VarNode* op) final { return GetRemappedVar(ffi::GetRef(op)); } Expr VisitExpr_(const LetNode* op) final { @@ -233,8 +233,15 @@ class IRConvertSSA final : public StmtExprMutator { } Stmt VisitStmt_(const DeclBufferNode* op) final { + Var v = op->buffer.var(); + if (defined_.count(v.get())) { + Var new_var = MakeNewVar(v); + PushVarRemap(v, new_var); + } else { + defined_.insert(v.get()); + } DeclBuffer decl = StmtExprMutator::VisitStmt_(op).as_or_throw(); - Buffer new_buffer = GetRemappedBuffer(decl->buffer); + BufferVar new_buffer = GetRemappedBuffer(decl->buffer); if (!new_buffer.same_as(decl->buffer)) { decl.CopyOnWrite()->buffer = std::move(new_buffer); } @@ -277,7 +284,7 @@ class IRConvertSSA final : public StmtExprMutator { template Node VisitBufferAccess(Node node) { - Buffer new_buf = GetRemappedBuffer(node->buffer); + BufferVar new_buf = GetRemappedBuffer(node->buffer); if (!new_buf.same_as(node->buffer)) { auto writer = node.CopyOnWrite(); writer->buffer = new_buf; @@ -297,11 +304,11 @@ class IRConvertSSA final : public StmtExprMutator { } } - Buffer GetRemappedBuffer(Buffer buf) { + BufferVar GetRemappedBuffer(BufferVar buf) { // Determine the buffer var that should be in the updated buffer, // given the current scope. If no redefines are present, then the // buffer var is unchanged. - Var new_buffer_var = GetRemappedVar(buf->data); + Var new_buffer_var = GetRemappedVar(buf.var()); PrimExpr elem_offset = VisitPrimExpr(buf->elem_offset); auto visit_expr = [this](const PrimExpr& expr) { return VisitPrimExpr(expr); }; ffi::Array shape = buf->shape.Map(visit_expr); @@ -334,7 +341,7 @@ class IRConvertSSA final : public StmtExprMutator { } // If no mapping is required, return the original buffer. - if (new_buffer_var.same_as(buf->data) && elem_offset.same_as(buf->elem_offset) && + if (new_buffer_var.same_as(buf.var()) && elem_offset.same_as(buf->elem_offset) && shape.same_as(buf->shape) && strides.same_as(buf->strides) && !layout_changed) { return buf; } @@ -342,25 +349,51 @@ class IRConvertSSA final : public StmtExprMutator { // If the current scope already has a mapping of this buffer, use // the mapped buffer. auto key = buf.get(); - std::vector& buffers = buf_remap_[key]; - if (buffers.size() && buffers.back()->data.same_as(new_buffer_var)) { + std::vector& buffers = buf_remap_[key]; + if (buffers.size() && buffers.back().same_as(new_buffer_var)) { return buffers.back(); } + // When only the buffer's identity changed, the remapped Var already has + // the desired BufferType. Reuse that exact Var so the definition and all + // subsequent uses remain in SSA. + if (const auto* type = new_buffer_var->ty.as()) { + BufferVar candidate(new_buffer_var); + if (shape.same_as(type->shape) && strides.same_as(type->strides) && + elem_offset.same_as(type->elem_offset) && !layout_changed) { + buffers.push_back(candidate); + return candidate; + } + } + // Otherwise, make and return a new buffer object that uses the // new buffer, pushing it onto the scoped stack of existing // buffers. This will be popped when the new_buffer_var // redefinition is popped. - Buffer new_buf = buf; - { - auto write_ptr = new_buf.CopyOnWrite(); - write_ptr->data = new_buffer_var; - write_ptr->shape = shape; - write_ptr->strides = strides; - write_ptr->elem_offset = elem_offset; - if (layout_changed) { - write_ptr->layout = std::move(new_layout); - } + auto type = CopyBufferType(buf); + type->shape = shape; + type->strides = strides; + type->elem_offset = elem_offset; + if (layout_changed) { + type->layout = std::move(new_layout); + } + BufferVar new_buf = + RebuildBufferVar(buf, std::move(type), new_buffer_var->name); + + // A BufferVar's metadata lives in its Var type. If rewriting the + // metadata required a fresh Var, make it the active remap as well. This + // keeps BufferLoad/BufferStore and ordinary Var uses (such as + // buffer_data) on the same identity. + auto it = var_remap_.find(buf.get()); + if (it != var_remap_.end() && it->second.size() && + it->second.back().same_as(new_buffer_var)) { + it->second.back() = new_buf.var(); + } else if (auto function_it = function_scope_var_remap_.find(buf.get()); + function_it != function_scope_var_remap_.end() && + function_it->second.same_as(new_buffer_var)) { + function_it->second = new_buf.var(); + } else { + PushVarRemap(buf.var(), new_buf.var()); } buffers.push_back(new_buf); return new_buf; @@ -418,25 +451,24 @@ class IRConvertSSA final : public StmtExprMutator { return scope_.WithNewScope([&]() -> Stmt { return StmtExprMutator::VisitStmt_(op); }); } Stmt VisitStmt_(const AllocBufferNode* op) final { - const Var& v = op->buffer->data; + Var v = op->buffer.var(); if (defined_.count(v.get())) { Var new_var = MakeNewVar(v); PushVarRemap(v, new_var); - Stmt stmt = StmtExprMutator::VisitStmt_(op); - op = stmt.as(); - // Use GetRemappedBuffer so that the AllocBuffer's buffer is the same - // object as the one used by BufferStore/BufferLoad in subsequent siblings. - Buffer new_buf = GetRemappedBuffer(op->buffer); - if (!new_buf.same_as(op->buffer)) { - auto node = stmt.as_or_throw(); - node.CopyOnWrite()->buffer = std::move(new_buf); - return node; - } - return stmt; } else { defined_.insert(v.get()); - return StmtExprMutator::VisitStmt_(op); } + Stmt stmt = StmtExprMutator::VisitStmt_(op); + op = stmt.as(); + // Use GetRemappedBuffer so that the AllocBuffer's buffer is the same + // object as the one used by BufferStore/BufferLoad in subsequent siblings. + BufferVar new_buf = GetRemappedBuffer(op->buffer); + if (!new_buf.same_as(op->buffer)) { + auto node = stmt.as_or_throw(); + node.CopyOnWrite()->buffer = std::move(new_buf); + return node; + } + return stmt; } Stmt VisitStmt_(const AttrStmtNode* op) final { if (const IterVarNode* iter_var = op->node.as()) { @@ -540,11 +572,9 @@ class IRConvertSSA final : public StmtExprMutator { /*! \brief Pop a single variable remap (used for expression-level Let scoping). */ void PopVarRemap(const Var& old_var, const Var& new_var) { var_remap_[old_var.get()].pop_back(); - for (auto& kv : buf_remap_) { - std::vector& buffers = kv.second; - if (buffers.size() && (buffers.back()->data.get() == new_var.get())) { - buffers.pop_back(); - } + if (auto it = buf_remap_.find(old_var.get()); + it != buf_remap_.end() && it->second.size()) { + it->second.pop_back(); } // Also remove from the current scope's tracking vector auto& current = scope_.Current(); @@ -559,11 +589,9 @@ class IRConvertSSA final : public StmtExprMutator { while (current.size()) { auto& remap = current.back(); var_remap_[remap.old_var.get()].pop_back(); - for (auto& kv : buf_remap_) { - std::vector& buffers = kv.second; - if (buffers.size() && (buffers.back()->data.get() == remap.new_var.get())) { - buffers.pop_back(); - } + if (auto it = buf_remap_.find(remap.old_var.get()); + it != buf_remap_.end() && it->second.size()) { + it->second.pop_back(); } current.pop_back(); } @@ -596,11 +624,9 @@ class IRConvertSSA final : public StmtExprMutator { while (remaps.size()) { auto& remap = remaps.back(); parent->var_remap_[remap.old_var.get()].pop_back(); - for (auto& kv : parent->buf_remap_) { - std::vector& buffers = kv.second; - if (buffers.size() && (buffers.back()->data.get() == remap.new_var.get())) { - buffers.pop_back(); - } + if (auto it = parent->buf_remap_.find(remap.old_var.get()); + it != parent->buf_remap_.end() && it->second.size()) { + it->second.pop_back(); } remaps.pop_back(); } @@ -620,11 +646,9 @@ class IRConvertSSA final : public StmtExprMutator { while (remaps.size()) { auto& remap = remaps.back(); parent->var_remap_[remap.old_var.get()].pop_back(); - for (auto& kv : parent->buf_remap_) { - std::vector& buffers = kv.second; - if (buffers.size() && (buffers.back()->data.get() == remap.new_var.get())) { - buffers.pop_back(); - } + if (auto it = parent->buf_remap_.find(remap.old_var.get()); + it != parent->buf_remap_.end() && it->second.size()) { + it->second.pop_back(); } remaps.pop_back(); } @@ -639,7 +663,7 @@ class IRConvertSSA final : public StmtExprMutator { std::unordered_map> var_remap_; std::unordered_set defined_; - std::unordered_map> buf_remap_; + std::unordered_map> buf_remap_; std::unordered_map function_scope_var_remap_; ScopeStack scope_; }; @@ -647,12 +671,16 @@ class IRConvertSSA final : public StmtExprMutator { Stmt ConvertSSA(Stmt stmt) { return IRConvertSSA()(std::move(stmt)); } ffi::String GetPtrStorageScope(Var buffer_var) { + if (const auto* buffer_type = buffer_var->ty.as()) { + return buffer_type->data_pointer_type->storage_scope; + } const auto* ptr_type = buffer_var->ty.as(); - TVM_FFI_ICHECK(ptr_type) << "The provided variable is not of pointer type"; + TVM_FFI_ICHECK(ptr_type) + << "The provided variable is neither a pointer nor a buffer-typed variable"; return ptr_type->storage_scope; } -ffi::Array GetBufferAllocationShape(const Buffer& buffer) { +ffi::Array GetBufferAllocationShape(const BufferVar& buffer) { ffi::Array alloc_shape = buffer->shape; if (buffer->strides.size()) { TVM_FFI_ICHECK_EQ(buffer->shape.size(), buffer->strides.size()); @@ -667,7 +695,7 @@ ffi::Array GetBufferAllocationShape(const Buffer& buffer) { ffi::Array ConvertIndices(const MatchBufferRegion& match_buffer, const ffi::Array& indices) { - const Buffer& target = match_buffer->buffer; + const BufferVar& target = match_buffer->buffer; const BufferRegion& source = match_buffer->source; TVM_FFI_ICHECK_EQ(indices.size(), target->shape.size()); @@ -689,7 +717,7 @@ ffi::Array ConvertIndices(const MatchBufferRegion& match_buffer, } Region ConvertRegion(const MatchBufferRegion& match_buffer, const Region& region) { - const Buffer& target = match_buffer->buffer; + const BufferVar& target = match_buffer->buffer; const BufferRegion& source = match_buffer->source; TVM_FFI_ICHECK_EQ(region.size(), target->shape.size()); @@ -871,8 +899,8 @@ class StorageAlignCollector : public StmtVisitor { auto storage_align_annotation = (*it).second.as_or_throw(); for (const auto& storage_align_tuple : storage_align_annotation) { int buffer_index = storage_align_tuple.get<0>(); - const Buffer& buffer = op->writes[buffer_index]->buffer; - storage_align_[buffer->data].push_back(storage_align_tuple); + const BufferVar& buffer = op->writes[buffer_index]->buffer; + storage_align_[buffer.var()].push_back(storage_align_tuple); } } StmtVisitor::VisitStmt_(op); @@ -888,7 +916,7 @@ class StorageAlignCollector : public StmtVisitor { // the first buffer idx info is meaningless for alloc // stmt and should set as negative intentionally. TVM_FFI_ICHECK_EQ(buffer_index, -1); - storage_align_[op->buffer->data].push_back(storage_align_tuple); + storage_align_[op->buffer.var()].push_back(storage_align_tuple); } } StmtVisitor::VisitStmt_(op); diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h index 716f26a5f1e5..df6d3368a652 100644 --- a/src/tirx/transform/ir_utils.h +++ b/src/tirx/transform/ir_utils.h @@ -114,7 +114,7 @@ inline PrimExpr TVMStructGet(PrimType type, Var handle, int index, inline Call AddressOffset(Var handle, PrimType dtype, int offset) { PrimExpr offset_expr = IntImm::Int32(offset * dtype.lanes()); ffi::Array shape = {offset_expr + 1}; - Buffer dummy_buf(handle, dtype, shape, {}, 0, handle->name, 0, 0, Span(), std::nullopt); + BufferVar dummy_buf(handle, dtype, shape, {}, 0, handle->name, 0, 0, Span(), std::nullopt); BufferLoad buf_load(dummy_buf, {offset_expr}); return Call(handle->ty, builtin::address_of(), {buf_load}); @@ -134,7 +134,7 @@ inline Call AddressOffset(Var handle, PrimType dtype, PrimExpr offset) { } ffi::Array shape = {offset + 1}; - Buffer dummy_buf(handle, dtype.WithLanes(1), shape, {}, 0, handle->name, 0, 0, Span(), + BufferVar dummy_buf(handle, dtype.WithLanes(1), shape, {}, 0, handle->name, 0, 0, Span(), std::nullopt); BufferLoad buf_load(dummy_buf, {offset}); @@ -243,7 +243,7 @@ Region ConvertRegion(const MatchBufferRegion& match_buffer, const Region& region * \param buffer The buffer object. * \return shape The shape considering buffer strides. */ -ffi::Array GetBufferAllocationShape(const Buffer& buffer); +ffi::Array GetBufferAllocationShape(const BufferVar& buffer); /*! * \brief Context helper to update domain map within conditional scope. diff --git a/src/tirx/transform/lower_intrin.cc b/src/tirx/transform/lower_intrin.cc index 4c93b16440e8..1af07c8b4a59 100644 --- a/src/tirx/transform/lower_intrin.cc +++ b/src/tirx/transform/lower_intrin.cc @@ -76,7 +76,7 @@ static Expr LowerAccessPtr(const CallNode* call) { offset = offset * IntImm(offset_ty, dtype.lanes()); offset = Ramp(offset, IntImm(offset_ty, 1), dtype.lanes()); } - Buffer dummy_buf(buffer_var, dtype.WithLanes(1), {offset + 1}, {}, 0, buffer_var->name, 0, 0); + BufferVar dummy_buf(buffer_var, dtype.WithLanes(1), {offset + 1}, {}, 0, buffer_var->name, 0, 0); BufferLoad buf_load(dummy_buf, {offset}); return Call(call->ty, builtin::address_of(), {buf_load}); } diff --git a/src/tirx/transform/lower_tirx_cleanup.cc b/src/tirx/transform/lower_tirx_cleanup.cc index 8a977a729d72..f3a5cc3cdfe6 100644 --- a/src/tirx/transform/lower_tirx_cleanup.cc +++ b/src/tirx/transform/lower_tirx_cleanup.cc @@ -44,28 +44,29 @@ namespace tirx { class LayoutApplier : public arith::IRMutatorWithAnalyzer { public: - static std::pair> Flatten( - const Stmt& stmt, const ffi::Map buffer_map, const Target& target) { + static std::pair> Flatten( + const Stmt& stmt, const ffi::Map buffer_map, const Target& target) { arith::Analyzer ana; LayoutApplier storage_lower(ana, target); - std::unordered_map new_buffer_map; - std::vector param_flattened_buffers; + std::unordered_map new_buffer_map; + std::vector> param_flattened_buffers; for (const auto& kv : buffer_map) { if (kv.second->layout.has_value()) { - param_flattened_buffers.push_back(storage_lower.GetFlattenedBuffer(kv.second)); - Buffer buffer = kv.second; - auto* writer = buffer.CopyOnWrite(); - writer->layout = std::nullopt; + BufferVar flattened = storage_lower.GetFlattenedBuffer(kv.second); + auto type = CopyBufferType(kv.second); + type->layout = std::nullopt; + BufferVar buffer = RebuildBufferVar(kv.second, std::move(type)); + param_flattened_buffers.emplace_back(flattened, buffer); new_buffer_map[kv.first] = buffer; } else { new_buffer_map[kv.first] = kv.second; } } auto new_stmt = storage_lower(stmt); - for (const auto& buf : param_flattened_buffers) { - new_stmt = SeqStmt::Flatten(DeclBuffer(buf), std::move(new_stmt)); + for (const auto& [buf, source] : param_flattened_buffers) { + new_stmt = SeqStmt::Flatten(DeclBuffer(buf, source.data()), std::move(new_stmt)); } - return std::make_pair(new_stmt, ffi::Map(new_buffer_map)); + return std::make_pair(new_stmt, ffi::Map(new_buffer_map)); } protected: @@ -79,7 +80,7 @@ class LayoutApplier : public arith::IRMutatorWithAnalyzer { if (any == nullptr) { return any; } - if (auto buffer = any.as()) { + if (auto buffer = any.as()) { return GetFlattenedBuffer(buffer.value()); } else if (auto prim_expr = any.as()) { return VisitPrimExpr(prim_expr.value()); @@ -90,7 +91,7 @@ class LayoutApplier : public arith::IRMutatorWithAnalyzer { } Stmt VisitStmt_(const AllocBufferNode* op) final { - auto mutate = [this](Buffer buf) { + auto mutate = [this](BufferVar buf) { if (target_->kind->name == "trn" && !buf->layout.has_value()) { return buf; } @@ -115,14 +116,14 @@ class LayoutApplier : public arith::IRMutatorWithAnalyzer { return Stmt(n); } - Buffer GetFlattenedBuffer(Buffer buf, bool is_alloc = false) { + BufferVar GetFlattenedBuffer(BufferVar buf, bool is_alloc = false) { auto it = buffer_remap_.find(buf); if (it != buffer_remap_.end()) { return it->second; } auto trn_layout = buf->layout.as(); - Buffer flattened; - tirx::BufferNode* writer; + BufferVar flattened; + ffi::ObjectPtr type; if (trn_layout && trn_layout->IsTrainium()) { ffi::Array new_shape = buf.scope() == "trn.psum" ? ffi::Array{trn_layout->GetSpan(ffi::String("Bank")), @@ -131,9 +132,9 @@ class LayoutApplier : public arith::IRMutatorWithAnalyzer { : ffi::Array{trn_layout->GetSize(ffi::String("P")), trn_layout->GetSpan(ffi::String("F"))}; flattened = buf; - writer = flattened.CopyOnWrite(); - writer->shape = new_shape; - writer->strides = {}; + type = CopyBufferType(flattened); + type->shape = new_shape; + type->strides = {}; } else if (is_alloc) { if (auto tile_layout = buf->layout.as(); tile_layout && tile_layout->HasThreadAxis()) { @@ -156,23 +157,24 @@ class LayoutApplier : public arith::IRMutatorWithAnalyzer { } } flattened = buf; - writer = flattened.CopyOnWrite(); - writer->shape = {ana->Simplify(mem_span)}; - writer->strides = {}; + type = CopyBufferType(flattened); + type->shape = {ana->Simplify(mem_span)}; + type->strides = {}; } else { flattened = buf.GetFlattenedBuffer(); - writer = flattened.CopyOnWrite(); + type = CopyBufferType(flattened); } } else { flattened = buf.GetFlattenedBuffer(); - writer = flattened.CopyOnWrite(); + type = CopyBufferType(flattened); } // canonicalize shape for (size_t i = 0; i < flattened->shape.size(); ++i) { - writer->shape.Set(i, analyzer_->canonical_simplify(flattened->shape[i])); + type->shape.Set(i, analyzer_->canonical_simplify(flattened->shape[i])); } - writer->layout = std::nullopt; - writer->elem_offset = StmtExprMutator::VisitPrimExpr(buf->elem_offset); + type->layout = std::nullopt; + type->elem_offset = StmtExprMutator::VisitPrimExpr(buf->elem_offset); + flattened = RebuildBufferVar(flattened, std::move(type)); buffer_remap_[buf] = flattened; return flattened; @@ -202,7 +204,7 @@ class LayoutApplier : public arith::IRMutatorWithAnalyzer { } } - ffi::Array GetSimplifiedElemOffset(const Buffer& buffer, + ffi::Array GetSimplifiedElemOffset(const BufferVar& buffer, const ffi::Array& indices) { if (buffer->layout.has_value()) { auto tile_layout = buffer->layout.value().as(); @@ -242,7 +244,7 @@ class LayoutApplier : public arith::IRMutatorWithAnalyzer { return node; } auto flattened_indices = GetSimplifiedElemOffset(node->buffer, node->indices); - Buffer flattened_buffer = GetFlattenedBuffer(node->buffer); + BufferVar flattened_buffer = GetFlattenedBuffer(node->buffer); auto writer = node.CopyOnWrite(); writer->buffer = flattened_buffer; writer->indices = flattened_indices; @@ -250,7 +252,7 @@ class LayoutApplier : public arith::IRMutatorWithAnalyzer { } /*! \brief Map of buffers being remapped. */ - std::unordered_map buffer_remap_; + std::unordered_map buffer_remap_; const Target& target_; }; @@ -274,11 +276,12 @@ class BufferOffsetRemover : public StmtExprMutator { if (elem_offset.same_as(buffer->elem_offset)) { return StmtExprMutator::VisitStmt_(op); } else { - auto n_buffer = buffer.CopyOnWrite(); - n_buffer->elem_offset = std::move(elem_offset); + auto type = CopyBufferType(buffer); + type->elem_offset = std::move(elem_offset); + buffer = RebuildBufferVar(buffer, std::move(type)); buffer_remap_[op->buffer] = buffer; auto n = CopyOnWrite(op); - n->buffer = ffi::GetRef(n_buffer); + n->buffer = buffer; return Stmt(n); } } @@ -310,7 +313,7 @@ class BufferOffsetRemover : public StmtExprMutator { return node; } - std::unordered_map buffer_remap_; + std::unordered_map buffer_remap_; }; namespace { diff --git a/src/tirx/transform/lower_tirx_opaque.cc b/src/tirx/transform/lower_tirx_opaque.cc index d2bfec19d868..12763457f231 100644 --- a/src/tirx/transform/lower_tirx_opaque.cc +++ b/src/tirx/transform/lower_tirx_opaque.cc @@ -83,22 +83,14 @@ class TIRxOpaqueLower : public StmtExprMutator { } Stmt VisitStmt_(const AllocBufferNode* op) final { - Stmt stmt = StmtExprMutator::VisitStmt_(op); - op = stmt.as(); - TVM_FFI_ICHECK(op); - - Buffer alloc_buf = op->buffer; - auto it = pool_sizes_.find(op->buffer->data); + auto it = pool_sizes_.find(op->buffer.var()); if (it != pool_sizes_.end()) { - auto* n = alloc_buf.CopyOnWrite(); - n->shape = {IntImm::Int64(it->second)}; - } - if (alloc_buf.same_as(op->buffer)) { - return stmt; + auto type = CopyBufferType(op->buffer); + type->shape = {IntImm::Int64(it->second)}; + BufferVar alloc_buf = RebuildBufferVar(op->buffer, std::move(type)); + buffer_remap_.Set(op->buffer, alloc_buf); } - auto n = CopyOnWrite(op); - n->buffer = std::move(alloc_buf); - return Stmt(n); + return StmtExprMutator::VisitStmt_(op); } Stmt VisitStmt_(const ForNode* op) final { diff --git a/src/tirx/transform/lower_tvm_builtin.cc b/src/tirx/transform/lower_tvm_builtin.cc index 55790e51a4bf..7998cca2a5de 100644 --- a/src/tirx/transform/lower_tvm_builtin.cc +++ b/src/tirx/transform/lower_tvm_builtin.cc @@ -106,7 +106,7 @@ class BuiltinLower : public StmtExprMutator { // Record stack frame for existing scope. struct AllocaScope { - Buffer stack_shape; + BufferVar stack_shape; Var stack_array = Var("stack_array", PointerType::VoidPointerTy()); Var stack_ffi_any = Var("stack_ffi_any", PointerType::VoidPointerTy()); @@ -161,7 +161,7 @@ class BuiltinLower : public StmtExprMutator { auto& scope = alloca_scope_.back(); // Initial check to identify maximum stack sizes. These are used - // to construct Buffer objects to hold the stack, which are then + // to construct BufferVar objects to hold the stack, which are then // used when mutating. scope.max_sizes = GetMaxStack(stmt); @@ -181,10 +181,11 @@ class BuiltinLower : public StmtExprMutator { if (scope.max_sizes.shape_stack != -1) { scope.stack_shape = decl_buffer({IntImm::Int64(scope.max_sizes.shape_stack)}, PrimType::Int(64), "stack_shape"); - alloca_stmts.push_back( - Bind(scope.stack_shape->data, - StackAlloca(scope.stack_shape->data->ty, "shape", scope.max_sizes.shape_stack))); - stmt = SeqStmt::Flatten(DeclBuffer(scope.stack_shape), stmt); + stmt = SeqStmt::Flatten( + DeclBuffer(scope.stack_shape, + StackAlloca(scope.stack_shape->data_pointer_type, "shape", + scope.max_sizes.shape_stack)), + stmt); } if (!alloca_stmts.empty()) { @@ -261,7 +262,7 @@ class BuiltinLower : public StmtExprMutator { int64_t nbytes = GetVectorBytes(op->buffer->dtype); if (const auto* dev_type = device_type_.as(); dev_type && dev_type->value == kDLCPU) { - auto storage_scope = op->buffer->data->ty.as_or_throw()->storage_scope; + auto storage_scope = op->buffer->data_pointer_type.as_or_throw()->storage_scope; if (storage_scope == "global") { auto constant_size = stmt.as_or_throw().ConstantAllocationSize(); if (constant_size.has_value() && constant_size.value() > 0 && @@ -280,26 +281,27 @@ class BuiltinLower : public StmtExprMutator { Call(PrimType::Int(32), builtin::tvm_throw_last_error(), {}).as_or_throw()); Stmt alloc_nullptr_check = IfThenElse( - Call(PrimType::Bool(), builtin::isnullptr(), {op->buffer->data}).as_or_throw(), + Call(PrimType::Bool(), builtin::isnullptr(), {op->buffer.data()}) + .as_or_throw(), throw_last_error); static const Op& free_workspace_op = Op::Get("tirx.TVMBackendFreeWorkspace"); static const Op& alloc_workspace_op = Op::Get("tirx.TVMBackendAllocWorkspace"); PrimExpr free_op = Call(PrimType::Int(32), free_workspace_op, - {cast(PrimType::Int(32), device_type_.value()), - cast(PrimType::Int(32), device_id_.value()), op->buffer->data}) + {cast(PrimType::Int(32), device_type_.value()), + cast(PrimType::Int(32), device_id_.value()), op->buffer.data()}) .as_or_throw(); Stmt free_stmt = IfThenElse(free_op != IntImm::Int32(0), throw_last_error); // Push free to enclosing scope's pending_frees (LIFO ordering preserved). scope_.Current().pending_frees.push_back(free_stmt); - Stmt alloc_bind = Bind( - op->buffer->data, - Call(op->buffer->data->ty, alloc_workspace_op, - {cast(PrimType::Int(32), device_type_.value()), - cast(PrimType::Int(32), device_id_.value()), total_bytes, - IntImm::Int32(op->buffer->dtype.code()), IntImm::Int32(op->buffer->dtype.bits())})); + Stmt alloc_bind = DeclBuffer( + op->buffer, Call(op->buffer->data_pointer_type, alloc_workspace_op, + {cast(PrimType::Int(32), device_type_.value()), + cast(PrimType::Int(32), device_id_.value()), total_bytes, + IntImm::Int32(op->buffer->dtype.code()), + IntImm::Int32(op->buffer->dtype.bits())})); return SeqStmt({alloc_bind, alloc_nullptr_check}); } @@ -488,7 +490,9 @@ class BuiltinLower : public StmtExprMutator { scope.stack_shape, cast(PrimType::Int(64), op->args[i].as_or_throw()), {ConstInt32(stack_begin + i)})); } - return AddressOffset(scope.stack_shape->data, PrimType::Int(64), stack_begin); + PrimExpr offset = ConstInt32(stack_begin); + BufferLoad load(scope.stack_shape, {offset}); + return Call(scope.stack_shape->data_pointer_type, builtin::address_of(), {load}); } // make array Expr MakeArray(const CallNode* op) { diff --git a/src/tirx/transform/lower_warp_memory.cc b/src/tirx/transform/lower_warp_memory.cc index 4981dee72558..0d88980037fa 100644 --- a/src/tirx/transform/lower_warp_memory.cc +++ b/src/tirx/transform/lower_warp_memory.cc @@ -146,7 +146,7 @@ class WarpStoreCoeffFinder : private StmtExprVisitor { } void VisitStmt_(const BufferStoreNode* op) final { - if (op->buffer->data.get() != buffer_) { + if (op->buffer.get() != buffer_) { StmtVisitor::VisitStmt_(op); return; } @@ -257,7 +257,7 @@ class WarpAccessRewriter : protected StmtExprMutator { // \param op The AllocBuffer node for warp memory. // \param body The remaining statements (siblings) that use this buffer. Stmt Rewrite(const AllocBufferNode* op, Stmt body) { - buffer_ = op->buffer->data.get(); + buffer_ = op->buffer.get(); int64_t alloc_size = 1; for (const auto& dim : op->buffer->shape) { if (const IntImmNode* int_size = dim.as()) { @@ -278,8 +278,13 @@ class WarpAccessRewriter : protected StmtExprMutator { warp_group_ = (alloc_size + (factor - 1)) / factor; alloc_size = warp_group_ * factor; - Buffer new_buf(op->buffer->data, op->buffer->dtype, {IntImm::Int32(alloc_size / width_)}, {}, - PrimExpr(), op->buffer->data->name, 0, 0); + auto type = CopyBufferType(op->buffer); + type->data_pointer_type = PointerType(type->data_pointer_type->element_type, "local"); + type->shape = {IntImm::Int32(alloc_size / width_)}; + type->strides = {}; + type->elem_offset = PrimExpr(); + BufferVar new_buf = RebuildBufferVar(op->buffer, std::move(type)); + new_buffer_ = new_buf; Stmt rewritten_body = this->VisitStmt(body); return SeqStmt::Flatten(AllocBuffer(new_buf, op->annotations), rewritten_body); } @@ -352,7 +357,7 @@ class WarpAccessRewriter : protected StmtExprMutator { Stmt VisitStmt_(const BufferStoreNode* op) override { auto store = StmtExprMutator::VisitStmt_(op).as_or_throw(); - if (store->buffer->data.get() == buffer_) { + if (store->buffer.get() == buffer_) { TVM_FFI_ICHECK_EQ(store->indices.size(), 1) << "Expected flat memory to use as warp memory. " << "Has FlattenBuffer been run?"; @@ -360,6 +365,7 @@ class WarpAccessRewriter : protected StmtExprMutator { (void)group; // https://gcc.gnu.org/bugzilla/show_bug.cgi?id=81767 auto writer = store.CopyOnWrite(); + writer->buffer = new_buffer_; writer->indices = {local_index}; } @@ -369,7 +375,7 @@ class WarpAccessRewriter : protected StmtExprMutator { Expr VisitExpr_(const BufferLoadNode* op) override { auto load = StmtExprMutator::VisitExpr_(op).as_or_throw(); - if (load->buffer->data.get() != buffer_) { + if (load->buffer.get() != buffer_) { return load; } @@ -384,6 +390,7 @@ class WarpAccessRewriter : protected StmtExprMutator { << " local_index=" << local_index; auto writer = load.CopyOnWrite(); + writer->buffer = new_buffer_; writer->indices = {local_index}; if (analyzer_->CanProveEqual(group, warp_index_.as_or_throw())) { @@ -433,6 +440,8 @@ class WarpAccessRewriter : protected StmtExprMutator { int warp_size_{0}; // The buffer variable const VarNode* buffer_; + // The fresh local buffer replacing the warp-scoped definition. + BufferVar new_buffer_; // number of threads involved in one shuffle int width_{0}; // Warp index @@ -500,8 +509,8 @@ class WarpMemoryRewriter : private StmtMutator { bool changed = false; for (size_t i = 0; i < op->seq.size(); ++i) { const auto* alloc = op->seq[i].as(); - if (alloc && GetPtrStorageScope(alloc->buffer->data) == "warp") { - new_storage_scopes_[alloc->buffer->data.get()] = "local"; + if (alloc && alloc->buffer.scope() == "warp") { + new_storage_scopes_[alloc->buffer.get()] = "local"; // Gather remaining siblings as the "body" for rewriting. ffi::Array remaining; for (size_t j = i + 1; j < op->seq.size(); ++j) { diff --git a/src/tirx/transform/make_packed_api.cc b/src/tirx/transform/make_packed_api.cc index 8beb6d918c58..8ad6dbaa03dc 100644 --- a/src/tirx/transform/make_packed_api.cc +++ b/src/tirx/transform/make_packed_api.cc @@ -279,7 +279,7 @@ PrimFunc MakePackedAPI(PrimFunc func) { << "In PrimFunc " << name_hint << " variables " << undefined << " are used, but are not passed in as API arguments"; - func_ptr->buffer_map = ffi::Map(); + func_ptr->buffer_map = ffi::Map(); func_ptr->ret_type = PrimType::Int(32); // return the function. diff --git a/src/tirx/transform/remove_no_op.cc b/src/tirx/transform/remove_no_op.cc index 87ab5d05abd7..9311a4e1ea96 100644 --- a/src/tirx/transform/remove_no_op.cc +++ b/src/tirx/transform/remove_no_op.cc @@ -203,7 +203,7 @@ class NoOpRemover : public arith::IRMutatorWithAnalyzer { // If the stored value is a load from the same location, the // statement is a no-op, regardless of contextual information. if (const BufferLoadNode* load = store->value.as()) { - if (load->buffer->data.same_as(store->buffer->data) && + if (load->buffer.same_as(store->buffer) && analyzer_->CanProveEqual(load->buffer->elem_offset, store->buffer->elem_offset) && ArrayValueEqual(load->buffer->shape, store->buffer->shape) && ArrayValueEqual(load->buffer->strides, store->buffer->strides) && diff --git a/src/tirx/transform/split_host_device.cc b/src/tirx/transform/split_host_device.cc index 5bf8c77584b5..54e3aebc8f84 100644 --- a/src/tirx/transform/split_host_device.cc +++ b/src/tirx/transform/split_host_device.cc @@ -131,7 +131,7 @@ class HostDeviceSplitter : public StmtMutator { private: Stmt SplitDeviceFunc(Stmt body, Target device_target) { - auto [params, buffers_to_declare] = [&]() -> std::tuple, ffi::Array> { + auto [params, buffers_to_declare] = [&]() -> std::tuple, ffi::Array> { VarUseDefAnalyzer use_def(/*defined_vars=*/{}, /*visit_thread_extent=*/true); use_def(body); @@ -151,7 +151,7 @@ class HostDeviceSplitter : public StmtMutator { } else { std::unordered_map param_order; for (size_t i = 0; i < cur_func_->params.size(); ++i) { - param_order[cur_func_->buffer_map[cur_func_->params[i]]->data] = i; + param_order[cur_func_->buffer_map[cur_func_->params[i]].var()] = i; } // sort by original order std::sort(params.begin(), params.end(), @@ -160,6 +160,29 @@ class HostDeviceSplitter : public StmtMutator { return {params, use_def.undefined_buffers_}; }(); + // Buffer Vars are compiler-side values, not ABI values. Thread their + // physical pointer projection through the kernel call and recover the + // typed buffer at the kernel entry with an explicit DeclBuffer source. + ffi::Array kernel_params; + ffi::Array call_args; + ffi::Map buffer_data_params; + ffi::Map kernel_buffer_remap; + for (const Var& param : params) { + if (param->ty.as()) { + BufferVar buffer(param); + BufferVar kernel_buffer(buffer.name(), buffer.type(), buffer.span()); + Var data_param(buffer.name() + "_data", buffer->data_pointer_type); + kernel_params.push_back(data_param); + call_args.push_back(buffer.data()); + buffer_data_params.Set(param, data_param); + kernel_buffer_remap.Set(param, kernel_buffer.var()); + } else { + kernel_params.push_back(param); + call_args.push_back(param); + } + } + body = Substitute(std::move(body), kernel_buffer_remap); + // CodeGenCPU is used for some device-side targets, such as // "ext_dev", and expects to be able to return a int32_t status // code. @@ -177,12 +200,19 @@ class HostDeviceSplitter : public StmtMutator { kernel_ret_type = VoidType(); } - for (Buffer buf : buffers_to_declare) { - body = SeqStmt::Flatten(DeclBuffer(buf), std::move(body)); + for (BufferVar buf : buffers_to_declare) { + auto data_param = buffer_data_params.Get(buf.var()); + auto kernel_buffer = kernel_buffer_remap.Get(buf.var()); + TVM_FFI_ICHECK(data_param.has_value()) + << "Undefined buffer " << buf.name() << " was not captured as a kernel parameter"; + TVM_FFI_ICHECK(kernel_buffer.has_value()); + body = SeqStmt::Flatten( + DeclBuffer(BufferVar(kernel_buffer.value().as_or_throw()), data_param.value()), + std::move(body)); } LaunchBoundsAttrExtractor launch_bounds_attr; body = launch_bounds_attr.Extract(std::move(body)); - PrimFunc device_func(params, body, kernel_ret_type); + PrimFunc device_func(kernel_params, body, kernel_ret_type); device_func = WithAttrs(std::move(device_func), {{tvm::attr::kTarget, device_target}, {tirx::attr::kNoAlias, true}, {tirx::attr::kIsGlobalFunc, true}}); @@ -200,11 +230,9 @@ class HostDeviceSplitter : public StmtMutator { } GlobalVar kernel_symbol_global = var_supply_(); (*device_mod_)->Add(kernel_symbol_global, device_func); - ffi::Array args = params.Map([](const Var& var) -> Expr { return var; }); - if (can_propagate_errors) { Var kernel_error_code("kernel_error_code", success.ty()); - Call kernel_call(success.ty(), kernel_symbol_global, args); + Call kernel_call(success.ty(), kernel_symbol_global, call_args); AssertStmt assert_success(kernel_error_code.as_or_throw() == success, StringImm("RuntimeError"), {StringImm("Error executing compute kernel")}); @@ -212,7 +240,8 @@ class HostDeviceSplitter : public StmtMutator { assert_success}); } else { - return Evaluate(Call(PrimType::Void(), kernel_symbol_global, args).as_or_throw()); + return Evaluate( + Call(PrimType::Void(), kernel_symbol_global, call_args).as_or_throw()); } } @@ -352,7 +381,7 @@ class DeviceInfoCollector : public StmtVisitor { } void VisitStmt_(const AllocBufferNode* op) final { - auto storage_scope = runtime::StorageScope::Create(GetPtrStorageScope(op->buffer->data)); + auto storage_scope = runtime::StorageScope::Create(op->buffer.scope()); if (storage_scope.rank == runtime::StorageRank::kShared && storage_scope.tag == ".dyn") { TVM_FFI_ICHECK(!dyn_shmem_size.has_value()) << "Only one dynamic shared memory allocation is allowed."; diff --git a/src/tirx/transform/stmt_simplify.cc b/src/tirx/transform/stmt_simplify.cc index bc23cae42f65..d37e59576aaf 100644 --- a/src/tirx/transform/stmt_simplify.cc +++ b/src/tirx/transform/stmt_simplify.cc @@ -129,7 +129,7 @@ class StmtSimplifier : public IRMutatorWithAnalyzer { // // Instead, we keep buffer definitions unchanged and rely on used_in_buffer_def_ // to prevent inlining LetStmt vars that appear in buffer definitions. - Buffer VisitBufferDef(const Buffer& buffer, bool alloc_data) override { return buffer; } + BufferVar VisitBufferDef(const BufferVar& buffer, bool alloc_data) override { return buffer; } Expr VisitExpr(const Expr& expr) final { if (auto prim_expr = expr.as()) { @@ -211,7 +211,7 @@ class StmtSimplifier : public IRMutatorWithAnalyzer { Stmt VisitStmt_(const BufferStoreNode* op) override { BufferStore store = Parent::VisitStmt_(op).as_or_throw(); if (const BufferLoadNode* load = store->value.as()) { - if (load->buffer->data.same_as(store->buffer->data) && + if (load->buffer.same_as(store->buffer) && ArrayDeepEqual(load->indices, store->indices) && tirx::ExprDeepEqual()(load->buffer->elem_offset, store->buffer->elem_offset) && ArrayDeepEqual(load->buffer->shape, store->buffer->shape) && diff --git a/src/tirx/transform/storage_rewrite.cc b/src/tirx/transform/storage_rewrite.cc index 947f787c1344..f174c776fa17 100644 --- a/src/tirx/transform/storage_rewrite.cc +++ b/src/tirx/transform/storage_rewrite.cc @@ -108,7 +108,7 @@ class LinearAccessPatternFinder final : public StmtExprVisitor { void VisitStmt_(const AllocBufferNode* op) final { size_t level = scope_.size(); - const VarNode* buf = op->buffer->data.get(); + const VarNode* buf = op->buffer.get(); AllocEntry entry; entry.alloc = op; @@ -126,14 +126,14 @@ class LinearAccessPatternFinder final : public StmtExprVisitor { all_buffers_accessed_.insert(op->buffer.get()); // Add write access. - const VarNode* buffer_var = op->buffer->data.get(); + const VarNode* buffer_var = op->buffer.get(); auto it = alloc_info_.find(buffer_var); if (it != alloc_info_.end() && it->second.alloc) { TVM_FFI_ICHECK_LT(it->second.level, scope_.size()); scope_[it->second.level].touched.push_back(buffer_var); TVM_FFI_ICHECK_EQ(1, it->second.num_physical_dimensions) - << "Buffer " << op->buffer->name << " is allocated with " + << "BufferVar " << op->buffer.name() << " is allocated with " << it->second.num_physical_dimensions << " physical dimensions, but is accessed as having " << "1 physical dimension" << std::endl; @@ -152,7 +152,7 @@ class LinearAccessPatternFinder final : public StmtExprVisitor { all_buffers_accessed_.insert(op->buffer.get()); - const VarNode* buffer_var = op->buffer->data.get(); + const VarNode* buffer_var = op->buffer.get(); auto it = alloc_info_.find(buffer_var); if (it != alloc_info_.end() && it->second.alloc) { TVM_FFI_ICHECK_LT(it->second.level, scope_.size()) @@ -160,7 +160,7 @@ class LinearAccessPatternFinder final : public StmtExprVisitor { scope_[it->second.level].touched.push_back(buffer_var); TVM_FFI_ICHECK_EQ(1, it->second.num_physical_dimensions) - << "Buffer " << op->buffer->name << " is allocated with " + << "BufferVar " << op->buffer.name() << " is allocated with " << it->second.num_physical_dimensions << " physical dimensions, but is accessed as having " << "1 physical dimension" << std::endl; @@ -259,9 +259,9 @@ class LinearAccessPatternFinder final : public StmtExprVisitor { std::vector linear_seq_; // The storage scope of each buffer std::unordered_map alloc_info_; - // A record of which Buffer objects have been accessed, to prune + // A record of which BufferVar objects have been accessed, to prune // unused DeclBuffer instances. - std::unordered_set all_buffers_accessed_; + std::unordered_set all_buffers_accessed_; private: // Whether already in thread env. @@ -343,7 +343,7 @@ class InplaceOpVerifier : public StmtExprVisitor { this->VisitExpr(index); } --mem_nest_; - if (op->buffer->data.get() == dst_) { + if (op->buffer.get() == dst_) { store_ = op; this->VisitExpr(op->value); store_ = nullptr; @@ -371,7 +371,7 @@ class InplaceOpVerifier : public StmtExprVisitor { } void VisitExpr_(const BufferLoadNode* op) final { - const VarNode* buf = op->buffer->data.get(); + const VarNode* buf = op->buffer.get(); // cannot read from dst_ (no reduction) if (buf == dst_) { result_ = false; @@ -451,9 +451,9 @@ class StoragePlanRewriter : public StmtExprMutator { template Node VisitBufferAccess(Node node) { - auto it = alloc_map_.find(node->buffer->data.get()); + auto it = alloc_map_.find(node->buffer.get()); if (it != alloc_map_.end()) { - Buffer buf = RemapBuffer(node->buffer, it->second->alloc_var); + BufferVar buf = RemapBuffer(node->buffer, it->second->alloc_var); ffi::Array indices = node->indices; indices.Set(indices.size() - 1, @@ -466,21 +466,24 @@ class StoragePlanRewriter : public StmtExprMutator { return node; } - Buffer RemapBuffer(Buffer buf, Var new_backing_array) { + BufferVar RemapBuffer(BufferVar buf, Var new_backing_array) { auto key = buf.get(); auto it = buffer_remap_.find(key); if (it != buffer_remap_.end()) { - TVM_FFI_ICHECK_EQ(it->second->data.get(), new_backing_array.get()) - << "Cannot remap buffer " << buf->name << " to use backing array " + TVM_FFI_ICHECK_EQ(buffer_backing_.at(it->second.get()).get(), new_backing_array.get()) + << "Cannot remap buffer " << buf.name() << " to use backing array " << new_backing_array->name << ", previously used backing array " - << it->second->data->name; + << buffer_backing_.at(it->second.get()).name(); return it->second; } - Buffer remapped = Buffer(new_backing_array, buf->dtype, buf->shape, buf->strides, - buf->elem_offset, new_backing_array->name, buf->data_alignment, - buf->offset_factor, buf->span, buf->layout, buf->allocated_addr); + BufferVar backing(new_backing_array); + BufferVar remapped = + buf.same_as(backing) + ? buf + : RebuildBufferVar(buf, CopyBufferType(buf), new_backing_array->name); buffer_remap_[key] = remapped; + buffer_backing_[remapped.get()] = backing; return remapped; } @@ -569,16 +572,16 @@ class StoragePlanRewriter : public StmtExprMutator { Stmt VisitStmt_(const AllocBufferNode* op) final { // AllocBuffer combines allocation and buffer declaration. // Storage rewrite may merge this allocation with others. - if (auto it = alloc_map_.find(op->buffer->data.get()); it != alloc_map_.end()) { - if (it->second->alloc_var.get() == op->buffer->data.get()) { + if (auto it = alloc_map_.find(op->buffer.get()); it != alloc_map_.end()) { + if (it->second->alloc_var.get() == op->buffer.get()) { // This is the "winner" allocation -- its AllocBuffer was already // hoisted by PrepareNewAlloc. Strip this one. return Evaluate(0); } // This allocation was merged into another. Emit a DeclBuffer // aliasing the winner's data variable. - Buffer buf = RemapBuffer(op->buffer, it->second->alloc_var); - return DeclBuffer(buf); + BufferVar buf = RemapBuffer(op->buffer, it->second->alloc_var); + return DeclBuffer(buf, BufferVar(it->second->alloc_var).data()); } // If not in alloc_map (e.g. unused), strip entirely. return Evaluate(0); @@ -591,9 +594,11 @@ class StoragePlanRewriter : public StmtExprMutator { } auto node = StmtExprMutator::VisitStmt_(op).as_or_throw(); - if (auto it = alloc_map_.find(op->buffer->data.get()); it != alloc_map_.end()) { - Buffer buf = RemapBuffer(op->buffer, it->second->alloc_var); - node.CopyOnWrite()->buffer = buf; + if (auto it = alloc_map_.find(op->buffer.get()); it != alloc_map_.end()) { + BufferVar buf = RemapBuffer(op->buffer, it->second->alloc_var); + auto writer = node.CopyOnWrite(); + writer->buffer = buf; + writer->data = BufferVar(it->second->alloc_var).data(); } return node; } @@ -702,8 +707,8 @@ class StoragePlanRewriter : public StmtExprMutator { if (e->allocs.size() == 1 && e->allocs[0]->buffer->dtype.IsScalableVector()) { // Scalable vector lanes are runtime-dependent. Keep these allocations exact rather // than trying to compare or merge their compile-time bit size. - e->alloc_var = e->allocs[0]->buffer->data; - Buffer buf = RemapBuffer(e->allocs[0]->buffer, e->alloc_var); + e->alloc_var = e->allocs[0]->buffer.var(); + BufferVar buf = RemapBuffer(e->allocs[0]->buffer, e->alloc_var); ffi::Map annotations; if (e->is_volatile) { annotations.Set(attr::kVolatile, true); @@ -712,7 +717,7 @@ class StoragePlanRewriter : public StmtExprMutator { continue; } // Get the allocation size; - e->alloc_var = e->allocs[0]->buffer->data; + e->alloc_var = e->allocs[0]->buffer.var(); PrimType alloc_type = e->allocs[0]->buffer->dtype; for (const AllocBufferNode* op : e->allocs) { if (op->buffer->dtype.lanes() > alloc_type.lanes()) { @@ -740,7 +745,7 @@ class StoragePlanRewriter : public StmtExprMutator { if (all_allocs_identical) { // Emit AllocBuffer for the hoisted allocation. - Buffer buf = RemapBuffer(e->allocs[0]->buffer, e->alloc_var); + BufferVar buf = RemapBuffer(e->allocs[0]->buffer, e->alloc_var); ffi::Map annotations; if (e->is_volatile) { annotations.Set(attr::kVolatile, true); @@ -751,7 +756,7 @@ class StoragePlanRewriter : public StmtExprMutator { PrimExpr combo_size; for (const AllocBufferNode* op : e->allocs) { TVM_FFI_ICHECK_EQ(op->buffer->shape.size(), 1) - << "Buffer var " << op->buffer->data->name + << "BufferVar var " << op->buffer.name() << " was identified as a re-usable allocation, but has " << op->buffer->shape.size() << " physical dimensions. " << "Currently, only flat 1-d memory spaces should be identified as re-usable " @@ -784,8 +789,10 @@ class StoragePlanRewriter : public StmtExprMutator { combo_size = combo_size + IntImm::Int32(1); } combo_size = analyzer_->Simplify(combo_size); - Buffer buf(e->alloc_var, alloc_type, {combo_size}, {}, PrimExpr(), e->alloc_var->name, 0, - 0); + BufferVar buf(e->alloc_var->name, + BufferType(PointerType(alloc_type, e->scope.to_string()), alloc_type, + {combo_size}, {}, PrimExpr(), 0, 0)); + e->alloc_var = buf.var(); ffi::Map annotations; if (e->is_volatile) { annotations.Set(attr::kVolatile, true); @@ -808,7 +815,7 @@ class StoragePlanRewriter : public StmtExprMutator { if (total_bits % align != 0) { total_bits += align - (total_bits % align); } - e->alloc_var = e->allocs[0]->buffer->data; + e->alloc_var = e->allocs[0]->buffer.var(); for (StorageEntry* child : e->merged_children) { TVM_FFI_ICHECK_NE(child->const_nbits, 0U); TVM_FFI_ICHECK_NE(total_bits, 0U); @@ -822,7 +829,10 @@ class StoragePlanRewriter : public StmtExprMutator { uint64_t type_bits = e->elem_type.bits() * e->elem_type.lanes(); PrimExpr alloc_size = MakeConst(e->allocs[0]->buffer->shape[0].ty(), (total_bits + type_bits - 1) / type_bits); - Buffer buf(e->alloc_var, e->elem_type, {alloc_size}, {}, PrimExpr(), e->alloc_var->name, 0, 0); + BufferVar buf(e->alloc_var->name, + BufferType(PointerType(e->elem_type, e->scope.to_string()), e->elem_type, + {alloc_size}, {}, PrimExpr(), 0, 0)); + e->alloc_var = buf.var(); bool any_volatile = e->is_volatile; for (StorageEntry* child : e->merged_children) { if (child->is_volatile) any_volatile = true; @@ -1115,12 +1125,14 @@ class StoragePlanRewriter : public StmtExprMutator { // The allocations std::vector> alloc_vec_; // The buffer objects being remapped - std::unordered_map buffer_remap_; + std::unordered_map buffer_remap_; + // Explicit physical backing for each remapped buffer view. + std::unordered_map buffer_backing_; // Buffers whose DeclBuffer has been hoisted to be adjacent to the new AllocBuffer location - std::unordered_set hoisted_buffer_decls_; + std::unordered_set hoisted_buffer_decls_; // Any buffers that is accessed at some point. DeclBuffer instances // that do not appear in this list may be removed. - std::unordered_set all_buffers_accessed_; + std::unordered_set all_buffers_accessed_; // Copy of the allocation info from LinearAccessPatternFinder. std::unordered_map alloc_info_; // analyzer @@ -1136,6 +1148,7 @@ struct BufferVarInfo { kPrimFuncBufferMap = (1 << 1), kAllocBufferNode = (1 << 2), kLetNode = (1 << 3), + kDeclBufferNode = (1 << 4), }; // The tirx::Var that represents this buffer. @@ -1227,7 +1240,7 @@ class VectorTypeAccessChecker : public StmtExprVisitor { * type as it is later accessed, with scalar element types. */ VectorTypeAccessChecker(const ffi::Array& params, - const ffi::Map& buffer_map, + const ffi::Map& buffer_map, bool allow_untyped_pointers = false, bool detect_scalar_read_patterns = true) : allow_untyped_pointers_(allow_untyped_pointers), @@ -1235,8 +1248,8 @@ class VectorTypeAccessChecker : public StmtExprVisitor { // If a parameter is in the buffer map, we want to track the // version in the map. for (auto it : buffer_map) { - Buffer& buffer = it.second; - Var buffer_var = buffer->data; + BufferVar& buffer = it.second; + Var buffer_var = buffer.var(); PrimType dtype = buffer->dtype; PrimExpr extent = buffer->shape.size() ? buffer->shape[buffer->shape.size() - 1] : 0; OnArrayDeclaration(buffer_var, dtype, extent, BufferVarInfo::kPrimFuncParam); @@ -1245,6 +1258,14 @@ class VectorTypeAccessChecker : public StmtExprVisitor { // If a pointer parameter isn't in the buffer map, then we want to // track the parameter itself. for (Var buffer_var : params) { + if (auto buffer_type = buffer_var->ty.as()) { + BufferVar buffer(buffer_var); + PrimExpr extent = + buffer->shape.size() ? buffer->shape[buffer->shape.size() - 1] : PrimExpr(0); + OnArrayDeclaration(buffer_var, buffer->dtype, extent, + BufferVarInfo::kPrimFuncParam); + continue; + } auto pointer_type = GetPointerType(buffer_var->ty); if (pointer_type.has_value() && !pointer_type.value().IsVoid() && (buffer_map.count(buffer_var) == 0)) { @@ -1259,13 +1280,13 @@ class VectorTypeAccessChecker : public StmtExprVisitor { } void VisitExpr_(const BufferLoadNode* op) final { - OnArrayAccess(op->ty.as_or_throw(), op->buffer->data.get(), op->indices, + OnArrayAccess(op->ty.as_or_throw(), op->buffer.get(), op->indices, /*is_buffer_load=*/true); StmtExprVisitor::VisitExpr_(op); } void VisitStmt_(const BufferStoreNode* op) final { - OnArrayAccess(op->value.ty(), op->buffer->data.get(), op->indices, /*is_buffer_load=*/false); + OnArrayAccess(op->value.ty(), op->buffer.get(), op->indices, /*is_buffer_load=*/false); StmtExprVisitor::VisitStmt_(op); } @@ -1282,7 +1303,7 @@ class VectorTypeAccessChecker : public StmtExprVisitor { } } else if (op->op.same_as(builtin::address_of())) { if (const auto* load = op->args[0].as()) { - OnArrayAccess(load->ty.as_or_throw(), load->buffer->data.get(), load->indices, + OnArrayAccess(load->ty.as_or_throw(), load->buffer.get(), load->indices, /*is_buffer_load=*/false); } } @@ -1292,12 +1313,40 @@ class VectorTypeAccessChecker : public StmtExprVisitor { void VisitStmt_(const AllocBufferNode* op) final { const ffi::Array& shape = op->buffer->shape; PrimExpr extent = shape.size() ? shape[shape.size() - 1] : PrimExpr(0); - OnArrayDeclaration(op->buffer->data, op->buffer->dtype, extent, + OnArrayDeclaration(op->buffer.var(), op->buffer->dtype, extent, BufferVarInfo::kAllocBufferNode); StmtExprVisitor::VisitStmt_(op); } + void VisitStmt_(const DeclBufferNode* op) final { + const ffi::Array& shape = op->buffer->shape; + PrimExpr extent = shape.size() ? shape[shape.size() - 1] : PrimExpr(0); + OnArrayDeclaration(op->buffer.var(), op->buffer->dtype, extent, + BufferVarInfo::kDeclBufferNode); + if (op->data.has_value()) { + if (auto source = op->data.value().as()) { + decl_buffer_sources_.emplace_back(op->buffer.get(), source.value().get()); + } + } + + StmtExprVisitor::VisitStmt_(op); + } + + void PropagateDeclBufferAccesses() { + for (const auto& [buffer, source] : decl_buffer_sources_) { + auto buffer_it = info_map_.find(buffer); + auto source_it = info_map_.find(source); + if (buffer_it == info_map_.end() || source_it == info_map_.end()) { + continue; + } + source_it->second.access_dtype.insert(buffer_it->second.access_dtype.begin(), + buffer_it->second.access_dtype.end()); + source_it->second.scalar_read_dtype.insert(buffer_it->second.scalar_read_dtype.begin(), + buffer_it->second.scalar_read_dtype.end()); + } + } + void VisitExpr_(const LetNode* op) final { HandleLetNode(op->var); StmtExprVisitor::VisitExpr_(op); @@ -1443,6 +1492,7 @@ class VectorTypeAccessChecker : public StmtExprVisitor { // Map of buffer variable information determined std::unordered_map info_map_; + std::vector> decl_buffer_sources_; // bool allow_untyped_pointers_{false}; @@ -1513,7 +1563,7 @@ class VectorTypeRewriter : public StmtExprMutator { rewrite_mask |= BufferVarInfo::kPrimFuncBufferMap; } if (rewrite_alloc_buffer_node) { - rewrite_mask |= BufferVarInfo::kAllocBufferNode; + rewrite_mask |= BufferVarInfo::kAllocBufferNode | BufferVarInfo::kDeclBufferNode; } if (rewrite_let_node) { rewrite_mask |= BufferVarInfo::kLetNode; @@ -1524,9 +1574,25 @@ class VectorTypeRewriter : public StmtExprMutator { PrimType preferred = var_info.get_preferred_dtype(); if (preferred != var_info.element_dtype && (rewrite_mask & var_info.declaration_location)) { Var old_buffer_var = var_info.var; - Var new_buffer_var(old_buffer_var->name, - PointerType(preferred, GetPtrStorageScope(old_buffer_var)), - old_buffer_var->span); + Var new_buffer_var = [&]() -> Var { + if (old_buffer_var->ty.as()) { + BufferVar old_buffer(old_buffer_var); + auto type = CopyBufferType(old_buffer); + type->data_pointer_type = PointerType(preferred, old_buffer.scope()); + type->dtype = preferred; + if (!type->shape.empty()) { + PrimExpr last_dim = type->shape.back(); + int factor = preferred.lanes() / var_info.element_dtype.lanes(); + type->shape.Set(type->shape.size() - 1, + last_dim / MakeConst(last_dim.ty(), factor)); + } + type->layout = std::nullopt; + return RebuildBufferVar(old_buffer, std::move(type)).var(); + } + return Var(old_buffer_var->name, + PointerType(preferred, GetPtrStorageScope(old_buffer_var)), + old_buffer_var->span); + }(); rewrite_map_.insert_or_assign( var_info.var.get(), @@ -1547,7 +1613,7 @@ class VectorTypeRewriter : public StmtExprMutator { return {node, shuffle_index}; } - auto it = rewrite_map_.find(node->buffer->data.get()); + auto it = rewrite_map_.find(node->buffer.get()); if (it == rewrite_map_.end()) { return {node, shuffle_index}; } @@ -1627,7 +1693,7 @@ class VectorTypeRewriter : public StmtExprMutator { } Stmt VisitStmt_(const AllocBufferNode* op) final { - Buffer new_buf = RemapBuffer(op->buffer); + BufferVar new_buf = RemapBuffer(op->buffer); if (new_buf.same_as(op->buffer)) { return ffi::GetRef(op); } @@ -1636,7 +1702,20 @@ class VectorTypeRewriter : public StmtExprMutator { return Stmt(n); } - Buffer RemapBuffer(Buffer buf) { + Stmt VisitStmt_(const DeclBufferNode* op) final { + Stmt stmt = StmtExprMutator::VisitStmt_(op); + op = stmt.as(); + TVM_FFI_ICHECK(op != nullptr); + BufferVar new_buf = RemapBuffer(op->buffer); + if (new_buf.same_as(op->buffer)) { + return stmt; + } + auto n = CopyOnWrite(op); + n->buffer = std::move(new_buf); + return Stmt(n); + } + + BufferVar RemapBuffer(BufferVar buf) { auto cache_key = buf.get(); auto cache_it = buffer_map_.find(cache_key); @@ -1644,19 +1723,10 @@ class VectorTypeRewriter : public StmtExprMutator { return cache_it->second; } - auto info_it = rewrite_map_.find(buf->data.get()); + auto info_it = rewrite_map_.find(buf.get()); if (info_it != rewrite_map_.end()) { auto& info = info_it->second; - - ffi::Array shape = buf->shape; - PrimExpr last_dim = shape[shape.size() - 1]; - shape.Set(shape.size() - 1, last_dim / MakeConst(last_dim.ty(), info.factor())); - - auto writer = buf.CopyOnWrite(); - writer->data = info.new_buffer_var; - writer->dtype = info.new_element_dtype; - writer->shape = shape; - writer->layout = std::nullopt; + buf = BufferVar(info.new_buffer_var); } buffer_map_[cache_key] = buf; @@ -1731,14 +1801,14 @@ class VectorTypeRewriter : public StmtExprMutator { return ffi::GetRef(op); } - Buffer VisitBufferDef(const Buffer& buffer, bool alloc_data) final { - Buffer new_buffer = StmtExprMutator::VisitBufferDef(buffer, alloc_data); - auto it = var_remap_.find(new_buffer->data.get()); + BufferVar VisitBufferDef(const BufferVar& buffer, bool alloc_data) final { + auto it = var_remap_.find(buffer.get()); if (it != var_remap_.end()) { - new_buffer.CopyOnWrite()->data = it->second; + BufferVar new_buffer(it->second); buffer_remap_.Set(buffer, new_buffer); + return new_buffer; } - return new_buffer; + return StmtExprMutator::VisitBufferDef(buffer, alloc_data); } Stmt VisitStmt_(const AttrStmtNode* op) final { @@ -1771,13 +1841,13 @@ class VectorTypeRewriter : public StmtExprMutator { } n->params = new_params; - // Remap the Buffer objects in PrimFunc::buffer_map so that the + // Remap the BufferVar objects in PrimFunc::buffer_map so that the // buffers use the new buffer variables - ffi::Map new_buffer_map; + ffi::Map new_buffer_map; for (const auto& pair : n->buffer_map) { Var key = pair.first; - Buffer old_buffer = pair.second; - Buffer new_buffer = RemapBuffer(old_buffer); + BufferVar old_buffer = pair.second; + BufferVar new_buffer = RemapBuffer(old_buffer); new_buffer_map.Set(key, new_buffer); } n->buffer_map = new_buffer_map; @@ -1800,7 +1870,7 @@ class VectorTypeRewriter : public StmtExprMutator { bool rewrite_indices_{true}; std::unordered_map rewrite_map_; - std::unordered_map buffer_map_; + std::unordered_map buffer_map_; arith::Analyzer analyzer_; }; @@ -1814,6 +1884,7 @@ PrimFunc PointerValueTypeRewrite(PrimFunc f, bool allow_untyped_pointers = false VectorTypeAccessChecker checker(f->params, f->buffer_map, allow_untyped_pointers, rewrite_scalar_read_to_vector_shuffle); checker(f->body); + checker.PropagateDeclBufferAccesses(); VectorTypeRewriter rewriter(checker.info_map_, rewrite_params, rewrite_buffer_map, rewrite_alloc_buffer_node, rewrite_indices, rewrite_let_node, diff --git a/src/tirx/transform/tile_primitive_dispatch.cc b/src/tirx/transform/tile_primitive_dispatch.cc index 5322c45b37ad..1f192999bd43 100644 --- a/src/tirx/transform/tile_primitive_dispatch.cc +++ b/src/tirx/transform/tile_primitive_dispatch.cc @@ -254,7 +254,7 @@ class TilePrimitiveDispatcher : public StmtExprMutator { private: class BufferRefRewriter : public StmtExprMutator { public: - static Stmt Rewrite(const Stmt& stmt, const Buffer& src, const Buffer& dst) { + static Stmt Rewrite(const Stmt& stmt, const BufferVar& src, const BufferVar& dst) { if (src.same_as(dst)) { return stmt; } @@ -262,25 +262,25 @@ class TilePrimitiveDispatcher : public StmtExprMutator { } private: - BufferRefRewriter(Buffer src, Buffer dst) : src_(std::move(src)), dst_(std::move(dst)) {} + BufferRefRewriter(BufferVar src, BufferVar dst) : src_(std::move(src)), dst_(std::move(dst)) {} - Buffer VisitBufferDef(const Buffer& buffer, bool alloc_data) final { - Buffer new_buffer = StmtExprMutator::VisitBufferDef(buffer, alloc_data); + BufferVar VisitBufferDef(const BufferVar& buffer, bool alloc_data) final { + BufferVar new_buffer = StmtExprMutator::VisitBufferDef(buffer, alloc_data); if (new_buffer.same_as(src_)) { return dst_; } return new_buffer; } - Buffer VisitBufferUse(const Buffer& buffer) final { + BufferVar VisitBufferUse(const BufferVar& buffer) final { if (buffer.same_as(src_)) { return dst_; } return StmtExprMutator::VisitBufferUse(buffer); } - Buffer src_; - Buffer dst_; + BufferVar src_; + BufferVar dst_; }; class KernelReplacePointSearcher : public StmtExprMutator { @@ -478,7 +478,7 @@ class TilePrimitiveDispatcher : public StmtExprMutator { } Stmt VisitStmt_(const AllocBufferNode* op) final { - Buffer old_buffer = op->buffer; + BufferVar old_buffer = op->buffer; Stmt stmt = StmtExprMutator::VisitStmt_(op); op = stmt.as(); TVM_FFI_ICHECK(op); @@ -489,7 +489,7 @@ class TilePrimitiveDispatcher : public StmtExprMutator { } Stmt VisitStmt_(const DeclBufferNode* op) final { - Buffer old_buffer = op->buffer; + BufferVar old_buffer = op->buffer; Stmt stmt = StmtExprMutator::VisitStmt_(op); op = stmt.as(); TVM_FFI_ICHECK(op); @@ -559,7 +559,7 @@ class TilePrimitiveDispatcher : public StmtExprMutator { TVM_FFI_ICHECK(res.defined()) << "TIRx dispatcher did not return a PrimFunc"; // Implementation found, handle callbacks if (auto bufs = sctx->callbacks.Get(tirx::callback::kPrivateAlloc)) { - auto buf_list = bufs.value().as>().value(); + auto buf_list = bufs.value().as>().value(); alloc_buffers_.insert(alloc_buffers_.end(), buf_list.begin(), buf_list.end()); } if (auto stmts = sctx->callbacks.Get(tirx::callback::kDeviceInitStmt)) { @@ -571,7 +571,7 @@ class TilePrimitiveDispatcher : public StmtExprMutator { host_init_stmts_.insert(host_init_stmts_.end(), stmt_list.begin(), stmt_list.end()); } if (auto mapping = sctx->callbacks.Get(tirx::callback::kPostBufferDefStmt)) { - auto map = mapping.value().as_or_throw>>(); + auto map = mapping.value().as_or_throw>>(); for (const auto& [buffer, stmts] : map) { auto& vec = post_buffer_def_stmts_[buffer]; vec.insert(vec.end(), stmts.begin(), stmts.end()); @@ -1431,10 +1431,10 @@ class TilePrimitiveDispatcher : public StmtExprMutator { std::vector> scope_id_defs_at_level_; std::vector ctx_stack_; std::unordered_map launch_params_; - std::vector alloc_buffers_; + std::vector alloc_buffers_; std::vector device_init_stmts_; std::vector host_init_stmts_; - std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> + std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> post_buffer_def_stmts_; ffi::Map shared_state_; std::vector> cluster_cta_axis_extents_; @@ -1442,10 +1442,10 @@ class TilePrimitiveDispatcher : public StmtExprMutator { bool is_first_block_{true}; bool is_first_thread_attr_{true}; - bool AppendPostBufferDefStmts(std::vector* seq, const Buffer& old_buffer, - const Buffer& new_buffer) { + bool AppendPostBufferDefStmts(std::vector* seq, const BufferVar& old_buffer, + const BufferVar& new_buffer) { auto append_with_remap = [this, seq, &new_buffer](auto it) -> bool { - Buffer src = it->first; + BufferVar src = it->first; for (const auto& stmt : it->second) { Stmt remapped = BufferRefRewriter::Rewrite(stmt, src, new_buffer); seq->push_back(KernelReplacePointSearcher::Seek(remapped, Evaluate(0))); diff --git a/src/tirx/transform/tvm_ffi_binder.cc b/src/tirx/transform/tvm_ffi_binder.cc index a2d44f6a4e9f..1ff7714c52c1 100644 --- a/src/tirx/transform/tvm_ffi_binder.cc +++ b/src/tirx/transform/tvm_ffi_binder.cc @@ -44,7 +44,7 @@ using ffi::reflection::AccessStep; // ============================================================ TVMFFIABIBuilder::TVMFFIABIBuilder(const ffi::String& func_name, const ffi::Array& params, - const ffi::Map& buffer_map, + const ffi::Map& buffer_map, const Var& v_packed_args, const Var& v_num_packed_args, const PrimExpr& device_type, const PrimExpr& device_id) : func_name_(func_name), @@ -60,8 +60,8 @@ TVMFFIABIBuilder::TVMFFIABIBuilder(const ffi::String& func_name, const ffi::Arra if (i > 0) os << ", "; Var param = params[i]; if (buffer_map.count(param)) { - Buffer buf = buffer_map[param]; - std::string buf_name = buf->name; + BufferVar buf = buffer_map[param]; + std::string buf_name = buf.name(); os << buf_name << ": Tensor(["; for (size_t j = 0; j < buf->shape.size(); ++j) { if (j > 0) os << ", "; @@ -388,12 +388,12 @@ void TVMFFIABIBuilder::BindArray(const ffi::Array& arg, const ffi::Arr // BindBuffer (buffer-to-buffer bind with ffi::reflection::AccessPath) // ============================================================ -void TVMFFIABIBuilder::BindBuffer(const Buffer& arg, const Buffer& value, +void TVMFFIABIBuilder::BindBuffer(const BufferVar& arg, const BufferVar& value, ffi::reflection::AccessPath base_path, bool fuzzy_match) { TVM_FFI_ICHECK_EQ(arg.scope(), value.scope()) - << "Argument " << arg->name << " Buffer bind scope mismatch"; + << "Argument " << arg.name() << " BufferVar bind scope mismatch"; TVM_FFI_ICHECK_EQ(arg->dtype, value->dtype) - << "Argument " << arg->name << " Buffer bind data type mismatch"; + << "Argument " << arg.name() << " BufferVar bind data type mismatch"; if (value->data_alignment % arg->data_alignment != 0) { LOG(WARNING) << "Trying to bind buffer to another one with lower alignment requirement " << " required alignment=" << arg->data_alignment @@ -403,12 +403,11 @@ void TVMFFIABIBuilder::BindBuffer(const Buffer& arg, const Buffer& value, if (value->elem_offset.defined()) { if (is_zero(arg->elem_offset)) { TVM_FFI_ICHECK(is_zero(value->elem_offset)) - << "Trying to bind a Buffer with offset into one without offset " + << "Trying to bind a BufferVar with offset into one without offset " << " required elem_offset=" << arg->elem_offset << ", provided elem_offset=" << value->elem_offset; } ffi::reflection::AccessPath data_path = base_path->Attr(ffi::String("data")); - BindPointer(arg->data, value->data, data_path, false); ffi::reflection::AccessPath offset_path = base_path->Attr(ffi::String("elem_offset")); if (BindScalar(arg->elem_offset, value->elem_offset, offset_path, false)) { if (arg->offset_factor > 1) { @@ -438,11 +437,11 @@ void TVMFFIABIBuilder::BindBuffer(const Buffer& arg, const Buffer& value, ffi::reflection::AccessPath strides_path = base_path->Attr(ffi::String("strides")); if (arg->shape.size() < value->shape.size()) { - TVM_FFI_ICHECK(fuzzy_match) << "Buffer size mismatch at " << RenderAccessPath(base_path); + TVM_FFI_ICHECK(fuzzy_match) << "BufferVar size mismatch at " << RenderAccessPath(base_path); size_t diff = value->shape.size() - arg->shape.size(); for (size_t i = 0; i < diff; ++i) { TVM_FFI_ICHECK(is_one(analyzer_->Simplify(value->shape[i]))) - << "Buffer shape mismatch at " << RenderAccessPath(base_path) << ": " << arg->shape + << "BufferVar shape mismatch at " << RenderAccessPath(base_path) << ": " << arg->shape << " vs " << value->shape; } for (size_t i = 0; i < arg->shape.size(); ++i) { @@ -599,13 +598,13 @@ void TVMFFIABIBuilder::DecodeAllParams() { for (int i = 0; i < num_args; ++i) { Var param = params_[i]; if (buffer_map_.count(param)) { - Buffer buffer = buffer_map_[param]; + BufferVar buffer = buffer_map_[param]; ffi::reflection::AccessPath param_path = ffi::reflection::AccessPath::Root() ->Extend(AccessStep::ArrayItem(i)) - ->Attr(ffi::String(buffer->name)); - DecodeParamDLTensor(buffer, device_type_, device_id_, param, func_name_ + "." + param->name, - param_path); - decl_buffers_.push_back(DeclBuffer(buffer)); + ->Attr(ffi::String(buffer.name())); + Expr data = DecodeParamDLTensor(buffer, device_type_, device_id_, param, + func_name_ + "." + param->name, param_path); + decl_buffers_.push_back(DeclBuffer(buffer, data)); } } } @@ -636,7 +635,7 @@ PrimExpr TVMFFIABIBuilder::LoadInt64ArrayElem(const Var& ptr, int index) { // Strides validation subfunctions // ============================================================ -void TVMFFIABIBuilder::BindCompactStrides(const Buffer& buffer, const Var& strides_ptr, +void TVMFFIABIBuilder::BindCompactStrides(const BufferVar& buffer, const Var& strides_ptr, const PrimExpr& v_strides_is_null, const ffi::reflection::AccessPath& param_path) { PrimType stype(buffer->DefaultIndexType()); @@ -654,7 +653,7 @@ void TVMFFIABIBuilder::BindCompactStrides(const Buffer& buffer, const Var& strid foldl([](PrimExpr a, PrimExpr b, Span span) { return logical_and(a, b, span); }, IntImm::Bool(true), conds), StringImm("ValueError"), - ffi::Array({StringImm("Mismatched "), StringImm(buffer->name), + ffi::Array({StringImm("Mismatched "), StringImm(buffer.name()), StringImm(".strides on argument #"), StringImm(std::to_string(param_index)), when_calling_imm_, sig_imm_, StringImm("`,\n expected to be compact array")})); @@ -663,7 +662,7 @@ void TVMFFIABIBuilder::BindCompactStrides(const Buffer& buffer, const Var& strid } } -void TVMFFIABIBuilder::BindRegularStrides(const Buffer& buffer, const Var& strides_ptr, +void TVMFFIABIBuilder::BindRegularStrides(const BufferVar& buffer, const Var& strides_ptr, const Var& shape_ptr, const PrimExpr& v_strides_is_null, const ffi::reflection::AccessPath& param_path) { PrimExpr stride_from_shape = 1; @@ -682,13 +681,13 @@ void TVMFFIABIBuilder::BindRegularStrides(const Buffer& buffer, const Var& strid // DecodeParamDLTensor (private) // ============================================================ -void TVMFFIABIBuilder::DecodeParamDLTensor(const Buffer& buffer, const PrimExpr& device_type, - const PrimExpr& device_id, const Var& handle, - const std::string& arg_name, - ffi::reflection::AccessPath base_path) { +Expr TVMFFIABIBuilder::DecodeParamDLTensor(const BufferVar& buffer, const PrimExpr& device_type, + const PrimExpr& device_id, const Var& handle, + const std::string& arg_name, + ffi::reflection::AccessPath base_path) { const PrimType tvm_ndim_type = PrimType::Int(32); - std::string buf_name = buffer->name; + std::string buf_name = buffer.name(); ffi::reflection::AccessPath param_path = base_path; int param_index = GetParamIndex(base_path); @@ -807,9 +806,9 @@ void TVMFFIABIBuilder::DecodeParamDLTensor(const Buffer& buffer, const PrimExpr& { ffi::reflection::AccessPath data_path = param_path->Attr(ffi::String("data")); Expr raw_data = TVMStructGet(PointerType::VoidPointerTy(), handle, 0, builtin::kDLTensorData); - Expr typed_data = Call(buffer->data->ty, builtin::reinterpret(), {raw_data}); - if (BindPointer(buffer->data, typed_data, data_path, true)) { - Var vptr(buffer->data); + Expr typed_data = Call(buffer->data_pointer_type, builtin::reinterpret(), {raw_data}); + { + Expr vptr = typed_data; auto alloc_size = [&]() -> PrimExpr { PrimExpr product = IntImm(PrimType(buffer->DefaultIndexType()), 1); @@ -859,6 +858,7 @@ void TVMFFIABIBuilder::DecodeParamDLTensor(const Buffer& buffer, const PrimExpr& IntImm::Int32(buffer->data_alignment), Evaluate(0))); } } + return typed_data; } } diff --git a/src/tirx/transform/tvm_ffi_binder.h b/src/tirx/transform/tvm_ffi_binder.h index a90a295f492d..9b1db11f87f2 100644 --- a/src/tirx/transform/tvm_ffi_binder.h +++ b/src/tirx/transform/tvm_ffi_binder.h @@ -100,7 +100,7 @@ class TVMFFIABIBuilder { std::vector init_nest; /*! \brief Validation checks (all AssertStmts). */ std::vector asserts; - /*! \brief Buffer declarations for buffer_map entries. */ + /*! \brief BufferVar declarations for buffer_map entries. */ std::vector decl_buffers; }; @@ -119,7 +119,7 @@ class TVMFFIABIBuilder { * \param device_id The device id variable (may be defined during buffer binding). */ TVMFFIABIBuilder(const ffi::String& func_name, const ffi::Array& params, - const ffi::Map& buffer_map, const Var& v_packed_args, + const ffi::Map& buffer_map, const Var& v_packed_args, const Var& v_num_packed_args, const PrimExpr& device_type, const PrimExpr& device_id); @@ -263,7 +263,7 @@ class TVMFFIABIBuilder { const ffi::reflection::AccessPath& base_path); /*! - * \brief Buffer-to-buffer bind with ffi::reflection::AccessPath. + * \brief BufferVar-to-buffer bind with ffi::reflection::AccessPath. * * Binds data, elem_offset, shape, and strides of \p arg against \p value, * emitting assertions for any mismatches. @@ -273,7 +273,7 @@ class TVMFFIABIBuilder { * \param base_path Base ffi::reflection::AccessPath for the buffer parameter. * \param fuzzy_match If true, allow value to have more dimensions than arg. */ - void BindBuffer(const Buffer& arg, const Buffer& value, ffi::reflection::AccessPath base_path, + void BindBuffer(const BufferVar& arg, const BufferVar& value, ffi::reflection::AccessPath base_path, bool fuzzy_match); /*! @@ -286,7 +286,7 @@ class TVMFFIABIBuilder { * \param arg_name Human-readable name for error messages. * \param base_path Base ffi::reflection::AccessPath for the tensor parameter. */ - void DecodeParamDLTensor(const Buffer& buffer, const PrimExpr& device_type, + Expr DecodeParamDLTensor(const BufferVar& buffer, const PrimExpr& device_type, const PrimExpr& device_id, const Var& handle, const std::string& arg_name, ffi::reflection::AccessPath base_path); @@ -321,7 +321,7 @@ class TVMFFIABIBuilder { * \param v_strides_is_null Expression checking if strides pointer is NULL. * \param param_path ffi::reflection::AccessPath for the tensor parameter. */ - void BindCompactStrides(const Buffer& buffer, const Var& strides_ptr, + void BindCompactStrides(const BufferVar& buffer, const Var& strides_ptr, const PrimExpr& v_strides_is_null, const ffi::reflection::AccessPath& param_path); @@ -334,7 +334,7 @@ class TVMFFIABIBuilder { * \param v_strides_is_null Expression checking if strides pointer is NULL. * \param param_path ffi::reflection::AccessPath for the tensor parameter. */ - void BindRegularStrides(const Buffer& buffer, const Var& strides_ptr, const Var& shape_ptr, + void BindRegularStrides(const BufferVar& buffer, const Var& strides_ptr, const Var& shape_ptr, const PrimExpr& v_strides_is_null, const ffi::reflection::AccessPath& param_path); @@ -391,7 +391,7 @@ class TVMFFIABIBuilder { std::vector init_nest_; /*! \brief Validation checks: all AssertStmts. */ std::vector asserts_; - /*! \brief Buffer declarations for buffer_map entries. */ + /*! \brief BufferVar declarations for buffer_map entries. */ std::vector decl_buffers_; /*! \brief Deferred constant-expression assertions for display-var substitution. */ std::vector pending_const_asserts_; @@ -406,7 +406,7 @@ class TVMFFIABIBuilder { /*! \brief The function parameters. */ ffi::Array params_; /*! \brief The buffer map from parameters to buffers. */ - ffi::Map buffer_map_; + ffi::Map buffer_map_; /*! \brief The packed args variable. */ Var v_packed_args_; /*! \brief The expected device type expression. */ diff --git a/src/tirx/transform/unroll_loop.cc b/src/tirx/transform/unroll_loop.cc index 7e29179a60be..82f2a73c2b0d 100644 --- a/src/tirx/transform/unroll_loop.cc +++ b/src/tirx/transform/unroll_loop.cc @@ -167,7 +167,7 @@ class LoopUnroller : public StmtExprMutator { Expr VisitExpr_(const BufferLoadNode* op) final { if (unroll_local_access_) { - auto storage_scope = runtime::StorageScope::Create(GetPtrStorageScope(op->buffer->data)); + auto storage_scope = runtime::StorageScope::Create(op->buffer.scope()); if (storage_scope.rank == runtime::StorageRank::kLocal || storage_scope.rank == runtime::StorageRank::kWarp) { VarLocalAccessMarker marker(&var_touched_local_); @@ -182,7 +182,7 @@ class LoopUnroller : public StmtExprMutator { Stmt VisitStmt_(const BufferStoreNode* op) final { ++step_count_; if (unroll_local_access_) { - auto storage_scope = runtime::StorageScope::Create(GetPtrStorageScope(op->buffer->data)); + auto storage_scope = runtime::StorageScope::Create(op->buffer.scope()); if (storage_scope.rank == runtime::StorageRank::kLocal || storage_scope.rank == runtime::StorageRank::kWarp) { VarLocalAccessMarker marker(&var_touched_local_); diff --git a/src/tirx/transform/unsupported_dtype_legalize.cc b/src/tirx/transform/unsupported_dtype_legalize.cc index 912bf189f0e9..2ab0093d936d 100644 --- a/src/tirx/transform/unsupported_dtype_legalize.cc +++ b/src/tirx/transform/unsupported_dtype_legalize.cc @@ -71,7 +71,7 @@ bool MatchPrimType(const Type& type, F f) { class ComputeLegalizePlanner : public StmtExprVisitor { public: ComputeLegalizePlanner( - std::unordered_map* buffer_remap, + std::unordered_map* buffer_remap, std::unordered_map* var_remap, PrimType promote_dtype) : buffer_remap_(buffer_remap), var_remap_(var_remap), promote_dtype_(promote_dtype) {} @@ -86,13 +86,13 @@ class ComputeLegalizePlanner : public StmtExprVisitor { var_remap_->erase(it); } } - ffi::Array drop_buffers; + ffi::Array drop_buffers; for (auto kv : *buffer_remap_) { - if (opaque_var_access_.count(kv.first->data)) { + if (opaque_var_access_.count(kv.first.var())) { drop_buffers.push_back(kv.first); } } - for (Buffer buffer : drop_buffers) { + for (BufferVar buffer : drop_buffers) { auto it = buffer_remap_->find(buffer); TVM_FFI_ICHECK(it != buffer_remap_->end()); buffer_remap_->erase(it); @@ -116,11 +116,14 @@ class ComputeLegalizePlanner : public StmtExprVisitor { if (MatchType(op->buffer->dtype)) { PrimType dtype = promote_dtype_.WithLanes(op->buffer->dtype.lanes()); ffi::String storage_scope = "global"; - if (auto* ptr_type = op->buffer->data->ty.as()) { + if (auto* ptr_type = op->buffer->data_pointer_type.as()) { storage_scope = ptr_type->storage_scope; } - Var buffer_var = Var(op->buffer->data->name, PointerType(dtype, storage_scope)); - (*var_remap_)[op->buffer->data] = buffer_var; + auto type = CopyBufferType(op->buffer); + type->data_pointer_type = PointerType(dtype, storage_scope); + type->dtype = dtype; + BufferVar buffer_var = RebuildBufferVar(op->buffer, std::move(type)); + (*var_remap_)[op->buffer.var()] = buffer_var.var(); } return StmtExprVisitor::VisitStmt_(op); } @@ -139,17 +142,15 @@ class ComputeLegalizePlanner : public StmtExprVisitor { } private: - void PopulateBufferRemap(Buffer buf) { - auto var_it = var_remap_->find(buf->data); + void PopulateBufferRemap(BufferVar buf) { + auto var_it = var_remap_->find(buf.var()); if (var_it == var_remap_->end()) return; - Buffer new_buffer(var_it->second, promote_dtype_.WithLanes(buf->dtype.lanes()), buf->shape, - buf->strides, buf->elem_offset, buf->name, buf->data_alignment, - buf->offset_factor, buf->span, buf->layout, buf->allocated_addr); + BufferVar new_buffer(var_it->second); (*buffer_remap_)[buf] = new_buffer; } - std::unordered_map* buffer_remap_; + std::unordered_map* buffer_remap_; std::unordered_map* var_remap_; std::unordered_set opaque_var_access_; PrimType promote_dtype_; @@ -158,7 +159,7 @@ class ComputeLegalizePlanner : public StmtExprVisitor { class BF16ComputeLegalizePlanner : public ComputeLegalizePlanner { public: explicit BF16ComputeLegalizePlanner( - std::unordered_map* buffer_remap, + std::unordered_map* buffer_remap, std::unordered_map* var_remap, PrimType promote_dtype) : ComputeLegalizePlanner(buffer_remap, var_remap, promote_dtype) {} bool MatchType(const Type& type) const { @@ -169,7 +170,7 @@ class BF16ComputeLegalizePlanner : public ComputeLegalizePlanner { class FP8ComputeLegalizePlanner : public ComputeLegalizePlanner { public: explicit FP8ComputeLegalizePlanner( - std::unordered_map* buffer_remap, + std::unordered_map* buffer_remap, std::unordered_map* var_remap, PrimType promote_dtype) : ComputeLegalizePlanner(buffer_remap, var_remap, promote_dtype) {} bool MatchType(const Type& type) const { @@ -363,7 +364,7 @@ class ComputeLegalizer : public StmtExprMutator { predicate = this->VisitPrimExpr(op->predicate.value()); } - Buffer new_buf = GetRemappedBuffer(op->buffer); + BufferVar new_buf = GetRemappedBuffer(op->buffer); if (value.same_as(op->value) && indices.same_as(op->indices) && predicate.same_as(op->predicate) && new_buf.same_as(op->buffer)) { @@ -388,7 +389,7 @@ class ComputeLegalizer : public StmtExprMutator { Stmt VisitStmt_(const AttrStmtNode* op) final { Stmt ret = StmtExprMutator::VisitStmt_(op); op = ret.as(); - if (auto buffer = op->node.as()) { + if (auto buffer = op->node.as()) { auto it = buffer_remap_.find(buffer.value()); if (it != buffer_remap_.end()) { return AttrStmt(it->second, op->attr_key, op->value, op->body); @@ -443,7 +444,7 @@ class ComputeLegalizer : public StmtExprMutator { Stmt ret = StmtExprMutator::VisitStmt_(op); op = ret.as(); - Buffer new_buf = GetRemappedBuffer(op->buffer); + BufferVar new_buf = GetRemappedBuffer(op->buffer); if (new_buf.same_as(op->buffer)) { return ret; } else { @@ -455,7 +456,7 @@ class ComputeLegalizer : public StmtExprMutator { Stmt ret = StmtExprMutator::VisitStmt_(op); op = ret.as(); - Buffer new_buf = GetRemappedBuffer(op->buffer); + BufferVar new_buf = GetRemappedBuffer(op->buffer); if (new_buf.same_as(op->buffer)) { return ret; } else { @@ -469,7 +470,7 @@ class ComputeLegalizer : public StmtExprMutator { PrimExpr ret = StmtExprMutator::VisitExpr_(op).as_or_throw(); op = ret.as(); - Buffer new_buf = GetRemappedBuffer(op->buffer); + BufferVar new_buf = GetRemappedBuffer(op->buffer); if (new_buf.same_as(op->buffer)) { return ret; } else { @@ -505,7 +506,7 @@ class ComputeLegalizer : public StmtExprMutator { return DTypeConversion(value, dtype); } - Buffer GetRemappedBuffer(Buffer buf) { + BufferVar GetRemappedBuffer(BufferVar buf) { auto buf_it = buffer_remap_.find(buf); if (buf_it != buffer_remap_.end()) { return buf_it->second; @@ -515,7 +516,7 @@ class ComputeLegalizer : public StmtExprMutator { protected: PrimType promote_dtype_; - std::unordered_map buffer_remap_; + std::unordered_map buffer_remap_; std::unordered_map var_remap_; }; @@ -572,21 +573,22 @@ class StorageLegalizer : public StmtExprMutator { } Stmt VisitStmt_(const AllocBufferNode* op) final { - Buffer buf = GetRemappedBuffer(op->buffer); + BufferVar buf = GetRemappedBuffer(op->buffer); // in a rare case the buffer didn't get remapped // because the original var is not bfloat* // force remap here if (MatchType(buf->dtype)) { PrimType new_dtype = GetStorageUIntDType(buf->dtype); ffi::String storage_scope = "global"; - if (auto* ptr_type = buf->data->ty.as()) { + if (auto* ptr_type = buf->data_pointer_type.as()) { storage_scope = ptr_type->storage_scope; } - Var new_data = Var(buf->data->name, PointerType(new_dtype, storage_scope)); - var_remap_[buf->data] = new_data; - buf = Buffer(new_data, new_dtype, buf->shape, buf->strides, buf->elem_offset, buf->name, - buf->data_alignment, buf->offset_factor, buf->span, buf->layout, - buf->allocated_addr); + auto type = CopyBufferType(buf); + type->data_pointer_type = PointerType(new_dtype, storage_scope); + type->dtype = new_dtype; + BufferVar new_buf = RebuildBufferVar(buf, std::move(type)); + var_remap_[buf.var()] = new_buf.var(); + buf = std::move(new_buf); buffer_remap_[op->buffer] = buf; } if (buf.same_as(op->buffer)) { @@ -599,20 +601,22 @@ class StorageLegalizer : public StmtExprMutator { } Stmt VisitStmt_(const DeclBufferNode* op) final { - Buffer buf = GetRemappedBuffer(op->buffer); + BufferVar buf = GetRemappedBuffer(op->buffer); // in a rare case the buffer didn't get remapped // because the original var is not bfloat* // force remap here if (MatchType(buf->dtype)) { - buf = Buffer(buf->data, GetStorageUIntDType(buf->dtype), buf->shape, buf->strides, - buf->elem_offset, buf->name, buf->data_alignment, buf->offset_factor, buf->span, - buf->layout, buf->allocated_addr); + PrimType new_dtype = GetStorageUIntDType(buf->dtype); + auto type = CopyBufferType(buf); + type->data_pointer_type = PointerType(new_dtype, buf.scope()); + type->dtype = new_dtype; + buf = RebuildBufferVar(buf, std::move(type)); buffer_remap_[op->buffer] = buf; } if (buf.same_as(op->buffer)) { return ffi::GetRef(op); } else { - return DeclBuffer(buf, op->span); + return DeclBuffer(buf, op->data, op->span); } } @@ -641,7 +645,7 @@ class StorageLegalizer : public StmtExprMutator { Stmt VisitStmt_(const BufferStoreNode* op) final { PrimExpr value = this->ChangeToUInt(VisitPrimExpr(op->value)); - Buffer new_buf = GetRemappedBuffer(op->buffer); + BufferVar new_buf = GetRemappedBuffer(op->buffer); auto indices = op->indices.Map([this](PrimExpr expr) { return this->VisitPrimExpr(expr); }); ffi::Optional predicate = std::nullopt; if (op->predicate.has_value()) { @@ -662,7 +666,7 @@ class StorageLegalizer : public StmtExprMutator { Stmt ret = StmtExprMutator::VisitStmt_(op); op = ret.as(); - if (auto buffer = op->node.as()) { + if (auto buffer = op->node.as()) { auto it = buffer_remap_.find(buffer.value()); if (it != buffer_remap_.end()) { return AttrStmt(it->second, op->attr_key, op->value, op->body); @@ -679,7 +683,7 @@ class StorageLegalizer : public StmtExprMutator { Expr VisitExpr_(const BufferLoadNode* op) final { PrimExpr ret = StmtExprMutator::VisitExpr_(op).as_or_throw(); op = ret.as(); - Buffer new_buf = GetRemappedBuffer(op->buffer); + BufferVar new_buf = GetRemappedBuffer(op->buffer); if (new_buf.same_as(op->buffer)) { return ret; } else { @@ -755,18 +759,15 @@ class StorageLegalizer : public StmtExprMutator { return var; } - Buffer GetRemappedBuffer(Buffer buf) { + BufferVar GetRemappedBuffer(BufferVar buf) { auto buf_it = buffer_remap_.find(buf); if (buf_it != buffer_remap_.end()) { return buf_it->second; } - Buffer new_buf = buf; - auto var_it = var_remap_.find(buf->data); + BufferVar new_buf = buf; + auto var_it = var_remap_.find(buf.var()); if (var_it != var_remap_.end()) { - PrimType dtype = MatchType(buf->dtype) ? GetStorageUIntDType(buf->dtype) : buf->dtype; - new_buf = Buffer(var_it->second, dtype, buf->shape, buf->strides, buf->elem_offset, buf->name, - buf->data_alignment, buf->offset_factor, buf->span, buf->layout, - buf->allocated_addr); + new_buf = BufferVar(var_it->second); } else { TVM_FFI_ICHECK(!MatchType(buf->dtype)) << "Cannot find var remap for " << buf; } @@ -776,7 +777,7 @@ class StorageLegalizer : public StmtExprMutator { return new_buf; } - std::unordered_map buffer_remap_; + std::unordered_map buffer_remap_; std::unordered_map var_remap_; }; diff --git a/src/tirx/transform/update_pointer_storage_scope.cc b/src/tirx/transform/update_pointer_storage_scope.cc index 0c796ff3646c..b531c208510a 100644 --- a/src/tirx/transform/update_pointer_storage_scope.cc +++ b/src/tirx/transform/update_pointer_storage_scope.cc @@ -48,7 +48,18 @@ Var WithStorageScope(const VarNode* buffer_var, ffi::String storage_scope) { UpdatePointerStorageScope::UpdatePointerStorageScope( const std::unordered_map& new_storage_scopes) { for (auto& kv : new_storage_scopes) { - new_var_remap_[kv.first] = WithStorageScope(kv.first, kv.second); + if (kv.first->ty.as()) { + BufferVar buffer = GetBufferVar(kv.first); + auto type = CopyBufferType(buffer); + const auto* pointer_type = type->data_pointer_type.as(); + TVM_FFI_ICHECK(pointer_type); + type->data_pointer_type = PointerType(pointer_type->element_type, kv.second); + BufferVar replacement = RebuildBufferVar(buffer, std::move(type)); + new_buffer_remap_[kv.first] = replacement; + new_var_remap_[kv.first] = replacement.var(); + } else { + new_var_remap_[kv.first] = WithStorageScope(kv.first, kv.second); + } } } @@ -70,7 +81,7 @@ Node UpdatePointerStorageScope::UpdateBufferAccess(Node node) { return node; } -Buffer UpdatePointerStorageScope::GetUpdatedBuffer(Buffer buf) { +BufferVar UpdatePointerStorageScope::GetUpdatedBuffer(BufferVar buf) { // Use the cached buffer, if it exists. auto key = buf.get(); auto it = new_buffer_remap_.find(key); @@ -78,13 +89,6 @@ Buffer UpdatePointerStorageScope::GetUpdatedBuffer(Buffer buf) { return it->second; } - // Update the buffer's var, if needed. - auto remapped = StmtExprMutator::VisitExpr(buf->data).as_or_throw(); - if (!remapped.same_as(buf->data)) { - auto writer = buf.CopyOnWrite(); - writer->data = remapped; - } - // Update the cache and return new_buffer_remap_[key] = buf; return buf; diff --git a/src/tirx/transform/update_pointer_storage_scope.h b/src/tirx/transform/update_pointer_storage_scope.h index 05002e675943..f913a35be71a 100644 --- a/src/tirx/transform/update_pointer_storage_scope.h +++ b/src/tirx/transform/update_pointer_storage_scope.h @@ -48,10 +48,10 @@ class UpdatePointerStorageScope : public StmtExprMutator { template Node UpdateBufferAccess(Node node); - Buffer GetUpdatedBuffer(Buffer buf); + BufferVar GetUpdatedBuffer(BufferVar buf); std::unordered_map new_var_remap_; - std::unordered_map new_buffer_remap_; + std::unordered_map new_buffer_remap_; }; } // namespace tirx diff --git a/src/tirx/transform/vectorize_loop.cc b/src/tirx/transform/vectorize_loop.cc index 9872a8c46217..81c0318a96f0 100644 --- a/src/tirx/transform/vectorize_loop.cc +++ b/src/tirx/transform/vectorize_loop.cc @@ -245,7 +245,7 @@ class TryPredicateBufferAccesses : public StmtExprMutator { num_accesses_rewritten_ += 1; auto writer = node.CopyOnWrite(); if (node->predicate.has_value() && allow_offset_predication_) { - // Buffer predicates are uint1 lane masks, so mask merging uses bitwise + // BufferVar predicates are uint1 lane masks, so mask merging uses bitwise // and rather than logical &&. writer->predicate = node->predicate.value() & lane_mask; } else { @@ -298,12 +298,12 @@ class VecAllocAccess : public StmtExprMutator { template Node UpdateBufferAccess(Node node) { // Only update the buffer that's being replaced. - if (node->buffer->data.get() != buf_) { + if (node->buffer.get() != buf_) { return node; } - // Find/make a Buffer object with the correct updated shape. - Buffer buf; + // Find/make a BufferVar object with the correct updated shape. + BufferVar buf; auto it = buffer_map_.find(node->buffer.get()); if (it != buffer_map_.end()) { buf = it->second; @@ -321,21 +321,21 @@ class VecAllocAccess : public StmtExprMutator { // are updated for consistency. // Update strides if defined. - ffi::Array strides; + ffi::Array strides = node->buffer->strides; for (size_t i = 0; i < strides.size(); i++) { PrimExpr stride = strides[i]; if (i != strides.size() - 1) { stride *= var_lanes_; } - strides.push_back(analyzer_->Simplify(stride)); + strides.Set(i, analyzer_->Simplify(stride)); } // Copy everything into the new buffer. - buf = node->buffer; - auto buf_writer = buf.CopyOnWrite(); - buf_writer->shape = shape; - buf_writer->strides = strides; - buffer_map_[buf.get()] = buf; + auto type = CopyBufferType(node->buffer); + type->shape = shape; + type->strides = strides; + buf = RebuildBufferVar(node->buffer, std::move(type)); + buffer_map_[node->buffer.get()] = buf; } // Extend the last index by the number of lanes in the vectorized @@ -353,7 +353,7 @@ class VecAllocAccess : public StmtExprMutator { // buffer var const VarNode* buf_; // Updated buffer objects. - std::unordered_map buffer_map_; + std::unordered_map buffer_map_; // variable to be replaced Var var_; // the lanes. @@ -864,7 +864,7 @@ class Vectorizer : public StmtMutator, public ExprFunctor { int total_lanes = std::max(index_lanes, value_dtype_lanes); TVM_FFI_ICHECK_EQ(total_lanes % other_index_lanes, 0) - << "When storing to buffer " << op->buffer->name << ", cannot produce " << total_lanes + << "When storing to buffer " << op->buffer.name() << ", cannot produce " << total_lanes << " lanes of storage location by changing the last index."; int last_index_lanes = total_lanes / other_index_lanes; diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc index 190a81d66574..45c0c257751e 100644 --- a/tests/cpp/ir_functor_test.cc +++ b/tests/cpp/ir_functor_test.cc @@ -154,7 +154,7 @@ TEST(IRF, StmtVisitor) { Stmt eval_body = Evaluate(z); PrimType dtype = PrimType::Float(32); tirx::Var data_var("b", PointerType(dtype)); - Buffer buf(data_var, dtype, {z, z}, {}, PrimExpr(), "b", 0, 0); + BufferVar buf(data_var, dtype, {z, z}, {}, PrimExpr(), "b", 0, 0); // AllocBuffer is flat (no body). Return as SeqStmt with eval. return SeqStmt({AllocBuffer(buf), eval_body}); }; @@ -168,7 +168,7 @@ TEST(IRF, StmtVisitor) { Stmt body = fmaketest(); PrimType dtype = PrimType::Float(32); tirx::Var buf_var("b", PointerType(dtype)); - Buffer buffer = decl_buffer({16}); + BufferVar buffer = decl_buffer({16}); body = SeqStmt({DeclBuffer(buffer), std::move(body)}); BufferRegion buffer_region(buffer, {Range::FromMinExtent(x + 1, 1)}); MatchBufferRegion match_buffer_region(decl_buffer({1}), buffer_region); @@ -208,7 +208,7 @@ TEST(IRF, StmtMutator) { auto z = x + 1; PrimType dtype = PrimType::Float(32); tirx::Var data_var("b", PointerType(dtype)); - Buffer buf(data_var, dtype, {1, z}, {}, PrimExpr(), "b", 0, 0); + BufferVar buf(data_var, dtype, {1, z}, {}, PrimExpr(), "b", 0, 0); return AllocBuffer(buf); }; @@ -299,7 +299,7 @@ TEST(IRF, StmtMutator) { // tests for block and block_realize // AllocBuffer and DeclBuffer are flat (no body), placed as siblings in SeqStmt Stmt eval_body = Evaluate(x + 1); - Buffer buffer = decl_buffer({16}); + BufferVar buffer = decl_buffer({16}); Stmt decl = DeclBuffer(buffer); Stmt alloc = fmakealloc(); // body is: DeclBuffer, AllocBuffer, Evaluate @@ -335,7 +335,7 @@ TEST(IRF, Substitute) { PrimVar n("n", PrimType::Int(32)); auto fmakebuffer = [&]() { - return Buffer{/*data=*/x, + return BufferVar{/*data=*/x, /*dtype=*/PrimType::Float(32), /*shape=*/{n}, /*strides=*/{}, @@ -346,12 +346,14 @@ TEST(IRF, Substitute) { }; { - // test substitute buffer data var and shape var via DeclBuffer + // Test substitution of an explicit DeclBuffer source and a dependent + // BufferType shape. Changing the type creates one fresh Var identity + // that is shared by the declaration and every use. tirx::Var y = x.CopyWithSuffix("subst"); PrimVar m("m", PrimType::Int(32)); - Buffer buffer = fmakebuffer(); + BufferVar buffer = fmakebuffer(); Stmt store = BufferStore(buffer, FloatImm(dtype, 0), {IntImm::Int32(0)}); - Stmt decl = SeqStmt({DeclBuffer(buffer), store}); + Stmt decl = SeqStmt({DeclBuffer(buffer, x), store}); auto f_subst = [&](const tirx::Var& var) -> ffi::Optional { if (var.same_as(x)) return Expr(y); if (var.same_as(n)) return Expr(m); @@ -362,13 +364,17 @@ TEST(IRF, Substitute) { TVM_FFI_ICHECK(seq_node != nullptr); auto* decl_node = seq_node->seq[0].as(); TVM_FFI_ICHECK(decl_node != nullptr); - TVM_FFI_ICHECK(decl_node->buffer->data.same_as(y)); + TVM_FFI_ICHECK(decl_node->data.value().same_as(y)); TVM_FFI_ICHECK(decl_node->buffer->shape[0].same_as(m)); + TVM_FFI_ICHECK(!decl_node->buffer.same_as(buffer)); + auto* store_node = seq_node->seq[1].as(); + TVM_FFI_ICHECK(store_node != nullptr); + TVM_FFI_ICHECK(store_node->buffer.same_as(decl_node->buffer)); } { // test identity substitution on expression - Buffer buffer = fmakebuffer(); + BufferVar buffer = fmakebuffer(); PrimExpr expr = BufferLoad(buffer, {IntImm::Int32(0)}); auto f_subst = [&](const tirx::Var& var) -> ffi::Optional { return Expr(var); }; PrimExpr new_expr = Substitute(expr, f_subst); diff --git a/tests/python/s_tir/analysis/test_sblock_access_region.py b/tests/python/s_tir/analysis/test_sblock_access_region.py index 1ecab58ab084..16c420263e73 100644 --- a/tests/python/s_tir/analysis/test_sblock_access_region.py +++ b/tests/python/s_tir/analysis/test_sblock_access_region.py @@ -209,7 +209,7 @@ def access_of_padding_pattern() -> None: def test_block_access_region_detector(): block = func.body.block.body.block alloc_buffers = func.body.block.alloc_buffers - buffer_var_map = {buf.data: buf for buf in alloc_buffers} + buffer_var_map = {buf: buf for buf in alloc_buffers} ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(block.reads, ret[0]) @@ -222,7 +222,7 @@ def test_block_access_region_detector(): def test_opaque_block(): alloc_buffers = opaque_block_func.body.block.alloc_buffers - buffer_var_map = {buf.data: buf for buf in alloc_buffers} + buffer_var_map = {buf: buf for buf in alloc_buffers} block0 = opaque_block_func.body.block.body.body.block ret = s_tir.analysis.get_sblock_access_region(block0, buffer_var_map) @@ -238,7 +238,7 @@ def test_opaque_block(): def test_opaque_access(): block = opaque_access_func.body.block.body.body.block alloc_buffers = opaque_access_func.body.block.alloc_buffers - buffer_var_map = {buf.data: buf for buf in alloc_buffers} + buffer_var_map = {buf: buf for buf in alloc_buffers} ret0 = s_tir.analysis.get_sblock_read_write_region(block, buffer_var_map) ret1 = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) @@ -251,7 +251,7 @@ def test_opaque_access(): def test_opaque_access_with_tvm_access_ptr(): block = opaque_access_with_tvm_access_ptr_func.body.block.body.block alloc_buffers = opaque_access_with_tvm_access_ptr_func.body.block.alloc_buffers - buffer_var_map = {buf.data: buf for buf in alloc_buffers} + buffer_var_map = {buf: buf for buf in alloc_buffers} ret0 = s_tir.analysis.get_sblock_read_write_region(block, buffer_var_map) ret1 = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) @@ -268,7 +268,7 @@ def test_match_buffer(): block = root_block.body.body.body.block block_inner = block.body[0].body.body.block alloc_buffers = match_buffer_func.body.block.alloc_buffers - buffer_var_map = {buf.data: buf for buf in alloc_buffers} + buffer_var_map = {buf: buf for buf in alloc_buffers} # Check block ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) @@ -284,7 +284,7 @@ def test_match_buffer(): # Check inner block AAA for match_buffer in block.match_buffers: target_buffer = match_buffer.buffer - buffer_var_map[target_buffer.data] = target_buffer + buffer_var_map[target_buffer] = target_buffer ret = s_tir.analysis.get_sblock_access_region(block_inner, buffer_var_map) tvm.ir.assert_structural_equal(block_inner.reads, ret[0]) @@ -294,7 +294,7 @@ def test_match_buffer(): def test_access_in_if_then_else_func(): block = access_in_if_then_else_func.body.block.body.block alloc_buffers = access_in_if_then_else_func.body.block.alloc_buffers - buffer_var_map = {buf.data: buf for buf in alloc_buffers} + buffer_var_map = {buf: buf for buf in alloc_buffers} ret0 = s_tir.analysis.get_sblock_read_write_region(block, buffer_var_map) ret1 = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(ret0[0], ret1[0]) @@ -304,7 +304,7 @@ def test_access_in_if_then_else_func(): def test_access_in_branch_func(): block = access_in_branch_func.body.block.body.block alloc_buffers = access_in_branch_func.body.block.alloc_buffers - buffer_var_map = {buf.data: buf for buf in alloc_buffers} + buffer_var_map = {buf: buf for buf in alloc_buffers} ret0 = s_tir.analysis.get_sblock_read_write_region(block, buffer_var_map) ret1 = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(ret0[0], ret1[0]) @@ -314,7 +314,7 @@ def test_access_in_branch_func(): def test_access_of_padding_pattern(): s = tvm.s_tir.schedule.Schedule(access_of_padding_pattern) alloc_buffers = s.get_sref(s.get_sblock("root")).stmt.alloc_buffers - buffer_var_map = {buf.data: buf for buf in alloc_buffers} + buffer_var_map = {buf: buf for buf in alloc_buffers} def do_compare_buffer_region(region, expect): assert region.buffer == expect.buffer @@ -340,7 +340,7 @@ def do_check_block(block_name): def test_access_of_reduction(): block = gemm.body.block.body.body.body.body.body.body.block alloc_buffers = gemm.body.block.alloc_buffers - buffer_var_map = {buf.data: buf for buf in alloc_buffers} + buffer_var_map = {buf: buf for buf in alloc_buffers} ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(block.reads, ret[0]) tvm.ir.assert_structural_equal(block.writes, ret[1]) @@ -350,7 +350,7 @@ def test_access_of_decompose_reduction(): init = decomposed_gemm.body.block.body.body.body[0].body.body.block update = decomposed_gemm.body.block.body.body.body[1].body.body.body.block alloc_buffers = decomposed_gemm.body.block.alloc_buffers - buffer_var_map = {buf.data: buf for buf in alloc_buffers} + buffer_var_map = {buf: buf for buf in alloc_buffers} for block in [init, update]: ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(block.reads, ret[0]) @@ -379,7 +379,7 @@ def func( output[vi, vs] = storage[seq_id, history_id, vs] block = func.body.block.body.body.body.block - buffer_var_map = {buf.data: buf for buf in func.buffer_map.values()} + buffer_var_map = {buf: buf for buf in func.buffer_map.values()} ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(block.reads, ret[0]) tvm.ir.assert_structural_equal(block.writes, ret[1]) @@ -405,7 +405,7 @@ def func( C[vi, vs1] = A[vi1, vs2] + B[vi2, vs3] block = func.body.block.body.body.body.block - buffer_var_map = {buf.data: buf for buf in func.buffer_map.values()} + buffer_var_map = {buf: buf for buf in func.buffer_map.values()} ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(block.reads, ret[0]) tvm.ir.assert_structural_equal(block.writes, ret[1]) diff --git a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py index ac707a35838d..29a7d9cace20 100644 --- a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py +++ b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py @@ -464,9 +464,8 @@ def test_error_undeclared_buffer_in_schedulable_tir(): ) prim_func = tvm.tirx.PrimFunc( - params=[A.data, B_data], + params=[A, B_data], body=tvm.tirx.For(i, 0, n, tvm.tirx.ForKind.SERIAL, block_realize), - buffer_map={A.data: A}, # Note: B is NOT in buffer_map, so its declaration scope is only # within a DeclBuffer node (which we intentionally omit here). ) diff --git a/tests/python/tirx-base/test_tir_buffer.py b/tests/python/tirx-base/test_tir_buffer.py index 42e24ce127a5..78f7ed3d17ae 100644 --- a/tests/python/tirx-base/test_tir_buffer.py +++ b/tests/python/tirx-base/test_tir_buffer.py @@ -32,9 +32,30 @@ def test_buffer(): Ab = tvm.tirx.decl_buffer((m, n), "float32") Bb = tvm.tirx.decl_buffer((n, l), "float32") - assert isinstance(Ab, tvm.tirx.Buffer) + assert type(Ab) is tvm.ir.Var + assert tvm.tirx.is_buffer(Ab) + assert isinstance(Ab.ty, tvm.tirx.BufferType) assert Ab.dtype == tvm.ir.PrimType("float32") assert tuple(Ab.shape) == (m, n) + assert not tvm.tirx.is_buffer(m) + + +def test_buffer_data_is_typed_projection(): + buffer = tvm.tirx.decl_buffer((8,), "bool", scope="shared") + + assert buffer.dtype == tvm.ir.PrimType("bool") + assert buffer.data_pointer_type == tvm.ir.PointerType(tvm.ir.PrimType("int8"), "shared") + assert buffer.data.op.name == "tirx.buffer_data" + assert buffer.data.args[0].same_as(buffer) + assert buffer.data.ty == buffer.data_pointer_type + + +def test_buffer_logical_dtype_independent_of_pointer_type(): + data = tvm.ir.Var("storage", tvm.ir.PointerType(tvm.ir.PrimType("uint8"), "local")) + buffer = tvm.tirx.decl_buffer((8,), "float16", data=data) + + assert buffer.dtype == tvm.ir.PrimType("float16") + assert buffer.data.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint8"), "local") def test_buffer_access_ptr(): @@ -191,7 +212,12 @@ def test_buffer_flatten(): """A buffer should flatten to a 1-d shape""" buf = tvm.tirx.decl_buffer([16, 32]) flat = buf.get_flattened_buffer() - assert buf.data.same_as(flat.data) + # A metadata-changing rewrite creates a fresh typed Var. The physical + # pointer is always derived from that Var instead of being stored as a + # second buffer identity. + assert not buf.same_as(flat) + assert flat.data.args[0].same_as(flat) + assert flat.data.op.name == "tirx.buffer_data" tvm.ir.assert_structural_equal(flat.shape, [T.int32(16 * 32)]) diff --git a/tests/python/tirx-base/test_tir_op_types.py b/tests/python/tirx-base/test_tir_op_types.py index 78e379b76524..388cb0e74a1c 100644 --- a/tests/python/tirx-base/test_tir_op_types.py +++ b/tests/python/tirx-base/test_tir_op_types.py @@ -264,7 +264,8 @@ def test_op_ptx_cp_async(): expr = _cuda_op.ptx_cp_async_legacy("float16", inner_dst, 3, inner_src, 5, 16) for access_ptr, expected_offset in zip(expr.args[:2], [5, 9]): assert access_ptr.op.name == "tirx.tvm_access_ptr" - assert isinstance(access_ptr.args[1], tirx.Var) + assert access_ptr.args[1].op.name == "tirx.buffer_data" + assert isinstance(access_ptr.args[1].args[0], tirx.Var) simplified_offset = tvm.arith.Analyzer().simplify(access_ptr.args[2]) assert int(simplified_offset) == expected_offset diff --git a/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py b/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py index ea2f9dbb3f18..8670c4424836 100644 --- a/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py +++ b/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py @@ -97,7 +97,9 @@ def main(input_A: T.Buffer((16, 16), "float32"), input_C: T.Buffer((16, 16), "fl C[((i * 16) + j)] = B_new[j] * 2.0 After = _transform()(Before) - tvm.ir.assert_structural_equal(After, Expected) + # This deliberately malformed input omits the DeclBuffer that would bind + # B_new. Compare the resulting free buffer variables by mapping them. + tvm.ir.assert_structural_equal(After, Expected, map_free_vars=True) def test_gpu(): diff --git a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py index 2e519c229e8a..3f3cc59dad09 100644 --- a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py +++ b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py @@ -399,9 +399,8 @@ def main() -> None: A_data: T.let[T.handle("int32x8")] = T.call_extern( "dummy_func", dtype=T.handle("int32x8").ty ) - A = T.decl_buffer([8], "int32", data=A_data) - A_1 = T.Buffer([1], "int32x8", data=A_data) - A_1[0] = T.broadcast(42, 8) + A = T.decl_buffer([1], "int32x8", data=A_data) + A[0] = T.broadcast(42, 8) After = tvm.tirx.transform.StorageRewrite()(Before) tvm.ir.assert_structural_equal(After, Expected) diff --git a/tests/python/tirx/test_parser_printer.py b/tests/python/tirx/test_parser_printer.py index e570f72200c2..02fc42b5cad5 100644 --- a/tests/python/tirx/test_parser_printer.py +++ b/tests/python/tirx/test_parser_printer.py @@ -1341,6 +1341,18 @@ def _visit(node): return bufs +def _collect_buffer_sources(func): + """Collect the explicit data source of each DeclBuffer.""" + sources = {} + + def _visit(node): + if isinstance(node, tvm.tirx.DeclBuffer): + sources[node.buffer.name] = node.data + + tvm.tirx.stmt_functor.post_order_visit(func.body, _visit) + return sources + + def test_buffer_local_ir(): """Verify .local() auto-infer: shape from storage shard extents, layout, shared data.""" @@ -1360,7 +1372,7 @@ def func() -> None: b_buf = bufs["B"] # Shared data pointer - assert b_local.data.same_as(b_buf.data) + assert_structural_equal(_collect_buffer_sources(func)["B_local"], b_buf.data) # Shape: single dim matching storage shard total assert len(b_local.shape) == 1 storage = b_buf.layout.storage() @@ -1394,7 +1406,7 @@ def func() -> None: b_buf = bufs["B"] # Shared data pointer - assert b_buf.data.same_as(a_buf.data) + assert_structural_equal(_collect_buffer_sources(func)["B"], a_buf.data) # Shape: [4, 8] from [8, 4] assert int(b_buf.shape[0]) == 4 assert int(b_buf.shape[1]) == 8 @@ -1422,7 +1434,7 @@ def func() -> None: b_buf = bufs["B"] # Shared data pointer - assert b_buf.data.same_as(a_buf.data) + assert_structural_equal(_collect_buffer_sources(func)["B"], a_buf.data) # dtype assert str(b_buf.dtype) == "float32" # Shape: [8, 4] (last dim halved since float32 is 2x float16) From dfa201210973af1ca20fdb651c6f5aab6e47c139 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Tue, 28 Jul 2026 10:48:44 +0000 Subject: [PATCH 2/2] [FIX][TIRX] Preserve behavior after typed-buffer migration Keep physical-buffer ownership explicit after Buffer becomes a typed Var. Storage planning, vector type rewriting, cleanup, and all-reduce lowering now maintain pass-local alias-to-root maps only where liveness or allocation remapping requires them; simpler passes decode buffer_data projections locally. Replace superseded DeclBuffer bindings instead of retaining stale declarations, reject aliases whose physical source has not been registered, and repair merged tagged-memory children after their parent allocation is rebuilt. Restore the alias-lifetime and original-output regressions that exercise these invariants. Preserve the Python compatibility surface deliberately: Buffer remains an annotation/import alias, is_buffer_var is the runtime discriminator, legacy metadata properties reject non-buffer Vars with AttributeError, and dtype continues to return DataType. Validated with a clean C++ build, 239 focused passing tests (3 skipped), and the complete TIRx/S-TIR Python suite; remaining broad-suite failures require unavailable LLVM, CUDA, NKI, Torch, or target-feature support. --- include/tvm/s_tir/schedule/schedule.h | 4 +- include/tvm/s_tir/transform.h | 2 +- include/tvm/tirx/buffer.h | 81 ++--- include/tvm/tirx/builtin.h | 5 +- include/tvm/tirx/function.h | 4 +- include/tvm/tirx/index_map.h | 4 +- include/tvm/tirx/script/builder/frame.h | 6 +- include/tvm/tirx/script/builder/ir.h | 16 +- include/tvm/tirx/stmt.h | 10 +- .../permute_layout/warp_xor_swizzle.py | 6 +- python/tvm/backend/cuda/script.py | 12 +- python/tvm/backend/metal/script.py | 8 +- .../operator/tile_primitive/binary/default.py | 2 +- .../operator/tile_primitive/binary/utils.py | 6 +- .../tile_primitive/compose_op/binary_chain.py | 2 +- .../compose_op/binary_reduce.py | 4 +- .../tile_primitive/compose_op/unary_reduce.py | 8 +- .../operator/tile_primitive/copy/default.py | 25 +- .../operator/tile_primitive/gemm/default.py | 28 +- .../tile_primitive/instruction_generator.py | 10 +- .../operator/tile_primitive/private_alloc.py | 14 +- .../tile_primitive/reduction/utils.py | 16 +- .../operator/tile_primitive/select/default.py | 8 +- .../operator/tile_primitive/unary/utils.py | 10 +- .../tile_primitive/workspace_utils.py | 10 +- .../backend/trn/transform/naive_allocator.py | 23 +- python/tvm/s_tir/schedule/schedule.py | 18 +- python/tvm/tirx/__init__.py | 11 +- python/tvm/tirx/buffer.py | 233 +++++++++------ python/tvm/tirx/function.py | 4 +- python/tvm/tirx/op.py | 36 +-- python/tvm/tirx/script/builder/ir.py | 22 +- python/tvm/tirx/script/builder/tirx.py | 80 +++-- python/tvm/tirx/script/builder/triton.py | 3 +- python/tvm/tirx/script/parser/parser.py | 56 ++-- python/tvm/tirx/script/tile.py | 6 +- python/tvm/tirx/stmt.py | 7 +- python/tvm/tirx/stmt_functor.py | 2 +- python/tvm/tirx/transform/common.py | 39 +-- python/tvm/tirx/transform/transform.py | 8 +- src/backend/cuda/codegen/codegen_cuda.cc | 8 +- .../hexagon/codegen/llvm/codegen_hexagon.cc | 3 +- src/backend/hexagon/runtime/hexagon_buffer.h | 2 +- src/backend/trn/codegen/codegen_trn.cc | 4 +- .../trn/transform/lower_trainium_layout.cc | 20 +- src/relax/analysis/layout_transformation.cc | 8 +- src/relax/analysis/tir_op_pattern_kind.cc | 3 +- .../backend/adreno/annotate_custom_storage.cc | 2 +- src/relax/backend/vm/vm_shape_lower.cc | 2 +- .../lower_global_view_to_local_view.cc | 9 +- src/relax/op/op.cc | 4 +- src/relax/op/tensor/inspect.cc | 2 +- src/relax/transform/fuse_tir.cc | 22 +- .../specialize_primfunc_based_on_callsite.cc | 8 +- .../transform/split_call_tir_by_pattern.cc | 4 +- .../analysis/calculate_allocated_memory.cc | 2 +- .../analysis/sblock_access_region_detector.cc | 10 +- .../sblock_buffer_access_lca_detector.cc | 2 +- .../backend/adreno/inject_texture_alloc.cc | 5 +- src/s_tir/backend/adreno/texture_flatten.cc | 6 +- .../feature_extractor/per_store_feature.cc | 58 ++-- .../multi_level_tiling_tensor_core.cc | 2 +- src/s_tir/schedule/analysis.h | 5 +- src/s_tir/schedule/analysis/analysis.cc | 4 +- src/s_tir/schedule/analysis/layout.cc | 15 +- src/s_tir/schedule/concrete_schedule.cc | 2 +- src/s_tir/schedule/concrete_schedule.h | 2 +- src/s_tir/schedule/ir_comparator.cc | 27 +- src/s_tir/schedule/primitive.h | 4 +- .../primitive/annotate_buffer_access.cc | 4 +- .../schedule/primitive/block_annotate.cc | 2 +- src/s_tir/schedule/primitive/cache_index.cc | 7 +- .../schedule/primitive/cache_read_write.cc | 14 +- src/s_tir/schedule/primitive/compute_at.cc | 4 +- .../schedule/primitive/compute_inline.cc | 23 +- .../primitive/layout_transformation.cc | 16 +- src/s_tir/schedule/primitive/read_write_at.cc | 8 +- src/s_tir/schedule/primitive/reduction.cc | 4 +- src/s_tir/schedule/schedule.cc | 2 +- src/s_tir/schedule/state.cc | 12 +- src/s_tir/schedule/trace.cc | 2 +- src/s_tir/schedule/traced_schedule.cc | 2 +- src/s_tir/schedule/traced_schedule.h | 2 +- src/s_tir/schedule/transform.cc | 16 +- src/s_tir/schedule/transform.h | 2 +- src/s_tir/transform/compact_buffer_region.cc | 46 ++- src/s_tir/transform/inject_double_buffer.cc | 24 +- src/s_tir/transform/inject_permuted_layout.cc | 30 +- src/s_tir/transform/inject_ptx_async_copy.cc | 13 +- .../transform/inject_software_pipeline.cc | 38 ++- src/s_tir/transform/inject_virtual_thread.cc | 53 +++- src/s_tir/transform/lower_async_dma.cc | 20 +- .../transform/lower_cross_thread_reduction.cc | 28 +- src/s_tir/transform/lower_match_buffer.cc | 13 + src/s_tir/transform/lower_thread_allreduce.cc | 94 ++++-- src/s_tir/transform/lower_vtcm_alloc.cc | 5 +- .../manifest_shared_memory_local_stage.cc | 5 +- src/s_tir/transform/memhammer_coalesce.cc | 2 +- .../transform/memhammer_lower_auto_copy.cc | 7 +- src/s_tir/transform/memhammer_rewrite_rule.h | 3 +- .../transform/memhammer_tensorcore_rewrite.cc | 76 ++--- .../merge_shared_memory_allocations.cc | 135 +++++++-- src/s_tir/transform/renew_defs.cc | 2 +- src/s_tir/transform/storage_access.cc | 30 +- .../transform/transform_mma_buffer_layout.cc | 9 +- .../using_assume_to_reduce_branches.cc | 3 +- src/target/llvm/codegen_cpu.cc | 15 +- src/target/llvm/codegen_llvm.cc | 39 ++- src/target/source/codegen_c.cc | 17 +- src/target/source/codegen_source_base.cc | 7 +- src/target/source/codegen_source_base.h | 7 + src/te/operation/create_primfunc.cc | 3 +- src/tirx/analysis/var_use_def_analysis.cc | 28 +- src/tirx/analysis/var_use_def_analysis.h | 2 - src/tirx/ir/async_structs.cc | 6 +- src/tirx/ir/buffer.cc | 187 ++++-------- src/tirx/ir/data_type_rewriter.cc | 8 +- src/tirx/ir/expr_functor.cc | 12 +- src/tirx/ir/specialize.cc | 4 + src/tirx/ir/stmt.cc | 9 +- src/tirx/ir/stmt_functor.cc | 29 +- src/tirx/ir/tir_visitor_with_path.h | 3 +- src/tirx/op/op.cc | 8 +- src/tirx/script/builder/frame.cc | 5 +- src/tirx/script/builder/ir.cc | 58 ++-- src/tirx/script/printer/buffer.cc | 2 +- src/tirx/script/printer/expr.cc | 5 +- src/tirx/script/printer/stmt.cc | 42 ++- src/tirx/transform/flatten_buffer.cc | 25 +- src/tirx/transform/ir_utils.cc | 11 +- src/tirx/transform/ir_utils.h | 10 +- src/tirx/transform/lower_intrin.cc | 20 +- src/tirx/transform/lower_tirx_cleanup.cc | 114 +++++-- src/tirx/transform/lower_tvm_builtin.cc | 26 +- src/tirx/transform/lower_warp_memory.cc | 26 +- src/tirx/transform/split_host_device.cc | 8 +- src/tirx/transform/stmt_simplify.cc | 3 +- src/tirx/transform/storage_rewrite.cc | 277 +++++++++++------- src/tirx/transform/tvm_ffi_binder.cc | 14 +- src/tirx/transform/tvm_ffi_binder.h | 4 +- .../transform/unsupported_dtype_legalize.cc | 66 ++--- .../transform/update_pointer_storage_scope.cc | 16 +- .../transform/update_pointer_storage_scope.h | 1 - tests/cpp/ir_functor_test.cc | 34 +-- .../codegen/test_target_codegen_llvm.py | 16 + ...test_analysis_suggest_layout_transforms.py | 2 +- .../test_tir_schedule_transform_layout.py | 15 +- ...t_s_tir_transform_inject_ptx_async_copy.py | 42 +-- ...t_s_tir_transform_inject_virtual_thread.py | 19 +- ...test_s_tir_transform_lower_match_buffer.py | 20 ++ ...s_tir_transform_lower_thread_all_reduce.py | 7 +- ...merge_dynamic_shared_memory_allocations.py | 34 ++- tests/python/te/test_te_tensor.py | 4 +- .../test_tir_analysis_undefined_vars.py | 20 +- tests/python/tirx-base/test_tir_buffer.py | 74 +++-- .../python/tirx-base/test_tir_constructor.py | 6 +- tests/python/tirx-base/test_tir_nodes.py | 3 +- tests/python/tirx-base/test_tir_op_types.py | 4 +- tests/python/tirx-base/test_tir_specialize.py | 14 +- .../python/tirx-base/test_tir_stmt_functor.py | 4 +- .../test_tir_inline_private_functions.py | 12 +- .../test_tir_transform_bf16_legalize.py | 4 +- .../test_tir_transform_flatten_buffer.py | 8 +- .../test_tir_transform_fp8_legalize.py | 11 + .../test_tir_transform_lower_intrin.py | 16 + .../test_tir_transform_lower_tvm_builtin.py | 30 +- .../test_tir_transform_make_packed_api.py | 25 ++ ...ir_transform_pointer_value_type_rewrite.py | 35 +++ .../test_tir_transform_split_host_device.py | 22 ++ .../test_tir_transform_storage_rewrite.py | 87 ++++-- .../tirx/codegen/test_codegen_hopper.py | 10 +- .../cuda/copy_async/test_tma.py | 2 +- .../tile_primitive/trn/test_copy_trn.py | 6 +- .../tile_primitive/trn/test_unary_trn.py | 2 +- tests/python/tirx/test_parser_printer.py | 28 +- .../tirx/transform/test_stmt_functor.py | 4 +- .../tvmscript/test_tvmscript_parser_tir.py | 4 +- 177 files changed, 2153 insertions(+), 1438 deletions(-) diff --git a/include/tvm/s_tir/schedule/schedule.h b/include/tvm/s_tir/schedule/schedule.h index f77d9b4edbd8..dd3b91877777 100644 --- a/include/tvm/s_tir/schedule/schedule.h +++ b/include/tvm/s_tir/schedule/schedule.h @@ -753,7 +753,7 @@ class ScheduleNode : public ffi::Object { /*! * \brief Apply a transformation represented by IndexMap to buffer * \details The indices and the access region to the target buffer is transformed by the given - * index_map. The index_map is used to infer the new shape of the buffer. BufferVar must be either + * index_map. The index_map is used to infer the new shape of the buffer. Buffer must be either * a function parameter, or allocated in a block (it cannot be a buffer subregion created via * 'match_buffer'). * \param block_rv The block that accesses the target buffer. @@ -825,7 +825,7 @@ class ScheduleNode : public ffi::Object { */ virtual void PadEinsum(const SBlockRV& block_rv, const ffi::Array& padding) = 0; - /******** Schedule: BufferVar transformation ********/ + /******** Schedule: Buffer transformation ********/ /*! * \brief Compute the target buffer via rolling buffering. * \details This primitive selects the outermost rollable axis with a positive bound overlap that diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h index c410ffbae512..36424ab6cc19 100644 --- a/include/tvm/s_tir/transform.h +++ b/include/tvm/s_tir/transform.h @@ -35,7 +35,7 @@ namespace tvm { namespace s_tir { /*! - * \brief Renew the definition nodes for a TIR, including Var, BufferVar and IterVar. + * \brief Renew the definition nodes for a TIR, including Var, Buffer and IterVar. * This pass works as a simple DeepCopy to duplicate a function with different Vars and * Buffers but the same behavior * \param func The input PrimFunc. diff --git a/include/tvm/tirx/buffer.h b/include/tvm/tirx/buffer.h index 66138d27fc4f..9413cd6b7f23 100644 --- a/include/tvm/tirx/buffer.h +++ b/include/tvm/tirx/buffer.h @@ -70,10 +70,10 @@ class Stmt; */ class BufferTypeNode : public TypeNode { public: - /*! \brief Type of the physical pointer projected by buffer_data. */ - PointerType data_pointer_type = PointerType::VoidPointerTy(); /*! \brief dtype in the content of the tensor */ PrimType dtype = PrimType::Void(); + /*! \brief Storage scope/address space of the buffer. */ + ffi::String storage_scope; /*! \brief The type of the buffer prior to flattening * * This contains the shape as it is accessed by @@ -110,8 +110,8 @@ class BufferTypeNode : public TypeNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() - .def_ro("data_pointer_type", &BufferTypeNode::data_pointer_type) .def_ro("dtype", &BufferTypeNode::dtype) + .def_ro("storage_scope", &BufferTypeNode::storage_scope) // TODO(tqchen): use SEqHashDefNonRecursive after the next pypi tvm-ffi release .def_ro("shape", &BufferTypeNode::shape, refl::AttachFieldFlag::SEqHashDefRecursive()) // TODO(tqchen): use SEqHashDefNonRecursive after the next pypi tvm-ffi release @@ -133,6 +133,9 @@ class BufferTypeNode : public TypeNode { /*! \return primitive element type for compiler-side uses. */ PrimType ElementType() const { return dtype; } + /*! \return type of the physical pointer projected by buffer_data. */ + PointerType DataPointerType() const { return PointerType(dtype, storage_scope); } + /*! \brief Determine the offset in the buffer of the given index. * * Returns the buffer offset, in number of elements of type dtype, @@ -152,7 +155,7 @@ class BufferTypeNode : public TypeNode { */ class BufferType : public Type { public: - TVM_DLL BufferType(PointerType data_pointer_type, PrimType dtype, ffi::Array shape, + TVM_DLL BufferType(ffi::String storage_scope, PrimType dtype, ffi::Array shape, ffi::Array strides, PrimExpr elem_offset, int data_alignment, int offset_factor, ffi::Optional layout = std::nullopt, ffi::Array allocated_addr = {}, Span span = Span()); @@ -175,21 +178,13 @@ class BufferType : public Type { */ class BufferVar : public Var { public: - // User can specify data_alignment and offset_factor to be 0 - // A default value will be picked. - TVM_DLL BufferVar(Var data, PrimType dtype, ffi::Array shape, - ffi::Array strides, PrimExpr elem_offset, ffi::String name, - int data_alignment, int offset_factor, Span span = Span(), - ffi::Optional layout = std::nullopt, - ffi::Array allocated_addr = {}); - /*! \brief Construct a fresh buffer variable from an explicit BufferType. */ TVM_DLL explicit BufferVar(ffi::String name, BufferType type, Span span = Span()); /*! \brief Create a checked buffer view over an existing ordinary Var. */ explicit BufferVar(Var var) : Var(std::move(var)) { - TVM_FFI_ICHECK(get() == nullptr || get()->ty.as()) - << "Expected a Var with BufferType"; + TVM_FFI_ICHECK(get() != nullptr && get()->ty.as()) + << "Expected a non-null Var with BufferType"; } /*! \brief Return the ordinary variable view over the same identity. */ @@ -204,11 +199,6 @@ class BufferVar : public Var { /*! \brief Return the source span carried by the ordinary Var. */ const Span& span() const { return get()->span; } - /*! \brief Return a fresh buffer variable with the same type and a new name. */ - BufferVar CopyWithName(ffi::String name) const { - return BufferVar(std::move(name), type(), span()); - } - /*! \brief Project the physical pointer established by the definition site. */ TVM_DLL Expr data() const; @@ -226,8 +216,7 @@ class BufferVar : public Var { * If stride is not needed in the slice, it won't be presented * \return the result buffer. */ - TVM_DLL BufferVar MakeSlice(ffi::Array begins, - ffi::Array extents) const; + TVM_DLL BufferVar MakeSlice(ffi::Array begins, ffi::Array extents) const; /*! * \brief Get access ptr to the entire buffer. * \param access_mask The access mask @@ -259,7 +248,11 @@ class BufferVar : public Var { ffi::Optional predicate = std::nullopt) const; /*! - * \brief Get a flattened version of the buffer + * \brief Get a flattened version of the buffer. + * + * If flattening changes the type, the result is a fresh BufferVar. Callers + * that use it as a view over this buffer must bind the returned variable with + * `DeclBuffer(flattened, this->data())`. */ BufferVar GetFlattenedBuffer() const; @@ -303,6 +296,9 @@ class BufferVar : public Var { /*! \return primitive element type for compiler-side uses. */ PrimType ElementType() const { return (*this)->ElementType(); } + /*! \return type of the physical pointer projected by buffer_data. */ + PointerType DataPointerType() const { return (*this)->DataPointerType(); } + BufferVar() = default; explicit BufferVar(ffi::ObjectPtr n) : Var(std::move(n)) {} explicit BufferVar(ffi::UnsafeInit tag) : Var(tag) {} @@ -313,32 +309,25 @@ class BufferVar : public Var { TVM_FFI_ICHECK(var_node != nullptr); const auto* type_node = var_node->ty.as(); TVM_FFI_ICHECK(type_node != nullptr) - << "Expected a Var with BufferType, but " << var_node->name << " has type " - << var_node->ty; + << "Expected a Var with BufferType, but " << var_node->name << " has type " << var_node->ty; return type_node; } const VarNode* get() const { return static_cast(data_.get()); } - [[maybe_unused]] static constexpr bool _type_is_nullable = true; + [[maybe_unused]] static constexpr bool _type_is_nullable = false; static constexpr bool _type_container_is_exact = false; using ContainerType = VarNode; }; -inline bool operator==(const BufferVar& lhs, const BufferVar& rhs) { - return lhs.same_as(rhs); -} +// Preserve ObjectRef-style identity comparison for exact BufferVar operands. +// Comparisons widened to Var or Expr continue to build symbolic expressions. +inline bool operator==(const BufferVar& lhs, const BufferVar& rhs) { return lhs.same_as(rhs); } -inline bool operator!=(const BufferVar& lhs, const BufferVar& rhs) { - return !lhs.same_as(rhs); -} +inline bool operator!=(const BufferVar& lhs, const BufferVar& rhs) { return !lhs.same_as(rhs); } /*! \brief Recover a checked buffer view from an ordinary VarNode pointer. */ -inline BufferVar GetBufferVar(const VarNode* var) { - return BufferVar(ffi::GetRef(var)); -} - -inline const BufferVar& GetBufferVar(const BufferVar& var) { return var; } +inline BufferVar GetBufferVar(const VarNode* var) { return BufferVar(ffi::GetRef(var)); } inline ffi::ObjectPtr CopyBufferType(const BufferVar& var) { return ffi::make_object(*var.operator->()); @@ -359,10 +348,9 @@ inline BufferVar RebuildBufferVar(const BufferVar& var, ffi::ObjectPtr shape, - PrimType dtype = PrimType::Float(32), - ffi::String name = "buffer", - ffi::String storage_scope = "", Span span = Span()); +TVM_DLL BufferVar decl_buffer(ffi::Array shape, PrimType dtype = PrimType::Float(32), + ffi::String name = "buffer", ffi::String storage_scope = "", + Span span = Span()); /*! * \brief Base node for data producers. @@ -408,7 +396,7 @@ class DataProducer : public PrimExprConvertible { }; /*! - * \brief Creates TIR BufferVar for provided parameters + * \brief Creates a TIR buffer for the provided parameters. * \param shape shape of the buffer * \param dtype data type * \param name buffer name @@ -421,8 +409,7 @@ class DataProducer : public PrimExprConvertible { */ TVM_DLL tirx::BufferVar BufferWithOffsetAlignment(ffi::Array shape, PrimType dtype, std::string name, int data_alignment, - int offset_factor, - std::string memory_scope = ""); + int offset_factor, std::string memory_scope = ""); } // namespace tirx } // namespace tvm @@ -444,7 +431,7 @@ struct TypeTraits : public ObjectRefTypeTraitsBasetype_index == TypeIndex::kTVMFFINone) { - return tirx::BufferVar::_type_is_nullable; + return false; } if (src->type_index != tirx::VarNode::RuntimeTypeIndex()) { return false; @@ -454,12 +441,8 @@ struct TypeTraits : public ObjectRefTypeTraitsBase(var->ExprNode::ty); } - TVM_FFI_INLINE static std::optional TryCastFromAnyView( - const TVMFFIAny* src) { + TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (CheckAnyStrict(src)) { - if (src->type_index == TypeIndex::kTVMFFINone) { - return details::ObjectUnsafe::ObjectRefFromObjectPtr(nullptr); - } return details::ObjectUnsafe::ObjectRefFromObjectPtr( details::ObjectUnsafe::ObjectPtrFromUnowned(src->v_obj)); } diff --git a/include/tvm/tirx/builtin.h b/include/tvm/tirx/builtin.h index 879196823463..30268940160f 100644 --- a/include/tvm/tirx/builtin.h +++ b/include/tvm/tirx/builtin.h @@ -783,8 +783,9 @@ TVM_DLL const Op& buffer_offset(); /*! * \brief Project the physical pointer associated with a BufferVar definition. * - * The result type is the BufferType::data_pointer_type of the sole BufferVar - * argument. This operation is consumed by TIRx lowering and code generation. + * The result pointer type is derived from the BufferType dtype and storage + * scope of the sole BufferVar argument. This operation is consumed by TIRx + * lowering and code generation. */ TVM_DLL const Op& buffer_data(); diff --git a/include/tvm/tirx/function.h b/include/tvm/tirx/function.h index 6b35c0bbb885..3e413d5ca8b4 100644 --- a/include/tvm/tirx/function.h +++ b/include/tvm/tirx/function.h @@ -53,13 +53,13 @@ class PrimFuncNode : public BaseFuncNode { /*! \brief The return type of the function. */ Type ret_type = Type::Missing(); /*! - * \brief Maps some parameters to specific BufferVar data structures. + * \brief Maps some parameters to specific buffer data structures. * * buffer_map provides a way to express data structure's field and shape * constraints. The provided information is used in the program analysis * and the code generation. * - * - It defines the vars in the BufferVar (m, n) in the cases below when + * - It defines the vars in the buffer (m, n) in the cases below when * they appears in the buffer_map for the first time. * - When a var appears multiple times, they translate into runtime * assertion to check the field constraint. diff --git a/include/tvm/tirx/index_map.h b/include/tvm/tirx/index_map.h index 776af5b5e885..7f33c4d26d92 100644 --- a/include/tvm/tirx/index_map.h +++ b/include/tvm/tirx/index_map.h @@ -21,7 +21,7 @@ * \file tvm/tirx/index_map.h * \brief Defines a remapping of buffer indices * - * For use with tvm::tirx::BufferVar. + * For use with TIR buffers. */ #ifndef TVM_TIR_INDEX_MAP_H_ #define TVM_TIR_INDEX_MAP_H_ @@ -46,7 +46,7 @@ namespace tirx { * \brief Defines a mapping between two representations of indices * into a buffer. * - * This is primarily used for layout transformations of BufferVar + * This is primarily used for layout transformations of buffers * objects. */ class IndexMapNode : public ffi::Object { diff --git a/include/tvm/tirx/script/builder/frame.h b/include/tvm/tirx/script/builder/frame.h index 1bd2d53981a8..16b351998875 100644 --- a/include/tvm/tirx/script/builder/frame.h +++ b/include/tvm/tirx/script/builder/frame.h @@ -77,7 +77,7 @@ class PrimFuncFrameNode : public TIRFrameNode { bool is_private; /*! \brief The return type of the function. */ ffi::Optional ret_type; - /*! \brief Maps some parameters to specific BufferVar data structures. */ + /*! \brief Maps some parameters to specific buffer data structures. */ ffi::Map buffer_map; /*! \brief Additional attributes storing the meta-data */ ffi::Map attrs; @@ -611,8 +611,8 @@ class DeclBufferFrameNode : public TIRFrameNode { public: /*! \brief The declared buffer. */ tvm::tirx::BufferVar buffer; - /*! \brief Optional physical pointer expression backing the declaration. */ - ffi::Optional data; + /*! \brief Physical pointer expression backing the declaration. */ + Expr data; /*! \brief The buffer allocated or not. */ bool allocated; diff --git a/include/tvm/tirx/script/builder/ir.h b/include/tvm/tirx/script/builder/ir.h index 80cd2aab8efb..40a92d194ab7 100644 --- a/include/tvm/tirx/script/builder/ir.h +++ b/include/tvm/tirx/script/builder/ir.h @@ -56,8 +56,7 @@ using tvm::tirx::Var; * \return The declared buffer. */ BufferVar BufferDecl(ffi::Array shape, PrimType dtype, ffi::String buffer_name, - ffi::Optional data, - ffi::Optional> strides, + ffi::Optional data, ffi::Optional> strides, ffi::Optional elem_offset, ffi::String storage_scope, int align, int offset_factor, ffi::Optional layout = std::nullopt, ffi::Array allocated_addr = {}); @@ -117,11 +116,9 @@ Type FuncRet(Type ret_type); * \return The matched buffer. */ BufferVar MatchBuffer(ffi::ObjectRef param, ffi::Array shape, - PrimType dtype = PrimType::Float(32), - ffi::Optional data = std::nullopt, + PrimType dtype = PrimType::Float(32), ffi::Optional data = std::nullopt, ffi::Array strides = {}, PrimExpr elem_offset = PrimExpr(), - ffi::String storage_scope = "global", int align = -1, - int offset_factor = 0, + ffi::String storage_scope = "global", int align = -1, int offset_factor = 0, ffi::Optional layout = std::nullopt); /*! @@ -409,8 +406,7 @@ ElseFrame Else(); * \return The declaration frame. */ DeclBufferFrame DeclBuffer(ffi::Array shape, PrimType dtype, ffi::String buffer_name, - ffi::Optional data, - ffi::Optional> strides, + ffi::Optional data, ffi::Optional> strides, ffi::Optional elem_offset, ffi::String storage_scope, int align, int offset_factor, ffi::Optional layout = std::nullopt, @@ -425,8 +421,8 @@ DeclBufferFrame DeclBuffer(ffi::Array shape, PrimType dtype, ffi::Stri * \return The allocated buffer. */ BufferVar AllocBuffer(ffi::Array shape, PrimType dtype = PrimType::Float(32), - ffi::String storage_scope = "global", - ffi::Optional> annotations = std::nullopt); + ffi::String storage_scope = "global", + ffi::Optional> annotations = std::nullopt); /*! * \brief Launch a thread. diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h index cf762abf0d05..9d64ac2e820c 100644 --- a/include/tvm/tirx/stmt.h +++ b/include/tvm/tirx/stmt.h @@ -240,14 +240,13 @@ class DeclBufferNode : public StmtNode { public: /*! \brief The buffer being declared */ BufferVar buffer; - /*! \brief Optional physical pointer expression backing the declaration. */ - ffi::Optional data; + /*! \brief Physical pointer expression backing the declaration. */ + Expr data; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() - .def_ro("buffer", &DeclBufferNode::buffer, - refl::AttachFieldFlag::SEqHashDefRecursive()) + .def_ro("buffer", &DeclBufferNode::buffer, refl::AttachFieldFlag::SEqHashDefRecursive()) .def_ro("data", &DeclBufferNode::data); } TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.DeclBuffer", DeclBufferNode, StmtNode); @@ -256,8 +255,7 @@ class DeclBufferNode : public StmtNode { /*! \brief Managed reference to DeclBufferNode */ class DeclBuffer : public Stmt { public: - TVM_DLL DeclBuffer(BufferVar buffer, ffi::Optional data = std::nullopt, - Span span = Span()); + TVM_DLL DeclBuffer(BufferVar buffer, Expr data, Span span = Span()); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DeclBuffer, Stmt, DeclBufferNode); TVM_DEFINE_OBJECT_REF_COW_METHOD(DeclBufferNode); }; diff --git a/python/tvm/backend/cuda/operator/tile_primitive/permute_layout/warp_xor_swizzle.py b/python/tvm/backend/cuda/operator/tile_primitive/permute_layout/warp_xor_swizzle.py index c2418ddd42e9..cf2d9b15561f 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/permute_layout/warp_xor_swizzle.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/permute_layout/warp_xor_swizzle.py @@ -75,7 +75,7 @@ from tvm.runtime import DataType from tvm.script import tirx as T -from tvm.tirx import Buffer, BufferRegion, IntImm, PrimFunc +from tvm.tirx import BufferRegion, IntImm, PrimFunc, is_buffer_var from tvm.tirx.layout import TileLayout, _flatten_coord from tvm.tirx.operator.tile_primitive import DispatchContext, fail, register_dispatch from tvm.tirx.stmt import TilePrimitiveCall @@ -87,9 +87,9 @@ def _as_buffer_and_region(arg): """Normalize a Buffer or BufferRegion to (buffer, start_list, extent_list).""" - if isinstance(arg, Buffer): + if is_buffer_var(arg): buf = arg - extent = list(buf.shape) + extent = list(buf.ty.shape) st = [0] * len(extent) elif isinstance(arg, BufferRegion): buf = arg.buffer diff --git a/python/tvm/backend/cuda/script.py b/python/tvm/backend/cuda/script.py index 46c7a5340903..062f27d38085 100644 --- a/python/tvm/backend/cuda/script.py +++ b/python/tvm/backend/cuda/script.py @@ -22,7 +22,7 @@ from typing import Any from tvm.backend.cuda import op as _cuda_op -from tvm.tirx import Buffer +from tvm.tirx import is_buffer_var from tvm.tirx import op as _tir_op from tvm.tirx.script.builder.ir import _dtype_forward, _op_wrapper @@ -30,7 +30,7 @@ def _ptx_ldg32(reg, guard, addr, local_addr): - if isinstance(addr, Buffer): + if is_buffer_var(addr): addr = addr[0] return _tir_op.call_intrin(reg.ty, "tirx.ptx.ldg32", reg, guard, addr, local_addr) @@ -485,25 +485,25 @@ def __init__(self): @staticmethod def _shfl_sync(mask, var, lane, width): - if isinstance(var, Buffer): + if is_buffer_var(var): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.cuda.__shfl_sync", mask, var, lane, width) @staticmethod def _shfl_up_sync(mask, var, delta, width): - if isinstance(var, Buffer): + if is_buffer_var(var): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.cuda.__shfl_up_sync", mask, var, delta, width) @staticmethod def _shfl_down_sync(mask, var, delta, width): - if isinstance(var, Buffer): + if is_buffer_var(var): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.cuda.__shfl_down_sync", mask, var, delta, width) @staticmethod def _shfl_xor_sync(mask, var, lane_mask, width): - if isinstance(var, Buffer): + if is_buffer_var(var): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.cuda.__shfl_xor_sync", mask, var, lane_mask, width) diff --git a/python/tvm/backend/metal/script.py b/python/tvm/backend/metal/script.py index 8f3fe337b465..e9913dac18db 100644 --- a/python/tvm/backend/metal/script.py +++ b/python/tvm/backend/metal/script.py @@ -19,7 +19,7 @@ from __future__ import annotations from tvm.backend.metal import op as _metal_op -from tvm.tirx import Buffer +from tvm.tirx import is_buffer_var from tvm.tirx import op as _tir_op from tvm.tirx.script.builder.ir import _op_wrapper @@ -35,19 +35,19 @@ def __init__(self): @staticmethod def simd_shuffle(var, lane): - if isinstance(var, Buffer): + if is_buffer_var(var): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.metal.simd_shuffle", var, lane) @staticmethod def simd_shuffle_up(var, delta): - if isinstance(var, Buffer): + if is_buffer_var(var): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.metal.simd_shuffle_up", var, delta) @staticmethod def simd_shuffle_down(var, delta): - if isinstance(var, Buffer): + if is_buffer_var(var): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.metal.simd_shuffle_down", var, delta) diff --git a/python/tvm/backend/trn/operator/tile_primitive/binary/default.py b/python/tvm/backend/trn/operator/tile_primitive/binary/default.py index 9546c2f6e43f..6b1baee43163 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/binary/default.py +++ b/python/tvm/backend/trn/operator/tile_primitive/binary/default.py @@ -56,7 +56,7 @@ def binary_trn( p_var = T.Var("P", "int32") b_var = T.Var("B", "int32") f_var = T.Var("F", "int32") - p_size = dst.layout.size("P") + p_size = dst.ty.layout.size("P") inst_size_limit = op.config.get("max_inst_size", 512) inst_repr.bound_inst_size(inst_size_limit, analyzer) inst_gen.bind_inst_iter(_dst, p_var, p_size, 1, False) diff --git a/python/tvm/backend/trn/operator/tile_primitive/binary/utils.py b/python/tvm/backend/trn/operator/tile_primitive/binary/utils.py index 891c5e9c4875..e3c35dc60332 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/binary/utils.py +++ b/python/tvm/backend/trn/operator/tile_primitive/binary/utils.py @@ -72,9 +72,9 @@ def try_find_inst_nary( valid_buffers = all( [ - dst.layout and all(src.layout for src in srcs if src is not None), - is_trainium_layout(dst.layout), - all(is_trainium_layout(src.layout) for src in srcs if src is not None), + dst.ty.layout and all(src.ty.layout for src in srcs if src is not None), + is_trainium_layout(dst.ty.layout), + all(is_trainium_layout(src.ty.layout) for src in srcs if src is not None), dst.scope() == "trn.sbuf", all(src.scope() in ["trn.sbuf", "trn.psum"] for src in srcs if src is not None), ] diff --git a/python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_chain.py b/python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_chain.py index daa64fe8adb9..e53651b12351 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_chain.py +++ b/python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_chain.py @@ -59,7 +59,7 @@ def binary_chain_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | p_var = T.Var("P", "int32") b_var = T.Var("B", "int32") f_var = T.Var("F", "int32") - p_size = output.buffer.layout.size("P") + p_size = output.buffer.ty.layout.size("P") inst_size_limit = op.config.get("max_inst_size", 512) inst_repr.bound_inst_size(inst_size_limit, analyzer) inst_gen.bind_inst_iter(output, p_var, p_size, 1, False) diff --git a/python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_reduce.py b/python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_reduce.py index d0c64d415331..35edbb31863c 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_reduce.py +++ b/python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_reduce.py @@ -42,7 +42,7 @@ def binary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc analyzer = init_analyzer(sctx) # Normalize negative axes - reduce_axes = [i if i >= 0 else len(binary_output.buffer.shape) + i for i in reduce_axes] + reduce_axes = [i if i >= 0 else len(binary_output.buffer.ty.shape) + i for i in reduce_axes] # Find instruction patterns inst_gen = InstructionGenerator( @@ -77,7 +77,7 @@ def binary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc f_var = T.Var("F", "int32") reduction_b_var = T.Var("rB", "int32") spatial_b_var = T.Var("sB", "int32") - p_size = binary_output.buffer.layout.size("P") + p_size = binary_output.buffer.ty.layout.size("P") inst_gen.bind_inst_iter(binary_output, p_var, p_size, 1, False) inst_gen.bind_inst_iter(binary_output, f_var, inst_repr.size, inst_repr.stride, True) reduction_b_extent = inst_gen.fill_in_block_dim(binary_output, reduction_b_var, reduce_axes) diff --git a/python/tvm/backend/trn/operator/tile_primitive/compose_op/unary_reduce.py b/python/tvm/backend/trn/operator/tile_primitive/compose_op/unary_reduce.py index a7c9f86c7b7a..54dd59fb7db2 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/compose_op/unary_reduce.py +++ b/python/tvm/backend/trn/operator/tile_primitive/compose_op/unary_reduce.py @@ -42,7 +42,7 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | analyzer = init_analyzer(sctx) # Normalize axes and default values - reduce_axes = [i if i >= 0 else len(unary_output.buffer.shape) + i for i in op.reduce_axes] + reduce_axes = [i if i >= 0 else len(unary_output.buffer.ty.shape) + i for i in op.reduce_axes] scale = 1.0 if scale is None else scale bias = 0.0 if bias is None else bias @@ -72,7 +72,7 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | f_var = T.Var("F", "int32") reduction_b_var = T.Var("rB", "int32") spatial_b_var = T.Var("sB", "int32") - p_size = unary_output.buffer.layout.size("P") + p_size = unary_output.buffer.ty.layout.size("P") inst_gen.bind_inst_iter(unary_output, p_var, p_size, 1, False) inst_gen.bind_inst_iter(unary_output, f_var, inst_repr.size, inst_repr.stride, True) reduction_b_extent = inst_gen.fill_in_block_dim(unary_output, reduction_b_var, reduce_axes) @@ -90,7 +90,9 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | bias_buffer = ( bias.buffer if isinstance(bias, BufferRegion) - else get_const_bias_tensor(bias, (p_size, inst_repr.size), dst1.dtype, op.workspace, sctx) + else get_const_bias_tensor( + bias, (p_size, inst_repr.size), dst1.ty.dtype, op.workspace, sctx + ) ) # Create appropriate implementation based on intermediate buffer requirement diff --git a/python/tvm/backend/trn/operator/tile_primitive/copy/default.py b/python/tvm/backend/trn/operator/tile_primitive/copy/default.py index 123698b62164..112b78ab6488 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/copy/default.py +++ b/python/tvm/backend/trn/operator/tile_primitive/copy/default.py @@ -47,8 +47,8 @@ def transpose_schedule( dst_f = T.Var("dst_F", "int32") b_var = T.Var("B", "int32") extend_b = T.Var("extend_B", "int32") - p_size = src_region.buffer.layout.size("P") - lhs_f_size = dst_region.buffer.layout.size("P") + p_size = src_region.buffer.ty.layout.size("P") + lhs_f_size = dst_region.buffer.ty.layout.size("P") rhs_f_size = p_size inst_gen.bind_inst_iter( src_region, lhs_f, inst_repr_src.size, inst_repr_src.stride, is_free_dim=True @@ -90,7 +90,10 @@ def transpose_schedule( "Identity tensor must be specified in workspace. Run tvm.tirx.trn.transform.TrnPrivateBufferAlloc first." # noqa: E501 ) identity_tensor = T.buffer( - (p_size, rhs_f_size), src_region.buffer.dtype, scope="trn.sbuf", buffer_name="identity" + (p_size, rhs_f_size), + src_region.buffer.ty.dtype, + scope="trn.sbuf", + buffer_name="identity", ) sctx.add_alloc_buffer(identity_tensor) @@ -158,7 +161,7 @@ def transpose_psum_output(): else: acc_psum = op.workspace["acc_psum"] check_workspace_buffer(acc_psum, (p_size, largest_psum_per_bank), "trn.psum") - max_psum_slots = acc_psum.shape[0] + max_psum_slots = acc_psum.ty.shape[0] # fmt: off @T.prim_func @@ -198,14 +201,14 @@ def copy_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None: # Check for valid buffer configurations valid_config = all( [ - src.layout and dst.layout, + src.ty.layout and dst.ty.layout, src.scope() in ["global", "trn.sbuf", "trn.psum"], dst.scope() in ["global", "trn.sbuf", "trn.psum"], src.scope() != "global" or dst.scope() != "global", - (src.scope() == "global" and isinstance(src.layout, T.TileLayout)) - or (src.scope() in ["trn.sbuf", "trn.psum"] and is_trainium_layout(src.layout)), - (dst.scope() == "global" and isinstance(dst.layout, T.TileLayout)) - or (dst.scope() in ["trn.sbuf", "trn.psum"] and is_trainium_layout(dst.layout)), + (src.scope() == "global" and isinstance(src.ty.layout, T.TileLayout)) + or (src.scope() in ["trn.sbuf", "trn.psum"] and is_trainium_layout(src.ty.layout)), + (dst.scope() == "global" and isinstance(dst.ty.layout, T.TileLayout)) + or (dst.scope() in ["trn.sbuf", "trn.psum"] and is_trainium_layout(dst.ty.layout)), ] ) @@ -233,7 +236,7 @@ def copy_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None: if not inst_gen.check_partition_dim_match(src_region, dst_region): return transpose_schedule(op, inst_gen, sctx) - if is_trainium_layout(src.layout): + if is_trainium_layout(src.ty.layout): inst = inst_gen.find_max_inst_size_from_one_region(src_region) inst = inst_gen.fit_inst_tile_to_region(inst, dst_region) src_to_dst = True @@ -262,7 +265,7 @@ def copy_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None: from_region, _to_region = src_region, dst_region else: from_region, _to_region = dst_region, src_region - p_size = from_region.buffer.layout.size("P") + p_size = from_region.buffer.ty.layout.size("P") inst_gen.bind_inst_iter(from_region, p_var, p_size, 1, is_free_dim=False) inst_gen.bind_inst_iter(from_region, f_var, inst.size, inst.stride, is_free_dim=True) b_extent = inst_gen.fill_in_block_dim(from_region, b_var) diff --git a/python/tvm/backend/trn/operator/tile_primitive/gemm/default.py b/python/tvm/backend/trn/operator/tile_primitive/gemm/default.py index eba660ab6d6a..d0ed02ba2bee 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/gemm/default.py +++ b/python/tvm/backend/trn/operator/tile_primitive/gemm/default.py @@ -55,12 +55,14 @@ def get_pf_dim_from_buffer_region( # Find non-unit dimensions non_unit_dims = [ i - for i in range(len(buffer_region.buffer.shape)) + for i in range(len(buffer_region.buffer.ty.shape)) if not analyzer.can_prove_equal(buffer_region.region[i].extent, 1) ] assert len(non_unit_dims) == 2, "Only 2D matrix is supported for gemm" - layout, seps = normalize_and_group(buffer_region.buffer.layout, buffer_region.buffer.shape) + layout, seps = normalize_and_group( + buffer_region.buffer.ty.layout, buffer_region.buffer.ty.shape + ) # Determine partition and free dimensions based on operator kind if operator_kind == OperatorKind.A: p_dim, f_dim = non_unit_dims[1], non_unit_dims[0] @@ -144,19 +146,19 @@ def matmul_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None: # Validate buffer properties assert all( [ - A.layout and B.layout and C.layout, - A.dtype == B.dtype, + A.ty.layout and B.ty.layout and C.ty.layout, + A.ty.dtype == B.ty.dtype, A.scope() == "trn.sbuf" and B.scope() == "trn.sbuf", C.scope() == "trn.psum" or C.scope() == "trn.sbuf", - is_trainium_layout(A.layout), - is_trainium_layout(B.layout), - is_trainium_layout(C.layout), - A.layout.size("P") == B.layout.size("P"), + is_trainium_layout(A.ty.layout), + is_trainium_layout(B.ty.layout), + is_trainium_layout(C.ty.layout), + A.ty.layout.size("P") == B.ty.layout.size("P"), ] ), "Invalid buffer layout and scope" - p_size = A.layout.size("P") - assert p_size == B.layout.size("P"), "Partition size mismatch" + p_size = A.ty.layout.size("P") + assert p_size == B.ty.layout.size("P"), "Partition size mismatch" # Get partition and free dimensions lhs_p_dim, lhs_f_dim = get_pf_dim_from_buffer_region( @@ -206,12 +208,12 @@ def matmul_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None: reduction_b = T.Var("reduction_b", "int32") lhs_b = T.Var("lhs_b", "int32") rhs_b = T.Var("rhs_b", "int32") - lhs_f_size = C.layout.size("P") + lhs_f_size = C.ty.layout.size("P") inst_gen.bind_inst_iter( B_buffer_region, rhs_f, inst_repr.size, inst_repr.stride, is_free_dim=True ) inst_gen.bind_inst_iter(C_buffer_region, lhs_f, lhs_f_size, 1, is_free_dim=False) - inst_gen.bind_inst_iter(A_buffer_region, p, A.layout.size("P"), 1, is_free_dim=False) + inst_gen.bind_inst_iter(A_buffer_region, p, A.ty.layout.size("P"), 1, is_free_dim=False) reduction_b_extent = inst_gen.fill_in_block_dim(A_buffer_region, reduction_b, [lhs_p_dim]) lhs_b_extent = inst_gen.fill_in_block_dim(A_buffer_region, lhs_b, [lhs_f_dim]) rhs_b_extent = inst_gen.fill_in_block_dim(B_buffer_region, rhs_b, [rhs_f_dim]) @@ -266,7 +268,7 @@ def impl_C_psum(): else: acc_psum = op.workspace["acc_psum"] check_workspace_buffer(acc_psum, (p_size, largest_psum_per_bank), "trn.psum") - max_psum_slots = acc_psum.shape[0] + max_psum_slots = acc_psum.ty.shape[0] @T.prim_func def impl_C_sbuf(): diff --git a/python/tvm/backend/trn/operator/tile_primitive/instruction_generator.py b/python/tvm/backend/trn/operator/tile_primitive/instruction_generator.py index 9b0c40f69684..4568bcaf2be3 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/instruction_generator.py +++ b/python/tvm/backend/trn/operator/tile_primitive/instruction_generator.py @@ -137,8 +137,8 @@ def _bound_buffer_region(self, buffer_region: BufferRegion): return buffer_region def _get_sub_layout(self, buffer_region: BufferRegion): - layout = buffer_region.buffer.layout - layout, seps = normalize_and_group(layout, buffer_region.buffer.shape) + layout = buffer_region.buffer.ty.layout + layout, seps = normalize_and_group(layout, buffer_region.buffer.ty.shape) tiled_range_infos_per_dim = [] new_shard = [] new_seps = [0] @@ -367,7 +367,7 @@ def fill_in_block_dim( # fixme: be cautious of the min of buffer region. This implementation is not correct. # we need to first take a view of sub-layout (keep strides, but reduce the extent # then we analyze the relationship between data iter of sub-layout - dims = dims or list(range(len(buffer_region.buffer.shape))) + dims = dims or list(range(len(buffer_region.buffer.ty.shape))) layout = self.split_layout_views[buffer_region] shards = layout.shard self._normalize_bind_iters() @@ -714,7 +714,9 @@ def restrict_inst_to_one_dim(self, inst_repr: InstructionRepr): indexed_selected_iters = [(i, iters[i]) for i in inst_repr.selected_data_iter_ids] indexed_selected_iters = sorted(indexed_selected_iters, key=lambda x: x[1].stride) iter_idx_to_dim = { - seps[j]: i for i in range(len(region.buffer.shape)) for j in range(seps[i], seps[i + 1]) + seps[j]: i + for i in range(len(region.buffer.ty.shape)) + for j in range(seps[i], seps[i + 1]) } last_dim = None inst_size = 1 diff --git a/python/tvm/backend/trn/operator/tile_primitive/private_alloc.py b/python/tvm/backend/trn/operator/tile_primitive/private_alloc.py index 88720f7cdccc..91f6f3b619c5 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/private_alloc.py +++ b/python/tvm/backend/trn/operator/tile_primitive/private_alloc.py @@ -49,16 +49,16 @@ def _scalar_dtype(scalar) -> str: def alloc_const_bias_trn( op: TilePrimitiveCall, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: DispatchContext ) -> dict[str, Any]: - bias = op.bias if op.bias is not None else FloatImm(op.dsts[0].buffer.dtype, 0.0) + bias = op.bias if op.bias is not None else FloatImm(op.dsts[0].buffer.ty.dtype, 0.0) if "const_bias" in op.workspace: return {} if not isinstance(bias, (FloatImm)): return {} - par_size = op.dsts[0].buffer.layout.size("P") + par_size = op.dsts[0].buffer.ty.layout.size("P") max_inst_size = op.config.get("max_inst_size", 512) if ("const_bias", bias.value) in buffer_dict: bias_buffer, bias_init_stmt = buffer_dict[("const_bias", bias.value)] - old_shape = bias_buffer.shape + old_shape = bias_buffer.ty.shape new_shape = [max(par_size, old_shape[0]), max(max_inst_size, old_shape[1])] if new_shape[0] == old_shape[0] and new_shape[1] == old_shape[1]: return {"const_bias": ("const_bias", bias.value)} @@ -105,17 +105,17 @@ def alloc_identity_trn( ) -> dict[str, Any]: if "identity" in op.workspace: return {} - par_size = op.srcs[0].buffer.layout.size("P") + par_size = op.srcs[0].buffer.ty.layout.size("P") if "identity" in buffer_dict: identity_buffer, identity_init_stmt = buffer_dict["identity"] - old_shape = identity_buffer.shape + old_shape = identity_buffer.ty.shape new_shape = [max(par_size, old_shape[0]), max(par_size, old_shape[1])] if new_shape[0] == old_shape[0] and new_shape[1] == old_shape[1]: return {"identity": "identity"} else: new_shape = (par_size, par_size) new_buffer = T.buffer( - new_shape, dtype=op.srcs[0].buffer.dtype, scope="trn.sbuf", buffer_name="identity" + new_shape, dtype=op.srcs[0].buffer.ty.dtype, scope="trn.sbuf", buffer_name="identity" ) @T.prim_func @@ -135,7 +135,7 @@ def alloc_acc_psum_trn( ) -> dict[str, Any]: if "acc_psum" in op.workspace or op.dsts[0].buffer.scope() == "trn.psum": return {} - par_size = op.dsts[0].buffer.layout.size("P") + par_size = op.dsts[0].buffer.ty.layout.size("P") acc_psum = T.buffer( (8, par_size, 512), "float32", diff --git a/python/tvm/backend/trn/operator/tile_primitive/reduction/utils.py b/python/tvm/backend/trn/operator/tile_primitive/reduction/utils.py index 43bd92d56d75..7fa26b11a1f6 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/reduction/utils.py +++ b/python/tvm/backend/trn/operator/tile_primitive/reduction/utils.py @@ -40,7 +40,7 @@ def generate_intermediate_buffer( Returns: Tuple[Optional[buffer], int]: The intermediate buffer and reduction factor size. """ - intermediate_shape = [dst_buffer_region.buffer.layout.size("P"), rfactor_size] + intermediate_shape = [dst_buffer_region.buffer.ty.layout.size("P"), rfactor_size] if "partial_reduce" in workspace: intermediate_buffer = workspace["partial_reduce"] @@ -51,7 +51,7 @@ def generate_intermediate_buffer( ) intermediate_buffer = T.buffer( intermediate_shape, - dtype=dst_buffer_region.buffer.dtype, + dtype=dst_buffer_region.buffer.ty.dtype, scope="trn.sbuf", buffer_name="partial_reduce", ) @@ -85,18 +85,18 @@ def reduction_trn( # Extract buffers dst = dst_buffer_region.buffer src = src_buffer_region.buffer - axes = [i if i >= 0 else len(src.shape) + i for i in axes] + axes = [i if i >= 0 else len(src.ty.shape) + i for i in axes] dim_map = get_reduction_dim_map(src_buffer_region, dst_buffer_region, axes, analyzer) # Layout validation assert all( [ - src.layout and dst.layout, + src.ty.layout and dst.ty.layout, src.scope() == "trn.sbuf" or src.scope() == "trn.psum", dst.scope() == "trn.sbuf", - is_trainium_layout(src.layout), - is_trainium_layout(dst.layout), - src.layout.size("P") == dst.layout.size("P"), + is_trainium_layout(src.ty.layout), + is_trainium_layout(dst.ty.layout), + src.ty.layout.size("P") == dst.ty.layout.size("P"), ] ), "Invalid layout" @@ -109,7 +109,7 @@ def reduction_trn( assert analyzer.can_prove(inst_repr.size > 1), "Instruction size must be greater than 1" # Get partition size and extents - p_size = src.layout.size("P") + p_size = src.ty.layout.size("P") f_var = T.Var("F", "int32") p_var = T.Var("P", "int32") spatial_b_var = T.Var("sB", "int32") diff --git a/python/tvm/backend/trn/operator/tile_primitive/select/default.py b/python/tvm/backend/trn/operator/tile_primitive/select/default.py index 07486e6ae7ab..ca44225ae782 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/select/default.py +++ b/python/tvm/backend/trn/operator/tile_primitive/select/default.py @@ -62,10 +62,10 @@ def select_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None: # Validate buffer layout and scope buffer_conditions = [ - dst.buffer.layout and true_value.buffer.layout, + dst.buffer.ty.layout and true_value.buffer.ty.layout, dst.buffer.scope() == "trn.sbuf" and true_value.buffer.scope() == "trn.sbuf", - is_trainium_layout(true_value.buffer.layout), - is_trainium_layout(dst.buffer.layout), + is_trainium_layout(true_value.buffer.ty.layout), + is_trainium_layout(dst.buffer.ty.layout), ] if not all(buffer_conditions): @@ -98,7 +98,7 @@ def select_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None: p_var = T.Var("p", "int32") b_var = T.Var("b", "int32") f_var = T.Var("f", "int32") - p_size = dst.buffer.layout.size("P") + p_size = dst.buffer.ty.layout.size("P") inst_gen.bind_inst_iter(dst, f_var, inst_repr.size, inst_repr.stride, True) inst_gen.bind_inst_iter(dst, p_var, p_size, 1, False) b_extent = inst_gen.fill_in_block_dim(dst, b_var) diff --git a/python/tvm/backend/trn/operator/tile_primitive/unary/utils.py b/python/tvm/backend/trn/operator/tile_primitive/unary/utils.py index dc58862bb4f9..a982f2fb2fbb 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/unary/utils.py +++ b/python/tvm/backend/trn/operator/tile_primitive/unary/utils.py @@ -54,11 +54,11 @@ def try_find_inst_unary( # Validate buffer layouts and scopes valid_layout_scope = all( [ - src.layout and dst.layout, + src.ty.layout and dst.ty.layout, src.scope() in ("trn.sbuf", "trn.psum"), dst.scope() == "trn.sbuf", - is_trainium_layout(src.layout), - is_trainium_layout(dst.layout), + is_trainium_layout(src.ty.layout), + is_trainium_layout(dst.ty.layout), ] ) @@ -136,7 +136,7 @@ def generate_unary_func( ): """Generate a function that implements a unary operation.""" # Prepare parameters - p_size = dst_buffer_region.buffer.layout.size("P") + p_size = dst_buffer_region.buffer.ty.layout.size("P") # Apply instruction size limits if specified inst_size_limit = config.get("max_inst_size", 512) @@ -159,7 +159,7 @@ def generate_unary_func( # Handle bias tensor if isinstance(bias, FloatImm | float): bias_buffer = get_const_bias_tensor( - bias, (p_size, inst_repr.size), dst.dtype, workspace, sctx + bias, (p_size, inst_repr.size), dst.ty.dtype, workspace, sctx ) elif isinstance(bias, BufferRegion): bias_buffer = bias.buffer diff --git a/python/tvm/backend/trn/operator/tile_primitive/workspace_utils.py b/python/tvm/backend/trn/operator/tile_primitive/workspace_utils.py index 26fb38933595..ec993a4b5eca 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/workspace_utils.py +++ b/python/tvm/backend/trn/operator/tile_primitive/workspace_utils.py @@ -41,14 +41,14 @@ def check_workspace_buffer(buffer: Buffer, shape: tuple[int], scope: str): If the buffer is invalid """ assert buffer.scope() == scope, f"workspace buffer must be a {scope} buffer" - assert buffer.layout is None, "workspace buffer must not have a layout" + assert buffer.ty.layout is None, "workspace buffer must not have a layout" if scope == "trn.psum": # the number of psum banks used is inferred from the shape # only check p and f dims - assert all(x >= y for x, y in zip(buffer.shape[1:], shape)), ( - f"workspace buffer must have enough size, {buffer.shape[1:]} cannot cover {shape}" + assert all(x >= y for x, y in zip(buffer.ty.shape[1:], shape)), ( + f"workspace buffer must have enough size, {buffer.ty.shape[1:]} cannot cover {shape}" ) else: - assert all(x >= y for x, y in zip(buffer.shape, shape)), ( - f"workspace buffer must have enough size, {buffer.shape} cannot cover {shape}" + assert all(x >= y for x, y in zip(buffer.ty.shape, shape)), ( + f"workspace buffer must have enough size, {buffer.ty.shape} cannot cover {shape}" ) diff --git a/python/tvm/backend/trn/transform/naive_allocator.py b/python/tvm/backend/trn/transform/naive_allocator.py index 72f628488660..1e7f0c4f76ca 100644 --- a/python/tvm/backend/trn/transform/naive_allocator.py +++ b/python/tvm/backend/trn/transform/naive_allocator.py @@ -25,7 +25,7 @@ def is_const_shape(buffer: Buffer) -> bool: - for i in buffer.shape: + for i in buffer.ty.shape: if not isinstance(i, IntImm): return False return True @@ -33,21 +33,21 @@ def is_const_shape(buffer: Buffer) -> bool: def get_buffer_size(buffer: Buffer) -> int: if buffer.scope() == "trn.sbuf": - if buffer.layout is None: + if buffer.ty.layout is None: # the first dimension is partition size - num_elem = functools.reduce(lambda x, y: x * y, buffer.shape[1:]) + num_elem = functools.reduce(lambda x, y: x * y, buffer.ty.shape[1:]) else: - par_size = buffer.layout.size("P") - num_elem = functools.reduce(lambda x, y: x * y, buffer.shape) // par_size + par_size = buffer.ty.layout.size("P") + num_elem = functools.reduce(lambda x, y: x * y, buffer.ty.shape) // par_size elif buffer.scope().startswith("shared"): - num_elem = functools.reduce(lambda x, y: x * y, buffer.shape) + num_elem = functools.reduce(lambda x, y: x * y, buffer.ty.shape) else: return None if not is_const_shape(buffer): raise ValueError( f"Buffer {buffer.name} has non-constant shape. Do not know how to allocate it." ) - return int(num_elem * buffer.dtype.dtype.itemsize) + return int(num_elem * buffer.ty.dtype.dtype.itemsize) class AllocInfoCollector(StmtVisitor): @@ -58,12 +58,14 @@ def __init__(self): def visit_alloc_buffer_(self, op: AllocBuffer): super().visit_alloc_buffer_(op) buffer = op.buffer - if len(buffer.allocated_addr) == 0: + if len(buffer.ty.allocated_addr) == 0: return op buffer_size = get_buffer_size(buffer) if buffer_size is None: return op - self.alloc_pool_start = max(self.alloc_pool_start, buffer.allocated_addr[-1] + buffer_size) + self.alloc_pool_start = max( + self.alloc_pool_start, buffer.ty.allocated_addr[-1] + buffer_size + ) class AllocMutator(BufferReplacer): @@ -75,11 +77,12 @@ def visit_alloc_buffer_(self, op: AllocBuffer): changed = False buffer = op.buffer buffer_size = get_buffer_size(buffer) - if len(buffer.allocated_addr) > 0 or buffer_size is None: + if len(buffer.ty.allocated_addr) > 0 or buffer_size is None: pass else: new_buffer = buffer.with_allocated_addr([self.alloc_offset]) self.buffer_map[buffer] = new_buffer + self.var_map[buffer] = new_buffer changed = True self.alloc_offset += buffer_size diff --git a/python/tvm/s_tir/schedule/schedule.py b/python/tvm/s_tir/schedule/schedule.py index a53b9eb655b4..36e590b49250 100644 --- a/python/tvm/s_tir/schedule/schedule.py +++ b/python/tvm/s_tir/schedule/schedule.py @@ -25,7 +25,7 @@ from tvm.error import register_error from tvm.ir import Expr, GlobalVar, IRModule, is_prim_expr from tvm.runtime import DataTypeCode, Object -from tvm.tirx import Buffer, FloatImm, For, IntImm, PrimFunc, SBlock, is_buffer +from tvm.tirx import Buffer, FloatImm, For, IntImm, PrimFunc, SBlock, is_buffer_var from tvm.tirx.function import IndexMap from . import _ffi_api @@ -3272,7 +3272,7 @@ def iter_buffers(): ) buffer_obj, (buffer_index_type, buffer_index) = next(iter(possible_buffers.items())) - elif is_buffer(buffer): + elif is_buffer_var(buffer): # Buffer lookup has unique id, can break out early found = False for buffer_index_type, buffer_index, buffer_obj in iter_buffers(): @@ -3437,7 +3437,7 @@ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> block = self._normalize_block_arg(block) buffer_index_type, buffer_index, buffer_obj = self._normalize_buffer_arg(block, buffer) - ndim = len(buffer_obj.shape) + ndim = len(buffer_obj.ty.shape) if callable(index_map): index_map = IndexMap.from_func( index_map, @@ -3458,14 +3458,14 @@ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> # buffer's type. If the default `tvm.runtime.convert` # behavior is applied, these would be converted to # int32/float32, which may not match the buffer's type. - if buffer_obj.dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT) and isinstance( + if buffer_obj.ty.dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT) and isinstance( pad_value, int ): - pad_value = IntImm(buffer_obj.dtype.dtype, pad_value) - elif buffer_obj.dtype.matches_code(DataTypeCode.FLOAT, DataTypeCode.BFLOAT) and ( - isinstance(pad_value, float) - ): - pad_value = FloatImm(buffer_obj.dtype.dtype, pad_value) + pad_value = IntImm(buffer_obj.ty.dtype.dtype, pad_value) + elif buffer_obj.ty.dtype.matches_code( + DataTypeCode.FLOAT, DataTypeCode.BFLOAT + ) and isinstance(pad_value, float): + pad_value = FloatImm(buffer_obj.ty.dtype.dtype, pad_value) pad_value = IndexMap.from_func( lambda *indices: pad_value, ndim=len(index_map.final_indices), diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py index b7dcf734bd65..e1c84297a311 100644 --- a/python/tvm/tirx/__init__.py +++ b/python/tvm/tirx/__init__.py @@ -26,7 +26,16 @@ from tvm.ir import Expr from tvm.runtime import const -from .buffer import Buffer, BufferType, DataProducer, decl_buffer, is_buffer +from .buffer import ( + Buffer, + BufferAccessKind, + BufferType, + DataProducer, + buffer_data, + buffer_data_pointer_type, + decl_buffer, + is_buffer_var, +) from .expr import convert from .expr import Var, Reduce, FloatImm, IntImm, StringImm, Cast from .expr import Add, Sub, Mul, Div, Mod, FloorDiv, FloorMod diff --git a/python/tvm/tirx/buffer.py b/python/tvm/tirx/buffer.py index a0e51bdef061..625926581b04 100644 --- a/python/tvm/tirx/buffer.py +++ b/python/tvm/tirx/buffer.py @@ -17,13 +17,14 @@ """Abstraction for array data structures.""" import functools +from enum import IntEnum from numbers import Integral import tvm_ffi import tvm from tvm.ir import PointerType, PrimType, Range, Type -from tvm.runtime import Object, Scriptable, convert +from tvm.runtime import Object, convert from . import _ffi_api @@ -32,8 +33,8 @@ class BufferType(Type): """The structural type carried by an ordinary buffer variable.""" - data_pointer_type: PointerType dtype: PrimType + storage_scope: str shape: list strides: list elem_offset: tvm.ir.Expr @@ -43,12 +44,24 @@ class BufferType(Type): allocated_addr: list -def is_buffer(value) -> bool: - """Return whether ``value`` is an ordinary Var carrying BufferType.""" +def is_buffer_var(value) -> bool: + """Return whether ``value`` is an ordinary Var carrying BufferType. + + Use this predicate instead of ``isinstance(value, Buffer)``. ``Buffer`` is + a source-compatibility alias for :class:`tvm.ir.Var` and therefore does not + discriminate buffer variables from scalar or pointer variables. + """ return isinstance(value, tvm.ir.Var) and isinstance(value.ty, BufferType) +class BufferAccessKind(IntEnum): + """Buffer access modes accepted by :func:`buffer_access_ptr`.""" + + READ = 1 + WRITE = 2 + + class _BufferMethods: """Symbolic data buffer in TVM. @@ -63,9 +76,6 @@ class _BufferMethods: decl_buffer : Declare a buffer """ - READ = 1 - WRITE = 2 - def access_ptr(self, access_mask, ptr_type="handle", content_lanes=1, offset=0, extent=None): """Get an access pointer to the head of buffer. @@ -100,7 +110,7 @@ def access_ptr(self, access_mask, ptr_type="handle", content_lanes=1, offset=0, # Get access ptr for read buffer.access_ptr("r") # Get access ptr for read/write with bitmask - buffer.access_ptr(Buffer.READ | Buffer.WRITE) + buffer.access_ptr(BufferAccessKind.READ | BufferAccessKind.WRITE) # Get access ptr for read/write with str flag buffer.access_ptr("rw") # Get access ptr for read with offset @@ -112,9 +122,9 @@ def access_ptr(self, access_mask, ptr_type="handle", content_lanes=1, offset=0, mask = 0 for value in access_mask: if value == "r": - mask = mask | Buffer.READ + mask = mask | BufferAccessKind.READ elif value == "w": - mask = mask | Buffer.WRITE + mask = mask | BufferAccessKind.WRITE else: raise ValueError(f"Unknown access_mask {access_mask}") access_mask = mask @@ -159,7 +169,7 @@ def vload(self, begin, dtype=None, predicate=None): The corresponding load expression. """ begin = (begin,) if isinstance(begin, int) or tvm.ir.is_prim_expr(begin) else begin - dtype = dtype if dtype else self.dtype + dtype = dtype if dtype else self.ty.dtype return _ffi_api.BufferVLoad(self, begin, dtype, predicate) # type: ignore def vstore(self, begin, value, predicate=None): @@ -233,7 +243,7 @@ def offset_of(self, indices): @property def byte_offset(self): """Get the byte offset of the buffer.""" - return self.elem_offset * tvm.DataType(self.dtype).bits // 8 + return self.ty.elem_offset * tvm.DataType(self.ty.dtype).bits // 8 def elem_offset_of(self, indices, inner=True): """Get the element offset of the buffer at the given indices. @@ -255,7 +265,7 @@ def elem_offset_of(self, indices, inner=True): """ if inner: return _ffi_api.BufferOffsetOfp(self, indices) - return self.elem_offset + _ffi_api.BufferOffsetOfp(self, indices) + return self.ty.elem_offset + _ffi_api.BufferOffsetOfp(self, indices) def byte_offset_of(self, indices, inner=True): """Get the byte offset of the buffer at the given indices. @@ -275,7 +285,7 @@ def byte_offset_of(self, indices, inner=True): offset: Expr The byte offset of the buffer at the given indices. """ - return self.elem_offset_of(indices, inner) * tvm.DataType(self.dtype).bits // 8 + return self.elem_offset_of(indices, inner) * tvm.DataType(self.ty.dtype).bits // 8 def is_scalar(self, alloc_or_decl=True): """Check if the buffer is a scalar. @@ -297,8 +307,9 @@ def ptr_to(self, indices): Note that the bufferload inside requires LowerTIPp pass to apply the layout to get the physical indices. """ # noqa: E501 - assert len(indices) == len(self.shape), ( - f"The number of indices {indices} does not match the shape of the buffer {self.shape}" + assert len(indices) == len(self.ty.shape), ( + f"The number of indices {indices} does not match " + f"the shape of the buffer {self.ty.shape}" ) return tvm.tirx.address_of(self[tuple(indices)]) @@ -318,7 +329,7 @@ def view(self, *args, **kwargs) -> "Buffer": def _infer_shape(shape): shape = list(shape) if -1 in shape and shape.count(-1) == 1: - size = functools.reduce(lambda x, y: x * y, self.shape) + size = functools.reduce(lambda x, y: x * y, self.ty.shape) n_size = functools.reduce(lambda x, y: x * y, [s for s in shape if s != -1], 1) shape[shape.index(-1)] = size // n_size else: @@ -326,13 +337,13 @@ def _infer_shape(shape): # are fully concrete: a Expr `==` returns an `EQ` node, not # a Python bool, and `assert ` raises (no __bool__). if all(isinstance(s, int) for s in shape) and all( - isinstance(s, int) for s in self.shape + isinstance(s, int) for s in self.ty.shape ): assert functools.reduce(lambda x, y: x * y, shape) == functools.reduce( - lambda x, y: x * y, self.shape + lambda x, y: x * y, self.ty.shape ), ( "The shape of the buffer " - + str(self.shape) + + str(self.ty.shape) + " and the new shape " + str(shape) + " are not compatible" @@ -341,31 +352,31 @@ def _infer_shape(shape): if len(args) == 1 and isinstance(args[0], str | tvm.DataType) and not kwargs: cast_dtype = tvm.DataType(args[0]) - cur_dtype = tvm.DataType(self.dtype) + cur_dtype = tvm.DataType(self.ty.dtype) if cast_dtype.bits > cur_dtype.bits: # cast up assert cast_dtype.bits % cur_dtype.bits == 0 ratio = cast_dtype.bits // cur_dtype.bits - layout = self.layout.pack(ratio) - shape = [s for s in self.shape[:-1]] + [self.shape[-1] // ratio] - new_elem_offset = self.elem_offset // ratio + layout = self.ty.layout.pack(ratio) + shape = [s for s in self.ty.shape[:-1]] + [self.ty.shape[-1] // ratio] + new_elem_offset = self.ty.elem_offset // ratio else: # cast down assert cur_dtype.bits % cast_dtype.bits == 0 ratio = cur_dtype.bits // cast_dtype.bits - layout = self.layout.unpack(ratio) - shape = [s for s in self.shape[:-1]] + [self.shape[-1] * ratio] - new_elem_offset = self.elem_offset * ratio + layout = self.ty.layout.unpack(ratio) + shape = [s for s in self.ty.shape[:-1]] + [self.ty.shape[-1] * ratio] + new_elem_offset = self.ty.elem_offset * ratio return tvm.tirx.script.builder.decl_buffer( shape, cast_dtype, - self.data, - self.strides, + buffer_data(self), + self.ty.strides, new_elem_offset, None, - self.scope(), - self.data_alignment, - self.offset_factor, + self.ty.storage_scope, + self.ty.data_alignment, + self.ty.offset_factor, layout, ) else: @@ -389,15 +400,15 @@ def _infer_shape(shape): return tvm.tirx.script.builder.decl_buffer( shape, - self.dtype, - self.data, - self.strides, - self.elem_offset, + self.ty.dtype, + buffer_data(self), + self.ty.strides, + self.ty.elem_offset, None, - self.scope(), - self.data_alignment, - self.offset_factor, - self.layout if layout is None else layout, + self.ty.storage_scope, + self.ty.data_alignment, + self.ty.offset_factor, + self.ty.layout if layout is None else layout, ) def local(self, *shape, layout=None) -> "Buffer": @@ -422,22 +433,22 @@ def local(self, *shape, layout=None) -> "Buffer": The corresponding local buffer. """ if not shape: - local_layout = self.layout.storage() + local_layout = self.ty.layout.storage() total = functools.reduce( lambda x, y: x * y, [it.extent for it in local_layout.shard], 1 ) shape = (total,) return tvm.tirx.script.builder.decl_buffer( shape, - self.dtype, - self.data, - self.strides, - self.elem_offset, + self.ty.dtype, + buffer_data(self), + self.ty.strides, + self.ty.elem_offset, None, - self.scope(), - self.data_alignment, - self.offset_factor, - self.layout.storage() if layout is None else layout, + self.ty.storage_scope, + self.ty.data_alignment, + self.ty.offset_factor, + self.ty.layout.storage() if layout is None else layout, ) def permute(self, *dims) -> "Buffer": @@ -453,7 +464,7 @@ def permute(self, *dims) -> "Buffer": permuted : DeclBufferFrame The buffer with permuted dimensions. """ - new_shape = [self.shape[d] for d in dims] + new_shape = [self.ty.shape[d] for d in dims] # Permute *logical* dims, not the layout's fine-grained shard iters: a # tcgen05/atom layout maps several shard iters to each logical axis, so # group by the current shape first and permute whole groups. ``group`` @@ -461,23 +472,23 @@ def permute(self, *dims) -> "Buffer": # plus seps over *that* layout — permute the regrouped one, not # ``self.layout``. For a simple layout (one shard iter per axis) this # reduces to ``permute_dims(dims)``. - grouped, seps = self.layout.group(list(self.shape)) + grouped, seps = self.ty.layout.group(list(self.ty.shape)) new_layout = grouped.permute_by_groups(seps, list(dims)) return tvm.tirx.script.builder.decl_buffer( new_shape, - self.dtype, - self.data, - self.strides, - self.elem_offset, + self.ty.dtype, + buffer_data(self), + self.ty.strides, + self.ty.elem_offset, None, - self.scope(), - self.data_alignment, - self.offset_factor, + self.ty.storage_scope, + self.ty.data_alignment, + self.ty.offset_factor, new_layout, ) def __getitem__(self, indices): - if not is_buffer(self): + if not is_buffer_var(self): return _ORIGINAL_VAR_GETITEM(self, indices) from ..arith import Analyzer # pylint: disable=import-outside-toplevel @@ -490,14 +501,14 @@ def __getitem__(self, indices): has_step = any( isinstance(i, slice) and (i.step is not None and i.step != 1) for i in indices ) - has_implicit_slice = len(indices) < len(self.shape) + has_implicit_slice = len(indices) < len(self.ty.shape) analyzer = Analyzer() if (has_slice and not has_step) or has_implicit_slice: region = [] for i, index in enumerate(indices): if isinstance(index, slice): start = 0 if index.start is None else index.start - stop = self.shape[i] if index.stop is None else index.stop + stop = self.ty.shape[i] if index.stop is None else index.stop region.append(Range.from_min_extent(start, analyzer.simplify(stop - start))) else: region.append( @@ -507,15 +518,15 @@ def __getitem__(self, indices): ) ) if has_implicit_slice: - for i in range(len(indices), len(self.shape)): - region.append(Range.from_min_extent(0, self.shape[i])) + for i in range(len(indices), len(self.ty.shape)): + region.append(Range.from_min_extent(0, self.ty.shape[i])) return BufferRegion(self, region) else: expr_indices = [] for i, index in enumerate(indices): if isinstance(index, slice): start = 0 if index.start is None else index.start - stop = self.shape[i] if index.stop is None else index.stop + stop = self.ty.shape[i] if index.stop is None else index.stop step = 1 if index.step is None else index.step # We should ensure the dtype of start is the same with that of step. if tvm.ir.is_prim_expr(start) and isinstance(step, int): @@ -557,36 +568,48 @@ def decl_buffer( if offset_factor != 0 and elem_offset is None: shape_ty = shape[0].ty if shape and tvm.ir.is_prim_expr(shape[0]) else "int32" elem_offset = Var(f"{name}_elem_offset", shape_ty) - if data is None: - # Bool is represented as uint1 in the IR, but stored as int8 - storage_type = dtype if isinstance(dtype, PrimType) else PrimType(dtype) - storage_type = PrimType("int8") if storage_type.dtype == "bool" else storage_type - data = Var(name, PointerType(storage_type, scope), span) - return _ffi_api.BufferVar( # type: ignore - data, + storage_scope = scope + if data is not None: + if not isinstance(data, tvm.ir.Var) or not isinstance(data.ty, PointerType): + raise TypeError("Buffer data must be a Var with PointerType") + if not isinstance(data.ty.element_type, PrimType): + raise TypeError("Buffer data must point to a primitive type") + storage_scope = data.ty.storage_scope + buffer_type = _ffi_api.BufferType( # type: ignore + storage_scope, dtype, shape, strides, elem_offset, - name, data_alignment, offset_factor, - span, layout, + (), + span, ) + return _ffi_api.BufferVar(name, buffer_type, span) # type: ignore -def _type_field(name): - def getter(self): - if not is_buffer(self): - raise AttributeError(f"{self.name} is not a Var with BufferType") - return getattr(self.ty, name) +def buffer_data(buffer): + """Project the physical pointer associated with a buffer variable.""" + + if not is_buffer_var(buffer): + raise TypeError("buffer_data expects a Var with BufferType") + return _ffi_api.BufferData(buffer) - return property(getter) +def buffer_data_pointer_type(buffer): + """Return the pointer type produced by :func:`buffer_data`.""" -# Buffer values intentionally retain runtime type key ``ir.Var``. Install the -# checked, type-directed convenience surface on that ordinary Python wrapper. + if not is_buffer_var(buffer): + raise TypeError("buffer_data_pointer_type expects a Var with BufferType") + return _ffi_api.BufferDataPointerType(buffer) + + +# Buffer values intentionally retain runtime type key ``ir.Var``. Importing +# ``tvm.tirx`` therefore augments ``tvm.ir.Var`` process-wide with the legacy +# buffer operation and metadata surface. Non-buffer Vars reject the metadata +# properties with AttributeError, preserving correct ``hasattr`` behavior. _ORIGINAL_VAR_GETITEM = tvm.ir.Var.__getitem__ for _name, _value in _BufferMethods.__dict__.items(): if _name.startswith("__") and _name != "__getitem__": @@ -594,9 +617,19 @@ def getter(self): if callable(_value) or isinstance(_value, property): setattr(tvm.ir.Var, _name, _value) + +def _buffer_type_field(name): + def getter(value): + if not is_buffer_var(value): + raise AttributeError(f"{name} is only available on a Var with BufferType") + return getattr(value.ty, name) + + return property(getter) + + +# Preserve Buffer's public metadata surface while keeping BufferType as the +# single source of truth. for _name in ( - "data_pointer_type", - "dtype", "shape", "strides", "elem_offset", @@ -605,18 +638,34 @@ def getter(self): "layout", "allocated_addr", ): - setattr(tvm.ir.Var, _name, _type_field(_name)) + setattr(tvm.ir.Var, _name, _buffer_type_field(_name)) + + +def _buffer_dtype_property(value): + if not is_buffer_var(value): + raise AttributeError("dtype is only available on a Var with BufferType") + # Preserve the pre-migration Python Buffer surface. BufferType stores a + # PrimType, while Python callers historically receive its runtime DataType. + return value.ty.dtype.dtype + + +tvm.ir.Var.dtype = property(_buffer_dtype_property) + + +# Keep the established ``A.data`` TVMScript surface as syntax sugar. Compiler +# and builder code calls ``buffer_data(A)`` directly. +def _buffer_data_property(value): + if not is_buffer_var(value): + raise AttributeError("data is only available on a Var with BufferType") + return buffer_data(value) + -tvm.ir.Var.data = property( - lambda self: _ffi_api.BufferData(self) - if is_buffer(self) - else (_ for _ in ()).throw(AttributeError(f"{self.name} is not a Var with BufferType")) -) -tvm.ir.Var.READ = _BufferMethods.READ -tvm.ir.Var.WRITE = _BufferMethods.WRITE +tvm.ir.Var.data = property(_buffer_data_property) # Source compatibility for annotations and imports only. There is no -# ``tirx.Buffer`` runtime object; constructors return ``tvm.ir.Var``. +# ``tirx.Buffer`` runtime object; constructors return ``tvm.ir.Var``. In +# particular, ``isinstance(value, Buffer)`` matches every Var. Runtime checks +# must use ``is_buffer_var(value)``. Buffer = tvm.ir.Var diff --git a/python/tvm/tirx/function.py b/python/tvm/tirx/function.py index 5fdc2c3792dd..bbacc0472fa1 100644 --- a/python/tvm/tirx/function.py +++ b/python/tvm/tirx/function.py @@ -31,7 +31,7 @@ from ..runtime._tensor import Tensor from . import _ffi_api -from .buffer import Buffer, is_buffer +from .buffer import Buffer, is_buffer_var from .expr import Expr, Var @@ -72,7 +72,7 @@ def __init__(self, params, body, ret_type=None, buffer_map=None, attrs=None, spa buffer_map = {} if buffer_map is None else buffer_map for x in params: x = tvm.runtime.convert(x) if not isinstance(x, Object) else x - if is_buffer(x): + if is_buffer_var(x): var = Var(x.name, ty="handle") param_list.append(var) buffer_map[var] = x diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py index 481b717c8f31..46b05c0bc622 100644 --- a/python/tvm/tirx/op.py +++ b/python/tvm/tirx/op.py @@ -30,7 +30,7 @@ from tvm.runtime import const from . import _ffi_api -from .buffer import Buffer, is_buffer +from .buffer import Buffer, buffer_data, is_buffer_var from .expr import BufferLoad, CommReducer, ExprOp, ExprWithOp, IntImm, Var tir = tirx # alias for backward compat with upstream tir.convert() calls @@ -79,27 +79,27 @@ def _pack_buffer(buf, span=None): """Build intrinsics that packs the buffer.""" shape = Call( "tirx.tvm_stack_make_shape", - buf.shape, + buf.ty.shape, span=span, ret_ty=PointerType(tvm.ir.PrimType("int64")), ) strides = ( Call( "tirx.tvm_stack_make_shape", - buf.strides, + buf.ty.strides, span=span, ret_ty=PointerType(tvm.ir.PrimType("int64")), ) - if buf.strides + if buf.ty.strides else 0 ) pack_args = [ - buf.data, + buffer_data(buf), shape, strides, - len(buf.shape), - const(0, dtype=buf.dtype), - buf.elem_offset, + len(buf.ty.shape), + const(0, dtype=buf.ty.dtype), + buf.ty.elem_offset, ] return Call(Op.get("tirx.tvm_stack_make_array"), pack_args, span=span, ret_ty="handle") @@ -129,7 +129,7 @@ def call_packed_lowered(*args, span=None): -------- te.extern : Create tensor with extern function call. """ - call_args = [_pack_buffer(x) if is_buffer(x) else x for x in args] + call_args = [_pack_buffer(x) if is_buffer_var(x) else x for x in args] return Call(Op.get("tirx.tvm_call_packed_lowered"), call_args, span=span, ret_ty="int32") @@ -155,7 +155,7 @@ def call_cpacked_lowered(*args, span=None): -------- te.extern : Create tensor with extern function call. """ - call_args = [_pack_buffer(x) if is_buffer(x) else x for x in args] + call_args = [_pack_buffer(x) if is_buffer_var(x) else x for x in args] return Call(Op.get("tirx.tvm_call_cpacked_lowered"), call_args, span=span, ret_ty="int32") @@ -186,7 +186,7 @@ def call_packed(*args, span=None): -------- te.extern : Create tensor with extern function call. """ - call_args = [_pack_buffer(x) if is_buffer(x) else x for x in args] + call_args = [_pack_buffer(x) if is_buffer_var(x) else x for x in args] return Call(Op.get("tirx.tvm_call_packed"), call_args, span=span, ret_ty="int32") @@ -213,7 +213,7 @@ def call_cpacked(*args, span=None): -------- te.extern : Create tensor with extern function call. """ - call_args = [_pack_buffer(x) if is_buffer(x) else x for x in args] + call_args = [_pack_buffer(x) if is_buffer_var(x) else x for x in args] return Call(Op.get("tirx.tvm_call_cpacked"), call_args, span=span, ret_ty="int32") @@ -674,10 +674,10 @@ def address_of(obj: Buffer | BufferLoad | Var, span: Span | None = None) -> Expr call : Expr The call expression. """ - if is_buffer(obj): - n_dim = len(obj.shape) + if is_buffer_var(obj): + n_dim = len(obj.ty.shape) buffer_load = BufferLoad(obj, [0] * n_dim) - return Call("tirx.address_of", [buffer_load], span=span, ret_ty=obj.data.ty) + return Call("tirx.address_of", [buffer_load], span=span, ret_ty=buffer_data(obj).ty) elif isinstance(obj, Var): if _is_tensormap_var(obj): return call_intrin("uint64", "tirx.address_of", obj, span=span) @@ -685,7 +685,7 @@ def address_of(obj: Buffer | BufferLoad | Var, span: Span | None = None) -> Expr raise TypeError(f"address_of expects a scalar or TensorMap Var, but got {obj.ty}") return Call("tirx.address_of", [obj], span=span, ret_ty=PointerType(obj.ty)) elif isinstance(obj, BufferLoad): - return Call("tirx.address_of", [obj], span=span, ret_ty=obj.buffer.data.ty) + return Call("tirx.address_of", [obj], span=span, ret_ty=buffer_data(obj.buffer).ty) else: raise ValueError(f"Invalid object type: {type(obj)}") @@ -1267,10 +1267,10 @@ def trace(args, trace_action="tvm.default_trace_action"): """ if not isinstance(args, list): raise Exception("tvm.tirx.trace consumes the args as list type") - call_args = [_pack_buffer(x) if is_buffer(x) else x for x in args] + call_args = [_pack_buffer(x) if is_buffer_var(x) else x for x in args] call_args.insert(0, tvm.tirx.StringImm(trace_action)) tracing_value = args[-1] - ret_ty = tracing_value.ty if isinstance(tracing_value, Expr) else tracing_value.dtype + ret_ty = tracing_value.ty if isinstance(tracing_value, Expr) else tracing_value.ty.dtype return tvm.ir.Call(Op.get("tirx.tvm_call_trace_packed"), call_args, ret_ty=ret_ty) diff --git a/python/tvm/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py index ee93a6404c15..5ad4d3a1f90b 100644 --- a/python/tvm/tirx/script/builder/ir.py +++ b/python/tvm/tirx/script/builder/ir.py @@ -44,7 +44,7 @@ # pylint: disable=unused-import from tvm.target.codegen import llvm_lookup_intrinsic_id -from tvm.tirx import Buffer, BufferRegion, Expr, IndexMap, is_buffer, type_annotation +from tvm.tirx import Buffer, BufferRegion, Expr, IndexMap, is_buffer_var, type_annotation from tvm.tirx import _ffi_api as _tirx_ffi_api from tvm.tirx import op as _tir_op from tvm.tirx.exec_scope import ExecScope, ScopeIdDef, Var @@ -516,7 +516,7 @@ def match_buffer( """ if shape is None: if isinstance(param, BufferRegion): - dtype = param.buffer.dtype + dtype = param.buffer.ty.dtype shape = [region.extent for region in param.region] else: raise ValueError("Shape must be specified when binding input param") @@ -1879,10 +1879,10 @@ def alloc_cast_frag(src, dtype): Buffer Fresh ``local`` frag, ``src.shape`` shaped, ``src.layout``, dtype-cast. """ - rows, cols = src.shape + rows, cols = src.ty.shape per_thread_elems = (rows * cols) // 128 flat = alloc_local((per_thread_elems,), dtype) - return flat.view(rows, cols, layout=src.layout) + return flat.view(rows, cols, layout=src.ty.layout) if TYPE_CHECKING: @@ -1986,7 +1986,7 @@ def __invert__(self): def alloc_scalar(dtype: str = "float32", scope: str = "global") -> BufferLoad: """Allocate a zero-dimensional buffer (scalar).""" buf = alloc_buffer(shape=(1,), dtype=dtype, scope=scope, layout=TileLayout(S[1])) - assert is_buffer(buf) + assert is_buffer_var(buf) scalar = buf[0] if _current_meta_construction_scope() is not None: return scalar @@ -2006,7 +2006,7 @@ def decl_scalar(dtype, data, scope, elem_offset=None, byte_offset=None) -> Buffe offset_factor=0, layout=TileLayout(S[1]), ) - assert is_buffer(buf) + assert is_buffer_var(buf) scalar = buf[0] if _current_meta_construction_scope() is not None: return scalar @@ -2042,7 +2042,7 @@ def _meta_resource_for_value(value: Any) -> Any | None: return value.scalar.buffer if isinstance(value, BufferLoad): return value.buffer - if is_buffer(value): + if is_buffer_var(value): return value return None @@ -2292,7 +2292,7 @@ def buffer_store( expr_indices.append(ramp(index.start, step, lanes)) else: expr_indices.append(index) - if isinstance(value, bool) and buffer.dtype == "bool": + if isinstance(value, bool) and buffer.ty.dtype == "bool": value = IntImm("bool", value) return _ffi_api.BufferStore( # type: ignore[attr-defined] # pylint: disable=no-member buffer, value, expr_indices, predicate @@ -2913,19 +2913,19 @@ class WebGPUNamespace: @staticmethod def subgroup_shuffle(var, lane): - if is_buffer(var): + if is_buffer_var(var): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.webgpu.subgroup_shuffle", var, lane) @staticmethod def subgroup_shuffle_up(var, delta): - if is_buffer(var): + if is_buffer_var(var): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.webgpu.subgroup_shuffle_up", var, delta) @staticmethod def subgroup_shuffle_down(var, delta): - if is_buffer(var): + if is_buffer_var(var): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.webgpu.subgroup_shuffle_down", var, delta) diff --git a/python/tvm/tirx/script/builder/tirx.py b/python/tvm/tirx/script/builder/tirx.py index 9167918fca07..5160cfb006c2 100644 --- a/python/tvm/tirx/script/builder/tirx.py +++ b/python/tvm/tirx/script/builder/tirx.py @@ -21,7 +21,7 @@ import tvm.tirx.operator as tirx_op from tvm.ir import Op -from tvm.tirx import Buffer, BufferRegion, Expr, is_buffer +from tvm.tirx import Buffer, BufferRegion, Expr, buffer_data, is_buffer_var from tvm.tirx.exec_scope import _SCOPE_KIND_TO_NAME, ExecScope from tvm.tirx.expr import FloatImm from tvm.tirx.lang.alloc_pool import SMEMPool, TMEMPool, TMEMStages @@ -123,12 +123,12 @@ def __getattr__(self, name: str): def _is_buffer_or_region(x): - return is_buffer(x) or isinstance(x, BufferRegion) + return is_buffer_var(x) or isinstance(x, BufferRegion) def _to_region(buffer: BufferRegion | Buffer): - if is_buffer(buffer): - return buffer[[slice(None, None, None) for _ in range(len(buffer.shape))]] + if is_buffer_var(buffer): + return buffer[[slice(None, None, None) for _ in range(len(buffer.ty.shape))]] assert isinstance(buffer, BufferRegion) return buffer @@ -222,7 +222,7 @@ def sqrt( config = kwargs or {} dst = _to_region(dst) src = _to_region(src) - if bias is not None and is_buffer(bias): + if bias is not None and is_buffer_var(bias): bias = _to_region(bias) return f_insert( tirx_op.Sqrt( @@ -268,9 +268,9 @@ def add( workspace = {} config = kwargs or {} dst = _to_region(dst) - if is_buffer(src1): + if is_buffer_var(src1): src1 = _to_region(src1) - if is_buffer(src2): + if is_buffer_var(src2): src2 = _to_region(src2) return f_insert( tirx_op.Add( @@ -309,9 +309,9 @@ def sub( workspace = {} config = kwargs or {} dst = _to_region(dst) - if is_buffer(src1): + if is_buffer_var(src1): src1 = _to_region(src1) - if is_buffer(src2): + if is_buffer_var(src2): src2 = _to_region(src2) return f_insert( tirx_op.Sub( @@ -350,9 +350,9 @@ def mul( workspace = {} config = kwargs or {} dst = _to_region(dst) - if is_buffer(src1): + if is_buffer_var(src1): src1 = _to_region(src1) - if is_buffer(src2): + if is_buffer_var(src2): src2 = _to_region(src2) return f_insert( tirx_op.Mul( @@ -392,7 +392,7 @@ def fdiv( config = kwargs or {} dst = _to_region(dst) src1 = _to_region(src1) - if is_buffer(src2): + if is_buffer_var(src2): src2 = _to_region(src2) return f_insert( tirx_op.FDiv( @@ -436,9 +436,9 @@ def fma( config = kwargs or {} dst = _to_region(dst) src = _to_region(src) - if is_buffer(scale): + if is_buffer_var(scale): scale = _to_region(scale) - if is_buffer(bias): + if is_buffer_var(bias): bias = _to_region(bias) return f_insert( tirx_op.FMA( @@ -1010,9 +1010,9 @@ def maximum( workspace = {} config = kwargs or {} dst = _to_region(dst) - if is_buffer(src1): + if is_buffer_var(src1): src1 = _to_region(src1) - if is_buffer(src2): + if is_buffer_var(src2): src2 = _to_region(src2) return f_insert( tirx_op.Maximum( @@ -1051,9 +1051,9 @@ def minimum( workspace = {} config = kwargs or {} dst = _to_region(dst) - if is_buffer(src1): + if is_buffer_var(src1): src1 = _to_region(src1) - if is_buffer(src2): + if is_buffer_var(src2): src2 = _to_region(src2) return f_insert( tirx_op.Minimum( @@ -1105,7 +1105,7 @@ def exp( config = kwargs or {} dst = _to_region(dst) src = _to_region(src) - if bias is not None and is_buffer(bias): + if bias is not None and is_buffer_var(bias): bias = _to_region(bias) return f_insert( tirx_op.Exp( @@ -1164,7 +1164,7 @@ def exp2( config = kwargs or {} dst = _to_region(dst) src = _to_region(src) - if bias is not None and is_buffer(bias): + if bias is not None and is_buffer_var(bias): bias = _to_region(bias) return f_insert( tirx_op.Exp2( @@ -1250,9 +1250,9 @@ def binary_reduce( workspace = {} binary_output = _to_region(binary_output) reduce_output = _to_region(reduce_output) - if is_buffer(binary_input1): + if is_buffer_var(binary_input1): binary_input1 = _to_region(binary_input1) - if is_buffer(binary_input2): + if is_buffer_var(binary_input2): binary_input2 = _to_region(binary_input2) reduce_axes = _wrap_elem_in_tuple(reduce_axes) @@ -1334,7 +1334,7 @@ def unary_reduce( reduce_output = _to_region(reduce_output) unary_input = _to_region(unary_input) - if bias is not None and is_buffer(bias): + if bias is not None and is_buffer_var(bias): bias = _to_region(bias) reduce_axes = _wrap_elem_in_tuple(reduce_axes) @@ -1418,9 +1418,9 @@ def binary_chain( output = _to_region(output) data = _to_region(data) - if is_buffer(operand0): + if is_buffer_var(operand0): operand0 = _to_region(operand0) - if is_buffer(operand1): + if is_buffer_var(operand1): operand1 = _to_region(operand1) if isinstance(op0, str): @@ -1533,9 +1533,9 @@ def select( The predicate to evaluate. The callable should take the same number of arguments as the dimensions of the destination buffer. """ # noqa: E501 dst = _to_region(dst) - if is_buffer(true_value): + if is_buffer_var(true_value): true_value = _to_region(true_value) - if is_buffer(false_value): + if is_buffer_var(false_value): false_value = _to_region(false_value) if not isinstance(pred, Predicate): pred = Predicate(pred) @@ -1547,15 +1547,15 @@ def reshape(buffer: Buffer, shape: list[Expr]): # for example, if buffer.shape is (1024, 1024) and shape is (128, -1, 2), then the new shape will be (128, 4, 2) # noqa: E501 shape = list(shape) if -1 in shape and shape.count(-1) == 1: - size = functools.reduce(lambda x, y: x * y, buffer.shape) + size = functools.reduce(lambda x, y: x * y, buffer.ty.shape) n_size = functools.reduce(lambda x, y: x * y, [s for s in shape if s != -1], 1) shape[shape.index(-1)] = size // n_size else: assert functools.reduce(lambda x, y: x * y, shape) == functools.reduce( - lambda x, y: x * y, buffer.shape + lambda x, y: x * y, buffer.ty.shape ), ( "The shape of the buffer " - + str(buffer.shape) + + str(buffer.ty.shape) + " and the new shape " + str(shape) + " are not compatible" @@ -1563,15 +1563,15 @@ def reshape(buffer: Buffer, shape: list[Expr]): return decl_buffer( shape, - buffer.dtype, - buffer.data, - buffer.strides, - buffer.elem_offset, + buffer.ty.dtype, + buffer_data(buffer), + buffer.ty.strides, + buffer.ty.elem_offset, None, buffer.scope(), - buffer.data_alignment, - buffer.offset_factor, - buffer.layout, + buffer.ty.data_alignment, + buffer.ty.offset_factor, + buffer.ty.layout, ) @@ -1604,11 +1604,9 @@ def permute_layout( # Promote Buffer to BufferRegion covering the full extent, matching the # convention used by ``Tx.`` fallback registration. - from tvm.tirx import Buffer as _TBuffer - def _to_region(b): - if isinstance(b, _TBuffer): - slices = [slice(None) for _ in range(len(b.shape))] + if is_buffer_var(b): + slices = [slice(None) for _ in range(len(b.ty.shape))] return b[slices] return b diff --git a/python/tvm/tirx/script/builder/triton.py b/python/tvm/tirx/script/builder/triton.py index 95690b7b1c5d..e9038eaf3331 100644 --- a/python/tvm/tirx/script/builder/triton.py +++ b/python/tvm/tirx/script/builder/triton.py @@ -24,7 +24,7 @@ from triton.runtime.jit import type_canonicalisation_dict from tvm import tirx -from tvm.ir import PointerType, PrimType, Var, is_prim_expr +from tvm.ir import PointerType, PrimType, is_prim_expr from tvm.runtime import Module from tvm.topi.utils import get_const_int @@ -120,7 +120,6 @@ def _generate_triton_kernel( kernel_args.append(arg) continue if isinstance(arg.ty, PointerType): - assert isinstance(arg, Var) assert isinstance(arg.ty.element_type, PrimType) elem_type = arg.ty.element_type.dtype pointer_type = "*" + type_canonicalisation_dict[elem_type] diff --git a/python/tvm/tirx/script/parser/parser.py b/python/tvm/tirx/script/parser/parser.py index 8026f07b9efc..93d9a9ab6a9c 100644 --- a/python/tvm/tirx/script/parser/parser.py +++ b/python/tvm/tirx/script/parser/parser.py @@ -29,7 +29,7 @@ from tvm.script.ir_builder.base import IRBuilderFrame as Frame from tvm.script.parser._core import Parser, dispatch, doc from tvm.script.parser.core.doc import from_doc -from tvm.tirx import Buffer, IterVar, Layout, is_buffer +from tvm.tirx import Buffer, IterVar, Layout, buffer_data, is_buffer_var from tvm.tirx.script import builder as T from tvm.tirx.script.builder.ir import name_meta_class_value from tvm.tirx.stmt import BufferRegion @@ -50,44 +50,44 @@ def slice_buffer_from_region(br: BufferRegion) -> Buffer: region = br.region new_shape = [r.extent for r in region] sliced_layout = None - if buf.layout is not None: + if buf.ty.layout is not None: range_pairs = [(r.min, r.min + r.extent) for r in region] - sliced_layout = buf.layout.slice(list(buf.shape), range_pairs) + sliced_layout = buf.ty.layout.slice(list(buf.ty.shape), range_pairs) if sliced_layout is not None: return T.decl_buffer( new_shape, - buf.dtype, - buf.data, - buf.strides, - buf.elem_offset, + buf.ty.dtype, + buffer_data(buf), + buf.ty.strides, + buf.ty.elem_offset, None, - buf.scope(), - buf.data_alignment, - buf.offset_factor, + buf.ty.storage_scope, + buf.ty.data_alignment, + buf.ty.offset_factor, sliced_layout, ) # Fallback: compute elem_offset for default/no layout strides = [] - for i in range(len(buf.shape)): + for i in range(len(buf.ty.shape)): stride = functools.reduce( - lambda x, y: x * y, buf.shape[i + 1 :], tvm.tirx.const(1, "int32") + lambda x, y: x * y, buf.ty.shape[i + 1 :], tvm.tirx.const(1, "int32") ) strides.append(stride) offset = tvm.tirx.const(0, "int32") for i, r in enumerate(region): offset = offset + r.min * strides[i] - new_elem_offset = buf.elem_offset + offset + new_elem_offset = buf.ty.elem_offset + offset return T.decl_buffer( new_shape, - buf.dtype, - buf.data, - buf.strides, + buf.ty.dtype, + buffer_data(buf), + buf.ty.strides, new_elem_offset, None, - buf.scope(), - buf.data_alignment, - buf.offset_factor, - buf.layout, + buf.ty.storage_scope, + buf.ty.data_alignment, + buf.ty.offset_factor, + buf.ty.layout, ) @@ -211,8 +211,10 @@ def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) - res = value.__enter__() IRBuilder.name(var_name, res) return res - elif is_buffer(value) or isinstance(value, IterVar | Layout) or ( - isinstance(value, tvm.ir.Var) and not self.var_table.exist(value) + elif ( + is_buffer_var(value) + or isinstance(value, IterVar | Layout) + or (isinstance(value, tvm.ir.Var) and not self.var_table.exist(value)) ): IRBuilder.name(var_name, value) return value @@ -407,12 +409,12 @@ def visit_assign(self: Parser, node: doc.Assign) -> None: # that genuine errors (e.g. wrong shape, bad store) are not swallowed. # Only TypeError from FFI type mismatch (e.g. rhs is a meta_var, not # a Expr or auto-convertible scalar) triggers fallthrough. - if isinstance(lhs_value, T.scalar_wrapper | T.BufferLoad) or is_buffer(lhs_value): + if isinstance(lhs_value, T.scalar_wrapper | T.BufferLoad) or is_buffer_var(lhs_value): if isinstance(lhs_value, T.scalar_wrapper): buffer = lhs_value.scalar.buffer else: buffer = lhs_value.buffer if isinstance(lhs_value, T.BufferLoad) else lhs_value - if len(buffer.shape) == 1 and bool(buffer.shape[0] == 1): + if len(buffer.ty.shape) == 1 and bool(buffer.ty.shape[0] == 1): # only 1-dim buffer with shape (1,) can be assigned directly # Note that shape can be a Expr, so we only judge by # bool(shape[0] == 1) rather than int(shape[0]) == 1. @@ -487,12 +489,12 @@ def visit_aug_assign(self: Parser, node: doc.AugAssign) -> None: lhs_value = self.eval_expr(lhs_copy) except Exception: # pylint: disable=broad-except pass - if isinstance(lhs_value, T.scalar_wrapper | T.BufferLoad) or is_buffer(lhs_value): + if isinstance(lhs_value, T.scalar_wrapper | T.BufferLoad) or is_buffer_var(lhs_value): if isinstance(lhs_value, T.scalar_wrapper): buffer = lhs_value.scalar.buffer else: buffer = lhs_value.buffer if isinstance(lhs_value, T.BufferLoad) else lhs_value - if len(buffer.shape) == 1 and bool(buffer.shape[0] == 1): + if len(buffer.ty.shape) == 1 and bool(buffer.ty.shape[0] == 1): try: T.buffer_store(buffer, rhs, [0]) return @@ -760,7 +762,7 @@ def visit_expr_stmt(self: Parser, node: doc.Expr) -> None: pass elif isinstance(res, tvm.tirx.stmt.BufferStore): T.buffer_store(res.buffer, res.value, res.indices, res.predicate) - elif is_buffer(res): + elif is_buffer_var(res): # ``T.match_buffer(...)`` used as a bare statement (no LHS) — the # buffer object is discarded; the underlying side effect (the # match_buffer node) has already been emitted into the frame. diff --git a/python/tvm/tirx/script/tile.py b/python/tvm/tirx/script/tile.py index 42fe3914bedc..49ff2c0315dd 100644 --- a/python/tvm/tirx/script/tile.py +++ b/python/tvm/tirx/script/tile.py @@ -18,12 +18,10 @@ import functools -from tvm.tirx import Buffer, BufferRegion +from tvm.tirx import BufferRegion, is_buffer_var from .builder import tirx as _builder -_TILE_ARG_TYPES = (Buffer, BufferRegion) - def _get_arg(args, kwargs, index, name): if len(args) > index: @@ -32,7 +30,7 @@ def _get_arg(args, kwargs, index, name): def _require_buffer_arg(op_name, arg_name, value): - if not isinstance(value, _TILE_ARG_TYPES): + if not (is_buffer_var(value) or isinstance(value, BufferRegion)): raise TypeError( f"Tx.{op_name} is tile-only and expects `{arg_name}` to be a Buffer " f"or BufferRegion; use T.{op_name} for expression/builtin calls" diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index c723964d747f..69fcf2c7542d 100644 --- a/python/tvm/tirx/stmt.py +++ b/python/tvm/tirx/stmt.py @@ -430,12 +430,15 @@ class DeclBuffer(Stmt): buffer: Buffer The buffer being declared. + data: Expr + The physical data expression bound to the buffer view. + span: Optional[Span] The location of this DeclBuffer in the source code. """ buffer: Buffer - data: Expr | None + data: Expr span: Span | None def __init__(self, buffer: Buffer, *args, **kwargs) -> None: @@ -481,6 +484,8 @@ def __init__(self, buffer: Buffer, *args, **kwargs) -> None: raise TypeError("DeclBuffer span specified by both args and kwargs") span = kw_span if kw_span is not None else span + if data is None: + raise TypeError("DeclBuffer requires a physical data binding") self.__init_handle_by_constructor__(_ffi_api.DeclBuffer, buffer, data, span) # Legacy compatibility. Body is carried on python side only. if body is not None: diff --git a/python/tvm/tirx/stmt_functor.py b/python/tvm/tirx/stmt_functor.py index 2f85a0d7ae3a..db418f18b6ca 100644 --- a/python/tvm/tirx/stmt_functor.py +++ b/python/tvm/tirx/stmt_functor.py @@ -539,7 +539,7 @@ def visit_decl_buffer_(self, op): body = self.visit_stmt(op.body) if body is op.body: return op - return tvm.tirx.DeclBuffer(op.buffer, body, op.span) + return tvm.tirx.DeclBuffer(op.buffer, body, op.span, data=op.data) return op def visit_buffer_store_(self, op): diff --git a/python/tvm/tirx/transform/common.py b/python/tvm/tirx/transform/common.py index 513f1417f0b9..d0deabf3961f 100644 --- a/python/tvm/tirx/transform/common.py +++ b/python/tvm/tirx/transform/common.py @@ -60,44 +60,44 @@ def mutate_buffer(self, buffer: Buffer): # unrelated buffers can be spuriously cloned and introduce alias buffers. prev_mutated = self.buffer_attr_var_mutated self.buffer_attr_var_mutated = False - new_shape = [self.visit_expr(expr) for expr in buffer.shape] - new_strides = [self.visit_expr(expr) for expr in buffer.strides] + new_shape = [self.visit_expr(expr) for expr in buffer.ty.shape] + new_strides = [self.visit_expr(expr) for expr in buffer.ty.strides] new_elem_offset = ( - self.visit_expr(buffer.elem_offset) if buffer.elem_offset is not None else None + self.visit_expr(buffer.ty.elem_offset) if buffer.ty.elem_offset is not None else None ) - if isinstance(buffer.layout, TileLayout): + if isinstance(buffer.ty.layout, TileLayout): new_shard = [] new_replicate = [] - for iter in buffer.layout.shard: + for iter in buffer.ty.layout.shard: new_iter = Iter( self.visit_expr(iter.extent), self.visit_expr(iter.stride), iter.axis ) new_shard.append(new_iter) - for iter in buffer.layout.replica: + for iter in buffer.ty.layout.replica: new_iter = Iter( self.visit_expr(iter.extent), self.visit_expr(iter.stride), iter.axis ) new_replicate.append(new_iter) new_layout = TileLayout.from_iters( - new_shard, new_replicate, offset=buffer.layout.offset + new_shard, new_replicate, offset=buffer.ty.layout.offset ) else: - new_layout = buffer.layout - new_allocated_addr = [self.visit_expr(expr) for expr in buffer.allocated_addr] + new_layout = buffer.ty.layout + new_allocated_addr = [self.visit_expr(expr) for expr in buffer.ty.allocated_addr] buffer_attr_mutated = self.buffer_attr_var_mutated self.buffer_attr_var_mutated = prev_mutated or buffer_attr_mutated if not buffer_attr_mutated: return None new_buffer = decl_buffer( new_shape, - buffer.dtype, + buffer.ty.dtype, buffer.name, - Var(buffer.name, buffer.data_pointer_type), + None, new_strides, new_elem_offset, buffer.scope(), - buffer.data_alignment, - buffer.offset_factor, + buffer.ty.data_alignment, + buffer.ty.offset_factor, layout=new_layout, ) if new_allocated_addr: @@ -107,11 +107,10 @@ def mutate_buffer(self, buffer: Buffer): return new_buffer def visit_var_(self, op: Var): - op = super().visit_var_(op) if op in self.var_map: self.buffer_attr_var_mutated = True return self.var_map[op] - return op + return super().visit_var_(op) def visit_buffer_load_(self, op: BufferLoad): new_buffer = self.mutate_buffer(op.buffer) @@ -136,9 +135,13 @@ def visit_buffer_region_(self, op: BufferRegion): def visit_decl_buffer_(self, op: DeclBuffer): new_buffer = self.mutate_buffer(op.buffer) - op = super().visit_decl_buffer_(op) - if new_buffer is not None: - return DeclBuffer(new_buffer, data=op.data, span=op.span) + data = self.visit_expr(op.data) + if new_buffer is not None or data is not op.data: + return DeclBuffer( + new_buffer if new_buffer is not None else op.buffer, + data=data, + span=op.span, + ) return op def visit_array_prim_expr_(self, op: list[Expr]): diff --git a/python/tvm/tirx/transform/transform.py b/python/tvm/tirx/transform/transform.py index 7eb26b722596..907324454916 100644 --- a/python/tvm/tirx/transform/transform.py +++ b/python/tvm/tirx/transform/transform.py @@ -253,15 +253,15 @@ def MakePackedAPI(): `buffer_map`, using it to generate arguments that implement the packed based TVM FFI API. - For static shapes, the `BufferNode::shape`, `BufferNode::strides`, - and `BufferNode::elem_offset` member variables are used to + For static shapes, the `BufferType::shape`, `BufferType::strides`, + and `BufferType::elem_offset` fields are used to generate runtime checks on the corresponding member variables in the user-provided `DLTensor*` or `tvm.runtime.tensor` argument. (e.g. A PrimFunc that accepts a buffer of shape `[16,32]` validates that the `DLTensor::shape` array is `[16,32]`.) - For dynamic Buffers, in which one or more of these `BufferNode` member - variables use `tirx.Var` that are not defined by other PrimFunc + For dynamic Buffers, in which one or more of these `BufferType` fields + use `tirx.Var` that are not defined by other PrimFunc parameters, these are instead used to define the variables based on the corresponding `DLTensor` members. (e.g. A PrimFunc that accepts a buffer of shape `[tirx.Var("n", "int64"), tirx.Var("m", "int64")]`, diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc index ba48f736ef03..0245de7a4297 100644 --- a/src/backend/cuda/codegen/codegen_cuda.cc +++ b/src/backend/cuda/codegen/codegen_cuda.cc @@ -1432,6 +1432,12 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op, std::ostream& os) { Expr arg = op->args[0]; const auto* var_node = arg.as(); + if (const auto* call = arg.as(); + call && call->op.same_as(tirx::builtin::buffer_data()) && call->args.size() == 1) { + var_node = call->args[0].as(); + TVM_FFI_ICHECK(var_node && var_node->ty.as()) + << "print_buffer expects buffer_data to project a BufferVar"; + } PrimType dtype_ty = op->ty.as_or_throw(); bool is_string = op->args[2].as()->value; bool is_scalar = op->args[3].as()->value; @@ -1632,7 +1638,7 @@ void CodeGenCUDA::VisitStmt_(const AttrStmtNode* op) { void CodeGenCUDA::VisitStmt_(const AllocBufferNode* op) { TVM_FFI_ICHECK(op->buffer.defined()); - std::string vid = AllocVarID(op->buffer.get()); + std::string vid = AllocVarID(op->buffer.get(), op->buffer.name() + "_ptr"); this->PrintIndent(); std::string scope = op->buffer.scope(); diff --git a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc index 2bd1a505c94b..07e6f384d743 100644 --- a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc +++ b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc @@ -107,7 +107,8 @@ class CodeGenHexagon final : public CodeGenCPU { bool IsQHLFunction(const std::string& func); - llvm::Value* VectorLookupLoad(BufferVar buffer, PrimType buffer_type, ffi::Array indices); + llvm::Value* VectorLookupLoad(BufferVar buffer, PrimType buffer_type, + ffi::Array indices); llvm::Value* Intrinsic(llvm::Intrinsic::ID, llvm::ArrayRef args); std::vector fqhl_list_ = { "tvm_vect_qhmath_hvx_cos_ahf", "tvm_vect_qhmath_hvx_tanh_ahf", diff --git a/src/backend/hexagon/runtime/hexagon_buffer.h b/src/backend/hexagon/runtime/hexagon_buffer.h index 0e578fccf477..d7384f2da4b9 100644 --- a/src/backend/hexagon/runtime/hexagon_buffer.h +++ b/src/backend/hexagon/runtime/hexagon_buffer.h @@ -85,7 +85,7 @@ class HexagonBuffer { /*! \brief Return data pointer into the buffer * * The returned pointer is intended for use as the runtime value - * corresponding to the `Var BufferNode::data` of a buffer. The + * corresponding to the runtime value of `buffer_data(BufferVar)`. The * return type depends on the dimensionality of the buffer being * accessed, and must be compatible with the usage defined in * `CodeGenHexagon::CreateBufferPtr`. diff --git a/src/backend/trn/codegen/codegen_trn.cc b/src/backend/trn/codegen/codegen_trn.cc index 466371ed9965..fefccbf28c0d 100644 --- a/src/backend/trn/codegen/codegen_trn.cc +++ b/src/backend/trn/codegen/codegen_trn.cc @@ -79,7 +79,7 @@ void CodeGenTrainium::AddFunction(const GlobalVar& gvar, const PrimFunc& func) { // We can switch to follow the flow with inter-function call process // after the Trainium function declaration is properly printed. // In Trainium, for PrimFuncs with signature - // def func(A: BufferVar, B: BufferVar, x: int, y: float) -> None + // def func(A: Buffer, B: Buffer, x: int, y: float) -> None // where there are trailing pod parameters, the codegen emits a struct // struct func_params{ x: int; y: float; } // for the function. In the flow of inter-function call process, @@ -103,7 +103,7 @@ void CodeGenTrainium::AddFunction(const GlobalVar& gvar, const PrimFunc& func) { // Function header. this->stream << "def " << static_cast(global_symbol.value()) << "("; - // BufferVar arguments + // Buffer arguments auto num_inputs = func->GetAttr(tvm::attr::kNumInputs); TVM_FFI_ICHECK(num_inputs.has_value()); std::vector output_vids; diff --git a/src/backend/trn/transform/lower_trainium_layout.cc b/src/backend/trn/transform/lower_trainium_layout.cc index 977865762a6e..2241d22f5911 100644 --- a/src/backend/trn/transform/lower_trainium_layout.cc +++ b/src/backend/trn/transform/lower_trainium_layout.cc @@ -115,13 +115,12 @@ class TrainiumLayoutApplier : public arith::IRMutatorWithAnalyzer { } Stmt VisitStmt_(const DeclBufferNode* op) final { + Expr data = VisitExpr(op->data); auto buffer = GetFlattenedBuffer(op->buffer); - if (buffer.same_as(op->buffer)) { + if (buffer.same_as(op->buffer) && data.same_as(op->data)) { return ffi::GetRef(op); } - auto n = CopyOnWrite(op); - n->buffer = buffer; - return Stmt(n); + return DeclBuffer(buffer, std::move(data), op->span); } BufferVar GetFlattenedBuffer(BufferVar buf, bool is_alloc = false) { @@ -177,7 +176,6 @@ class TrainiumLayoutApplier : public arith::IRMutatorWithAnalyzer { } if (flattened->dtype->dtype == DLDataType{kDLBool, 8, 1}) { type->dtype = PrimType::Int(8); - type->data_pointer_type = PointerType(PrimType::Int(8), flattened.scope()); } for (size_t i = 0; i < flattened->shape.size(); ++i) { type->shape.Set(i, analyzer_->canonical_simplify(flattened->shape[i])); @@ -300,17 +298,17 @@ class TrainiumBufferOffsetRemover : public StmtExprMutator { Stmt VisitStmt_(const DeclBufferNode* op) { auto buffer = op->buffer; auto elem_offset = this->VisitPrimExpr(buffer->elem_offset); - if (elem_offset.same_as(buffer->elem_offset)) { - return StmtExprMutator::VisitStmt_(op); - } else { + Expr data = VisitExpr(op->data); + if (!elem_offset.same_as(buffer->elem_offset)) { auto type = CopyBufferType(buffer); type->elem_offset = std::move(elem_offset); buffer = RebuildBufferVar(buffer, std::move(type)); buffer_remap_[op->buffer] = buffer; - auto n = CopyOnWrite(op); - n->buffer = buffer; - return Stmt(n); } + if (buffer.same_as(op->buffer) && data.same_as(op->data)) { + return ffi::GetRef(op); + } + return DeclBuffer(buffer, std::move(data), op->span); } using StmtExprMutator::VisitExpr_; diff --git a/src/relax/analysis/layout_transformation.cc b/src/relax/analysis/layout_transformation.cc index 6b48cd41eb37..d3aef0ee1829 100644 --- a/src/relax/analysis/layout_transformation.cc +++ b/src/relax/analysis/layout_transformation.cc @@ -273,7 +273,7 @@ static ffi::Optional InferLayoutTransformation(const SpatialLayout& sr }); if (depends_on_initial_indices) { LOG(WARNING) - << "[LayoutInference] BufferVar access is dependent on both defined and undefined vars"; + << "[LayoutInference] Buffer access is dependent on both defined and undefined vars"; return {}; } // It is ok to erase this final index expression as it only depends on undefined vars. @@ -411,7 +411,7 @@ class BlockAnalyzer : public StmtExprVisitor { IndexMap read_transformation = maybe_read_transformation.value(); if (buffer_transformation_cache_.count(r->buffer) != 0) { if (!AreIdenticalTransforms(read_transformation, buffer_transformation_cache_[r->buffer])) - LOG(WARNING) << "[LayoutInference] BufferVar: " << r->buffer + LOG(WARNING) << "[LayoutInference] Buffer: " << r->buffer << " has conflicting transform proposals -- (preferred) " << buffer_transformation_cache_[r->buffer] << " vs. " << read_transformation; continue; @@ -526,7 +526,9 @@ class BlockAnalyzer : public StmtExprVisitor { public: bool CanBeTransformed() { return can_transform_block_; } IndexMap GetSBlockTransformation() { return block_transformation_; } - ffi::Map GetReadBufferTransformations() { return read_buffer_transformations_; } + ffi::Map GetReadBufferTransformations() { + return read_buffer_transformations_; + } private: bool can_transform_block_; diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc index 5dc990bd1df9..9290f0056490 100644 --- a/src/relax/analysis/tir_op_pattern_kind.cc +++ b/src/relax/analysis/tir_op_pattern_kind.cc @@ -436,7 +436,8 @@ bool HasReshapePattern(const PrimFunc& func) { // This check requires at least one of the src/dst side is a trivial buffer // access (e.g., buf[ax0, ax1, ax2]). - auto f_calc_flattened_idx = [&](const BufferVar& buffer, const ffi::Array& indices) { + auto f_calc_flattened_idx = [&](const BufferVar& buffer, + const ffi::Array& indices) { TVM_FFI_ICHECK_EQ(indices.size(), buffer->shape.size()); int ndim = indices.size(); PrimExpr idx = 0; diff --git a/src/relax/backend/adreno/annotate_custom_storage.cc b/src/relax/backend/adreno/annotate_custom_storage.cc index 19e02ed573fd..237fe8ca6016 100644 --- a/src/relax/backend/adreno/annotate_custom_storage.cc +++ b/src/relax/backend/adreno/annotate_custom_storage.cc @@ -231,7 +231,7 @@ * - Fusion * - FoldVDeviceScopeChange: There existed some ToVDevice copies from texture to buffer * This pass removes the copes and updates producer scope to global. - * - SpecializePrimFuncBasedOnCallSite: Finally we updates the BufferVar Var maps according to + * - SpecializePrimFuncBasedOnCallSite: Finally we update the buffer maps according to * VDevice scopes. * */ diff --git a/src/relax/backend/vm/vm_shape_lower.cc b/src/relax/backend/vm/vm_shape_lower.cc index f671c1290440..cd4b1e952c21 100644 --- a/src/relax/backend/vm/vm_shape_lower.cc +++ b/src/relax/backend/vm/vm_shape_lower.cc @@ -263,7 +263,7 @@ class PrimExprSlotCollector : public ExprVisitor, public TypeVisitor { * \code * * @T.prim_func - * def shape_func(H: T.BufferVar([3], "int64")): + * def shape_func(H: T.Buffer([3], "int64")): * H[1] = H[2] + 1 * * \endcode 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 cac73cd97cc7..7f9628398b43 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 @@ -45,7 +45,8 @@ class DistBufferReplacer : public StmtExprMutator { } private: - explicit DistBufferReplacer(ffi::Map buffer_map) : buffer_map_(buffer_map) {} + explicit DistBufferReplacer(ffi::Map buffer_map) + : buffer_map_(buffer_map) {} Stmt VisitStmt_(const BufferStoreNode* _store) final { BufferStore store = StmtExprMutator::VisitStmt_(_store).as_or_throw(); @@ -270,9 +271,9 @@ class DistributedBufferCompactor : StmtExprMutator { shape.push_back(buffer->shape[i]); } } - BufferType new_type(buffer->data_pointer_type, buffer->dtype, std::move(shape), - buffer->strides, buffer->elem_offset, buffer->data_alignment, - buffer->offset_factor, buffer->layout, buffer->allocated_addr); + BufferType new_type(buffer->storage_scope, buffer->dtype, std::move(shape), buffer->strides, + buffer->elem_offset, buffer->data_alignment, buffer->offset_factor, + buffer->layout, buffer->allocated_addr); return BufferVar(buffer.name(), std::move(new_type), buffer.span()); } diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc index e6f31492a4bf..80de3c8f8eda 100644 --- a/src/relax/op/op.cc +++ b/src/relax/op/op.cc @@ -263,8 +263,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { * * For dynamic shapes, it is not always possible to infer the output * of a TIR PrimFunc from its inputs. For example, a PrimFunc that - * accepts input buffer `T.BufferVar([16], "float32")` and output buffer - * `T.BufferVar([M, N], "float32")` infers the values of `M` and `N` from + * accepts input buffer `T.Buffer([16], "float32")` and output buffer + * `T.Buffer([M, N], "float32")` infers the values of `M` and `N` from * the shape of the provided output buffer. * * If the arguments provided are not compatible with the PrimFunc's diff --git a/src/relax/op/tensor/inspect.cc b/src/relax/op/tensor/inspect.cc index 5a15dab73518..d4aec119bb0c 100644 --- a/src/relax/op/tensor/inspect.cc +++ b/src/relax/op/tensor/inspect.cc @@ -278,7 +278,7 @@ Expr LegalizeTensorShape(const BlockBuilder& bb, const Call& call) { "Specified axis may not be larger than the tensor's dimensionality")}), tirx::DeclBuffer( shape_buffer, - tvm::Call(shape_buffer->data_pointer_type, tirx::builtin::tvm_struct_get(), + tvm::Call(shape_buffer.DataPointerType(), tirx::builtin::tvm_struct_get(), {dlpack_handle, IntImm::Int32(0), IntImm::Int32(tirx::builtin::TVMStructFieldKind::kDLTensorShape)})), tirx::Bind(extent, tirx::BufferLoad(shape_buffer, {axis.as_or_throw()})), diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc index 3a03c1c186bc..28f7178a3e13 100644 --- a/src/relax/transform/fuse_tir.cc +++ b/src/relax/transform/fuse_tir.cc @@ -190,10 +190,9 @@ class FuseTIRBufferSubstitutor : private StmtExprMutator { elem_offset.same_as(buffer->elem_offset)) { return buffer; } else { - BufferType new_type( - buffer->data_pointer_type, buffer->dtype, std::move(shape), - std::move(strides), std::move(elem_offset), buffer->data_alignment, - buffer->offset_factor, buffer->layout, buffer->allocated_addr); + BufferType new_type(buffer->storage_scope, buffer->dtype, std::move(shape), + std::move(strides), std::move(elem_offset), buffer->data_alignment, + buffer->offset_factor, buffer->layout, buffer->allocated_addr); BufferVar new_buffer(buffer.name(), std::move(new_type), buffer.span()); this->buffer_remap_.Set(buffer, new_buffer); return new_buffer; @@ -303,7 +302,7 @@ class FuseTIRBufferSubstitutor : private StmtExprMutator { ffi::Map var_remap_; ffi::Array UnionAccessRegion(const ffi::Array& regions) const { - // For now we only allow BufferVar access the same elements. + // For now we only allow buffers to access the same elements. // e.g. `[A[vi, vj], A[vi, vj]]` is a legal pattern but need to union to `A[vi, vj]` // However, `A[vi, vj], A[vi, vj + 1]` is not allow for now. // Note: the order of return region should remain the same as the first occurrence of the region @@ -610,7 +609,8 @@ class FusedTIRConstructor : public ExprVisitor { const ffi::Array& buffers = (*it).second; // map of input buffers to indices (helpful for detecting in-place inputs) - std::unordered_map buffer_to_idx; + std::unordered_map + buffer_to_idx; std::unordered_map input_to_idx; for (size_t i = 0; i < func_info_.params.size(); i++) { input_to_idx[func_info_.params[i]] = i; @@ -901,12 +901,10 @@ class FusedTIRConstructor : public ExprVisitor { return unique_name; }; // Update buffer with new symbolic shape according to the ty - tirx::BufferType new_type( - buffer->data_pointer_type, buffer->dtype, output_shapes[i], - buffer->strides, buffer->elem_offset, buffer->data_alignment, - buffer->offset_factor, buffer->layout, buffer->allocated_addr); - tirx::BufferVar new_buffer(unify_name_hints(), std::move(new_type), - buffer.span()); + tirx::BufferType new_type(buffer->storage_scope, buffer->dtype, output_shapes[i], + buffer->strides, buffer->elem_offset, buffer->data_alignment, + buffer->offset_factor, buffer->layout, buffer->allocated_addr); + tirx::BufferVar new_buffer(unify_name_hints(), std::move(new_type), buffer.span()); func_info_.alloc_buffers.push_back(new_buffer); output_buffers.push_back(new_buffer); diff --git a/src/relax/transform/specialize_primfunc_based_on_callsite.cc b/src/relax/transform/specialize_primfunc_based_on_callsite.cc index f6a99c9efd9a..8bd4706c510e 100644 --- a/src/relax/transform/specialize_primfunc_based_on_callsite.cc +++ b/src/relax/transform/specialize_primfunc_based_on_callsite.cc @@ -100,7 +100,7 @@ class SpecializeTIRCallArgs : ExprMutator { } const BufferVar& buffer = tirx::decl_buffer(GetShapeFromTensorType(tensor_ty), - tensor_ty->dtype.value(), name, scope); + tensor_ty->dtype.value(), name, scope); param_map.Set(pfunc->params[i], buffer); } ffi::String scope = "global"; @@ -133,17 +133,13 @@ class SpecializeTIRCallArgs : ExprMutator { } const BufferVar& buffer = tirx::decl_buffer(GetShapeFromTensorType(ty), ty->dtype.value(), - "ret_val_" + std::to_string(index), scope); + "ret_val_" + std::to_string(index), scope); param_map.Set(pfunc->params[args.size() + index], buffer); index++; } } auto new_pfunc = Specialize(pfunc, param_map); - for (const auto& [var, buffer] : new_pfunc->buffer_map) { - auto* ptr = buffer->data_pointer_type.as(); - TVM_FFI_ICHECK(ptr) << "BufferVar Var's type annotation must be of PointerType"; - } auto new_prim_func = WithAttr(new_pfunc, "scoped", static_cast(1)); updates_->Add(gv, new_prim_func); return call; diff --git a/src/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc index a385f0a7e03e..18618840d98c 100644 --- a/src/relax/transform/split_call_tir_by_pattern.cc +++ b/src/relax/transform/split_call_tir_by_pattern.cc @@ -336,8 +336,8 @@ class ForMatcher : public TensorizeComparator { for (size_t i = 0; i < lhs->shape.size(); ++i) { if (!VisitExpr(lhs->shape[i], rhs->shape[i])) return false; } - equal = DefEqual(lhs.var(), rhs.var()) && lhs->dtype == rhs->dtype && - lhs.scope() == rhs.scope(); + equal = + DefEqual(lhs.var(), rhs.var()) && lhs->dtype == rhs->dtype && lhs.scope() == rhs.scope(); if (equal) { rhs_buffer_map_[rhs] = lhs; } diff --git a/src/s_tir/analysis/calculate_allocated_memory.cc b/src/s_tir/analysis/calculate_allocated_memory.cc index 09fa73a5e1c2..e5600b70bbe7 100644 --- a/src/s_tir/analysis/calculate_allocated_memory.cc +++ b/src/s_tir/analysis/calculate_allocated_memory.cc @@ -41,7 +41,7 @@ using namespace tvm::tirx; std::string GetStorageScope(const Var& var) { auto* ptr = var->ty.as(); - TVM_FFI_ICHECK(ptr) << "BufferVar Var's type annotation must be of PointerType"; + TVM_FFI_ICHECK(ptr) << "Buffer Var's type annotation must be of PointerType"; return ptr->storage_scope; } diff --git a/src/s_tir/analysis/sblock_access_region_detector.cc b/src/s_tir/analysis/sblock_access_region_detector.cc index 516078f569ef..4eb3f32627f0 100644 --- a/src/s_tir/analysis/sblock_access_region_detector.cc +++ b/src/s_tir/analysis/sblock_access_region_detector.cc @@ -44,7 +44,8 @@ class BlockReadWriteDetector : public StmtExprVisitor { public: explicit BlockReadWriteDetector(const ffi::Map& buffer_var_map) : buffer_var_map_(buffer_var_map) { - for (const auto& [raw_pointer, buffer] : buffer_var_map) { + for (const auto& item : buffer_var_map) { + const BufferVar& buffer = item.second; buffer_var_map_.Set(buffer.var(), buffer); } } @@ -317,8 +318,8 @@ std::vector BlockReadWriteDetector::ConvertMatchedRegion( } void BlockReadWriteDetector::Update(std::vector* buffers, - std::vector>* regions, BufferVar buffer, - std::vector region) { + std::vector>* regions, + BufferVar buffer, std::vector region) { if (buffer_var_map_.find(buffer.var()) == buffer_var_map_.end()) return; // Handle match_buffer remap auto it = match_buffers_.find(buffer.get()); @@ -343,7 +344,8 @@ void BlockReadWriteDetector::Update(std::vector* buffers, } ffi::Array BlockReadWriteDetector::CollectRegions( - const std::vector& buffers, const std::vector>& regions, + const std::vector& buffers, + const std::vector>& regions, const std::unordered_set* excluded_buffers) { TVM_FFI_ICHECK_EQ(buffers.size(), regions.size()); ffi::Array res; 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 9a0002dd6a3c..b0de1d8f2115 100644 --- a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc +++ b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc @@ -34,7 +34,7 @@ namespace tvm { namespace tirx { /*! - * \brief Detect the lowest common ancestor(LCA) position of BufferVar access. + * \brief Detect the lowest common ancestor(LCA) position of buffer access. * \note * - Only consider BlockNode and ForNode to be the LCA nodes. * - In the LCA locator, we are aware of the buffer scope and CUDA hierarchy so that any buffer in diff --git a/src/s_tir/backend/adreno/inject_texture_alloc.cc b/src/s_tir/backend/adreno/inject_texture_alloc.cc index 2876df779915..d60589e9ef30 100644 --- a/src/s_tir/backend/adreno/inject_texture_alloc.cc +++ b/src/s_tir/backend/adreno/inject_texture_alloc.cc @@ -83,8 +83,7 @@ class TextureAllocInjector : public arith::IRMutatorWithAnalyzer { {texture.width, texture.height, texture.depth})); args.push_back(IntImm::Int64(channel_size)); stmt = DeclBuffer( - op->buffer, - Call(op->buffer->data_pointer_type, builtin::nd_mem_alloc_with_scope(), args)); + op->buffer, Call(op->buffer.DataPointerType(), builtin::nd_mem_alloc_with_scope(), args)); } return stmt; } @@ -92,7 +91,7 @@ class TextureAllocInjector : public arith::IRMutatorWithAnalyzer { protected: std::string GetStorageScope(const Var& buffer_var) { auto* ptr = buffer_var->ty.as(); - TVM_FFI_ICHECK(ptr) << "BufferVar Var's type annotation must be of PointerType"; + TVM_FFI_ICHECK(ptr) << "Buffer Var's type annotation must be of PointerType"; return ptr->storage_scope; } }; diff --git a/src/s_tir/backend/adreno/texture_flatten.cc b/src/s_tir/backend/adreno/texture_flatten.cc index 1d222d5b6dd4..5d14e314bdd2 100644 --- a/src/s_tir/backend/adreno/texture_flatten.cc +++ b/src/s_tir/backend/adreno/texture_flatten.cc @@ -71,11 +71,7 @@ class TextureLoweringBase : public StmtExprMutator { } protected: - std::string GetStorageScope(const BufferVar& buffer) { - auto* ptr = buffer->data_pointer_type.as(); - TVM_FFI_ICHECK(ptr) << "BufferVar Var's type annotation must be of PointerType"; - return ptr->storage_scope; - } + std::string GetStorageScope(const BufferVar& buffer) { return buffer->storage_scope; } // Set of all external input and output buffers std::unordered_set extern_buf_; 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 73bbc7eac617..a8578f6567b7 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 @@ -49,7 +49,9 @@ using ForVec = std::vector; * \tparam V The value type */ template -using ForBufferMap = std::unordered_map>; +using ForBufferMap = + std::unordered_map>; /*! \brief Given x, compute log2(|x| + 1) */ inline double slog(double x) { return x >= 0 ? std::log2(x + 1) : std::log2(-x + 1); } @@ -63,10 +65,10 @@ namespace utils { * \return The shape of the buffer */ std::vector GetBufferShape(const BufferVar& buffer, arith::AnalyzerObj* analyzer) { - int ndim = GetBufferVar(buffer)->shape.size(); + int ndim = buffer->shape.size(); std::vector result; result.reserve(ndim); - for (const PrimExpr& i : GetBufferVar(buffer)->shape) { + for (const PrimExpr& i : buffer->shape) { if (const IntImmNode* int_imm = i.as()) { result.push_back(int_imm->value); continue; @@ -665,14 +667,15 @@ struct Feature { struct SubFeature { /*! \brief The buffer this feature is for */ - const VarNode* buffer = nullptr; + BufferVar buffer; /*! \brief The access type of the buffer */ AccessType access_type = AccessType::kUnknownRW; /*! \brief A list of multi-dimensonal indices used to access the buffer */ std::vector multi_indices = {}; // Access information /*! \brief loop_accessed_numel[i][...] means the number of elements accessed by loops[i] */ - std::vector> loop_accessed_numel = {}; + std::vector> + loop_accessed_numel = {}; /*! \brief The shape of the data access */ IntVec access_shape; /*! \brief The bytes that are continuously accessed */ @@ -750,9 +753,9 @@ struct Feature { void SetFeature(const LoopNest& loop_nest, int64_t cache_line_bytes); - explicit SubFeature(const VarNode* buffer, AccessType access_type, + explicit SubFeature(BufferVar buffer, AccessType access_type, std::vector multi_indices, int n_loops) - : buffer(buffer), + : buffer(std::move(buffer)), access_type(access_type), multi_indices(multi_indices), loop_accessed_numel(n_loops) {} @@ -788,15 +791,15 @@ void Feature::Init(const BufferStoreNode* store, int n_loops) { AccessType access_type = AccessType::kUnknownRW; std::vector multi_indices; }; - std::unordered_map buffer_info; + std::unordered_map buffer_info; { - Info& info = buffer_info[store->buffer.get()]; + Info& info = buffer_info[store->buffer]; info.access_type = AccessType::kWrite; info.multi_indices.push_back({store->indices.begin(), store->indices.end()}); } PostOrderVisit(store->value, [&buffer_info](const ffi::ObjectRef& obj) -> void { if (const BufferLoadNode* load = obj.as()) { - const VarNode* buffer = load->buffer.get(); + BufferVar buffer = load->buffer; Info& info = buffer_info[buffer]; switch (info.access_type) { case AccessType::kRead: @@ -838,7 +841,7 @@ void Feature::SetRegion(const LoopNest& loop_nest, IntVec* for_touched_bytes, if (n_loops == 0) { // In this case, the `access_shape` is not calculated for (SubFeature& feature : sub_features) { - feature.access_shape = IntVec(GetBufferVar(feature.buffer)->shape.size(), 1); + feature.access_shape = IntVec(feature.buffer->shape.size(), 1); } return; } @@ -850,14 +853,14 @@ void Feature::SetRegion(const LoopNest& loop_nest, IntVec* for_touched_bytes, /*allow_override=*/true); int64_t& touched_bytes = (*for_touched_bytes)[i] = 0; for (SubFeature& feature : sub_features) { - const VarNode* buffer = feature.buffer; + BufferVar buffer = feature.buffer; // Note: `feature.access_shape` for `i == 0` is the only one preserved, // while others are discarded int64_t numel; feature.access_shape = utils::RelaxAndUnion(feature.multi_indices, &numel, analyzer); numel = std::max(0, numel); feature.loop_accessed_numel[i][buffer] = numel; - touched_bytes += numel * static_cast(GetBufferVar(buffer)->dtype.StorageBytes()); + touched_bytes += numel * static_cast(buffer->dtype.StorageBytes()); (*buffer_touched_under_loop)[loop][buffer].push_back(numel); } } @@ -867,9 +870,9 @@ void Feature::SubFeature::SetStride(const LoopNest& loop_nest, arith::AnalyzerOb int n_loops = loop_nest.loops.size(); const std::vector& loops = loop_nest.loops; // For each buffer, we find the loop stride on it - const VarNode* buffer = this->buffer; - int ndim = GetBufferVar(this->buffer)->shape.size(); - IntVec buffer_shape = utils::GetBufferShape(BufferVar(ffi::GetRef(buffer)), analyzer); + BufferVar buffer = this->buffer; + int ndim = buffer->shape.size(); + IntVec buffer_shape = utils::GetBufferShape(buffer, analyzer); // Calculate the buffer's stride from its shape IntVec buffer_stride(ndim); if (ndim >= 1) { @@ -885,7 +888,7 @@ void Feature::SubFeature::SetStride(const LoopNest& loop_nest, arith::AnalyzerOb TVM_FFI_ICHECK_EQ(access_shape.size(), buffer_shape.size()); for (int i = ndim - 1; i >= 0; --i) { if (access_shape[i] == buffer_shape[i]) { - num_continuous_bytes = buffer_shape[i] * static_cast(GetBufferVar(buffer)->dtype.StorageBytes()); + num_continuous_bytes = buffer_shape[i] * static_cast(buffer->dtype.StorageBytes()); break; } } @@ -913,7 +916,7 @@ void Feature::SubFeature::SetStride(const LoopNest& loop_nest, arith::AnalyzerOb void Feature::SubFeature::SetReuse(const LoopNest& loop_nest, int64_t top_loop_touch_bytes, const ForBufferMap& buffer_touched_under_loop) { - const VarNode* buffer = this->buffer; + BufferVar buffer = this->buffer; // Step 3.1. Collect all `Var`s that appears in the buffer region std::unordered_set region_vars; for (const MultiIndex& multi_index : this->multi_indices) { @@ -955,10 +958,10 @@ void Feature::SubFeature::SetReuse(const LoopNest& loop_nest, int64_t top_loop_t reuse_dis_bytes = top_loop_touch_bytes; } else { for (const auto& iter : buffer_touched_under_loop.at(loops[i + 1])) { - const VarNode* buffer = iter.first; + BufferVar buffer = iter.first; const IntVec& numels = iter.second; int64_t numel = std::accumulate(numels.begin(), numels.end(), int64_t(0)); - reuse_dis_bytes += numel * static_cast(GetBufferVar(buffer)->dtype.StorageBytes()); + reuse_dis_bytes += numel * static_cast(buffer->dtype.StorageBytes()); } } break; @@ -975,10 +978,10 @@ void Feature::SubFeature::SetReuse(const LoopNest& loop_nest, int64_t top_loop_t reuse_dis_iter = *std::min_element(touched.begin(), touched.end()); reuse_dis_bytes = 0.0; for (const auto& iter : buffer_touched_under_loop.at(loop)) { - const VarNode* buffer = iter.first; + BufferVar buffer = iter.first; const IntVec& numels = iter.second; int64_t numel = std::accumulate(numels.begin(), numels.end(), int64_t(0)); - reuse_dis_bytes += numel * static_cast(GetBufferVar(buffer)->dtype.StorageBytes()); + reuse_dis_bytes += numel * static_cast(buffer->dtype.StorageBytes()); } reuse_dis_iter /= extent; reuse_dis_bytes /= extent; @@ -988,7 +991,7 @@ void Feature::SubFeature::SetReuse(const LoopNest& loop_nest, int64_t top_loop_t } void Feature::SubFeature::SetFeature(const LoopNest& loop_nest, int64_t cache_line_bytes) { - int64_t dtype_bytes = static_cast(GetBufferVar(this->buffer)->dtype.StorageBytes()); + int64_t dtype_bytes = static_cast(this->buffer->dtype.StorageBytes()); this->stride = this->innermost_stride; this->bytes = dtype_bytes * loop_nest.prod; if (loop_nest.loops.empty()) { @@ -1028,7 +1031,7 @@ Feature::Feature(const BufferStoreNode* store, const LoopNest& loop_nest, int64_ int64_t top_loop_touch_bytes = 0.0; if (n_loops > 0) { for (const SubFeature& feature : sub_features) { - int64_t bytes = static_cast(GetBufferVar(feature.buffer)->dtype.StorageBytes()); + int64_t bytes = static_cast(feature.buffer->dtype.StorageBytes()); int64_t n_buffer = feature.loop_accessed_numel[0].size(); top_loop_touch_bytes += bytes * n_buffer; } @@ -1048,7 +1051,7 @@ Feature::Feature(const BufferStoreNode* store, const LoopNest& loop_nest, int64_ if (a.bytes != b.bytes) { return a.bytes > b.bytes; } - return GetBufferVar(a.buffer).name() < GetBufferVar(b.buffer).name(); + return a.buffer.name() < b.buffer.name(); }); } @@ -1160,13 +1163,14 @@ struct Feature { Feature() = default; - explicit Feature(const LoopNest& loop_nest, const BufferVar& buffer, arith::AnalyzerObj* analyzer) { + explicit Feature(const LoopNest& loop_nest, const BufferVar& buffer, + arith::AnalyzerObj* analyzer) { std::vector shape = utils::GetBufferShape(buffer, analyzer); int64_t numel = 1; for (int64_t x : shape) { numel *= x; } - alloc_size = numel * static_cast(GetBufferVar(buffer)->dtype.StorageBytes()); + alloc_size = numel * static_cast(buffer->dtype.StorageBytes()); alloc_prod = numel * loop_nest.prod; alloc_outer_prod = loop_nest.prod; } diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc index 1d5fc2c45e6b..9f5c5a3fee46 100644 --- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc +++ b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc @@ -850,7 +850,7 @@ ffi::Optional MultiLevelTilingTensorCoreNode::TransformWithTensorIntrin( std::unordered_set visited_buffers; ffi::Map buffer_sub_index_map; // cache of the sub index map - // associated with each buffer + // associated with each buffer auto f_transform_buffer_layout = [&](s_tir::BufferIndexType index_type, int buffer_index) { const tirx::BufferVar& lhs_buffer = s_tir::GetNthAccessBuffer( diff --git a/src/s_tir/schedule/analysis.h b/src/s_tir/schedule/analysis.h index 4f1065bfe585..81454d0dc2c4 100644 --- a/src/s_tir/schedule/analysis.h +++ b/src/s_tir/schedule/analysis.h @@ -451,7 +451,7 @@ struct ProducerConsumerSplit { * \throw ScheduleError If the buffer index is out of bound. */ BufferVar GetNthAccessBuffer(const ScheduleState& self, const SBlock& block, int n, - BufferIndexType index_type); + BufferIndexType index_type); /*! * \brief Get the n-th read or write buffer of the given block. @@ -600,7 +600,8 @@ bool CanReverseComputeAt(const ScheduleState& self, const StmtSRef& block_sref, * \param predicate The predicate of the access * \param analyzer Arithmetic analyzer */ -ffi::Optional SuggestIndexMap(const BufferVar& buffer, const ffi::Array& indices, +ffi::Optional SuggestIndexMap(const BufferVar& buffer, + const ffi::Array& indices, const ffi::Array& loops, const PrimExpr& predicate, arith::AnalyzerObj* analyzer); diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc index 83f7c31c3c09..31ef367cb786 100644 --- a/src/s_tir/schedule/analysis/analysis.cc +++ b/src/s_tir/schedule/analysis/analysis.cc @@ -1264,7 +1264,7 @@ BufferRegion GetNthAccessBufferRegion(const ScheduleState& self, const SBlock& b } BufferVar GetNthAccessBuffer(const ScheduleState& self, const SBlock& block, int n, - BufferIndexType index_type) { + BufferIndexType index_type) { return GetNthAccessBufferRegion(self, block, n, index_type)->buffer; } @@ -2006,7 +2006,7 @@ class AutoTensorizeMappingProposer { update_mask(var.value().get(), &rhs_buffer_masks, rhs_buffer_index.at(rhs_buffer)); } else { TVM_FFI_THROW(ValueError) - << "BufferVar index " << rhs_index + << "Buffer index " << rhs_index << " other that variables in tensor intrinsics is not supported."; } } diff --git a/src/s_tir/schedule/analysis/layout.cc b/src/s_tir/schedule/analysis/layout.cc index 12efe4122431..93d258eddac5 100644 --- a/src/s_tir/schedule/analysis/layout.cc +++ b/src/s_tir/schedule/analysis/layout.cc @@ -129,7 +129,8 @@ class SplitExprCollector { std::vector exprs_; }; -ffi::Optional SuggestIndexMap(const BufferVar& buffer, const ffi::Array& indices, +ffi::Optional SuggestIndexMap(const BufferVar& buffer, + const ffi::Array& indices, const ffi::Array& loops, const PrimExpr& predicate, arith::AnalyzerObj* analyzer) { int ndim = buffer->shape.size(); @@ -247,12 +248,12 @@ ffi::Optional SuggestIndexMap(const BufferVar& buffer, const ffi::Arra TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def( - "s_tir.schedule.SuggestIndexMap", - [](BufferVar buffer, ffi::Array indices, ffi::Array loops, PrimExpr predicate) { - arith::Analyzer analyzer; - return SuggestIndexMap(buffer, indices, loops, predicate, analyzer.get()); - }); + refl::GlobalDef().def("s_tir.schedule.SuggestIndexMap", + [](BufferVar buffer, ffi::Array indices, ffi::Array loops, + PrimExpr predicate) { + arith::Analyzer analyzer; + return SuggestIndexMap(buffer, indices, loops, predicate, analyzer.get()); + }); } } // namespace s_tir diff --git a/src/s_tir/schedule/concrete_schedule.cc b/src/s_tir/schedule/concrete_schedule.cc index b22b41d25f85..588822efb36e 100644 --- a/src/s_tir/schedule/concrete_schedule.cc +++ b/src/s_tir/schedule/concrete_schedule.cc @@ -1067,7 +1067,7 @@ void ConcreteScheduleNode::PadEinsum(const SBlockRV& block_rv, const ffi::Array< this->state_->DebugVerify(); } -/******** Schedule: BufferVar Transformation ********/ +/******** Schedule: Buffer Transformation ********/ void ConcreteScheduleNode::RollingBuffer(const SBlockRV& block_rv, int write_buffer_index) { TVM_TIR_SCHEDULE_BEGIN(); diff --git a/src/s_tir/schedule/concrete_schedule.h b/src/s_tir/schedule/concrete_schedule.h index 899c15b3a24b..ab78f8af923d 100644 --- a/src/s_tir/schedule/concrete_schedule.h +++ b/src/s_tir/schedule/concrete_schedule.h @@ -183,7 +183,7 @@ class ConcreteScheduleNode : public ScheduleNode { void TransformBlockLayout(const SBlockRV& block_rv, const IndexMap& index_map) override; /******** Schedule: Padding decomposition ********/ SBlockRV DecomposePadding(const SBlockRV& block_rv, const LoopRV& loop_rv) override; - /******** Schedule: BufferVar transformation ********/ + /******** Schedule: Buffer transformation ********/ void RollingBuffer(const SBlockRV& block_rv, int write_buffer_index) override; /******** Schedule: Misc ********/ void EnterPostproc() override {} diff --git a/src/s_tir/schedule/ir_comparator.cc b/src/s_tir/schedule/ir_comparator.cc index 269952e0cd0c..748743a0f856 100644 --- a/src/s_tir/schedule/ir_comparator.cc +++ b/src/s_tir/schedule/ir_comparator.cc @@ -487,9 +487,18 @@ bool TensorizeComparator::CompareBuffer(const BufferVar& lhs, const BufferVar& r if (it != rhs_buffer_map_.end()) { equal = (*it).second.same_as(lhs); } else { - // Remap the buffer variable definition, skipping buffer shape. - equal = - DefEqual(lhs.var(), rhs.var()) && lhs->dtype == rhs->dtype && lhs.scope() == rhs.scope(); + // Remap the buffer variable definition without recursively comparing its + // BufferType. Tensorization intentionally matches a region of a larger + // workload buffer against the intrinsic's smaller descriptor buffer. + auto data_it = equal_map_.find(lhs.var()); + if (data_it != equal_map_.end()) { + equal = data_it->second.same_as(rhs.var()); + } else { + equal = lhs->dtype == rhs->dtype && lhs.scope() == rhs.scope(); + if (equal) { + equal_map_[lhs.var()] = rhs.var(); + } + } if (equal) { rhs_buffer_map_[rhs] = lhs; } else { @@ -741,22 +750,18 @@ bool AutoTensorizeComparator::CompareBuffer(const BufferVar& lhs, const BufferVa if (it != rhs_buffer_map_.end()) { equal = (*it).second.same_as(lhs); } else { - // Remap both buffer itself and buffer data, skipping buffer shape and storage scope. Auto - // tensorization inserts the cache stages that move workload buffers into an intrinsic's - // required scope, while the pointer element type must still agree. + // Remap the buffer itself, skipping buffer shape and storage scope. Auto + // tensorization inserts the cache stages that move workload buffers into + // an intrinsic's required scope, while the element dtype must still agree. auto data_it = equal_map_.find(lhs.var()); if (data_it != equal_map_.end()) { equal = data_it->second.same_as(rhs.var()); } else { - const auto* lhs_ptr = lhs->data_pointer_type.as(); - const auto* rhs_ptr = rhs->data_pointer_type.as(); - equal = lhs_ptr && rhs_ptr && - ffi::StructuralEqual()(lhs_ptr->element_type, rhs_ptr->element_type); + equal = lhs->dtype == rhs->dtype; if (equal) { equal_map_[lhs.var()] = rhs.var(); } } - equal = equal && lhs->dtype == rhs->dtype; if (equal) { rhs_buffer_map_[rhs] = lhs; lhs_buffer_map_[lhs] = rhs; diff --git a/src/s_tir/schedule/primitive.h b/src/s_tir/schedule/primitive.h index 1d1f9cf4ce60..0c2c02af86e3 100644 --- a/src/s_tir/schedule/primitive.h +++ b/src/s_tir/schedule/primitive.h @@ -640,7 +640,7 @@ TVM_DLL void Unannotate(ScheduleState self, const StmtSRef& sref, const ffi::Str /*! * \brief Apply a transformation represented by IndexMap to buffer * \details The indices and the access region to the target buffer is transformed by the given - * index_map. The index_map is also used to infer the new shape of the buffer. BufferVar must be + * index_map. The index_map is also used to infer the new shape of the buffer. Buffer must be * one of the parameter of the function, or allocated in some blocks (it cannot be a buffer * subregion created via match_buffer). * \param self The state of the schedule @@ -691,7 +691,7 @@ TVM_DLL StmtSRef DecomposePadding(ScheduleState self, const StmtSRef& block_sref */ TVM_DLL void PadEinsum(ScheduleState self, const StmtSRef& block_sref, const ffi::Array& padding); -/******** Schedule: BufferVar transformation ********/ +/******** Schedule: Buffer transformation ********/ /*! * \brief Compute the target buffer via rolling buffering. * \details This primitive selects the outermost rollable axis with a positive bound overlap that diff --git a/src/s_tir/schedule/primitive/annotate_buffer_access.cc b/src/s_tir/schedule/primitive/annotate_buffer_access.cc index cbe626cac8fe..c8640290f180 100644 --- a/src/s_tir/schedule/primitive/annotate_buffer_access.cc +++ b/src/s_tir/schedule/primitive/annotate_buffer_access.cc @@ -39,9 +39,9 @@ class AnnotateRegionRewriter : public StmtExprMutator { ffi::Array regions = buffer_index_type_ == BufferIndexType::kWrite ? block->writes : block->reads; - TVM_FFI_ICHECK_GE(buffer_index_, 0) << "BufferVar index must be non-negative"; + TVM_FFI_ICHECK_GE(buffer_index_, 0) << "Buffer index must be non-negative"; TVM_FFI_ICHECK_LT(buffer_index_, static_cast(regions.size())) - << "BufferVar index out of range"; + << "Buffer index out of range"; regions.Set(buffer_index_, new_region_); ffi::ObjectPtr n = CopyOnWrite(block.get()); diff --git a/src/s_tir/schedule/primitive/block_annotate.cc b/src/s_tir/schedule/primitive/block_annotate.cc index 5ec65c97b2f2..4c59f1b2d46f 100644 --- a/src/s_tir/schedule/primitive/block_annotate.cc +++ b/src/s_tir/schedule/primitive/block_annotate.cc @@ -227,7 +227,7 @@ void StorageAlign(ScheduleState self, const StmtSRef& block_sref, int buffer_ind int factor, int offset) { const SBlockNode* block_ptr = TVM_SREF_TO_SBLOCK(block_sref); BufferVar buffer = GetNthAccessBuffer(self, ffi::GetRef(block_ptr), buffer_index, - BufferIndexType::kWrite); + BufferIndexType::kWrite); StorageAlignInvalidFactorError::Check(self->mod, factor); axis = StorageAlignAxisOutOfRangeError::CheckAndUpdate(self->mod, buffer, axis); NonAllocatedBufferError::CheckAndGetBufferAllocationSite(self->mod, block_sref, buffer); diff --git a/src/s_tir/schedule/primitive/cache_index.cc b/src/s_tir/schedule/primitive/cache_index.cc index e089a58135ba..f7a3837bfa73 100644 --- a/src/s_tir/schedule/primitive/cache_index.cc +++ b/src/s_tir/schedule/primitive/cache_index.cc @@ -262,15 +262,14 @@ ffi::Array MakeIndexCacheStage(IndexInfo* info, const ffi::String& stora } PrimType data_ty = index_expr.ty(); - Var index_buffer_var("index_var_" + std::to_string(expr_index), - PointerType(data_ty, storage_scope)); + ffi::String index_buffer_name = "index_var_" + std::to_string(expr_index); ffi::Array buffer_shape; for (const Var& it : info->origin_block_vars[expr_index]) { buffer_shape.push_back( arith::EvalSet(info->var_binding.at(it), arith::AsIntSet(info->range_map)).max() + 1); } - info->cache_buffer.push_back( - BufferVar(index_buffer_var, data_ty, buffer_shape, {1}, {0}, index_buffer_var->name, 0, 0)); + info->cache_buffer.push_back(BufferVar( + index_buffer_name, BufferType(storage_scope, data_ty, buffer_shape, {1}, {0}, 0, 0))); // Create loop vars and block vars' binding_value std::vector loop_vars; diff --git a/src/s_tir/schedule/primitive/cache_read_write.cc b/src/s_tir/schedule/primitive/cache_read_write.cc index f5483119e6c7..d03959178c51 100644 --- a/src/s_tir/schedule/primitive/cache_read_write.cc +++ b/src/s_tir/schedule/primitive/cache_read_write.cc @@ -1440,7 +1440,7 @@ class ReindexCacheWriteRewriter : public CacheWriteRewriter { * \return The new buffer with target shape. */ BufferVar CreateReindexBuffer(const BufferVar& buffer, const ffi::Array& block_iters, - const std::unordered_set& covered) { + const std::unordered_set& covered) { ffi::ObjectPtr new_buffer = CopyBufferType(buffer); std::vector new_shape; std::vector new_strides; @@ -1732,7 +1732,7 @@ StmtSRef CacheRead(ScheduleState self, const StmtSRef& block_sref, int read_buff // Step 1. Check index, getting the target buffer and the parent scope const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); BufferVar read_buffer = GetNthAccessBuffer(self, ffi::GetRef(block), read_buffer_index, - BufferIndexType::kRead); + BufferIndexType::kRead); StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false); // Check required region cover for cache_read CheckRegionCover(self, scope_sref, read_buffer); @@ -1842,7 +1842,7 @@ StmtSRef CacheWrite(ScheduleState self, const StmtSRef& block_sref, int write_bu // Step 1. Checking index, getting the target buffer and the parent scope const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); BufferVar write_buffer = GetNthAccessBuffer(self, ffi::GetRef(block), write_buffer_index, - BufferIndexType::kWrite); + BufferIndexType::kWrite); StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false); // Step 2. Creating CacheStageInfo @@ -2035,8 +2035,7 @@ void CollectReindexCacheStageInfoAndCreateBuffer( // Create new buffer ffi::ObjectPtr new_buffer = CopyBufferType(old_buffer); - const auto* ptr_type = TVM_TYPE_AS(old_buffer->data_pointer_type, PointerTypeNode); - new_buffer->data_pointer_type = PointerType(ptr_type->element_type, storage_scope); + new_buffer->storage_scope = storage_scope; new_buffer->shape = new_shape; BufferVar rebuilt = RebuildBufferVar(old_buffer, std::move(new_buffer), old_buffer.name() + "_" + storage_scope); @@ -2084,7 +2083,8 @@ StmtSRef ReindexCacheRead(ScheduleState self, const StmtSRef& block_sref, int re // Step 1. Check index, getting the target buffer and the parent scope SBlock block = ffi::GetRef(TVM_SREF_TO_SBLOCK(block_sref)); SBlockRealize realize = GetSBlockRealize(self, block_sref); - BufferVar read_buffer = GetNthAccessBuffer(self, block, read_buffer_index, BufferIndexType::kRead); + BufferVar read_buffer = + GetNthAccessBuffer(self, block, read_buffer_index, BufferIndexType::kRead); StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/true); // Step 2. Create CacheStageInfo @@ -2235,7 +2235,7 @@ ffi::Array CacheInplace(ScheduleState self, const StmtSRef& block_sref // Check 1. Check index, get the target buffer and the parent scope const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); BufferVar buffer = GetNthAccessBuffer(self, ffi::GetRef(block), read_buffer_index, - BufferIndexType::kRead); + BufferIndexType::kRead); StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false); // Check 3. Check required region cover for cache_read diff --git a/src/s_tir/schedule/primitive/compute_at.cc b/src/s_tir/schedule/primitive/compute_at.cc index 535541ddbd3d..5d8894c4f98c 100644 --- a/src/s_tir/schedule/primitive/compute_at.cc +++ b/src/s_tir/schedule/primitive/compute_at.cc @@ -579,8 +579,8 @@ bool UpdateBlockVarDomainAffine(const VarNode* buffer, const ffi::Array // calculate backward mapping (required region point -> block vars) NDIntSet required_bound; for (size_t i = 0; i < ndim; ++i) { - required_bound.push_back( - arith::IntSet::Interval(IntImm(GetBufferVar(buffer)->shape[i].ty(), 0), max(GetBufferVar(buffer)->shape[i] - 1, 0))); + required_bound.push_back(arith::IntSet::Interval(IntImm(GetBufferVar(buffer)->shape[i].ty(), 0), + max(GetBufferVar(buffer)->shape[i] - 1, 0))); } ffi::Map var_dom = InverseAffineIterMap(res->indices, required_region, analyzer); diff --git a/src/s_tir/schedule/primitive/compute_inline.cc b/src/s_tir/schedule/primitive/compute_inline.cc index c6997a19bbfb..c707026efe38 100644 --- a/src/s_tir/schedule/primitive/compute_inline.cc +++ b/src/s_tir/schedule/primitive/compute_inline.cc @@ -92,9 +92,10 @@ class NotSingleReadWriteBuffer : public ScheduleError { SBlock block_; static BufferVar GetSingleRead(const ScheduleState& self, const SBlock& block, - const StmtSRef& scope_root_sref) { - const std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>& - buffer_writers = self->block_info.at(scope_root_sref).scope->buffer_writers; + const StmtSRef& scope_root_sref) { + const std::unordered_map, ffi::ObjectPtrHash, + ffi::ObjectPtrEqual>& buffer_writers = + self->block_info.at(scope_root_sref).scope->buffer_writers; const VarNode* read_buffer = nullptr; for (const BufferRegion& read_region : block->reads) { const VarNode* buffer = read_region->buffer.get(); @@ -184,8 +185,8 @@ class NonSingleProducerError : public ScheduleError { class ProducerFinder : public StmtVisitor { public: static std::vector GetProducer(const ScheduleState& self, - const StmtSRef& scope_root_sref, const BufferVar& buffer, - const SBlock& scope_block) { + const StmtSRef& scope_root_sref, + const BufferVar& buffer, const SBlock& scope_block) { ProducerFinder finder(self, scope_root_sref, buffer); finder(scope_block); return finder.producer_across_scope_.back(); @@ -996,7 +997,8 @@ void ReverseComputeInline(ScheduleState self, const StmtSRef& consumer_block_sre */ class ReductionEpilogueFuser : public BaseInliner { public: - explicit ReductionEpilogueFuser(const BufferVar& reduction_buffer, const SBlockNode* reduction_block, + explicit ReductionEpilogueFuser(const BufferVar& reduction_buffer, + const SBlockNode* reduction_block, const SBlockRealize& epilogue_block_realize, const StmtSRef& scope_root_sref) : BaseInliner(reduction_buffer, epilogue_block_realize->block, scope_root_sref), @@ -1053,11 +1055,11 @@ class ReductionEpilogueFuser : public BaseInliner { PrimExpr epilogue_expression_{ nullptr}; // The entire epilogue expression (e.g., temp + C, max(temp + C, 0)) const BufferLoadNode* reduction_buffer_load_{ - nullptr}; // The reduction buffer load in epilogue expression + nullptr}; // The reduction buffer load in epilogue expression BufferVar epilogue_output_buffer_{nullptr}; // Output buffer D ffi::Array epilogue_output_indices_{nullptr}; // Indices of D[vi, vj] BufferRegion epilogue_output_region_{nullptr}; // Write region of D - BufferVar epilogue_addend_buffer_{nullptr}; // Additional buffer (e.g., bias buffer C) + BufferVar epilogue_addend_buffer_{nullptr}; // Additional buffer (e.g., bias buffer C) BufferRegion epilogue_addend_region_{nullptr}; // Read region of additional buffer }; @@ -1312,8 +1314,9 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( // remove that operand (bias addend) from update expression class UpdateSubstituter : public StmtExprMutator { public: - UpdateSubstituter(const BufferVar& old_buf, const BufferVar& new_buf, const BufferVar& reduction_buf, - const PrimExpr& epilogue_expr, const std::unordered_map& var_map) + UpdateSubstituter(const BufferVar& old_buf, const BufferVar& new_buf, + const BufferVar& reduction_buf, const PrimExpr& epilogue_expr, + const std::unordered_map& var_map) : old_buffer_(old_buf), new_buffer_(new_buf), reduction_buffer_(reduction_buf), diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc b/src/s_tir/schedule/primitive/layout_transformation.cc index c9de8f6fbe76..9c8f835a0a55 100644 --- a/src/s_tir/schedule/primitive/layout_transformation.cc +++ b/src/s_tir/schedule/primitive/layout_transformation.cc @@ -95,8 +95,8 @@ class TransformLayoutPlanner : private StmtExprVisitor { using TransformPlan = std::variant; - static TransformPlan Plan(SBlock block, BufferVar old_buffer, BufferVar new_buffer, IndexMap index_map, - IndexMap inverse, PrimExpr padding_predicate, + static TransformPlan Plan(SBlock block, BufferVar old_buffer, BufferVar new_buffer, + IndexMap index_map, IndexMap inverse, PrimExpr padding_predicate, ffi::Optional pad_value, arith::AnalyzerObj* analyzer) { TVM_FFI_ICHECK(!pad_value.has_value() || pad_value.value()->final_indices.size() == 1) << "Internal error: Should be caught by ScheduleError checks prior to this point"; @@ -222,8 +222,9 @@ class TransformLayoutPlanner : private StmtExprVisitor { class BufferStoreReplacer : public StmtExprMutator { public: - BufferStoreReplacer(const WriteInfo& info, const BufferVar& new_buffer, PrimExpr padding_predicate, - const IndexMap& inverse, const ffi::Optional& pad_value, + BufferStoreReplacer(const WriteInfo& info, const BufferVar& new_buffer, + PrimExpr padding_predicate, const IndexMap& inverse, + const ffi::Optional& pad_value, ffi::Map* new_block_to_old, arith::AnalyzerObj* analyzer) : info(info), new_buffer(new_buffer), @@ -951,7 +952,8 @@ class BufferIsSubregionError : public ScheduleError { ffi::String DetailRenderTemplate() const final { std::ostringstream os; - os << "ScheduleError: The input buffer " << buffer_.name() << " is defined in `match_buffer` of " + os << "ScheduleError: The input buffer " << buffer_.name() + << " is defined in `match_buffer` of " << "a block, it is expected to be a function parameter or allocated by a block."; return os.str(); } @@ -1007,7 +1009,7 @@ class TransformationPaddingTypeError : public ScheduleError { ffi::String DetailRenderTemplate() const final { std::ostringstream ss; - ss << "ScheduleError: BufferVar " << buffer_.name() << " has elements of type " << buffer_->dtype + ss << "ScheduleError: Buffer " << buffer_.name() << " has elements of type " << buffer_->dtype << ", but the transformation fills padding with " << pad_value_ << ", which is of type " << pad_value_dtype_; return ss.str(); @@ -1127,7 +1129,7 @@ IndexMap LegalizeIndexMapDType(const IndexMap& index_map, const ffi::Arraydtype; if (index_dtype.has_value()) { TVM_FFI_ICHECK_EQ(*index_dtype, arg_dtype) - << "BufferVar index " << args[i] << " has dtype " << arg_dtype + << "Buffer index " << args[i] << " has dtype " << arg_dtype << ", but previous index for the same buffer access used index type " << *index_dtype; } else { index_dtype = arg_dtype; diff --git a/src/s_tir/schedule/primitive/read_write_at.cc b/src/s_tir/schedule/primitive/read_write_at.cc index cd0e3656517a..0d2345796dd7 100644 --- a/src/s_tir/schedule/primitive/read_write_at.cc +++ b/src/s_tir/schedule/primitive/read_write_at.cc @@ -40,7 +40,7 @@ bool HasBuffer(const ffi::Array& buffer_regions, const BufferVar& } void RelaxBufferRegions(const ffi::Array& buffer_regions, - const BufferVar& buffer, // + const BufferVar& buffer, // const ffi::Map& var_dom, // const ffi::Map& bindings, // std::vector* relaxed_regions) { @@ -55,8 +55,8 @@ void RelaxBufferRegions(const ffi::Array& buffer_regions, class ScopeReplacer : public StmtMutator { public: - static SBlock Replace(const SBlockNode* scope_block, const BufferVar& dst, const ForNode* old_loop, - const ForNode* new_loop) { + static SBlock Replace(const SBlockNode* scope_block, const BufferVar& dst, + const ForNode* old_loop, const ForNode* new_loop) { ffi::ObjectPtr new_scope_block = ffi::make_object(*scope_block); new_scope_block->body = ScopeReplacer(old_loop, new_loop)(std::move(new_scope_block->body)); new_scope_block->alloc_buffers.push_back(dst); @@ -131,7 +131,7 @@ struct ReadWriteAtImpl { ffi::Map annotations) { const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); BufferVar src = GetNthAccessBuffer(self, ffi::GetRef(block), buffer_index, - is_read ? BufferIndexType::kRead : BufferIndexType::kWrite); + is_read ? BufferIndexType::kRead : BufferIndexType::kWrite); BufferVar dst = WithScope(src, storage_scope); ReadWriteAtImpl impl(self, loop_sref, src, dst, annotations); std::pair new_loop_block = diff --git a/src/s_tir/schedule/primitive/reduction.cc b/src/s_tir/schedule/primitive/reduction.cc index b2f6fcfecd00..e19abf6c39a6 100644 --- a/src/s_tir/schedule/primitive/reduction.cc +++ b/src/s_tir/schedule/primitive/reduction.cc @@ -654,8 +654,8 @@ std::unordered_map GetLoopVar2LoopMap(const ffi::Array * \param rf_loop The rfactor loop * \return The new created intermediate rfactor buffer */ -ffi::Array CreateRFactorBuffers(const ffi::Array& buf_stores, int factor_axis, - const ForNode* rf_loop) { +ffi::Array CreateRFactorBuffers(const ffi::Array& buf_stores, + int factor_axis, const ForNode* rf_loop) { ffi::Array rf_buffers; rf_buffers.reserve(buf_stores.size()); for (const BufferStore& buf_store : buf_stores) { diff --git a/src/s_tir/schedule/schedule.cc b/src/s_tir/schedule/schedule.cc index 6275b8463332..e8e666672f08 100644 --- a/src/s_tir/schedule/schedule.cc +++ b/src/s_tir/schedule/schedule.cc @@ -333,7 +333,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def_method("s_tir.schedule.ScheduleDecomposePadding", &ScheduleNode::DecomposePadding) .def_method("s_tir.schedule.SchedulePadEinsum", &ScheduleNode::PadEinsum); } -/******** (FFI) BufferVar transformation ********/ +/******** (FFI) Buffer transformation ********/ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def_method("s_tir.schedule.ScheduleRollingBuffer", diff --git a/src/s_tir/schedule/state.cc b/src/s_tir/schedule/state.cc index 9f6528fc7b86..afdf1b34c322 100644 --- a/src/s_tir/schedule/state.cc +++ b/src/s_tir/schedule/state.cc @@ -272,11 +272,12 @@ class SBlockInfoCollector : private StmtVisitor { continue; } // For each buffer, record the regions generated under this loop - std::unordered_map>> + std::unordered_map>, ffi::ObjectPtrHash, + ffi::ObjectPtrEqual> touched_regions; // Step 2.3.1. Find all the regions read by the consumer that we care about for (const BufferRegion& region : block_reads_unbound.at(consumer_block_sref.get())) { - const VarNode* buffer = region->buffer.get(); + BufferVar buffer = region->buffer; touched_regions[buffer] = {}; } // Step 2.3.2. Find all the regions written by each producer @@ -284,7 +285,7 @@ class SBlockInfoCollector : private StmtVisitor { const SBlockRealize& producer_realize = block2realize_.at(producer_block_sref->stmt); StmtSRef parent_sref = ffi::GetRef(producer_block_sref->parent); for (const BufferRegion& region : block_writes_unbound.at(producer_block_sref)) { - const VarNode* buffer = region->buffer.get(); + BufferVar buffer = region->buffer; auto it = touched_regions.find(buffer); // Skip the regions that is not read by the consumer if (it != touched_regions.end()) { @@ -306,7 +307,7 @@ class SBlockInfoCollector : private StmtVisitor { { StmtSRef parent_sref = ffi::GetRef(consumer_block_sref->parent); for (const BufferRegion& region : block_reads_unbound.at(consumer_block_sref.get())) { - const VarNode* buffer = region->buffer.get(); + BufferVar buffer = region->buffer; const std::vector>& touched_region = touched_regions.at(buffer); if (!touched_region.empty()) { @@ -318,8 +319,7 @@ class SBlockInfoCollector : private StmtVisitor { /*dom_low_inclusive=*/parent_sref, /*dom_high_exclusive=*/lca, /*analyzer=*/analyzer_.get()); - if (!ProducerCoversConsumer(GetBufferVar(buffer)->shape, produced_region, - consumed_region, + if (!ProducerCoversConsumer(buffer->shape, produced_region, consumed_region, analyzer_.get())) { region_cover = false; self_->block_info.at(consumer_block_sref).region_cover = region_cover; diff --git a/src/s_tir/schedule/trace.cc b/src/s_tir/schedule/trace.cc index 516a44d2872d..b2424079af4f 100644 --- a/src/s_tir/schedule/trace.cc +++ b/src/s_tir/schedule/trace.cc @@ -72,7 +72,7 @@ ffi::Array TranslateInputRVs( const ffi::Object* dst = it->second; TVM_FFI_CHECK(dst->IsInstance(), TypeError) << "Expect 'tirx.Var', but gets: " << dst->GetTypeKey(); - return ffi::GetRef(static_cast(dst)).as_or_throw(); + return ffi::GetRef(static_cast(dst)); }; auto f_subst_with_rv_map_prim = [&](const Var& var) -> ffi::Optional { if (auto replacement = f_subst_with_rv_map(var)) { diff --git a/src/s_tir/schedule/traced_schedule.cc b/src/s_tir/schedule/traced_schedule.cc index 2849183bc43a..646269abfb7e 100644 --- a/src/s_tir/schedule/traced_schedule.cc +++ b/src/s_tir/schedule/traced_schedule.cc @@ -759,7 +759,7 @@ void TracedScheduleNode::PadEinsum(const SBlockRV& block_rv, const ffi::Array& padding) final; - /******** Schedule: BufferVar transformation ********/ + /******** Schedule: Buffer transformation ********/ void RollingBuffer(const SBlockRV& block_rv, int write_buffer_index) final; /******** Schedule: Misc ********/ void EnterPostproc() final; diff --git a/src/s_tir/schedule/transform.cc b/src/s_tir/schedule/transform.cc index 976c39f55d1a..964f0a515ae3 100644 --- a/src/s_tir/schedule/transform.cc +++ b/src/s_tir/schedule/transform.cc @@ -38,20 +38,18 @@ SBlock WithAnnotation(const SBlockNode* block, const ffi::String& attr_key, return SBlock(new_block); } -/******** BufferVar Related ********/ +/******** Buffer Related ********/ BufferVar WithScope(const BufferVar& buffer, const ffi::String& scope) { - const auto* ptr_type = TVM_TYPE_AS(buffer->data_pointer_type, PointerTypeNode); - BufferType new_type(PointerType(ptr_type->element_type, scope), buffer->dtype, buffer->shape, - buffer->strides, buffer->elem_offset, buffer->data_alignment, - buffer->offset_factor, buffer->layout, buffer->allocated_addr); + BufferType new_type(scope, buffer->dtype, buffer->shape, buffer->strides, buffer->elem_offset, + buffer->data_alignment, buffer->offset_factor, buffer->layout, + buffer->allocated_addr); return BufferVar(buffer.name() + "_" + scope, new_type, buffer.span()); } BufferVar WithDType(const BufferVar& buffer, PrimType dtype) { - const auto* ptr_type = TVM_TYPE_AS(buffer->data_pointer_type, PointerTypeNode); - BufferType new_type(PointerType(dtype, ptr_type->storage_scope), dtype, buffer->shape, - buffer->strides, buffer->elem_offset, buffer->data_alignment, - buffer->offset_factor, buffer->layout, buffer->allocated_addr); + BufferType new_type(buffer->storage_scope, dtype, buffer->shape, buffer->strides, + buffer->elem_offset, buffer->data_alignment, buffer->offset_factor, + buffer->layout, buffer->allocated_addr); return BufferVar(buffer.name(), new_type, buffer.span()); } diff --git a/src/s_tir/schedule/transform.h b/src/s_tir/schedule/transform.h index 57e422b1d947..f52fbf0994a1 100644 --- a/src/s_tir/schedule/transform.h +++ b/src/s_tir/schedule/transform.h @@ -45,7 +45,7 @@ using namespace tvm::tirx; SBlock WithAnnotation(const SBlockNode* block, const ffi::String& attr_key, const ffi::ObjectRef& attr_value); -/******** BufferVar Related ********/ +/******** Buffer Related ********/ /*! * \brief Create a new buffer by changing the storage scope. diff --git a/src/s_tir/transform/compact_buffer_region.cc b/src/s_tir/transform/compact_buffer_region.cc index c119956ca296..598e4364adb1 100644 --- a/src/s_tir/transform/compact_buffer_region.cc +++ b/src/s_tir/transform/compact_buffer_region.cc @@ -243,7 +243,8 @@ class BufferAccessRegionCollector : public StmtExprVisitor { TVM_FFI_ICHECK(!op->init.has_value()); TVM_FFI_ICHECK_EQ(op->iter_vars.size(), 0) << "CompactBufferRegion only works on opaque blocks"; // Step 1. Record and update current read/write region annotations - std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> + std::unordered_map, ffi::ObjectPtrHash, + ffi::ObjectPtrEqual> cur_access_annotations; for (const BufferRegion& region : op->reads) { cur_access_annotations[region->buffer].push_back(region); @@ -526,14 +527,16 @@ class BufferAccessRegionCollector : public StmtExprVisitor { /*! \brief The analyzer aware of loop domains. */ arith::Analyzer dom_analyzer_; /*! \brief The map from BufferVar to it's relaxed access set. */ - std::unordered_map relaxed_accesses_; + std::unordered_map + relaxed_accesses_; /*! * \brief The map from BufferVar to it entire access region, used for returning. * The entire access region should get updated on the buffer's define point * and we sanity check that every buffer is defined only once. */ - std::unordered_map buffer_access_region_; + std::unordered_map + buffer_access_region_; /*! \brief The map from BufferVar to it's access regions annotated by current block. */ std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> @@ -567,47 +570,35 @@ struct BufferAllocInfo { class BufferCompactor : public StmtExprMutator { public: explicit BufferCompactor(std::unordered_map buffer_info) - : buffer_info_(std::move(buffer_info)) { - std::vector> remapped_entries; - for (const auto& [old_buffer, info] : buffer_info_) { - if (!old_buffer.same_as(info.new_buffer.var())) { - remapped_entries.emplace_back(info.new_buffer.var(), info); - } - } - for (auto& [new_buffer, info] : remapped_entries) { - buffer_info_.emplace(std::move(new_buffer), std::move(info)); - } - } + : buffer_info_(std::move(buffer_info)) {} Stmt VisitStmt_(const BufferStoreNode* _op) final { BufferStore store = StmtExprMutator::VisitStmt_(_op).as_or_throw(); BufferStoreNode* op = store.CopyOnWrite(); - RewriteBufferAccess(&op->buffer, &op->indices); + RewriteBufferAccess(_op->buffer, &op->buffer, &op->indices); return store; } Expr VisitExpr_(const BufferLoadNode* _op) final { BufferLoad load = StmtExprMutator::VisitExpr_(_op).as_or_throw(); BufferLoadNode* op = load.CopyOnWrite(); - RewriteBufferAccess(&op->buffer, &op->indices); + RewriteBufferAccess(_op->buffer, &op->buffer, &op->indices); return load; } Stmt VisitStmt_(const SBlockNode* op) final { // Step 0. Check there is no Init part. TVM_FFI_ICHECK(!op->init.has_value()); - // Step 1. Reallocate and rewrite alloc_buffers, also update BufferAllocInfo. - ffi::Array alloc_buffers = - op->alloc_buffers.Map([this](const BufferVar& buf) { return RewriteAllocBuffer(buf); }); - // Step 2. Recursively rewrite BufferLoad/BufferStore. - SBlock block = StmtExprMutator::VisitStmt_(op).as_or_throw(); - // Step 3. Update block signature. + // Rewrite the signature while its buffer identities still match buffer_info_. + SBlock block = ffi::GetRef(op); SBlockNode* n = block.CopyOnWrite(); RewriteBufferRegions(&n->reads); RewriteBufferRegions(&n->writes); RewriteMatchBuffers(&n->match_buffers); - n->alloc_buffers = std::move(alloc_buffers); - return block; + n->alloc_buffers = + op->alloc_buffers.Map([this](const BufferVar& buf) { return RewriteAllocBuffer(buf); }); + // Recursively rewrite the body after installing the allocation remaps. + return StmtExprMutator::VisitStmt_(block.get()); } Stmt VisitStmt_(const DeclBufferNode* op) final { @@ -618,7 +609,7 @@ class BufferCompactor : public StmtExprMutator { Stmt VisitStmt_(const AllocBufferNode* op) final { RewriteAllocBuffer(op->buffer); AllocBuffer alloc_buf = StmtExprMutator::VisitStmt_(op).as_or_throw(); - auto it = buffer_info_.find(alloc_buf->buffer.var()); + auto it = buffer_info_.find(op->buffer.var()); if (it == buffer_info_.end()) { return alloc_buf; } @@ -642,8 +633,9 @@ class BufferCompactor : public StmtExprMutator { return buffer; } - void RewriteBufferAccess(BufferVar* buffer, ffi::Array* indices) const { - auto it = buffer_info_.find((*buffer).var()); + void RewriteBufferAccess(const BufferVar& original_buffer, BufferVar* buffer, + ffi::Array* indices) const { + auto it = buffer_info_.find(original_buffer.var()); if (it == buffer_info_.end()) { return; } diff --git a/src/s_tir/transform/inject_double_buffer.cc b/src/s_tir/transform/inject_double_buffer.cc index 63cf54062659..456d4ec17513 100644 --- a/src/s_tir/transform/inject_double_buffer.cc +++ b/src/s_tir/transform/inject_double_buffer.cc @@ -36,6 +36,21 @@ namespace tvm { namespace s_tir { using namespace tvm::tirx; +namespace { + +ffi::Optional GetBufferDataVar(const ffi::Any& data) { + if (auto var = data.as()) { + return var; + } + if (const auto* call = data.as(); + call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) { + return call->args[0].as(); + } + return std::nullopt; +} + +} // namespace + struct InjectDoubleBufferConfigNode : public ffi::Object { int split_loop; @@ -64,7 +79,9 @@ class DoubleBufferDetector : public StmtExprVisitor { public: void VisitStmt_(const AttrStmtNode* op) final { if (op->attr_key == s_tir::attr::double_buffer_scope) { - touched_.insert(op->node.as()); + if (auto buffer = GetBufferDataVar(op->node)) { + touched_.insert(buffer.value().get()); + } StmtExprVisitor::VisitStmt_(op); } else { StmtExprVisitor::VisitStmt_(op); @@ -269,7 +286,7 @@ class DoubleBufferInjector : public StmtExprMutator { private: Stmt MakeProducer(const AttrStmtNode* op) { - const Var buffer = op->node.as_or_throw(); + const Var buffer = GetBufferDataVar(op->node).value(); TVM_FFI_ICHECK_NE(loop_nest_.size(), 0U) << "Double buffer scope must be inside a loop"; auto it = dbuffer_info_.find(buffer.get()); if (it == dbuffer_info_.end()) { @@ -294,7 +311,8 @@ class DoubleBufferInjector : public StmtExprMutator { vmap[e.loop->loop_var.get()] = loop_shift; vmap[e.switch_write_var.get()] = indexmod(loop_shift, two); body = Substitute(body, vmap); - body = AttrStmt(buffer, s_tir::attr::double_buffer_write, 1, body); + body = AttrStmt(GetRemappedBuffer(BufferVar(buffer), e.stride).data(), + s_tir::attr::double_buffer_write, 1, body); body = IfThenElse(loop_shift < e.loop->extent, body); return body; } diff --git a/src/s_tir/transform/inject_permuted_layout.cc b/src/s_tir/transform/inject_permuted_layout.cc index b5686e77b148..6c2ad82710a7 100644 --- a/src/s_tir/transform/inject_permuted_layout.cc +++ b/src/s_tir/transform/inject_permuted_layout.cc @@ -41,6 +41,21 @@ using namespace tvm::tirx; using namespace arith; using namespace runtime; +namespace { + +ffi::Optional GetBufferDataVar(const ffi::Any& data) { + if (auto var = data.as()) { + return var; + } + if (const auto* call = data.as(); + call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) { + return call->args[0].as(); + } + return std::nullopt; +} + +} // namespace + class PermutedLayoutInjector : private IRMutatorWithAnalyzer { public: static PrimFunc Transform(PrimFunc func) { @@ -56,6 +71,9 @@ class PermutedLayoutInjector : private IRMutatorWithAnalyzer { explicit PermutedLayoutInjector(PrimFunc func, const Analyzer& analyzer) : IRMutatorWithAnalyzer(analyzer) { buffer_map_.insert(func->buffer_map.begin(), func->buffer_map.end()); + for (const auto& [_, buffer] : func->buffer_map) { + buffer_map_.insert({buffer.var(), buffer}); + } } using IRMutatorWithAnalyzer::VisitExpr_; @@ -156,10 +174,11 @@ class PermutedLayoutInjector : private IRMutatorWithAnalyzer { if (buffer_row_size % 64 != 0) { TVM_FFI_ICHECK(buffer_row_size % 32 == 0) - << "Permuted SLayout for BufferVar \"" << buffer.name() << "\" with shape " << buffer->shape - << " is not supported since its second dimension is not divisible by 32"; + << "Permuted SLayout for BufferVar \"" << buffer.name() << "\" with shape " + << buffer->shape << " is not supported since its second dimension is not divisible by 32"; TVM_FFI_ICHECK(buffer_col_size % 2 == 0) - << "Permuted SLayout for BufferVar \"" << buffer.name() << "\" with shape " << buffer->shape + << "Permuted SLayout for BufferVar \"" << buffer.name() << "\" with shape " + << buffer->shape << " is not supported since its first dimension is not divisible by 2 and second " "dimension is not divisible by 64"; } @@ -227,7 +246,10 @@ class PermutedLayoutInjector : private IRMutatorWithAnalyzer { TVM_FFI_ICHECK(access_ptr_call->op.same_as(builtin::tvm_access_ptr())) << "Invalid access ptr for permuted layout: " << access_ptr; - auto buffer_map_iter = buffer_map_.find(access_ptr_call->args[1].as_or_throw()); + auto data_var = GetBufferDataVar(access_ptr_call->args[1]); + TVM_FFI_ICHECK(data_var.has_value()) + << "Expected a buffer data expression, but received " << access_ptr_call->args[1]; + auto buffer_map_iter = buffer_map_.find(data_var.value()); TVM_FFI_ICHECK(buffer_map_iter != buffer_map_.end()) << "The buffer corresponding to data Var " << access_ptr_call->args[1] << " is not found"; int buffer_row_size = CheckAndGetBufferRowSize(buffer_map_iter->second); diff --git a/src/s_tir/transform/inject_ptx_async_copy.cc b/src/s_tir/transform/inject_ptx_async_copy.cc index 93c9261c532b..1cbbd388afe0 100644 --- a/src/s_tir/transform/inject_ptx_async_copy.cc +++ b/src/s_tir/transform/inject_ptx_async_copy.cc @@ -62,8 +62,8 @@ class PTXAsyncCopyInjector : public StmtMutator { const int bytes = indices_lanes * static_cast(load->buffer->dtype.StorageBytes()); if (bytes == 4 || bytes == 8 || bytes == 16) { - auto dst_elem_type = GetPointerType(store->buffer->data_pointer_type); - auto src_elem_type = GetPointerType(load->buffer->data_pointer_type); + auto dst_elem_type = GetPointerType(store->buffer.DataPointerType()); + auto src_elem_type = GetPointerType(load->buffer.DataPointerType()); TVM_FFI_ICHECK(dst_elem_type.has_value() && src_elem_type.has_value()) << "Both store and load buffer should have a pointer type annotation."; @@ -152,10 +152,11 @@ class PTXAsyncCopyInjector : public StmtMutator { if (src_offset.defined() && dst_offset.defined()) { static const Op& ptx_cp_async_op = Op::Get("tirx.ptx.cp_async_raw"); - return Evaluate(Call(store->buffer->dtype, ptx_cp_async_op, - {store->buffer.data(), mul(dst_offset, PrimExpr(index_factor)), - load->buffer.data(), src_offset, PrimExpr(bytes), predicate_value}) - .as_or_throw()); + return Evaluate( + Call(store->buffer->dtype, ptx_cp_async_op, + {store->buffer.data(), mul(dst_offset, PrimExpr(index_factor)), + load->buffer.data(), src_offset, PrimExpr(bytes), predicate_value}) + .as_or_throw()); } } } diff --git a/src/s_tir/transform/inject_software_pipeline.cc b/src/s_tir/transform/inject_software_pipeline.cc index 598fb876e7fa..2f20f10303ec 100644 --- a/src/s_tir/transform/inject_software_pipeline.cc +++ b/src/s_tir/transform/inject_software_pipeline.cc @@ -43,6 +43,21 @@ using namespace tvm::tirx; namespace software_pipeline { +namespace { + +ffi::Optional GetBufferDataVar(const ffi::Any& data) { + if (auto var = data.as()) { + return var; + } + if (const auto* call = data.as(); + call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) { + return call->args[0].as(); + } + return std::nullopt; +} + +} // namespace + /*! * \brief Create a block and infer the access region with the given body. * @@ -114,11 +129,12 @@ class PipelineOpaqueAccessRewriter { static const Op& ptx_ldmatrix_legacy = Op::Get("tirx.ptx.ldmatrix_legacy"); static const Op& ptx_mma_legacy = Op::Get("tirx.ptx.mma_legacy"); if (call->op.same_as(load_matrix_sync) || call->op.same_as(store_matrix_sync)) { - const BufferVar& buffer = buffer_data_to_buffer_.at(call->args[0].as_or_throw()); + const BufferVar& buffer = buffer_data_to_buffer_.at(GetBufferDataVar(call->args[0]).value()); auto it = buffer_remap_.find(buffer); if (it != buffer_remap_.end()) { ffi::Array new_args = call->args; const BufferVar& new_buffer = (*it).second; + new_args.Set(0, new_buffer.data()); new_args.Set( 4, RewriteWmmaFragmentIndex(buffer, new_buffer, call->args[4].as_or_throw())); return Call(call->ty, call->op, new_args, call->attrs, {}, call->span); @@ -126,12 +142,14 @@ class PipelineOpaqueAccessRewriter { } else if (call->op.same_as(mma_sync)) { ffi::Array new_args = call->args; for (int i = 0; i < 4; i++) { - const Var& buffer_var = call->args[i * 2].as_or_throw(); + const Var& buffer_var = GetBufferDataVar(call->args[i * 2]).value(); PrimExpr index = call->args[i * 2 + 1].as_or_throw(); const BufferVar& buffer = buffer_data_to_buffer_.at(buffer_var); auto it = buffer_remap_.find(buffer); if (it != buffer_remap_.end()) { + const BufferVar& new_buffer = (*it).second; PrimExpr new_index = RewriteWmmaFragmentIndex(buffer, (*it).second, index); + new_args.Set(i * 2, new_buffer.data()); new_args.Set(i * 2 + 1, new_index); } } @@ -175,10 +193,11 @@ class PipelineOpaqueAccessRewriter { }; ffi::Array new_args = call->args; for (int i : arg_indices) { - const BufferVar& buffer = buffer_data_to_buffer_.at(call->args[i].as_or_throw()); + const BufferVar& buffer = buffer_data_to_buffer_.at(GetBufferDataVar(call->args[i]).value()); auto it = buffer_remap_.find(buffer); if (it != buffer_remap_.end()) { const BufferVar& new_buffer = (*it).second; + new_args.Set(i, new_buffer.data()); PrimExpr old_index = call->args[i + 1].as_or_throw(); PrimExpr offset; if (new_buffer->strides.empty()) { @@ -233,7 +252,11 @@ class PipelineBodyRewriter : public StmtExprMutator { pipeline_loop_(pipeline_loop), access_all_versions_(access_all_versions), opaque_access_rewriter_(buffer_data_to_buffer_, buffer_remap_, pipeline_loop_, - fragment_info) {} + fragment_info) { + for (const auto& [_, remapped] : buffer_remap_) { + buffer_data_to_buffer_.Set(remapped.var(), remapped); + } + } private: BufferRegion RewritePipelineBufferRegion(const BufferRegion& buffer_region) const { @@ -360,6 +383,9 @@ class PipelineRewriter : public StmtExprMutator { buffer_remap_.Set(buffer, RewriteAllocBuffer(buffer, num_versions)); } } + for (const auto& [_, remapped] : buffer_remap_) { + buffer_data_to_buffer_.Set(remapped.var(), remapped); + } ordered_stmts_.resize(pipeline_info_.size()); for (const auto& pair : pipeline_info_) { @@ -394,8 +420,10 @@ class PipelineRewriter : public StmtExprMutator { // Step 3: Make a new block that contains new buffer allocations after pipeline rewriting. ffi::Array alloc_buffers; for (const auto& alloc : pipeline_allocs_) { - alloc_buffers.push_back(buffer_remap_.Get(alloc).value_or(alloc)); + BufferVar remapped = buffer_remap_.Get(alloc).value_or(alloc); + alloc_buffers.push_back(remapped); buffer_data_to_buffer_.erase(alloc.var()); + buffer_data_to_buffer_.erase(remapped.var()); } SBlock block = MakeSBlock(stmt, buffer_data_to_buffer_); block.CopyOnWrite()->alloc_buffers = std::move(alloc_buffers); diff --git a/src/s_tir/transform/inject_virtual_thread.cc b/src/s_tir/transform/inject_virtual_thread.cc index f6e4c1a7add9..e6b0396190af 100644 --- a/src/s_tir/transform/inject_virtual_thread.cc +++ b/src/s_tir/transform/inject_virtual_thread.cc @@ -38,6 +38,21 @@ namespace tvm { namespace s_tir { using namespace tvm::tirx; +namespace { + +ffi::Optional GetBufferDataVar(const ffi::Any& data) { + if (auto var = data.as()) { + return var; + } + if (const auto* call = data.as(); + call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) { + return call->args[0].as(); + } + return std::nullopt; +} + +} // namespace + // If expression is touched by var. class ExprTouched final : public StmtExprVisitor { public: @@ -62,8 +77,8 @@ class ExprTouched final : public StmtExprVisitor { void VisitExpr_(const CallNode* op) final { if (op->op.same_as(builtin::tvm_access_ptr())) { const auto* rw_mask = op->args[4].as(); - const VarNode* buffer_var = op->args[1].as(); - if (buffer_var == nullptr) { + auto buffer = GetBufferDataVar(op->args[1]); + if (!buffer.has_value()) { // Nested access pointers are valid pointer expressions. Visit the // inner pointer and this access's offset instead of assuming a raw // buffer Var at every level. @@ -71,6 +86,7 @@ class ExprTouched final : public StmtExprVisitor { this->VisitExpr(op->args[2].as_or_throw()); return; } + const VarNode* buffer_var = buffer.value().get(); TVM_FFI_ICHECK(rw_mask); // read if (rw_mask->value & 1) { @@ -212,7 +228,8 @@ class VTInjector : public arith::IRMutatorWithAnalyzer { } // Variable Expr VisitExpr_(const VarNode* op) final { - TVM_FFI_ICHECK(!alloc_remap_.count(op)) << "BufferVar address may get rewritten in virtual thread"; + TVM_FFI_ICHECK(!alloc_remap_.count(op)) + << "BufferVar address may get rewritten in virtual thread"; if (touched_var_.count(op)) { visit_touched_var_ = true; } @@ -223,22 +240,33 @@ class VTInjector : public arith::IRMutatorWithAnalyzer { } // Expression. Expr VisitExpr_(const CallNode* op) final { - if (op->op.same_as(builtin::tvm_access_ptr())) { + if (op->op.same_as(builtin::buffer_data())) { + auto buffer = GetBufferDataVar(ffi::GetRef(op)).value(); + auto it = alloc_remap_.find(buffer.get()); + if (it == alloc_remap_.end()) { + return StmtExprMutator::VisitExpr_(op); + } + visit_touched_var_ = true; + return GetRemappedBuffer(BufferVar(buffer), it->second).data(); + } else if (op->op.same_as(builtin::tvm_access_ptr())) { TVM_FFI_ICHECK_EQ(op->args.size(), 5U); PrimType dtype = op->args[0].as_or_throw().ty(); - const VarNode* buffer = op->args[1].as(); - if (buffer == nullptr) { + auto buffer = GetBufferDataVar(op->args[1]); + if (!buffer.has_value()) { return StmtExprMutator::VisitExpr_(op); } - auto it = alloc_remap_.find(buffer); + auto it = alloc_remap_.find(buffer.value().get()); if (it == alloc_remap_.end()) return StmtExprMutator::VisitExpr_(op); visit_touched_var_ = true; PrimExpr offset = this->VisitPrimExpr(op->args[2].as_or_throw()); PrimExpr extent = this->VisitPrimExpr(op->args[3].as_or_throw()); PrimExpr stride = it->second / MakeConst(offset.ty(), dtype.lanes()); offset = RewriteIndex(offset, stride); + Expr data = buffer.value()->ty.as() + ? GetRemappedBuffer(BufferVar(buffer.value()), it->second).data() + : op->args[1]; - return Call(op->ty, op->op, {op->args[0], op->args[1], offset, extent, op->args[4]}); + return Call(op->ty, op->op, {op->args[0], data, offset, extent, op->args[4]}); } else if (op->op.same_as(builtin::tvm_context_id())) { return allow_share_ ? Expr(ffi::GetRef(op)) : Expr(var_); } else { @@ -280,11 +308,11 @@ class VTInjector : public arith::IRMutatorWithAnalyzer { } BufferVar GetRemappedBuffer(BufferVar buf, PrimExpr alloc_extent) { - auto key = buf.get(); - auto it = buf_remap_.find(key); + auto it = buf_remap_.find(buf); if (it != buf_remap_.end()) { return it->second; } + BufferVar original = buf; TVM_FFI_ICHECK_EQ(buf->shape.size(), 1) << "Expected buffers being rewritten to already be flattened."; @@ -292,7 +320,7 @@ class VTInjector : public arith::IRMutatorWithAnalyzer { writer->shape = {buf->shape[0] * alloc_extent}; buf = RebuildBufferVar(buf, std::move(writer)); - buf_remap_[key] = buf; + buf_remap_[original] = buf; return buf; } @@ -455,6 +483,7 @@ class VTInjector : public arith::IRMutatorWithAnalyzer { auto type = CopyBufferType(op->buffer); type->shape = shape; BufferVar new_buffer = RebuildBufferVar(op->buffer, std::move(type)); + buf_remap_[op->buffer] = new_buffer; return AllocBuffer(new_buffer, op->annotations); } } @@ -522,7 +551,7 @@ class VTInjector : public arith::IRMutatorWithAnalyzer { * the allocated buffer size, then modifying the indices at which * each virtual thread accesses the buffer. */ - std::unordered_map buf_remap_; + std::unordered_map buf_remap_; }; class VirtualThreadInjector : public arith::IRMutatorWithAnalyzer { diff --git a/src/s_tir/transform/lower_async_dma.cc b/src/s_tir/transform/lower_async_dma.cc index b4ecd26b3ac1..15ad74203595 100644 --- a/src/s_tir/transform/lower_async_dma.cc +++ b/src/s_tir/transform/lower_async_dma.cc @@ -77,16 +77,16 @@ class AsyncDMALowerer : public arith::IRMutatorWithAnalyzer { auto src = BufferLoad(mem_copy->source->buffer, {src_min}); auto dst = BufferLoad(mem_copy->dest->buffer, {dst_min}); PrimExpr dst_nbytes = dst_extent * static_cast(src.ty().StorageBytes()); - return Evaluate( - Call(PrimType::Int(32), builtin::dma_copy(), - ffi::Array{PrimExpr(async_queue_id_.value()), - Call(mem_copy->dest->buffer->data_pointer_type, builtin::address_of(), - ffi::Array{dst}, Attrs(), {}, Span()), - Call(mem_copy->source->buffer->data_pointer_type, builtin::address_of(), - ffi::Array{src}, Attrs(), {}, Span()), - dst_nbytes, PrimExpr(dma_bypass_cache_)}, - Attrs(), {}, Span()) - .as_or_throw()); + return Evaluate(Call(PrimType::Int(32), builtin::dma_copy(), + ffi::Array{ + PrimExpr(async_queue_id_.value()), + Call(mem_copy->dest->buffer.DataPointerType(), builtin::address_of(), + ffi::Array{dst}, Attrs(), {}, Span()), + Call(mem_copy->source->buffer.DataPointerType(), builtin::address_of(), + ffi::Array{src}, Attrs(), {}, Span()), + dst_nbytes, PrimExpr(dma_bypass_cache_)}, + Attrs(), {}, Span()) + .as_or_throw()); } Stmt VisitStmt_(const AttrStmtNode* op) final { diff --git a/src/s_tir/transform/lower_cross_thread_reduction.cc b/src/s_tir/transform/lower_cross_thread_reduction.cc index 5564a1eb2132..620457d7e5b2 100644 --- a/src/s_tir/transform/lower_cross_thread_reduction.cc +++ b/src/s_tir/transform/lower_cross_thread_reduction.cc @@ -141,20 +141,19 @@ bool IsReductionBlock(const SBlockRealize& realize, const ffi::Map& * \return The created buffers */ ffi::Array MakeScratchpads(const ffi::Array& reduction_buffers, - bool is_cross_thread_buffer) { + bool is_cross_thread_buffer) { ffi::Array new_buffers; new_buffers.reserve(reduction_buffers.size()); for (const BufferVar& buffer : reduction_buffers) { ffi::String name = is_cross_thread_buffer ? "cross" : "in"; name = name + "_thread_" + buffer.name(); - new_buffers.push_back(BufferVar(/*ptr=*/Var(name, PointerType(buffer->dtype, "local")), - /*dtype=*/buffer->dtype, - /*shape=*/{IntImm::Int32(1)}, - /*strides=*/{IntImm::Int32(1)}, - /*elem_offset=*/PrimExpr{nullptr}, - /*name=*/name, - /*data_alignment=*/0, - /*offset_factor=*/0)); + new_buffers.push_back(BufferVar(name, BufferType(/*storage_scope=*/"local", + /*dtype=*/buffer->dtype, + /*shape=*/{IntImm::Int32(1)}, + /*strides=*/{IntImm::Int32(1)}, + /*elem_offset=*/PrimExpr{nullptr}, + /*data_alignment=*/0, + /*offset_factor=*/0))); } return new_buffers; } @@ -297,13 +296,13 @@ class InThreadReducerMaker : private StmtMutator { * \param combiner_rhs The RHS values of the combiner * \param reduction_loops The reduction loops */ -Stmt TransformReductionBlock(const SBlockRealizeNode* realize, // +Stmt TransformReductionBlock(const SBlockRealizeNode* realize, // const ffi::Optional>& it_buffers, // const ffi::Array& ct_buffers, // const ffi::Array& wb_buffers, // - const ffi::Array& old_wb_indices, // - const CommReducer& reducer, // - const ffi::Array& combiner_rhs, // + const ffi::Array& old_wb_indices, // + const CommReducer& reducer, // + const ffi::Array& combiner_rhs, // const std::vector& reduction_loops) { int n_buffers = wb_buffers.size(); const SBlockNode* block = realize->block.get(); @@ -928,8 +927,7 @@ class CrossThreadReductionTransformer : public StmtMutator { int block_idx_depth = 0; int thread_idx_depth = 0; - std::unordered_map>> - crt_buf2threads_; + std::unordered_map>> crt_buf2threads_; }; namespace transform { diff --git a/src/s_tir/transform/lower_match_buffer.cc b/src/s_tir/transform/lower_match_buffer.cc index 7163a5f7f2bb..98343099ea1b 100644 --- a/src/s_tir/transform/lower_match_buffer.cc +++ b/src/s_tir/transform/lower_match_buffer.cc @@ -106,6 +106,19 @@ class MatchBufferLower : public StmtExprMutator { } } + Expr VisitExpr_(const CallNode* op) final { + if (op->op.same_as(builtin::buffer_data()) && op->args.size() == 1) { + if (auto var = op->args[0].as(); + var.has_value() && var.value()->ty.as()) { + auto it = match_buffers_.find(BufferVar(var.value())); + if (it != match_buffers_.end()) { + return (*it).second->buffer.data(); + } + } + } + return StmtExprMutator::VisitExpr_(op); + } + Stmt VisitStmt_(const BufferStoreNode* op) final { // Save the original buffer before base class mutation may remap it BufferVar orig_buffer = op->buffer; diff --git a/src/s_tir/transform/lower_thread_allreduce.cc b/src/s_tir/transform/lower_thread_allreduce.cc index 06dde28c34af..1b288725bb5e 100644 --- a/src/s_tir/transform/lower_thread_allreduce.cc +++ b/src/s_tir/transform/lower_thread_allreduce.cc @@ -41,12 +41,37 @@ namespace tvm { namespace s_tir { using namespace tvm::tirx; +namespace { + +ffi::Optional GetBufferDataVar(const ffi::Any& data) { + if (auto var = data.as()) { + return var; + } + if (const auto* call = data.as(); + call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) { + return call->args[0].as(); + } + return std::nullopt; +} + +} // namespace + class ThreadAllreduceBuilder final : public StmtExprMutator { public: - explicit ThreadAllreduceBuilder(const TargetNode* target) + explicit ThreadAllreduceBuilder(const TargetNode* target, const ffi::Array& params, + const ffi::Map& buffer_map) : target_(target), warp_size_(target->GetAttr("thread_warp_size", 1).value()), - max_num_threads_(target->GetAttr("max_num_threads", -1).value()) {} + max_num_threads_(target->GetAttr("max_num_threads", -1).value()) { + for (const auto& [_, buffer] : buffer_map) { + buffer_aliases_.Set(buffer.var(), buffer.var()); + } + for (const Var& param : params) { + if (param->ty.as()) { + buffer_aliases_.Set(param, param); + } + } + } Stmt VisitStmt_(const AttrStmtNode* op) final { if (op->attr_key == tirx::attr::thread_extent) { @@ -76,6 +101,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { } } Stmt VisitStmt_(const AllocBufferNode* op) final { + buffer_aliases_.Set(op->buffer.var(), op->buffer.var()); // In flat IR, alloc_remap_ may not yet be populated when this AllocBuffer is visited // (the remap is set up by MakeAllreduce which runs during AttrStmt/Evaluate visit // that appears later in the sequence). We record the original data pointer and @@ -109,20 +135,15 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { } ffi::Optional GetRemappedBuffer(const BufferVar& buf) { - if (auto it = buf_remap_.find(buf.get()); it != buf_remap_.end()) { - return it->second; - } - if (auto it = var_remap_.find(buf.get()); it != var_remap_.end()) { - BufferVar new_buf(it->second); - buf_remap_[buf.get()] = new_buf; - return new_buf; + return BufferVar(it->second); } return std::nullopt; } Stmt VisitStmt_(const DeclBufferNode* op) final { + RegisterBufferAlias(op->buffer, op->data); auto node = StmtExprMutator::VisitStmt_(op).as_or_throw(); if (auto buf = GetRemappedBuffer(node->buffer)) { node.CopyOnWrite()->buffer = buf.value(); @@ -401,9 +422,10 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { // The AllocBuffer doesn't need to be emitted here since alloc_remap_ // will cause the existing allocation to be rewritten in VisitStmt_(AllocBufferNode*). - alloc_remap_[buffers[i].get()] = buf; + const VarNode* alloc_key = GetAllocationKey(buffers[i].get()); + alloc_remap_[alloc_key] = buf; + var_remap_[alloc_key] = buf.var(); var_remap_[buffers[i].get()] = buf.var(); - buf_remap_[buffers[i].get()] = buf; } } else { std::vector shared_bufs(size); @@ -434,9 +456,10 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { {BufIndex(IntImm(reduce_index.ty(), 0), group_index, reduce_extent)}); TVM_FFI_ICHECK_EQ(load.ty(), dtypes[idx]); load_remap_[buffers[idx].get()] = load; - alloc_remap_[buffers[idx].get()] = shared_bufs[idx]; + const VarNode* alloc_key = GetAllocationKey(buffers[idx].get()); + alloc_remap_[alloc_key] = shared_bufs[idx]; + var_remap_[alloc_key] = shared_bufs[idx].var(); var_remap_[buffers[idx].get()] = shared_bufs[idx].var(); - buf_remap_[buffers[idx].get()] = shared_bufs[idx]; } } @@ -801,6 +824,19 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { } } + private: + void RegisterBufferAlias(BufferVar buffer, const Expr& data) { + Var root = buffer.var(); + if (auto source = GetBufferDataVar(data); + source.has_value() && source.value()->ty.as()) { + auto source_root = buffer_aliases_.Get(source.value()); + TVM_FFI_ICHECK(source_root.has_value()) << "Buffer alias source " << source.value()->name + << " must be registered before its DeclBuffer alias"; + root = source_root.value(); + } + buffer_aliases_.Set(buffer.var(), root); + } + // The target. const TargetNode* target_ = nullptr; @@ -820,15 +856,23 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { arith::Analyzer analyzer_; public: + const VarNode* GetAllocationKey(const VarNode* buffer) const { + if (buffer->ty.as()) { + Var var = ffi::GetRef(buffer); + return buffer_aliases_.Get(var).value_or(var).get(); + } + return buffer; + } + // These members are public for post-processing by DeferredRemapper. // Allocate remap std::unordered_map alloc_remap_; // BufferVar remap std::unordered_map var_remap_; - // BufferVar remap - std::unordered_map buf_remap_; // Pending AllocBuffer original data pointers (for flat IR deferred remapping) std::vector pending_alloc_buffers_; + // Physical roots of buffer aliases, flattened at each declaration. + ffi::Map buffer_aliases_; }; namespace transform { @@ -845,9 +889,9 @@ class DeferredRemapper : public StmtExprMutator { public: DeferredRemapper(const std::unordered_map& alloc_remap, const std::unordered_map& var_remap, - const std::unordered_map& buf_remap, + const ffi::Map& buffer_aliases, const std::vector& pending) - : alloc_remap_(alloc_remap), var_remap_(var_remap), buf_remap_(buf_remap) { + : alloc_remap_(alloc_remap), var_remap_(var_remap), buffer_aliases_(buffer_aliases) { for (const VarNode* ptr : pending) { pending_set_.insert(ptr); } @@ -879,11 +923,8 @@ class DeferredRemapper : public StmtExprMutator { } Stmt VisitStmt_(const DeclBufferNode* op) final { - // If the DeclBuffer's original data var was remapped by alloc_remap_, - // the corresponding AllocBuffer was also remapped, making this DeclBuffer - // redundant. Remove it by replacing with a no-op. - const VarNode* orig_data = op->buffer.get(); - if (pending_set_.count(orig_data) && alloc_remap_.count(orig_data)) { + const VarNode* root = buffer_aliases_.Get(op->buffer.var()).value_or(op->buffer.var()).get(); + if (pending_set_.count(root) && alloc_remap_.count(root)) { return Evaluate(0); } auto node = StmtExprMutator::VisitStmt_(op).as_or_throw(); @@ -895,9 +936,6 @@ class DeferredRemapper : public StmtExprMutator { private: ffi::Optional GetRemappedBuffer(const BufferVar& buf) { - if (auto it = buf_remap_.find(buf.get()); it != buf_remap_.end()) { - return it->second; - } if (auto it = var_remap_.find(buf.get()); it != var_remap_.end()) { return BufferVar(it->second); } @@ -906,7 +944,7 @@ class DeferredRemapper : public StmtExprMutator { const std::unordered_map& alloc_remap_; const std::unordered_map& var_remap_; - const std::unordered_map& buf_remap_; + const ffi::Map& buffer_aliases_; std::unordered_set pending_set_; }; @@ -916,11 +954,11 @@ Pass LowerThreadAllreduce() { auto target = f->GetAttr(tvm::attr::kTarget); TVM_FFI_ICHECK(target.has_value()) << "LowerThreadAllreduce: Require the target attribute"; const TargetNode* target_node = target.as(); - ThreadAllreduceBuilder thread_all_reduce(target_node); + ThreadAllreduceBuilder thread_all_reduce(target_node, f->params, f->buffer_map); n->body = thread_all_reduce(n->body); // Post-process: apply deferred remappings for flat IR DeferredRemapper remapper(thread_all_reduce.alloc_remap_, thread_all_reduce.var_remap_, - thread_all_reduce.buf_remap_, + thread_all_reduce.buffer_aliases_, thread_all_reduce.pending_alloc_buffers_); if (remapper.HasPendingRemaps()) { n->body = remapper(n->body); diff --git a/src/s_tir/transform/lower_vtcm_alloc.cc b/src/s_tir/transform/lower_vtcm_alloc.cc index 002ef491dc8f..b2459ba5027a 100644 --- a/src/s_tir/transform/lower_vtcm_alloc.cc +++ b/src/s_tir/transform/lower_vtcm_alloc.cc @@ -46,8 +46,7 @@ class VtcmAllocator : public StmtExprMutator { args.push_back( Call(PointerType(PrimType::Int(64)), builtin::tvm_stack_make_shape(), op->buffer->shape)); return DeclBuffer( - op->buffer, - Call(op->buffer->data_pointer_type, builtin::nd_mem_alloc_with_scope(), args)); + op->buffer, Call(op->buffer.DataPointerType(), builtin::nd_mem_alloc_with_scope(), args)); } return StmtExprMutator::VisitStmt_(op); } @@ -55,7 +54,7 @@ class VtcmAllocator : public StmtExprMutator { protected: std::string GetStorageScope(const Var& var) { auto* ptr = var->ty.as(); - TVM_FFI_ICHECK(ptr) << "BufferVar Var's type annotation must be of PointerType"; + TVM_FFI_ICHECK(ptr) << "Buffer Var's type annotation must be of PointerType"; return ptr->storage_scope; } }; 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 a515a6583013..5dd428dece42 100644 --- a/src/s_tir/transform/manifest_shared_memory_local_stage.cc +++ b/src/s_tir/transform/manifest_shared_memory_local_stage.cc @@ -269,8 +269,9 @@ class SharedMemoryLocalStageInserter : public StmtMutator { std::vector ancestor_loop_or_blocks_; // ancestor loops or block realize ffi::Map buffer_remap_; // mapping from the target buffer to the intermediate buffer - ffi::Map buffer_local_stage_; // mapping from the target buffer to the local stage - ffi::Array target_buffers_; // the target buffers for rewriting + ffi::Map + buffer_local_stage_; // mapping from the target buffer to the local stage + ffi::Array target_buffers_; // the target buffers for rewriting }; namespace transform { diff --git a/src/s_tir/transform/memhammer_coalesce.cc b/src/s_tir/transform/memhammer_coalesce.cc index 728bcace443a..a628c9574838 100644 --- a/src/s_tir/transform/memhammer_coalesce.cc +++ b/src/s_tir/transform/memhammer_coalesce.cc @@ -51,7 +51,7 @@ Stmt FuseNestLoops(Stmt body) { if (auto replacement = subst_map.Get(v)) { return replacement.value(); } - return v.as_or_throw(); + return std::nullopt; }; PrimExpr fused_extent = 1; for (int i = 0; i < n; i++) { diff --git a/src/s_tir/transform/memhammer_lower_auto_copy.cc b/src/s_tir/transform/memhammer_lower_auto_copy.cc index 4caa23e1c075..b018b4ef7c61 100644 --- a/src/s_tir/transform/memhammer_lower_auto_copy.cc +++ b/src/s_tir/transform/memhammer_lower_auto_copy.cc @@ -194,7 +194,8 @@ class AutoPadder { Stmt RewriteBufferAccess(const Stmt& stmt) { class Rewriter : public StmtExprMutator { public: - explicit Rewriter(const ffi::Map& buffer_map) : buffer_map_(buffer_map) {} + explicit Rewriter(const ffi::Map& buffer_map) + : buffer_map_(buffer_map) {} private: Expr VisitExpr_(const BufferLoadNode* _op) final { @@ -467,14 +468,14 @@ class AutoPadder { if (v.same_as(var)) { return IntImm::Int32(0); } else { - return v.as_or_throw(); + return std::nullopt; } }); PrimExpr e2 = Substitute(e, [var](const Var& v) -> ffi::Optional { if (v.same_as(var)) { return IntImm::Int32(1); } else { - return v.as_or_throw(); + return std::nullopt; } }); arith::Analyzer analyzer; diff --git a/src/s_tir/transform/memhammer_rewrite_rule.h b/src/s_tir/transform/memhammer_rewrite_rule.h index 124f2ec58c6c..cadb5ec1bd72 100644 --- a/src/s_tir/transform/memhammer_rewrite_rule.h +++ b/src/s_tir/transform/memhammer_rewrite_rule.h @@ -251,7 +251,8 @@ class WmmaToShared : public RewriteRule { */ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::String storage_scope, ffi::Optional compute_location, - const ffi::Array& outer_loops, BufferVar* alloc_buffer); + const ffi::Array& outer_loops, + BufferVar* alloc_buffer); } // namespace s_tir } // namespace tvm diff --git a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc index ab24b0483a22..09906d5e8b86 100644 --- a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc +++ b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc @@ -135,23 +135,21 @@ Stmt RewriteWmmaLoad(Stmt stmt) { BufferVar tgt_buffer = buf_store->buffer; std::string layout = tgt_buffer.scope() == "wmma.matrix_a" ? "row_major" : "col_major"; BufferVar new_src_buffer( - /*data=*/Var("src", PointerType(dtype_ty, src_buffer.scope())), - /*dtype=*/dtype, - /*shape=*/{IntImm::Int32(16), IntImm::Int32(16)}, - /*strides=*/{PrimVar("s1", int32_ty), PrimVar("s0", int32_ty)}, - /*elem_offset=*/PrimVar("src_elem_offset", int32_ty), - /*name=*/"src", - /*data_alignment=*/64, - /*offset_factor=*/16); + /*name=*/"src", BufferType(/*storage_scope=*/src_buffer.scope(), + /*dtype=*/dtype, + /*shape=*/{IntImm::Int32(16), IntImm::Int32(16)}, + /*strides=*/{PrimVar("s1", int32_ty), PrimVar("s0", int32_ty)}, + /*elem_offset=*/PrimVar("src_elem_offset", int32_ty), + /*data_alignment=*/64, + /*offset_factor=*/16)); BufferVar new_tgt_buffer( - /*data=*/Var("tgt", PointerType(dtype_ty, tgt_buffer.scope())), - /*dtype=*/dtype, - /*shape=*/{IntImm::Int32(16), IntImm::Int32(16)}, - /*strides=*/{}, - /*elem_offset=*/PrimVar("tgt_elem_offset", int32_ty), - /*name=*/"tgt", - /*data_alignment=*/64, - /*offset_factor=*/16); + /*name=*/"tgt", BufferType(/*storage_scope=*/tgt_buffer.scope(), + /*dtype=*/dtype, + /*shape=*/{IntImm::Int32(16), IntImm::Int32(16)}, + /*strides=*/{}, + /*elem_offset=*/PrimVar("tgt_elem_offset", int32_ty), + /*data_alignment=*/64, + /*offset_factor=*/16)); ffi::Array read_region = RelaxIndices(buf_load->indices, src_buffer->shape, var_dom); ffi::Array write_region = RelaxIndices(buf_store->indices, tgt_buffer->shape, var_dom); static const Op& tvm_load_matrix_sync_op = Op::Get("tirx.tvm_load_matrix_sync"); @@ -244,22 +242,13 @@ Stmt RewriteWmmaStore(Stmt stmt) { PrimType dtype_ty = src_buffer->dtype; const PrimType& dtype = dtype_ty; - BufferVar new_src_buffer(/*data=*/Var("src", PointerType(dtype_ty, src_buffer.scope())), - /*dtype=*/dtype, - /*shape=*/{IntImm::Int32(16), IntImm::Int32(16)}, - /*strides=*/{}, - /*elem_offset=*/PrimVar("src_elem_offset", int32_ty), - /*name=*/"src", - /*data_alignment=*/64, - /*offset_factor=*/16); - BufferVar new_tgt_buffer(/*data=*/Var("tgt", PointerType(dtype_ty, tgt_buffer.scope())), - /*dtype=*/dtype, - /*shape=*/{IntImm::Int32(16), IntImm::Int32(16)}, - /*strides=*/{PrimVar("s1", int32_ty), PrimVar("s0", int32_ty)}, - /*elem_offset=*/PrimVar("tgt_elem_offset", int32_ty), - /*name=*/"tgt", - /*data_alignment=*/64, - /*offset_factor=*/16); + BufferVar new_src_buffer( + "src", BufferType(src_buffer.scope(), dtype, {IntImm::Int32(16), IntImm::Int32(16)}, {}, + PrimVar("src_elem_offset", int32_ty), 64, 16)); + BufferVar new_tgt_buffer( + "tgt", BufferType(tgt_buffer.scope(), dtype, {IntImm::Int32(16), IntImm::Int32(16)}, + {PrimVar("s1", int32_ty), PrimVar("s0", int32_ty)}, + PrimVar("tgt_elem_offset", int32_ty), 64, 16)); ffi::Array read_region = RelaxIndices(buf_load->indices, src_buffer->shape, var_dom); ffi::Array write_region = RelaxIndices(buf_store->indices, tgt_buffer->shape, var_dom); @@ -470,22 +459,13 @@ Stmt RewriteMmaStore(Stmt stmt) { BufferVar tgt_buffer = buf_store->buffer; PrimType dtype_ty = src_buffer->dtype; const PrimType& dtype = dtype_ty; - BufferVar new_src_buffer(/*data=*/Var("src", PointerType(dtype_ty, src_buffer.scope())), - /*dtype=*/dtype, - /*shape=*/{IntImm::Int32(8), IntImm::Int32(8)}, - /*strides=*/{}, - /*elem_offset=*/PrimVar("src_elem_offset", int32_ty), - /*name=*/"src", - /*data_alignment=*/64, - /*offset_factor=*/8); - BufferVar new_tgt_buffer(/*data=*/Var("tgt", PointerType(dtype_ty, tgt_buffer.scope())), - /*dtype=*/dtype, - /*shape=*/{IntImm::Int32(8), IntImm::Int32(8)}, - /*strides=*/{PrimVar("s1", int32_ty), PrimVar("s0", int32_ty)}, - /*elem_offset=*/PrimVar("tgt_elem_offset", int32_ty), - /*name=*/"tgt", - /*data_alignment=*/64, - /*offset_factor=*/8); + BufferVar new_src_buffer( + "src", BufferType(src_buffer.scope(), dtype, {IntImm::Int32(8), IntImm::Int32(8)}, {}, + PrimVar("src_elem_offset", int32_ty), 64, 8)); + BufferVar new_tgt_buffer( + "tgt", BufferType(tgt_buffer.scope(), dtype, {IntImm::Int32(8), IntImm::Int32(8)}, + {PrimVar("s1", int32_ty), PrimVar("s0", int32_ty)}, + PrimVar("tgt_elem_offset", int32_ty), 64, 8)); // Step 3.2. Generate new r/w region ffi::Array read_region = RelaxIndices(buf_load->indices, src_buffer->shape, var_dom); diff --git a/src/s_tir/transform/merge_shared_memory_allocations.cc b/src/s_tir/transform/merge_shared_memory_allocations.cc index 93ccb33ee647..128ec729e283 100644 --- a/src/s_tir/transform/merge_shared_memory_allocations.cc +++ b/src/s_tir/transform/merge_shared_memory_allocations.cc @@ -50,6 +50,21 @@ using namespace tvm::tirx; using runtime::StorageRank; using runtime::StorageScope; +namespace { + +ffi::Optional GetBufferDataVar(const ffi::Any& data) { + if (auto var = data.as()) { + return var; + } + if (const auto* call = data.as(); + call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) { + return call->args[0].as(); + } + return std::nullopt; +} + +} // namespace + bool IsDynamicSharedMemory(Var buffer_var) { StorageScope storage_scope = runtime::StorageScope::Create(GetPtrStorageScope(buffer_var)); return storage_scope.rank == runtime::StorageRank::kShared && storage_scope.tag == ".dyn"; @@ -159,7 +174,7 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { // visit subexpr StmtExprVisitor::VisitStmt_(op); // Add write access. - const VarNode* buf = op->buffer.get(); + const VarNode* buf = ResolveAlias(op->buffer.get()); auto it = alloc_info_.find(buf); if (it != alloc_info_.end() && it->second.buffer.defined()) { TVM_FFI_ICHECK_LT(it->second.level, scope_.size()); @@ -187,10 +202,21 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { } } + void VisitStmt_(const DeclBufferNode* op) final { + if (auto source = GetBufferDataVar(op->data)) { + const VarNode* allocation = ResolveAlias(source.value().get()); + if (alloc_info_.count(allocation)) { + buffer_alias_sources_[op->buffer.get()] = allocation; + return; + } + } + StmtExprVisitor::VisitStmt_(op); + } + void VisitExpr_(const BufferLoadNode* op) final { // Add read access. StmtExprVisitor::VisitExpr_(op); - const VarNode* buf = op->buffer.get(); + const VarNode* buf = ResolveAlias(op->buffer.get()); auto it = alloc_info_.find(buf); if (it != alloc_info_.end() && it->second.buffer.defined()) { TVM_FFI_ICHECK_LT(it->second.level, scope_.size()) @@ -217,6 +243,7 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { void VisitExpr_(const VarNode* buf) final { // Directly reference to the variable count as a read. + buf = ResolveAlias(buf); auto it = alloc_info_.find(buf); if (it != alloc_info_.end() && it->second.buffer.defined()) { TVM_FFI_ICHECK_LT(it->second.level, scope_.size()); @@ -274,8 +301,19 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { std::vector linear_seq_; // The storage scope of each buffer std::unordered_map alloc_info_; + // Typed views of an allocated shared-memory buffer. + std::unordered_map buffer_alias_sources_; private: + const VarNode* ResolveAlias(const VarNode* buffer) const { + while (true) { + auto it = buffer_alias_sources_.find(buffer); + if (it == buffer_alias_sources_.end()) break; + buffer = it->second; + } + return buffer; + } + // Wrapper function to determine if the shared memory allocation for a variable is appropriate. bool IsAppropriateSharedMemory(const Var& var) { return is_dynamic_ ? IsDynamicSharedMemory(var) : IsStaticSharedMemory(var); @@ -336,6 +374,10 @@ class SharedMemoryRewriter : public StmtExprMutator { std::unordered_map buffer_byte_offsets; // BufferVar-object remap: original BufferVar -> merged-data-var BufferVar. std::unordered_map buffer_remap; + // Typed views whose physical source is one of shmem_allocs. + std::unordered_map buffer_alias_sources; + // Remapped buffers in first-use order, for deterministic alias emission. + std::vector buffer_remap_order; // Has any original alloc in this scope been marked volatile? bool has_volatile_alloc{false}; // Liveness data (event_map, alloc_map, const_free_map, sym_free_list) — all per-scope. @@ -390,7 +432,9 @@ class SharedMemoryRewriter : public StmtExprMutator { // 5. Recursively mutate the body — reads scope_stack_.back() for all rewrites. Stmt visited_body = StmtExprMutator::VisitStmt(op->body); - for (const auto& [_, remapped] : scope.buffer_remap) { + for (const BufferVar& remapped : scope.buffer_remap_order) { + // The uint8 merged allocation intentionally supplies storage for + // typed views; target codegen emits the required pointer cast. visited_body = SeqStmt::Flatten(DeclBuffer(remapped, scope.merged_buffer.data()), visited_body); } @@ -437,6 +481,16 @@ class SharedMemoryRewriter : public StmtExprMutator { } Stmt VisitStmt_(const DeclBufferNode* op) final { + if (!scope_stack_.empty()) { + if (auto source = GetBufferDataVar(op->data)) { + KernelScope& scope = scope_stack_.back(); + if (const VarNode* allocation = ResolveAllocation(source.value().get(), scope)) { + scope.buffer_alias_sources[op->buffer.get()] = allocation; + return Evaluate(0); + } + } + } + auto node = StmtExprMutator::VisitStmt_(op).as_or_throw(); if (auto new_buf = GetUpdatedBuffer(node->buffer); !new_buf.same_as(node->buffer)) { node.CopyOnWrite()->buffer = new_buf; @@ -457,14 +511,13 @@ class SharedMemoryRewriter : public StmtExprMutator { template Node VisitBufferAccess(Node node) { if (IsAppropriateSharedMemory(node->buffer) && !scope_stack_.empty() && - scope_stack_.back().shmem_allocs.count(node->buffer.get())) { + ResolveAllocation(node->buffer.get(), scope_stack_.back())) { TVM_FFI_ICHECK_EQ(node->indices.size(), 1) << "MergeSharedMemoryAllocations expects flat memory buffers, " << "and is to be run after " << "FlattenBuffer"; ffi::Array indices = { - node->indices[0] + - this->GetBufferOffset(node->buffer.var(), node->buffer->dtype->dtype)}; + node->indices[0] + this->GetBufferOffset(node->buffer.var(), node->buffer->dtype->dtype)}; auto writer = node.CopyOnWrite(); writer->buffer = GetUpdatedBuffer(node->buffer); @@ -477,7 +530,7 @@ class SharedMemoryRewriter : public StmtExprMutator { BufferVar GetUpdatedBuffer(BufferVar buffer) { if (scope_stack_.empty()) return buffer; KernelScope& scope = scope_stack_.back(); - if (!scope.shmem_allocs.count(buffer.get())) return buffer; + if (!ResolveAllocation(buffer.get(), scope)) return buffer; auto key = buffer.get(); auto it = scope.buffer_remap.find(key); @@ -487,7 +540,7 @@ class SharedMemoryRewriter : public StmtExprMutator { if (IsAppropriateSharedMemory(buffer)) { TVM_FFI_ICHECK_EQ(buffer->shape.size(), 1) - << "BufferVar " << buffer << " has shape " << buffer->shape << ". " + << "Buffer " << buffer << " has shape " << buffer->shape << ". " << "MergeSharedMemoryAllocations expects flat memory buffers, " << "and is to be run after " << "FlattenBuffer"; @@ -495,6 +548,7 @@ class SharedMemoryRewriter : public StmtExprMutator { } scope.buffer_remap[key] = buffer; + scope.buffer_remap_order.push_back(buffer); return buffer; } @@ -503,37 +557,61 @@ class SharedMemoryRewriter : public StmtExprMutator { if (op->op.same_as(builtin::tvm_access_ptr())) { TVM_FFI_ICHECK_EQ(op->args.size(), 5U); DLDataType dtype = op->args[0].as_or_throw().ty()->dtype; - auto buffer_opt = op->args[1].as(); + auto buffer_opt = GetBufferDataVar(op->args[1]); if (!buffer_opt.has_value()) { return StmtExprMutator::VisitExpr_(op); } Var buffer = buffer_opt.value(); - if (!IsAppropriateSharedMemory(buffer) || scope_stack_.empty() || - !scope_stack_.back().shmem_allocs.count(buffer.get())) { + bool is_shared = buffer->ty.as() + ? IsAppropriateSharedMemory(BufferVar(buffer)) + : IsAppropriateSharedMemory(buffer); + if (!is_shared || scope_stack_.empty() || + !ResolveAllocation(buffer.get(), scope_stack_.back())) { return StmtExprMutator::VisitExpr_(op); } PrimExpr extra_offset = GetBufferOffset(buffer, dtype); + Expr merged_data = buffer->ty.as() + ? GetUpdatedBuffer(BufferVar(buffer)).data() + : scope_stack_.back().merged_buffer.data(); PrimExpr offset = this->VisitPrimExpr(op->args[2].as_or_throw()); PrimExpr extent = this->VisitPrimExpr(op->args[3].as_or_throw()); return Call(op->ty, op->op, - {op->args[0], scope_stack_.back().merged_buffer.data(), - extra_offset + offset, extent, - op->args[4]}); + {op->args[0], merged_data, extra_offset + offset, extent, op->args[4]}); } else if (op->op.same_as(ptx_cp_async_op)) { TVM_FFI_ICHECK((op->args.size() == 5U) || (op->args.size() == 6U)); - Var buffer = op->args[0].as_or_throw(); - const auto* ptr_type = buffer->ty.as(); - TVM_FFI_ICHECK(ptr_type) << "The buffer should be a pointer type."; - const auto* prim_type = ptr_type->element_type.as(); - TVM_FFI_ICHECK(prim_type) << "The buffer should be a pointer to a primitive type."; - DLDataType dtype = prim_type->dtype; - if (!IsAppropriateSharedMemory(buffer) || scope_stack_.empty() || - !scope_stack_.back().shmem_allocs.count(buffer.get())) { + auto buffer_opt = GetBufferDataVar(op->args[0]); + if (!buffer_opt.has_value()) { + return StmtExprMutator::VisitExpr_(op); + } + Var buffer = buffer_opt.value(); + DLDataType dtype; + bool is_shared; + if (buffer->ty.as()) { + BufferVar typed_buffer(buffer); + dtype = typed_buffer->dtype->dtype; + is_shared = IsAppropriateSharedMemory(typed_buffer); + } else { + const auto* ptr_type = buffer->ty.as(); + TVM_FFI_ICHECK(ptr_type) << "The buffer should be a pointer type."; + const auto* prim_type = ptr_type->element_type.as(); + TVM_FFI_ICHECK(prim_type) << "The buffer should be a pointer to a primitive type."; + dtype = prim_type->dtype; + is_shared = IsAppropriateSharedMemory(buffer); + } + if (!is_shared || scope_stack_.empty() || + !ResolveAllocation(buffer.get(), scope_stack_.back())) { return StmtExprMutator::VisitExpr_(op); } PrimExpr extra_offset = GetBufferOffset(buffer, dtype); PrimExpr offset = this->VisitPrimExpr(op->args[1].as_or_throw()); + if (buffer->ty.as()) { + Expr merged_data = GetUpdatedBuffer(BufferVar(buffer)).data(); + ffi::Array args = op->args; + args.Set(0, merged_data); + args.Set(1, extra_offset + offset); + return Call(op->ty, op->op, args, op->attrs, {}, op->span); + } // the dst shared memory is a byte buffer generated by merging shared memory. // we need to multiply the offset index by the byte size of the original value dtype, to get // the correct offset of merged shared buffer. @@ -559,12 +637,23 @@ class SharedMemoryRewriter : public StmtExprMutator { PrimExpr GetBufferOffset(Var buffer_var, DLDataType dtype) { TVM_FFI_ICHECK(!scope_stack_.empty()); KernelScope& scope = scope_stack_.back(); - auto it = scope.buffer_byte_offsets.find(buffer_var.get()); + const VarNode* allocation = ResolveAllocation(buffer_var.get(), scope); + TVM_FFI_ICHECK(allocation); + auto it = scope.buffer_byte_offsets.find(allocation); TVM_FFI_ICHECK(it != scope.buffer_byte_offsets.end()); int elem_bytes = (static_cast(dtype.bits) * static_cast(dtype.lanes) + 7) / 8; return indexdiv(it->second, elem_bytes); } + const VarNode* ResolveAllocation(const VarNode* buffer, const KernelScope& scope) const { + while (true) { + auto it = scope.buffer_alias_sources.find(buffer); + if (it == scope.buffer_alias_sources.end()) break; + buffer = it->second; + } + return scope.shmem_allocs.count(buffer) ? buffer : nullptr; + } + // Wrapper function to determine if the shared memory allocation for a variable is appropriate. bool IsAppropriateSharedMemory(const Var& var) { return is_dynamic_ ? IsDynamicSharedMemory(var) : IsStaticSharedMemory(var); diff --git a/src/s_tir/transform/renew_defs.cc b/src/s_tir/transform/renew_defs.cc index 0a88920b4e5a..519a82d8be35 100644 --- a/src/s_tir/transform/renew_defs.cc +++ b/src/s_tir/transform/renew_defs.cc @@ -19,7 +19,7 @@ /*! * \file renew_defs.cc - * \brief Renew the definition nodes for a TIR, including Var, BufferVar and IterVar. + * \brief Renew the definition nodes for a TIR, including Var, Buffer and IterVar. */ #include diff --git a/src/s_tir/transform/storage_access.cc b/src/s_tir/transform/storage_access.cc index 57b3ce33cbdd..4db74ee41b88 100644 --- a/src/s_tir/transform/storage_access.cc +++ b/src/s_tir/transform/storage_access.cc @@ -35,6 +35,21 @@ namespace tvm { namespace s_tir { using namespace tvm::tirx; +namespace { + +ffi::Optional GetBufferDataVar(const ffi::Any& data) { + if (auto var = data.as()) { + return var; + } + if (const auto* call = data.as(); + call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) { + return call->args[0].as(); + } + return std::nullopt; +} + +} // namespace + void StorageAccessVisitor::VisitExpr_(const BufferLoadNode* op) { Var buf = op->buffer.var(); StorageScope scope = StorageScope::Create(op->buffer.scope()); @@ -111,7 +126,10 @@ void StorageAccessVisitor::VisitStmt_(const BindNode* op) { void StorageAccessVisitor::VisitStmt_(const AttrStmtNode* op) { if (op->attr_key == s_tir::attr::double_buffer_write) { TVM_FFI_ICHECK(double_buffer_write_ == nullptr); - double_buffer_write_ = op->node.as(); + auto buffer = GetBufferDataVar(op->node); + TVM_FFI_ICHECK(buffer.has_value()) + << "Expected a buffer data expression for double-buffer writes, but received " << op->node; + double_buffer_write_ = buffer.value().get(); scope_.push_back(std::vector()); StmtExprVisitor::VisitStmt_(op); StmtEntry s; @@ -248,8 +266,8 @@ void StorageAccessVisitor::VisitExpr_(const CallNode* op) { } else if (op->op.same_as(builtin::tvm_access_ptr())) { TVM_FFI_ICHECK_EQ(op->args.size(), 5U); PrimType dtype = op->args[0].as_or_throw().ty(); - const VarNode* buffer = op->args[1].as(); - if (buffer == nullptr) { + auto buffer_var = GetBufferDataVar(op->args[1]); + if (!buffer_var.has_value()) { // args[1] is not a raw Var — e.g. a nested tvm_access_ptr or some // other PrimExpr. Recurse into sub-exprs so any inner buffer var // refs still get visited, but don't try to record an access entry @@ -257,10 +275,11 @@ void StorageAccessVisitor::VisitExpr_(const CallNode* op) { StmtExprVisitor::VisitExpr_(op); return; } + const VarNode* buffer = buffer_var.value().get(); PrimExpr offset = op->args[2].as_or_throw(); PrimExpr extent = op->args[3].as_or_throw(); const IntImmNode* flag = op->args[4].as(); - StorageScope scope = GetScope(ffi::GetRef(buffer)); + StorageScope scope = GetScope(buffer_var.value()); // The buffer scope. if (Enabled(buffer, scope)) { TVM_FFI_ICHECK(allow_append_); @@ -297,6 +316,9 @@ void StorageAccessVisitor::VisitExpr_(const CallNode* op) { } StorageScope StorageAccessVisitor::GetScope(Var buffer_var) const { + if (auto buffer_type = buffer_var->ty.as()) { + return StorageScope::Create(buffer_type.value()->storage_scope); + } if (buffer_var->ty.as()) { return StorageScope::Create(GetPtrStorageScope(buffer_var)); } diff --git a/src/s_tir/transform/transform_mma_buffer_layout.cc b/src/s_tir/transform/transform_mma_buffer_layout.cc index 6ce119f40572..35ee0e6f0321 100644 --- a/src/s_tir/transform/transform_mma_buffer_layout.cc +++ b/src/s_tir/transform/transform_mma_buffer_layout.cc @@ -70,7 +70,8 @@ class MmaBufferLayoutTransformer : public StmtExprMutator { new_shape.insert(new_shape.end(), {IntImm::Int32(dim0->value / 16), IntImm::Int32(dim1->value / 8), 2, 2}); - BufferVar new_buffer = decl_buffer(std::move(new_shape), buffer->dtype, buffer.name(), "local"); + BufferVar new_buffer = + decl_buffer(std::move(new_shape), buffer->dtype, buffer.name(), "local"); this->buffer_map_.insert({buffer, new_buffer}); this->buffer_var_map_.insert({buffer.var(), new_buffer.var()}); return new_buffer; @@ -91,7 +92,8 @@ class MmaBufferLayoutTransformer : public StmtExprMutator { new_shape.insert(new_shape.end(), {IntImm::Int32(dim0->value / 32), IntImm::Int32(dim1->value / 8), 4, 2}); - BufferVar new_buffer = decl_buffer(std::move(new_shape), buffer->dtype, buffer.name(), "local"); + BufferVar new_buffer = + decl_buffer(std::move(new_shape), buffer->dtype, buffer.name(), "local"); this->buffer_map_.insert({buffer, new_buffer}); this->buffer_var_map_.insert({buffer.var(), new_buffer.var()}); return new_buffer; @@ -112,7 +114,8 @@ class MmaBufferLayoutTransformer : public StmtExprMutator { new_shape.insert(new_shape.end(), {IntImm::Int32(dim0->value / 8), IntImm::Int32(dim1->value / 32), 1, 8}); - BufferVar new_buffer = decl_buffer(std::move(new_shape), buffer->dtype, buffer.name(), "local"); + BufferVar new_buffer = + decl_buffer(std::move(new_shape), buffer->dtype, buffer.name(), "local"); this->buffer_map_.insert({buffer, new_buffer}); this->buffer_var_map_.insert({buffer.var(), new_buffer.var()}); return new_buffer; 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 7e0a2f4edfc5..51b6456285ad 100644 --- a/src/s_tir/transform/using_assume_to_reduce_branches.cc +++ b/src/s_tir/transform/using_assume_to_reduce_branches.cc @@ -135,8 +135,7 @@ class ParseAssumeAndOvercompute : public IRMutatorWithAnalyzer { std::vector conditions_; // Storing all the buffer assumptions data in map - std::unordered_map + std::unordered_map map_buffer_assumption; tirx::BufferVar current_bufferstorenode_name; diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc index 232ef19ca03f..7b634bfbf82d 100644 --- a/src/target/llvm/codegen_cpu.cc +++ b/src/target/llvm/codegen_cpu.cc @@ -573,7 +573,8 @@ void CodeGenCPU::CreateComputeScope(const AttrStmtNode* op) { llvm::Argument* v = &(*it); const Var& var = vargs[idx]; var_map_[var.get()] = v; - if (var->ty.as() && !alias_var_set_.count(var.get())) { + if ((var->ty.as() || var->ty.as()) && + !alias_var_set_.count(var.get())) { // set non alias. fcompute->addParamAttr(idx, llvm::Attribute::NoAlias); // always not inline compute function to make the code structure clean @@ -591,16 +592,22 @@ void CodeGenCPU::CreateComputeScope(const AttrStmtNode* op) { } function_ = fcompute; + ffi::Array debug_param_types = vargs.Map([](const Var& var) -> Type { + if (const auto* buffer_type = var->ty.as()) { + // Compute-scope captures use their physical LLVM pointer values. + return buffer_type->DataPointerType(); + } + return var->ty; + }); di_subprogram_ = - CreateDebugFunction(MakeStringRef(value->value), - vargs.Map([](const Var& var) { return var->ty; }), PrimType::Int(32)); + CreateDebugFunction(MakeStringRef(value->value), debug_param_types, PrimType::Int(32)); auto* compute_entry = llvm::BasicBlock::Create(*ctx, "entry", function_); builder_->SetInsertPoint(compute_entry); this->VisitStmt(op->body); builder_->CreateRet(ConstInt32(0)); builder_->SetInsertPoint(compute_call_end); - AddDebugInformation(fcompute, vargs.Map([](const Var& var) { return var->ty; })); + AddDebugInformation(fcompute, debug_param_types); } CodeGenLLVM::TypedPointer CodeGenCPU::PackClosureData(const ffi::Array& vfields, diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc index 6da359bcbed3..cc4e9472e5b1 100644 --- a/src/target/llvm/codegen_llvm.cc +++ b/src/target/llvm/codegen_llvm.cc @@ -109,7 +109,7 @@ PrimType WithScalableVScaleFactor(const PrimType& dtype, int vscale_factor) { return PrimType::ScalableVector(dtype.code(), dtype.bits(), vscale_factor); } -// Underlying access type for a BufferVar: bool is backed by int8 so vectorized +// Underlying access type for a buffer: bool is backed by int8 so vectorized // accesses lower to real loads/stores instead of i1 predicate registers. PrimType BufferAccessType(const PrimType& dtype) { if (!dtype.MatchesCode(DLDataTypeCode::kDLBool)) return dtype; @@ -1007,7 +1007,7 @@ CodeGenLLVM::TypedPointer CodeGenLLVM::CreateBufferPtr(llvm::Value* buffer_ptr, "no padding for alignment."; } else { TVM_FFI_ICHECK(buffer_element_type.as()) - << "BufferVar elements must have primitive or pointer type, but got " << buffer_element_type; + << "Buffer elements must have primitive or pointer type, but got " << buffer_element_type; } llvm::Value* value_ptr = builder_->CreateInBoundsGEP(llvm_element_type, buffer_ptr, index); @@ -1732,7 +1732,7 @@ void CodeGenLLVM::BufferAccessHelper( PrimType buffer_element_dtype = BufferAccessType(buffer->dtype); TVM_FFI_ICHECK_GE(indices.size(), 1) - << "BufferVar " << buffer.name() << " is accessed with no indices. " + << "Buffer " << buffer.name() << " is accessed with no indices. " << "0-d scalar buffers are expected to be flattened to 1-d buffers prior to codegen."; // Only the last index is allowed to be multi-lane. All earlier @@ -1742,7 +1742,7 @@ void CodeGenLLVM::BufferAccessHelper( std::vector earlier_index_values; for (size_t i = 0; i < indices.size() - 1; i++) { TVM_FFI_ICHECK_EQ(PrimType(indices[i].ty()->dtype).lanes(), 1) - << "BufferVar " << buffer.name() << " is accessed with a multi-lane index at position " << i + << "Buffer " << buffer.name() << " is accessed with a multi-lane index at position " << i << ". Multi-lane indices are only supported as the last index."; earlier_index_values.push_back(MakeValue(indices[i])); } @@ -2220,7 +2220,28 @@ void CodeGenLLVM::VisitStmt_(const SeqStmtNode* op) { } } -void CodeGenLLVM::VisitStmt_(const DeclBufferNode* op) { EmitDebugLocation(op); } +void CodeGenLLVM::VisitStmt_(const DeclBufferNode* op) { + EmitDebugLocation(op); + const VarNode* buffer = op->buffer.get(); + TVM_FFI_ICHECK(!var_map_.count(buffer)); + if (!is_restricted_) { + alias_var_set_.insert(buffer); + } + + llvm::Value* value = MakeValue(op->data); + llvm::Type* expected_type = GetLLVMType(op->buffer.DataPointerType()); + if (value->getType() != expected_type) { + value->setName((op->buffer.name() + "_source_ptr").c_str()); + value = builder_->CreatePointerCast(value, expected_type); + } + + AddDebugInformation(value, op->buffer.var()); + var_map_[buffer] = value; + if (alloc_storage_info_.count(buffer) && alloc_storage_info_[buffer].alignment > 1) { + builder_->CreateAlignmentAssumption(*data_layout_, GetVarValue(buffer), + alloc_storage_info_[buffer].alignment); + } +} void CodeGenLLVM::VisitStmt_(const EvaluateNode* op) { EmitDebugLocation(op); @@ -2309,7 +2330,13 @@ void CodeGenLLVM::AddDebugInformation(llvm::Value* llvm_value, const Var& tir_va if (!di_subprogram_) return; - auto dbg_dtype = GetDebugType(tir_var->ty); + Type debug_type = tir_var->ty; + if (const auto* buffer_type = debug_type.as()) { + // A BufferVar is a compiler-side identity. Its LLVM value is the physical + // data pointer installed by AllocBuffer or DeclBuffer. + debug_type = buffer_type->DataPointerType(); + } + auto dbg_dtype = GetDebugType(debug_type); // no invalid dtypes if (!dbg_dtype) return; auto local_var = dbg_info_->di_builder_->createAutoVariable( diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc index a813d4fb5e34..4408ec0f835b 100644 --- a/src/target/source/codegen_c.cc +++ b/src/target/source/codegen_c.cc @@ -905,13 +905,22 @@ void CodeGenC::PrintVecBinaryOp(const std::string& op, const PrimType& t, PrimEx } void CodeGenC::VisitStmt_(const DeclBufferNode* op) { - if (!op->data.has_value()) { + const VarNode* source = op->data.as(); + if (const auto* call = op->data.as(); + call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) { + source = call->args[0].as(); + } + if (source && var_idmap_.count(source)) { + TVM_FFI_ICHECK(!var_idmap_.count(op->buffer.get())); + var_idmap_[op->buffer.get()] = GetVarID(source); + RegisterHandleType(op->buffer.get(), op->buffer->dtype); return; } + this->PrintIndent(); - PrintType(op->buffer->data_pointer_type, stream); + PrintType(op->buffer.DataPointerType(), stream); stream << ' ' << AllocVarID(op->buffer.get()) << " = "; - PrintExpr(op->data.value(), stream); + PrintExpr(Call(op->buffer.DataPointerType(), builtin::reinterpret(), {op->data}), stream); stream << ";\n"; RegisterHandleType(op->buffer.get(), op->buffer->dtype); } @@ -1204,7 +1213,7 @@ void CodeGenC::VisitStmt_(const BindNode* op) { void CodeGenC::VisitStmt_(const AllocBufferNode* op) { TVM_FFI_ICHECK(op->buffer.defined()); - std::string vid = AllocVarID(op->buffer.get()); + std::string vid = AllocVarID(op->buffer.get(), op->buffer.name() + "_ptr"); this->PrintIndent(); const auto& shape = op->buffer->shape; diff --git a/src/target/source/codegen_source_base.cc b/src/target/source/codegen_source_base.cc index 9efb3815a35a..f83240d6d87c 100644 --- a/src/target/source/codegen_source_base.cc +++ b/src/target/source/codegen_source_base.cc @@ -54,10 +54,11 @@ std::string CodeGenSourceBase::SSAGetID(std::string src, const Type& t) { return e.vid; } -std::string CodeGenSourceBase::AllocVarID(const tirx::VarNode* v) { +std::string CodeGenSourceBase::AllocVarID(const tirx::VarNode* v) { return AllocVarID(v, v->name); } + +std::string CodeGenSourceBase::AllocVarID(const tirx::VarNode* v, std::string name_hint) { TVM_FFI_ICHECK(!var_idmap_.count(v)) << "Need input to be in SSA form dup " << v->name; - std::string key = v->name; - std::string vid = name_supply_->FreshName(key); + std::string vid = name_supply_->FreshName(name_hint); std::replace(vid.begin(), vid.end(), ':', '_'); std::replace(vid.begin(), vid.end(), '-', '_'); std::replace(vid.begin(), vid.end(), '.', '_'); diff --git a/src/target/source/codegen_source_base.h b/src/target/source/codegen_source_base.h index 95b6c00e7002..531dae37ef67 100644 --- a/src/target/source/codegen_source_base.h +++ b/src/target/source/codegen_source_base.h @@ -84,6 +84,13 @@ class CodeGenSourceBase { * \return the variable name. */ std::string AllocVarID(const tirx::VarNode* v); + /*! + * \brief Allocate a variable name using an explicit diagnostic hint. + * \param v The variable. + * \param name_hint The name hint. + * \return the variable name. + */ + std::string AllocVarID(const tirx::VarNode* v, std::string name_hint); /*! * \brief Get a variable name. * \param v The variable. diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc index 41b28c116d4f..410dc4c56844 100644 --- a/src/te/operation/create_primfunc.cc +++ b/src/te/operation/create_primfunc.cc @@ -48,7 +48,8 @@ namespace tirx { /*! \brief The helper mutator that transforms ProducerLoad to BufferLoad */ class ProducerToBufferTransformer : public StmtExprMutator { public: - explicit ProducerToBufferTransformer(const std::unordered_map& tensor2buffers) + explicit ProducerToBufferTransformer( + const std::unordered_map& tensor2buffers) : tensor2buffers_(tensor2buffers) {} Expr VisitExpr_(const ProducerLoadNode* op) final { diff --git a/src/tirx/analysis/var_use_def_analysis.cc b/src/tirx/analysis/var_use_def_analysis.cc index 6d5c860f9fd3..82857ae73bf3 100644 --- a/src/tirx/analysis/var_use_def_analysis.cc +++ b/src/tirx/analysis/var_use_def_analysis.cc @@ -91,7 +91,12 @@ void VarUseDefAnalyzer::VisitExpr_(const LetNode* op) { } void VarUseDefAnalyzer::VisitExpr_(const VarNode* op) { - this->HandleUse(ffi::GetRef(op)); + Var var = ffi::GetRef(op); + if (var->ty.as()) { + this->VisitBufferUse(BufferVar(var)); + } else { + this->HandleUse(var); + } StmtExprVisitor::VisitExpr_(op); } @@ -109,9 +114,6 @@ void VarUseDefAnalyzer::VisitBufferDef(const BufferVar& buffer, bool alloc_data) auto it = use_count_.find(buffer.get()); if (it == use_count_.end()) { HandleDef(buffer.var()); - } else { - TVM_FFI_ICHECK_GE(it->second, 0) - << "variable " << buffer.name() << " has been used before definition!"; } } // Visit shape/strides/elem_offset as uses of vars from the enclosing scope. @@ -125,17 +127,6 @@ void VarUseDefAnalyzer::VisitBufferUse(const BufferVar& buffer) { HandleUse(buffer.var()); } -void VarUseDefAnalyzer::VisitBuffer(const BufferVar& buffer) { - auto visit_arr = [&](ffi::Array arr) { - for (const auto& element : arr) { - this->VisitExpr(element); - } - }; - - visit_arr(buffer->shape); - visit_arr(buffer->strides); -} - void VarUseDefAnalyzer::HandleDef(const Var& var) { auto v = var.get(); TVM_FFI_ICHECK(!def_count_.count(v)) @@ -165,12 +156,11 @@ void VarUseDefAnalyzer::HandleDef(const BufferVar& buf) { // reference the same BufferVar object. Treat repeated definition of the same // buffer object as idempotent. if (buffer_def_count_.count(ptr)) { - VisitBuffer(buf); return; } - TVM_FFI_ICHECK(!buffer_use_count_.count(ptr)) - << "buffer " << ptr->name << " has been used before definition!"; - buffer_use_count_[ptr] = 0; + if (!buffer_use_count_.count(ptr)) { + buffer_use_count_[ptr] = 0; + } buffer_def_count_[ptr] = 1; // BufferVar fields (data, shape, strides) are visited by the caller // (VisitBufferDef) via the base class, not here. diff --git a/src/tirx/analysis/var_use_def_analysis.h b/src/tirx/analysis/var_use_def_analysis.h index 293f5574a0d5..bf4cbffc0dee 100644 --- a/src/tirx/analysis/var_use_def_analysis.h +++ b/src/tirx/analysis/var_use_def_analysis.h @@ -80,8 +80,6 @@ class VarUseDefAnalyzer : public StmtExprVisitor { void HandleDef(const BufferVar& buf); void HandleUse(const BufferVar& buf); - - void VisitBuffer(const BufferVar& buffer); }; } // namespace tirx diff --git a/src/tirx/ir/async_structs.cc b/src/tirx/ir/async_structs.cc index 0259af5e211f..0213468fb100 100644 --- a/src/tirx/ir/async_structs.cc +++ b/src/tirx/ir/async_structs.cc @@ -51,9 +51,9 @@ Pipeline::Pipeline(ExecScope thread_scope, size_t depth, bool separate_pc, ffi:: TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def( - "tirx.Pipeline", - [](ExecScope thread_scope, size_t depth, bool separate_pc, ffi::String name_hint, - ffi::Map workspace, ffi::Map schedule_config) { + "tirx.Pipeline", [](ExecScope thread_scope, size_t depth, bool separate_pc, + ffi::String name_hint, ffi::Map workspace, + ffi::Map schedule_config) { return Pipeline(thread_scope, depth, separate_pc, name_hint, workspace, schedule_config); }); } diff --git a/src/tirx/ir/buffer.cc b/src/tirx/ir/buffer.cc index c0fa6e065ca4..c0a29ff8baa8 100644 --- a/src/tirx/ir/buffer.cc +++ b/src/tirx/ir/buffer.cc @@ -44,15 +44,14 @@ TVM_FFI_STATIC_INIT_BLOCK() { BufferTypeNode::RegisterReflection(); } using IndexMod = tirx::FloorModNode; using IndexDiv = tirx::FloorDivNode; -BufferType::BufferType(PointerType data_pointer_type, PrimType dtype, - ffi::Array shape, ffi::Array strides, - PrimExpr elem_offset, int data_alignment, int offset_factor, - ffi::Optional layout, ffi::Array allocated_addr, - Span span) +BufferType::BufferType(ffi::String storage_scope, PrimType dtype, ffi::Array shape, + ffi::Array strides, PrimExpr elem_offset, int data_alignment, + int offset_factor, ffi::Optional layout, + ffi::Array allocated_addr, Span span) : Type(ffi::UnsafeInit{}) { auto n = ffi::make_object(); - n->data_pointer_type = std::move(data_pointer_type); n->dtype = std::move(dtype); + n->storage_scope = storage_scope.empty() ? ffi::String("global") : std::move(storage_scope); n->shape = std::move(shape); n->strides = std::move(strides); if (!elem_offset.defined()) { @@ -70,25 +69,8 @@ BufferType::BufferType(PointerType data_pointer_type, PrimType dtype, namespace { -BufferType MakeBufferType(Var data, PrimType dtype, ffi::Array shape, - ffi::Array strides, PrimExpr elem_offset, - int data_alignment, int offset_factor, - ffi::Optional layout, - ffi::Array allocated_addr) { - TVM_FFI_ICHECK(!data->ty.IsMissing()) - << "Variable " << data->name << " is missing a type annotation."; - const auto* pointer_type = data->ty.as(); - TVM_FFI_ICHECK(pointer_type) << "Variable " << data->name << " is not a pointer."; - TVM_FFI_ICHECK(pointer_type->element_type.as()) - << "Variable " << data->name << " does not point to a primitive."; - return BufferType(ffi::GetRef(pointer_type), std::move(dtype), - std::move(shape), std::move(strides), std::move(elem_offset), - data_alignment, offset_factor, std::move(layout), - std::move(allocated_addr)); -} - -BufferVar RebuildBufferVar(const BufferVar& buffer, BufferType type, - ffi::String name_suffix = "") { +BufferVar RebuildBufferVarFromType(const BufferVar& buffer, BufferType type, + ffi::String name_suffix = "") { return BufferVar(buffer.name() + name_suffix, std::move(type), buffer.span()); } @@ -103,10 +85,7 @@ ffi::Array SimplifyArray(arith::AnalyzerObj* ana, ffi::Array BufferVar decl_buffer(ffi::Array shape, PrimType dtype, ffi::String name, ffi::String storage_scope, Span span) { - PrimType storage_type = (dtype == PrimType::Bool() ? PrimType::Int(8) : dtype); - return BufferVar(Var(name, PointerType(storage_type, storage_scope), span), dtype, shape, - ffi::Array(), PrimExpr(), name, 0, 0, span, std::nullopt, - ffi::Array()); + return BufferVar(name, BufferType(storage_scope, dtype, shape, {}, PrimExpr(), 0, 0), span); } // Split the given expression w.r.t the add operator @@ -308,11 +287,11 @@ ffi::Array BufferVar::OffsetOf(ffi::Array input_indices) con // The buffer offset in convention of number of elements of // original data ignoring number of lanes. // We also perform optimization to simplify the indexing expression. -ffi::Array BufferTypeNode::ElemOffset(ffi::Array input_indices, bool inner) const { +ffi::Array BufferTypeNode::ElemOffset(ffi::Array input_indices, + bool inner) const { TVM_FFI_ICHECK_EQ(shape.size(), input_indices.size()) << "BufferType is " << shape.size() << "-dimensional, cannot be indexed with the " - << input_indices.size() - << "-dimensional indices provided."; + << input_indices.size() << "-dimensional indices provided."; if (strides.size()) { TVM_FFI_ICHECK_EQ(this->strides.size(), input_indices.size()) @@ -389,16 +368,15 @@ BufferVar BufferVar::GetFlattenedBuffer() const { // structural compares against a freshly-decl'd 1-D buffer would diff // (see test_tir_transform_flatten_buffer). Reset to the default layout // for the new shape so the buffer stays internally consistent. - return RebuildBufferVar( - *this, - BufferType(self->data_pointer_type, self->dtype, output_shape, {}, - self->elem_offset, self->data_alignment, self->offset_factor, - TileLayoutNode::DefaultLayout(output_shape), self->allocated_addr)); + return RebuildBufferVarFromType( + *this, BufferType(self->storage_scope, self->dtype, output_shape, {}, self->elem_offset, + self->data_alignment, self->offset_factor, + TileLayoutNode::DefaultLayout(output_shape), self->allocated_addr)); } } PrimExpr BufferVar::vload(ffi::Array begin, PrimType value_dtype, - ffi::Optional predicate) const { + ffi::Optional predicate) const { const BufferTypeNode* n = operator->(); TVM_FFI_ICHECK(n != nullptr); PrimType buffer_dtype(n->dtype); @@ -423,7 +401,7 @@ PrimExpr BufferVar::vload(ffi::Array begin, PrimType value_dtype, } Stmt BufferVar::vstore(ffi::Array begin, PrimExpr value, - ffi::Optional predicate) const { + ffi::Optional predicate) const { const BufferTypeNode* n = operator->(); TVM_FFI_ICHECK(n != nullptr); PrimType value_dtype = value.ty(); @@ -448,12 +426,7 @@ Stmt BufferVar::vstore(ffi::Array begin, PrimExpr value, return BufferStore(*this, value, indices, predicate); } -ffi::String BufferVar::scope() const { - if ((*this)->data_pointer_type->storage_scope.empty()) { - return "global"; - } - return (*this)->data_pointer_type->storage_scope; -} +ffi::String BufferVar::scope() const { return (*this)->storage_scope; } BufferVar BufferVar::MakeStrideView() const { if ((*this)->strides.size() != 0) return *this; @@ -470,10 +443,10 @@ BufferVar BufferVar::MakeStrideView() const { for (size_t i = temp.size(); i != 0; --i) { strides.push_back(temp[i - 1]); } - return RebuildBufferVar( - *this, BufferType(self->data_pointer_type, self->dtype, self->shape, - std::move(strides), self->elem_offset, self->data_alignment, - self->offset_factor, self->layout, self->allocated_addr)); + return RebuildBufferVarFromType( + *this, BufferType(self->storage_scope, self->dtype, self->shape, std::move(strides), + self->elem_offset, self->data_alignment, self->offset_factor, self->layout, + self->allocated_addr)); } BufferVar BufferVar::MakeSlice(ffi::Array begins, ffi::Array extents) const { @@ -502,22 +475,21 @@ BufferVar BufferVar::MakeSlice(ffi::Array begins, ffi::Array return MakeStrideView().MakeSlice(begins, extents); } } - return RebuildBufferVar( + return RebuildBufferVarFromType( *this, - BufferType(n->data_pointer_type, n->dtype, extents, strides, elem_offset[0], - n->data_alignment, 0, TileLayoutNode::DefaultLayout(extents)), + BufferType(n->storage_scope, n->dtype, extents, strides, elem_offset[0], n->data_alignment, 0, + TileLayoutNode::DefaultLayout(extents)), "_slice"); } -Expr BufferVar::access_ptr(int access_mask, PointerType ptr_type, int content_lanes, PrimExpr offset, - ffi::Optional input_extent) const { +Expr BufferVar::access_ptr(int access_mask, PointerType ptr_type, int content_lanes, + PrimExpr offset, ffi::Optional input_extent) const { const BufferTypeNode* self = operator->(); TVM_FFI_ICHECK(self != nullptr); // An access pointer addresses the same allocation as the buffer data. The // requested type controls its pointee, while the buffer controls its address // space (for example, shared or local memory). - ptr_type = - PointerType(ptr_type->element_type, self->data_pointer_type->storage_scope); + ptr_type = PointerType(ptr_type->element_type, self->storage_scope); PrimExpr e_dtype; PrimExpr extent; if (self->shape.size() == 0) { @@ -542,36 +514,18 @@ Expr BufferVar::access_ptr(int access_mask, PointerType ptr_type, int content_la if (input_extent.has_value()) { extent = input_extent.value(); } - ffi::Array acc_args{e_dtype, data(), elem_offset, extent, - IntImm::Int32(access_mask)}; + ffi::Array acc_args{e_dtype, data(), elem_offset, extent, IntImm::Int32(access_mask)}; return Call(ptr_type, tirx::builtin::tvm_access_ptr(), acc_args); } BufferVar::BufferVar(ffi::String name, BufferType type, Span span) : Var(Var(std::move(name), std::move(type), std::move(span))) {} -BufferVar::BufferVar(Var data, PrimType dtype, ffi::Array shape, - ffi::Array strides, PrimExpr elem_offset, - ffi::String name, int data_alignment, int offset_factor, - Span span, ffi::Optional layout, - ffi::Array allocated_addr) - : BufferVar( - std::move(name), - MakeBufferType(std::move(data), std::move(dtype), std::move(shape), - std::move(strides), std::move(elem_offset), data_alignment, - offset_factor, std::move(layout), std::move(allocated_addr)), - std::move(span)) {} - -Expr BufferVar::data() const { - return Call((*this)->data_pointer_type, builtin::buffer_data(), {var()}); -} - -tirx::BufferVar BufferWithOffsetAlignment(ffi::Array shape, PrimType dtype, std::string name, - int data_alignment, int offset_factor, - std::string memory_scope) { - PrimType storage_type = (dtype == PrimType::Bool() ? PrimType::Int(8) : dtype); - auto data = tirx::Var(name, PointerType(storage_type, memory_scope)); +Expr BufferVar::data() const { return Call(DataPointerType(), builtin::buffer_data(), {var()}); } +tirx::BufferVar BufferWithOffsetAlignment(ffi::Array shape, PrimType dtype, + std::string name, int data_alignment, int offset_factor, + std::string memory_scope) { PrimExpr elem_offset; if (offset_factor != 0) { elem_offset = tirx::PrimVar(name + "_elem_offset", shape[0].ty()); @@ -579,25 +533,24 @@ tirx::BufferVar BufferWithOffsetAlignment(ffi::Array shape, PrimType d elem_offset = PrimExpr(); } - return tirx::BufferVar(data, dtype, shape, ffi::Array(), elem_offset, name, data_alignment, - offset_factor); + return tirx::BufferVar( + name, BufferType(memory_scope, dtype, shape, {}, elem_offset, data_alignment, offset_factor)); } BufferVar BufferVar::with_allocated_addr(ffi::Array allocated_addr) const { const auto* self = operator->(); - return RebuildBufferVar( - *this, BufferType(self->data_pointer_type, self->dtype, self->shape, - self->strides, self->elem_offset, self->data_alignment, - self->offset_factor, self->layout, + return RebuildBufferVarFromType( + *this, BufferType(self->storage_scope, self->dtype, self->shape, self->strides, + self->elem_offset, self->data_alignment, self->offset_factor, self->layout, std::move(allocated_addr))); } BufferVar BufferVar::with_dtype(PrimType dtype) const { const auto* self = operator->(); - return RebuildBufferVar( - *this, BufferType(self->data_pointer_type, std::move(dtype), self->shape, - self->strides, self->elem_offset, self->data_alignment, - self->offset_factor, self->layout, self->allocated_addr)); + return RebuildBufferVarFromType( + *this, BufferType(self->storage_scope, std::move(dtype), self->shape, self->strides, + self->elem_offset, self->data_alignment, self->offset_factor, self->layout, + self->allocated_addr)); } PrimExpr BufferVar::OffsetOf_p(const Array& indices) const { @@ -617,49 +570,37 @@ bool BufferVar::IsScalar(bool alloc_or_decl) const { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() - .def_packed("tirx.BufferVar", - [](ffi::PackedArgs args, ffi::Any* ret) { - TVM_FFI_ICHECK_EQ(args.size(), 10); - auto data = args[0].cast(); - auto dtype = args[1].cast(); - auto shape = args[2].cast>(); - auto strides = args[3].cast>(); - auto elem_offset = args[4].cast(); - auto name = args[5].cast(); - auto data_alignment = args[6].cast(); - auto offset_factor = args[7].cast(); - auto span = args[8].cast(); - auto layout = args[9].cast(); - *ret = BufferVar(data, dtype, shape, strides, elem_offset, name, data_alignment, - offset_factor, span, layout); - }) - .def_method("tirx.BufferAccessPtr", - static_cast) const>(&BufferVar::access_ptr)) + .def("tirx.BufferVar", + [](ffi::String name, BufferType type, Span span) { + return BufferVar(std::move(name), std::move(type), std::move(span)); + }) + .def_method( + "tirx.BufferAccessPtr", + static_cast) + const>(&BufferVar::access_ptr)) .def_method("tirx.BufferGetFlattenedBuffer", &BufferVar::GetFlattenedBuffer) .def_method("tirx.BufferOffsetOf", &BufferVar::OffsetOf) .def_method("tirx.BufferOffsetOfp", &BufferVar::OffsetOf_p) - .def_method("tirx.BufferVLoad", - static_cast, PrimType, - ffi::Optional) const>(&BufferVar::vload)) + .def_method( + "tirx.BufferVLoad", + static_cast, PrimType, + ffi::Optional) const>(&BufferVar::vload)) .def_method("tirx.BufferVStore", &BufferVar::vstore) .def_method("tirx.BufferStorageScope", &BufferVar::scope) .def_method("tirx.BufferWithAllocatedAddr", &BufferVar::with_allocated_addr) .def_method("tirx.BufferWithDtype", &BufferVar::with_dtype) .def_method("tirx.BufferIsScalar", &BufferVar::IsScalar) .def_method("tirx.BufferData", &BufferVar::data) - .def("tirx.BufferType", - [](PointerType data_pointer_type, PrimType dtype, - ffi::Array shape, ffi::Array strides, - PrimExpr elem_offset, int data_alignment, int offset_factor, - ffi::Optional layout, - ffi::Array allocated_addr, Span span) { - return BufferType(std::move(data_pointer_type), std::move(dtype), - std::move(shape), std::move(strides), - std::move(elem_offset), data_alignment, offset_factor, - std::move(layout), std::move(allocated_addr), - std::move(span)); - }); + .def_method("tirx.BufferDataPointerType", &BufferVar::DataPointerType) + .def("tirx.BufferType", [](ffi::String storage_scope, PrimType dtype, + ffi::Array shape, ffi::Array strides, + PrimExpr elem_offset, int data_alignment, int offset_factor, + ffi::Optional layout, ffi::Array allocated_addr, + Span span) { + return BufferType(std::move(storage_scope), std::move(dtype), std::move(shape), + std::move(strides), std::move(elem_offset), data_alignment, offset_factor, + std::move(layout), std::move(allocated_addr), std::move(span)); + }); } } // namespace tirx diff --git a/src/tirx/ir/data_type_rewriter.cc b/src/tirx/ir/data_type_rewriter.cc index c8bd53f100ef..04abe0dc5ee1 100644 --- a/src/tirx/ir/data_type_rewriter.cc +++ b/src/tirx/ir/data_type_rewriter.cc @@ -336,11 +336,13 @@ Stmt IndexDataTypeRewriter::VisitStmt_(const SBlockRealizeNode* op) { } Stmt IndexDataTypeRewriter::VisitStmt_(const SBlockNode* op) { - ffi::Array new_alloc_buffers = op->alloc_buffers.Map( - [this](const BufferVar& buffer) { return this->VisitBufferDef(buffer, /*alloc_data=*/true); }); + ffi::Array new_alloc_buffers = op->alloc_buffers.Map([this](const BufferVar& buffer) { + return this->VisitBufferDef(buffer, /*alloc_data=*/true); + }); ffi::Array new_match_buffers = op->match_buffers.Map([this](const MatchBufferRegion& match_buffer_region) { - BufferVar new_buffer = this->VisitBufferDef(match_buffer_region->buffer, /*alloc_data=*/true); + BufferVar new_buffer = + this->VisitBufferDef(match_buffer_region->buffer, /*alloc_data=*/true); BufferRegion new_buffer_region = this->VisitBufferRegion(match_buffer_region->source); if (!new_buffer.same_as(match_buffer_region->buffer) || !new_buffer_region.same_as(match_buffer_region->source)) { diff --git a/src/tirx/ir/expr_functor.cc b/src/tirx/ir/expr_functor.cc index 721077004580..2fbd99f62eab 100644 --- a/src/tirx/ir/expr_functor.cc +++ b/src/tirx/ir/expr_functor.cc @@ -20,6 +20,7 @@ * \file expr_functor.cc */ #include +#include #include #include "functor_common.h" @@ -147,7 +148,16 @@ Expr ExprMutator::VisitExpr_(const CallNode* op) { if (args.same_as(op->args)) { return ffi::GetRef(op); } else { - return Call(op->ExprNode::ty, op->op, args, op->attrs, op->ty_args, op->span); + Type result_type = op->ExprNode::ty; + if (op->op.same_as(builtin::buffer_data())) { + TVM_FFI_ICHECK_EQ(args.size(), 1); + const auto* buffer_var = args[0].as(); + TVM_FFI_ICHECK(buffer_var); + const auto* buffer_type = buffer_var->ty.as(); + TVM_FFI_ICHECK(buffer_type); + result_type = buffer_type->DataPointerType(); + } + return Call(result_type, op->op, args, op->attrs, op->ty_args, op->span); } } diff --git a/src/tirx/ir/specialize.cc b/src/tirx/ir/specialize.cc index fcea0f7bf349..804eec2e64c5 100644 --- a/src/tirx/ir/specialize.cc +++ b/src/tirx/ir/specialize.cc @@ -364,6 +364,10 @@ void UpdateSpecializeVarMap(const PrimFunc& func, const Var& param, const Buffer build_var_mapping(specific_buf->strides[i], buf_to_specialize->strides[i]); } build_var_mapping(specific_buf->elem_offset, buf_to_specialize->elem_offset); + // The buffer identity owns the pointer projection instead of storing a + // separate data Var. Remap the typed Var itself so buffer_data uses retain + // the alias relationship established by specialization. + build_var_mapping(specific_buf.var(), buf_to_specialize.var()); // Check data_alignment and offset_factor. // These two signatures are int, so we do not need map them. diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index c7b494c4b2c0..14ca7f7cdace 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -277,7 +277,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } // DeclBuffer -DeclBuffer::DeclBuffer(BufferVar buffer, ffi::Optional data, Span span) { +DeclBuffer::DeclBuffer(BufferVar buffer, Expr data, Span span) { // Enforce storage scope rules for DeclBuffer. std::string scope = static_cast(buffer.scope()); if (scope.empty()) { @@ -300,10 +300,9 @@ DeclBuffer::DeclBuffer(BufferVar buffer, ffi::Optional data, Span span) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def( - "tirx.DeclBuffer", [](BufferVar buffer, ffi::Optional data, Span span) { - return DeclBuffer(buffer, data, span); - }); + refl::GlobalDef().def("tirx.DeclBuffer", [](BufferVar buffer, Expr data, Span span) { + return DeclBuffer(buffer, data, span); + }); } // AllocBuffer diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc index 904d11f0b8a7..171f1b7eb305 100644 --- a/src/tirx/ir/stmt_functor.cc +++ b/src/tirx/ir/stmt_functor.cc @@ -100,9 +100,7 @@ void StmtVisitor::VisitStmt_(const AllocBufferNode* op) { } void StmtVisitor::VisitStmt_(const DeclBufferNode* op) { - if (op->data.has_value()) { - this->VisitExpr(op->data.value()); - } + this->VisitExpr(op->data); this->VisitBufferDef(op->buffer, /*alloc_data=*/false); } @@ -417,15 +415,13 @@ BufferVar StmtMutator::VisitBufferDef(const BufferVar& buffer, bool alloc_data) } if (shape.same_as(buffer->shape) && strides.same_as(buffer->strides) && - elem_offset.same_as(buffer->elem_offset) && - allocated_addr.same_as(buffer->allocated_addr) && !layout_changed) { + elem_offset.same_as(buffer->elem_offset) && allocated_addr.same_as(buffer->allocated_addr) && + !layout_changed) { return buffer; } - BufferType new_type( - buffer->data_pointer_type, buffer->dtype, std::move(shape), - std::move(strides), std::move(elem_offset), buffer->data_alignment, - buffer->offset_factor, std::move(new_layout), std::move(allocated_addr), - buffer->span); + BufferType new_type(buffer->storage_scope, buffer->dtype, std::move(shape), std::move(strides), + std::move(elem_offset), buffer->data_alignment, buffer->offset_factor, + std::move(new_layout), std::move(allocated_addr), buffer->span); BufferVar new_buf(buffer.name(), std::move(new_type), buffer.span()); buffer_remap_.Set(buffer, new_buf); return new_buf; @@ -472,10 +468,7 @@ Stmt StmtMutator::VisitStmt_(const AllocBufferNode* op) { } Stmt StmtMutator::VisitStmt_(const DeclBufferNode* op) { - ffi::Optional data = std::nullopt; - if (op->data.has_value()) { - data = this->VisitExpr(op->data.value()); - } + Expr data = this->VisitExpr(op->data); BufferVar new_buf = this->VisitBufferDef(op->buffer, /*alloc_data=*/false); if (new_buf.same_as(op->buffer) && data.same_as(op->data)) { @@ -836,8 +829,8 @@ class IRSubstitute : public StmtExprMutator { if (auto mapped = vmap_(new_buf.var())) { auto mapped_var = mapped.value().as(); TVM_FFI_ICHECK(mapped_var && mapped_var.value()->ty.as()) - << "BufferVar " << new_buf << " was substituted into " - << mapped.value() << ", which is not a Var with BufferType"; + << "BufferVar " << new_buf << " was substituted into " << mapped.value() + << ", which is not a Var with BufferType"; new_buf = BufferVar(mapped_var.value()); buffer_remap_.Set(buffer, new_buf); } @@ -849,8 +842,8 @@ class IRSubstitute : public StmtExprMutator { if (auto mapped = vmap_(new_buf.var())) { auto mapped_var = mapped.value().as(); TVM_FFI_ICHECK(mapped_var && mapped_var.value()->ty.as()) - << "BufferVar " << new_buf << " was substituted into " - << mapped.value() << ", which is not a Var with BufferType"; + << "BufferVar " << new_buf << " was substituted into " << mapped.value() + << ", which is not a Var with BufferType"; new_buf = BufferVar(mapped_var.value()); } return new_buf; diff --git a/src/tirx/ir/tir_visitor_with_path.h b/src/tirx/ir/tir_visitor_with_path.h index 860759dc4b5b..85c42570f600 100644 --- a/src/tirx/ir/tir_visitor_with_path.h +++ b/src/tirx/ir/tir_visitor_with_path.h @@ -242,7 +242,8 @@ class TIRVisitorWithPath : protected ExprFunctor> WithMatchBufferDefs(BufferVar buf, ffi::reflection::AccessPath path) { + std::vector> WithMatchBufferDefs(BufferVar buf, + ffi::reflection::AccessPath path) { std::vector> context; auto try_visit_implicit_var_def = [this, &context](const Expr& expr, diff --git a/src/tirx/op/op.cc b/src/tirx/op/op.cc index 416a0085b8e3..68230f419a53 100644 --- a/src/tirx/op/op.cc +++ b/src/tirx/op/op.cc @@ -1429,7 +1429,7 @@ int ExtractInt(const ffi::PackedArgs& args, int index) { } } -PrimExpr PrintOpPacked(Var data, DLDataType dtype, bool is_string, bool is_scalar, int dim_num, +PrimExpr PrintOpPacked(Expr data, DLDataType dtype, bool is_string, bool is_scalar, int dim_num, ffi::Array shape) { PrimType value_ty(dtype); PrimType u32_ty = PrimType::UInt(32); @@ -1449,7 +1449,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def_packed("tirx.print_buffer", [](ffi::PackedArgs args, ffi::Any* ret) { // Expected arguments: - // args[0]: buffer_var (Var) + // args[0]: buffer data expression // args[1]: dtype (DLDataType) // args[2]: is_string (bool or IntImm) // args[3]: is_scalar (bool or IntImm) @@ -1458,7 +1458,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { TVM_FFI_ICHECK_GE(args.size(), 5) << "print_buffer expects at least 5 arguments"; - Var buffer_var = args[0].cast(); + Expr buffer_data = args[0].cast(); DLDataType dtype = args[1].cast(); bool is_string = ExtractBool(args, 2); bool is_scalar = ExtractBool(args, 3); @@ -1469,7 +1469,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { shape.push_back(args[i].cast()); } - *ret = PrintOpPacked(buffer_var, dtype, is_string, is_scalar, dim_num, shape); + *ret = PrintOpPacked(buffer_data, dtype, is_string, is_scalar, dim_num, shape); }); } diff --git a/src/tirx/script/builder/frame.cc b/src/tirx/script/builder/frame.cc index 63647ee1ada4..a253ee8cc3b6 100644 --- a/src/tirx/script/builder/frame.cc +++ b/src/tirx/script/builder/frame.cc @@ -37,7 +37,7 @@ namespace tirx { namespace { // In s_tir functions, buffer-typed parameters must not carry a layout (the -// s_tir IR doesn't track per-buffer layouts on params). When `T.BufferVar(...)` is +// s_tir IR doesn't track per-buffer layouts on params). When `T.Buffer(...)` is // used as a parameter annotation, the parser evaluates the annotation outside // the PrimFunc frame; if the annotation captures an outer-scope variable (e.g. // `dtype` in a closure-based generator), the evaluation happens *before* @@ -124,8 +124,7 @@ void PrimFuncFrameNode::ExitWithScope() { if (buf->layout.has_value()) { ffi::ObjectPtr type = tvm::tirx::CopyBufferType(buf); type->layout = std::nullopt; - tvm::tirx::BufferVar new_buf = - tvm::tirx::RebuildBufferVar(buf, std::move(type)); + tvm::tirx::BufferVar new_buf = tvm::tirx::RebuildBufferVar(buf, std::move(type)); normalizer.Register(buf, new_buf); new_buffer_map.Set(kv.first, new_buf); } else { diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc index fa973fced3ba..60556ed647b8 100644 --- a/src/tirx/script/builder/ir.cc +++ b/src/tirx/script/builder/ir.cc @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -42,8 +43,7 @@ using tvm::tirx::IterVar; using tvm::tirx::Layout; BufferVar BufferDecl(ffi::Array shape, PrimType dtype, ffi::String buffer_name, - ffi::Optional data, - ffi::Optional> strides, + ffi::Optional data, ffi::Optional> strides, ffi::Optional elem_offset, ffi::String storage_scope, int align, int offset_factor, ffi::Optional layout, ffi::Array allocated_addr) { @@ -52,26 +52,17 @@ BufferVar BufferDecl(ffi::Array shape, PrimType dtype, ffi::String buf << "ValueError: `allocated_addr` can only be used with `data`, `elem_offset`, and " "`offset_factor` undefined"; } - PointerType data_pointer_type = PointerType::VoidPointerTy(); - if (!data.has_value()) { - DLDataType storage_dtype = dtype->dtype; - if (storage_dtype == DLDataType{kDLBool, 8, 1}) { - storage_dtype = DLDataType{kDLInt, 8, 1}; - } - data_pointer_type = PointerType(PrimType(storage_dtype), storage_scope); - } else { - data_pointer_type = data.value()->ty.as_or_throw(); + if (data.has_value()) { + storage_scope = data.value()->ty.as_or_throw()->storage_scope; } if (!elem_offset.has_value() && offset_factor) { PrimType shape_dtype = shape.empty() ? PrimType::Int(32) : shape[0].ty(); elem_offset = tvm::tirx::PrimVar("elem_offset", shape_dtype); } - return BufferVar( - buffer_name, - tvm::tirx::BufferType(data_pointer_type, dtype, shape, - strides.value_or(ffi::Array()), - elem_offset.value_or(PrimExpr()), align, offset_factor, layout, - allocated_addr)); + return BufferVar(buffer_name, tvm::tirx::BufferType(storage_scope, dtype, shape, + strides.value_or(ffi::Array()), + elem_offset.value_or(PrimExpr()), align, + offset_factor, layout, allocated_addr)); } PrimFuncFrame PrimFunc(bool is_private, bool s_tir, bool persistent) { @@ -99,7 +90,7 @@ Var Arg(ffi::String name, Var var) { BufferVar Arg(ffi::String name, BufferVar buffer) { PrimFuncFrame frame = FindPrimFuncFrame("T.Arg"); details::Namer::Name(buffer, name); - // A BufferVar parameter is an opaque ABI handle. The BufferVar's data Var + // A buffer parameter is an opaque ABI handle. The buffer's data pointer // carries the exact pointee type used within the function body. Var handle(buffer.name() + "_handle", PointerType::VoidPointerTy()); frame->args.push_back(handle); @@ -150,11 +141,11 @@ tvm::Type FuncRet(tvm::Type ret_type) { } BufferVar MatchBuffer(ffi::ObjectRef param, ffi::Array shape, PrimType dtype, - ffi::Optional data, ffi::Array strides, - PrimExpr elem_offset, ffi::String storage_scope, int align, - int offset_factor, ffi::Optional layout) { + ffi::Optional data, ffi::Array strides, PrimExpr elem_offset, + ffi::String storage_scope, int align, int offset_factor, + ffi::Optional layout) { BufferVar buffer = BufferDecl(shape, dtype, "", data, strides, elem_offset, storage_scope, align, - offset_factor, layout, {}); + offset_factor, layout, {}); if (auto var = param.as()) { PrimFuncFrame frame = FindPrimFuncFrame("T.match_buffer"); Var v = var.value(); @@ -386,7 +377,7 @@ ffi::Variant SBlockAllocBuffer( ffi::Optional opt_elem_offset = elem_offset.defined() ? ffi::Optional(elem_offset) : std::nullopt; BufferVar buffer = BufferDecl(shape, dtype, "", std::nullopt, strides, opt_elem_offset, - storage_scope, align, offset_factor, layout, allocated_addr); + storage_scope, align, offset_factor, layout, allocated_addr); IRBuilder builder = IRBuilder::Current(); auto opt_func_frame = builder->FindFrame(); if (opt_func_frame.has_value()) { @@ -812,8 +803,7 @@ void BufferStore(BufferVar buffer, PrimExpr value, ffi::Array indices, } DeclBufferFrame DeclBuffer(ffi::Array shape, PrimType dtype, ffi::String buffer_name, - ffi::Optional data, - ffi::Optional> strides, + ffi::Optional data, ffi::Optional> strides, ffi::Optional elem_offset, ffi::String storage_scope, int align, int offset_factor, ffi::Optional layout, ffi::Optional allocated_addr) { @@ -849,6 +839,12 @@ DeclBufferFrame DeclBuffer(ffi::Array shape, PrimType dtype, ffi::Stri align, offset_factor, layout, allocated_addr_arr); if (data.has_value()) { n->data = data.value(); + } else if (scope == "tmem") { + // Tensor memory is an externally allocated address space. Make that + // address-to-pointer relationship explicit so every DeclBuffer has a + // physical data binding. + n->data = Call(n->buffer.DataPointerType(), tvm::tirx::builtin::reinterpret(), + {allocated_addr.value()}); } // For tmem, even without `data`, we should not emit an Allocate node. n->allocated = (scope == "tmem") || data.has_value(); @@ -856,9 +852,9 @@ DeclBufferFrame DeclBuffer(ffi::Array shape, PrimType dtype, ffi::Stri } BufferVar AllocBuffer(ffi::Array shape, PrimType dtype, ffi::String storage_scope, - ffi::Optional> annotations) { + ffi::Optional> annotations) { BufferVar buffer = BufferDecl(shape, dtype, "", std::nullopt, std::nullopt, std::nullopt, - storage_scope, 0, 0, std::nullopt, {}); + storage_scope, 0, 0, std::nullopt, {}); AddToParent( tvm::tirx::AllocBuffer(buffer, annotations.value_or(ffi::Map()))); return buffer; @@ -898,10 +894,10 @@ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() .def("script.ir_builder.tirx.Buffer", - static_cast, PrimType, ffi::String, ffi::Optional, - ffi::Optional>, ffi::Optional, ffi::String, int, - int, ffi::Optional, ffi::Array)>(BufferDecl)) + static_cast, PrimType, ffi::String, + ffi::Optional, ffi::Optional>, + ffi::Optional, ffi::String, int, int, + ffi::Optional, ffi::Array)>(BufferDecl)) .def("script.ir_builder.tirx.PrimFunc", PrimFunc) .def("script.ir_builder.tirx.Arg", [](ffi::String name, ffi::ObjectRef obj) -> ffi::ObjectRef { diff --git a/src/tirx/script/printer/buffer.cc b/src/tirx/script/printer/buffer.cc index b251bc4eb448..b90a0f569019 100644 --- a/src/tirx/script/printer/buffer.cc +++ b/src/tirx/script/printer/buffer.cc @@ -211,7 +211,7 @@ ffi::Map BufferAttrs(tirx::BufferVar buffer, const AccessP PrimExpr addr = buffer->allocated_addr[0]; AccessPath addr_p = buffer_p->Attr("allocated_addr")->ArrayItem(0); if (const auto* bl = addr.as()) { - // Ensure the buffer variable is defined (may emit a T.BufferVar(...) statement). + // Ensure the buffer variable is defined (may emit a T.Buffer(...) statement). d->AsDoc(bl->buffer, addr_p->Attr("buffer")); // Get the variable name bound to this buffer. ffi::Optional buf_var = d->GetVarDoc(bl->buffer); diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc index 23c2ce4f4dd5..87915fa1451d 100644 --- a/src/tirx/script/printer/expr.cc +++ b/src/tirx/script/printer/expr.cc @@ -46,7 +46,7 @@ ExprDoc PrintVarCreation(const tirx::Var& var, const AccessPath& var_p, const IR } else { ExprDoc element_type = LiteralDoc::DataType(prim_type->dtype, type_p->Attr("element_type")->Attr("dtype")); - if (ptr_type->storage_scope == "") { + if (ptr_type->storage_scope == "global") { rhs = rhs->Call({element_type}, kwargs_keys, kwargs_values); } else { rhs = rhs->Call({element_type, @@ -287,8 +287,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) Doc PrintTIRCall(Call call, AccessPath call_p, IRDocsifier d) { if (call->op.same_as(tirx::builtin::buffer_data())) { TVM_FFI_ICHECK_EQ(call->args.size(), 1); - return d->AsDoc(call->args[0], call_p->Attr("args")->ArrayItem(0)) - ->Attr("data"); + return d->AsDoc(call->args[0], call_p->Attr("args")->ArrayItem(0))->Attr("data"); } ffi::Optional call_prim_type = call->ty.as(); auto get_call_type_doc = [&](AccessPath type_p) -> ExprDoc { diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc index 661685a3b281..5b6c30e11185 100644 --- a/src/tirx/script/printer/stmt.cc +++ b/src/tirx/script/printer/stmt.cc @@ -251,13 +251,32 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) namespace { /*! - * \brief Find all parent buffers that share the same data pointer with the given child buffer. + * \brief Find the parent buffer named by a child's explicit data projection. * \param child The child buffer. + * \param data The child's explicit source pointer, if any. * \param d The IRDocsifier. * \return A list of candidate parent buffers. */ -std::vector FindParentBuffers(const tirx::BufferVar& child, const IRDocsifier& d) { - return {}; +std::vector FindParentBuffers(const tirx::BufferVar& child, + const ffi::Optional& data, + const IRDocsifier& d) { + if (!data.has_value()) { + return {}; + } + const auto* call = data.value().as(); + if (call == nullptr || !call->op.same_as(tirx::builtin::buffer_data()) || + call->args.size() != 1) { + return {}; + } + auto parent_var = call->args[0].as(); + if (!parent_var.has_value() || !parent_var.value()->ty.as()) { + return {}; + } + tirx::BufferVar parent(parent_var.value()); + if (parent.same_as(child) || !d->GetVarDoc(parent).has_value()) { + return {}; + } + return {parent}; } /*! @@ -274,8 +293,8 @@ bool IsDefaultLayout(const ffi::Optional& layout, const ffi::Array * * Returns std::nullopt if no sugar pattern matches. */ -ffi::Optional TryDeclBufferSugarWithParent(const tirx::BufferVar& child, const AccessPath& p, - const IRDocsifier& d, +ffi::Optional TryDeclBufferSugarWithParent(const tirx::BufferVar& child, + const AccessPath& p, const IRDocsifier& d, const tirx::BufferVar& parent) { ffi::Optional parent_doc = d->GetVarDoc(parent); if (!parent_doc.has_value()) return std::nullopt; @@ -605,7 +624,7 @@ ffi::Optional TryDeclBufferSugarWithParent(const tirx::BufferVar& child // --- (f) View(*shape, layout=L): different shape/layout, same dtype and elem_offset --- if (same_elem_offset && same_dtype && !same_shape) { - // BufferVar.view(...) copies the parent's strides onto the child (see + // Buffer.view(...) copies the parent's strides onto the child (see // python/tvm/tirx/buffer.py:view). If parent has strides but child // doesn't (or vice versa), the sugar can't faithfully round-trip // through view — fall back to T.decl_buffer where strides is an @@ -650,8 +669,8 @@ ffi::Optional TryDeclBufferSugarWithParent(const tirx::BufferVar& child * \brief Try to produce a DeclBuffer sugar expression, trying all parent buffer candidates. */ ffi::Optional TryDeclBufferSugar(const tirx::BufferVar& child, const AccessPath& p, - const IRDocsifier& d) { - auto parents = FindParentBuffers(child, d); + const ffi::Optional& data, const IRDocsifier& d) { + auto parents = FindParentBuffers(child, data, d); for (const auto& parent : parents) { if (auto sugar = TryDeclBufferSugarWithParent(child, p, d, parent)) { return sugar; @@ -664,14 +683,13 @@ Doc DeclBufferDoc(tirx::DeclBuffer stmt, AccessPath p, IRDocsifier d, BufferVarDefinition var_definitions) { // Try sugar detection when syntax_sugar is enabled if (d->cfg->syntax_sugar) { - if (auto sugar = TryDeclBufferSugar(stmt->buffer, p, d)) { + if (auto sugar = TryDeclBufferSugar(stmt->buffer, p, stmt->data, d)) { ExprDoc lhs = DefineBuffer(stmt->buffer, d->frames.back(), d); return AssignDoc(lhs, sugar.value(), std::nullopt); } } - ExprDoc rhs = - BufferDecl(stmt->buffer, "decl_buffer", {}, p->Attr("buffer"), d->frames.back(), d, - var_definitions, stmt->data); + ExprDoc rhs = BufferDecl(stmt->buffer, "decl_buffer", {}, p->Attr("buffer"), d->frames.back(), d, + var_definitions, stmt->data); ExprDoc lhs = DefineBuffer(stmt->buffer, d->frames.back(), d); return AssignDoc(lhs, rhs, std::nullopt); } diff --git a/src/tirx/transform/flatten_buffer.cc b/src/tirx/transform/flatten_buffer.cc index e525c2e3feff..72ea51f40c7a 100644 --- a/src/tirx/transform/flatten_buffer.cc +++ b/src/tirx/transform/flatten_buffer.cc @@ -121,14 +121,12 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { } Stmt VisitStmt_(const DeclBufferNode* op) final { - auto node = StmtExprMutator::VisitStmt_(op).as_or_throw(); - - auto new_buf = GetFlattenedBuffer(node->buffer); - if (!node->buffer.same_as(new_buf)) { - node.CopyOnWrite()->buffer = new_buf; + Expr data = VisitExpr(op->data); + BufferVar flattened = GetFlattenedBuffer(op->buffer); + if (flattened.same_as(op->buffer) && data.same_as(op->data)) { + return ffi::GetRef(op); } - - return std::move(node); + return DeclBuffer(flattened, std::move(data), op->span); } BufferVar GetFlattenedBuffer(BufferVar buf) { @@ -163,6 +161,19 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { return load; } + Expr VisitExpr_(const CallNode* op) final { + if (op->op.same_as(builtin::buffer_data()) && op->args.size() == 1) { + if (auto var = op->args[0].as()) { + if (var.value()->ty.as()) { + BufferVar original(var.value()); + buffers_used_.insert(original); + return GetFlattenedBuffer(original).data(); + } + } + } + return IRMutatorWithAnalyzer::VisitExpr_(op); + } + ffi::Array GetSimplifiedElemOffset(const BufferVar& buffer, const ffi::Array& indices) { auto flattened_indices = buffer->ElemOffset(indices); diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc index dfb359c2225d..11dd0ef1cea0 100644 --- a/src/tirx/transform/ir_utils.cc +++ b/src/tirx/transform/ir_utils.cc @@ -377,16 +377,14 @@ class IRConvertSSA final : public StmtExprMutator { if (layout_changed) { type->layout = std::move(new_layout); } - BufferVar new_buf = - RebuildBufferVar(buf, std::move(type), new_buffer_var->name); + BufferVar new_buf = RebuildBufferVar(buf, std::move(type), new_buffer_var->name); // A BufferVar's metadata lives in its Var type. If rewriting the // metadata required a fresh Var, make it the active remap as well. This // keeps BufferLoad/BufferStore and ordinary Var uses (such as // buffer_data) on the same identity. auto it = var_remap_.find(buf.get()); - if (it != var_remap_.end() && it->second.size() && - it->second.back().same_as(new_buffer_var)) { + if (it != var_remap_.end() && it->second.size() && it->second.back().same_as(new_buffer_var)) { it->second.back() = new_buf.var(); } else if (auto function_it = function_scope_var_remap_.find(buf.get()); function_it != function_scope_var_remap_.end() && @@ -572,8 +570,7 @@ class IRConvertSSA final : public StmtExprMutator { /*! \brief Pop a single variable remap (used for expression-level Let scoping). */ void PopVarRemap(const Var& old_var, const Var& new_var) { var_remap_[old_var.get()].pop_back(); - if (auto it = buf_remap_.find(old_var.get()); - it != buf_remap_.end() && it->second.size()) { + if (auto it = buf_remap_.find(old_var.get()); it != buf_remap_.end() && it->second.size()) { it->second.pop_back(); } // Also remove from the current scope's tracking vector @@ -672,7 +669,7 @@ Stmt ConvertSSA(Stmt stmt) { return IRConvertSSA()(std::move(stmt)); } ffi::String GetPtrStorageScope(Var buffer_var) { if (const auto* buffer_type = buffer_var->ty.as()) { - return buffer_type->data_pointer_type->storage_scope; + return buffer_type->storage_scope; } const auto* ptr_type = buffer_var->ty.as(); TVM_FFI_ICHECK(ptr_type) diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h index df6d3368a652..04da364c90a4 100644 --- a/src/tirx/transform/ir_utils.h +++ b/src/tirx/transform/ir_utils.h @@ -45,6 +45,7 @@ namespace tvm { namespace tirx { + /*! * \brief combine the nest stmt, whose body is not defined. * \param nest A list of For and Bind, whose body is not defined. @@ -114,7 +115,9 @@ inline PrimExpr TVMStructGet(PrimType type, Var handle, int index, inline Call AddressOffset(Var handle, PrimType dtype, int offset) { PrimExpr offset_expr = IntImm::Int32(offset * dtype.lanes()); ffi::Array shape = {offset_expr + 1}; - BufferVar dummy_buf(handle, dtype, shape, {}, 0, handle->name, 0, 0, Span(), std::nullopt); + auto pointer_type = handle->ty.as_or_throw(); + BufferVar dummy_buf(handle->name, + BufferType(pointer_type->storage_scope, dtype, shape, {}, 0, 0, 0)); BufferLoad buf_load(dummy_buf, {offset_expr}); return Call(handle->ty, builtin::address_of(), {buf_load}); @@ -134,8 +137,9 @@ inline Call AddressOffset(Var handle, PrimType dtype, PrimExpr offset) { } ffi::Array shape = {offset + 1}; - BufferVar dummy_buf(handle, dtype.WithLanes(1), shape, {}, 0, handle->name, 0, 0, Span(), - std::nullopt); + auto pointer_type = handle->ty.as_or_throw(); + BufferVar dummy_buf(handle->name, BufferType(pointer_type->storage_scope, dtype.WithLanes(1), + shape, {}, 0, 0, 0)); BufferLoad buf_load(dummy_buf, {offset}); return Call(handle->ty, builtin::address_of(), {buf_load}); diff --git a/src/tirx/transform/lower_intrin.cc b/src/tirx/transform/lower_intrin.cc index 1af07c8b4a59..db9ca6da79ef 100644 --- a/src/tirx/transform/lower_intrin.cc +++ b/src/tirx/transform/lower_intrin.cc @@ -66,6 +66,12 @@ static Expr LowerAccessPtr(const CallNode* call) { buffer = inner->args[1]; } + const auto* buffer_data = buffer.as(); + if (buffer_data && buffer_data->op.same_as(builtin::buffer_data())) { + TVM_FFI_ICHECK_EQ(buffer_data->args.size(), 1U); + buffer = buffer_data->args[0]; + } + const auto* buffer_node = buffer.as(); TVM_FFI_ICHECK(buffer_node) << "tvm_access_ptr expects a buffer Var or nested tvm_access_ptr as args[1], but got " @@ -76,8 +82,18 @@ static Expr LowerAccessPtr(const CallNode* call) { offset = offset * IntImm(offset_ty, dtype.lanes()); offset = Ramp(offset, IntImm(offset_ty, 1), dtype.lanes()); } - BufferVar dummy_buf(buffer_var, dtype.WithLanes(1), {offset + 1}, {}, 0, buffer_var->name, 0, 0); - BufferLoad buf_load(dummy_buf, {offset}); + BufferVar access_buffer{nullptr}; + if (buffer_var->ty.as()) { + access_buffer = BufferVar(buffer_var); + TVM_FFI_ICHECK_EQ(access_buffer->dtype, dtype.WithLanes(1)) + << "tvm_access_ptr element type must match the source buffer"; + } else { + auto pointer_type = buffer_var->ty.as_or_throw(); + access_buffer = BufferVar( + buffer_var->name, + BufferType(pointer_type->storage_scope, dtype.WithLanes(1), {offset + 1}, {}, 0, 0, 0)); + } + BufferLoad buf_load(access_buffer, {offset}); return Call(call->ty, builtin::address_of(), {buf_load}); } diff --git a/src/tirx/transform/lower_tirx_cleanup.cc b/src/tirx/transform/lower_tirx_cleanup.cc index f3a5cc3cdfe6..d89687b954e5 100644 --- a/src/tirx/transform/lower_tirx_cleanup.cc +++ b/src/tirx/transform/lower_tirx_cleanup.cc @@ -34,10 +34,12 @@ #include #include +#include #include #include #include "../../arith/ir_mutator_with_analyzer.h" +#include "ir_utils.h" namespace tvm { namespace tirx { @@ -45,12 +47,19 @@ namespace tirx { class LayoutApplier : public arith::IRMutatorWithAnalyzer { public: static std::pair> Flatten( - const Stmt& stmt, const ffi::Map buffer_map, const Target& target) { + const Stmt& stmt, const ffi::Array& params, + const ffi::Map buffer_map, const Target& target) { arith::Analyzer ana; LayoutApplier storage_lower(ana, target); + for (const Var& param : params) { + if (param->ty.as()) { + storage_lower.buffer_aliases_.Set(param, param); + } + } std::unordered_map new_buffer_map; std::vector> param_flattened_buffers; for (const auto& kv : buffer_map) { + storage_lower.buffer_aliases_.Set(kv.second.var(), kv.second.var()); if (kv.second->layout.has_value()) { BufferVar flattened = storage_lower.GetFlattenedBuffer(kv.second); auto type = CopyBufferType(kv.second); @@ -90,7 +99,30 @@ class LayoutApplier : public arith::IRMutatorWithAnalyzer { return any; } + Expr VisitExpr_(const VarNode* op) final { + Var var = ffi::GetRef(op); + if (auto it = var_remap_.find(var); it != var_remap_.end()) { + return it->second; + } + return IRMutatorWithAnalyzer::VisitExpr_(op); + } + + Expr VisitExpr_(const CallNode* op) final { + if (op->op.same_as(builtin::buffer_data()) && op->args.size() == 1) { + if (auto var = op->args[0].as(); + var.has_value() && var.value()->ty.as()) { + Var root = buffer_aliases_.Get(var.value()).value_or(var.value()); + if (auto it = var_remap_.find(root); it != var_remap_.end()) { + root = it->second; + } + return BufferVar(root).data(); + } + } + return IRMutatorWithAnalyzer::VisitExpr_(op); + } + Stmt VisitStmt_(const AllocBufferNode* op) final { + buffer_aliases_.Set(op->buffer.var(), op->buffer.var()); auto mutate = [this](BufferVar buf) { if (target_->kind->name == "trn" && !buf->layout.has_value()) { return buf; @@ -107,19 +139,19 @@ class LayoutApplier : public arith::IRMutatorWithAnalyzer { } Stmt VisitStmt_(const DeclBufferNode* op) final { + RegisterBufferAlias(op->buffer, op->data); + Expr data = VisitExpr(op->data); auto buffer = GetFlattenedBuffer(op->buffer); - if (buffer.same_as(op->buffer)) { + if (buffer.same_as(op->buffer) && data.same_as(op->data)) { return ffi::GetRef(op); } - auto n = CopyOnWrite(op); - n->buffer = buffer; - return Stmt(n); + return DeclBuffer(buffer, std::move(data), op->span); } BufferVar GetFlattenedBuffer(BufferVar buf, bool is_alloc = false) { - auto it = buffer_remap_.find(buf); - if (it != buffer_remap_.end()) { - return it->second; + auto it = var_remap_.find(buf.var()); + if (it != var_remap_.end()) { + return BufferVar(it->second); } auto trn_layout = buf->layout.as(); BufferVar flattened; @@ -169,14 +201,17 @@ class LayoutApplier : public arith::IRMutatorWithAnalyzer { type = CopyBufferType(flattened); } // canonicalize shape - for (size_t i = 0; i < flattened->shape.size(); ++i) { - type->shape.Set(i, analyzer_->canonical_simplify(flattened->shape[i])); + for (size_t i = 0; i < type->shape.size(); ++i) { + type->shape.Set(i, analyzer_->canonical_simplify(type->shape[i])); } type->layout = std::nullopt; type->elem_offset = StmtExprMutator::VisitPrimExpr(buf->elem_offset); + if (ffi::StructuralEqual()(buf.type(), BufferType(type))) { + return buf; + } flattened = RebuildBufferVar(flattened, std::move(type)); - buffer_remap_[buf] = flattened; + var_remap_[buf.var()] = flattened.var(); return flattened; } @@ -251,8 +286,28 @@ class LayoutApplier : public arith::IRMutatorWithAnalyzer { return node; } - /*! \brief Map of buffers being remapped. */ - std::unordered_map buffer_remap_; + /*! \brief Map of variables being remapped, including buffer variables. */ + std::unordered_map var_remap_; + + private: + void RegisterBufferAlias(BufferVar buffer, const Expr& data) { + Var root = buffer.var(); + if (const auto* call = data.as(); + call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) { + if (auto source = call->args[0].as(); + source.has_value() && source.value()->ty.as()) { + auto source_root = buffer_aliases_.Get(source.value()); + TVM_FFI_ICHECK(source_root.has_value()) + << "Buffer alias source " << source.value()->name + << " must be registered before its DeclBuffer alias"; + root = source_root.value(); + } + } + buffer_aliases_.Set(buffer.var(), root); + } + + /*! \brief Physical roots of buffer aliases, flattened at each declaration. */ + ffi::Map buffer_aliases_; const Target& target_; }; @@ -261,6 +316,14 @@ class BufferOffsetRemover : public StmtExprMutator { static Stmt Remove(const Stmt& stmt) { return BufferOffsetRemover()(stmt); } private: + Expr VisitExpr_(const VarNode* op) final { + Var var = ffi::GetRef(op); + if (auto it = var_remap_.find(var); it != var_remap_.end()) { + return it->second; + } + return StmtExprMutator::VisitExpr_(op); + } + Expr VisitExpr_(const CallNode* call) final { if (call->op.same_as(tirx::builtin::buffer_offset())) { auto buffer_load = call->args[0].as_or_throw(); @@ -272,18 +335,18 @@ class BufferOffsetRemover : public StmtExprMutator { Stmt VisitStmt_(const DeclBufferNode* op) { auto buffer = op->buffer; + Expr data = VisitExpr(op->data); auto elem_offset = this->VisitPrimExpr(buffer->elem_offset); - if (elem_offset.same_as(buffer->elem_offset)) { - return StmtExprMutator::VisitStmt_(op); - } else { + if (!elem_offset.same_as(buffer->elem_offset)) { auto type = CopyBufferType(buffer); type->elem_offset = std::move(elem_offset); buffer = RebuildBufferVar(buffer, std::move(type)); - buffer_remap_[op->buffer] = buffer; - auto n = CopyOnWrite(op); - n->buffer = buffer; - return Stmt(n); + var_remap_[op->buffer.var()] = buffer.var(); + } + if (buffer.same_as(op->buffer) && data.same_as(op->data)) { + return ffi::GetRef(op); } + return DeclBuffer(buffer, std::move(data), op->span); } using StmtExprMutator::VisitExpr_; @@ -304,16 +367,16 @@ class BufferOffsetRemover : public StmtExprMutator { template Node VisitBufferAccess(Node node) { TVM_FFI_ICHECK(node->buffer.defined()); - auto it = buffer_remap_.find(node->buffer); - if (it != buffer_remap_.end()) { + auto it = var_remap_.find(node->buffer.var()); + if (it != var_remap_.end()) { auto writer = node.CopyOnWrite(); - writer->buffer = it->second; + writer->buffer = BufferVar(it->second); return node; } return node; } - std::unordered_map buffer_remap_; + std::unordered_map var_remap_; }; namespace { @@ -332,7 +395,8 @@ Pass LowerTIRxCleanup() { auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) { Target target = ResolveTarget(f); auto* n = f.CopyOnWrite(); - std::tie(n->body, n->buffer_map) = LayoutApplier::Flatten(n->body, n->buffer_map, target); + std::tie(n->body, n->buffer_map) = + LayoutApplier::Flatten(n->body, n->params, n->buffer_map, target); n->body = BufferOffsetRemover::Remove(n->body); return f; }; diff --git a/src/tirx/transform/lower_tvm_builtin.cc b/src/tirx/transform/lower_tvm_builtin.cc index 7998cca2a5de..3020d58e5fe1 100644 --- a/src/tirx/transform/lower_tvm_builtin.cc +++ b/src/tirx/transform/lower_tvm_builtin.cc @@ -161,7 +161,7 @@ class BuiltinLower : public StmtExprMutator { auto& scope = alloca_scope_.back(); // Initial check to identify maximum stack sizes. These are used - // to construct BufferVar objects to hold the stack, which are then + // to construct buffers to hold the stack, which are then // used when mutating. scope.max_sizes = GetMaxStack(stmt); @@ -182,9 +182,8 @@ class BuiltinLower : public StmtExprMutator { scope.stack_shape = decl_buffer({IntImm::Int64(scope.max_sizes.shape_stack)}, PrimType::Int(64), "stack_shape"); stmt = SeqStmt::Flatten( - DeclBuffer(scope.stack_shape, - StackAlloca(scope.stack_shape->data_pointer_type, "shape", - scope.max_sizes.shape_stack)), + DeclBuffer(scope.stack_shape, StackAlloca(scope.stack_shape.DataPointerType(), "shape", + scope.max_sizes.shape_stack)), stmt); } @@ -262,7 +261,7 @@ class BuiltinLower : public StmtExprMutator { int64_t nbytes = GetVectorBytes(op->buffer->dtype); if (const auto* dev_type = device_type_.as(); dev_type && dev_type->value == kDLCPU) { - auto storage_scope = op->buffer->data_pointer_type.as_or_throw()->storage_scope; + auto storage_scope = op->buffer->storage_scope; if (storage_scope == "global") { auto constant_size = stmt.as_or_throw().ConstantAllocationSize(); if (constant_size.has_value() && constant_size.value() > 0 && @@ -281,14 +280,13 @@ class BuiltinLower : public StmtExprMutator { Call(PrimType::Int(32), builtin::tvm_throw_last_error(), {}).as_or_throw()); Stmt alloc_nullptr_check = IfThenElse( - Call(PrimType::Bool(), builtin::isnullptr(), {op->buffer.data()}) - .as_or_throw(), + Call(PrimType::Bool(), builtin::isnullptr(), {op->buffer.data()}).as_or_throw(), throw_last_error); static const Op& free_workspace_op = Op::Get("tirx.TVMBackendFreeWorkspace"); static const Op& alloc_workspace_op = Op::Get("tirx.TVMBackendAllocWorkspace"); PrimExpr free_op = Call(PrimType::Int(32), free_workspace_op, - {cast(PrimType::Int(32), device_type_.value()), + {cast(PrimType::Int(32), device_type_.value()), cast(PrimType::Int(32), device_id_.value()), op->buffer.data()}) .as_or_throw(); Stmt free_stmt = IfThenElse(free_op != IntImm::Int32(0), throw_last_error); @@ -297,11 +295,11 @@ class BuiltinLower : public StmtExprMutator { scope_.Current().pending_frees.push_back(free_stmt); Stmt alloc_bind = DeclBuffer( - op->buffer, Call(op->buffer->data_pointer_type, alloc_workspace_op, - {cast(PrimType::Int(32), device_type_.value()), - cast(PrimType::Int(32), device_id_.value()), total_bytes, - IntImm::Int32(op->buffer->dtype.code()), - IntImm::Int32(op->buffer->dtype.bits())})); + op->buffer, + Call(op->buffer.DataPointerType(), alloc_workspace_op, + {cast(PrimType::Int(32), device_type_.value()), + cast(PrimType::Int(32), device_id_.value()), total_bytes, + IntImm::Int32(op->buffer->dtype.code()), IntImm::Int32(op->buffer->dtype.bits())})); return SeqStmt({alloc_bind, alloc_nullptr_check}); } @@ -492,7 +490,7 @@ class BuiltinLower : public StmtExprMutator { } PrimExpr offset = ConstInt32(stack_begin); BufferLoad load(scope.stack_shape, {offset}); - return Call(scope.stack_shape->data_pointer_type, builtin::address_of(), {load}); + return Call(scope.stack_shape.DataPointerType(), builtin::address_of(), {load}); } // make array Expr MakeArray(const CallNode* op) { diff --git a/src/tirx/transform/lower_warp_memory.cc b/src/tirx/transform/lower_warp_memory.cc index 0d88980037fa..00f4303e3b4e 100644 --- a/src/tirx/transform/lower_warp_memory.cc +++ b/src/tirx/transform/lower_warp_memory.cc @@ -105,6 +105,17 @@ namespace tirx { // Visitor to find m in pattern // store warp_mem[m * warp_index + (width * m) * y + x] +const VarNode* GetBufferVar(const Expr& expr) { + if (const auto* var = expr.as()) { + return var; + } + if (const auto* call = expr.as(); + call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) { + return call->args[0].as(); + } + return nullptr; +} + class WarpStoreCoeffFinder : private StmtExprVisitor { public: WarpStoreCoeffFinder(const VarNode* buffer, Var warp_index, arith::AnalyzerObj* analyzer) @@ -122,18 +133,18 @@ class WarpStoreCoeffFinder : private StmtExprVisitor { static const Op& mma_fill_op = Op::Get("tirx.mma_fill"); static const Op& ptx_ldmatrix_legacy_op = Op::Get("tirx.ptx.ldmatrix_legacy"); static const Op& mma_fill_legacy_op = Op::Get("tirx.mma_fill_legacy"); - if (op->op.same_as(ptx_ldmatrix_op) && op->args[3].as() == buffer_) { + if (op->op.same_as(ptx_ldmatrix_op) && GetBufferVar(op->args[3]) == buffer_) { UpdatePattern(op->args[4].as_or_throw()); - } else if (op->op.same_as(mma_fill_op) && op->args[1].as() == buffer_) { + } else if (op->op.same_as(mma_fill_op) && GetBufferVar(op->args[1]) == buffer_) { auto* local_size = op->args[0].as(); TVM_FFI_ICHECK(local_size) << "Integer expected for the first argument of mma_fill"; warp_coeff_ = local_size->value; - } else if (op->op.same_as(ptx_ldmatrix_legacy_op) && op->args[3].as() == buffer_) { + } else if (op->op.same_as(ptx_ldmatrix_legacy_op) && GetBufferVar(op->args[3]) == buffer_) { // ldmatrix writes the warp buffer; its local_offset carries // ``... + lift(local_size) * tx`` from which the warp coefficient // is derived. UpdatePattern(op->args[4].as_or_throw()); - } else if (op->op.same_as(mma_fill_legacy_op) && op->args[1].as() == buffer_) { + } else if (op->op.same_as(mma_fill_legacy_op) && GetBufferVar(op->args[1]) == buffer_) { auto* local_size = op->args[0].as(); TVM_FFI_ICHECK(local_size) << "Integer expected for the first argument of mma_fill_legacy"; warp_coeff_ = local_size->value; @@ -279,10 +290,10 @@ class WarpAccessRewriter : protected StmtExprMutator { alloc_size = warp_group_ * factor; auto type = CopyBufferType(op->buffer); - type->data_pointer_type = PointerType(type->data_pointer_type->element_type, "local"); + type->storage_scope = "local"; type->shape = {IntImm::Int32(alloc_size / width_)}; type->strides = {}; - type->elem_offset = PrimExpr(); + type->elem_offset = IntImm(op->buffer->elem_offset.ty(), 0); BufferVar new_buf = RebuildBufferVar(op->buffer, std::move(type)); new_buffer_ = new_buf; Stmt rewritten_body = this->VisitStmt(body); @@ -294,8 +305,9 @@ class WarpAccessRewriter : protected StmtExprMutator { ffi::Array new_args = op->args; for (int i : indices) { // Preserve the pointer operand as an Expr and narrow only its scalar index. - if (op->args[i].as() == buffer_) { + if (GetBufferVar(op->args[i]) == buffer_) { PrimExpr local_index = SplitIndexByGroup(op->args[i + 1].as_or_throw()).first; + new_args.Set(i, new_buffer_.data()); new_args.Set(i + 1, local_index); } } diff --git a/src/tirx/transform/split_host_device.cc b/src/tirx/transform/split_host_device.cc index 54e3aebc8f84..1a895cbdd516 100644 --- a/src/tirx/transform/split_host_device.cc +++ b/src/tirx/transform/split_host_device.cc @@ -131,7 +131,8 @@ class HostDeviceSplitter : public StmtMutator { private: Stmt SplitDeviceFunc(Stmt body, Target device_target) { - auto [params, buffers_to_declare] = [&]() -> std::tuple, ffi::Array> { + auto [params, + buffers_to_declare] = [&]() -> std::tuple, ffi::Array> { VarUseDefAnalyzer use_def(/*defined_vars=*/{}, /*visit_thread_extent=*/true); use_def(body); @@ -140,7 +141,8 @@ class HostDeviceSplitter : public StmtMutator { if (device_target->kind->name != "trn") { std::sort(params.begin(), params.end(), [](const Var& a, const Var& b) { auto sort_key = [](const Var& var) { - bool is_handle = var->ty.as() != nullptr; + bool is_handle = + var->ty.as() != nullptr || var->ty.as() != nullptr; return std::tuple{ !is_handle, var->name, @@ -171,7 +173,7 @@ class HostDeviceSplitter : public StmtMutator { if (param->ty.as()) { BufferVar buffer(param); BufferVar kernel_buffer(buffer.name(), buffer.type(), buffer.span()); - Var data_param(buffer.name() + "_data", buffer->data_pointer_type); + Var data_param(buffer.name() + "_ptr", buffer.DataPointerType()); kernel_params.push_back(data_param); call_args.push_back(buffer.data()); buffer_data_params.Set(param, data_param); diff --git a/src/tirx/transform/stmt_simplify.cc b/src/tirx/transform/stmt_simplify.cc index d37e59576aaf..1be9222ec293 100644 --- a/src/tirx/transform/stmt_simplify.cc +++ b/src/tirx/transform/stmt_simplify.cc @@ -211,8 +211,7 @@ class StmtSimplifier : public IRMutatorWithAnalyzer { Stmt VisitStmt_(const BufferStoreNode* op) override { BufferStore store = Parent::VisitStmt_(op).as_or_throw(); if (const BufferLoadNode* load = store->value.as()) { - if (load->buffer.same_as(store->buffer) && - ArrayDeepEqual(load->indices, store->indices) && + if (load->buffer.same_as(store->buffer) && ArrayDeepEqual(load->indices, store->indices) && tirx::ExprDeepEqual()(load->buffer->elem_offset, store->buffer->elem_offset) && ArrayDeepEqual(load->buffer->shape, store->buffer->shape) && ArrayDeepEqual(load->buffer->strides, store->buffer->strides)) { diff --git a/src/tirx/transform/storage_rewrite.cc b/src/tirx/transform/storage_rewrite.cc index f174c776fa17..ba5e893c9093 100644 --- a/src/tirx/transform/storage_rewrite.cc +++ b/src/tirx/transform/storage_rewrite.cc @@ -54,6 +54,17 @@ using runtime::StorageScope; namespace { +ffi::Optional GetBufferDataVar(const ffi::Any& data) { + if (auto var = data.as()) { + return var; + } + if (const auto* call = data.as(); + call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) { + return call->args[0].as(); + } + return std::nullopt; +} + struct PrimTypeHash { size_t operator()(const PrimType& ty) const { DLDataType dtype = ty->dtype; @@ -84,6 +95,18 @@ struct PrimTypeEqual { // class LinearAccessPatternFinder final : public StmtExprVisitor { public: + LinearAccessPatternFinder(const ffi::Array& params, + const ffi::Map& buffer_map) { + for (const auto& [_, buffer] : buffer_map) { + buffer_aliases_.Set(buffer.var(), buffer.var()); + } + for (const Var& param : params) { + if (param->ty.as()) { + buffer_aliases_.Set(param, param); + } + } + } + /*! \brief record the touch hist of statment. */ struct StmtEntry { // The statment @@ -109,6 +132,7 @@ class LinearAccessPatternFinder final : public StmtExprVisitor { void VisitStmt_(const AllocBufferNode* op) final { size_t level = scope_.size(); const VarNode* buf = op->buffer.get(); + buffer_aliases_.Set(op->buffer.var(), op->buffer.var()); AllocEntry entry; entry.alloc = op; @@ -119,14 +143,15 @@ class LinearAccessPatternFinder final : public StmtExprVisitor { StmtExprVisitor::VisitStmt_(op); } + void VisitStmt_(const DeclBufferNode* op) final { RegisterBufferAlias(op->buffer, op->data); } + void VisitStmt_(const BufferStoreNode* op) final { scope_.push_back(StmtEntry()); // visit subexpr StmtExprVisitor::VisitStmt_(op); - all_buffers_accessed_.insert(op->buffer.get()); - // Add write access. - const VarNode* buffer_var = op->buffer.get(); + const VarNode* buffer_var = + buffer_aliases_.Get(op->buffer.var()).value_or(op->buffer.var()).get(); auto it = alloc_info_.find(buffer_var); if (it != alloc_info_.end() && it->second.alloc) { TVM_FFI_ICHECK_LT(it->second.level, scope_.size()); @@ -150,9 +175,8 @@ class LinearAccessPatternFinder final : public StmtExprVisitor { // Add write access. StmtExprVisitor::VisitExpr_(op); - all_buffers_accessed_.insert(op->buffer.get()); - - const VarNode* buffer_var = op->buffer.get(); + const VarNode* buffer_var = + buffer_aliases_.Get(op->buffer.var()).value_or(op->buffer.var()).get(); auto it = alloc_info_.find(buffer_var); if (it != alloc_info_.end() && it->second.alloc) { TVM_FFI_ICHECK_LT(it->second.level, scope_.size()) @@ -192,6 +216,10 @@ class LinearAccessPatternFinder final : public StmtExprVisitor { void VisitExpr_(const VarNode* buf) final { // Directly reference to the variable count as a read. + if (buf->ty.as()) { + Var var = ffi::GetRef(buf); + buf = buffer_aliases_.Get(var).value_or(var).get(); + } auto it = alloc_info_.find(buf); if (it != alloc_info_.end() && it->second.alloc) { TVM_FFI_ICHECK_LT(it->second.level, scope_.size()) << " buf=" << buf->name; @@ -259,11 +287,22 @@ class LinearAccessPatternFinder final : public StmtExprVisitor { std::vector linear_seq_; // The storage scope of each buffer std::unordered_map alloc_info_; - // A record of which BufferVar objects have been accessed, to prune - // unused DeclBuffer instances. - std::unordered_set all_buffers_accessed_; + // Physical roots of buffer aliases, flattened when each DeclBuffer is visited. + ffi::Map buffer_aliases_; private: + void RegisterBufferAlias(BufferVar buffer, const Expr& data) { + Var root = buffer.var(); + if (auto source = GetBufferDataVar(data); + source.has_value() && source.value()->ty.as()) { + auto source_root = buffer_aliases_.Get(source.value()); + TVM_FFI_ICHECK(source_root.has_value()) << "Buffer alias source " << source.value()->name + << " must be registered before its DeclBuffer alias"; + root = source_root.value(); + } + buffer_aliases_.Set(buffer.var(), root); + } + // Whether already in thread env. bool in_thread_env_{false}; // The scope stack. @@ -429,17 +468,16 @@ class StoragePlanRewriter : public StmtExprMutator { using StmtEntry = LinearAccessPatternFinder::StmtEntry; using AllocEntry = LinearAccessPatternFinder::AllocEntry; - Stmt Rewrite(Stmt stmt, bool detect_inplace, bool enable_reuse, - bool reuse_require_exact_matched_dtype) { + Stmt Rewrite(Stmt stmt, const ffi::Array& params, const ffi::Map& buffer_map, + bool detect_inplace, bool enable_reuse, bool reuse_require_exact_matched_dtype) { detect_inplace_ = detect_inplace; // plan the rewrite - LinearAccessPatternFinder finder; + LinearAccessPatternFinder finder(params, buffer_map); finder(stmt); this->LivenessAnalysis(finder.linear_seq_); this->PlanMemory(finder.linear_seq_, finder.alloc_info_, enable_reuse, reuse_require_exact_matched_dtype); - all_buffers_accessed_ = finder.all_buffers_accessed_; - alloc_info_ = finder.alloc_info_; + buffer_aliases_ = std::move(finder.buffer_aliases_); this->PrepareNewAlloc(); // start rewrite stmt = operator()(std::move(stmt)); @@ -451,7 +489,9 @@ class StoragePlanRewriter : public StmtExprMutator { template Node VisitBufferAccess(Node node) { - auto it = alloc_map_.find(node->buffer.get()); + const VarNode* root = + buffer_aliases_.Get(node->buffer.var()).value_or(node->buffer.var()).get(); + auto it = alloc_map_.find(root); if (it != alloc_map_.end()) { BufferVar buf = RemapBuffer(node->buffer, it->second->alloc_var); @@ -470,20 +510,19 @@ class StoragePlanRewriter : public StmtExprMutator { auto key = buf.get(); auto it = buffer_remap_.find(key); if (it != buffer_remap_.end()) { - TVM_FFI_ICHECK_EQ(buffer_backing_.at(it->second.get()).get(), new_backing_array.get()) + TVM_FFI_ICHECK_EQ(remapped_backing_.at(it->second.get()).get(), new_backing_array.get()) << "Cannot remap buffer " << buf.name() << " to use backing array " << new_backing_array->name << ", previously used backing array " - << buffer_backing_.at(it->second.get()).name(); + << remapped_backing_.at(it->second.get()).name(); return it->second; } BufferVar backing(new_backing_array); - BufferVar remapped = - buf.same_as(backing) - ? buf - : RebuildBufferVar(buf, CopyBufferType(buf), new_backing_array->name); + BufferVar remapped = buf.same_as(backing) + ? buf + : RebuildBufferVar(buf, CopyBufferType(buf), new_backing_array->name); buffer_remap_[key] = remapped; - buffer_backing_[remapped.get()] = backing; + remapped_backing_[remapped.get()] = backing; return remapped; } @@ -498,7 +537,12 @@ class StoragePlanRewriter : public StmtExprMutator { } Expr VisitExpr_(const VarNode* op) final { - auto it = alloc_map_.find(op); + const VarNode* root = op; + if (op->ty.as()) { + Var var = ffi::GetRef(op); + root = buffer_aliases_.Get(var).value_or(var).get(); + } + auto it = alloc_map_.find(root); if (it != alloc_map_.end()) { if (it->second->bits_offset != 0) { LOG(WARNING) << "Use a merged buffer variable address, could cause error"; @@ -513,10 +557,15 @@ class StoragePlanRewriter : public StmtExprMutator { TVM_FFI_ICHECK_EQ(op->args.size(), 5U); PrimExpr dtype_marker = op->args[0].as_or_throw(); PrimType dtype = dtype_marker.ty(); - const VarNode* buffer = op->args[1].as(); - if (buffer == nullptr) { + auto buffer_var = GetBufferDataVar(op->args[1]); + if (!buffer_var.has_value()) { return StmtExprMutator::VisitExpr_(op); } + const VarNode* buffer = buffer_var.value().get(); + if (buffer->ty.as()) { + Var var = buffer_var.value(); + buffer = buffer_aliases_.Get(var).value_or(var).get(); + } auto it = alloc_map_.find(buffer); if (it == alloc_map_.end()) { return StmtExprMutator::VisitExpr_(op); @@ -588,16 +637,12 @@ class StoragePlanRewriter : public StmtExprMutator { } Stmt VisitStmt_(const DeclBufferNode* op) final { - if (hoisted_buffer_decls_.count(op->buffer.get()) || - !all_buffers_accessed_.count(op->buffer.get())) { - return Evaluate(0); - } + const VarNode* root = buffer_aliases_.Get(op->buffer.var()).value_or(op->buffer.var()).get(); auto node = StmtExprMutator::VisitStmt_(op).as_or_throw(); - - if (auto it = alloc_map_.find(op->buffer.get()); it != alloc_map_.end()) { - BufferVar buf = RemapBuffer(op->buffer, it->second->alloc_var); + auto it = alloc_map_.find(root); + if (it != alloc_map_.end()) { auto writer = node.CopyOnWrite(); - writer->buffer = buf; + writer->buffer = RemapBuffer(node->buffer, it->second->alloc_var); writer->data = BufferVar(it->second->alloc_var).data(); } return node; @@ -789,9 +834,8 @@ class StoragePlanRewriter : public StmtExprMutator { combo_size = combo_size + IntImm::Int32(1); } combo_size = analyzer_->Simplify(combo_size); - BufferVar buf(e->alloc_var->name, - BufferType(PointerType(alloc_type, e->scope.to_string()), alloc_type, - {combo_size}, {}, PrimExpr(), 0, 0)); + BufferVar buf(e->alloc_var->name, BufferType(e->scope.to_string(), alloc_type, + {combo_size}, {}, PrimExpr(), 0, 0)); e->alloc_var = buf.var(); ffi::Map annotations; if (e->is_volatile) { @@ -829,10 +873,12 @@ class StoragePlanRewriter : public StmtExprMutator { uint64_t type_bits = e->elem_type.bits() * e->elem_type.lanes(); PrimExpr alloc_size = MakeConst(e->allocs[0]->buffer->shape[0].ty(), (total_bits + type_bits - 1) / type_bits); - BufferVar buf(e->alloc_var->name, - BufferType(PointerType(e->elem_type, e->scope.to_string()), e->elem_type, - {alloc_size}, {}, PrimExpr(), 0, 0)); + BufferVar buf(e->alloc_var->name, BufferType(e->scope.to_string(), e->elem_type, {alloc_size}, + {}, PrimExpr(), 0, 0)); e->alloc_var = buf.var(); + for (StorageEntry* child : e->merged_children) { + child->alloc_var = e->alloc_var; + } bool any_volatile = e->is_volatile; for (StorageEntry* child : e->merged_children) { if (child->is_volatile) any_volatile = true; @@ -1126,15 +1172,10 @@ class StoragePlanRewriter : public StmtExprMutator { std::vector> alloc_vec_; // The buffer objects being remapped std::unordered_map buffer_remap_; - // Explicit physical backing for each remapped buffer view. - std::unordered_map buffer_backing_; - // Buffers whose DeclBuffer has been hoisted to be adjacent to the new AllocBuffer location - std::unordered_set hoisted_buffer_decls_; - // Any buffers that is accessed at some point. DeclBuffer instances - // that do not appear in this list may be removed. - std::unordered_set all_buffers_accessed_; - // Copy of the allocation info from LinearAccessPatternFinder. - std::unordered_map alloc_info_; + // Physical backing chosen for each remapped buffer view. + std::unordered_map remapped_backing_; + // Physical roots of buffer aliases, flattened by LinearAccessPatternFinder. + ffi::Map buffer_aliases_; // analyzer arith::Analyzer analyzer_; }; @@ -1250,6 +1291,7 @@ class VectorTypeAccessChecker : public StmtExprVisitor { for (auto it : buffer_map) { BufferVar& buffer = it.second; Var buffer_var = buffer.var(); + buffer_aliases_.Set(buffer_var, buffer_var); PrimType dtype = buffer->dtype; PrimExpr extent = buffer->shape.size() ? buffer->shape[buffer->shape.size() - 1] : 0; OnArrayDeclaration(buffer_var, dtype, extent, BufferVarInfo::kPrimFuncParam); @@ -1260,10 +1302,10 @@ class VectorTypeAccessChecker : public StmtExprVisitor { for (Var buffer_var : params) { if (auto buffer_type = buffer_var->ty.as()) { BufferVar buffer(buffer_var); + buffer_aliases_.Set(buffer_var, buffer_var); PrimExpr extent = buffer->shape.size() ? buffer->shape[buffer->shape.size() - 1] : PrimExpr(0); - OnArrayDeclaration(buffer_var, buffer->dtype, extent, - BufferVarInfo::kPrimFuncParam); + OnArrayDeclaration(buffer_var, buffer->dtype, extent, BufferVarInfo::kPrimFuncParam); continue; } auto pointer_type = GetPointerType(buffer_var->ty); @@ -1293,13 +1335,13 @@ class VectorTypeAccessChecker : public StmtExprVisitor { void VisitExpr_(const CallNode* op) final { if (op->op.same_as(builtin::tvm_access_ptr())) { PrimType dtype = op->args[0].as_or_throw().ty(); - const VarNode* buffer = op->args[1].as(); + auto buffer_var = GetBufferDataVar(op->args[1]); PrimExpr index = op->args[2].as_or_throw(); // args[1] may be a nested Call (e.g. another tvm_access_ptr) rather // than a raw Var; OnArrayAccess derefs `buffer` so skip the record // here and let the recursive visit handle any inner buffer var. - if (buffer != nullptr) { - OnArrayAccess(dtype, buffer, {index}, false); + if (buffer_var.has_value()) { + OnArrayAccess(dtype, buffer_var.value().get(), {index}, false); } } else if (op->op.same_as(builtin::address_of())) { if (const auto* load = op->args[0].as()) { @@ -1311,6 +1353,7 @@ class VectorTypeAccessChecker : public StmtExprVisitor { } void VisitStmt_(const AllocBufferNode* op) final { + buffer_aliases_.Set(op->buffer.var(), op->buffer.var()); const ffi::Array& shape = op->buffer->shape; PrimExpr extent = shape.size() ? shape[shape.size() - 1] : PrimExpr(0); OnArrayDeclaration(op->buffer.var(), op->buffer->dtype, extent, @@ -1320,33 +1363,13 @@ class VectorTypeAccessChecker : public StmtExprVisitor { } void VisitStmt_(const DeclBufferNode* op) final { + RegisterBufferAlias(op->buffer, op->data); const ffi::Array& shape = op->buffer->shape; - PrimExpr extent = shape.size() ? shape[shape.size() - 1] : PrimExpr(0); - OnArrayDeclaration(op->buffer.var(), op->buffer->dtype, extent, - BufferVarInfo::kDeclBufferNode); - if (op->data.has_value()) { - if (auto source = op->data.value().as()) { - decl_buffer_sources_.emplace_back(op->buffer.get(), source.value().get()); - } - } - + PrimExpr extent = shape.size() ? shape.back() : PrimExpr(0); + OnArrayDeclaration(op->buffer.var(), op->buffer->dtype, extent, BufferVarInfo::kDeclBufferNode); StmtExprVisitor::VisitStmt_(op); } - void PropagateDeclBufferAccesses() { - for (const auto& [buffer, source] : decl_buffer_sources_) { - auto buffer_it = info_map_.find(buffer); - auto source_it = info_map_.find(source); - if (buffer_it == info_map_.end() || source_it == info_map_.end()) { - continue; - } - source_it->second.access_dtype.insert(buffer_it->second.access_dtype.begin(), - buffer_it->second.access_dtype.end()); - source_it->second.scalar_read_dtype.insert(buffer_it->second.scalar_read_dtype.begin(), - buffer_it->second.scalar_read_dtype.end()); - } - } - void VisitExpr_(const LetNode* op) final { HandleLetNode(op->var); StmtExprVisitor::VisitExpr_(op); @@ -1407,6 +1430,11 @@ class VectorTypeAccessChecker : public StmtExprVisitor { */ void OnArrayAccess(PrimType value_dtype, const VarNode* buffer, const ffi::Array& indices, bool is_buffer_load) { + Var buffer_var = ffi::GetRef(buffer); + const VarNode* root = buffer_aliases_.Get(buffer_var).value_or(buffer_var).get(); + if (info_map_.count(root)) { + buffer = root; + } auto it = info_map_.find(buffer); TVM_FFI_ICHECK(it != info_map_.end()) << "Load/Store of buffer " << buffer->name << " (" << buffer << ") occurred before its declaration."; @@ -1490,10 +1518,26 @@ class VectorTypeAccessChecker : public StmtExprVisitor { var_info.access_dtype.insert(access_dtype.WithLanes(lanes_used)); } + private: + void RegisterBufferAlias(BufferVar buffer, const Expr& data) { + Var root = buffer.var(); + if (auto source = GetBufferDataVar(data); + source.has_value() && source.value()->ty.as()) { + auto source_root = buffer_aliases_.Get(source.value()); + TVM_FFI_ICHECK(source_root.has_value()) << "Buffer alias source " << source.value()->name + << " must be registered before its DeclBuffer alias"; + root = source_root.value(); + } + buffer_aliases_.Set(buffer.var(), root); + } + + public: // Map of buffer variable information determined std::unordered_map info_map_; - std::vector> decl_buffer_sources_; + // Physical roots of buffer aliases, flattened at each declaration. + ffi::Map buffer_aliases_; + public: // bool allow_untyped_pointers_{false}; // Whether to detect scalar read patterns for rewriting to vector shuffle @@ -1550,11 +1594,11 @@ class VectorTypeRewriter : public StmtExprMutator { * should be re-written. */ VectorTypeRewriter(const std::unordered_map& info_map, - bool rewrite_params = true, bool rewrite_buffer_map = true, - bool rewrite_alloc_buffer_node = true, bool rewrite_indices = true, - bool rewrite_let_node = true, + const ffi::Map& buffer_aliases, bool rewrite_params = true, + bool rewrite_buffer_map = true, bool rewrite_alloc_buffer_node = true, + bool rewrite_indices = true, bool rewrite_let_node = true, bool rewrite_scalar_read_to_vector_shuffle = true) - : rewrite_indices_(rewrite_indices) { + : rewrite_indices_(rewrite_indices), buffer_aliases_(buffer_aliases) { int rewrite_mask = 0; if (rewrite_params) { rewrite_mask |= BufferVarInfo::kPrimFuncParam; @@ -1563,7 +1607,7 @@ class VectorTypeRewriter : public StmtExprMutator { rewrite_mask |= BufferVarInfo::kPrimFuncBufferMap; } if (rewrite_alloc_buffer_node) { - rewrite_mask |= BufferVarInfo::kAllocBufferNode | BufferVarInfo::kDeclBufferNode; + rewrite_mask |= BufferVarInfo::kAllocBufferNode; } if (rewrite_let_node) { rewrite_mask |= BufferVarInfo::kLetNode; @@ -1578,13 +1622,11 @@ class VectorTypeRewriter : public StmtExprMutator { if (old_buffer_var->ty.as()) { BufferVar old_buffer(old_buffer_var); auto type = CopyBufferType(old_buffer); - type->data_pointer_type = PointerType(preferred, old_buffer.scope()); type->dtype = preferred; if (!type->shape.empty()) { PrimExpr last_dim = type->shape.back(); int factor = preferred.lanes() / var_info.element_dtype.lanes(); - type->shape.Set(type->shape.size() - 1, - last_dim / MakeConst(last_dim.ty(), factor)); + type->shape.Set(type->shape.size() - 1, last_dim / MakeConst(last_dim.ty(), factor)); } type->layout = std::nullopt; return RebuildBufferVar(old_buffer, std::move(type)).var(); @@ -1613,7 +1655,8 @@ class VectorTypeRewriter : public StmtExprMutator { return {node, shuffle_index}; } - auto it = rewrite_map_.find(node->buffer.get()); + Var root = buffer_aliases_.Get(node->buffer.var()).value_or(node->buffer.var()); + auto it = rewrite_map_.find(root.get()); if (it == rewrite_map_.end()) { return {node, shuffle_index}; } @@ -1703,16 +1746,12 @@ class VectorTypeRewriter : public StmtExprMutator { } Stmt VisitStmt_(const DeclBufferNode* op) final { - Stmt stmt = StmtExprMutator::VisitStmt_(op); - op = stmt.as(); - TVM_FFI_ICHECK(op != nullptr); - BufferVar new_buf = RemapBuffer(op->buffer); - if (new_buf.same_as(op->buffer)) { - return stmt; + auto node = StmtExprMutator::VisitStmt_(op).as_or_throw(); + BufferVar new_buf = RemapBuffer(node->buffer); + if (!new_buf.same_as(node->buffer)) { + node.CopyOnWrite()->buffer = new_buf; } - auto n = CopyOnWrite(op); - n->buffer = std::move(new_buf); - return Stmt(n); + return node; } BufferVar RemapBuffer(BufferVar buf) { @@ -1723,10 +1762,23 @@ class VectorTypeRewriter : public StmtExprMutator { return cache_it->second; } - auto info_it = rewrite_map_.find(buf.get()); + Var root = buffer_aliases_.Get(buf.var()).value_or(buf.var()); + auto info_it = rewrite_map_.find(root.get()); if (info_it != rewrite_map_.end()) { auto& info = info_it->second; - buf = BufferVar(info.new_buffer_var); + if (root.same_as(buf.var())) { + buf = BufferVar(info.new_buffer_var); + } else { + auto type = CopyBufferType(buf); + type->dtype = info.new_element_dtype; + if (!type->shape.empty()) { + PrimExpr last_dim = type->shape.back(); + type->shape.Set(type->shape.size() - 1, + last_dim / MakeConst(last_dim.ty(), info.factor())); + } + type->layout = std::nullopt; + buf = RebuildBufferVar(buf, std::move(type)); + } } buffer_map_[cache_key] = buf; @@ -1734,6 +1786,12 @@ class VectorTypeRewriter : public StmtExprMutator { } Expr VisitExpr_(const CallNode* op) final { + if (op->op.same_as(builtin::buffer_data()) && op->args.size() == 1) { + if (auto var = op->args[0].as(); + var.has_value() && var.value()->ty.as()) { + return RemapBuffer(BufferVar(var.value())).data(); + } + } if (op->op.same_as(builtin::tvm_access_ptr())) { Expr expr = StmtExprMutator::VisitExpr_(op); op = expr.as(); @@ -1742,11 +1800,13 @@ class VectorTypeRewriter : public StmtExprMutator { return expr; } - const VarNode* buffer_var = op->args[1].as(); - if (buffer_var == nullptr) { + auto buffer = GetBufferDataVar(op->args[1]); + if (!buffer.has_value()) { return expr; } - auto it = rewrite_map_.find(buffer_var); + Var var = buffer.value(); + Var root = buffer_aliases_.Get(var).value_or(var); + auto it = rewrite_map_.find(root.get()); if (it == rewrite_map_.end()) { return expr; } @@ -1760,7 +1820,10 @@ class VectorTypeRewriter : public StmtExprMutator { int factor = info.factor(); extent = extent / MakeConst(extent.ty(), factor); index = index / MakeConst(index.ty(), factor); - ffi::Array acc_args{e_dtype, info.new_buffer_var, index, extent, flag}; + Expr data = info.new_buffer_var->ty.as() + ? BufferVar(info.new_buffer_var).data() + : Expr(info.new_buffer_var); + ffi::Array acc_args{e_dtype, data, index, extent, flag}; auto old_pointer_type = op->ty.as_or_throw(); Type new_pointer_type = PointerType(info.new_element_dtype, old_pointer_type->storage_scope); return Call(new_pointer_type, builtin::tvm_access_ptr(), acc_args); @@ -1871,6 +1934,7 @@ class VectorTypeRewriter : public StmtExprMutator { bool rewrite_indices_{true}; std::unordered_map rewrite_map_; std::unordered_map buffer_map_; + const ffi::Map& buffer_aliases_; arith::Analyzer analyzer_; }; @@ -1884,11 +1948,10 @@ PrimFunc PointerValueTypeRewrite(PrimFunc f, bool allow_untyped_pointers = false VectorTypeAccessChecker checker(f->params, f->buffer_map, allow_untyped_pointers, rewrite_scalar_read_to_vector_shuffle); checker(f->body); - checker.PropagateDeclBufferAccesses(); - VectorTypeRewriter rewriter(checker.info_map_, rewrite_params, rewrite_buffer_map, - rewrite_alloc_buffer_node, rewrite_indices, rewrite_let_node, - rewrite_scalar_read_to_vector_shuffle); + VectorTypeRewriter rewriter(checker.info_map_, checker.buffer_aliases_, rewrite_params, + rewrite_buffer_map, rewrite_alloc_buffer_node, rewrite_indices, + rewrite_let_node, rewrite_scalar_read_to_vector_shuffle); PrimFuncNode* n = f.CopyOnWrite(); n->body = rewriter(std::move(n->body)); rewriter.Finalize(&f); @@ -1917,8 +1980,8 @@ Pass StorageRewrite() { reuse_require_exact_matched_dtype = true; } auto* n = f.CopyOnWrite(); - n->body = StoragePlanRewriter().Rewrite(std::move(n->body), true, enable_reuse, - reuse_require_exact_matched_dtype); + n->body = StoragePlanRewriter().Rewrite(std::move(n->body), n->params, n->buffer_map, true, + enable_reuse, reuse_require_exact_matched_dtype); // Parameters may not be rewritten, but internal allocations may. return PointerValueTypeRewrite(std::move(f), true, false, false, true, true, true, false); }; diff --git a/src/tirx/transform/tvm_ffi_binder.cc b/src/tirx/transform/tvm_ffi_binder.cc index 1ff7714c52c1..2404cd3ae20c 100644 --- a/src/tirx/transform/tvm_ffi_binder.cc +++ b/src/tirx/transform/tvm_ffi_binder.cc @@ -603,7 +603,7 @@ void TVMFFIABIBuilder::DecodeAllParams() { ->Extend(AccessStep::ArrayItem(i)) ->Attr(ffi::String(buffer.name())); Expr data = DecodeParamDLTensor(buffer, device_type_, device_id_, param, - func_name_ + "." + param->name, param_path); + func_name_ + "." + param->name, param_path); decl_buffers_.push_back(DeclBuffer(buffer, data)); } } @@ -682,9 +682,9 @@ void TVMFFIABIBuilder::BindRegularStrides(const BufferVar& buffer, const Var& st // ============================================================ Expr TVMFFIABIBuilder::DecodeParamDLTensor(const BufferVar& buffer, const PrimExpr& device_type, - const PrimExpr& device_id, const Var& handle, - const std::string& arg_name, - ffi::reflection::AccessPath base_path) { + const PrimExpr& device_id, const Var& handle, + const std::string& arg_name, + ffi::reflection::AccessPath base_path) { const PrimType tvm_ndim_type = PrimType::Int(32); std::string buf_name = buffer.name(); @@ -806,7 +806,7 @@ Expr TVMFFIABIBuilder::DecodeParamDLTensor(const BufferVar& buffer, const PrimEx { ffi::reflection::AccessPath data_path = param_path->Attr(ffi::String("data")); Expr raw_data = TVMStructGet(PointerType::VoidPointerTy(), handle, 0, builtin::kDLTensorData); - Expr typed_data = Call(buffer->data_pointer_type, builtin::reinterpret(), {raw_data}); + Expr typed_data = Call(buffer.DataPointerType(), builtin::reinterpret(), {raw_data}); { Expr vptr = typed_data; @@ -850,11 +850,11 @@ Expr TVMFFIABIBuilder::DecodeParamDLTensor(const BufferVar& buffer, const PrimEx } // mark alignment of external bufs — must be after the alignment assertion // so the compiler does not emit aligned loads before the check fires. - asserts_.emplace_back(AttrStmt(vptr, tirx::attr::storage_alignment, + asserts_.emplace_back(AttrStmt(buffer.var(), tirx::attr::storage_alignment, IntImm::Int32(buffer->data_alignment), Evaluate(0))); } else { // Even without alignment check, mark alignment for the compiler. - init_nest_.emplace_back(AttrStmt(vptr, tirx::attr::storage_alignment, + init_nest_.emplace_back(AttrStmt(buffer.var(), tirx::attr::storage_alignment, IntImm::Int32(buffer->data_alignment), Evaluate(0))); } } diff --git a/src/tirx/transform/tvm_ffi_binder.h b/src/tirx/transform/tvm_ffi_binder.h index 9b1db11f87f2..c6b95291c49d 100644 --- a/src/tirx/transform/tvm_ffi_binder.h +++ b/src/tirx/transform/tvm_ffi_binder.h @@ -273,8 +273,8 @@ class TVMFFIABIBuilder { * \param base_path Base ffi::reflection::AccessPath for the buffer parameter. * \param fuzzy_match If true, allow value to have more dimensions than arg. */ - void BindBuffer(const BufferVar& arg, const BufferVar& value, ffi::reflection::AccessPath base_path, - bool fuzzy_match); + void BindBuffer(const BufferVar& arg, const BufferVar& value, + ffi::reflection::AccessPath base_path, bool fuzzy_match); /*! * \brief DLTensor bind: ndim/dtype/shape/strides/data/device assertions. diff --git a/src/tirx/transform/unsupported_dtype_legalize.cc b/src/tirx/transform/unsupported_dtype_legalize.cc index 2ab0093d936d..1bca19fbc7a7 100644 --- a/src/tirx/transform/unsupported_dtype_legalize.cc +++ b/src/tirx/transform/unsupported_dtype_legalize.cc @@ -70,9 +70,9 @@ bool MatchPrimType(const Type& type, F f) { // populate the buffer_remap and var_remap accordingly. class ComputeLegalizePlanner : public StmtExprVisitor { public: - ComputeLegalizePlanner( - std::unordered_map* buffer_remap, - std::unordered_map* var_remap, PrimType promote_dtype) + ComputeLegalizePlanner(std::unordered_map* buffer_remap, + std::unordered_map* var_remap, PrimType promote_dtype) : buffer_remap_(buffer_remap), var_remap_(var_remap), promote_dtype_(promote_dtype) {} // run planning to populate buffer remap and var remap. @@ -115,12 +115,7 @@ class ComputeLegalizePlanner : public StmtExprVisitor { // remap all intermediate constant buffer to promote data types (fp16/fp32) if (MatchType(op->buffer->dtype)) { PrimType dtype = promote_dtype_.WithLanes(op->buffer->dtype.lanes()); - ffi::String storage_scope = "global"; - if (auto* ptr_type = op->buffer->data_pointer_type.as()) { - storage_scope = ptr_type->storage_scope; - } auto type = CopyBufferType(op->buffer); - type->data_pointer_type = PointerType(dtype, storage_scope); type->dtype = dtype; BufferVar buffer_var = RebuildBufferVar(op->buffer, std::move(type)); (*var_remap_)[op->buffer.var()] = buffer_var.var(); @@ -133,6 +128,15 @@ class ComputeLegalizePlanner : public StmtExprVisitor { this->PopulateBufferRemap(op->buffer); } + void VisitExpr_(const CallNode* op) final { + if (op->op.same_as(builtin::buffer_data()) && op->args.size() == 1) { + if (auto buffer = op->args[0].as()) { + opaque_var_access_.insert(buffer.value()); + } + } + StmtExprVisitor::VisitExpr_(op); + } + void VisitExpr_(const VarNode* op) final { StmtExprVisitor::VisitExpr_(op); Var buffer_var = ffi::GetRef(op); @@ -158,9 +162,10 @@ class ComputeLegalizePlanner : public StmtExprVisitor { class BF16ComputeLegalizePlanner : public ComputeLegalizePlanner { public: - explicit BF16ComputeLegalizePlanner( - std::unordered_map* buffer_remap, - std::unordered_map* var_remap, PrimType promote_dtype) + explicit BF16ComputeLegalizePlanner(std::unordered_map* buffer_remap, + std::unordered_map* var_remap, + PrimType promote_dtype) : ComputeLegalizePlanner(buffer_remap, var_remap, promote_dtype) {} bool MatchType(const Type& type) const { return MatchPrimType(type, [](const PrimType& prim_type) { return IsBFloat16Type(prim_type); }); @@ -169,9 +174,10 @@ class BF16ComputeLegalizePlanner : public ComputeLegalizePlanner { class FP8ComputeLegalizePlanner : public ComputeLegalizePlanner { public: - explicit FP8ComputeLegalizePlanner( - std::unordered_map* buffer_remap, - std::unordered_map* var_remap, PrimType promote_dtype) + explicit FP8ComputeLegalizePlanner(std::unordered_map* buffer_remap, + std::unordered_map* var_remap, + PrimType promote_dtype) : ComputeLegalizePlanner(buffer_remap, var_remap, promote_dtype) {} bool MatchType(const Type& type) const { return MatchPrimType(type, [](const PrimType& prim_type) { return IsFloat8Type(prim_type); }); @@ -441,15 +447,12 @@ class ComputeLegalizer : public StmtExprMutator { } Stmt VisitStmt_(const DeclBufferNode* op) final { - Stmt ret = StmtExprMutator::VisitStmt_(op); - op = ret.as(); - + Expr data = VisitExpr(op->data); BufferVar new_buf = GetRemappedBuffer(op->buffer); - if (new_buf.same_as(op->buffer)) { - return ret; - } else { - return DeclBuffer(new_buf); + if (new_buf.same_as(op->buffer) && data.same_as(op->data)) { + return ffi::GetRef(op); } + return DeclBuffer(new_buf, std::move(data), op->span); } Stmt VisitStmt_(const AllocBufferNode* op) final { @@ -573,18 +576,13 @@ class StorageLegalizer : public StmtExprMutator { } Stmt VisitStmt_(const AllocBufferNode* op) final { - BufferVar buf = GetRemappedBuffer(op->buffer); + BufferVar buf = GetRemappedBuffer(op->buffer, /*allow_definition=*/true); // in a rare case the buffer didn't get remapped // because the original var is not bfloat* // force remap here if (MatchType(buf->dtype)) { PrimType new_dtype = GetStorageUIntDType(buf->dtype); - ffi::String storage_scope = "global"; - if (auto* ptr_type = buf->data_pointer_type.as()) { - storage_scope = ptr_type->storage_scope; - } auto type = CopyBufferType(buf); - type->data_pointer_type = PointerType(new_dtype, storage_scope); type->dtype = new_dtype; BufferVar new_buf = RebuildBufferVar(buf, std::move(type)); var_remap_[buf.var()] = new_buf.var(); @@ -601,23 +599,23 @@ class StorageLegalizer : public StmtExprMutator { } Stmt VisitStmt_(const DeclBufferNode* op) final { - BufferVar buf = GetRemappedBuffer(op->buffer); + BufferVar buf = GetRemappedBuffer(op->buffer, /*allow_definition=*/true); + Expr data = VisitExpr(op->data); // in a rare case the buffer didn't get remapped // because the original var is not bfloat* // force remap here if (MatchType(buf->dtype)) { PrimType new_dtype = GetStorageUIntDType(buf->dtype); auto type = CopyBufferType(buf); - type->data_pointer_type = PointerType(new_dtype, buf.scope()); type->dtype = new_dtype; buf = RebuildBufferVar(buf, std::move(type)); + var_remap_[op->buffer.var()] = buf.var(); buffer_remap_[op->buffer] = buf; } - if (buf.same_as(op->buffer)) { + if (buf.same_as(op->buffer) && data.same_as(op->data)) { return ffi::GetRef(op); - } else { - return DeclBuffer(buf, op->data, op->span); } + return DeclBuffer(buf, std::move(data), op->span); } Expr VisitExpr_(const LetNode* op) final { @@ -759,7 +757,7 @@ class StorageLegalizer : public StmtExprMutator { return var; } - BufferVar GetRemappedBuffer(BufferVar buf) { + BufferVar GetRemappedBuffer(BufferVar buf, bool allow_definition = false) { auto buf_it = buffer_remap_.find(buf); if (buf_it != buffer_remap_.end()) { return buf_it->second; @@ -768,7 +766,7 @@ class StorageLegalizer : public StmtExprMutator { auto var_it = var_remap_.find(buf.var()); if (var_it != var_remap_.end()) { new_buf = BufferVar(var_it->second); - } else { + } else if (!allow_definition) { TVM_FFI_ICHECK(!MatchType(buf->dtype)) << "Cannot find var remap for " << buf; } diff --git a/src/tirx/transform/update_pointer_storage_scope.cc b/src/tirx/transform/update_pointer_storage_scope.cc index b531c208510a..da4c2b140010 100644 --- a/src/tirx/transform/update_pointer_storage_scope.cc +++ b/src/tirx/transform/update_pointer_storage_scope.cc @@ -51,11 +51,8 @@ UpdatePointerStorageScope::UpdatePointerStorageScope( if (kv.first->ty.as()) { BufferVar buffer = GetBufferVar(kv.first); auto type = CopyBufferType(buffer); - const auto* pointer_type = type->data_pointer_type.as(); - TVM_FFI_ICHECK(pointer_type); - type->data_pointer_type = PointerType(pointer_type->element_type, kv.second); + type->storage_scope = kv.second; BufferVar replacement = RebuildBufferVar(buffer, std::move(type)); - new_buffer_remap_[kv.first] = replacement; new_var_remap_[kv.first] = replacement.var(); } else { new_var_remap_[kv.first] = WithStorageScope(kv.first, kv.second); @@ -82,15 +79,10 @@ Node UpdatePointerStorageScope::UpdateBufferAccess(Node node) { } BufferVar UpdatePointerStorageScope::GetUpdatedBuffer(BufferVar buf) { - // Use the cached buffer, if it exists. - auto key = buf.get(); - auto it = new_buffer_remap_.find(key); - if (it != new_buffer_remap_.end()) { - return it->second; + auto it = new_var_remap_.find(buf.get()); + if (it != new_var_remap_.end()) { + return BufferVar(it->second); } - - // Update the cache and return - new_buffer_remap_[key] = buf; return buf; } diff --git a/src/tirx/transform/update_pointer_storage_scope.h b/src/tirx/transform/update_pointer_storage_scope.h index f913a35be71a..8003b443a438 100644 --- a/src/tirx/transform/update_pointer_storage_scope.h +++ b/src/tirx/transform/update_pointer_storage_scope.h @@ -51,7 +51,6 @@ class UpdatePointerStorageScope : public StmtExprMutator { BufferVar GetUpdatedBuffer(BufferVar buf); std::unordered_map new_var_remap_; - std::unordered_map new_buffer_remap_; }; } // namespace tirx diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc index 45c0c257751e..6ca291ce43be 100644 --- a/tests/cpp/ir_functor_test.cc +++ b/tests/cpp/ir_functor_test.cc @@ -153,8 +153,7 @@ TEST(IRF, StmtVisitor) { auto z = x + 1; Stmt eval_body = Evaluate(z); PrimType dtype = PrimType::Float(32); - tirx::Var data_var("b", PointerType(dtype)); - BufferVar buf(data_var, dtype, {z, z}, {}, PrimExpr(), "b", 0, 0); + BufferVar buf("b", BufferType("global", dtype, {z, z}, {}, PrimExpr(), 0, 0)); // AllocBuffer is flat (no body). Return as SeqStmt with eval. return SeqStmt({AllocBuffer(buf), eval_body}); }; @@ -169,7 +168,7 @@ TEST(IRF, StmtVisitor) { PrimType dtype = PrimType::Float(32); tirx::Var buf_var("b", PointerType(dtype)); BufferVar buffer = decl_buffer({16}); - body = SeqStmt({DeclBuffer(buffer), std::move(body)}); + body = SeqStmt({DeclBuffer(buffer, buf_var), std::move(body)}); BufferRegion buffer_region(buffer, {Range::FromMinExtent(x + 1, 1)}); MatchBufferRegion match_buffer_region(decl_buffer({1}), buffer_region); @@ -183,8 +182,10 @@ TEST(IRF, StmtVisitor) { // x visited in: reads range (1), writes range (1), match_buffers range (1), // init DeclBuffer(0) + AllocBuffer shape(2) + Evaluate(1) = 3, // body DeclBuffer(0) + AllocBuffer shape(2) + Evaluate(1) = 3. - // Total: 1 + 1 + 1 + 3 + 3 = 9. - TVM_FFI_ICHECK_EQ(v.count, 9); + // The block's read/write BufferTypes each visit their dependent shape once, + // in addition to the ranges, match buffer, init, and body. + // Total: 2 + 2 + 1 + 3 + 3 = 11. + TVM_FFI_ICHECK_EQ(v.count, 11); } } @@ -207,8 +208,7 @@ TEST(IRF, StmtMutator) { auto fmakealloc = [&]() { auto z = x + 1; PrimType dtype = PrimType::Float(32); - tirx::Var data_var("b", PointerType(dtype)); - BufferVar buf(data_var, dtype, {1, z}, {}, PrimExpr(), "b", 0, 0); + BufferVar buf("b", BufferType("global", dtype, {1, z}, {}, PrimExpr(), 0, 0)); return AllocBuffer(buf); }; @@ -300,7 +300,8 @@ TEST(IRF, StmtMutator) { // AllocBuffer and DeclBuffer are flat (no body), placed as siblings in SeqStmt Stmt eval_body = Evaluate(x + 1); BufferVar buffer = decl_buffer({16}); - Stmt decl = DeclBuffer(buffer); + tirx::Var buffer_data("buffer_data", buffer.DataPointerType()); + Stmt decl = DeclBuffer(buffer, buffer_data); Stmt alloc = fmakealloc(); // body is: DeclBuffer, AllocBuffer, Evaluate Stmt body = SeqStmt({decl, alloc, eval_body}); @@ -335,14 +336,13 @@ TEST(IRF, Substitute) { PrimVar n("n", PrimType::Int(32)); auto fmakebuffer = [&]() { - return BufferVar{/*data=*/x, - /*dtype=*/PrimType::Float(32), - /*shape=*/{n}, - /*strides=*/{}, - /*elem_offset=*/PrimExpr(), - /*name=*/"buf", - /*data_alignment=*/1, - /*offset_factor=*/1}; + return BufferVar("buf", BufferType(/*storage_scope=*/"global", + /*dtype=*/PrimType::Float(32), + /*shape=*/{n}, + /*strides=*/{}, + /*elem_offset=*/PrimExpr(), + /*data_alignment=*/1, + /*offset_factor=*/1)); }; { @@ -364,7 +364,7 @@ TEST(IRF, Substitute) { TVM_FFI_ICHECK(seq_node != nullptr); auto* decl_node = seq_node->seq[0].as(); TVM_FFI_ICHECK(decl_node != nullptr); - TVM_FFI_ICHECK(decl_node->data.value().same_as(y)); + TVM_FFI_ICHECK(decl_node->data.same_as(y)); TVM_FFI_ICHECK(decl_node->buffer->shape[0].same_as(m)); TVM_FFI_ICHECK(!decl_node->buffer.same_as(buffer)); auto* store_node = seq_node->seq[1].as(); diff --git a/tests/python/codegen/test_target_codegen_llvm.py b/tests/python/codegen/test_target_codegen_llvm.py index 64f77c4acca9..ed5364646e9c 100644 --- a/tests/python/codegen/test_target_codegen_llvm.py +++ b/tests/python/codegen/test_target_codegen_llvm.py @@ -1142,6 +1142,22 @@ def main(a: T.handle("float64"), b: T.handle("float64"), n: T.int64): tvm.compile(Module, target="llvm") +@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") +def test_debug_symbol_for_buffer_var(): + """BufferVars use their physical data pointer type in LLVM debug info.""" + + @I.ir_module(s_tir=True) + class Module: + @T.prim_func(s_tir=True) + def main(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")): + C = T.alloc_buffer((16,), "float32") + for i in T.parallel(16): + C[i] = A[i] + B[i] = C[i] + + tvm.compile(Module, target="llvm") + + @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") def test_subroutine_call(): @I.ir_module(s_tir=True) diff --git a/tests/python/relax/test_analysis_suggest_layout_transforms.py b/tests/python/relax/test_analysis_suggest_layout_transforms.py index fb382acc914d..cd4e4e3d8c43 100644 --- a/tests/python/relax/test_analysis_suggest_layout_transforms.py +++ b/tests/python/relax/test_analysis_suggest_layout_transforms.py @@ -34,7 +34,7 @@ def apply_transformations(func, suggested_transfoms, print_transformation=False) print("Block transformation: ", block_name, " :: ", index_map) sch.transform_block_layout(block_name, index_map) else: - assert isinstance(obj, tirx.Buffer) + assert tirx.is_buffer_var(obj) buffer = obj if print_transformation: print("Buffer transformation: ", buffer, " :: ", index_map) diff --git a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py index 1948814a506d..5a6a97f618c3 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py @@ -304,8 +304,13 @@ def ref(B: T.Buffer((8, 8, 16, 16), "float32"), C: T.Buffer((128, 128), "float32 # T.reads(B[vi // 16 + vi_o, vj // 16 + vj_o, vi % 16, vj % 16]) # C[...] = B[vi // 16 + vi_o, vj // 16 + vj_o, vi % 16, vj % 16] + T.float32(1) - # not comparing PrimFuncs - tvm.ir.assert_structural_equal(ref.body.block.body, sch.get(sch.get_loops(block_outer)[0])) + expected = tvm.tirx.PrimFunc(list(ref.buffer_map.values()), ref.body.block.body) + actual_block = sch.get(block_outer) + actual = tvm.tirx.PrimFunc( + [actual_block.reads[0].buffer, actual_block.writes[0].buffer], + sch.get(sch.get_loops(block_outer)[0]), + ) + tvm.ir.assert_structural_equal(expected, actual) def test_var_args_sugar(): @@ -1021,7 +1026,7 @@ def main(A: T.Buffer(14, "int32")): sch = tvm.s_tir.Schedule(Before) A = sch.get(sch.get_sblock("block")).reads[0].buffer - other = tirx.decl_buffer(1, A.dtype, name="other") + other = tirx.decl_buffer(1, A.ty.dtype, name="other") with pytest.raises(tvm.s_tir.schedule.schedule.ScheduleError): sch.transform_layout( "block", @@ -1171,7 +1176,7 @@ def after(a: T.handle, b: T.handle, c: T.handle): # pylint: enable=invalid-name,line-too-long,too-many-locals # fmt: on # pylint: disable=invalid-name - _, _, n, _ = before.buffer_map[before.params[1]].shape + _, _, n, _ = before.buffer_map[before.params[1]].ty.shape sch = tvm.s_tir.Schedule(before) block = sch.get_sblock("NT_matmul") sch.transform_layout( @@ -1221,7 +1226,7 @@ def after(a: T.handle, b: T.handle, c: T.handle): # pylint: enable=invalid-name,line-too-long,too-many-locals # fmt: on # pylint: disable=invalid-name - _, _, n, _ = before.buffer_map[before.params[1]].shape + _, _, n, _ = before.buffer_map[before.params[1]].ty.shape sch = tvm.s_tir.Schedule(before) block = sch.get_sblock("NT_matmul") sch.transform_block_layout( diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py index a5f5a2c4e71f..aa7e7883e1d9 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# ruff: noqa: E501, F401, F841 +# ruff: noqa: E501, F401 import numpy as np import pytest @@ -495,8 +495,8 @@ def complex_compute( 16, 16, 16, - C.elem_offset // C_s0 // 16 * (C_s0 // 16) - + C.elem_offset % C_s0 // 16, + C.ty.elem_offset // C_s0 // 16 * (C_s0 // 16) + + C.ty.elem_offset % C_s0 // 16, T.float32(0), ) for k_0_0 in T.serial( @@ -687,12 +687,12 @@ def complex_compute( 16, 16, 16, - C.elem_offset // C_s0 // 16 * (C_s0 // 16) - + C.elem_offset % C_s0 // 16, + C.ty.elem_offset // C_s0 // 16 * (C_s0 // 16) + + C.ty.elem_offset % C_s0 // 16, T.tvm_access_ptr( T.type_annotation("float16"), A_1.data, - A_1.elem_offset, + A_1.ty.elem_offset, A_s0 * 16, 1, ), @@ -748,12 +748,12 @@ def complex_compute( 16, 16, 16, - C.elem_offset // C_s0 // 16 * (C_s0 // 16) - + C.elem_offset % C_s0 // 16, + C.ty.elem_offset // C_s0 // 16 * (C_s0 // 16) + + C.ty.elem_offset % C_s0 // 16, T.tvm_access_ptr( T.type_annotation("float16"), A_1.data, - A_1.elem_offset, + A_1.ty.elem_offset, A_s0 * 16, 1, ), @@ -826,17 +826,17 @@ def complex_compute( ) T.tvm_mma_sync( C.data, - C.elem_offset // C_s0 // 16 * (C_s0 // 16) - + C.elem_offset % C_s0 // 16, + C.ty.elem_offset // C_s0 // 16 * (C_s0 // 16) + + C.ty.elem_offset % C_s0 // 16, A_1.data, - A_1.elem_offset // A_s0 // 16 * (A_s0 // 16) - + A_1.elem_offset % A_s0 // 16, + A_1.ty.elem_offset // A_s0 // 16 * (A_s0 // 16) + + A_1.ty.elem_offset % A_s0 // 16, B.data, - B.elem_offset // B_s0 // 16 * (B_s0 // 16) - + B.elem_offset % B_s0 // 16, + B.ty.elem_offset // B_s0 // 16 * (B_s0 // 16) + + B.ty.elem_offset % B_s0 // 16, C.data, - C.elem_offset // C_s0 // 16 * (C_s0 // 16) - + C.elem_offset % C_s0 // 16, + C.ty.elem_offset // C_s0 // 16 * (C_s0 // 16) + + C.ty.elem_offset % C_s0 // 16, ) for ax0_0, ax1_0 in T.grid(2, 2): with T.sblock("Conv_reindex_wmma.accumulator_o"): @@ -876,12 +876,12 @@ def complex_compute( 16, 16, 16, - A_1.elem_offset // A_s0 // 16 * (A_s0 // 16) - + A_1.elem_offset % A_s0 // 16, + A_1.ty.elem_offset // A_s0 // 16 * (A_s0 // 16) + + A_1.ty.elem_offset % A_s0 // 16, T.tvm_access_ptr( T.type_annotation("float16"), C.data, - C.elem_offset, + C.ty.elem_offset, C_s0 * 16, 2, ), @@ -928,7 +928,7 @@ def main(A: T.Buffer((32, 128), "float16")): "float16", A_shared.data, tx * T.int64(128) + cse_v1 * T.int64(8), - A.data, + A_flattened.data, tx * T.int64(128) + cse_v1 * T.int64(8), 16, ) diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py index 3d1a63f49982..5eaf1f5d3792 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py @@ -62,7 +62,7 @@ def find_allocates(node): tvm.tirx.stmt_functor.post_order_visit(stmt.body, find_allocates) assert len(allocates) == 1 - assert list(allocates[0].buffer.shape) == [B_expected_alloc] + assert list(allocates[0].buffer.ty.shape) == [B_expected_alloc] def test_vthread_extern(): @@ -112,7 +112,7 @@ def find_allocates(node): tvm.tirx.stmt_functor.post_order_visit(stmt.body, find_allocates) assert len(allocates) == 3 # Check that we have the expected extents (order may vary) - extents = sorted([int(a.buffer.shape[0]) for a in allocates]) + extents = sorted([int(a.buffer.ty.shape[0]) for a in allocates]) assert extents == sorted([A_expected_alloc, A_expected_alloc, C_expected_alloc]) @@ -167,16 +167,15 @@ def before_func(): B = T.alloc_buffer((4,), "int32", scope="shared") B[0:4] = T.broadcast(vthread, 4) - @T.prim_func(check_well_formed=False, s_tir=True) + @T.prim_func(s_tir=True) def expected_func(): B = T.alloc_buffer((16,), "int32", scope="shared") - B_1 = T.Buffer([16], "int32", data=B.data, scope="shared") # The indices for B should each be a single Ramp node, and # should not be the sum of a Ramp and Broadcast node. - B_1[T.Mul(0, 4) : T.Mul(0, 4) + 4] = T.broadcast(0, 4) - B_1[T.Mul(1, 4) : T.Mul(1, 4) + 4] = T.broadcast(1, 4) - B_1[T.Mul(2, 4) : T.Mul(2, 4) + 4] = T.broadcast(2, 4) - B_1[T.Mul(3, 4) : T.Mul(3, 4) + 4] = T.broadcast(3, 4) + B[T.Mul(0, 4) : T.Mul(0, 4) + 4] = T.broadcast(0, 4) + B[T.Mul(1, 4) : T.Mul(1, 4) + 4] = T.broadcast(1, 4) + B[T.Mul(2, 4) : T.Mul(2, 4) + 4] = T.broadcast(2, 4) + B[T.Mul(3, 4) : T.Mul(3, 4) + 4] = T.broadcast(3, 4) before_mod = tvm.IRModule.from_expr(before_func.with_attr("global_symbol", "main")) after_mod = tvm.s_tir.transform.InjectVirtualThread()(before_mod) @@ -210,8 +209,8 @@ def visitor(op): tvm.tirx.stmt_functor.post_order_visit(after_func.body, visitor) assert allocate_node is not None - assert list(allocate_node.buffer.shape) == [4] - assert allocate_node.buffer.dtype == "int32x4" + assert list(allocate_node.buffer.ty.shape) == [4] + assert allocate_node.buffer.ty.dtype == "int32x4" if __name__ == "__main__": diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py index eaa78cf626a1..601931ab52a3 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py @@ -154,6 +154,25 @@ def transformed_opaque_access(a: T.handle, b: T.handle) -> None: ) +@T.prim_func(s_tir=True) +def opaque_buffer_data_projection(a: T.handle) -> None: + A = T.match_buffer(a, (16,)) + with T.sblock(): + T.reads([]) + T.writes(A[4:8]) + sub_A = T.match_buffer(A[4:8], (4,), offset_factor=1) + T.evaluate(T.call_extern("consume", sub_A.data, sub_A.elem_offset, dtype="int32")) + + +@T.prim_func(s_tir=True) +def transformed_opaque_buffer_data_projection(a: T.handle) -> None: + A = T.match_buffer(a, (16,)) + with T.sblock(): + T.reads([]) + T.writes(A[4:8]) + T.evaluate(T.call_extern("consume", A.data, 4, dtype="int32")) + + @T.prim_func(s_tir=True) def high_dim_opaque_access(a: T.handle) -> None: A = T.match_buffer(a, (16, 32, 64)) @@ -501,6 +520,7 @@ def test_buffer_load_store(): def test_opaque_access(): _check(opaque_access, transformed_opaque_access) + _check(opaque_buffer_data_projection, transformed_opaque_buffer_data_projection) def test_high_dim_opaque_access(): diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py index 3685740432c7..e334efbe2584 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py @@ -107,6 +107,7 @@ def main(A: T.Buffer((128, 32), "float32"), B: T.Buffer(128, "float32")): After = transform(Before) assert After is not None + assert tvm.tirx.analysis.verify_well_formed(After) After_script = After.script() assert "tvm_warp_shuffle" in After_script @@ -443,6 +444,7 @@ def main(A: T.Buffer((128, 32), "float32"), B: T.Buffer(128, "float32")): reduce_data = T.alloc_buffer((1,), "float32", scope="local") reduce = T.decl_buffer(1, data=reduce_data.data, scope="local") + reduce_alias = T.decl_buffer(1, data=reduce.data, scope="local") with T.attr( T.comm_reducer(lambda x, y: x + y, [T.float32(0)]), @@ -453,14 +455,15 @@ def main(A: T.Buffer((128, 32), "float32"), B: T.Buffer(128, "float32")): T.uint32(1), A_flat[0], T.bool(True), - reduce[0], + reduce_alias[0], threadIdx_x, ) if threadIdx_x == 0: - B[i] = reduce[0] + B[i] = reduce_alias[0] After = transform(Before) assert After is not None + assert tvm.tirx.analysis.verify_well_formed(After) After_script = After.script() assert "tvm_warp_shuffle_down" in After_script assert "tvm_warp_shuffle(" in After_script diff --git a/tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py b/tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py index 86ff73d273d2..2da847042970 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py @@ -268,9 +268,37 @@ def main(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")): script = After["main"].script() # Verify merged allocation (512 bytes - A_sh and B_sh can be reused) assert '"uint8"' in script and '"shared.dyn"' in script and "(512,)" in script - # Verify cp_async uses the merged buffer - assert "buf_dyn_shmem" in script - assert "threadIdx_x * 4" in script + # Verify cp_async uses typed views of the merged byte allocation. Its + # offsets remain in float32 elements and are scaled by the intrinsic + # lowering, rather than being pre-scaled as byte offsets here. + assert 'A_sh = buf_dyn_shmem.view("float32")' in script + assert 'B_sh = buf_dyn_shmem.view("float32")' in script + assert 'T.ptx.cp_async("float32", A_sh.data, threadIdx_x' in script + assert 'T.ptx.cp_async("float32", B_sh.data, threadIdx_x' in script + + +def test_decl_buffer_alias_extends_allocation_lifetime(): + """Access through a typed view keeps its source allocation live.""" + transform = tvm.s_tir.transform.MergeSharedMemoryAllocations() + + @I.ir_module + class Before: + @T.prim_func(s_tir=True) + def main(C: T.Buffer((128,), "float32")): + threadIdx_x = T.launch_thread("threadIdx.x", 128) + A_sh = T.alloc_buffer((128,), "float32", scope="shared.dyn") + B_sh = T.alloc_buffer((128,), "float32", scope="shared.dyn") + A_view = T.decl_buffer((128,), "float32", data=A_sh.data, scope="shared.dyn") + B_view = T.decl_buffer((128,), "float32", data=B_sh.data, scope="shared.dyn") + A_view[threadIdx_x] = T.float32(1) + B_view[threadIdx_x] = T.float32(2) + C[threadIdx_x] = A_view[threadIdx_x] + B_view[threadIdx_x] + + After = transform(Before) + script = After["main"].script() + assert 'alloc_buffer((1024,), "uint8", scope="shared.dyn")' in script + assert "B_view[threadIdx_x + 128]" in script + assert "A_view[threadIdx_x]" in script def test_multi_thread_extent_blocks(): diff --git a/tests/python/te/test_te_tensor.py b/tests/python/te/test_te_tensor.py index 62fc5a529a60..959af34c95bb 100644 --- a/tests/python/te/test_te_tensor.py +++ b/tests/python/te/test_te_tensor.py @@ -168,7 +168,7 @@ def test_extern(): A = te.placeholder((m,), name="A") def extern_func(ins, outs): - assert isinstance(ins[0], tvm.tirx.Buffer) + assert tvm.tirx.is_buffer_var(ins[0]) return tvm.tirx.call_packed("myadd", ins[0].data, outs[0].data, m) B = te.extern((m,), [A], extern_func) @@ -181,7 +181,7 @@ def test_extern_multi_out(): B = te.compute((m,), lambda i: A[i] * 10) def extern_func(ins, outs): - assert isinstance(ins[0], tvm.tirx.Buffer) + assert tvm.tirx.is_buffer_var(ins[0]) return tvm.tirx.call_packed("myadd", ins[0].data, outs[0].data, outs[1].data, m) res = te.extern([A.shape, A.shape], [A, B], extern_func) diff --git a/tests/python/tirx-analysis/test_tir_analysis_undefined_vars.py b/tests/python/tirx-analysis/test_tir_analysis_undefined_vars.py index fc080ce2ee7d..1a636f495754 100644 --- a/tests/python/tirx-analysis/test_tir_analysis_undefined_vars.py +++ b/tests/python/tirx-analysis/test_tir_analysis_undefined_vars.py @@ -35,7 +35,7 @@ def test_decl_buffer_data_is_use(): buf = tirx.decl_buffer((n,), "float32", "buf", data=data_ptr) body = tirx.Evaluate(tirx.BufferLoad(buf, [0])) - decl = tirx.DeclBuffer(buf) + decl = tirx.DeclBuffer(buf, data=data_ptr) stmt = tirx.SeqStmt([decl, body]) undef = tvm.tirx.analysis.undefined_vars(stmt, []) @@ -58,7 +58,7 @@ def test_decl_buffer_elem_offset_is_use(): buf = tirx.decl_buffer((n,), "float32", "buf", data=data_ptr, elem_offset=elem_off) body = tirx.Evaluate(tirx.BufferLoad(buf, [0])) - decl = tirx.DeclBuffer(buf) + decl = tirx.DeclBuffer(buf, data=data_ptr) stmt = tirx.SeqStmt([decl, body]) undef = tvm.tirx.analysis.undefined_vars(stmt, []) @@ -84,13 +84,21 @@ def test_alloc_buffer_data_is_def(): undef = tvm.tirx.analysis.undefined_vars(stmt, []) undef_names = {v.name for v in undef} - # data should NOT be undefined — AllocBuffer defines it - assert buf.data.name not in undef_names, ( - f"AllocBuffer data should be defined, but found {buf.data.name} in {undef_names}" - ) + # The buffer Var itself is defined by AllocBuffer. + assert buf.name not in undef_names # shape var n should be undefined (comes from enclosing scope) assert "n" in undef_names, f"Expected shape var 'n' in undefined vars, got {undef_names}" +def test_buffer_data_projection_is_buffer_use(): + """An opaque data projection must retain the BufferVar identity.""" + buf = tirx.decl_buffer((16,), "float32", "buf") + stmt = tirx.Evaluate(buf.data) + + undef = tvm.tirx.analysis.undefined_vars(stmt, []) + assert len(undef) == 1 + assert undef[0].same_as(buf) + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx-base/test_tir_buffer.py b/tests/python/tirx-base/test_tir_buffer.py index 78f7ed3d17ae..9db789c1484b 100644 --- a/tests/python/tirx-base/test_tir_buffer.py +++ b/tests/python/tirx-base/test_tir_buffer.py @@ -22,7 +22,7 @@ import tvm import tvm.testing from tvm.script import tirx as T -from tvm.tirx import Buffer +from tvm.tirx import BufferAccessKind def test_buffer(): @@ -33,29 +33,63 @@ def test_buffer(): Bb = tvm.tirx.decl_buffer((n, l), "float32") assert type(Ab) is tvm.ir.Var - assert tvm.tirx.is_buffer(Ab) + assert tvm.tirx.is_buffer_var(Ab) assert isinstance(Ab.ty, tvm.tirx.BufferType) - assert Ab.dtype == tvm.ir.PrimType("float32") - assert tuple(Ab.shape) == (m, n) - assert not tvm.tirx.is_buffer(m) + assert Ab.ty.dtype == tvm.ir.PrimType("float32") + assert tuple(Ab.ty.shape) == (m, n) + assert not tvm.tirx.is_buffer_var(m) + + +def test_buffer_compatibility_alias_and_global_var_properties(): + scalar = tvm.ir.Var("scalar", tvm.ir.PrimType("int32")) + buffer = tvm.tirx.decl_buffer((8,), "float32") + + assert tvm.tirx.Buffer is tvm.ir.Var + assert isinstance(scalar, tvm.tirx.Buffer) + assert not tvm.tirx.is_buffer_var(scalar) + assert tvm.tirx.is_buffer_var(buffer) + + assert tuple(buffer.shape) == (8,) + assert buffer.dtype == tvm.DataType("float32") + assert buffer.data.args[0].same_as(buffer) + + for name in ("shape", "dtype", "data"): + assert not hasattr(scalar, name) + with pytest.raises(AttributeError, match="only available on a Var with BufferType"): + getattr(scalar, name) def test_buffer_data_is_typed_projection(): buffer = tvm.tirx.decl_buffer((8,), "bool", scope="shared") - assert buffer.dtype == tvm.ir.PrimType("bool") - assert buffer.data_pointer_type == tvm.ir.PointerType(tvm.ir.PrimType("int8"), "shared") + assert buffer.ty.dtype == tvm.ir.PrimType("bool") + assert tvm.tirx.buffer_data_pointer_type(buffer) == tvm.ir.PointerType( + tvm.ir.PrimType("bool"), "shared" + ) assert buffer.data.op.name == "tirx.buffer_data" assert buffer.data.args[0].same_as(buffer) - assert buffer.data.ty == buffer.data_pointer_type + assert buffer.data.ty == tvm.tirx.buffer_data_pointer_type(buffer) -def test_buffer_logical_dtype_independent_of_pointer_type(): +def test_buffer_pointer_type_derived_from_dtype_and_scope(): data = tvm.ir.Var("storage", tvm.ir.PointerType(tvm.ir.PrimType("uint8"), "local")) buffer = tvm.tirx.decl_buffer((8,), "float16", data=data) - assert buffer.dtype == tvm.ir.PrimType("float16") - assert buffer.data.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint8"), "local") + assert buffer.ty.dtype == tvm.ir.PrimType("float16") + assert buffer.ty.storage_scope == "local" + assert buffer.data.ty == tvm.ir.PointerType(tvm.ir.PrimType("float16"), "local") + + +def test_decl_buffer_requires_physical_data_binding(): + buffer = tvm.tirx.decl_buffer((8,), "float32") + data = tvm.tirx.Var("data", buffer.data.ty) + + with pytest.raises(TypeError, match="requires a physical data binding"): + tvm.tirx.DeclBuffer(buffer) + + decl = tvm.tirx.DeclBuffer(buffer, data=data) + assert decl.buffer.same_as(buffer) + assert decl.data.same_as(data) def test_buffer_access_ptr(): @@ -65,9 +99,9 @@ def test_buffer_access_ptr(): aptr = Ab.access_ptr("rw") assert isinstance(aptr.ty, tvm.ir.PointerType) assert aptr.ty.element_type == tvm.ir.PrimType("void") - tvm.ir.assert_structural_equal(aptr.args[3], Ab.strides[0] * m) - assert aptr.args[0].ty == Ab.dtype - assert aptr.args[4].value == Buffer.READ | Buffer.WRITE + tvm.ir.assert_structural_equal(aptr.args[3], Ab.ty.strides[0] * m) + assert aptr.args[0].ty == Ab.ty.dtype + assert aptr.args[4].value == BufferAccessKind.READ | BufferAccessKind.WRITE typed_ptr = Ab.access_ptr("r", ptr_type="uint8") assert typed_ptr.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint8")) shared = tvm.tirx.decl_buffer((m, n), "float32", scope="shared") @@ -76,7 +110,7 @@ def test_buffer_access_ptr(): tvm.ir.PrimType("uint8"), "shared" ) aptr = Ab.access_ptr("w") - assert aptr.args[4].value == Buffer.WRITE + assert aptr.args[4].value == BufferAccessKind.WRITE def test_buffer_access_ptr_offset(): @@ -85,16 +119,16 @@ def test_buffer_access_ptr_offset(): Ab = tvm.tirx.decl_buffer((m, n), "float32") aptr = Ab.access_ptr("rw", offset=100) tvm.testing.assert_prim_expr_equal(aptr.args[2], 100) - assert aptr.args[4].value == Buffer.READ | Buffer.WRITE + assert aptr.args[4].value == BufferAccessKind.READ | BufferAccessKind.WRITE v = tvm.tirx.Var("int32", "int32") aptr = Ab.access_ptr("rw", offset=100 + 100 + v) tvm.testing.assert_prim_expr_equal(aptr.args[2], 200 + v) - assert aptr.args[4].value == Buffer.READ | Buffer.WRITE + assert aptr.args[4].value == BufferAccessKind.READ | BufferAccessKind.WRITE aptr = Ab.access_ptr("rw", offset=tvm.tirx.call_extern("int32", "test_call", 100 + 100 + v)) tvm.testing.assert_prim_expr_equal( aptr.args[2], tvm.tirx.call_extern("int32", "test_call", 200 + v) ) - assert aptr.args[4].value == Buffer.READ | Buffer.WRITE + assert aptr.args[4].value == BufferAccessKind.READ | BufferAccessKind.WRITE def test_buffer_access_ptr_extent(): @@ -107,7 +141,7 @@ def test_buffer_access_ptr_extent(): tvm.ir.assert_structural_equal(aptr.args[3], m * n - 100) Ab = tvm.tirx.decl_buffer((m, n), "float32", strides=[n + 1, 1]) aptr = Ab.access_ptr("rw", offset=100) - tvm.ir.assert_structural_equal(aptr.args[3], Ab.strides[0] * m - 100) + tvm.ir.assert_structural_equal(aptr.args[3], Ab.ty.strides[0] * m - 100) # Test extent from input params aptr = Ab.access_ptr("rw", extent=200) @@ -218,7 +252,7 @@ def test_buffer_flatten(): assert not buf.same_as(flat) assert flat.data.args[0].same_as(flat) assert flat.data.op.name == "tirx.buffer_data" - tvm.ir.assert_structural_equal(flat.shape, [T.int32(16 * 32)]) + tvm.ir.assert_structural_equal(flat.ty.shape, [T.int32(16 * 32)]) def test_buffer_flatten_preserves_identity(): diff --git a/tests/python/tirx-base/test_tir_constructor.py b/tests/python/tirx-base/test_tir_constructor.py index 11e6d72f3332..4be31c01748c 100644 --- a/tests/python/tirx-base/test_tir_constructor.py +++ b/tests/python/tirx-base/test_tir_constructor.py @@ -111,7 +111,8 @@ def test_expr_constructor(): assert isinstance(x, tvm.tirx.BufferLoad) assert x.ty == tvm.ir.PrimType("float32") assert x.buffer == buffer - assert x.buffer.data == buffer_var + assert x.buffer.data.args[0].same_as(buffer) + assert x.buffer.data.ty == tvm.tirx.buffer_data_pointer_type(buffer) assert list(x.indices) == [1] x = tvm.tirx.Ramp(1, 2, 10) @@ -228,7 +229,8 @@ def test_stmt_constructor(): x = tvm.tirx.BufferStore(buffer, tvm.tirx.IntImm("bool", 1), [10]) assert isinstance(x, tvm.tirx.BufferStore) assert x.buffer == buffer - assert x.buffer.data == buffer_var + assert x.buffer.data.args[0].same_as(buffer) + assert x.buffer.data.ty == tvm.tirx.buffer_data_pointer_type(buffer) assert list(x.indices) == [10] assert x.value.value == 1 diff --git a/tests/python/tirx-base/test_tir_nodes.py b/tests/python/tirx-base/test_tir_nodes.py index f2ecdd32d036..48b92eb4850b 100644 --- a/tests/python/tirx-base/test_tir_nodes.py +++ b/tests/python/tirx-base/test_tir_nodes.py @@ -85,7 +85,8 @@ def test_ir2(): st = tvm.tirx.BufferStore(buf, x + 1, [1]) assert isinstance(st, tvm.tirx.BufferStore) assert st.buffer == buf - assert st.buffer.data == array + assert st.buffer.data.args[0].same_as(buf) + assert st.buffer.data.ty == array.ty def test_let(): diff --git a/tests/python/tirx-base/test_tir_op_types.py b/tests/python/tirx-base/test_tir_op_types.py index 388cb0e74a1c..a3e7e0401abb 100644 --- a/tests/python/tirx-base/test_tir_op_types.py +++ b/tests/python/tirx-base/test_tir_op_types.py @@ -224,7 +224,7 @@ def test_tir_op_mma_store(): 16, buffer.access_ptr("w"), buffer_w.data, - buffer_w.elem_offset, + buffer_w.ty.elem_offset, x, ) assert expr.op.name == "tirx.mma_store" @@ -232,7 +232,7 @@ def test_tir_op_mma_store(): def test_tir_op_mma_fill(): buffer_w = tirx.decl_buffer([16, 8], dtype="int32", scope="warp", offset_factor=1) - expr = _cuda_op.mma_fill("int32", 8, buffer_w.data, buffer_w.elem_offset) + expr = _cuda_op.mma_fill("int32", 8, buffer_w.data, buffer_w.ty.elem_offset) assert expr.op.name == "tirx.mma_fill" diff --git a/tests/python/tirx-base/test_tir_specialize.py b/tests/python/tirx-base/test_tir_specialize.py index a1428348d287..efbe9babd6b6 100644 --- a/tests/python/tirx-base/test_tir_specialize.py +++ b/tests/python/tirx-base/test_tir_specialize.py @@ -199,7 +199,7 @@ def test_specialize_elemwise(): func = element_wise.specialize({a: tvm.tirx.decl_buffer((128, 64))}) assert_structural_equal_ignore_global_symbol(func, element_wise_128_64) # partially specialized - func = element_wise.specialize({c: tvm.tirx.decl_buffer((128, C.shape[1]))}) + func = element_wise.specialize({c: tvm.tirx.decl_buffer((128, C.ty.shape[1]))}) assert_structural_equal_ignore_global_symbol(func, element_wise_128_n) @@ -299,14 +299,7 @@ def expected(A: T.Buffer([16, 16], "float32"), B_handle: T.handle): def test_specialize_buffer_var_to_expr(): - """Handle specialization of buffer var - - The `tirx::Buffer::data` field must be an explicit `tirx::Var`, and - cannot be replaced with a handle-typed `tirx::Expr`. However, - these substitutions are useful - when lowering. If these occur, a binding of the `tirx::Var` is - included in the specialized function. - """ + """A DeclBuffer source expression may be specialized directly.""" @T.prim_func(private=True, s_tir=True) def before(A_data: T.handle("float32"), B_data: T.handle("float32")): @@ -318,8 +311,7 @@ def before(A_data: T.handle("float32"), B_data: T.handle("float32")): @T.prim_func(private=True, s_tir=True) def expected(A_data: T.handle("float32")): A_buf = T.decl_buffer(32, "float32", data=A_data) - B_data: T.let[T.Ptr[T.float32]] = T.address_of(A_buf[16]) - B_buf = T.decl_buffer(16, "float32", data=B_data) + B_buf = T.decl_buffer(16, "float32", data=T.address_of(A_buf[16])) for i in range(16): B_buf[i] = A_buf[i] * 2.0 diff --git a/tests/python/tirx-base/test_tir_stmt_functor.py b/tests/python/tirx-base/test_tir_stmt_functor.py index d71151028434..e3862b062b67 100644 --- a/tests/python/tirx-base/test_tir_stmt_functor.py +++ b/tests/python/tirx-base/test_tir_stmt_functor.py @@ -666,7 +666,9 @@ def func(A: T.Buffer((10,), "int32")): continue # DeclBuffer - buffer_decl = tir.DeclBuffer(T.buffer((10,), "int32"), evaluate_stmt) + buffer = T.buffer((10,), "int32") + buffer_data = tir.Var("buffer_data", buffer.data.ty) + buffer_decl = tir.DeclBuffer(buffer, evaluate_stmt, data=buffer_data) # OpCall @T.prim_func(s_tir=True) diff --git a/tests/python/tirx-transform/test_tir_inline_private_functions.py b/tests/python/tirx-transform/test_tir_inline_private_functions.py index fc6c060e9b43..903c84913606 100644 --- a/tests/python/tirx-transform/test_tir_inline_private_functions.py +++ b/tests/python/tirx-transform/test_tir_inline_private_functions.py @@ -57,10 +57,8 @@ class Expected: @T.prim_func(s_tir=True) def main(A: T.Buffer([80, 16], "float32"), B: T.Buffer([64, 16], "float32")): for i in range(64): - A_view_data: T.let[T.handle("float32")] = T.address_of(A[i, 0]) - Aview = T.decl_buffer([16, 16], "float32", data=A_view_data) - B_view_data: T.let[T.handle("float32")] = T.address_of(B[i, 0]) - Bview = T.decl_buffer([16], "float32", data=B_view_data) + Aview = T.decl_buffer([16, 16], "float32", data=T.address_of(A[i, 0])) + Bview = T.decl_buffer([16], "float32", data=T.address_of(B[i, 0])) for j in range(16): Bview[j] = 0.0 for k in range(16): @@ -229,11 +227,11 @@ def main(A: T.Buffer(16, "float32")): Before.subroutine( T.tvm_stack_make_array( A.data, - T.tvm_stack_make_shape(*A.shape, dtype="handle"), + T.tvm_stack_make_shape(*A.ty.shape, dtype="handle"), 0, - len(A.shape), + len(A.ty.shape), 0.0, - A.elem_offset, + A.ty.elem_offset, dtype="handle", ) ) diff --git a/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py b/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py index 3abf3e198006..a75b87ec566a 100644 --- a/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py +++ b/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py @@ -316,7 +316,7 @@ class After: def main( Aptr: T.handle("uint16", storage_scope="shared"), ): - A_flat_1 = T.decl_buffer(4096, "uint16", data=Aptr) + A_flat = T.decl_buffer(4096, "uint16", data=Aptr) for i in range(128): threadIdx_x = T.launch_thread("threadIdx.x", 32) @@ -333,7 +333,7 @@ def main( T.reinterpret( "float32", T.shift_left( - T.Cast("uint32", T.reinterpret("uint16", A_flat_1[0])), + T.Cast("uint32", T.reinterpret("uint16", A_flat[0])), T.uint32(16), ), ), diff --git a/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py b/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py index 8670c4424836..d01400d3a5b2 100644 --- a/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py +++ b/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py @@ -76,7 +76,7 @@ class Before: def main(A: T.Buffer((16, 16), "float32"), C: T.Buffer((16, 16), "float32")): for i in T.serial(0, 16): B_new_buf = T.alloc_buffer((1, 16), "float32") - B_new = T.Buffer([1, 16], "float32", data=B_new_buf.data) + B_new = T.decl_buffer([1, 16], "float32", data=B_new_buf.data) for j in T.serial(0, 16): B_new[0, j] = A[i, j] + 1.0 for j in T.serial(0, 16): @@ -90,16 +90,14 @@ def main(input_A: T.Buffer((16, 16), "float32"), input_C: T.Buffer((16, 16), "fl C = T.decl_buffer(256, dtype="float32", data=input_C.data) for i in T.serial(0, 16): B_new_buf = T.alloc_buffer((16,), "float32") - B_new = T.Buffer(16, "float32", data=B_new_buf.data) + B_new = T.decl_buffer(16, "float32", data=B_new_buf.data) for j in T.serial(0, 16): B_new[j] = A[((i * 16) + j)] + 1.0 for j in T.serial(0, 16): C[((i * 16) + j)] = B_new[j] * 2.0 After = _transform()(Before) - # This deliberately malformed input omits the DeclBuffer that would bind - # B_new. Compare the resulting free buffer variables by mapping them. - tvm.ir.assert_structural_equal(After, Expected, map_free_vars=True) + tvm.ir.assert_structural_equal(After, Expected) def test_gpu(): diff --git a/tests/python/tirx-transform/test_tir_transform_fp8_legalize.py b/tests/python/tirx-transform/test_tir_transform_fp8_legalize.py index cc28cee0841e..6e8721f08b65 100644 --- a/tests/python/tirx-transform/test_tir_transform_fp8_legalize.py +++ b/tests/python/tirx-transform/test_tir_transform_fp8_legalize.py @@ -216,6 +216,17 @@ def test_fp8_compute_legalize(dtype, promote_dtype): tvm.ir.assert_structural_equal(after, expected) +def test_fp8_compute_legalize_preserves_opaque_buffer_access(dtype, promote_dtype): + @T.prim_func(s_tir=True) + def before(): + buffer = T.alloc_buffer((16,), dtype) + T.evaluate(T.call_extern("void", "consume", buffer.data)) + + before_mod = tvm.IRModule.from_expr(before) + after = tvm.tirx.transform.FP8ComputeLegalize(promote_dtype)(before_mod) + tvm.ir.assert_structural_equal(after, before_mod) + + def test_fp8_storage_legalize(dtype, promote_dtype): target = Target("nvidia/nvidia-a100") before = BindTarget(target)(get_after_compute_legalize(dtype, promote_dtype)) 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 bde693fbceba..7f670d90716e 100644 --- a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py +++ b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py @@ -114,6 +114,22 @@ def collect(node): assert int(tvm.arith.Analyzer().simplify(load.indices[0])) == 5 +def test_lower_buffer_data_access_ptr_preserves_buffer_identity(): + buffer = tvm.tirx.decl_buffer((16,), "float32", "buffer") + access = tvm.tirx.tvm_access_ptr("float32", buffer.data, 3, 8, 1) + + func = tvm.tirx.PrimFunc([buffer], tvm.tirx.Evaluate(access)).with_attr( + "target", tvm.target.Target("llvm") + ) + lowered = tvm.tirx.transform.LowerIntrin()(tvm.IRModule.from_expr(func))["main"].body.value + assert isinstance(lowered, tvm.ir.Call) + assert lowered.op.name == "tirx.address_of" + load = lowered.args[0] + assert isinstance(load, tvm.tirx.BufferLoad) + assert load.buffer.same_as(buffer) + assert int(load.indices[0]) == 3 + + def get_ref_data(): """Get reference data for every pairs""" import itertools 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 d58d993d2a03..bacbe70d42f2 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 @@ -54,13 +54,11 @@ def main( T.func_attr({"target": tvm.target.Target("llvm")}) stack_ffi_any: T.let[T.handle] = T.tvm_stack_alloca("tvm_ffi_any", 4) stack_array: T.let[T.handle] = T.tvm_stack_alloca("array", 3) - stack_shape: T.let[T.handle("int64")] = T.tvm_stack_alloca("shape", 6) - stack_shape_1 = T.decl_buffer((T.int64(6),), "int64", data=stack_shape) - stack_shape_1[0] = T.int64(64) - stack_shape_1[1] = T.int64(64) + stack_shape = T.decl_buffer((T.int64(6),), "int64", data=T.tvm_stack_alloca("shape", 6)) + stack_shape[0] = T.int64(64) + stack_shape[1] = T.int64(64) T.tvm_struct_set(stack_array, 0, 1, A.data) - stack_shape_2 = T.Buffer((1,), "int64", data=stack_shape) - T.tvm_struct_set(stack_array, 0, 2, T.address_of(stack_shape_2[0])) + T.tvm_struct_set(stack_array, 0, 2, T.address_of(stack_shape[0])) T.tvm_struct_set(stack_array, 0, 3, T.reinterpret("handle", T.uint64(0))) T.tvm_struct_set(stack_array, 0, 4, 2) T.tvm_struct_set(stack_array, 0, 5, T.uint8(2)) @@ -69,11 +67,10 @@ def main( T.tvm_struct_set(stack_array, 0, 8, T.uint64(0)) T.tvm_struct_set(stack_array, 0, 9, 0) T.tvm_struct_set(stack_array, 0, 10, 1) - stack_shape_1[2] = T.int64(64) - stack_shape_1[3] = T.int64(64) + stack_shape[2] = T.int64(64) + stack_shape[3] = T.int64(64) T.tvm_struct_set(stack_array, 1, 1, B.data) - stack_shape_3 = T.Buffer((3,), "int64", data=stack_shape) - T.tvm_struct_set(stack_array, 1, 2, T.address_of(stack_shape_3[2])) + T.tvm_struct_set(stack_array, 1, 2, T.address_of(stack_shape[2])) T.tvm_struct_set(stack_array, 1, 3, T.reinterpret("handle", T.uint64(0))) T.tvm_struct_set(stack_array, 1, 4, 2) T.tvm_struct_set(stack_array, 1, 5, T.uint8(2)) @@ -82,11 +79,10 @@ def main( T.tvm_struct_set(stack_array, 1, 8, T.uint64(0)) T.tvm_struct_set(stack_array, 1, 9, 0) T.tvm_struct_set(stack_array, 1, 10, 1) - stack_shape_1[4] = T.int64(64) - stack_shape_1[5] = T.int64(64) + stack_shape[4] = T.int64(64) + stack_shape[5] = T.int64(64) T.tvm_struct_set(stack_array, 2, 1, C.data) - stack_shape_4 = T.Buffer((5,), "int64", data=stack_shape) - T.tvm_struct_set(stack_array, 2, 2, T.address_of(stack_shape_4[4])) + T.tvm_struct_set(stack_array, 2, 2, T.address_of(stack_shape[4])) T.tvm_struct_set(stack_array, 2, 3, T.reinterpret("handle", T.uint64(0))) T.tvm_struct_set(stack_array, 2, 4, 2) T.tvm_struct_set(stack_array, 2, 5, T.uint8(2)) @@ -187,9 +183,9 @@ def variance4(rxplaceholder: T.Buffer((T.int64(1), T.int64(32), T.int64(25690112 T.func_attr({"global_symbol": "variance4", "tirx.noalias": True}) rxplaceholder_red = T.alloc_buffer((32,), "float32") T_subtract = T.alloc_buffer((822083584,), "float32") - rxplaceholder_red_1 = T.Buffer((T.int64(32),), data=rxplaceholder_red.data) - rxplaceholder_1 = T.Buffer((T.int64(822083584),), data=rxplaceholder.data) - T_subtract_1 = T.Buffer((T.int64(822083584),), data=T_subtract.data) + rxplaceholder_red_1 = T.decl_buffer((T.int64(32),), data=rxplaceholder_red.data) + rxplaceholder_1 = T.decl_buffer((T.int64(822083584),), data=rxplaceholder.data) + T_subtract_1 = T.decl_buffer((T.int64(822083584),), data=T_subtract.data) for ax1, ax2 in T.grid(32, 25690112): cse_v1: T.let[T.int32] = ax1 * 25690112 + ax2 T_subtract_1[cse_v1] = rxplaceholder_1[cse_v1] - rxplaceholder_red_1[ax1] diff --git a/tests/python/tirx-transform/test_tir_transform_make_packed_api.py b/tests/python/tirx-transform/test_tir_transform_make_packed_api.py index 4b69f933358a..81d46061ff9d 100644 --- a/tests/python/tirx-transform/test_tir_transform_make_packed_api.py +++ b/tests/python/tirx-transform/test_tir_transform_make_packed_api.py @@ -462,5 +462,30 @@ def main(a: T.handle, b: T.handle): assert len(After["main"].params) == 4 +def test_buffer_alignment_attached_to_buffer_var(): + """Packed ABI alignment metadata remains keyed by the logical BufferVar.""" + + @I.ir_module + class Before: + @T.prim_func(s_tir=True) + def main(A: T.Buffer((16,), "float32", align=64)): + T.func_attr({"global_symbol": "main", "target": T.target("llvm", host="llvm")}) + T.evaluate(A[0]) + + after = tvm.tirx.transform.MakePackedAPI()(Before)["main"] + alignment_nodes = [] + declared_buffers = [] + + def collect(node): + if isinstance(node, tirx.AttrStmt) and node.attr_key == "storage_alignment": + alignment_nodes.append(node.node) + if isinstance(node, tirx.DeclBuffer): + declared_buffers.append(node.buffer) + + tirx.stmt_functor.post_order_visit(after.body, collect) + assert len(alignment_nodes) == 1 + assert any(alignment_nodes[0].same_as(buffer) for buffer in declared_buffers) + + if __name__ == "__main__": tvm.testing.main() 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 98a903a09cc8..c11fdf402444 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 @@ -142,5 +142,40 @@ def main(A: T.Buffer((16,), "float32")): tvm.ir.assert_structural_equal(After, Expected) +def test_decl_buffer_alias_chain_uses_flat_root_map(): + transform = tvm.tirx.transform.PointerValueTypeRewrite() + + @I.ir_module + class Before: + @T.prim_func(s_tir=True) + def main(A: T.Buffer((16,), "float32")): + A_view = T.decl_buffer((16,), "float32", data=A.data) + A_view_2 = T.decl_buffer((16,), "float32", data=A_view.data) + for i in range(4): + A_view_2[i * 4 : i * 4 + 4] = T.broadcast(T.float32(1), 4) + + After = transform(Before) + assert tvm.tirx.analysis.verify_well_formed(After) + func = After["main"] + assert next(iter(func.buffer_map.values())).ty.dtype == tvm.ir.PrimType("float32x4") + + decl_buffers = [] + buffer_stores = [] + tvm.tirx.stmt_functor.post_order_visit( + func.body, + lambda node: ( + decl_buffers.append(node) + if isinstance(node, tvm.tirx.DeclBuffer) + else buffer_stores.append(node) + if isinstance(node, tvm.tirx.BufferStore) + else None + ), + ) + assert len(decl_buffers) == 2 + assert all(decl.buffer.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") + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx-transform/test_tir_transform_split_host_device.py b/tests/python/tirx-transform/test_tir_transform_split_host_device.py index 8a026c7a8ab5..69bcf5a95f0d 100644 --- a/tests/python/tirx-transform/test_tir_transform_split_host_device.py +++ b/tests/python/tirx-transform/test_tir_transform_split_host_device.py @@ -345,6 +345,28 @@ def main(var_A: T.handle, var_B: T.handle): assert isinstance(after["main_kernel"].params[2], tvm.tirx.Var) +def test_buffer_used_only_through_data_projection(): + @I.ir_module + class Before: + @T.prim_func(s_tir=True) + def main(A: T.Buffer((16,), "float32")): + T.func_attr({"target": T.target("cuda", host="llvm")}) + with T.attr(T.target("cuda"), "target", 0): + T.evaluate(T.call_extern("consume", A.data, dtype="int32")) + + after = tvm.tirx.transform.SplitHostDevice()(Before) + kernel = after["main_kernel"] + declared_buffers = [] + + def collect(node): + if isinstance(node, tvm.tirx.DeclBuffer): + declared_buffers.append(node.buffer) + + tvm.tirx.stmt_functor.post_order_visit(kernel.body, collect) + assert len(declared_buffers) == 1 + assert not tvm.tirx.analysis.undefined_vars(kernel.body, kernel.params) + + def test_thread_extent_region_extracted_as_device_kernel(): """A bare thread_extent is annotated and extracted as a device kernel.""" diff --git a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py index 3f3cc59dad09..47473737ccb9 100644 --- a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py +++ b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py @@ -46,7 +46,7 @@ def func(n: T.int32): def verify(n): if isinstance(n, tvm.tirx.AllocBuffer): num_alloc[0] += 1 - assert n.buffer.shape[0].value == 200 + assert n.buffer.ty.shape[0].value == 200 tvm.tirx.stmt_functor.post_order_visit(body, verify) assert num_alloc[0] == 1 @@ -100,7 +100,7 @@ def offset_generater(dtype_list, length): def dtype_test(dtype_list, length): def verify(n): if isinstance(n, tvm.tirx.AllocBuffer): - assert n.buffer.shape[0].value == offset + assert n.buffer.ty.shape[0].value == offset mod = make_mod(dtype_list, length) offset = offset_generater(dtype_list, length) @@ -157,7 +157,7 @@ def before(A: T.Buffer(8, "float32"), E: T.Buffer(8, "float32")): def verify(n): if isinstance(n, tvm.tirx.AllocBuffer): - total_alloc[0] += n.buffer.shape[0].value + total_alloc[0] += n.buffer.ty.shape[0].value total_alloc = [0] mod = tvm.IRModule.from_expr(before.with_attr("global_symbol", "main")) @@ -280,7 +280,7 @@ def func(n: T.int32): def verify(n): if isinstance(n, tvm.tirx.AllocBuffer): num_alloc[0] += 1 - assert n.buffer.shape[0].value == 500 + assert n.buffer.ty.shape[0].value == 500 tvm.tirx.stmt_functor.post_order_visit(body, verify) assert num_alloc[0] == 1 @@ -310,7 +310,7 @@ def func(n: T.int32): def verify(n): if isinstance(n, tvm.tirx.AllocBuffer): num_alloc[0] += 1 - assert n.buffer.shape[0].value == 200 + assert n.buffer.ty.shape[0].value == 200 tvm.tirx.stmt_functor.post_order_visit(body, verify) assert num_alloc[0] == 1 @@ -342,7 +342,7 @@ def func(n: T.int32): def verify(n): if isinstance(n, tvm.tirx.AllocBuffer): num_alloc[0] += 1 - assert n.buffer.shape[0].value == 800 + assert n.buffer.ty.shape[0].value == 800 tvm.tirx.stmt_functor.post_order_visit(body, verify) assert num_alloc[0] == 1 @@ -371,15 +371,11 @@ def func_rewritten(A: T.Buffer((8,), "float32")) -> None: tvm.ir.assert_structural_equal(mod["main"], func_rewritten.with_attr("global_symbol", "main")) -def test_let_buffer_rewrite(): - """StorageRewrite replaces the bound var of backing allocations +def test_decl_buffer_is_not_vectorized(): + """StorageRewrite leaves explicit DeclBuffer views unchanged. - If StorageRewrite replaces the backing variable of an array, such - as when vectorizing the storage type, the variable must be - replaced in the Bind that defines it. Currently, StmtMutator - only visits usage of variables, and does not visit definitions of - variables, so the definition in a Bind must be explicitly - handled. + Vectorization of DeclBuffer views was dropped because the rewritten result + violates the immutable BufferVar binding invariants. """ @I.ir_module @@ -392,18 +388,8 @@ def main() -> None: A = T.decl_buffer([8], "int32", data=A_data) A[0:8] = T.broadcast(42, 8) - @I.ir_module(check_well_formed=False) - class Expected: - @T.prim_func(s_tir=True) - def main() -> None: - A_data: T.let[T.handle("int32x8")] = T.call_extern( - "dummy_func", dtype=T.handle("int32x8").ty - ) - A = T.decl_buffer([1], "int32x8", data=A_data) - A[0] = T.broadcast(42, 8) - After = tvm.tirx.transform.StorageRewrite()(Before) - tvm.ir.assert_structural_equal(After, Expected) + tvm.ir.assert_structural_equal(After, Before) def test_rewrite_decl_buffer(): @@ -442,9 +428,60 @@ def main(A: T.Buffer(16, "float32"), D: T.Buffer(16, "float32")): D[i] = C[i] After = tvm.tirx.transform.StorageRewrite()(Before) + assert tvm.tirx.analysis.verify_well_formed(After) tvm.ir.assert_structural_equal(After, Expected) +def test_decl_buffer_alias_chain_uses_flat_root(): + """StorageRewrite resolves every alias in a chain to the same flat root.""" + + @I.ir_module + class Before: + @T.prim_func(s_tir=True) + def main(D: T.Buffer(1, "float32")): + A = T.decl_buffer(16, dtype="float32") + B = T.decl_buffer(16, dtype="float32", data=A.data) + C = T.decl_buffer(16, dtype="float32", data=B.data) + A[0] = 1.0 + D[0] = C[0] + + @I.ir_module + class Expected: + @T.prim_func(s_tir=True) + def main(D: T.Buffer(1, "float32")): + A = T.decl_buffer(16, dtype="float32") + B = T.decl_buffer(16, dtype="float32", data=A.data) + C = T.decl_buffer(16, dtype="float32", data=A.data) + A[0] = 1.0 + D[0] = C[0] + + After = tvm.tirx.transform.StorageRewrite()(Before) + assert tvm.tirx.analysis.verify_well_formed(After) + tvm.ir.assert_structural_equal(After, Expected) + + +def test_decl_buffer_alias_extends_source_lifetime(): + """An access through an alias prevents reuse of its source allocation.""" + + @T.prim_func(s_tir=True) + def func(D: T.Buffer(1, "float32")): + A = T.decl_buffer(16, dtype="float32") + B = T.decl_buffer(16, dtype="float32", data=A.data) + A[0] = 1.0 + + C = T.decl_buffer(16, dtype="float32") + C[0] = 2.0 + D[0] = B[0] + C[0] + + after = tvm.tirx.transform.StorageRewrite()(tvm.IRModule.from_expr(func))["func"] + allocations = [] + tvm.tirx.stmt_functor.post_order_visit( + after.body, + lambda node: allocations.append(node) if isinstance(node, tvm.tirx.AllocBuffer) else None, + ) + assert len(allocations) == 2 + + def test_no_orphaned_decl_buffer(): """A DeclBuffer of an unused Allocate should be removed diff --git a/tests/python/tirx/codegen/test_codegen_hopper.py b/tests/python/tirx/codegen/test_codegen_hopper.py index 59b93ecf2a32..e3405ce3fe28 100644 --- a/tests/python/tirx/codegen/test_codegen_hopper.py +++ b/tests/python/tirx/codegen/test_codegen_hopper.py @@ -24,7 +24,7 @@ import tvm.testing from tvm.script import tirx as T from tvm.testing import env -from tvm.tirx import Buffer +from tvm.tirx import BufferAccessKind def _get_source(func: tvm.tirx.PrimFunc) -> tuple[str, tvm.IRModule]: @@ -775,16 +775,16 @@ def main(A_ptr: T.handle, B_ptr: T.handle): T.ptx.fence.proxy_async("shared::cta") T.ptx.mbarrier.arrive.expect_tx(T.address_of(bar), total_bytes) if clusterCtaIdx == 0: - T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.access_ptr(Buffer.WRITE, offset=A_smem.elem_offset_of(coord0[::-1])), # noqa: E501 + T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.access_ptr(BufferAccessKind.WRITE, offset=A_smem.elem_offset_of(coord0[::-1])), # noqa: E501 T.address_of(bar), T.address_of(A_map), int("1111", 2), 1, "", *coord0) # noqa: E501 if clusterCtaIdx == 1: - T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.access_ptr(Buffer.WRITE, offset=A_smem.elem_offset_of(coord1[::-1])), # noqa: E501 + T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.access_ptr(BufferAccessKind.WRITE, offset=A_smem.elem_offset_of(coord1[::-1])), # noqa: E501 T.address_of(bar), T.address_of(A_map), int("1111", 2), 1, "", *coord1) # noqa: E501 if clusterCtaIdx == 2: - T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.access_ptr(Buffer.WRITE, offset=A_smem.elem_offset_of(coord2[::-1])), # noqa: E501 + T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.access_ptr(BufferAccessKind.WRITE, offset=A_smem.elem_offset_of(coord2[::-1])), # noqa: E501 T.address_of(bar), T.address_of(A_map), int("1111", 2), 1, "", *coord2) # noqa: E501 if clusterCtaIdx == 3: - T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.access_ptr(Buffer.WRITE, offset=A_smem.elem_offset_of(coord3[::-1])), # noqa: E501 + T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.access_ptr(BufferAccessKind.WRITE, offset=A_smem.elem_offset_of(coord3[::-1])), # noqa: E501 T.address_of(bar), T.address_of(A_map), int("1111", 2), 1, "", *coord3) # noqa: E501 # wait for the copy to finish T.ptx.mbarrier.try_wait(T.address_of(bar), phase) diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py index 7cbce484ba45..6cf9de6381cb 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py @@ -266,7 +266,7 @@ def _build_expected_impl(direction, dtype, s_shape, s_layout, impl_spec): # Wrap: DeclBuffer -> nested For loops (skipped when total extent is 1, # matching the implementation's always-unroll single-loop emission). - body = DeclBuffer(s_buf, eval_stmt) + body = DeclBuffer(s_buf, eval_stmt, data=s_buf_ptr) for i in range(n_loops - 1, -1, -1): body = tvm.tirx.For( loop_vars[i], diff --git a/tests/python/tirx/operator/tile_primitive/trn/test_copy_trn.py b/tests/python/tirx/operator/tile_primitive/trn/test_copy_trn.py index dab83844bb31..308048a081a0 100644 --- a/tests/python/tirx/operator/tile_primitive/trn/test_copy_trn.py +++ b/tests/python/tirx/operator/tile_primitive/trn/test_copy_trn.py @@ -172,7 +172,7 @@ def expected(A_ptr: T.handle): _A_flat = T.decl_buffer((262144,), data=A.data, layout=None) A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf") A_sbuf_view = T.decl_buffer((128, 2048), data=A_sbuf.data, scope="trn.sbuf", layout=None) - A_view = T.decl_buffer((262144,), data=A.data, layout=None) + A_view = T.decl_buffer((262144,), data=_A_flat.data, layout=None) for i, b_loop in T.grid(4, 1): T.attr(0, "tensorized_nki_instruction", 1) for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}): @@ -337,13 +337,13 @@ def expected(): T.func_attr({"global_symbol": "copy"}) A_sbuf = T.alloc_buffer((128, 256), scope="trn.sbuf") B_sbuf = T.alloc_buffer((128, 16), scope="trn.sbuf") - _B_sbuf_view = T.decl_buffer((128, 16), data=B_sbuf.data, scope="trn.sbuf", layout=None) + B_sbuf_view = T.decl_buffer((128, 16), data=B_sbuf.data, scope="trn.sbuf", layout=None) for b_loop in T.serial(0, 4): T.attr(0, "tensorized_nki_instruction", 1) for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}): for f_loop in T.serial(0, 4, annotations={"nki_dim": "F"}): T.nki.tensor_copy( - B_sbuf[p_loop, b_loop * 4 + f_loop], + B_sbuf_view[p_loop, b_loop * 4 + f_loop], A_sbuf[p_loop, b_loop * 64 + f_loop], ) diff --git a/tests/python/tirx/operator/tile_primitive/trn/test_unary_trn.py b/tests/python/tirx/operator/tile_primitive/trn/test_unary_trn.py index a774b4c9e447..a6557c346a1a 100644 --- a/tests/python/tirx/operator/tile_primitive/trn/test_unary_trn.py +++ b/tests/python/tirx/operator/tile_primitive/trn/test_unary_trn.py @@ -130,7 +130,7 @@ def expected(): if op_type == "reciprocal": T.nki.reciprocal(B_sbuf_view[p_loop, i * 512 + f_loop], A_sbuf_view[p_loop, i * 1024 + f_loop]) # noqa: E501 elif op_type == "memset": - T.nki.memset(B_sbuf[p_loop, i * 512 + f_loop], 0.0) + T.nki.memset(B_sbuf_view[p_loop, i * 512 + f_loop], 0.0) # fmt: on with target: mod = tvm.IRModule({"main": unary}) diff --git a/tests/python/tirx/test_parser_printer.py b/tests/python/tirx/test_parser_printer.py index 02fc42b5cad5..4b4ea96437d0 100644 --- a/tests/python/tirx/test_parser_printer.py +++ b/tests/python/tirx/test_parser_printer.py @@ -1374,14 +1374,14 @@ def func() -> None: # Shared data pointer assert_structural_equal(_collect_buffer_sources(func)["B_local"], b_buf.data) # Shape: single dim matching storage shard total - assert len(b_local.shape) == 1 - storage = b_buf.layout.storage() + assert len(b_local.ty.shape) == 1 + storage = b_buf.ty.layout.storage() expected_total = 1 for it in storage.shard: expected_total *= int(it.extent) - assert int(b_local.shape[0]) == expected_total + assert int(b_local.ty.shape[0]) == expected_total # Layout: storage layout (parent layout with thread axes removed) - assert_structural_equal(b_local.layout, storage) + assert_structural_equal(b_local.ty.layout, storage) # Round-trip code = func.script() @@ -1408,10 +1408,10 @@ def func() -> None: # Shared data pointer assert_structural_equal(_collect_buffer_sources(func)["B"], a_buf.data) # Shape: [4, 8] from [8, 4] - assert int(b_buf.shape[0]) == 4 - assert int(b_buf.shape[1]) == 8 + assert int(b_buf.ty.shape[0]) == 4 + assert int(b_buf.ty.shape[1]) == 8 # Layout: permuted - assert_structural_equal(b_buf.layout, a_buf.layout.permute_dims([1, 0])) + assert_structural_equal(b_buf.ty.layout, a_buf.ty.layout.permute_dims([1, 0])) code = func.script() assert from_source(code).script() == code @@ -1436,10 +1436,10 @@ def func() -> None: # Shared data pointer assert_structural_equal(_collect_buffer_sources(func)["B"], a_buf.data) # dtype - assert str(b_buf.dtype) == "float32" + assert str(b_buf.ty.dtype) == "float32" # Shape: [8, 4] (last dim halved since float32 is 2x float16) - assert int(b_buf.shape[0]) == 8 - assert int(b_buf.shape[1]) == 4 + assert int(b_buf.ty.shape[0]) == 8 + assert int(b_buf.ty.shape[1]) == 4 code = func.script() assert from_source(code).script() == code @@ -1811,6 +1811,14 @@ def func(): code = func.script() assert from_source(code).script() == code assert_structural_equal(func, from_source(code)) + decls = [] + tvm.tirx.stmt_functor.post_order_visit( + func.body, + lambda node: decls.append(node) if isinstance(node, tvm.tirx.DeclBuffer) else None, + ) + assert len(decls) == 1 + tmem_decl = next(decl for decl in decls if decl.buffer.scope() == "tmem") + assert tmem_decl.data.op.name == "tirx.reinterpret" def test_roundtrip_cuda_func_call_source_code(): diff --git a/tests/python/tirx/transform/test_stmt_functor.py b/tests/python/tirx/transform/test_stmt_functor.py index 393cfa47a0ef..764431e1ae13 100644 --- a/tests/python/tirx/transform/test_stmt_functor.py +++ b/tests/python/tirx/transform/test_stmt_functor.py @@ -681,7 +681,9 @@ def func(A: T.Buffer((10,), "int32")): continue # DeclBuffer - buffer_decl = tir.DeclBuffer(T.buffer((10,), "int32"), evaluate_stmt) + buffer = T.buffer((10,), "int32") + buffer_data = tir.Var("buffer_data", buffer.data.ty) + buffer_decl = tir.DeclBuffer(buffer, evaluate_stmt, data=buffer_data) # TilePrimitiveCall — extract the TilePrimitiveCall from the kernel body, then wrap in an SBlock @T.prim_func diff --git a/tests/python/tvmscript/test_tvmscript_parser_tir.py b/tests/python/tvmscript/test_tvmscript_parser_tir.py index 157e59f59e92..04d111d25054 100644 --- a/tests/python/tvmscript/test_tvmscript_parser_tir.py +++ b/tests/python/tvmscript/test_tvmscript_parser_tir.py @@ -27,14 +27,14 @@ def test_tir_buffer_proxy(): buffer_0 = T.Buffer((128, 128), "float32") assert ( - isinstance(buffer_0, tirx.Buffer) + tirx.is_buffer_var(buffer_0) and list(buffer_0.shape) == [128, 128] and buffer_0.dtype == ir.PrimType("float32") ) buffer_1 = T.Buffer((64, 64, 64), "int32") assert ( - isinstance(buffer_1, tirx.Buffer) + tirx.is_buffer_var(buffer_1) and list(buffer_1.shape) == [64, 64, 64] and buffer_1.dtype == ir.PrimType("int32") )