From f7ffe5ea4d3e1f70cf02ffc71a41e1a4df1816ab Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 4 Nov 2026 17:05:07 +0100 Subject: [PATCH] [TIR][CUDA] Require aligned addresses for cp.async --- src/s_tir/transform/inject_ptx_async_copy.cc | 29 +++++++++++- ...t_s_tir_transform_inject_ptx_async_copy.py | 47 +++++++++++++++++++ 2 files changed, 74 insertions(+), 2 deletions(-) diff --git a/src/s_tir/transform/inject_ptx_async_copy.cc b/src/s_tir/transform/inject_ptx_async_copy.cc index 38b2e0acef22..153ea4b0c1ef 100644 --- a/src/s_tir/transform/inject_ptx_async_copy.cc +++ b/src/s_tir/transform/inject_ptx_async_copy.cc @@ -33,6 +33,7 @@ #include #include +#include #include #include #include @@ -132,9 +133,31 @@ class PTXAsyncCopyInjector : public StmtExprMutator { index_factor = static_cast(src_elem_type.StorageBytes()); } + // cp.async requires both source and destination addresses to be + // aligned to the transaction size. The offsets are in element units; + // index_factor accounts for the byte buffer used by merged shared memory. + auto aligned = [this, bytes](const TensorVar& tensor, const PrimExpr& offset, + int offset_unit_bytes) { + int base_alignment = tensor.ty()->data_alignment; + if (base_alignment < bytes || base_alignment % bytes != 0) return false; + const int modulus = bytes / std::gcd(bytes, offset_unit_bytes); + sym::ModularSet indices = analyzer_->modular_set(offset); + return indices->coeff % modulus == 0 && indices->base % modulus == 0; + }; + auto can_copy = [&](const PrimExpr& src_offset, const PrimExpr& dst_offset) { + return aligned(load->source.as_or_throw(), src_offset, + static_cast(src_elem_type.StorageBytes())) && + aligned(store->dest.as_or_throw(), dst_offset, + static_cast(dst_elem_type.StorageBytes()) * index_factor); + }; + if (indices_lanes == 1) { auto src_offset = load->indices[0]; auto dst_offset = store->indices[0]; + if (!can_copy(src_offset, dst_offset)) { + return StmtExprMutator::Mutate_(store, InplaceMode::kDisallow) + .ValueOrUnchanged(ffi::GetRef(store)); + } ffi::Array args = { store->dest.as_or_throw().data(), mul(dst_offset, PrimExpr(index_factor)), load->source.as_or_throw().data(), src_offset, PrimExpr(bytes)}; @@ -172,7 +195,8 @@ class PTXAsyncCopyInjector : public StmtExprMutator { } return std::nullopt; }(); - if (src_offset.has_value() && dst_offset.has_value()) { + if (src_offset.has_value() && dst_offset.has_value() && + can_copy(src_offset.value(), dst_offset.value())) { static const Op ptx_cp_async_op = Op::Get("tirx.s_tir.cp_async_raw"); return Evaluate(Call(store->dest.as_or_throw().dtype(), ptx_cp_async_op, {store->dest.as_or_throw().data(), @@ -206,7 +230,8 @@ class PTXAsyncCopyInjector : public StmtExprMutator { return std::nullopt; }(); - if (src_offset.has_value() && dst_offset.has_value()) { + if (src_offset.has_value() && dst_offset.has_value() && + can_copy(src_offset.value(), dst_offset.value())) { static const Op ptx_cp_async_op = Op::Get("tirx.s_tir.cp_async_raw"); return Evaluate(Call(store->dest.as_or_throw().dtype(), ptx_cp_async_op, {store->dest.as_or_throw().data(), 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 792fdea019af..38d03e7d8b39 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 @@ -943,5 +943,52 @@ def main(A: T.Tensor((32, 128), "float16")): tvm_tirx.ir.assert_structural_equal(After, Expected) +@pytest.mark.parametrize( + "src_stride,dst_stride,src_align,dst_align,expect_async", + [ + (16, 16, 16, 16, 1), + (17, 16, 16, 16, 0), + (16, 17, 16, 16, 0), + (17, 17, 16, 16, 0), + (16, 16, 4, 16, 0), + (16, 16, 16, 4, 0), + ], +) +def test_inject_async_copy_requires_alignment( + src_stride, dst_stride, src_align, dst_align, expect_async +): + @Ts.function + def copy( + A: T.Tensor((32, src_stride), "float32", data_alignment=src_align), + B: T.Tensor((32, 4), "float32"), + ) -> None: + T.func_attr({"global_symbol": "main", "tirx.noalias": True}) + T.launch_thread("blockIdx.x", 1) + tx = T.launch_thread("threadIdx.x", 32) + with Ts.sblock(): + shared = Ts.sblock_alloc_tensor( + [32, dst_stride], "float32", scope="shared", align=dst_align + ) + Ts.reads(A[0:32, 0:4]) + Ts.writes(B[0:32, 0:4]) + with Ts.async_copy_scope(): + shared[tx, T.Ramp(0, 1, 4)] = A[tx, T.Ramp(0, 1, 4)] + T.evaluate(T.ptx.cp.async_.commit_group()) + T.evaluate(T.ptx.cp.async_.wait_group(0)) + for j in range(4): + B[tx, j] = shared[tx, j] + + mod = tvm_tirx.IRModule.from_expr(copy) + for transform in ( + tvm.s_tir.transform.LowerOpaqueBlock(), + tvm.s_tir.transform.LowerThreadBinding(), + tvm_tirx.tirx.transform.FlattenTensor(), + tvm_tirx.tirx.transform.VectorizeLoop(), + tvm.s_tir.transform.InjectPTXAsyncCopy(), + ): + mod = transform(mod) + assert count_cp_async(mod["main"].body) == expect_async + + if __name__ == "__main__": tvm_tirx.testing.main()