fix(runtime): address each cp.async vector in the destination's own layout

include/tilefoundry/runtime/cuda/ops/copy/copy_impl.h

@@ -62,12 +62,12 @@ struct CopyAsync {
             same_slice_size<decltype(s), decltype(d)>(),
             "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<cute::uint_bit_t<bytes * 8>>(s);
+        auto d_v = cute::recast<cute::uint_bit_t<bytes * 8>>(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<TSrc>,
                       "ops::copy_async: cp.async requires sm_80 or newer");

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

tests/ops/tir/test_copy_async.py

@@ -17,14 +17,15 @@
 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
 from tilefoundry.ir.tir.prim_function import PrimFunction
 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)