Skip to content

[S-TIR][CUDA] Fix legacy predicated cp.async zero fill #19741

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tqchen merged 2 commits into from
Jun 13, 2026
Merged

[S-TIR][CUDA] Fix legacy predicated cp.async zero fill #19741

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8896e22
Fix legacy predicated cp.async zero fill
tlopex Jun 11, 2026
93efb06
[TIRX][S-TIR] Make CSE planning deterministic via insertion-ordered t…
tlopex Jun 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 12 additions & 2 deletions python/tvm/tirx/operator/intrinsics/cuda/cp_async.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,7 +205,8 @@ def codegen_ptx_cp_async(*args):
bytes; offsets are pre-scaled by the pass) and the call is
forwarded with default cache / predicate / fill_mode.
* 6 args ``(dst_ptr, dst_offset, src_ptr, src_offset, cp_size,
predicate)`` — same as 5-arg form with an explicit predicate.
predicate)`` — same as 5-arg form with an explicit predicate,
zero-filling the destination when the predicate is false.
* 8 args ``(dst_ptr, src_ptr, cp_size, cache_policy, has_cache_hint,
prefetch_size, predicate, fill_mode)`` — the fork-native wrapper
API.
Expand Down Expand Up @@ -269,6 +270,14 @@ def _scale(n):
func_name = (
f"ptx_cp_async_legacy_pred_{ca_or_cg}_{cp_size_v}_{dst_elem_bytes}_{src_elem_bytes}"
)
if cp_size_v == 4:
zero_fill = ' " @!p st.shared.u32 [%0], {%4};\\n"\n'
elif cp_size_v == 8:
zero_fill = ' " @!p st.shared.v2.u32 [%0], {%4, %4};\\n"\n'
elif cp_size_v == 16:
zero_fill = ' " @!p st.shared.v4.u32 [%0], {%4, %4, %4, %4};\\n"\n'
else:
raise ValueError(f"unsupported legacy predicated cp.async size: {cp_size_v}")
body = (
f" uint8_t* dst_p = (uint8_t*)dst + dst_off{dst_scale};\n"
f" uint8_t* src_p = (uint8_t*)src + src_off{src_scale};\n"
Expand All @@ -279,8 +288,9 @@ def _scale(n):
' " setp.eq.u32 p, %3, 1;\\n"\n'
f' " @p cp.async.{ca_or_cg}.shared.global'
' [%0], [%1], %2;\\n"\n'
f"{zero_fill}"
' "}\\n"\n'
f' :: "r"(dst_addr), "l"(src_p), "n"({cp_size_v}), "r"(predicate)\n'
f' :: "r"(dst_addr), "l"(src_p), "n"({cp_size_v}), "r"(predicate), "r"(0)\n'
" );"
)
source_code = (
Expand Down
26 changes: 17 additions & 9 deletions src/tirx/transform/common_subexpr_elim.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@
#include <utility>
#include <vector>

#include "../../support/ordered_map.h"
#include "../analysis/check_contains.h"

namespace tvm {
Expand Down Expand Up @@ -239,8 +240,16 @@ class CSEPlanner : public StmtExprVisitor {
int consumed{0};
};

/*! \brief Expression table keyed by structural equality (ExprDeepEqual). */
using ExprTable = std::unordered_map<PrimExpr, ExprEntry, ffi::StructuralHash, ExprDeepEqual>;
/*!
* \brief Expression table keyed by structural equality (ExprDeepEqual).
*
* An insertion-ordered map so that iteration visits entries in discovery
* (program) order. This makes the plan — and hence cse_v numbering —
* deterministic. A plain unordered_map iterates in hash order, and
* StructuralHash hashes free variables by object identity, which varies
* between processes (ASLR).
*/
using ExprTable = support::OrderedMap<PrimExpr, ExprEntry, ffi::StructuralHash, ExprDeepEqual>;

// ------------------------------------------------------------------
// Eligibility predicates
Expand Down Expand Up @@ -592,8 +601,9 @@ class CSEPlanner : public StmtExprVisitor {
* \brief Convert the accumulated expression table into InsertBefore + ExprRemap tables.
*
* Algorithm (shallower-first with repr propagation):
* 1. Collect all entries and sort by expr_depth ascending (shallower first),
* with structural hash as tie-breaker for determinism.
* 1. Collect all entries and sort by expr_depth ascending (shallower first).
* The stable sort over the insertion-ordered table keeps entries of
* equal depth in discovery (program) order, so the plan is deterministic.
* 2. Compute independent occurrence counts from the DAG children.
* For each parent P with count >= 2, its children's consumed counts
* are incremented by `(P.count - 1) * multiplicity` (the Bind value
Expand All @@ -610,7 +620,8 @@ class CSEPlanner : public StmtExprVisitor {
* \return A pair of (InsertBeforeTable, ExprRemapTable).
*/
std::pair<InsertBeforeTable, ExprRemapTable> ComputePlan() {
// Step 1: Sort entries by depth ascending (shallower first), hash for determinism
// Step 1: Sort entries by depth ascending (shallower first). table_ iterates
// in discovery order, which the stable sort preserves among equal depths.
std::vector<std::pair<PrimExpr, ExprEntry*>> all_entries;
for (auto& kv : table_) {
all_entries.push_back({kv.first, &kv.second});
Expand All @@ -619,10 +630,7 @@ class CSEPlanner : public StmtExprVisitor {
std::stable_sort(
all_entries.begin(), all_entries.end(),
[](const std::pair<PrimExpr, ExprEntry*>& a, const std::pair<PrimExpr, ExprEntry*>& b) {
if (a.second->expr_depth != b.second->expr_depth)
return a.second->expr_depth < b.second->expr_depth;
ffi::StructuralHash hasher;
return hasher(a.first) < hasher(b.first);
return a.second->expr_depth < b.second->expr_depth;
});

// Step 2: Compute consumed counts in ExprEntry from the DAG.
Expand Down
211 changes: 83 additions & 128 deletions tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -194,19 +194,8 @@ def test_inject_async_copy_shared_dyn():
# `ptx_mbarrier_*` family instead.


# Note: the expected output contains a dead CSE variable `cse_v1 = (i < 12)`.
# CSE extracts (i < 12) before inject_ptx_async_copy runs, but the latter
# replaces the original IfThenElse guards with new cast(int32, i < 12)
# expressions for predicated async copies, leaving cse_v1 unused.
expected_cuda_script = r"""#include <cuda.h>
__forceinline__ __device__ unsigned int
cast_smem_ptr_to_int(const void* const smem_ptr)
{
unsigned int smem_int;
asm volatile ("{ .reg .u64 smem_int; cvta.to.shared.u64 smem_int, %1; cvt.u32.u64 %0, smem_int; }"
: "=r"(smem_int) : "l"(smem_ptr));
return smem_int;
}
#endif

#if (((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 4)) || \
(__CUDACC_VER_MAJOR__ > 11))
Expand All @@ -215,132 +204,101 @@ def test_inject_async_copy_shared_dyn():
#define TVM_ENABLE_L2_PREFETCH 0
#endif

#ifdef __CUDACC_RTC__
using int64_t = long long;
using uint64_t = unsigned long long;
#ifdef _WIN32
using uint = unsigned int;
using uchar = unsigned char;
using ushort = unsigned short;
using int64_t = long long;
using uint64_t = unsigned long long;
#else
#include <cstdint>
#define uint unsigned int
#define uchar unsigned char
#define ushort unsigned short
#endif
using uint = unsigned int;
using uchar = unsigned char;
using ushort = unsigned short;

extern "C" __global__ void __launch_bounds__(16) main_kernel(float* __restrict__ A, float* __restrict__ B, float* __restrict__ C);
extern "C" __global__ void __launch_bounds__(16) main_kernel(float* __restrict__ A, float* __restrict__ B, float* __restrict__ C) {
__shared__ float A_shared[64];
__shared__ float B_shared[64];
A_shared[((int)threadIdx.x)] = 0x0p+0f/*0.000000e+00*/;
B_shared[((int)threadIdx.x)] = 0x0p+0f/*0.000000e+00*/;
__asm__ __volatile__("cp.async.commit_group;");


{
unsigned int addr = cast_smem_ptr_to_int(A_shared + (((int)threadIdx.x) + 16));
__asm__ __volatile__(
#if TVM_ENABLE_L2_PREFETCH
"cp.async.ca.shared.global.L2::128B [%0], [%1], %2;"
#else
"cp.async.ca.shared.global [%0], [%1], %2;"
#endif
:: "r"(addr), "l"((void*)(A + (((int)threadIdx.x) * 14))), "n"(4)
);
}

{
unsigned int addr = cast_smem_ptr_to_int(B_shared + (((int)threadIdx.x) + 16));
__asm__ __volatile__(
#if TVM_ENABLE_L2_PREFETCH
"cp.async.ca.shared.global.L2::128B [%0], [%1], %2;"
#else
"cp.async.ca.shared.global [%0], [%1], %2;"
#endif
:: "r"(addr), "l"((void*)(B + (((int)threadIdx.x) * 14))), "n"(4)
);
}
__asm__ __volatile__("cp.async.commit_group;");


{
unsigned int addr = cast_smem_ptr_to_int(A_shared + (((int)threadIdx.x) + 32));
__asm__ __volatile__(
#if TVM_ENABLE_L2_PREFETCH
"cp.async.ca.shared.global.L2::128B [%0], [%1], %2;"
#else
"cp.async.ca.shared.global [%0], [%1], %2;"
#endif
:: "r"(addr), "l"((void*)(A + ((((int)threadIdx.x) * 14) + 1))), "n"(4)
);
}
__forceinline__ __device__ void tvm_builtin_ptx_cp_async_wait_group_0() {
asm volatile("cp.async.wait_group 0;");
}

{
unsigned int addr = cast_smem_ptr_to_int(B_shared + (((int)threadIdx.x) + 32));
__asm__ __volatile__(
#if TVM_ENABLE_L2_PREFETCH
"cp.async.ca.shared.global.L2::128B [%0], [%1], %2;"
#else
"cp.async.ca.shared.global [%0], [%1], %2;"
#endif
:: "r"(addr), "l"((void*)(B + ((((int)threadIdx.x) * 14) + 1))), "n"(4)
);
}
__asm__ __volatile__("cp.async.commit_group;");
__forceinline__ __device__ void tvm_builtin_ptx_cp_async_wait_group_1() {
asm volatile("cp.async.wait_group 1;");
}

for (int i = 0; i < 13; ++i) {
bool cse_v1 = (i < 12);

{
unsigned int addr = cast_smem_ptr_to_int(A_shared + ((((i + 3) & 3) * 16) + ((int)threadIdx.x)));
int pred_guard = (int)(i < 12);
__asm__ __volatile__(
"{ .reg .pred p;"
" setp.ne.b32 p, %0, 0;"
#if TVM_ENABLE_L2_PREFETCH
" @p cp.async.ca.shared.global.L2::128B [%1], [%2], %3;"
#else
" @p cp.async.ca.shared.global [%1], [%2], %3;"
#endif
" @!p st.shared.u32 [%1], {%4};}"
:: "r"(pred_guard), "r"(addr), "l"((void*)(A + (((((int)threadIdx.x) * 14) + i) + 2))), "n"(4), "r"(0)
);
}
__asm__ __volatile__("cp.async.commit_group;");
__forceinline__ __device__ void tvm_builtin_ptx_cp_async_wait_group_2() {
asm volatile("cp.async.wait_group 2;");
}

__forceinline__ __device__ void tvm_builtin_ptx_cp_async_wait_group_5() {
asm volatile("cp.async.wait_group 5;");
}

__forceinline__ __device__ void ptx_cp_async_legacy_pred_ca_4_4_4(void* dst, int dst_off, void* src, int src_off, int predicate) {
uint8_t* dst_p = (uint8_t*)dst + dst_off * 4;
uint8_t* src_p = (uint8_t*)src + src_off * 4;
unsigned int dst_addr = __cvta_generic_to_shared(dst_p);
__asm__ __volatile__(
"{\n"
" .reg .pred p;\n"
" setp.eq.u32 p, %3, 1;\n"
" @p cp.async.ca.shared.global [%0], [%1], %2;\n"
" @!p st.shared.u32 [%0], {%4};\n"
"}\n"
:: "r"(dst_addr), "l"(src_p), "n"(4), "r"(predicate), "r"(0)
);
}

__asm__ __volatile__("cp.async.wait_group 5;");
__forceinline__ __device__ void ptx_cp_async_legacy_ca_4_4_4(void* dst, int dst_off, void* src, int src_off) {
uint8_t* dst_p = (uint8_t*)dst + dst_off * 4;
uint8_t* src_p = (uint8_t*)src + src_off * 4;
unsigned int dst_addr = __cvta_generic_to_shared(dst_p);
asm volatile("cp.async.ca.shared.global [%0], [%1], %2;"
:: "r"(dst_addr), "l"(src_p), "n"(4));
}

__forceinline__ __device__ void tvm_builtin_ptx_cp_async_commit_group() {
asm volatile("cp.async.commit_group;");
}
extern "C" __global__ void __launch_bounds__(16) main_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr, float* __restrict__ C_ptr);
extern "C" __global__ void __launch_bounds__(16) main_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr, float* __restrict__ C_ptr) {
__shared__ alignas(64) float A_shared_ptr[64];
__shared__ alignas(64) float B_shared_ptr[64];
A_shared_ptr[((int)threadIdx.x)] = 0x0p+0f/*0.000000e+00*/;
B_shared_ptr[((int)threadIdx.x)] = 0x0p+0f/*0.000000e+00*/;
tvm_builtin_ptx_cp_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(A_shared_ptr, (((int)threadIdx.x) + 16), A_ptr, (((int)threadIdx.x) * 14));
ptx_cp_async_legacy_ca_4_4_4(B_shared_ptr, (((int)threadIdx.x) + 16), B_ptr, (((int)threadIdx.x) * 14));
tvm_builtin_ptx_cp_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(A_shared_ptr, (((int)threadIdx.x) + 32), A_ptr, ((((int)threadIdx.x) * 14) + 1));
ptx_cp_async_legacy_ca_4_4_4(B_shared_ptr, (((int)threadIdx.x) + 32), B_ptr, ((((int)threadIdx.x) * 14) + 1));
tvm_builtin_ptx_cp_async_commit_group();
int cse_v4 = (((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(A_shared_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();
tvm_builtin_ptx_cp_async_wait_group_5();
__syncthreads();
C[((((int)threadIdx.x) * 16) + i)] = (A_shared[(((i & 3) * 16) + ((int)threadIdx.x))] + B_shared[(((i & 3) * 16) + ((int)threadIdx.x))]);
int cse_v8 = (((i & 3) * 16) + ((int)threadIdx.x));
C_ptr[((((int)threadIdx.x) * 16) + i)] = (A_shared_ptr[(((i & 3) * 16) + ((int)threadIdx.x))] + B_shared_ptr[(((i & 3) * 16) + ((int)threadIdx.x))]);
__syncthreads();

{
unsigned int addr = cast_smem_ptr_to_int(B_shared + ((((i + 3) & 3) * 16) + ((int)threadIdx.x)));
int pred_guard = (int)(i < 12);
__asm__ __volatile__(
"{ .reg .pred p;"
" setp.ne.b32 p, %0, 0;"
#if TVM_ENABLE_L2_PREFETCH
" @p cp.async.ca.shared.global.L2::128B [%1], [%2], %3;"
#else
" @p cp.async.ca.shared.global [%1], [%2], %3;"
#endif
" @!p st.shared.u32 [%1], {%4};}"
:: "r"(pred_guard), "r"(addr), "l"((void*)(B + (((((int)threadIdx.x) * 14) + i) + 2))), "n"(4), "r"(0)
);
}
__asm__ __volatile__("cp.async.commit_group;");

ptx_cp_async_legacy_pred_ca_4_4_4(B_shared_ptr, ((((i + 3) & 3) * 16) + ((int)threadIdx.x)), B_ptr, (((((int)threadIdx.x) * 14) + i) + 2), (i < 12));
tvm_builtin_ptx_cp_async_commit_group();
}
__asm__ __volatile__("cp.async.wait_group 2;");

tvm_builtin_ptx_cp_async_wait_group_2();
__syncthreads();
C[((((int)threadIdx.x) * 16) + 13)] = (A_shared[(((int)threadIdx.x) + 16)] + B_shared[(((int)threadIdx.x) + 16)]);
__asm__ __volatile__("cp.async.wait_group 1;");

C_ptr[((((int)threadIdx.x) * 16) + 13)] = (A_shared_ptr[(((int)threadIdx.x) + 16)] + B_shared_ptr[(((int)threadIdx.x) + 16)]);
tvm_builtin_ptx_cp_async_wait_group_1();
__syncthreads();
C[((((int)threadIdx.x) * 16) + 14)] = (A_shared[(((int)threadIdx.x) + 32)] + B_shared[(((int)threadIdx.x) + 32)]);
__asm__ __volatile__("cp.async.wait_group 0;");

C_ptr[((((int)threadIdx.x) * 16) + 14)] = (A_shared_ptr[(((int)threadIdx.x) + 32)] + B_shared_ptr[(((int)threadIdx.x) + 32)]);
tvm_builtin_ptx_cp_async_wait_group_0();
__syncthreads();
C[((((int)threadIdx.x) * 16) + 15)] = (A_shared[(((int)threadIdx.x) + 48)] + B_shared[(((int)threadIdx.x) + 48)]);
int cse_v5 = (((int)threadIdx.x) + 48);
C_ptr[((((int)threadIdx.x) * 16) + 15)] = (A_shared_ptr[(((int)threadIdx.x) + 48)] + B_shared_ptr[(((int)threadIdx.x) + 48)]);
}

"""
Expand Down Expand Up @@ -392,9 +350,6 @@ def tvm_callback_cuda_postproc(code, _):
tvm.register_global_func(func_name, prev_postproc, override=True)


# TODO(tlopex): fix CSE determinism (change unordered map to ordered map) and
# remove this xfail; see #19741.
@pytest.mark.xfail
@tvm.testing.requires_cuda
def test_cp_async_in_if_then_else(postproc_if_missing_async_support):
@T.prim_func(s_tir=True)
Expand Down
Loading