This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tvm-buffer-type-var-refactor in repository https://gitbox.apache.org/repos/asf/tvm.git
commit db0029c1a63cbac9a90bd19757dd42133e18a181 Author: Tianqi Chen <[email protected]> AuthorDate: Tue Jul 28 23:12:06 2026 +0000 [FIX][TIRX] Repair typed-buffer CI regressions --- include/tvm/tirx/expr.h | 2 +- include/tvm/tirx/stmt.h | 4 ++-- .../tvm/s_tir/dlight/analysis/common_analysis.py | 2 +- python/tvm/s_tir/dlight/cpu/reduction.py | 4 +--- python/tvm/tirx/buffer.py | 4 ++-- python/tvm/tirx/op.py | 2 +- src/backend/trn/codegen/codegen_trn.cc | 26 ++++++++++++++++++++-- src/backend/trn/codegen/codegen_trn.h | 2 ++ src/relax/script/printer/expr.cc | 3 ++- src/relax/transform/fuse_tir.cc | 2 +- .../analysis/sblock_access_region_detector.cc | 7 ++++++ src/s_tir/transform/tensorcore_infer_fragment.cc | 23 ++++++++++++++----- src/target/source/codegen_c.cc | 4 +++- src/te/operation/create_primfunc.cc | 1 + src/tirx/ir/data_type_rewriter.cc | 7 ++++-- src/tirx/ir/specialize.cc | 17 +++++++++++++- src/tirx/ir/stmt.cc | 3 +++ src/tirx/script/printer/expr.cc | 2 +- src/tirx/script/printer/stmt.cc | 3 ++- src/tirx/transform/flatten_buffer.cc | 19 +++++++++++++--- tests/python/relax/test_blockbuilder_core.py | 2 +- .../s_tir/analysis/test_sblock_access_region.py | 22 ++++++++++++++++++ .../test_tvmscript_printer_structural_equal.py | 18 +++++++++++++-- .../python/tvmscript/test_tvmscript_printer_tir.py | 6 ++--- 24 files changed, 149 insertions(+), 36 deletions(-) diff --git a/include/tvm/tirx/expr.h b/include/tvm/tirx/expr.h index d4a63b62ee..96500d475f 100644 --- a/include/tvm/tirx/expr.h +++ b/include/tvm/tirx/expr.h @@ -553,7 +553,7 @@ class BufferLoadNode : public ExprNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<BufferLoadNode>() - .def_ro("buffer", &BufferLoadNode::buffer) + .def_ro("buffer", &BufferLoadNode::buffer, refl::AttachFieldFlag::SEqHashDefRecursive()) .def_ro("indices", &BufferLoadNode::indices) .def_ro("predicate", &BufferLoadNode::predicate); } diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h index 9d64ac2e82..d95c1af501 100644 --- a/include/tvm/tirx/stmt.h +++ b/include/tvm/tirx/stmt.h @@ -213,7 +213,7 @@ class BufferStoreNode : public StmtNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<BufferStoreNode>() - .def_ro("buffer", &BufferStoreNode::buffer) + .def_ro("buffer", &BufferStoreNode::buffer, refl::AttachFieldFlag::SEqHashDefRecursive()) .def_ro("value", &BufferStoreNode::value) .def_ro("indices", &BufferStoreNode::indices) .def_ro("predicate", &BufferStoreNode::predicate); @@ -783,7 +783,7 @@ class BufferRegionNode : public PrimExprConvertibleNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<BufferRegionNode>() - .def_ro("buffer", &BufferRegionNode::buffer) + .def_ro("buffer", &BufferRegionNode::buffer, refl::AttachFieldFlag::SEqHashDefRecursive()) .def_ro("region", &BufferRegionNode::region); } diff --git a/python/tvm/s_tir/dlight/analysis/common_analysis.py b/python/tvm/s_tir/dlight/analysis/common_analysis.py index 17061895e3..9c29d49d79 100644 --- a/python/tvm/s_tir/dlight/analysis/common_analysis.py +++ b/python/tvm/s_tir/dlight/analysis/common_analysis.py @@ -158,7 +158,7 @@ class BufferInfo: ) vbuf_extent = int(self.shape[-1]) & ~(int(self.shape[-1]) - 1) - return min(vlp_extent, vbuf_extent, vbits // self.buf_region.buffer.dtype.dtype.bits) + return min(vlp_extent, vbuf_extent, vbits // self.buf_region.buffer.dtype.bits) def __str__(self) -> str: return f"BufferInfo({self.buf_region})" diff --git a/python/tvm/s_tir/dlight/cpu/reduction.py b/python/tvm/s_tir/dlight/cpu/reduction.py index cf02a1f1cb..bbce215098 100644 --- a/python/tvm/s_tir/dlight/cpu/reduction.py +++ b/python/tvm/s_tir/dlight/cpu/reduction.py @@ -85,9 +85,7 @@ class Reduction(CPUScheduleRule): # Infer dtype from the last block's write buffer. last_block_stmt = sch.get(block_infos[-1].block_rv) - dtype_bits = ( - last_block_stmt.writes[0].buffer.dtype.dtype.bits if last_block_stmt.writes else 32 - ) + dtype_bits = last_block_stmt.writes[0].buffer.dtype.bits if last_block_stmt.writes else 32 # Determine vector lanes from target VLEN. vlen_bits = llvm_get_vector_width(target) diff --git a/python/tvm/tirx/buffer.py b/python/tvm/tirx/buffer.py index 625926581b..ef9cee6ee0 100644 --- a/python/tvm/tirx/buffer.py +++ b/python/tvm/tirx/buffer.py @@ -570,8 +570,8 @@ def decl_buffer( elem_offset = Var(f"{name}_elem_offset", shape_ty) 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, tvm.ir.Expr) or not isinstance(data.ty, PointerType): + raise TypeError("Buffer data must be an Expr 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 diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py index 46b05c0bc6..397b44a10d 100644 --- a/python/tvm/tirx/op.py +++ b/python/tvm/tirx/op.py @@ -1270,7 +1270,7 @@ def trace(args, trace_action="tvm.default_trace_action"): 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.ty.dtype + ret_ty = tracing_value.ty if isinstance(tracing_value, Expr) else tracing_value.dtype return tvm.ir.Call(Op.get("tirx.tvm_call_trace_packed"), call_args, ret_ty=ret_ty) diff --git a/src/backend/trn/codegen/codegen_trn.cc b/src/backend/trn/codegen/codegen_trn.cc index fefccbf28c..fbad6be2a3 100644 --- a/src/backend/trn/codegen/codegen_trn.cc +++ b/src/backend/trn/codegen/codegen_trn.cc @@ -90,6 +90,7 @@ void CodeGenTrainium::AddFunction(const GlobalVar& gvar, const PrimFunc& func) { // clear previous generated state. this->InitFuncState(func); buffer_idmap_.clear(); + buffer_data_varmap_.clear(); data_buffer_idmap_.clear(); data_decl_buffer_map_.clear(); // skip the first underscore, so SSA variable starts from _1 @@ -114,6 +115,9 @@ void CodeGenTrainium::AddFunction(const GlobalVar& gvar, const PrimFunc& func) { LOG(FATAL) << "Trainium codegen currently only support buffer arguments"; }; std::string vid = AllocVarID(v.get()); + if (auto buffer = func->buffer_map.Get(v)) { + var_idmap_[buffer.value().get()] = vid; + } if (i >= static_cast<size_t>(num_inputs.value())) { this->stream << vid << ": nt.mutable_tensor, "; output_vids.push_back(vid); @@ -209,7 +213,7 @@ 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.get()); + std::string vid = AllocVarID(op->buffer.get(), op->buffer.name() + "_ptr"); this->PrintIndent(); auto scope = op->buffer.scope(); @@ -607,7 +611,24 @@ void CodeGenTrainium::VisitStmt_(const DeclBufferNode* op) { if (op->buffer.scope() == "trn.psum" || op->buffer.scope() == "trn.sbuf") { return; } - const VarNode* data = op->buffer.get(); + const VarNode* data = op->data.as<VarNode>(); + if (const auto* call = op->data.as<CallNode>(); + call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) { + data = call->args[0].as<VarNode>(); + } + TVM_FFI_ICHECK(data) << "Trainium codegen expects DeclBuffer data to be a buffer variable"; + if (data->ty.as<PointerTypeNode>()) { + buffer_idmap_[op->buffer] = GetVarID(data); + buffer_data_varmap_[op->buffer] = data; + return; + } + TVM_FFI_ICHECK(data->ty.as<BufferTypeNode>()); + BufferVar source_buffer(ffi::GetRef<Var>(data)); + auto source_it = buffer_data_varmap_.find(source_buffer); + TVM_FFI_ICHECK(source_it != buffer_data_varmap_.end()) + << "Trainium codegen expects the source buffer to be declared before its alias"; + data = source_it->second; + auto it = data_buffer_idmap_.find(data); if (it != data_buffer_idmap_.end()) { const BufferVar& prev_buffer = data_decl_buffer_map_.at(data); @@ -620,6 +641,7 @@ void CodeGenTrainium::VisitStmt_(const DeclBufferNode* op) { std::string data_vid = GetVarID(data); std::string buffer_vid = name_supply_->FreshName(data_vid + "_buffer"); buffer_idmap_[op->buffer] = buffer_vid; + buffer_data_varmap_[op->buffer] = data; data_buffer_idmap_[data] = buffer_vid; data_decl_buffer_map_[data] = op->buffer; PrintIndent(); diff --git a/src/backend/trn/codegen/codegen_trn.h b/src/backend/trn/codegen/codegen_trn.h index 63d6712926..cb05ff988b 100644 --- a/src/backend/trn/codegen/codegen_trn.h +++ b/src/backend/trn/codegen/codegen_trn.h @@ -81,6 +81,8 @@ class CodeGenTrainium final : public CodeGenC { NKIInstructionCtx ctx_; std::unordered_map<std::string, std::string> opcode_map_; std::unordered_map<BufferVar, std::string, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> buffer_idmap_; + std::unordered_map<BufferVar, const VarNode*, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> + buffer_data_varmap_; std::unordered_map<const VarNode*, std::string> data_buffer_idmap_; std::unordered_map<const VarNode*, BufferVar> data_decl_buffer_map_; bool is_outermost_loop_ = true; diff --git a/src/relax/script/printer/expr.cc b/src/relax/script/printer/expr.cc index 3a633259e6..34ce7b0886 100644 --- a/src/relax/script/printer/expr.cc +++ b/src/relax/script/printer/expr.cc @@ -158,7 +158,8 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable).set_dispatch<relax::DataflowVar>("rel std::string ReprPrintVar(const ffi::ObjectRef& obj, const PrinterConfig& cfg) { Var var = obj.as_or_throw<Var>(); - if (var->ty.as<PrimTypeNode>() || var->ty.as<PointerTypeNode>()) { + if (var->ty.as<PrimTypeNode>() || var->ty.as<PointerTypeNode>() || + var->ty.as<tirx::BufferTypeNode>()) { return ReprPrintTIR(obj, cfg); } return ReprPrintRelax(obj, cfg); diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc index 28f7178a3e..a7145144b8 100644 --- a/src/relax/transform/fuse_tir.cc +++ b/src/relax/transform/fuse_tir.cc @@ -497,7 +497,7 @@ class RelaxToTIRVarMapCollector : public ExprVisitor { // structurally equal to the `new_buf` passed 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)) + TVM_FFI_ICHECK(ffi::StructuralEqual()((*it).second.type(), new_buf.type())) << "Inconsistent buffers " << (*it).second << " and " << new_buf << " mapped to the same relax var: " << expr; } diff --git a/src/s_tir/analysis/sblock_access_region_detector.cc b/src/s_tir/analysis/sblock_access_region_detector.cc index 4eb3f32627..3feb2c42ea 100644 --- a/src/s_tir/analysis/sblock_access_region_detector.cc +++ b/src/s_tir/analysis/sblock_access_region_detector.cc @@ -122,6 +122,7 @@ class BlockReadWriteDetector : public StmtExprVisitor { void VisitStmt_(const ForNode* op) override; void VisitStmt_(const IfThenElseNode* op) override; void VisitStmt_(const SBlockRealizeNode* op) override; + void VisitStmt_(const DeclBufferNode* op) override; void VisitStmt_(const BufferStoreNode* op) override; void VisitStmt_(const BindNode* op) override; void VisitExpr_(const BufferLoadNode* op) override; @@ -195,6 +196,12 @@ void BlockReadWriteDetector::VisitStmt_(const IfThenElseNode* op) { } } +void BlockReadWriteDetector::VisitStmt_(const DeclBufferNode* op) { + // A DeclBuffer data expression defines the alias source. It is not an + // opaque buffer access by the containing block. + VisitBufferDef(op->buffer, /*alloc_data=*/false); +} + void BlockReadWriteDetector::VisitStmt_(const BindNode* op) { if (auto value = op->value.as<PrimExpr>()) { let_bindings_[op->var.get()] = value.value(); diff --git a/src/s_tir/transform/tensorcore_infer_fragment.cc b/src/s_tir/transform/tensorcore_infer_fragment.cc index a97d5a1668..402848ce5b 100644 --- a/src/s_tir/transform/tensorcore_infer_fragment.cc +++ b/src/s_tir/transform/tensorcore_infer_fragment.cc @@ -41,6 +41,17 @@ namespace tvm { namespace s_tir { using namespace tvm::tirx; +const VarNode* GetBufferVarFromData(const Expr& data) { + if (const auto* var = data.as<VarNode>()) { + return var; + } + if (const auto* call = data.as<CallNode>(); + call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) { + return call->args[0].as<VarNode>(); + } + return nullptr; +} + // Get fragment information from tensor intrinsics class FragmentGetter : public StmtExprVisitor { public: @@ -53,7 +64,7 @@ class FragmentGetter : public StmtExprVisitor { if (op->op.same_as(tvm_load_matrix_sync_op) || op->op.same_as(tvm_store_matrix_sync_op)) { // Get shape and layout information from load and store intrinsic TVM_FFI_ICHECK_EQ(op->args.size(), 8U); - const VarNode* buffer_var = op->args[0].as<VarNode>(); + const VarNode* buffer_var = GetBufferVarFromData(op->args[0]); TVM_FFI_ICHECK(buffer_var); // Get shape const IntImmNode* m = op->args[1].as<IntImmNode>(); @@ -88,7 +99,7 @@ class FragmentGetter : public StmtExprVisitor { } else if (op->op.same_as(tvm_fill_fragment_op)) { // Get shape information from fill intrinsic TVM_FFI_ICHECK_EQ(op->args.size(), 6U); - const VarNode* buffer_var = op->args[0].as<VarNode>(); + const VarNode* buffer_var = GetBufferVarFromData(op->args[0]); TVM_FFI_ICHECK(buffer_var); // Get shape const IntImmNode* m = op->args[1].as<IntImmNode>(); @@ -143,10 +154,10 @@ class FragmentChecker : public StmtExprVisitor { static const Op& tvm_bmma_sync_op = Op::Get("tirx.tvm_bmma_sync"); if (op->op.same_as(tvm_mma_sync_op) || op->op.same_as(tvm_bmma_sync_op)) { TVM_FFI_ICHECK_EQ(op->args.size(), 8U); - const VarNode* buffer_var_d = op->args[0].as<VarNode>(); - const VarNode* buffer_var_a = op->args[2].as<VarNode>(); - const VarNode* buffer_var_b = op->args[4].as<VarNode>(); - const VarNode* buffer_var_c = op->args[6].as<VarNode>(); + const VarNode* buffer_var_d = GetBufferVarFromData(op->args[0]); + const VarNode* buffer_var_a = GetBufferVarFromData(op->args[2]); + const VarNode* buffer_var_b = GetBufferVarFromData(op->args[4]); + const VarNode* buffer_var_c = GetBufferVarFromData(op->args[6]); TVM_FFI_ICHECK(buffer_var_d); TVM_FFI_ICHECK(buffer_var_a); TVM_FFI_ICHECK(buffer_var_b); diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc index 4408ec0f83..5006014cef 100644 --- a/src/target/source/codegen_c.cc +++ b/src/target/source/codegen_c.cc @@ -913,7 +913,9 @@ void CodeGenC::VisitStmt_(const DeclBufferNode* op) { 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); + auto it = handle_data_type_.find(source); + RegisterHandleType(op->buffer.get(), + it == handle_data_type_.end() ? op->buffer->dtype : it->second); return; } diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc index 410dc4c568..90c8922c3c 100644 --- a/src/te/operation/create_primfunc.cc +++ b/src/te/operation/create_primfunc.cc @@ -655,6 +655,7 @@ Stmt GenerateStmtFromExternOp(const te::ExternOp& extern_op, CreateFuncInfo* inf input_buffer_map[placeholder.get()] = output_buffer; info->root_alloc.push_back(output_buffer); } + var_map[placeholder.get()] = output_buffer.var(); info->tensor2buffers[output_tensor] = output_buffer; } diff --git a/src/tirx/ir/data_type_rewriter.cc b/src/tirx/ir/data_type_rewriter.cc index 04abe0dc5e..a8cb4f62e1 100644 --- a/src/tirx/ir/data_type_rewriter.cc +++ b/src/tirx/ir/data_type_rewriter.cc @@ -164,6 +164,9 @@ Stmt DataTypeLegalizer::VisitStmt_(const BindNode* op) { } Expr DataTypeLegalizer::VisitExpr_(const VarNode* op) { + if (op->ty.as<BufferTypeNode>()) { + return VisitBufferUse(GetBufferVar(op)).var(); + } if (auto it = var_remap_.find(op); it != var_remap_.end()) { return it->second; } @@ -392,8 +395,8 @@ ffi::Map<ffi::String, ffi::Any> IndexDataTypeRewriter::VisitBlockAnnotations( if (obj == nullptr) { return obj; } - if (obj.as<BufferTypeNode>()) { - BufferVar buffer = obj.as_or_throw<BufferVar>(); + if (auto var = obj.as<Var>(); var && var.value()->ty.as<BufferTypeNode>()) { + BufferVar buffer(var.value()); if (BufferVar new_buffer = VisitBufferUse(buffer); !new_buffer.same_as(buffer)) { return new_buffer; } diff --git a/src/tirx/ir/specialize.cc b/src/tirx/ir/specialize.cc index 804eec2e64..687b37dd3a 100644 --- a/src/tirx/ir/specialize.cc +++ b/src/tirx/ir/specialize.cc @@ -189,6 +189,16 @@ class PrimFuncSpecializer : public StmtExprMutator { private: BufferVar MutateBuffer(const BufferVar& buffer) { + ffi::Optional<ffi::String> specialized_storage_scope; + if (auto it = var_map_.find(buffer.var()); it != var_map_.end()) { + if (const auto* new_var = it->second.as<VarNode>()) { + if (new_var->ty.as<BufferTypeNode>()) { + BufferVar replacement(ffi::GetRef<Var>(new_var)); + specialized_storage_scope = replacement->storage_scope; + } + } + } + ffi::Array<PrimExpr> shape = buffer->shape.Map([this](const PrimExpr& e) { return VisitPrimExpr(e); }); ffi::Array<PrimExpr> strides = @@ -220,8 +230,10 @@ class PrimFuncSpecializer : public StmtExprMutator { } } + bool storage_scope_changed = specialized_storage_scope.has_value() && + specialized_storage_scope.value() != buffer->storage_scope; if (buffer->elem_offset.same_as(elem_offset) && buffer->shape.same_as(shape) && - buffer->strides.same_as(strides) && !layout_changed) { + buffer->strides.same_as(strides) && !layout_changed && !storage_scope_changed) { return buffer; } else { auto n = CopyBufferType(buffer); @@ -231,6 +243,9 @@ class PrimFuncSpecializer : public StmtExprMutator { if (layout_changed) { n->layout = std::move(layout); } + if (storage_scope_changed) { + n->storage_scope = specialized_storage_scope.value(); + } return RebuildBufferVar(buffer, std::move(n)); } } diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index 14ca7f7cda..817da4efed 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -382,6 +382,9 @@ TVM_FFI_STATIC_INIT_BLOCK() { // Evaluate Evaluate::Evaluate(Expr value, Span span) { TVM_FFI_ICHECK(value.defined()); + TVM_FFI_ICHECK(!(value->IsInstance<VarNode>() && value->ty.as<BufferTypeNode>())) + << "A buffer variable cannot be used as a scalar Evaluate value; " + << "use buffer.data to evaluate its physical pointer"; ffi::ObjectPtr<EvaluateNode> node = ffi::make_object<EvaluateNode>(); node->value = std::move(value); diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc index 87915fa145..9d79955e28 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 == "global") { + if (ptr_type->storage_scope.empty()) { rhs = rhs->Call({element_type}, kwargs_keys, kwargs_values); } else { rhs = rhs->Call({element_type, diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc index 5b6c30e111..9aa9822c10 100644 --- a/src/tirx/script/printer/stmt.cc +++ b/src/tirx/script/printer/stmt.cc @@ -176,7 +176,8 @@ TVM_SCRIPT_REPR(tirx::TilePrimitiveCallNode, ReprPrintTIR); TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) .set_dispatch<tirx::Evaluate>("", [](tirx::Evaluate eval, AccessPath p, IRDocsifier d) -> Doc { ExprDoc value = d->AsDoc<ExprDoc>(eval->value, p->Attr("value")); - if (eval->value->IsInstance<CallNode>()) { + const auto* call = eval->value.as<CallNode>(); + if (call && !call->op.same_as(tirx::builtin::buffer_data())) { return ExprStmtDoc(value); } return ExprStmtDoc(TIR(d, "evaluate")->Call({value})); diff --git a/src/tirx/transform/flatten_buffer.cc b/src/tirx/transform/flatten_buffer.cc index 72ea51f40c..8650099a21 100644 --- a/src/tirx/transform/flatten_buffer.cc +++ b/src/tirx/transform/flatten_buffer.cc @@ -48,6 +48,9 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { arith::Analyzer ana; auto pass = BufferFlattener(ana); pass.MarkBufferMapShapes(func); + for (const auto& [param, buffer] : func->buffer_map) { + pass.extern_buffers_.insert(buffer); + } auto body = pass.VisitStmt(func->body); // The buffers in func->buffer_map are deliberately left @@ -121,7 +124,17 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { } Stmt VisitStmt_(const DeclBufferNode* op) final { - Expr data = VisitExpr(op->data); + Expr data = op->data; + bool is_extern_buffer_source = false; + if (const auto* call = op->data.as<CallNode>(); + call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) { + if (const auto* var = call->args[0].as<VarNode>(); var && var->ty.as<BufferTypeNode>()) { + is_extern_buffer_source = extern_buffers_.count(BufferVar(ffi::GetRef<Var>(var))); + } + } + if (!is_extern_buffer_source) { + data = VisitExpr(op->data); + } BufferVar flattened = GetFlattenedBuffer(op->buffer); if (flattened.same_as(op->buffer) && data.same_as(op->data)) { return ffi::GetRef<Stmt>(op); @@ -223,8 +236,8 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { */ std::unordered_set<BufferVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> buffers_used_; - /*! \brief The updated external buffer map. */ - ffi::Map<Var, BufferVar> updated_extern_buffer_map_; + /*! \brief Buffers whose storage is supplied by a PrimFunc parameter. */ + std::unordered_set<BufferVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> extern_buffers_; }; PrimFunc FlattenBuffer(PrimFunc f) { return BufferFlattener::Flatten(f); } diff --git a/tests/python/relax/test_blockbuilder_core.py b/tests/python/relax/test_blockbuilder_core.py index bdee939d59..8df1ca822d 100644 --- a/tests/python/relax/test_blockbuilder_core.py +++ b/tests/python/relax/test_blockbuilder_core.py @@ -389,7 +389,7 @@ def test_call_te_unique_tensor_name(): buffer_B = f_matmul.buffer_map[param_B] assert param_A.name != param_B.name assert buffer_A.name != buffer_B.name - assert buffer_A.data.name != buffer_B.data.name + assert not buffer_A.same_as(buffer_B) def test_call_te_with_unsupported_shape_arg(): 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 16c420263e..daf9430198 100644 --- a/tests/python/s_tir/analysis/test_sblock_access_region.py +++ b/tests/python/s_tir/analysis/test_sblock_access_region.py @@ -120,6 +120,18 @@ def opaque_access_with_tvm_access_ptr_func() -> None: T.evaluate(C.access_ptr("rw")) [email protected]_func(s_tir=True) +def decl_buffer_alias_func( + A: T.Buffer((16,), "float32"), + B: T.Buffer((16,), "float32"), +) -> None: + with T.sblock("alias"): + T.reads(A[0]) + T.writes(B[0]) + A_view = T.decl_buffer((16,), "float32", data=A.data) + B[0] = A[0] + A_view[0] + + @T.prim_func(s_tir=True) def access_in_if_then_else_func() -> None: A = T.sblock_alloc_buffer([8]) @@ -263,6 +275,16 @@ def test_opaque_access_with_tvm_access_ptr(): tvm.ir.assert_structural_equal(ret0[1], ret1[1]) +def test_decl_buffer_alias_is_not_an_opaque_access(): + block = decl_buffer_alias_func.body.block + buffer_var_map = {buf: buf for buf in decl_buffer_alias_func.buffer_map.values()} + + reads, writes, opaque = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) + tvm.ir.assert_structural_equal(block.reads, reads) + tvm.ir.assert_structural_equal(block.writes, writes) + tvm.ir.assert_structural_equal([], opaque) + + def test_match_buffer(): root_block = match_buffer_func.body.block block = root_block.body.body.body.block diff --git a/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py b/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py index 6441cd6745..f25f53c67c 100644 --- a/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py +++ b/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py @@ -76,12 +76,14 @@ def test_prim_func_buffer_map(): AccessPath.root() .attr("buffer_map") .map_item(func1.params[1]) + .attr("ty") .attr("shape") .array_item(1) .attr("value"), AccessPath.root() .attr("buffer_map") .map_item(func2.params[1]) + .attr("ty") .attr("shape") .array_item(1) .attr("value"), @@ -139,8 +141,20 @@ def test_allocate(): assert _error_message(ve.value) == _expected_result( func1, func2, - AccessPath.root().attr("body").attr("buffer").attr("shape").array_item(0).attr("value"), - AccessPath.root().attr("body").attr("buffer").attr("shape").array_item(0).attr("value"), + AccessPath.root() + .attr("body") + .attr("buffer") + .attr("ty") + .attr("shape") + .array_item(0) + .attr("value"), + AccessPath.root() + .attr("body") + .attr("buffer") + .attr("ty") + .attr("shape") + .array_item(0) + .attr("value"), ) diff --git a/tests/python/tvmscript/test_tvmscript_printer_tir.py b/tests/python/tvmscript/test_tvmscript_printer_tir.py index c29b99eaa3..ac0d761667 100644 --- a/tests/python/tvmscript/test_tvmscript_printer_tir.py +++ b/tests/python/tvmscript/test_tvmscript_printer_tir.py @@ -90,7 +90,7 @@ def main(a: T.handle, B: T.Buffer((256, 256), "float32")): ) -def test_prim_func_no_sugar_shared_buffer_data(): +def test_prim_func_buffer_data_argument_is_scope_hint(): a = tirx.Var("a", "handle") b = tirx.Var("b", "handle") buffer_data = tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A").data @@ -114,9 +114,7 @@ def test_prim_func_no_sugar_shared_buffer_data(): # from tvm.tirx.layout import Axis @T.prim_func(s_tir=True) -def main(a: T.handle, b: T.handle): - A = T.match_buffer(a, (128, 128)) - B = T.match_buffer(b, (256, 256), data=A.data) +def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")): T.evaluate(0) """, )
