Answer a cross-boundary copy with a Memcpy, or refuse it

dace/libraries/standard/nodes/copy/select.py

@@ -54,8 +54,19 @@ def select_copy_implementation(node: "CopyLibraryNode", parent_state: dace.SDFGS
             return 'MemcpyCUDA1D'
         return 'Tasklet'
 
-    # cudaMemcpyAsync can't issue from device code, so in-kernel multi-element copies map instead.
+    # cudaMemcpyAsync can't issue from device code, so in-kernel multi-element copies map instead --
+    # but only where the kernel can dereference both ends. Crossing the boundary it cannot: there is
+    # no implementation of a host-to-device copy issued from inside a kernel, neither a mapped
+    # tasklet nor a Memcpy variant, so the answer is that the copy does not belong in the kernel.
+    # Saying so here keeps the refusal at the point the choice is made, instead of returning an
+    # implementation that then reports the boundary as its own limitation and suggests a Memcpy
+    # variant that is just as impossible in device code.
     if is_devicelevel_gpu(parent_state.sdfg, parent_state, node):
+        if _is_cross_cpu_gpu(inp.storage, out.storage, node, parent_state):
+            raise ValueError(f"No copy implementation crosses the CPU/GPU boundary inside a kernel "
+                             f"(got {inp.storage} -> {out.storage} for '{node.label}' in state "
+                             f"'{parent_state.label}'). Device code cannot address host memory and "
+                             f"cannot issue a Memcpy; place the copy outside the kernel.")
         return 'MappedTasklet'
 
     # Host CPU-resident: same-shape/contiguous/same-layout below the parallel-transfer threshold
@@ -71,10 +82,16 @@ def select_copy_implementation(node: "CopyLibraryNode", parent_state: dace.SDFGS
                      and not is_in_parallel_scope(node, parent_state))):
         return 'MemcpyCPU'
 
+    # Anything that crosses the boundary is a Memcpy: no mapped tasklet can dereference both ends,
+    # so a copy that gets here strided or rank-mismatched belongs to the refinement below, which
+    # answers with the pitched 2-D form or a loop of ``cudaMemcpyAsync`` per contiguous chunk.
+    # ``Register`` is outside ``allowed`` yet reaches here at host level, where it IS host memory.
     gpu = dtypes.StorageType.GPU_Global
     allowed = CPU_RESIDENT_STORAGES | {dtypes.StorageType.Default, gpu}
-    impl = ('MemcpyCUDA1D' if ((inp.storage == gpu or out.storage == gpu) and inp.storage in allowed
-                               and out.storage in allowed) else None)
+    crosses_boundary = _is_cross_cpu_gpu(inp.storage, out.storage, node, parent_state)
+    both_gpu_or_host = inp.storage in allowed and out.storage in allowed
+    impl = ('MemcpyCUDA1D' if
+            (crosses_boundary or ((inp.storage == gpu or out.storage == gpu) and both_gpu_or_host)) else None)
 
     if impl == 'MemcpyCUDA1D':
         refined = _refine_cuda_impl_for_subsets(node, parent_state)

tests/library/copy_node_test.py

@@ -2000,5 +2000,61 @@ def test_symbolic_extent_expansions_keep_their_ranges_symbolic():
                         f'{expanded.name}: C++ spelling leaked into a memlet subset: {e.data.subset}'
 
 
+def _make_in_kernel_copy_sdfg(src_storage: dace.dtypes.StorageType,
+                              dst_storage: dace.dtypes.StorageType) -> Tuple[dace.SDFG, CopyLibraryNode]:
+    """A multi-element ``CopyLibraryNode`` sitting inside a ``GPU_Device`` map."""
+    sdfg = dace.SDFG("in_kernel_copy")
+    sdfg.add_array("src", [4, 8], dace.float64, storage=src_storage)
+    sdfg.add_array("dst", [4, 8], dace.float64, storage=dst_storage)
+    state = sdfg.add_state()
+    entry, exit_ = state.add_map("kern", dict(i="0:4"), schedule=dace.dtypes.ScheduleType.GPU_Device)
+    libnode = CopyLibraryNode(name="cp")
+    state.add_memlet_path(state.add_read("src"),
+                          entry,
+                          libnode,
+                          dst_conn=CopyLibraryNode.INPUT_CONNECTOR_NAME,
+                          memlet=dace.memlet.Memlet("src[0:4, 0:8]"))
+    state.add_memlet_path(libnode,
+                          exit_,
+                          state.add_write("dst"),
+                          src_conn=CopyLibraryNode.OUTPUT_CONNECTOR_NAME,
+                          memlet=dace.memlet.Memlet("dst[0:4, 0:8]"))
+    return sdfg, libnode
+
+
+def test_in_kernel_cross_boundary_copy_is_refused_by_selection():
+    """Device code can neither address host memory nor issue a Memcpy, so no implementation fits.
+
+    Auto used to answer ``MappedTasklet`` for every in-kernel multi-element copy, boundary
+    unchecked; the expansion then rejected it and pointed at ``MemcpyCUDA1D``, which device code
+    cannot issue either. The refusal belongs where the choice is made, and it has to name the
+    kernel, because moving the copy out of it is the fix.
+    """
+    sdfg, libnode = _make_in_kernel_copy_sdfg(dace.dtypes.StorageType.CPU_Heap, dace.dtypes.StorageType.GPU_Global)
+    state = sdfg.start_state
+    with pytest.raises(ValueError, match="inside a kernel"):
+        select_copy_implementation(libnode, state)
+
+
+def test_in_kernel_device_to_device_copy_still_maps():
+    """The refusal is only for the boundary -- an in-kernel device copy has nothing else to be."""
+    sdfg, libnode = _make_in_kernel_copy_sdfg(dace.dtypes.StorageType.GPU_Global, dace.dtypes.StorageType.GPU_Global)
+    assert select_copy_implementation(libnode, sdfg.start_state) == "MappedTasklet"
+
+
+def test_a_host_level_cross_boundary_copy_never_falls_back_to_a_mapped_tasklet():
+    """``Register`` is outside the storage set the CUDA branch tested, so it reached the fallback.
+
+    At host level a ``Register`` endpoint IS host memory, so the copy crosses the boundary and has
+    to be a ``cudaMemcpy`` -- a mapped tasklet cannot dereference the device side, and answering
+    with one only moves the failure into the expansion.
+    """
+    sdfg, libnode = _make_copy_sdfg(
+        _ArraySpec(shape=(4, 8), storage=dace.dtypes.StorageType.Register, transient=True),
+        _ArraySpec(shape=(4, 8), storage=dace.dtypes.StorageType.GPU_Global, transient=True),
+    )
+    assert select_copy_implementation(libnode, sdfg.start_state) == "MemcpyCUDA1D"
+
+
 if __name__ == "__main__":
     pytest.main([__file__])