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 90659b83e2 [REFACTOR][S-TIR] Represent async and manual
synchronization explicitly (#20562)
90659b83e2 is described below
commit 90659b83e261a7c5b42775c89b843bac2375dfd1
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 6 13:25:44 2026 -0400
[REFACTOR][S-TIR] Represent async and manual synchronization explicitly
(#20562)
Replace async-copy and manual-synchronization attributes with
S-TIR-owned RegionStmt operations, and represent queue commits and waits
as leaf calls. Preserve software-pipeline predicates, commit grouping,
and shared-memory barriers while lowering the temporary operations
before codegen.
Migrate synchronization planning, async-copy lowering, meta-schedule
inspection, TVMScript construction, and the TOPI GPU sort producer
together. Remove the superseded attribute encoding and CUDA codegen
handlers. Keep variable remapping active inside opaque regions during
CSE so rebuilt buffer definitions remain consistent with their uses.
---
include/tvm/s_tir/stmt.h | 26 +--
include/tvm/s_tir/transform.h | 7 +
include/tvm/tirx/stmt.h | 3 -
python/tvm/s_tir/backend/adreno/pipeline.py | 3 +-
python/tvm/s_tir/pipeline.py | 3 +-
python/tvm/s_tir/script/__init__.py | 4 +
python/tvm/s_tir/script/ir_builder/ir.py | 25 +++
python/tvm/s_tir/transform/transform.py | 5 +
python/tvm/topi/gpu/sort.py | 3 +-
src/backend/cuda/codegen/codegen_cuda.cc | 39 ----
src/backend/cuda/codegen/codegen_cuda.h | 1 -
.../postproc/disallow_async_strided_mem_copy.cc | 74 ++------
src/s_tir/script/printer/expr.cc | 24 +++
src/s_tir/stmt.cc | 36 ++++
src/s_tir/transform/inject_ptx_async_copy.cc | 8 +-
src/s_tir/transform/inject_software_pipeline.cc | 72 ++++----
src/s_tir/transform/stmt_simplify.cc | 13 ++
src/s_tir/transform/storage_access.cc | 15 +-
src/s_tir/transform/storage_access.h | 1 -
src/s_tir/transform/thread_storage_sync.cc | 87 +++++++--
src/tirx/transform/common_subexpr_elim.cc | 22 ++-
src/tirx/transform/ir_utils.cc | 8 -
src/tirx/transform/ir_utils.h | 4 -
src/tirx/transform/remove_no_op.cc | 17 --
.../test_s_tir_transform_inject_ptx_async_copy.py | 78 ++++----
...est_s_tir_transform_inject_software_pipeline.py | 196 ++++++++++-----------
26 files changed, 398 insertions(+), 376 deletions(-)
diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h
index c60da8af42..8f45f5dc22 100644
--- a/include/tvm/s_tir/stmt.h
+++ b/include/tvm/s_tir/stmt.h
@@ -216,19 +216,16 @@ class SBlockRealize : public tirx::Stmt {
TVM_DEFINE_OBJECT_REF_COW_METHOD(SBlockRealizeNode);
};
-namespace attr {
+/*! \brief Region marking copies eligible for asynchronous lowering. */
+TVM_DLL const Op& async_copy_scope();
+/*! \brief Commit the preceding asynchronous copies to a queue. */
+TVM_DLL const Op& async_commit();
+/*! \brief Wait until at most the given number of committed groups remain in
flight. */
+TVM_DLL const Op& async_wait();
+/*! \brief Region whose author manages cross-thread synchronization
explicitly. */
+TVM_DLL const Op& manual_sync();
-/*!
- * \brief Annotations for invoking and synchronizing asynchronous operations.
- */
-constexpr const char* async_commit_queue_scope = "async_commit_queue_scope";
-constexpr const char* async_wait_queue_scope =
tirx::attr::async_wait_queue_scope;
-constexpr const char* async_wait_inflight_count =
tirx::attr::async_wait_inflight_count;
-
-/*!
- * \brief Mark that the attached statement runs asynchronously.
- */
-constexpr const char* async_scope = "async_scope";
+namespace attr {
/*!
* \brief SBlock annotation selecting a write-buffer index for double
buffering in
@@ -383,11 +380,6 @@ constexpr const char* warp_execution = "warp_execution";
*/
constexpr const char* layout_transforms = "layout_transforms";
-/*!
- * \brief Mark that the kernel is hand threaded and doesn't need syncs inserted
- */
-constexpr const char* hand_threaded = "hand_threaded";
-
} // namespace attr
} // namespace s_tir
} // namespace tvm
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index 19bf6d23f5..80378db1fc 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -285,6 +285,13 @@ TVM_DLL Pass LowerVtcmAlloc();
*/
TVM_DLL Pass ThreadSync(tvm::ffi::String storage_scope);
+/*!
+ * \brief Lower asynchronous queue operations and remove synchronization
regions.
+ * Run after ThreadSync and optional InjectPTXAsyncCopy.
+ * \return The pass.
+ */
+TVM_DLL Pass LowerSynchronization();
+
/*!
* \brief Infer the TensorCore fragment information using tensor intrinsics.
* \return The pass.
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 4894b8b55f..48fbb8620b 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -864,9 +864,6 @@ constexpr const char* auto_unroll_max_step =
"auto_unroll_max_step";
* Integer policy inherited by nested loops unless they override it. Consumed
by UnrollLoop.
*/
constexpr const char* unroll_explicit = "unroll_explicit";
-/*! \brief Shared execution attributes consumed before and after block
lowering. */
-constexpr const char* async_wait_queue_scope = "async_wait_queue_scope";
-constexpr const char* async_wait_inflight_count = "async_wait_inflight_count";
/*! \brief Annotation key on AllocTensor marking the allocation as volatile. */
constexpr const char* kVolatile = "tirx.volatile";
/*! \brief Mark buffer initial addr alignment in bytes */
diff --git a/python/tvm/s_tir/backend/adreno/pipeline.py
b/python/tvm/s_tir/backend/adreno/pipeline.py
index 5cac658c12..b7910477be 100644
--- a/python/tvm/s_tir/backend/adreno/pipeline.py
+++ b/python/tvm/s_tir/backend/adreno/pipeline.py
@@ -65,7 +65,7 @@ def default_tir_pipeline(*, prepare_only=False):
s_tir.transform.HoistIfThenElse(),
tirx.transform.UnrollLoop(),
s_tir.transform.RenormalizeSplitPattern(),
- tirx.transform.StmtSimplify(),
+ s_tir.transform.StmtSimplify(),
tirx.transform.RemoveNoOp(),
s_tir.transform.RewriteUnsafeSelect(),
]
@@ -94,6 +94,7 @@ def default_tir_pipeline(*, prepare_only=False):
)
if bool(config.get("tirx.use_async_copy", False)):
passes.append(s_tir.transform.InjectPTXAsyncCopy())
+ passes.append(s_tir.transform.LowerSynchronization())
passes.append(s_tir.transform.MergeSharedMemoryAllocations())
if not prepare_only:
passes.extend(
diff --git a/python/tvm/s_tir/pipeline.py b/python/tvm/s_tir/pipeline.py
index df7814ba3b..206d300317 100644
--- a/python/tvm/s_tir/pipeline.py
+++ b/python/tvm/s_tir/pipeline.py
@@ -65,7 +65,7 @@ def default_s_tir_pipeline(*, prepare_only=False):
s_tir.transform.HoistIfThenElse(),
tirx.transform.UnrollLoop(),
s_tir.transform.RenormalizeSplitPattern(),
- tirx.transform.StmtSimplify(),
+ s_tir.transform.StmtSimplify(),
tirx.transform.RemoveNoOp(),
s_tir.transform.RewriteUnsafeSelect(),
]
@@ -94,6 +94,7 @@ def default_s_tir_pipeline(*, prepare_only=False):
)
if bool(config.get("tirx.use_async_copy", False)):
passes.append(s_tir.transform.InjectPTXAsyncCopy())
+ passes.append(s_tir.transform.LowerSynchronization())
passes.append(s_tir.transform.MergeSharedMemoryAllocations())
if not prepare_only:
passes.extend(
diff --git a/python/tvm/s_tir/script/__init__.py
b/python/tvm/s_tir/script/__init__.py
index 4bdf2bb228..78cdf8c94a 100644
--- a/python/tvm/s_tir/script/__init__.py
+++ b/python/tvm/s_tir/script/__init__.py
@@ -41,6 +41,10 @@ def _initialize():
if name not in ("builder", "ir_builder", "jit")
)
for name in (
+ "async_copy_scope",
+ "async_commit",
+ "async_wait",
+ "manual_sync",
"function_",
"arg_",
"bind_",
diff --git a/python/tvm/s_tir/script/ir_builder/ir.py
b/python/tvm/s_tir/script/ir_builder/ir.py
index 2a14db2a0f..592b8bbdf9 100644
--- a/python/tvm/s_tir/script/ir_builder/ir.py
+++ b/python/tvm/s_tir/script/ir_builder/ir.py
@@ -29,10 +29,31 @@ from tvm.script.parser.protocol_registry import
register_mutable_decl as _regist
from tvm.tirx import Expr, IntImm, Var
from tvm.tirx.layout import Layout
from tvm.tirx.script.ir_builder.ir import _get_layout, _record_meta_resource
+from tvm.tirx.script.ir_builder.parser_protocol import region
from . import _ffi_api
+def async_copy_scope():
+ """Mark eligible memory copies for asynchronous-copy lowering."""
+ return region("s_tir.async_copy_scope", [])
+
+
+def async_commit(queue_id):
+ """Commit the preceding asynchronous copies to a queue."""
+ return tir.call_intrin("void", "s_tir.async_commit", queue_id)
+
+
+def async_wait(queue_id, inflight_count):
+ """Wait until at most inflight_count committed groups remain in flight."""
+ return tir.call_intrin("void", "s_tir.async_wait", queue_id,
inflight_count)
+
+
+def manual_sync():
+ """Use explicitly authored synchronization within this region."""
+ return region("s_tir.manual_sync", [])
+
+
@_register_mutable_decl("s_tir.match_buffer")
def match_buffer(
param: TensorLoad | TensorRegion,
@@ -428,8 +449,12 @@ class axis: # pylint: disable=invalid-name
__all__ = [
+ "async_commit",
+ "async_copy_scope",
+ "async_wait",
"axis",
"block_name_suffix_context",
+ "manual_sync",
"match_buffer",
"reads",
"sblock_alloc_buffer",
diff --git a/python/tvm/s_tir/transform/transform.py
b/python/tvm/s_tir/transform/transform.py
index e89513371a..6e691ef2fe 100644
--- a/python/tvm/s_tir/transform/transform.py
+++ b/python/tvm/s_tir/transform/transform.py
@@ -357,6 +357,11 @@ def ThreadSync(storage_scope):
return _ffi_api.ThreadSync(storage_scope) # type: ignore
+def LowerSynchronization():
+ """Lower async queue operations and erase regions after synchronization
planning."""
+ return _ffi_api.LowerSynchronization() # type: ignore
+
+
def InferFragment():
"""Infer the TensorCore fragment information using tensor intrinsics.
diff --git a/python/tvm/topi/gpu/sort.py b/python/tvm/topi/gpu/sort.py
index 3683d19421..c072581049 100644
--- a/python/tvm/topi/gpu/sort.py
+++ b/python/tvm/topi/gpu/sort.py
@@ -20,6 +20,7 @@
import tvm
from tvm import te
+from tvm.s_tir.script import ir_builder as Ts
from tvm.script.ir_builder import IRBuilder
from tvm.tirx.script import ir_builder as T
@@ -100,7 +101,7 @@ def _odd_even_sort(
with T.frame_scope(
[
- T.attr(tvm.tirx.const(0), "hand_threaded", 0),
+ Ts.manual_sync(),
*_thread_launches(nthread_tx, nthread_bx, nthread_by),
]
) as (_, tx, bx, by):
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index 9affe0ec57..e68687cd67 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -1552,45 +1552,6 @@ void CodeGenCUDA::Dispatch_(const CallNode* op,
std::ostream& os) {
}
}
-void CodeGenCUDA::Dispatch_(const AttrStmtNode* op) {
- if (op->attr_key == s_tir::attr::async_commit_queue_scope) {
- const IntImmNode* queue_id = op->value.as<IntImmNode>();
- TVM_FFI_ICHECK(queue_id && queue_id->value == 0)
- << "For CUDA, the index of an async queue must be 0.";
- this->Dispatch(op->body);
- static const Op ptx_cp_async_commit_group_op =
Op::Get("tirx.ptx.cp_async_commit_group");
- // ptx Call layout: [operands...] [slot tokens] [pred marker ""].
- auto commit_group = Call(PrimType::Void(), ptx_cp_async_commit_group_op,
- {StringImm("async"), StringImm("commit_group"),
StringImm("")})
- .as_or_throw<PrimExpr>();
- this->PrintIndent();
- this->Dispatch(commit_group, this->stream);
- this->stream << ";\n";
- return;
- } else if (op->attr_key == s_tir::attr::async_wait_queue_scope) {
- auto wait_attrs = GetAsyncWaitAttributes(op);
- auto queue_id = wait_attrs.first.as<IntImmNode>();
- TVM_FFI_ICHECK(queue_id && queue_id->value == 0)
- << "For CUDA, the index of an async queue must be 0.";
- auto wait_cnt = wait_attrs.second;
- static const Op ptx_cp_async_wait_group_op =
Op::Get("tirx.ptx.cp_async_wait_group");
- // ptx Call layout: [operands...] [slot tokens] [pred marker ""]. The group
- // count is a role="imm" operand baked into the instruction text, so it
must
- // already be a compile-time constant here (the pipeline pass guarantees
it).
- auto wait_group = Call(PrimType::Void(), ptx_cp_async_wait_group_op,
- {wait_cnt, StringImm("async"),
StringImm("wait_group"), StringImm("")})
- .as_or_throw<PrimExpr>();
- this->PrintIndent();
- this->Dispatch(wait_group, this->stream);
- this->stream << ";\n";
- auto inner = op->body.as<AttrStmtNode>();
- TVM_FFI_ICHECK(inner);
- this->Dispatch(inner->body);
- return;
- }
- CodeGenC::Dispatch_(op);
-}
-
void CodeGenCUDA::Dispatch_(const BindNode* op) {
if (const auto* call = op->value.as<CallNode>(); call) {
if (call->op.same_as(tirx::builtin::alloc_tensor())) return
DispatchAllocTensor(op, call);
diff --git a/src/backend/cuda/codegen/codegen_cuda.h
b/src/backend/cuda/codegen/codegen_cuda.h
index bfcca6c2a9..907eff672a 100644
--- a/src/backend/cuda/codegen/codegen_cuda.h
+++ b/src/backend/cuda/codegen/codegen_cuda.h
@@ -80,7 +80,6 @@ class CodeGenCUDA final : public CodeGenC {
void Dispatch_(const ReturnNode* op) final;
void Dispatch_(const BindNode* op) final;
void DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call);
- void Dispatch_(const AttrStmtNode* op) final;
// Target
Target target;
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 f5ce36ad95..fd488714c0 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
@@ -53,70 +53,30 @@ struct AsyncStridedMemCopyFinder : public StmtExprVisitor {
return std::nullopt;
}
- ffi::Optional<VisitInterrupt> Visit_(const AttrStmtNode* attrStmt) final {
- if (!found_) {
- if (attrStmt->attr_key == s_tir::attr::async_commit_queue_scope) {
- auto async_scope = attrStmt->body.as<AttrStmtNode>();
- if (!async_scope) {
-
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(attrStmt));
- }
-
- auto for_loop = async_scope->body.as<ForNode>();
- if (!for_loop) {
-
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(attrStmt));
- }
-
- input_iters.Set(for_loop->loop_var, Range(for_loop->min,
for_loop->extent));
-
- auto bufferstorenode = for_loop->body.as<BufferStoreNode>();
- if (!bufferstorenode) {
-
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(attrStmt));
- }
-
- auto bufferloadnode = bufferstorenode->value.as<TensorLoadNode>();
- if (!bufferloadnode) {
-
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(attrStmt));
- }
-
- // get store buffer; assert it exists and is contiguous given it uses
a single index
- auto bufferstore = bufferstorenode->buffer.as<TensorTypeNode>();
-
- // get load buffer; assert it exists and is contiguous given it uses a
single index
- TensorVar load_buffer =
bufferloadnode->source.as_or_throw<TensorVar>();
- auto bufferload = load_buffer.as<TensorTypeNode>();
-
- if (!bufferstore || !bufferload) {
-
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(attrStmt));
- }
-
- // map loop variable to zero for the store index & simplify
- ffi::Array<PrimExpr> store_index = bufferstorenode->indices;
+ ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op) final {
+ bool previous = in_async_copy_;
+ in_async_copy_ |= op->op.same_as(s_tir::async_copy_scope());
+ auto result = StmtExprVisitor::Visit_(op);
+ in_async_copy_ = previous;
+ return result;
+ }
- // Use DetectIterMap to detect whether store index is non-contiguous.
+ ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
+ if (!found_ && in_async_copy_) {
+ if (const auto* load = op->value.as<TensorLoadNode>()) {
+ // Inspect each copy, including copies grouped under one commit or
predicate.
sym::Analyzer analyzer;
- auto store_iter_map = DetectIterMap(store_index, input_iters, 1,
- sym::IterMapLevel::Surjective,
analyzer, false);
- if (!store_iter_map->errors.empty()) {
- found_ = true;
- }
-
- // map loop variable to zero for the load index & simplify
- ffi::Array<PrimExpr> load_index = bufferloadnode->indices;
-
- // Use DetectIterMap to detect whether load index is non-contiguous.
- auto load_iter_map = DetectIterMap(load_index, input_iters, 1,
- sym::IterMapLevel::Surjective,
analyzer, false);
- if (!load_iter_map->errors.empty()) {
- found_ = true;
- }
- }
- if (!found_) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(attrStmt));
+ auto store_map = DetectIterMap(op->indices, input_iters, 1,
sym::IterMapLevel::Surjective,
+ analyzer, false);
+ auto load_map = DetectIterMap(load->indices, input_iters, 1,
sym::IterMapLevel::Surjective,
+ analyzer, false);
+ found_ = !store_map->errors.empty() || !load_map->errors.empty();
}
}
return std::nullopt;
}
+ bool in_async_copy_ = false;
bool found_ = false;
ffi::Map<PrimVar, Range> input_iters = ffi::Map<PrimVar, Range>();
};
diff --git a/src/s_tir/script/printer/expr.cc b/src/s_tir/script/printer/expr.cc
index e624638e8a..1faf0947fd 100644
--- a/src/s_tir/script/printer/expr.cc
+++ b/src/s_tir/script/printer/expr.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/ir/op.h>
+#include <tvm/s_tir/stmt.h>
#include "../../../tirx/script/printer/utils.h"
@@ -42,7 +43,30 @@ ffi::Optional<ExprDoc>
CpAsyncRawDocTranslate(DocTranslatorObj* d, ffi::AnyView
return NamespaceDoc("tirx")->Attr("s_tir")->Attr("cp_async_raw")->Call(args);
}
+ffi::Optional<ExprDoc> AsyncQueueDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* call =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
CallNode>(input);
+ if (!CanTranslateExplicitResultCall(call) ||
+ !ffi::StructuralEqual()(call->ty, PrimType::Void())) {
+ return RawCall(d, call);
+ }
+ ffi::Array<ExprDoc> args;
+ for (const Expr& arg : call->args) {
+ args.push_back(MaterializeCallArgument(d, arg, d->Translate(arg).value()));
+ }
+ return NamespaceDoc("s_tir")
+ ->Attr(call->op.same_as(s_tir::async_commit()) ? "async_commit" :
"async_wait")
+ ->Call(args);
+}
+
TVM_FFI_STATIC_INIT_BLOCK() {
+ OpDef("s_tir.async_commit")
+ .set_attr<FDocTranslate>(kOpCallDocTranslate,
+
FDocTranslate::FromNative<&AsyncQueueDocTranslate>());
+ OpDef("s_tir.async_wait")
+ .set_attr<FDocTranslate>(kOpCallDocTranslate,
+
FDocTranslate::FromNative<&AsyncQueueDocTranslate>());
OpDef("tirx.s_tir.cp_async_raw")
.set_attr<FDocTranslate>(kOpCallDocTranslate,
FDocTranslate::FromNative<&CpAsyncRawDocTranslate>());
diff --git a/src/s_tir/stmt.cc b/src/s_tir/stmt.cc
index 8f42fe56f9..7b17a48eaa 100644
--- a/src/s_tir/stmt.cc
+++ b/src/s_tir/stmt.cc
@@ -33,6 +33,42 @@ namespace s_tir {
using namespace tvm::tirx;
using namespace tvm::prim;
+const Op& async_copy_scope() {
+ static const Op op = Op::Get("s_tir.async_copy_scope");
+ return op;
+}
+
+const Op& async_commit() {
+ static const Op op = Op::Get("s_tir.async_commit");
+ return op;
+}
+
+const Op& async_wait() {
+ static const Op op = Op::Get("s_tir.async_wait");
+ return op;
+}
+
+const Op& manual_sync() {
+ static const Op op = Op::Get("s_tir.manual_sync");
+ return op;
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ OpDef("s_tir.async_copy_scope", "Mark eligible copies for asynchronous
lowering.").signature();
+ OpDef("s_tir.manual_sync", "Use explicitly authored synchronization within
the body.")
+ .signature();
+ OpDef("s_tir.async_commit", "Commit asynchronous copies to a queue.")
+ .signature(sig::arg<IntImm>("queue_id"))
+ .set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void())
+ .set_attr<TCallEffectKind>("TCallEffectKind",
+
static_cast<int64_t>(CallEffectKind::kUpdateState));
+ OpDef("s_tir.async_wait", "Wait for committed asynchronous copies.")
+ .signature(sig::arg<IntImm>("queue_id"),
sig::arg<PrimExpr>("inflight_count"))
+ .set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void())
+ .set_attr<TCallEffectKind>("TCallEffectKind",
+
static_cast<int64_t>(CallEffectKind::kUpdateState));
+}
+
namespace {
TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>>
MatchBufferRegionVisit(
diff --git a/src/s_tir/transform/inject_ptx_async_copy.cc
b/src/s_tir/transform/inject_ptx_async_copy.cc
index 581b8fbc13..7d1937d98b 100644
--- a/src/s_tir/transform/inject_ptx_async_copy.cc
+++ b/src/s_tir/transform/inject_ptx_async_copy.cc
@@ -50,15 +50,15 @@ class PTXAsyncCopyInjector : public StmtExprMutator {
return StmtExprMutator::Mutate(value, inplace_mode);
}
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* attr, InplaceMode
inplace_mode) {
- if (attr->attr_key == s_tir::attr::async_scope) {
+ UnchangedOr<Stmt> Mutate_(const RegionStmtNode* region, InplaceMode
inplace_mode) {
+ if (region->op.same_as(s_tir::async_copy_scope())) {
TVM_FFI_ICHECK(in_async == false) << "Nested async scopes not supported";
in_async = true;
- auto body = this->Mutate(attr->body,
inplace_mode).ValueOrUnchanged(attr->body);
+ auto body = this->Mutate(region->body,
inplace_mode).ValueOrUnchanged(region->body);
in_async = false;
return body;
}
- return StmtExprMutator::Mutate_(attr, inplace_mode);
+ return StmtExprMutator::Mutate_(region, inplace_mode);
}
Stmt InjectPTX(const TensorLoadNode* load, const BufferStoreNode* store,
bool predicated = false,
diff --git a/src/s_tir/transform/inject_software_pipeline.cc
b/src/s_tir/transform/inject_software_pipeline.cc
index a168ac1919..e23d368138 100644
--- a/src/s_tir/transform/inject_software_pipeline.cc
+++ b/src/s_tir/transform/inject_software_pipeline.cc
@@ -680,19 +680,19 @@ class PipelineRewriter : public StmtExprMutator {
//
// for i in range(13):
// # Stage 0
- // async_commit_queue(0):
- // async_scope:
- // A_shared[(i + 3) % 4] = A[...]
+ // with async_copy_scope():
+ // A_shared[(i + 3) % 4] = A[...]
+ // async_commit(0)
//
//
// # Stage 1
- // async_wait_queue(0, 5):
- // compute(A_shared[i], B_shared[i])
+ // async_wait(0, 5)
+ // compute(A_shared[i], B_shared[i])
//
// # Stage 0
- // async_commit_queue(0)
- // async_scope:
- // B_shared[(i + 3) % 4] = B[...]
+ // with async_copy_scope():
+ // B_shared[(i + 3) % 4] = B[...]
+ // async_commit(0)
//
//
// Here, multiple async producers in the same stage are interleaved with
their consumer in
@@ -701,14 +701,14 @@ class PipelineRewriter : public StmtExprMutator {
//
// for i in range(13):
// # Stage 0
- // async_commit_queue(0):
- // async_scope:
- // A_shared[(i + 3) % 4] = A[...]
- // B_shared[(i + 3) % 4] = B[...]
+ // with async_copy_scope():
+ // A_shared[(i + 3) % 4] = A[...]
+ // B_shared[(i + 3) % 4] = B[...]
+ // async_commit(0)
//
// # Stage 1
- // async_wait_queue(0, 3):
- // compute(A_shared[i], B_shared[i])
+ // async_wait(0, 3)
+ // compute(A_shared[i], B_shared[i])
//
// The correct wait_count can be determined by considering each commit
group separately, and
// summing "per-commit" wait_counts.
@@ -799,12 +799,12 @@ class PipelineRewriter : public StmtExprMutator {
}
if (state.pending_wait.valid()) {
- auto attach_wait_scope = [&new_blocks](int i, int stage_id, PrimExpr
wait_count) {
+ auto attach_wait = [&new_blocks](int i, int stage_id, PrimExpr
wait_count) {
auto& block = new_blocks[i].block;
SBlockNode* n = block.CopyOnWrite();
- n->body =
- AttrStmt(0, s_tir::attr::async_wait_queue_scope,
IntImm::Int32(stage_id),
- AttrStmt(0, s_tir::attr::async_wait_inflight_count,
wait_count, n->body));
+ n->body = SeqStmt({Evaluate(Call(PrimType::Void(),
s_tir::async_wait(),
+ {IntImm::Int32(stage_id),
wait_count})),
+ n->body});
};
if (state.predicate &&
!ana_normalized->CanProve(state.predicate.value())) {
@@ -815,10 +815,10 @@ class PipelineRewriter : public StmtExprMutator {
ffi::Array<PrimExpr>{state.predicate.value(),
state.pending_wait.wait_count.value(), 0})
.as_or_throw<PrimExpr>();
- attach_wait_scope(state.pending_wait.insert_before, stage_id,
wait_count);
+ attach_wait(state.pending_wait.insert_before, stage_id, wait_count);
} else {
- attach_wait_scope(state.pending_wait.insert_before, stage_id,
- state.pending_wait.wait_count.value());
+ attach_wait(state.pending_wait.insert_before, stage_id,
+ state.pending_wait.wait_count.value());
}
}
}
@@ -840,18 +840,10 @@ class PipelineRewriter : public StmtExprMutator {
group_bodies.push_back(new_blocks[i].block->body);
}
- if (group_bodies.size() > 1) {
- auto merged_bodies = SeqStmt(group_bodies);
- group_bodies.clear();
- group_bodies.push_back(merged_bodies);
- }
-
- for (auto body : group_bodies) {
- auto commit_queue_scope =
- AttrStmt(0, s_tir::attr::async_commit_queue_scope,
IntImm::Int32(stage_id), body);
- auto new_block = MakeSBlock(commit_queue_scope,
buffer_data_to_buffer_);
- stmts.push_back(SBlockRealize({}, predicate, new_block));
- }
+ group_bodies.push_back(
+ Evaluate(Call(PrimType::Void(), s_tir::async_commit(),
{IntImm::Int32(stage_id)})));
+ auto new_block = MakeSBlock(SeqStmt(group_bodies),
buffer_data_to_buffer_);
+ stmts.push_back(SBlockRealize({}, predicate, new_block));
}
}
@@ -960,12 +952,12 @@ class PipelineRewriter : public StmtExprMutator {
commit_group_id = local_state.commit_groups.size();
local_state.commit_groups.push_back({new_blocks.size()});
} else {
- // This is the case when one commit_queue groups multiple async
blocks.
- // with commit_queue(stage):
- // async_scope:
- // A_shared[...] = ...
- // async_scope:
- // B_shared[...] = ...
+ // One commit can group multiple async-copy regions.
+ // with async_copy_scope():
+ // A_shared[...] = ...
+ // with async_copy_scope():
+ // B_shared[...] = ...
+ // async_commit(stage)
commit_group_id = local_state.commit_groups.size() - 1;
local_state.commit_groups.back().push_back(new_blocks.size());
@@ -987,7 +979,7 @@ class PipelineRewriter : public StmtExprMutator {
}
SBlockNode* n = new_block.CopyOnWrite();
- n->body = AttrStmt(0, s_tir::attr::async_scope, IntImm::Int32(1),
n->body);
+ n->body = RegionStmt(s_tir::async_copy_scope(), {}, {}, DictAttrs(),
n->body);
}
new_blocks.push_back(
diff --git a/src/s_tir/transform/stmt_simplify.cc
b/src/s_tir/transform/stmt_simplify.cc
index 2e1246c6b2..f9f152caa8 100644
--- a/src/s_tir/transform/stmt_simplify.cc
+++ b/src/s_tir/transform/stmt_simplify.cc
@@ -54,6 +54,19 @@ class StmtSimplifier final : public tirx::StmtSimplifier {
return s_tir::StmtExprMutator::MutateBlockRealize(this, op, inplace_mode);
}
+ UnchangedOr<Stmt> Mutate_(const EvaluateNode* op, InplaceMode inplace_mode)
final {
+ Stmt stmt = Parent::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
+ if (const auto* evaluate = stmt.as<EvaluateNode>()) {
+ if (const auto* call = evaluate->value.as<CallNode>();
+ call && call->op.same_as(s_tir::async_wait()) &&
+ analyzer_->CanProve(call->args[1].as_or_throw<PrimExpr>() < 0)) {
+ // Unrolling may make a wait count negative; such a wait is a no-op.
+ return Evaluate(0);
+ }
+ }
+ return stmt;
+ }
+
protected:
static void InitVTable(VTable* vtable) {
Parent::InitVTable(vtable);
diff --git a/src/s_tir/transform/storage_access.cc
b/src/s_tir/transform/storage_access.cc
index 2a9d8eb360..3b48b36399 100644
--- a/src/s_tir/transform/storage_access.cc
+++ b/src/s_tir/transform/storage_access.cc
@@ -138,18 +138,11 @@ ffi::Optional<VisitInterrupt>
StorageAccessVisitor::Visit_(const BindNode* op) {
return std::nullopt;
}
-ffi::Optional<VisitInterrupt> StorageAccessVisitor::Visit_(const AttrStmtNode*
op) {
- if (op->attr_key == s_tir::attr::hand_threaded) {
- // skip this pass on blocks that were hand_threaded
- // this avoids control flow and read/write conflicts
- // between hand-threaded kernels and automatic threading
- } else {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
- }
- return std::nullopt;
-}
-
ffi::Optional<VisitInterrupt> StorageAccessVisitor::Visit_(const
RegionStmtNode* op) {
+ if (op->op.same_as(s_tir::manual_sync())) {
+ // Trust the region's explicit barriers instead of planning access-based
synchronization.
+ return std::nullopt;
+ }
if (op->op.same_as(tirx::builtin::launch_thread()) &&
std::string(op->args[0].as_or_throw<StringImm>()->value).rfind("vthread", 0) !=
0) {
PrimExpr extent = op->args[1].as_or_throw<PrimExpr>();
diff --git a/src/s_tir/transform/storage_access.h
b/src/s_tir/transform/storage_access.h
index 7ab53970c9..795b70f2f4 100644
--- a/src/s_tir/transform/storage_access.h
+++ b/src/s_tir/transform/storage_access.h
@@ -86,7 +86,6 @@ class StorageAccessVisitor : public StmtExprVisitor {
ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final;
ffi::Optional<VisitInterrupt> Visit_(const EvaluateNode* op) final;
ffi::Optional<VisitInterrupt> Visit_(const BindNode* op) final;
- ffi::Optional<VisitInterrupt> Visit_(const AttrStmtNode* op) final;
ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op) final;
ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) final;
ffi::Optional<VisitInterrupt> Visit_(const IfThenElseNode* op) final;
diff --git a/src/s_tir/transform/thread_storage_sync.cc
b/src/s_tir/transform/thread_storage_sync.cc
index 237988e7ee..a25cca7826 100644
--- a/src/s_tir/transform/thread_storage_sync.cc
+++ b/src/s_tir/transform/thread_storage_sync.cc
@@ -265,23 +265,23 @@ class ThreadSyncPlanner : public StorageAccessVisitor {
};
// There are cases where necessary syncthreads is not inserted by
ThreadSyncInserter.
-// For example, syncthreads is needed after async_wait_queue in the second
loop below,
+// For example, syncthreads is needed after async_wait in the second loop
below,
// but since ThreadSyncInserter is not aware of the asynchronous semantics, it
cannot tell
// that the syncthreads is needed there.
//
// // Pipeline prologue
// for i in range(125):
-// async_commit_queue(0):
-// async_scope:
-// shared[(i + 3) % 4] = ...
+// with async_copy_scope():
+// shared[(i + 3) % 4] = ...
+// async_commit(0)
// ...
//
// // Pipeline Epilogue
// for i in range(3):
-// async_wait_queue(0, 2 - i):
-// local[...] = shared[(i + 125) % 4]
+// async_wait(0, 2 - i)
+// local[...] = shared[(i + 125) % 4]
-// This class adds syncthreads after all async_wait_queue. That includes
syncthreads that
+// This class adds syncthreads after all async_wait operations. That includes
syncthreads that
// can be inserted by ThreadSyncInserter as well, but ThreadSyncInserter will
not insert
// duplicate syncthreads if it finds an existing one at the synchronization
point.
class ThreadSyncAfterWaitQueueInserter : public StmtExprMutator {
@@ -291,16 +291,12 @@ class ThreadSyncAfterWaitQueueInserter : public
StmtExprMutator {
explicit ThreadSyncAfterWaitQueueInserter(StorageScope sync_scope) :
sync_scope_(sync_scope) {}
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- if (op->attr_key == s_tir::attr::async_wait_queue_scope) {
+ UnchangedOr<Stmt> Mutate_(const EvaluateNode* op, InplaceMode inplace_mode)
final {
+ if (const auto* call = op->value.as<CallNode>();
+ call && call->op.same_as(s_tir::async_wait())) {
auto sync = Evaluate(Call(PrimType::Int(32),
tirx::builtin::tvm_storage_sync(),
- {StringImm(sync_scope_.to_string())})
- .as_or_throw<PrimExpr>());
- auto inner = op->body.as<AttrStmtNode>();
- TVM_FFI_ICHECK(inner && inner->attr_key ==
s_tir::attr::async_wait_inflight_count);
- auto new_body = SeqStmt({sync, inner->body});
- return AttrStmt(0, s_tir::attr::async_wait_queue_scope, op->value,
- AttrStmt(0, s_tir::attr::async_wait_inflight_count,
inner->value, new_body));
+ {StringImm(sync_scope_.to_string())}));
+ return SeqStmt({ffi::GetRef<Stmt>(op), sync});
}
return StmtExprMutator::Mutate_(op, inplace_mode);
}
@@ -351,6 +347,46 @@ Stmt ThreadSync(Stmt stmt, std::string storage_scope) {
.ValueOrUnchanged(std::move(stmt));
}
+// The synchronization markers have served their purpose after all ThreadSync
passes.
+// Only CUDA has backend queue operations; other targets retain synchronous
copies.
+class SynchronizationLowerer : public StmtExprMutator {
+ public:
+ using StmtExprMutator::Mutate;
+ using StmtExprMutator::Mutate_;
+ explicit SynchronizationLowerer(bool is_cuda) : is_cuda_(is_cuda) {}
+
+ UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) final {
+ if (op->op.same_as(s_tir::manual_sync()) ||
op->op.same_as(s_tir::async_copy_scope())) {
+ return Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body);
+ }
+ return StmtExprMutator::Mutate_(op, inplace_mode);
+ }
+
+ UnchangedOr<Stmt> Mutate_(const EvaluateNode* op, InplaceMode inplace_mode)
final {
+ const auto* call = op->value.as<CallNode>();
+ if (!call ||
+ (!call->op.same_as(s_tir::async_commit()) &&
!call->op.same_as(s_tir::async_wait()))) {
+ return StmtExprMutator::Mutate_(op, inplace_mode);
+ }
+ if (!is_cuda_) return Evaluate(0);
+ TVM_FFI_ICHECK_EQ(call->args[0].as_or_throw<IntImm>()->value, 0)
+ << "For CUDA, the index of an async queue must be 0.";
+ if (call->op.same_as(s_tir::async_commit())) {
+ static const Op commit = Op::Get("tirx.ptx.cp_async_commit_group");
+ return Evaluate(Call(PrimType::Void(), commit,
+ {StringImm("async"), StringImm("commit_group"),
StringImm("")}));
+ }
+ static const Op wait = Op::Get("tirx.ptx.cp_async_wait_group");
+ // PTX's immediate operand retains the compile-time wait-count requirement.
+ return Evaluate(
+ Call(PrimType::Void(), wait,
+ {call->args[1], StringImm("async"), StringImm("wait_group"),
StringImm("")}));
+ }
+
+ private:
+ bool is_cuda_;
+};
+
namespace transform {
Pass ThreadSync(ffi::String storage_scope) {
@@ -363,10 +399,25 @@ Pass ThreadSync(ffi::String storage_scope) {
return CreatePrimFuncPass(pass_func, 0, "s_tir.ThreadSync", {});
}
+Pass LowerSynchronization() {
+ auto pass_func = [](PrimFunc f, IRModule, PassContext) {
+ if (!f->body.has_value()) return f;
+ auto target = f->GetAttr<Target>(tvm::attr::kTarget);
+ bool is_cuda = target && target.value()->kind->name == "cuda";
+ auto* n = f.CopyOnWrite();
+ n->body = ffi::make_object<SynchronizationLowerer>(is_cuda)
+ ->Mutate(n->body.value(), InplaceMode::kAllow)
+ .ValueOrUnchanged(n->body.value());
+ return f;
+ };
+ return CreatePrimFuncPass(pass_func, 0, "s_tir.LowerSynchronization", {});
+}
+
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("s_tir.transform.ThreadSync",
- static_cast<Pass (*)(ffi::String)>(ThreadSync));
+ refl::GlobalDef()
+ .def("s_tir.transform.ThreadSync", static_cast<Pass
(*)(ffi::String)>(ThreadSync))
+ .def("s_tir.transform.LowerSynchronization", LowerSynchronization);
}
} // namespace transform
diff --git a/src/tirx/transform/common_subexpr_elim.cc
b/src/tirx/transform/common_subexpr_elim.cc
index f34d358dbc..d6cb658dab 100644
--- a/src/tirx/transform/common_subexpr_elim.cc
+++ b/src/tirx/transform/common_subexpr_elim.cc
@@ -802,6 +802,7 @@ class CSERewriter : public StmtExprMutator {
* values and their copy of the subtree.
*/
UnchangedOr<ffi::Any> Mutate(ffi::AnyView input, InplaceMode inplace_mode)
override {
+ if (opaque_body_) return StmtExprMutator::Mutate(input, inplace_mode);
if (auto prim_expr = input.as<PrimExpr>()) {
auto it = expr_remap_.find(prim_expr.value());
if (it != expr_remap_.end()) return it->second;
@@ -843,22 +844,29 @@ class CSERewriter : public StmtExprMutator {
}
UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) override {
- if (op->op.same_as(tirx::builtin::launch_thread()))
+ if (opaque_body_ || op->op.same_as(tirx::builtin::launch_thread()))
return StmtExprMutator::Mutate_(op, inplace_mode);
// Match the planner's opaque boundary: only operands outside the body
// participate in this plan and may reference its CSE bindings.
auto args = Mutate(op->args,
inplace_mode).as_or_throw<UnchangedOr<ffi::Array<Expr>>>();
auto attrs = Mutate(op->attrs,
inplace_mode).as_or_throw<UnchangedOr<DictAttrs>>();
- if (args.UnchangedOrSameAs(op->args) &&
attrs.UnchangedOrSameAs(op->attrs)) {
- return ffi::Unchanged();
+ RegionStmt region = ffi::GetRef<RegionStmt>(op);
+ if (!args.UnchangedOrSameAs(op->args) ||
!attrs.UnchangedOrSameAs(op->attrs)) {
+ auto* n = region.CopyOnWrite();
+ if (!args.IsUnchanged()) n->args = std::move(args).ValueUnchecked();
+ if (!attrs.IsUnchanged()) n->attrs = std::move(attrs).ValueUnchecked();
}
- auto copy = ffi::make_object<RegionStmtNode>(*op);
- if (!args.IsUnchanged()) copy->args = std::move(args).ValueUnchecked();
- if (!attrs.IsUnchanged()) copy->attrs = std::move(attrs).ValueUnchecked();
- return Stmt(std::move(copy));
+ // Definitions outside the region may have been rebuilt, including tensor
+ // variables whose shapes changed. Remap their uses without applying CSE.
+ opaque_body_ = true;
+ auto result = StmtExprMutator::Mutate_(region.get(),
inplace_mode).ValueOrUnchanged(region);
+ opaque_body_ = false;
+ return result;
}
private:
+ /*! \brief Traverse opaque bodies only to remap variables. */
+ bool opaque_body_ = false;
/*! \brief Plan: stmts to insert each target (keyed by object identity). */
InsertBeforeTable insert_before_;
/*! \brief Plan: expressions to replace with CSE vars. */
diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc
index a13c7130c6..caec09540d 100644
--- a/src/tirx/transform/ir_utils.cc
+++ b/src/tirx/transform/ir_utils.cc
@@ -405,14 +405,6 @@ ffi::Array<PrimExpr> GetBufferAllocationShape(const
TensorVar& buffer) {
return alloc_shape;
}
-// Attribute strings are the metadata protocol shared by lowered and
schedulable statements.
-std::pair<PrimExpr, PrimExpr> GetAsyncWaitAttributes(const AttrStmtNode* op) {
- TVM_FFI_ICHECK(op && op->attr_key ==
tvm::tirx::attr::async_wait_queue_scope);
- auto inner = op->body.as<AttrStmtNode>();
- TVM_FFI_ICHECK(inner && inner->attr_key ==
tvm::tirx::attr::async_wait_inflight_count);
- return std::make_pair(op->value.as_or_throw<PrimExpr>(),
inner->value.as_or_throw<PrimExpr>());
-}
-
int Stoi(const std::string& str) {
try {
return std::stoi(str);
diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h
index 8dccdeed66..db29ba0fb5 100644
--- a/src/tirx/transform/ir_utils.h
+++ b/src/tirx/transform/ir_utils.h
@@ -358,10 +358,6 @@ struct FragmentInfo {
*/
std::unordered_map<const VarNode*, FragmentInfo>
GetTensorCoreFragmentInfo(const Stmt& stmt);
-// Return the queue id and the in-flight count associated with the given
-// tvm::tirx::attr::async_wait_queue_scope annotation.
-std::pair<PrimExpr, PrimExpr> GetAsyncWaitAttributes(const AttrStmtNode* op);
-
/*!
* \brief Split string separated by "," to get wmma fragment dimension size.
* \param shape_str The string to split.
diff --git a/src/tirx/transform/remove_no_op.cc
b/src/tirx/transform/remove_no_op.cc
index c07ed6eb25..5d9033d77a 100644
--- a/src/tirx/transform/remove_no_op.cc
+++ b/src/tirx/transform/remove_no_op.cc
@@ -92,23 +92,6 @@ class NoOpRemover : public IRMutatorWithAnalyzer {
private:
UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- if (op->attr_key == tvm::tirx::attr::async_wait_queue_scope) {
- auto wait_attrs = GetAsyncWaitAttributes(op);
- auto wait_cnt = wait_attrs.second;
- sym::Analyzer ana;
- if (ana->CanProve(wait_cnt < 0)) {
- // A negative wait count can arise if it depends on a loop variable.
- // For example, a wait count 1 - i can be negative after loop
unrolling.
- // We assume that such wait is a nop.
- auto inner = op->body.as<AttrStmtNode>();
- TVM_FFI_ICHECK(inner);
- return Parent::Mutate(ffi::AnyView(inner->body),
- inner->unique() ? inplace_mode :
InplaceMode::kDisallow)
- .ValueOrUnchanged(inner->body)
- .as_or_throw<Stmt>();
- }
- }
-
Stmt stmt = Parent::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
op = stmt.as<AttrStmtNode>();
return is_no_op(op->body) ? MakeEvaluate(op->value) : stmt;
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 57928cc917..f35b5ead7e 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
@@ -73,16 +73,16 @@ def generate_global_to_shared_vectorized_copy(dtype,
vector_size):
Ts.reads(A[0:32, 0:128])
Ts.writes(B[0:32, 0:128])
- T.attr("default", "async_scope", 1)
- for i in T.serial(num_iters):
- for j in T.vectorized(vector_size):
- A_shared[tx, i * vector_size_expr + j] = A[tx, i *
vector_size_expr + j]
+ with Ts.async_copy_scope():
+ for i in T.serial(num_iters):
+ for j in T.vectorized(vector_size):
+ A_shared[tx, i * vector_size_expr + j] = A[tx, i *
vector_size_expr + j]
- T.evaluate(T.ptx.cp.async_.commit_group())
- T.evaluate(T.ptx.cp.async_.wait_group(0))
+ T.evaluate(T.ptx.cp.async_.commit_group())
+ T.evaluate(T.ptx.cp.async_.wait_group(0))
- for i in range(128):
- B[tx, i] = A_shared[tx, i]
+ for i in range(128):
+ B[tx, i] = A_shared[tx, i]
return ptx_global_to_shared_copy
@@ -99,15 +99,15 @@ def ptx_global_to_shared_copy_fp32x1(
Ts.reads(A[0:32, 0:128])
Ts.writes(B[0:32, 0:128])
- T.attr("default", "async_scope", 1)
- for i in T.serial(128):
- A_shared[tx, i] = A[tx, i]
+ with Ts.async_copy_scope():
+ for i in T.serial(128):
+ A_shared[tx, i] = A[tx, i]
- T.evaluate(T.ptx.cp.async_.commit_group())
- T.evaluate(T.ptx.cp.async_.wait_group(0))
+ T.evaluate(T.ptx.cp.async_.commit_group())
+ T.evaluate(T.ptx.cp.async_.wait_group(0))
- for i in range(128):
- B[tx, i] = A_shared[tx, i]
+ for i in range(128):
+ B[tx, i] = A_shared[tx, i]
@Ts.prim_func
@@ -125,17 +125,17 @@ def ptx_global_to_shared_dyn_copy_fp16x8(
Ts.reads(A[0:32, 0:128], B[0:32, 0:128])
Ts.writes(C[0:32, 0:128])
- T.attr("default", "async_scope", 1)
- for i in T.serial(16):
- for j in T.vectorized(8):
- A_shared[tx, i * 8 + j] = A[tx, i * 8 + j]
- B_shared[tx, i * 8 + j] = B[tx, i * 8 + j]
+ with Ts.async_copy_scope():
+ for i in T.serial(16):
+ for j in T.vectorized(8):
+ A_shared[tx, i * 8 + j] = A[tx, i * 8 + j]
+ B_shared[tx, i * 8 + j] = B[tx, i * 8 + j]
- T.evaluate(T.ptx.cp.async_.commit_group())
- T.evaluate(T.ptx.cp.async_.wait_group(0))
+ T.evaluate(T.ptx.cp.async_.commit_group())
+ T.evaluate(T.ptx.cp.async_.wait_group(0))
- for i in range(128):
- C[tx, i] = A_shared[tx, i] + B_shared[tx, i]
+ for i in range(128):
+ C[tx, i] = A_shared[tx, i] + B_shared[tx, i]
@pytest.mark.gpu
@@ -284,25 +284,19 @@ extern "C" __global__ void __launch_bounds__(16)
main_kernel(float* __restrict__
_ptr[((int)threadIdx.x)] = 0x0p+0f/*0.000000e+00*/;
_ptr_1[((int)threadIdx.x)] = 0x0p+0f/*0.000000e+00*/;
tvm_builtin_ptx_cp_async_commit_group_async_commit_group();
- int cse_v1 = (((int)threadIdx.x) * 14);
- int cse_v2 = (((int)threadIdx.x) + 16);
ptx_cp_async_legacy_ca_4_4_4(_ptr, (((int)threadIdx.x) + 16), A_ptr,
(((int)threadIdx.x) * 14));
ptx_cp_async_legacy_ca_4_4_4(_ptr_1, (((int)threadIdx.x) + 16), B_ptr,
(((int)threadIdx.x) * 14));
tvm_builtin_ptx_cp_async_commit_group_async_commit_group();
- int cse_v3 = (((int)threadIdx.x) + 32);
- int cse_v6 = ((((int)threadIdx.x) * 14) + 1);
ptx_cp_async_legacy_ca_4_4_4(_ptr, (((int)threadIdx.x) + 32), A_ptr,
((((int)threadIdx.x) * 14) + 1));
ptx_cp_async_legacy_ca_4_4_4(_ptr_1, (((int)threadIdx.x) + 32), B_ptr,
((((int)threadIdx.x) * 14) + 1));
tvm_builtin_ptx_cp_async_commit_group_async_commit_group();
- int cse_v4 = (((int)threadIdx.x) * 16);
+ int cse_v1 = (((int)threadIdx.x) * 16);
for (int i = 0; i < 13; ++i) {
- int cse_v7 = (((((int)threadIdx.x) * 14) + i) + 2);
- int cse_v9 = ((((i + 3) & 3) * 16) + ((int)threadIdx.x));
ptx_cp_async_legacy_pred_ca_4_4_4(_ptr, ((((i + 3) & 3) * 16) +
((int)threadIdx.x)), A_ptr, (((((int)threadIdx.x) * 14) + i) + 2), (i < 12));
tvm_builtin_ptx_cp_async_commit_group_async_commit_group();
tvm_builtin_ptx_cp_async_wait_group_async_wait_group_5();
__syncthreads();
- int cse_v8 = (((i & 3) * 16) + ((int)threadIdx.x));
+ int cse_v5 = (((i & 3) * 16) + ((int)threadIdx.x));
C_ptr[((((int)threadIdx.x) * 16) + i)] = (_ptr[(((i & 3) * 16) +
((int)threadIdx.x))] + _ptr_1[(((i & 3) * 16) + ((int)threadIdx.x))]);
__syncthreads();
ptx_cp_async_legacy_pred_ca_4_4_4(_ptr_1, ((((i + 3) & 3) * 16) +
((int)threadIdx.x)), B_ptr, (((((int)threadIdx.x) * 14) + i) + 2), (i < 12));
@@ -310,13 +304,15 @@ extern "C" __global__ void __launch_bounds__(16)
main_kernel(float* __restrict__
}
tvm_builtin_ptx_cp_async_wait_group_async_wait_group_2();
__syncthreads();
+ int cse_v2 = (((int)threadIdx.x) + 16);
C_ptr[((((int)threadIdx.x) * 16) + 13)] = (_ptr[(((int)threadIdx.x) + 16)] +
_ptr_1[(((int)threadIdx.x) + 16)]);
tvm_builtin_ptx_cp_async_wait_group_async_wait_group_1();
__syncthreads();
+ int cse_v3 = (((int)threadIdx.x) + 32);
C_ptr[((((int)threadIdx.x) * 16) + 14)] = (_ptr[(((int)threadIdx.x) + 32)] +
_ptr_1[(((int)threadIdx.x) + 32)]);
tvm_builtin_ptx_cp_async_wait_group_async_wait_group_0();
__syncthreads();
- int cse_v5 = (((int)threadIdx.x) + 48);
+ int cse_v4 = (((int)threadIdx.x) + 48);
C_ptr[((((int)threadIdx.x) * 16) + 15)] = (_ptr[(((int)threadIdx.x) + 48)] +
_ptr_1[(((int)threadIdx.x) + 48)]);
}
@@ -904,14 +900,14 @@ def test_multiplication_nodes_are_inlined():
A_flattened = T.decl_tensor((4096,), "float16", data=A.data)
A_shared = T.decl_tensor([4096], "float16", scope="shared")
- T.attr("default", "async_scope", 1)
- for i in range(16):
- cse_v1: T.int64 = T.Cast("int64", i)
- A_shared[T.Ramp(tx * T.int64(128) + cse_v1 * T.int64(8),
T.int64(1), 8)] = (
- A_flattened[T.Ramp(tx * T.int64(128) + cse_v1 *
T.int64(8), T.int64(1), 8)]
- )
- T.ptx.cp.async_.commit_group()
- T.ptx.cp.async_.wait_group(0)
+ with Ts.async_copy_scope():
+ for i in range(16):
+ cse_v1: T.int64 = T.Cast("int64", i)
+ A_shared[T.Ramp(tx * T.int64(128) + cse_v1 * T.int64(8),
T.int64(1), 8)] = (
+ A_flattened[T.Ramp(tx * T.int64(128) + cse_v1 *
T.int64(8), T.int64(1), 8)]
+ )
+ T.ptx.cp.async_.commit_group()
+ T.ptx.cp.async_.wait_group(0)
@I.ir_module
class Expected:
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
b/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
index d413ab84db..419edea422 100644
---
a/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
+++
b/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
@@ -1203,9 +1203,9 @@ def test_simple_compute_async():
with Ts.sblock():
Ts.reads(A[tx, 0])
Ts.writes(B[T.FloorMod(0, 2), tx, 0])
- with T.attr(0, "async_commit_queue_scope", 0):
- with T.attr(0, "async_scope", 1):
- B[T.FloorMod(0, 2), tx, 0] = A[tx, 0] *
T.float32(2)
+ with Ts.async_copy_scope():
+ B[T.FloorMod(0, 2), tx, 0] = A[tx, 0] * T.float32(2)
+ Ts.async_commit(0)
with Ts.sblock():
Ts.reads(A[tx, 1:16], B[0:2, tx, 0])
Ts.writes(B[0:2, tx, 0], C[tx, 0:15])
@@ -1214,22 +1214,20 @@ def test_simple_compute_async():
Ts.where(i + 1 < 16)
Ts.reads(A[tx, i + 1])
Ts.writes(B[(i + 1) % 2, tx, 0])
- with T.attr(0, "async_commit_queue_scope", 0):
- with T.attr(0, "async_scope", 1):
- B[(i + 1) % 2, tx, 0] = A[tx, i + 1] *
T.float32(2)
+ with Ts.async_copy_scope():
+ B[(i + 1) % 2, tx, 0] = A[tx, i + 1] *
T.float32(2)
+ Ts.async_commit(0)
with Ts.sblock():
Ts.where(i + 1 - 1 < 16)
Ts.reads(B[(i - 1 + 1) % 2, tx, 0])
Ts.writes(C[tx, i - 1 + 1])
- with T.attr(0, "async_wait_queue_scope", 0):
- with T.attr(0, "async_wait_inflight_count", 1):
- C[tx, i - 1 + 1] = B[(i - 1 + 1) % 2, tx,
0] + T.float32(1)
+ Ts.async_wait(0, 1)
+ C[tx, i - 1 + 1] = B[(i - 1 + 1) % 2, tx, 0] +
T.float32(1)
with Ts.sblock():
Ts.reads(B[T.FloorMod(15, 2), tx, 0])
Ts.writes(C[tx, 15])
- with T.attr(0, "async_wait_queue_scope", 0):
- with T.attr(0, "async_wait_inflight_count", 0):
- C[tx, 15] = B[T.FloorMod(15, 2), tx, 0] +
T.float32(1)
+ Ts.async_wait(0, 0)
+ C[tx, 15] = B[T.FloorMod(15, 2), tx, 0] + T.float32(1)
tvm.ir.assert_structural_equal(mod["main"], ref.with_attr("global_symbol",
"main"), True)
@@ -1255,9 +1253,9 @@ def test_simple_compute_async():
Ts.where(i < 16)
Ts.reads(A[tx, i])
Ts.writes(B[i % 4, tx, 0])
- T.attr(0, "async_commit_queue_scope", 0)
- T.attr(0, "async_scope", 1)
- B[i % 4, tx, 0] = A[tx, i] * T.float32(2)
+ with Ts.async_copy_scope():
+ B[i % 4, tx, 0] = A[tx, i] * T.float32(2)
+ Ts.async_commit(0)
with Ts.sblock():
Ts.reads(A[tx, 3:16], B[0:4, tx, 0])
Ts.writes(B[0:4, tx, 0], C[tx, 0:13])
@@ -1266,16 +1264,15 @@ def test_simple_compute_async():
Ts.where(i + 3 < 16)
Ts.reads(A[tx, i + 3])
Ts.writes(B[(i + 3) % 4, tx, 0])
- T.attr(0, "async_commit_queue_scope", 0)
- T.attr(0, "async_scope", 1)
- B[(i + 3) % 4, tx, 0] = A[tx, i + 3] * T.float32(2)
+ with Ts.async_copy_scope():
+ B[(i + 3) % 4, tx, 0] = A[tx, i + 3] *
T.float32(2)
+ Ts.async_commit(0)
with Ts.sblock():
Ts.where(i + 3 - 3 < 16)
Ts.reads(B[0:4, tx, 0])
Ts.writes(C[tx, i - 3 + 3])
- with T.attr(0, "async_wait_queue_scope", 0):
- with T.attr(0, "async_wait_inflight_count", 3):
- C[tx, i - 3 + 3] = B[(i - 3 + 3) % 4, tx,
0] + T.float32(1)
+ Ts.async_wait(0, 3)
+ C[tx, i - 3 + 3] = B[(i - 3 + 3) % 4, tx, 0] +
T.float32(1)
with Ts.sblock():
Ts.reads(B[0:4, tx, 0])
Ts.writes(C[tx, 13:16])
@@ -1284,9 +1281,8 @@ def test_simple_compute_async():
Ts.where(i + 16 - 3 < 16)
Ts.reads(B[0:4, tx, 0])
Ts.writes(C[tx, i - 3 + 16])
- with T.attr(0, "async_wait_queue_scope", 0):
- with T.attr(0, "async_wait_inflight_count", 2
- i):
- C[tx, i - 3 + 16] = B[(i - 3 + 16) % 4,
tx, 0] + T.float32(1)
+ Ts.async_wait(0, 2 - i)
+ C[tx, i - 3 + 16] = B[(i - 3 + 16) % 4, tx, 0] +
T.float32(1)
tvm.ir.assert_structural_equal(mod["main"], ref.with_attr("global_symbol",
"main"), True)
@@ -1347,11 +1343,11 @@ def test_async_producer_interleaving():
Ts.where(i < 16)
Ts.reads(A[tx, i], B[tx, i])
Ts.writes(A_shared[i % 4, tx, 0], B_shared[i % 4,
tx, 0])
- with T.attr(0, "async_commit_queue_scope", 0):
- with T.attr(0, "async_scope", 1):
- A_shared[i % 4, tx, 0] = A[tx, i]
- with T.attr(0, "async_scope", 1):
- B_shared[i % 4, tx, 0] = B[tx, i]
+ with Ts.async_copy_scope():
+ A_shared[i % 4, tx, 0] = A[tx, i]
+ with Ts.async_copy_scope():
+ B_shared[i % 4, tx, 0] = B[tx, i]
+ Ts.async_commit(0)
with Ts.sblock():
Ts.reads(A[tx, 3:16], A_shared[0:4, tx, 0], B_shared[0:4,
tx, 0], B[tx, 3:16])
Ts.writes(A_shared[0:4, tx, 0], C[tx, 0:13], B_shared[0:4,
tx, 0])
@@ -1360,26 +1356,24 @@ def test_async_producer_interleaving():
Ts.where(i + 3 < 16)
Ts.reads(A[tx, i + 3])
Ts.writes(A_shared[(i + 3) % 4, tx, 0])
- with T.attr(0, "async_commit_queue_scope", 0):
- with T.attr(0, "async_scope", 1):
- A_shared[(i + 3) % 4, tx, 0] = A[tx, i + 3]
+ with Ts.async_copy_scope():
+ A_shared[(i + 3) % 4, tx, 0] = A[tx, i + 3]
+ Ts.async_commit(0)
with Ts.sblock():
Ts.where(i + 3 - 3 < 16)
Ts.reads(A_shared[0:4, tx, 0], B_shared[0:4, tx,
0])
Ts.writes(C[tx, i - 3 + 3])
- with T.attr(0, "async_wait_queue_scope", 0):
- with T.attr(0, "async_wait_inflight_count", 5):
- C[tx, i - 3 + 3] = (
- A_shared[(i - 3 + 3) % 4, tx, 0]
- + B_shared[(i - 3 + 3) % 4, tx, 0]
- )
+ Ts.async_wait(0, 5)
+ C[tx, i - 3 + 3] = (
+ A_shared[(i - 3 + 3) % 4, tx, 0] + B_shared[(i
- 3 + 3) % 4, tx, 0]
+ )
with Ts.sblock():
Ts.where(i + 3 < 16)
Ts.reads(B[tx, i + 3])
Ts.writes(B_shared[(i + 3) % 4, tx, 0])
- with T.attr(0, "async_commit_queue_scope", 0):
- with T.attr(0, "async_scope", 1):
- B_shared[(i + 3) % 4, tx, 0] = B[tx, i + 3]
+ with Ts.async_copy_scope():
+ B_shared[(i + 3) % 4, tx, 0] = B[tx, i + 3]
+ Ts.async_commit(0)
with Ts.sblock():
Ts.reads(A_shared[0:4, tx, 0], B_shared[0:4, tx, 0])
Ts.writes(C[tx, 13:16])
@@ -1388,12 +1382,11 @@ def test_async_producer_interleaving():
Ts.where(i + 16 - 3 < 16)
Ts.reads(A_shared[0:4, tx, 0], B_shared[0:4, tx,
0])
Ts.writes(C[tx, i - 3 + 16])
- with T.attr(0, "async_wait_queue_scope", 0):
- with T.attr(0, "async_wait_inflight_count", 2
- i):
- C[tx, i - 3 + 16] = (
- A_shared[(i - 3 + 16) % 4, tx, 0]
- + B_shared[(i - 3 + 16) % 4, tx, 0]
- )
+ Ts.async_wait(0, 2 - i)
+ C[tx, i - 3 + 16] = (
+ A_shared[(i - 3 + 16) % 4, tx, 0]
+ + B_shared[(i - 3 + 16) % 4, tx, 0]
+ )
tvm.ir.assert_structural_equal(mod["main"], ref.with_attr("global_symbol",
"main"), True)
@@ -1423,20 +1416,17 @@ def test_three_stage_compute_two_stage_async():
Ts.where(i < 16)
Ts.reads(A[tx, i])
Ts.writes(B[i % 2, tx, 0])
- with T.attr(0, "async_commit_queue_scope", 0):
- with T.attr(0, "async_scope", 1):
- B[i % 2, tx, 0] = A[tx, i] * T.float32(2)
+ with Ts.async_copy_scope():
+ B[i % 2, tx, 0] = A[tx, i] * T.float32(2)
+ Ts.async_commit(0)
with Ts.sblock():
Ts.where(i == 1 and i - 1 < 16)
Ts.reads(B[(i - 1) % 2, tx, 0])
Ts.writes(C[(i - 1) % 2, tx, 0])
- with T.attr(0, "async_commit_queue_scope", 1):
- with T.attr(0, "async_wait_queue_scope", 0):
- with T.attr(0,
"async_wait_inflight_count", 1):
- with T.attr(0, "async_scope", 1):
- C[(i - 1) % 2, tx, 0] = B[
- (i - 1) % 2, tx, 0
- ] + T.float32(2)
+ Ts.async_wait(0, 1)
+ with Ts.async_copy_scope():
+ C[(i - 1) % 2, tx, 0] = B[(i - 1) % 2, tx, 0]
+ T.float32(2)
+ Ts.async_commit(1)
with Ts.sblock():
Ts.reads(A[tx, 2:16], B[0:2, tx, 0], C[0:2, tx, 0])
Ts.writes(B[0:2, tx, 0], C[0:2, tx, 0], D[tx, 0:14])
@@ -1445,27 +1435,23 @@ def test_three_stage_compute_two_stage_async():
Ts.where(i + 2 < 16)
Ts.reads(A[tx, i + 2])
Ts.writes(B[(i + 2) % 2, tx, 0])
- with T.attr(0, "async_commit_queue_scope", 0):
- with T.attr(0, "async_scope", 1):
- B[(i + 2) % 2, tx, 0] = A[tx, i + 2] *
T.float32(2)
+ with Ts.async_copy_scope():
+ B[(i + 2) % 2, tx, 0] = A[tx, i + 2] *
T.float32(2)
+ Ts.async_commit(0)
with Ts.sblock():
Ts.where(i + 2 - 1 < 16)
Ts.reads(B[(i - 1 + 2) % 2, tx, 0])
Ts.writes(C[(i - 1 + 2) % 2, tx, 0])
- with T.attr(0, "async_commit_queue_scope", 1):
- with T.attr(0, "async_wait_queue_scope", 0):
- with T.attr(0,
"async_wait_inflight_count", 1):
- with T.attr(0, "async_scope", 1):
- C[(i - 1 + 2) % 2, tx, 0] = B[
- (i - 1 + 2) % 2, tx, 0
- ] + T.float32(2)
+ Ts.async_wait(0, 1)
+ with Ts.async_copy_scope():
+ C[(i - 1 + 2) % 2, tx, 0] = B[(i - 1 + 2) % 2,
tx, 0] + T.float32(2)
+ Ts.async_commit(1)
with Ts.sblock():
Ts.where(i + 2 - 2 < 16)
Ts.reads(C[0:2, tx, 0])
Ts.writes(D[tx, i - 2 + 2])
- with T.attr(0, "async_wait_queue_scope", 1):
- with T.attr(0, "async_wait_inflight_count", 1):
- D[tx, i - 2 + 2] = C[(i - 2 + 2) % 2, tx,
0] + T.float32(1)
+ Ts.async_wait(1, 1)
+ D[tx, i - 2 + 2] = C[(i - 2 + 2) % 2, tx, 0] +
T.float32(1)
with Ts.sblock():
Ts.reads(B[0:2, tx, 0], C[0:2, tx, 0])
Ts.writes(C[0:2, tx, 0], D[tx, 14:16])
@@ -1474,24 +1460,18 @@ def test_three_stage_compute_two_stage_async():
Ts.where(i + 16 - 1 < 16)
Ts.reads(B[(i - 1 + 16) % 2, tx, 0])
Ts.writes(C[(i - 1 + 16) % 2, tx, 0])
- with T.attr(0, "async_commit_queue_scope", 1):
- with T.attr(0, "async_wait_queue_scope", 0):
- with T.attr(0,
"async_wait_inflight_count", 0 - i):
- with T.attr(0, "async_scope", 1):
- C[(i - 1 + 16) % 2, tx, 0] = B[
- (i - 1 + 16) % 2, tx, 0
- ] + T.float32(2)
+ Ts.async_wait(0, 0 - i)
+ with Ts.async_copy_scope():
+ C[(i - 1 + 16) % 2, tx, 0] = B[(i - 1 + 16) %
2, tx, 0] + T.float32(
+ 2
+ )
+ Ts.async_commit(1)
with Ts.sblock():
Ts.where(i + 16 - 2 < 16)
Ts.reads(C[0:2, tx, 0])
Ts.writes(D[tx, i - 2 + 16])
- with T.attr(0, "async_wait_queue_scope", 1):
- with T.attr(
- 0,
- "async_wait_inflight_count",
- T.if_then_else(i + 16 - 1 < 16, 1, 0),
- ):
- D[tx, i - 2 + 16] = C[(i - 2 + 16) % 2,
tx, 0] + T.float32(1)
+ Ts.async_wait(1, T.if_then_else(i + 16 - 1 < 16,
1, 0))
+ D[tx, i - 2 + 16] = C[(i - 2 + 16) % 2, tx, 0] +
T.float32(1)
tvm.ir.assert_structural_equal(mod["main"], ref.with_attr("global_symbol",
"main"), True)
@@ -1577,19 +1557,21 @@ def test_async_pipelined_mma_gemm_simple():
pipeline =
mod["main"].body.block.body.body.body.body.block.body[1].block.body
prologue, body, epilogue = pipeline
- commit_queue_scope = prologue.block.body.body.block.body
- assert len(commit_queue_scope.body) == 2
- assert commit_queue_scope.value == 0
+ commit_group = prologue.block.body.body.block.body
+ assert len(commit_group) == 3
+ assert commit_group[-1].value.args[0] == 0
- commit_queue_scope = body.block.body.body[0].block.body
- assert len(commit_queue_scope.body) == 2
- assert commit_queue_scope.value == 0
+ commit_group = body.block.body.body[0].block.body
+ assert len(commit_group) == 3
+ assert commit_group[-1].value.args[0] == 0
- assert body.block.body.body[1].block.body.body.attr_key ==
"async_wait_inflight_count"
- assert body.block.body.body[1].block.body.body.value == 3
+ assert
body.block.body.body[1].block.body[0].value.op.same_as(tvm.ir.Op.get("s_tir.async_wait"))
+ assert body.block.body.body[1].block.body[0].value.args[1] == 3
- assert epilogue.block.body.body.block.body.body.attr_key ==
"async_wait_inflight_count"
- assert str(epilogue.block.body.body.block.body.body.value) == "2 - k_0_0"
+ assert epilogue.block.body.body.block.body[0].value.op.same_as(
+ tvm.ir.Op.get("s_tir.async_wait")
+ )
+ assert str(epilogue.block.body.body.block.body[0].value.args[1]) == "2 -
k_0_0"
build_and_run(sch)
@@ -1622,21 +1604,25 @@ def
test_async_nested_pipeline_mma_gemm_ideal_annotation():
pipeline =
mod["main"].body.block.body.body.body.body.block.body[1].block.body
prologue, body, epilogue = pipeline
- commit_queue_scope = prologue.block.body.body[0].block.body
- assert len(commit_queue_scope.body) == 2
- assert commit_queue_scope.value == 0
+ commit_group = prologue.block.body.body[0].block.body
+ assert len(commit_group) == 3
+ assert commit_group[-1].value.args[0] == 0
- assert prologue.block.body.body[1].block.body.body.attr_key ==
"async_wait_inflight_count"
- assert prologue.block.body.body[1].block.body.body.value == 2
+ assert (
+ prologue.block.body.body[1]
+ .block.body[0]
+ .value.op.same_as(tvm.ir.Op.get("s_tir.async_wait"))
+ )
+ assert prologue.block.body.body[1].block.body[0].value.args[1] == 2
- commit_queue_scope = body.block.body.body[0].block.body
- assert len(commit_queue_scope.body) == 2
- assert commit_queue_scope.value == 0
+ commit_group = body.block.body.body[0].block.body
+ assert len(commit_group) == 3
+ assert commit_group[-1].value.args[0] == 0
- assert body.block.body.body[1].block.body.body.attr_key ==
"async_wait_inflight_count"
- assert body.block.body.body[1].block.body.body.value == 2
+ assert
body.block.body.body[1].block.body[0].value.op.same_as(tvm.ir.Op.get("s_tir.async_wait"))
+ assert body.block.body.body[1].block.body[0].value.args[1] == 2
- assert str(epilogue.block.body.body[0].block.body.body.value) == "1 -
k_0_0"
+ assert str(epilogue.block.body.body[0].block.body[0].value.args[1]) == "1
- k_0_0"
build_and_run(sch)
