This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tqchen/printer-refinement in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 10b4a14e9a000ac17a31dd200b413c103a1b87b2 Author: Tianqi Chen <[email protected]> AuthorDate: Thu Oct 1 22:09:03 2026 +0000 [FIX][TIRx] Preserve storage synchronization calls in named syntax --- python/tvm/tirx/op.py | 21 ++++++++++--- src/tirx/script/printer/expr.cc | 35 ++++++++++++++++++---- .../transform/test_s_tir_transform_thread_sync.py | 9 ++---- 3 files changed, 49 insertions(+), 16 deletions(-) diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py index d583d8930f..4eff21210a 100644 --- a/python/tvm/tirx/op.py +++ b/python/tvm/tirx/op.py @@ -833,7 +833,7 @@ def tvm_thread_invariant(cond): return call_intrin(_primexpr_ty(cond), "tirx.tvm_thread_invariant", cond) -def tvm_storage_sync(storage_scope, is_load=False, num_blocks=-1): +def tvm_storage_sync(storage_scope, is_load=False, num_blocks=-1, *, dtype="void"): """Perform synchronization in specified scope. Parameters @@ -841,18 +841,31 @@ def tvm_storage_sync(storage_scope, is_load=False, num_blocks=-1): storage_scope : str The storage scope to perform synchronization. - is_load : bool + is_load : bool or Expr or None Whether to perform load synchronization. (for global sync only) + Set both ``is_load`` and ``num_blocks`` to None to omit these operands. - num_blocks : int + num_blocks : int or Expr or None The number of blocks to synchronize. (for global sync only) + Set to None to omit this operand. + + dtype : str or tvm.ir.Type + The stored result type. Defaults to void. Returns ------- call : Expr The call expression. """ - return call_intrin("void", "tirx.tvm_storage_sync", storage_scope, is_load, num_blocks) + args = [storage_scope] + if is_load is None: + if num_blocks is not None: + raise ValueError("num_blocks must be None when is_load is omitted") + else: + args.append(is_load) + if num_blocks is not None: + args.append(num_blocks) + return call_intrin(dtype, "tirx.tvm_storage_sync", *args) def tvm_kernel_replace_point(): diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc index 45c4437ccf..87f07f2fd5 100644 --- a/src/tirx/script/printer/expr.cc +++ b/src/tirx/script/printer/expr.cc @@ -181,6 +181,35 @@ bool CanTranslateExplicitResultCall(const CallNode* call) { }); } +ffi::Optional<ExprDoc> StorageSyncDocTranslate(DocTranslatorObj* d, ffi::AnyView input, + const ffi::Object*) { + const auto* call = + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const CallNode>(input); + if (!CanTranslateExplicitResultCall(call) || call->args.empty() || call->args.size() > 3) { + return RawCall(d, call, false); + } + ffi::Array<ExprDoc> args; + for (const Expr& arg : call->args) { + args.push_back(MaterializeCallArgument(d, arg, d->Translate(arg).value())); + } + // Explicit None operands suppress the helper's defaults without changing + // the stored argument list of native one- and two-operand calls. + while (args.size() < 3) args.push_back(LiteralDoc::None(std::nullopt)); + ffi::Array<ffi::String> keys; + ffi::Array<ExprDoc> values; + if (!ffi::StructuralEqual()(call->ty, PrimType::Void())) { + keys.push_back("dtype"); + values.push_back(TypeValue(d, call->ty)); + } + return NamespaceDoc("tirx")->Attr("tvm_storage_sync")->Call(args, keys, values); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + OpDef("tirx.tvm_storage_sync") + .set_attr<FDocTranslate>(kOpCallDocTranslate, + FDocTranslate::FromNative<&StorageSyncDocTranslate>()); +} + ffi::Optional<ExprDoc> CallExternDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { const auto* call = @@ -446,12 +475,6 @@ ffi::Optional<ExprDoc> TIRCallPrefixDocTranslate(DocTranslatorObj* d, const Call return RawCall(d, call, false); } } - // The named storage-sync helper supplies two default operands. A shorter - // native Call must keep its actual operand list and explicit stored type. - if (call->op.same_as(tirx::builtin::tvm_storage_sync()) && - (call->args.size() != 3 || !ffi::StructuralEqual()(call->ty, PrimType::Void()))) { - return RawCall(d, call, false); - } if (auto op = call->op.as<Op>(); op && op.value()->name == "tirx.isnan") { auto input_type = call->args.size() == 1 ? call->args[0]->ty.as<PrimType>() : std::nullopt; // The named helper folds constants and widens float16 operands, and its diff --git a/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py b/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py index 93587b8022..5d5a52e5f6 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py @@ -59,7 +59,7 @@ def test_sync_read_thread_id_independent_location(): result_local[0] = result_local[0] + temp_shared[0] * p1[1] mod = run_passes(func) - assert 'I.Call("tirx.tvm_storage_sync", ["shared"], ty="int32")' in str(mod) + assert 'T.tvm_storage_sync("shared", None, None, dtype="int32")' in str(mod) def test_sync_inside_condition(): @@ -101,7 +101,7 @@ def test_sync_inside_condition(): for func in (func2, func3): mod = s_tir.transform.ThreadSync("shared")(tvm.IRModule.from_expr(func)) - assert 'I.Call("tirx.tvm_storage_sync", ["shared"], ty="int32")' in mod.script() + assert 'T.tvm_storage_sync("shared", None, None, dtype="int32")' in mod.script() def test_sync_shared_dyn(): @@ -165,10 +165,7 @@ def test_sync_shared_aliasing_buffer_views(): # In addition to the explicit write-to-read barrier, the shared physical # storage needs a read-to-next-write barrier across loop iterations. - assert ( - str(mod["main"]).count("T.tvm_storage_sync") - + str(mod["main"]).count('I.Call("tirx.tvm_storage_sync"') - ) == 2 + assert str(mod["main"]).count("T.tvm_storage_sync") == 2 @pytest.mark.gpu
