diff --git a/include/tilefoundry/runtime/cuda/ops/copy/copy_impl.h b/include/tilefoundry/runtime/cuda/ops/copy/copy_impl.h index 94f04296a..dcf4e5761 100644 --- a/include/tilefoundry/runtime/cuda/ops/copy/copy_impl.h +++ b/include/tilefoundry/runtime/cuda/ops/copy/copy_impl.h @@ -62,12 +62,12 @@ struct CopyAsync { same_slice_size(), "ops::copy_async: the two projected slices must hold the same " "number of elements"); - constexpr int V = bytes / int(sizeof(value_type)); #if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800) auto s_v = cute::recast>(s); + auto d_v = cute::recast>(d); const int nv = int(cute::size(s_v)); for (int i = 0; i < nv; ++i) - __pipeline_memcpy_async(&d(i * V), &s_v(i), bytes); + __pipeline_memcpy_async(&d_v(i), &s_v(i), bytes); #else static_assert(dependent_false_v, "ops::copy_async: cp.async requires sm_80 or newer"); diff --git a/tests/fixtures/tir/split_stage.py b/tests/fixtures/tir/split_stage.py new file mode 100644 index 000000000..1c35e8bcd --- /dev/null +++ b/tests/fixtures/tir/split_stage.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from tilefoundry.dsl import * +from tilefoundry.target import CpuTarget, CudaTarget + + +@module( + entry="split_stage_host", + target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("thread", 128),), +) +class SplitStage: + @prim_func(target=CudaTarget("nvidia.h200_sxm")) + def split_stage_device(a: Tensor[(2, 128, 4), "f32"], b: Tensor[(2, 128, 4), "f32"]): + with Mesh((Topology("thread", 128),), layout=(128,), names=("t",)) as m: + a_view = T.tensor_view(T.ptr_of(a), layout=((2, 128 @ m.t, 4), (512, 4, 1))) + shared = T.alloc_tensor(tensor_type=Tensor[(2, 128, 4), "f32", (2, 128, 4), "smem"]) + shared_view = T.tensor_view(T.ptr_of(shared), layout=((2, 128 @ m.t, 4), (512, 4, 1))) + b_view = T.tensor_view(T.ptr_of(b), layout=((2, 128 @ m.t, 4), (512, 4, 1))) + T.copy_async(a_view, shared_view) + T.cp_async_commit() + T.cp_async_wait(n=0) + T.sync(m) + T.copy(shared_view, b_view) + + @prim_func(target=CpuTarget()) + def split_stage_host(a: Tensor[(2, 128, 4), "f32"], b: Tensor[(2, 128, 4), "f32"]): + launch(split_stage_device, a, b, grid=( + 1, + 1, + 1, + ), block=( + 128, + 1, + 1, + )) # noqa: F821 diff --git a/tests/ops/tir/test_copy_async.py b/tests/ops/tir/test_copy_async.py index 34dcedf9c..f72ece973 100644 --- a/tests/ops/tir/test_copy_async.py +++ b/tests/ops/tir/test_copy_async.py @@ -17,6 +17,7 @@ import tilefoundry import tilefoundry.codegen.cuda # noqa: F401 — trigger emitter autodiscovery from tests.fixtures.tir.async_sync import AsyncStage +from tests.fixtures.tir.split_stage import SplitStage from tests.ops.tir.targets import CUDA_TARGETS, retarget from tilefoundry.ir.core import Var, VerifyError from tilefoundry.ir.tir.async_copy import CopyAsync, CpAsyncWait @@ -24,7 +25,7 @@ from tilefoundry.ir.tir.stmts import Evaluate, MeshScope, Return, Sequential from tilefoundry.ir.types import DType, Layout, Mesh, TensorType, Topology, make_tensor_type from tilefoundry.target import CudaTarget -from tilefoundry.visitor_registry.verify import verify_prim_function +from tilefoundry.visitor_registry.verify import verify_module, verify_prim_function def _copy_async_pf(src_ty: TensorType, dst_ty: TensorType) -> PrimFunction: @@ -110,3 +111,18 @@ def test_async_copy_stage_matches_input(target: CudaTarget) -> None: runtime(source, out) torch.cuda.synchronize() torch.testing.assert_close(out, source, rtol=0, atol=0) + + +@pytest.mark.parametrize("target", CUDA_TARGETS) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_async_copy_stages_two_separate_runs_per_thread(target: CudaTarget) -> None: + """Each thread's second run lands 512 elements after its first, as the destination states.""" + module = retarget(SplitStage, target) + verify_module(module) + runtime = tilefoundry.compile(module, target=target) + torch.manual_seed(4) + source = torch.randn(2, 128, 4, dtype=torch.float32, device="cuda") + out = torch.empty_like(source) + runtime(source, out) + torch.cuda.synchronize() + torch.testing.assert_close(out, source, rtol=0, atol=0)