This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new a4bfa7e116 [REFACTOR][s_tir] Retire legacy double-buffer injection
(#20553)
a4bfa7e116 is described below
commit a4bfa7e1168f6b3e81edba07b1b4a2d1a58c5a67
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 5 15:44:58 2026 -0400
[REFACTOR][s_tir] Retire legacy double-buffer injection (#20553)
Remove InjectDoubleBuffer and its attribute-region protocol, including
the pass API and configuration, pipeline callers, and legacy
synchronization bookkeeping.
Keep the SBlock double_buffer_scope annotation used by
InjectSoftwarePipeline and its scheduling producers, and clarify the
annotation's write-buffer index semantics.
---
include/tvm/s_tir/stmt.h | 8 +-
include/tvm/s_tir/transform.h | 7 -
python/tvm/s_tir/backend/adreno/pipeline.py | 1 -
python/tvm/s_tir/pipeline.py | 1 -
python/tvm/s_tir/transform/transform.py | 16 -
.../postproc/disallow_async_strided_mem_copy.cc | 1 -
.../meta_schedule/postproc/verify_gpu_code.cc | 1 -
src/s_tir/transform/inject_double_buffer.cc | 475 ---------------------
src/s_tir/transform/storage_access.cc | 23 +-
src/s_tir/transform/storage_access.h | 4 -
src/s_tir/transform/thread_storage_sync.cc | 27 +-
.../test_s_tir_transform_inject_double_buffer.py | 177 --------
12 files changed, 10 insertions(+), 731 deletions(-)
diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h
index 379a1ba005..09d66aab2a 100644
--- a/include/tvm/s_tir/stmt.h
+++ b/include/tvm/s_tir/stmt.h
@@ -234,15 +234,11 @@ constexpr const char* async_scope = "async_scope";
constexpr const char* buffer_bound = "buffer_bound";
/*!
- * \brief Marks production of double buffer data
+ * \brief SBlock annotation selecting a write-buffer index for double
buffering in
+ * InjectSoftwarePipeline.
*/
constexpr const char* double_buffer_scope = "double_buffer_scope";
-/*!
- * \brief Marks region used by double buffer write
- */
-constexpr const char* double_buffer_write = "double_buffer_write";
-
/*!
* \brief String-valued allocation Call attribute containing the TensorCore
fragment shape
*/
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index 90cdf5cba1..fcaa2c530b 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -242,13 +242,6 @@ TVM_DLL Pass LoopPartition();
*/
TVM_DLL Pass InjectVirtualThread();
-/*!
- * \brief Inject double buffer statements.
- *
- * \return The pass.
- */
-TVM_DLL Pass InjectDoubleBuffer();
-
/*!
* \brief Hoist loop-invariant IfThenElse nodes to
* outside the eligible loops.
diff --git a/python/tvm/s_tir/backend/adreno/pipeline.py
b/python/tvm/s_tir/backend/adreno/pipeline.py
index e5a8dd0774..ddd15a2820 100644
--- a/python/tvm/s_tir/backend/adreno/pipeline.py
+++ b/python/tvm/s_tir/backend/adreno/pipeline.py
@@ -57,7 +57,6 @@ def default_tir_pipeline(*, prepare_only=False):
tirx.transform.NarrowDataType(32),
tirx.transform.VectorizeLoop(not
bool(config.get("tirx.disable_vectorize", False))),
s_tir.transform.InjectVirtualThread(),
- s_tir.transform.InjectDoubleBuffer(),
]
if not bool(config.get("tirx.disable_storage_rewrite", False)):
passes.append(tirx.transform.StorageRewrite())
diff --git a/python/tvm/s_tir/pipeline.py b/python/tvm/s_tir/pipeline.py
index d6f159f976..b0c5624a08 100644
--- a/python/tvm/s_tir/pipeline.py
+++ b/python/tvm/s_tir/pipeline.py
@@ -57,7 +57,6 @@ def default_s_tir_pipeline(*, prepare_only=False):
tirx.transform.NarrowDataType(32),
tirx.transform.VectorizeLoop(not
bool(config.get("tirx.disable_vectorize", False))),
s_tir.transform.InjectVirtualThread(),
- s_tir.transform.InjectDoubleBuffer(),
]
if not bool(config.get("tirx.disable_storage_rewrite", False)):
passes.append(tirx.transform.StorageRewrite())
diff --git a/python/tvm/s_tir/transform/transform.py
b/python/tvm/s_tir/transform/transform.py
index 65eb20b4f9..e0ac0a60cc 100644
--- a/python/tvm/s_tir/transform/transform.py
+++ b/python/tvm/s_tir/transform/transform.py
@@ -255,22 +255,6 @@ def InjectVirtualThread():
return _ffi_api.InjectVirtualThread() # type: ignore
-@_ffi.register_object("s_tir.transform.InjectDoubleBufferConfig")
-class InjectDoubleBufferConfig(_ffi.Object):
- """Config for inject double buffer pass"""
-
-
-def InjectDoubleBuffer():
- """Inject double buffer statements.
-
- Returns
- -------
- fpass : tvm.transform.Pass
- The result pass
- """
- return _ffi_api.InjectDoubleBuffer() # type: ignore
-
-
def HoistIfThenElse(variant=None):
"""Hoist loop-invariant IfThenElse nodes to outside the eligible loops.
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 2069b478ba..762a66457a 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
@@ -158,7 +158,6 @@ class DisallowAsyncStridedMemCopyNode : public PostprocNode
{
pass_list.push_back(tirx::transform::NarrowDataType(32));
pass_list.push_back(tirx::transform::StmtSimplify());
pass_list.push_back(s_tir::transform::InjectVirtualThread());
- pass_list.push_back(s_tir::transform::InjectDoubleBuffer());
pass_list.push_back(tirx::transform::VectorizeLoop(true));
pass_list.push_back(tirx::transform::StorageRewrite());
tirx::PrimFunc f = WithAttr(ffi::GetRef<tirx::PrimFunc>(prim_func),
"global_symbol",
diff --git a/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc
b/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc
index 9a6a6c2061..76b0251a79 100644
--- a/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc
+++ b/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc
@@ -190,7 +190,6 @@ class VerifyGPUCodeNode : public PostprocNode {
// Phase 2
pass_list.push_back(tirx::transform::VectorizeLoop(true));
pass_list.push_back(s_tir::transform::InjectVirtualThread());
- pass_list.push_back(s_tir::transform::InjectDoubleBuffer());
pass_list.push_back(tirx::transform::StorageRewrite());
pass_list.push_back(s_tir::transform::MergeSharedMemoryAllocations());
pass_list.push_back(tirx::transform::LowerIntrin());
diff --git a/src/s_tir/transform/inject_double_buffer.cc
b/src/s_tir/transform/inject_double_buffer.cc
deleted file mode 100644
index 36cea3567f..0000000000
--- a/src/s_tir/transform/inject_double_buffer.cc
+++ /dev/null
@@ -1,475 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \brief Inject double buffering optimization for data fetch.
- * \file inject_double_buffer.cc
- */
-#include <tvm/ffi/cast.h>
-#include <tvm/ffi/extra/structural_mutate.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/attrs.h>
-#include <tvm/runtime/logging.h>
-#include <tvm/s_tir/stmt.h>
-#include <tvm/s_tir/stmt_functor.h>
-#include <tvm/s_tir/transform.h>
-#include <tvm/tirx/builtin.h>
-#include <tvm/tirx/op.h>
-
-#include "ir_utils.h"
-
-namespace tvm {
-namespace s_tir {
-using namespace tvm::prim;
-using namespace tvm::tirx;
-
-namespace {
-
-ffi::Optional<Var> GetBufferDataVar(const ffi::Any& data) {
- if (auto var = data.as<Var>()) {
- return var;
- }
- if (const auto* call = data.as<CallNode>();
- call && call->op.same_as(tirx::builtin::buffer_data()) &&
call->args.size() == 1) {
- return call->args[0].as<Var>();
- }
- return std::nullopt;
-}
-
-} // namespace
-
-struct InjectDoubleBufferConfigNode : public ffi::Object {
- int split_loop;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<InjectDoubleBufferConfigNode>().def_ro(
- "split_loop", &InjectDoubleBufferConfigNode::split_loop, "Split loop
factors",
- refl::DefaultValue(1));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.transform.InjectDoubleBufferConfig",
- InjectDoubleBufferConfigNode, ffi::Object);
-};
-
-class InjectDoubleBufferConfig : public ffi::ObjectRef {
- public:
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(InjectDoubleBufferConfig,
ffi::ObjectRef,
- InjectDoubleBufferConfigNode);
-};
-
-TVM_FFI_STATIC_INIT_BLOCK() {
InjectDoubleBufferConfigNode::RegisterReflection(); }
-
-TVM_REGISTER_PASS_CONFIG_OPTION("s_tir.InjectDoubleBuffer",
InjectDoubleBufferConfig);
-
-// Detect double buffer variables.
-class DoubleBufferDetector : public StmtExprVisitor {
- public:
- using StmtExprVisitor::Visit_;
- ffi::Optional<VisitInterrupt> Visit_(const AttrStmtNode* op) final {
- if (op->attr_key == s_tir::attr::double_buffer_scope) {
- if (auto buffer = GetBufferDataVar(op->node)) {
- touched_.insert(buffer.value().get());
- }
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
- } else {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
- }
- return std::nullopt;
- }
-
- // Known loads and stores are not opaque escapes of the buffer variable.
- ffi::Optional<VisitInterrupt> Visit_(const TensorLoadNode* op) final {
- for (const auto& index : op->indices) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(index));
- }
- return std::nullopt;
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->value));
- for (const auto& index : op->indices) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(index));
- }
- return std::nullopt;
- }
-
- // Declared regions carry bounds, not opaque runtime accesses.
- ffi::Optional<VisitInterrupt> Visit_(const TensorRegionNode* op) final {
- if (!op->source.as<TensorVar>()) return StmtExprVisitor::Visit_(op);
- for (const Range& range : op->region) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(range->min));
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(range->extent));
- }
- return std::nullopt;
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const VarNode* op) final {
- if (def_region_kind() != kTVMFFIDefRegionKindNone) return
StmtExprVisitor::Visit_(op);
- if (touched_.count(op)) {
- touched_.erase(op);
- }
- return std::nullopt;
- }
- // The set of touched variable.
- std::unordered_set<const VarNode*> touched_;
-};
-
-class StripDoubleBufferWrite : public StmtExprMutator {
- public:
- using StmtExprMutator::Mutate;
- using StmtExprMutator::Mutate_;
- UnchangedOr<ffi::Any> Mutate(ffi::AnyView value, InplaceMode inplace_mode)
override {
- if (value.as<ExprNode>()) return ffi::Unchanged();
- return StmtExprMutator::Mutate(value, inplace_mode);
- }
-
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- if (op->attr_key == s_tir::attr::double_buffer_write) {
- return Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body);
- } else {
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
- }
-};
-
-class DoubleBufferInjector : public StmtExprMutator {
- public:
- using StmtExprMutator::Mutate;
- using StmtExprMutator::Mutate_;
-
- explicit DoubleBufferInjector(int split_loop) : split_loop_(split_loop) {}
-
- Stmt Inject(Stmt stmt) {
- auto detector = ffi::make_object<DoubleBufferDetector>();
- detector->Visit(stmt);
- if (detector->touched_.empty()) return stmt;
- for (const VarNode* v : detector->touched_) {
- dbuffer_info_[v] = StorageEntry();
- }
- return s_tir::ConvertSSA(Mutate(stmt,
InplaceMode::kAllow).ValueOrUnchanged(std::move(stmt)));
- }
-
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- if (op->attr_key == s_tir::attr::double_buffer_scope) {
- return MakeProducer(op, inplace_mode);
- } else {
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
- }
-
- UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
final {
- if (const auto* call = op->value.as<CallNode>();
- call && call->op.same_as(tirx::builtin::alloc_tensor())) {
- return Mutate_AllocTensor(op, call, inplace_mode);
- }
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
-
- UnchangedOr<Stmt> Mutate_AllocTensor(const BindNode* op, const CallNode*
call,
- InplaceMode inplace_mode) {
- const VarNode* buf = op->var.as_or_throw<TensorVar>().get();
- auto it = dbuffer_info_.find(buf);
- if (it != dbuffer_info_.end()) {
- StorageEntry& entry = it->second;
- tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
- entry.scope = call->args[2].as_or_throw<StringImm>()->value;
-
- TVM_FFI_ICHECK_EQ(shape->fields.size(), 1) << "InjectDoubleBuffer
expects flat 1-d buffers. "
- << "Has FlattenBuffer been
run?";
- entry.stride = shape->fields[0].as_or_throw<PrimExpr>();
-
- // In flat IR, AllocTensor appears before its usage in the SeqStmt,
- // so entry.loop may not be set yet. Defer double-buffer allocation
- // processing to be handled in VisitStmt_(ForNode*).
- pending_dbuffer_allocs_.insert_or_assign(buf, ffi::GetRef<Bind>(op));
- // Remove the original AllocTensor (will be re-emitted in ForNode
visitor)
- return Evaluate(0);
- } else {
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
- }
-
- UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
- loop_nest_.push_back(op);
- Stmt stmt = StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- // Process any pending double-buffer allocations that were deferred
- // from allocation binding visitor -- now entry.loop should be set.
- for (auto pend_it = pending_dbuffer_allocs_.begin();
- pend_it != pending_dbuffer_allocs_.end();) {
- auto db_it = dbuffer_info_.find(pend_it->first);
- if (db_it != dbuffer_info_.end() && db_it->second.loop != nullptr) {
- StorageEntry& entry = db_it->second;
- const Bind& alloc = pend_it->second;
- auto new_buf = GetRemappedBuffer(alloc->var.as_or_throw<TensorVar>(),
entry.stride.value());
- auto& alloc_nest = loop_allocs_[entry.loop];
- const auto* call = alloc->value.as<CallNode>();
- alloc_nest.emplace_back(Bind(new_buf.var(),
- Call(new_buf.type(),
tirx::builtin::alloc_tensor(),
- {tvm::Tuple(new_buf->shape,
call->args[0]->span),
- DataTypeImm(new_buf->dtype->dtype,
call->args[1]->span),
- StringImm(new_buf.scope(),
call->args[2]->span)},
- call->attrs, call->ty_args,
call->span),
- alloc->span));
- pend_it = pending_dbuffer_allocs_.erase(pend_it);
- } else {
- ++pend_it;
- }
- }
- auto it = loop_pre_.find(op);
- if (it != loop_pre_.end()) {
- const ForNode* old_loop = stmt.as<ForNode>();
- if (split_loop_ != 0) {
- // Explicitly unroll the loop
- TVM_FFI_ICHECK(split_loop_ % 2 == 0 || split_loop_ == 1)
- << "It is better to split with multiple of 2";
- TVM_FFI_ICHECK(is_zero(old_loop->min));
- PrimExpr zero = old_loop->min;
- PrimExpr new_ext = old_loop->extent - IntImm(old_loop->loop_var.ty(),
1);
- PrimExpr factor = IntImm(new_ext.ty(), split_loop_);
- PrimExpr outer_ext = new_ext / factor;
- PrimExpr tail_base = outer_ext * factor;
- Var outer_var(old_loop->loop_var->name + ".outer",
old_loop->loop_var.ty());
- std::unordered_map<const VarNode*, PrimExpr> vmap;
- auto map_var = [&vmap](const Var& var) ->
ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
- if (auto it = vmap.find(var.get()); it != vmap.end()) return
ffi::Any(it->second);
- return ffi::Unchanged();
- };
- std::vector<Stmt> loop_seq;
- for (int32_t i = 0; i < split_loop_; ++i) {
- vmap.insert_or_assign(
- old_loop->loop_var.get(),
- outer_var.as_or_throw<PrimExpr>() * factor + IntImm(factor.ty(),
i));
- loop_seq.emplace_back(
- ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(old_loop->body,
map_var)
- .as_or_throw<Stmt>());
- }
- Stmt loop = For(outer_var.as_or_throw<PrimVar>(), zero, outer_ext,
old_loop->kind,
- SeqStmt::Flatten(loop_seq));
- // tail
- std::vector<Stmt> tail_seq;
- Stmt tail_body = ffi::make_object<StripDoubleBufferWrite>()
- ->Mutate(old_loop->body)
- .ValueOrUnchanged(old_loop->body);
- for (int32_t i = 0; i < split_loop_; ++i) {
- PrimExpr idx = tail_base + IntImm(tail_base.ty(), i);
- vmap.insert_or_assign(old_loop->loop_var.get(), idx);
- tail_seq.emplace_back(
- IfThenElse(idx < old_loop->extent,
-
ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(tail_body, map_var)
- .as_or_throw<Stmt>()));
- }
- stmt = SeqStmt::Flatten(loop, tail_seq);
- }
- stmt = SeqStmt::Flatten(it->second, stmt);
- }
- it = loop_allocs_.find(op);
- if (it != loop_allocs_.end()) {
- stmt = MergeNest(it->second, stmt);
- }
- loop_nest_.pop_back();
- return stmt;
- }
-
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode
inplace_mode) final {
- auto value = Mutate(op->value);
- auto indices =
Mutate(op->indices).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
- BufferStore node = ffi::GetRef<BufferStore>(op);
- if (!value.UnchangedOrSameAs(op->value) ||
!indices.UnchangedOrSameAs(op->indices)) {
- auto* n = node.CopyOnWrite();
- n->value = std::move(value).ValueOrUnchanged(op->value);
- n->indices = std::move(indices).ValueOrUnchanged(op->indices);
- }
-
- 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_);
- TVM_FFI_ICHECK(e.switch_write_var.has_value());
-
- TVM_FFI_ICHECK_EQ(node->indices.size(), 1) << "InjectDoubleBuffer
expects flat 1-d buffers. "
- << "Has FlattenBuffer been
run?";
-
- auto writer = node.CopyOnWrite();
- writer->buffer = GetRemappedBuffer(node->buffer, e.stride.value());
- writer->indices = {e.switch_write_var.value().as_or_throw<PrimExpr>() *
e.stride.value() +
- node->indices[0]};
- }
-
- return node;
- }
-
- UnchangedOr<PrimExpr> Mutate_(const TensorLoadNode* op, InplaceMode
inplace_mode) final {
- auto indices =
Mutate(op->indices).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
- TensorLoad node = ffi::GetRef<TensorLoad>(op);
- if (!indices.UnchangedOrSameAs(op->indices)) {
- node.CopyOnWrite()->indices = std::move(indices).ValueUnchecked();
- }
- TensorVar buffer = node->source.as_or_throw<tvm::tirx::TensorVar>();
-
- auto it = dbuffer_info_.find(buffer.get());
- if (it != dbuffer_info_.end()) {
- const StorageEntry& e = it->second;
- TVM_FFI_ICHECK(e.switch_read_var.has_value());
-
- TVM_FFI_ICHECK_EQ(node->indices.size(), 1) << "InjectDoubleBuffer
expects flat 1-d buffers. "
- << "Has FlattenBuffer been
run?";
-
- auto* writer = node.CopyOnWrite();
- writer->source = GetRemappedBuffer(buffer, e.stride.value());
- writer->indices = {e.switch_read_var.value() * e.stride.value() +
node->indices[0]};
- return node;
- }
-
- return node;
- }
-
- TensorVar GetRemappedBuffer(TensorVar buf, PrimExpr stride) {
- TensorVar original = buf;
- if (auto replacement = VarRemapGet(buf).as<TensorVar>()) return
replacement.value();
-
- TVM_FFI_ICHECK(stride.defined());
- // TODO(Lunderberg): Move this pass to before
- // FlattenBuffer. That will simplify the
- // implementation, to be the insertion of a new dimension for the
- // buffer, rather than adjusting the other indices.
- TVM_FFI_ICHECK_EQ(buf->shape.size(), 1) << "InjectDoubleBuffer expects
flat 1-d buffers. "
- << "Has FlattenBuffer been run?";
-
- // Stride gives the distance between the two halves of the
- // double-buffer, not the stride of the buffer's index.
- auto type = CopyTensorType(buf);
- type->shape = {buf->shape[0] + stride};
- buf = RebuildTensorVar(buf, std::move(type));
-
- VarRemapSet(original, buf);
- return buf;
- }
-
- UnchangedOr<Expr> Mutate_(const TensorRegionNode* op, InplaceMode
inplace_mode) final {
- if (!op->source.as<TensorVar>()) {
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
- auto region =
Mutate(op->region).as_or_throw<UnchangedOr<ffi::Array<Range>>>();
- if (region.UnchangedOrSameAs(op->region)) return ffi::Unchanged();
- TensorRegion node = ffi::GetRef<TensorRegion>(op);
- node.CopyOnWrite()->region = std::move(region).ValueUnchecked();
- return node;
- }
-
- UnchangedOr<Expr> Mutate_(const VarNode* op, InplaceMode inplace_mode) final
{
- if (def_region_kind() == kTVMFFIDefRegionKindNone) {
- TVM_FFI_ICHECK(!dbuffer_info_.count(op));
- }
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
-
- private:
- Stmt MakeProducer(const AttrStmtNode* op, InplaceMode inplace_mode) {
- 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()) {
- LOG(WARNING) << "Skip double buffer scope " << op->node;
- return this->Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body);
- }
- StorageEntry& e = it->second;
- e.loop = loop_nest_.back();
- PrimExpr zero = IntImm(e.loop->loop_var.ty(), 0);
- PrimExpr one = IntImm(e.loop->loop_var.ty(), 1);
- PrimExpr two = IntImm(e.loop->loop_var.ty(), 2);
- PrimExpr loop_shift = e.loop->loop_var + one;
- e.switch_write_var = Var(e.loop->loop_var->name + ".db",
e.loop->loop_var.ty());
- e.switch_read_var = indexmod(e.loop->loop_var, two);
- in_double_buffer_scope_ = true;
- Stmt body = this->Mutate(op->body,
inplace_mode).ValueOrUnchanged(op->body);
- in_double_buffer_scope_ = false;
- std::unordered_map<const VarNode*, PrimExpr> vmap;
- auto map_var = [&vmap](const Var& var) ->
ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
- if (auto it = vmap.find(var.get()); it != vmap.end()) return
ffi::Any(it->second);
- return ffi::Unchanged();
- };
- vmap.insert_or_assign(e.switch_write_var.value().get(), zero);
- vmap.insert_or_assign(e.loop->loop_var.get(), zero);
- loop_pre_[e.loop].emplace_back(
- ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(body,
map_var).as_or_throw<Stmt>());
- vmap.insert_or_assign(e.loop->loop_var.get(), loop_shift);
- vmap.insert_or_assign(e.switch_write_var.value().get(),
indexmod(loop_shift, two));
- body = ffi::StructuralMap<ffi::WalkOrder::kPostOrder>(body,
map_var).as_or_throw<Stmt>();
- body = AttrStmt(GetRemappedBuffer(buffer.as_or_throw<TensorVar>(),
e.stride.value()).data(),
- s_tir::attr::double_buffer_write, IntImm::Int32(1), body);
- body = IfThenElse(loop_shift < e.loop->extent, body);
- return body;
- }
- // Storage entry for those who need double buffering.
- struct StorageEntry {
- // The size of the buffer
- ffi::Optional<PrimExpr> stride;
- // The loop we need
- const ForNode* loop{nullptr};
- // The switch variable.
- ffi::Optional<Var> switch_write_var;
- // The switch variable for reading.
- ffi::Optional<PrimExpr> switch_read_var;
- // The storage scope.
- std::string scope;
- };
- // Whether split loop
- int32_t split_loop_;
- // Whether we are inside double buffer scope.
- bool in_double_buffer_scope_{false};
- // The current loop next
- std::vector<const ForNode*> loop_nest_;
- // The allocs to be appended before the loop
- std::unordered_map<const ForNode*, std::vector<Stmt>> loop_allocs_;
- // The stmt to be appended before the loop
- std::unordered_map<const ForNode*, std::vector<Stmt>> loop_pre_;
- // The allocation size of the buffer
- std::unordered_map<const VarNode*, StorageEntry> dbuffer_info_;
- // The updated TensorVar objects
- // Pending double-buffer AllocTensor nodes (deferred from flat AllocTensor
visit)
- std::unordered_map<const VarNode*, Bind> pending_dbuffer_allocs_;
-};
-
-namespace transform {
-
-Pass InjectDoubleBuffer() {
- auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
- auto* n = f.CopyOnWrite();
- auto cfg =
ctx->GetConfig<InjectDoubleBufferConfig>("s_tir.InjectDoubleBuffer");
- if (!cfg.has_value()) {
- cfg = tvm::transform::PassConfigWithDefaults<InjectDoubleBufferConfig>();
- }
- n->body = ffi::make_object<DoubleBufferInjector>(cfg.value()->split_loop)
- ->Inject(std::move(n->body).value());
- return f;
- };
- return CreatePrimFuncPass(pass_func, 0, "s_tir.InjectDoubleBuffer", {});
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("s_tir.transform.InjectDoubleBuffer",
InjectDoubleBuffer);
-}
-
-} // namespace transform
-
-} // namespace s_tir
-} // namespace tvm
diff --git a/src/s_tir/transform/storage_access.cc
b/src/s_tir/transform/storage_access.cc
index e390ed3266..e491800e12 100644
--- a/src/s_tir/transform/storage_access.cc
+++ b/src/s_tir/transform/storage_access.cc
@@ -139,28 +139,7 @@ ffi::Optional<VisitInterrupt>
StorageAccessVisitor::Visit_(const BindNode* op) {
}
ffi::Optional<VisitInterrupt> StorageAccessVisitor::Visit_(const AttrStmtNode*
op) {
- if (op->attr_key == s_tir::attr::double_buffer_write) {
- TVM_FFI_ICHECK(double_buffer_write_ == nullptr);
- 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_ = ResolveBuffer(buffer.value()).get();
- scope_.push_back(std::vector<StmtEntry>());
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
- StmtEntry s;
- s.stmt = op;
- s.access = Summarize(std::move(scope_.back()), nullptr);
- scope_.pop_back();
- if (!s.access.empty()) {
- for (AccessEntry& e : s.access) {
- if (e.type == kWrite && e.buffer.get() == double_buffer_write_) {
- e.double_buffer_write = true;
- }
- }
- scope_.back().emplace_back(std::move(s));
- }
- double_buffer_write_ = nullptr;
- } else if (op->attr_key == tirx::attr::thread_extent) {
+ if (op->attr_key == tirx::attr::thread_extent) {
IterVar iv = op->node.as_or_throw<IterVar>();
env_threads_.push_back(iv);
if (!in_device_env_) {
diff --git a/src/s_tir/transform/storage_access.h
b/src/s_tir/transform/storage_access.h
index afa121643d..b78a8bb378 100644
--- a/src/s_tir/transform/storage_access.h
+++ b/src/s_tir/transform/storage_access.h
@@ -73,8 +73,6 @@ class StorageAccessVisitor : public StmtExprVisitor {
AccessType type;
/*! \brief The storage scope */
StorageScope scope;
- /*! \brief Whether the access is double buffer write */
- bool double_buffer_write = false;
};
/*! \brief Access pattern about a single statement */
struct StmtEntry {
@@ -138,8 +136,6 @@ class StorageAccessVisitor : public StmtExprVisitor {
bool in_device_env_{false};
// Whether we are inside condition.
int condition_counter_{0};
- // The current double buffer write scope.
- const VarNode* double_buffer_write_{nullptr};
// the current free stmt entry.
StmtEntry curr_stmt_;
// The involving threads
diff --git a/src/s_tir/transform/thread_storage_sync.cc
b/src/s_tir/transform/thread_storage_sync.cc
index fe653711d1..237988e7ee 100644
--- a/src/s_tir/transform/thread_storage_sync.cc
+++ b/src/s_tir/transform/thread_storage_sync.cc
@@ -88,12 +88,12 @@ class ThreadSyncPlanner : public StorageAccessVisitor {
}
for (const AccessEntry& acc : s.access) {
if (acc.type == kRead) {
- if (FindConflict(writes, acc, false)) {
+ if (FindConflict(writes, acc)) {
sync_before_stmt = true;
break;
}
} else if (acc.type == kWrite) {
- if (FindConflict(reads, acc, false)) {
+ if (FindConflict(reads, acc)) {
sync_before_stmt = true;
break;
}
@@ -131,12 +131,12 @@ class ThreadSyncPlanner : public StorageAccessVisitor {
bool sync_before_stmt = false;
for (const AccessEntry& acc : s.access) {
if (acc.type == kRead) {
- if (FindConflict(writes, acc, true)) {
+ if (FindConflict(writes, acc)) {
sync_before_stmt = true;
break;
}
} else if (acc.type == kWrite) {
- if (FindConflict(reads, acc, true)) {
+ if (FindConflict(reads, acc)) {
sync_before_stmt = true;
break;
}
@@ -188,28 +188,21 @@ class ThreadSyncPlanner : public StorageAccessVisitor {
}
}
head.insert(head.end(), tail.begin(), tail.end());
- if (loop != nullptr) {
- // clear double buffer flag after a loop is finished.
- for (AccessEntry& e : head) {
- e.double_buffer_write = false;
- }
- }
return head;
}
private:
// find conflicting entry in vec.
- bool FindConflict(const std::vector<AccessEntry>& prev, const AccessEntry&
curr,
- bool loop_carry) {
+ bool FindConflict(const std::vector<AccessEntry>& prev, const AccessEntry&
curr) {
for (const AccessEntry& x : prev) {
- if (FindConflict(x, curr, loop_carry)) {
+ if (FindConflict(x, curr)) {
return true;
}
}
return false;
}
- bool FindConflict(const AccessEntry& prev, const AccessEntry& curr, bool
loop_carry) {
+ bool FindConflict(const AccessEntry& prev, const AccessEntry& curr) {
// Access to different buffers does not conflict.
if (!prev.buffer.same_as(curr.buffer)) {
return false;
@@ -261,12 +254,6 @@ class ThreadSyncPlanner : public StorageAccessVisitor {
return false;
}
- // If this is a read into a double buffer that was previously
- // swapped out, then it doesn't conflict.
- if (prev.double_buffer_write && curr.type == kRead && !loop_carry) {
- return false;
- }
-
// If nothing else allows sharing the same buffer, then they are
// in conflict.
return true;
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_inject_double_buffer.py
b/tests/python/s_tir/transform/test_s_tir_transform_inject_double_buffer.py
deleted file mode 100644
index a9d49efa0d..0000000000
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_double_buffer.py
+++ /dev/null
@@ -1,177 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# ruff: noqa: F841
-
-import tvm_ffi
-
-import tvm
-import tvm.testing
-from tvm.script import ir as I
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-
-
-def _is_buffer_binding(node, *op_names):
- return (
- isinstance(node, tvm.tirx.Bind)
- and isinstance(node.value, tvm.ir.Call)
- and isinstance(node.value.op, tvm.ir.Op)
- and node.value.op.name in op_names
- )
-
-
-def test_double_buffer():
- n = 100
- m = 4
-
- @I.ir_module
- class Module:
- @Ts.prim_func
- def db(A: T.handle("float32"), C: T.handle("float32")):
- A_buf = T.decl_tensor((n * m,), "float32", data=A)
- C_buf = T.decl_tensor((m,), "float32", data=C)
- tx = T.launch_thread("threadIdx.x", 1)
- for i in range(n):
- B = T.alloc_tensor((m,), "float32", scope="shared")
- with T.attr(B.data, "double_buffer_scope", 1):
- for j in range(m):
- B[j] = A_buf[i * 4 + j]
- for j in range(m):
- C_buf[j] = B[j] + T.float32(1.0)
-
- mod = Module
-
- opt = tvm.transform.Sequential(
- [tvm.s_tir.transform.InjectDoubleBuffer(),
tvm.tirx.transform.StmtSimplify()]
- )
-
- with tvm.transform.PassContext(config={"s_tir.InjectDoubleBuffer":
{"split_loop": 2}}):
- mod = opt(mod)
- stmt = mod["db"].body
-
- # After transformation, the buffer allocation should be doubled
- allocate_node = None
-
- def visitor(op):
- nonlocal allocate_node
- if _is_buffer_binding(op, "tirx.alloc_tensor") and "B" in
str(op.var.data):
- allocate_node = op
-
- tvm_ffi.structural_walk(stmt, visitor)
- assert allocate_node is not None
- assert list(allocate_node.var.shape) == [m * 2]
-
- f = tvm.s_tir.transform.ThreadSync("shared")(mod)["db"]
- count = [0]
-
- def count_sync(op):
- if isinstance(op, tvm.ir.Call) and
op.op.same_as(tvm.ir.Op.get("tirx.tvm_storage_sync")):
- count[0] += 1
-
- tvm_ffi.structural_walk(f.body, count_sync)
- assert count[0] == 4
-
-
-def test_double_buffer_transform():
- transform = tvm.ir.transform.Sequential(
- [
- tvm.s_tir.transform.InjectDoubleBuffer(),
- tvm.tirx.transform.StmtSimplify(),
- ]
- )
-
- @I.ir_module
- class Before:
- @Ts.prim_func
- def main(A: T.Tensor([16, 32], "float32"), B: T.Tensor(16, "float32")):
- for i in range(16):
- cache = T.alloc_tensor((32,), "float32")
-
- T.attr(cache.data, "double_buffer_scope", 1)
-
- for j in range(32):
- cache[j] = A[i, j]
-
- B[i] = 0.0
- for j in range(32):
- B[i] = B[i] + cache[j]
-
- After = transform(Before)
- # Verify the double buffer transformation: the allocation should be
doubled (32 -> 64)
- allocate_node = None
-
- def visitor(op):
- nonlocal allocate_node
- if _is_buffer_binding(op, "tirx.alloc_tensor"):
- allocate_node = op
-
- tvm_ffi.structural_walk(After["main"].body, visitor)
- assert allocate_node is not None
- assert list(allocate_node.var.shape) == [64]
-
-
-def test_double_buffer_with_decl_buffer():
- """Like test_double_buffer_transform, but with a declared buffer object"""
-
- transform = tvm.ir.transform.Sequential(
- [
- tvm.s_tir.transform.InjectDoubleBuffer(),
- tvm.tirx.transform.StmtSimplify(),
- ]
- )
-
- @I.ir_module
- class Before:
- @Ts.prim_func
- def main(A: T.Tensor((16, 32), "float32"), B: T.Tensor(16, "float32")):
- for i in range(16):
- cache = T.decl_tensor(32, "float32")
- T.attr(cache.data, "double_buffer_scope", 1)
-
- for j in range(32):
- cache[j] = A[i, j]
-
- B[i] = 0.0
- for j in range(32):
- B[i] = B[i] + cache[j]
-
- @I.ir_module
- class Expected:
- @Ts.prim_func
- def main(A: T.Tensor((16, 32), "float32"), B: T.Tensor(16, "float32")):
- cache = T.decl_tensor(64, "float32")
- for j in range(32):
- cache[j] = A[0, j]
-
- B[0] = T.float32(0.0)
- for j in range(32):
- B[0] = B[0] + cache[j]
-
- for i_outer in range(15):
- T.attr(cache.data, "double_buffer_write", 1)
- for j in range(32):
- cache[(i_outer + 1) % 2 * 32 + j] = A[i_outer + 1, j]
- B[i_outer + 1] = T.float32(0.0)
- for j in range(32):
- B[i_outer + 1] = B[i_outer + 1] + cache[(i_outer + 1) % 2
* 32 + j]
-
- After = transform(Before)
- tvm.ir.assert_structural_equal(After, Expected)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
