Skip to content

[TIR][CUDA] Require aligned addresses for cp.async #20913

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

Open
0z5a wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[TIR][CUDA] Require aligned addresses for cp.async #20913

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
29 changes: 27 additions & 2 deletions src/s_tir/transform/inject_ptx_async_copy.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op/memory.h>

#include <numeric>
#include <unordered_map>
#include <utility>
#include <vector>
Expand Down Expand Up @@ -132,9 +133,31 @@ class PTXAsyncCopyInjector : public StmtExprMutator {
index_factor = static_cast<int>(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<TensorVar>(), src_offset,
static_cast<int>(src_elem_type.StorageBytes())) &&
aligned(store->dest.as_or_throw<TensorVar>(), dst_offset,
static_cast<int>(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<Stmt>(store));
}
ffi::Array<Expr> args = {
store->dest.as_or_throw<TensorVar>().data(), mul(dst_offset, PrimExpr(index_factor)),
load->source.as_or_throw<tvm::tirx::TensorVar>().data(), src_offset, PrimExpr(bytes)};
Expand Down Expand Up @@ -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<TensorVar>().dtype(), ptx_cp_async_op,
{store->dest.as_or_throw<TensorVar>().data(),
Expand Down Expand Up @@ -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<TensorVar>().dtype(), ptx_cp_async_op,
{store->dest.as_or_throw<TensorVar>().data(),
Expand Down
47 changes: 47 additions & 0 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 @@ -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()