diff --git a/src/s_tir/transform/inject_ptx_async_copy.cc b/src/s_tir/transform/inject_ptx_async_copy.cc index 38b2e0acef22..160c097a2ab7 100644 --- a/src/s_tir/transform/inject_ptx_async_copy.cc +++ b/src/s_tir/transform/inject_ptx_async_copy.cc @@ -150,22 +150,25 @@ class PTXAsyncCopyInjector : public StmtExprMutator { // Predicated load don't support vectorized indexing. if (!predicated) { // Only some vectorized indexing patterns are supported for now. - auto src_offset = [=]() -> ffi::Optional { - if (load->indices[0]->IsInstance()) { - return load->indices[0].as()->base; + auto src_offset = [this, load]() -> ffi::Optional { + const auto* ramp = load->indices[0].as(); + if (ramp && analyzer_->CanProveEqual(ramp->stride, 1)) { + return ramp->base; } return std::nullopt; }(); - auto dst_offset = [=]() -> ffi::Optional { - if (store->indices[0].as()) { - return store->indices[0].as()->base; + auto dst_offset = [this, store]() -> ffi::Optional { + if (const auto* ramp = store->indices[0].as()) { + if (analyzer_->CanProveEqual(ramp->stride, 1)) return ramp->base; + return std::nullopt; } else if (store->indices[0].as()) { // The case where the dst tensor is a byte tensor generated by merging dynamic // shared memory. // A_shared.dyn[(ramp(...), 1, 8) + x8(17408))] = A_global[ramp(...),1, 8)] auto* add = store->indices[0].as(); - if (!add->a->IsInstance()) return std::nullopt; + const auto* ramp = add->a.as(); + if (!ramp || !analyzer_->CanProveEqual(ramp->stride, 1)) return std::nullopt; if (!add->b->IsInstance()) return std::nullopt; return prim::Add(add->a.as()->base, add->b.as()->value); @@ -183,22 +186,25 @@ class PTXAsyncCopyInjector : public StmtExprMutator { } } else { // Only some vectorized indexing patterns are supported for now. - auto src_offset = [=]() -> ffi::Optional { - if (load->indices[0]->IsInstance()) { - return load->indices[0].as()->base; + auto src_offset = [this, load]() -> ffi::Optional { + const auto* ramp = load->indices[0].as(); + if (ramp && analyzer_->CanProveEqual(ramp->stride, 1)) { + return ramp->base; } return std::nullopt; }(); - auto dst_offset = [=]() -> ffi::Optional { - if (store->indices[0].as()) { - return store->indices[0].as()->base; + auto dst_offset = [this, store]() -> ffi::Optional { + if (const auto* ramp = store->indices[0].as()) { + if (analyzer_->CanProveEqual(ramp->stride, 1)) return ramp->base; + return std::nullopt; } else if (store->indices[0].as()) { // The case where the dst tensor is a byte tensor generated by merging dynamic // shared memory. // A_shared.dyn[(ramp(...), 1, 8) + x8(17408))] = A_global[ramp(...),1, 8)] auto* add = store->indices[0].as(); - if (!add->a->IsInstance()) return std::nullopt; + const auto* ramp = add->a.as(); + if (!ramp || !analyzer_->CanProveEqual(ramp->stride, 1)) return std::nullopt; if (!add->b->IsInstance()) return std::nullopt; return prim::Add(add->a.as()->base, add->b.as()->value); diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy_stride.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy_stride.py new file mode 100644 index 000000000000..6427aaa6abbf --- /dev/null +++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy_stride.py @@ -0,0 +1,93 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import numpy as np +import pytest +import tvm_ffi + +import tvm +import tvm_tirx +from tvm.s_tir import script as Ts +from tvm_tirx.script import tirx as T +from tvm_tirx.testing import env + + +def make_copy(src_stride, dst_stride): + @Ts.function + def copy(A: T.Tensor((32, 16), "float32"), B: T.Tensor((32, 4), "float32")): + 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, 16], "float32", scope="shared") + Ts.reads(A[0:32, 0:16]) + Ts.writes(B[0:32, 0:4]) + with Ts.async_copy_scope(): + shared[tx, T.Ramp(0, dst_stride, 4)] = A[tx, T.Ramp(0, src_stride, 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 * dst_stride] + + return copy + + +@pytest.mark.parametrize( + "src_stride,dst_stride,expected_count", + [(1, 1, 1), (2, 1, 0), (1, 2, 0), (2, 2, 0)], +) +def test_async_copy_requires_contiguous_vectors(src_stride, dst_stride, expected_count): + mod = tvm_tirx.IRModule.from_expr(make_copy(src_stride, dst_stride)) + 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) + + calls = [] + + def visit(node): + if isinstance(node, tvm_tirx.ir.Call) and node.op.name == "tirx.s_tir.cp_async_raw": + calls.append(node) + + tvm_ffi.structural_walk(mod["main"].body, visit) + assert len(calls) == expected_count + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") +@pytest.mark.parametrize("src_stride,dst_stride", [(2, 1), (1, 2), (2, 2)]) +def test_noncontiguous_async_copy_values(src_stride, dst_stride): + with tvm_tirx.ir.transform.PassContext(config={"tirx.use_async_copy": 1}): + exe = tvm.compile( + tvm_tirx.IRModule.from_expr(make_copy(src_stride, dst_stride)), target="cuda" + ) + + data = np.arange(32 * 16, dtype="float32").reshape(32, 16) + expected = data[:, ::src_stride][:, :4] + + def run_and_check(): + dev = tvm_tirx.cuda() + source = tvm_tirx.runtime.tensor(data, dev) + result = tvm_tirx.runtime.tensor(np.full((32, 4), -1, dtype="float32"), dev) + exe(source, result) + np.testing.assert_array_equal(result.numpy(), expected) + + tvm_tirx.testing.run_with_gpu_lock(run_and_check)