diff --git a/dace/codegen/codegen.py b/dace/codegen/codegen.py index 78a9499f50..b8a614244d 100644 --- a/dace/codegen/codegen.py +++ b/dace/codegen/codegen.py @@ -204,6 +204,16 @@ def generate_code(sdfg: SDFG, validate=True) -> List[CodeObject]: # Set default storage/schedule types in SDFG infer_types.set_default_schedule_and_storage_types(sdfg, None) + # Give every implicit copy a node of its own, before the expansion below lowers it. An implicit + # copy is a write no node performs, so an empty memlet ordering that write has nothing to point + # at and the copy is free to move ahead of a write it must follow. + if config.Config.get_bool('compiler', 'cpu', 'explicit_copy'): + from dace.transformation.passes.insert_explicit_copies import InsertExplicitCopies + InsertExplicitCopies().apply_pass(sdfg, {}) + # The nodes just inserted carry no inferred connector types, and the expansion below reads + # them to decide pointer vs. value. Storage defaults above already hold. + infer_types.infer_connector_types(sdfg) + # Recursively expand library nodes that have not yet been expanded sdfg.expand_library_nodes() diff --git a/dace/codegen/instrumentation/gpu_tx_markers.py b/dace/codegen/instrumentation/gpu_tx_markers.py index 7377fd042e..a3b87d4ce8 100644 --- a/dace/codegen/instrumentation/gpu_tx_markers.py +++ b/dace/codegen/instrumentation/gpu_tx_markers.py @@ -152,6 +152,24 @@ def on_copy_end(self, sdfg: SDFG, cfg: ControlFlowRegion, state: SDFGState, src_ return self.print_range_pop(local_stream) + def on_node_begin(self, sdfg: SDFG, cfg: ControlFlowRegion, state: SDFGState, node: nodes.Node, + outer_stream: CodeIOStream, inner_stream: CodeIOStream, global_stream: CodeIOStream) -> None: + if not isinstance(node, nodes.CodeNode) or node.instrument != dtypes.InstrumentationType.GPU_TX_MARKERS: + return + if is_devicelevel_gpu_kernel(sdfg, state, node): + # Don't instrument device code + return + self.print_range_push(node.label, sdfg, outer_stream) + + def on_node_end(self, sdfg: SDFG, cfg: ControlFlowRegion, state: SDFGState, node: nodes.Node, + outer_stream: CodeIOStream, inner_stream: CodeIOStream, global_stream: CodeIOStream) -> None: + if not isinstance(node, nodes.CodeNode) or node.instrument != dtypes.InstrumentationType.GPU_TX_MARKERS: + return + if is_devicelevel_gpu_kernel(sdfg, state, node): + # Don't instrument device code + return + self.print_range_pop(outer_stream) + def on_scope_entry(self, sdfg: SDFG, cfg: ControlFlowRegion, state: SDFGState, node: nodes.EntryNode, outer_stream: CodeIOStream, inner_stream: CodeIOStream, global_stream: CodeIOStream) -> None: if node.map.instrument != dtypes.InstrumentationType.GPU_TX_MARKERS: diff --git a/dace/codegen/targets/cpp.py b/dace/codegen/targets/cpp.py index eae90d882c..8546f274fe 100644 --- a/dace/codegen/targets/cpp.py +++ b/dace/codegen/targets/cpp.py @@ -813,11 +813,13 @@ def unparse_cr(sdfg, wcr_ast, dtype): def connected_to_gpu_memory(node: nodes.Node, state: SDFGState, sdfg: SDFG): + # Both ends of the path count: a host tasklet that only WRITES GPU memory needs the stream just + # as much as one that reads it. Same rule as the stream-retention walk in ``cuda.py``. for e in state.all_edges(node): path = state.memlet_path(e) - if ((isinstance(path[0].src, nodes.AccessNode) - and path[0].src.desc(sdfg).storage is dtypes.StorageType.GPU_Global)): - return True + for endpoint in (path[0].src, path[-1].dst): + if isinstance(endpoint, nodes.AccessNode) and endpoint.desc(sdfg).storage is dtypes.StorageType.GPU_Global: + return True return False @@ -915,7 +917,17 @@ def unparse_tasklet(sdfg, cfg, state_id, dfg, node, function_stream, callsite_st gpu_codegen = next(cg for cg in codegen._dispatcher.used_targets if isinstance(cg, cuda.CUDACodeGen)) except StopIteration: return - synchronize_streams(sdfg, cfg, state_dfg, state_id, node, node, callsite_stream, gpu_codegen) + # The tasklet's own code names the stream through the local defined above, so the + # synchronization it may need must name the same expression. + synchronize_streams(sdfg, + cfg, + state_dfg, + state_id, + node, + node, + callsite_stream, + gpu_codegen, + stream_expr='__dace_current_stream') return body = node.code.code @@ -1346,10 +1358,12 @@ def presynchronize_streams(sdfg: SDFG, cfg: ControlFlowRegion, dfg: StateSubgrap # TODO: This should be in the CUDA code generator. Add appropriate conditions to node dispatch predicate -def synchronize_streams(sdfg, cfg, dfg, state_id, node, scope_exit, callsite_stream, codegen): +def synchronize_streams(sdfg, cfg, dfg, state_id, node, scope_exit, callsite_stream, codegen, stream_expr=None): # Post-kernel stream synchronization (with host or other streams) max_streams = int(Config.get("compiler", "cuda", "max_concurrent_streams")) - if max_streams >= 0: + if stream_expr is not None: + cudastream = stream_expr + elif max_streams >= 0: cudastream = common.gpu_stream_expr(node._cuda_stream) else: # Only default stream is used cudastream = 'nullptr' @@ -1393,8 +1407,30 @@ def synchronize_streams(sdfg, cfg, dfg, state_id, node, scope_exit, callsite_str if max_streams >= 0 and hasattr(node, "_cuda_stream"): backend = common.get_gpu_backend() + synced_host = False for edge in dfg.out_edges(scope_exit): + # A host-located destination is read by plain host code as soon as the asynchronous + # work is issued -- a kernel launch packing a by-value argument counts -- and neither + # events nor consumer stream stamps order the host. Wait on the issuing stream once + # (copy-edge analog: ``_emit_copy`` in cuda.py). + hostnode = edge.dst + while (isinstance(hostnode, nodes.AccessNode) and hostnode.data is not None + and isinstance(sdfg.arrays[hostnode.data], data.View)): + hostnode = dfg.out_edges(hostnode)[0].dst + if (isinstance(hostnode, nodes.AccessNode) and hostnode.data is not None + and sdfg.arrays[hostnode.data].storage + not in (dtypes.StorageType.GPU_Global, dtypes.StorageType.GPU_Shared)): + if not synced_host: + callsite_stream.write( + "DACE_GPU_CHECK(%sStreamSynchronize(%s));" % (backend, cudastream), + cfg, + state_id, + [edge.src, edge.dst], + ) + synced_host = True + continue + if (isinstance(edge.dst, nodes.AccessNode) and hasattr(edge.dst, '_cuda_stream') and edge.dst._cuda_stream != node._cuda_stream): # Stream assignment gives a cross-stream edge its own event. Event 0 belongs to some diff --git a/dace/codegen/targets/cpu.py b/dace/codegen/targets/cpu.py index 9e87fee31b..050da2d00d 100644 --- a/dace/codegen/targets/cpu.py +++ b/dace/codegen/targets/cpu.py @@ -1305,6 +1305,17 @@ def memlet_definition(self, # constexpr arrays if memlet.data in self._frame.symbols_and_constants(sdfg): result += "const {} {} = {};".format(memlet_type, local_name, expr) + elif (var_type == DefinedType.Scalar and isinstance(conntype, dtypes.pointer) + and not isinstance(desc.dtype, dtypes.opaque)): + # Scalar source feeding a pointer-typed connector (e.g. CopyLibraryNode + # -> cudaMemcpyAsync from a host scalar argument). The connector's + # pointer type wins over the source's scalar ctypedef, and the address + # of the variable is what the callee wants; `define_out_memlet` already + # does this on the write side. Skip opaque dtypes (MPI_Comm / + # MPI_Request / GPU handles) -- the value is already a pointer-like + # handle, so address-of adds an indirection the callee rejects + # (``MPI_Bcast`` expects ``MPI_Comm``, not ``MPI_Comm *``). + result += "{}* {} = &{};".format(ctypedef, local_name, expr) else: # Pointer reference result += "{} {} = {};".format(ctypedef, local_name, expr) diff --git a/dace/codegen/targets/cuda.py b/dace/codegen/targets/cuda.py index 4a28efc0f4..7aa3cd2197 100644 --- a/dace/codegen/targets/cuda.py +++ b/dace/codegen/targets/cuda.py @@ -99,6 +99,9 @@ def __init__(self, frame_codegen: 'DaCeCodeGenerator', sdfg: SDFG): self._exitcode = CodeIOStream() self._global_sdfg: SDFG = sdfg self._toplevel_schedule = None + # True while generating an SDFG nested below the one whose schedule established the current + # device-level scope, i.e. no longer the kernel's own state machine. + self._below_toplevel_sdfg = False self._arglists: Dict[nodes.MapEntry, Dict[str, dt.Data]] = {} """ # Keep track of which kernels got a threadBlock map inserted @@ -1633,7 +1636,15 @@ def generate_devicelevel_state(self, sdfg: SDFG, cfg: ControlFlowRegion, state: callsite_stream.write("} // subgraph end", cfg, state.block_id) - callsite_stream.write('__gbar.Sync();', cfg, state.block_id) + # A state machine needs a barrier between its states wherever it runs: a nested SDFG + # with several states (or control flow) below the kernel still synchronizes between + # them. An SDFG that is one lone state has no state transition to order, so below the + # kernel's own SDFG it emits no barrier -- it may run inside a single-thread-guarded + # component, where a barrier is reached by one thread and never releases. Its writes + # are ordered by the enclosing state's own barrier instead. + lone_state = sdfg.number_of_nodes() == 1 and isinstance(sdfg.nodes()[0], SDFGState) + if not (self._below_toplevel_sdfg and lone_state): + callsite_stream.write('__gbar.Sync();', cfg, state.block_id) # done here, code is generated return @@ -1722,9 +1733,12 @@ def generate_scope(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg_scope: StateSub self._in_device_code = oldval self.extra_nsdfg_args.append((desc.as_arg(name=''), inner_name, outer_name)) + # A DefinedType.Pointer registers the POINTER ctype, as every other allocation + # site does; the element type here makes consumers that adopt the defined ctype + # (``emit_memlet_reference``) declare the parameter one indirection short. self._dispatcher.defined_vars.add(inner_name, DefinedType.Pointer, - desc.dtype.ctype, + dtypes.pointer(desc.dtype).ctype, allow_shadowing=True) extra_call_args.append(outer_name) extra_call_args_typed.append(desc.as_arg(name=inner_name)) @@ -2936,9 +2950,13 @@ def _generate_NestedSDFG(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg: StateSub node: nodes.NestedSDFG, function_stream: CodeIOStream, callsite_stream: CodeIOStream) -> None: old_schedule = self._toplevel_schedule + old_below_toplevel = self._below_toplevel_sdfg nested_schedule = get_node_schedule(sdfg, dfg, node) if nested_schedule != dtypes.ScheduleType.Default: self._toplevel_schedule = nested_schedule + # A device-level scope is already open, so this SDFG sits inside one of its components + # rather than being the state machine that scope is made of. + self._below_toplevel_sdfg = old_schedule in dtypes.GPU_SCHEDULES old_codegen = self._cpu_codegen.calling_codegen self._cpu_codegen.calling_codegen = self @@ -2946,6 +2964,7 @@ def _generate_NestedSDFG(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg: StateSub self._cpu_codegen.calling_codegen = old_codegen self._toplevel_schedule = old_schedule + self._below_toplevel_sdfg = old_below_toplevel def _generate_MapExit(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg: StateSubgraphView, state_id: int, node: nodes.MapExit, function_stream: CodeIOStream, callsite_stream: CodeIOStream) -> None: diff --git a/dace/config_schema.yml b/dace/config_schema.yml index 83bcb9458d..fee30470b3 100644 --- a/dace/config_schema.yml +++ b/dace/config_schema.yml @@ -349,6 +349,24 @@ required: generate "#pragma omp parallel sections" code around them. + parallel_transfer_min_elements: + type: int + default: 262144 + title: Parallel copy/fill element threshold + description: > + Minimum element count (symbolic size counts as large) for a + copy/fill to lower to an OpenMP element map instead of memcpy/memset. + + explicit_copy: + type: bool + default: true + title: Lower implicit copies to explicit copy nodes + description: > + If set to true, codegen lifts every implicit copy edge to a + CopyLibraryNode before expansion. An implicit copy is a write no node + performs, so an empty memlet ordering that write has nothing to point + at; giving the copy a node is what makes it orderable. + ############################################# # GPU (CUDA/HIP) compiler cuda: diff --git a/dace/libraries/standard/environments/__init__.py b/dace/libraries/standard/environments/__init__.py index a47c7755f7..92bc55d6d8 100644 --- a/dace/libraries/standard/environments/__init__.py +++ b/dace/libraries/standard/environments/__init__.py @@ -1,2 +1,3 @@ # Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +from .cpu import CPU from .cuda import CUDA diff --git a/dace/libraries/standard/environments/cpu.py b/dace/libraries/standard/environments/cpu.py new file mode 100644 index 0000000000..6f8ab27977 --- /dev/null +++ b/dace/libraries/standard/environments/cpu.py @@ -0,0 +1,23 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""DaCe library environment exposing the C++ standard headers used by CPU-side libnode expansions.""" +import dace.library + + +@dace.library.environment +class CPU: + """Minimal library environment that pulls in ```` for plain CPU expansions.""" + + cmake_minimum_version = None + cmake_packages = [] + cmake_variables = {} + cmake_includes = [] + cmake_libraries = [] + cmake_compile_flags = [] + cmake_link_flags = [] + cmake_files = [] + + headers = {'frame': ["cstring"]} + state_fields = [] + init_code = "" + finalize_code = "" + dependencies = [] diff --git a/dace/libraries/standard/helper.py b/dace/libraries/standard/helper.py new file mode 100644 index 0000000000..2409b229bd --- /dev/null +++ b/dace/libraries/standard/helper.py @@ -0,0 +1,96 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Shared helpers for CopyLibraryNode and FillLibraryNode expansions.""" +from typing import Callable, List, Tuple + +import dace +from dace import dtypes +from dace.sdfg import nodes +from dace.sdfg.scope import is_in_scope + +# Both legacy and experimental codegens consume this exact name for stream wiring. +CURRENT_STREAM_NAME = "__dace_current_stream" + +# Register is intentionally in neither set: resolves by scope (GPU register vs. host stack slot). +GPU_RESIDENT_STORAGES = frozenset({ + dtypes.StorageType.GPU_Global, + dtypes.StorageType.GPU_Shared, +}) +CPU_RESIDENT_STORAGES = frozenset({ + dtypes.StorageType.CPU_Heap, + dtypes.StorageType.CPU_Pinned, + dtypes.StorageType.CPU_ThreadLocal, +}) + + +def collapse_shape_and_strides( + subset: dace.subsets.Range, + strides: List[dace.symbolic.SymExpr]) -> Tuple[List[dace.symbolic.SymExpr], List[dace.symbolic.SymExpr]]: + """Drop length-1 dims from a (subset, strides) pair; surviving strides scale by the subset step. + + A tiled dimension (``b:e:step:tile``) addresses ``tile`` contiguous elements per step, which no + single (length, stride) pair expresses -- it expands into two dims: the step count at + ``stride * step``, then the tile at ``stride``. + + :param subset: The access range, one ``(begin, end, step)`` per dimension. + :param strides: The parent array strides, aligned with ``subset``. + :returns: ``(collapsed_shape, collapsed_strides)`` with singletons removed. + """ + collapsed_shape = [] + collapsed_strides = [] + # ``Range.size()`` already folds the tile in (``tile * ceiling((e + 1 - b) / step)``); dividing + # it back out is exact and avoids re-deriving a per-dim count formula that could drift from it. + for (_, _, s), stride, tile, dim_size in zip(subset, strides, subset.tile_sizes, subset.size()): + length = dim_size / tile + if length != 1: + collapsed_shape.append(length) + collapsed_strides.append(stride * s) + if tile != 1: + collapsed_shape.append(tile) + collapsed_strides.append(stride) + return collapsed_shape, collapsed_strides + + +def is_parallel_cpu_transfer_size(num_elements: dace.symbolic.SymbolicType) -> bool: + """False only when ``num_elements`` is a compile-time constant below + ``compiler.cpu.parallel_transfer_min_elements``; a symbolic (unknown-at-compile-time) size + is assumed large and takes the parallel path too. + + :param num_elements: total contiguous element count (constant or symbolic). + :returns: ``True`` to route to the mapped expansion, ``False`` to keep the single libc call. + """ + threshold = int(dace.Config.get('compiler', 'cpu', 'parallel_transfer_min_elements')) + try: + return int(dace.symbolic.simplify(num_elements)) >= threshold + except (TypeError, ValueError): + return True + + +def is_in_parallel_scope(node: nodes.LibraryNode, parent_state: dace.SDFGState) -> bool: + """True when a multi-threaded map encloses this transfer, so the mapped form would open one + OpenMP region per entry instead of one for the whole transfer. + + ``Default`` counts: an unresolved enclosing map becomes ``CPU_Multicore`` at the top level. + + :param node: the transfer library node. + :param parent_state: state containing ``node``. + :returns: ``True`` if a parallel map scope encloses the node, at any nesting depth. + """ + return is_in_scope(parent_state.sdfg, parent_state, node, + [dtypes.ScheduleType.CPU_Multicore, dtypes.ScheduleType.Default]) + + +def auto_dispatch(node: nodes.LibraryNode, parent_state: dace.SDFGState, + select_fn: Callable[[nodes.LibraryNode, dace.SDFGState], str], library_cls: type): + """Dispatch a library node's ``'Auto'`` implementation to the one ``select_fn`` picks, setting + ``node.implementation`` so introspection reflects what was chosen. + + :param node: the library node being expanded. + :param parent_state: state containing ``node`` (owning SDFG is ``parent_state.sdfg``). + :param select_fn: callable returning a concrete implementation name (not ``'Auto'``). + :param library_cls: the library node class with the ``implementations`` dict. + :returns: whatever the resolved expansion returns. + """ + impl_name = select_fn(node, parent_state) + assert impl_name != 'Auto', f"{select_fn.__name__} must not return 'Auto'." + node.implementation = impl_name + return library_cls.implementations[impl_name].expansion(node, parent_state, parent_state.sdfg) diff --git a/dace/libraries/standard/nodes/__init__.py b/dace/libraries/standard/nodes/__init__.py index 762e77760c..bbd83a5d5d 100644 --- a/dace/libraries/standard/nodes/__init__.py +++ b/dace/libraries/standard/nodes/__init__.py @@ -1,4 +1,6 @@ # Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. from .code import CodeLibraryNode +from .copy import CopyLibraryNode +from .fill import FillLibraryNode from .gearbox import Gearbox from .reduce import Reduce diff --git a/dace/libraries/standard/nodes/copy/__init__.py b/dace/libraries/standard/nodes/copy/__init__.py new file mode 100644 index 0000000000..d983431a15 --- /dev/null +++ b/dace/libraries/standard/nodes/copy/__init__.py @@ -0,0 +1,8 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +from dace.libraries.standard.nodes.copy.common import CopyExpansion, INPUT_CONNECTOR_NAME, OUTPUT_CONNECTOR_NAME +from dace.libraries.standard.nodes.copy.node import CopyLibraryNode +from dace.libraries.standard.nodes.copy.select import select_copy_implementation +from dace.libraries.standard.nodes.copy.expansions import (ExpandAuto, ExpandMappedTasklet, ExpandMemcpyCPU, + ExpandMemcpyCUDA1D, ExpandMemcpyCUDA2D, + ExpandMemcpyCUDANDStrided, ExpandSharedMemoryCollective, + ExpandTasklet) diff --git a/dace/libraries/standard/nodes/copy/common.py b/dace/libraries/standard/nodes/copy/common.py new file mode 100644 index 0000000000..94afe70b0a --- /dev/null +++ b/dace/libraries/standard/nodes/copy/common.py @@ -0,0 +1,399 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Shared pieces of the ``CopyLibraryNode`` expansions. + +Imported by both the node and its expansions, so it must not import either. +""" +import functools +import operator +from dataclasses import dataclass +from typing import List, Optional, Tuple, TYPE_CHECKING + +import dace +from dace import data, nodes, dtypes, subsets, symbolic +from dace.codegen.common import sym2cpp, get_gpu_backend +from dace.libraries.standard.helper import (CURRENT_STREAM_NAME, CPU_RESIDENT_STORAGES, GPU_RESIDENT_STORAGES, + collapse_shape_and_strides) +from dace.sdfg.scope import devicelevel_block_size, is_devicelevel_gpu + +if TYPE_CHECKING: + from dace.libraries.standard.nodes.copy.node import CopyLibraryNode + +INPUT_CONNECTOR_NAME = "_cpy_in" +OUTPUT_CONNECTOR_NAME = "_cpy_out" + + +@dataclass +class CopyExpansion: + """Inputs + collapsed-shape state shared across :class:`CopyLibraryNode` + expansions that build a wrapper SDFG. Returned by :func:`_make_expansion_sdfg`.""" + sdfg: dace.SDFG + state: dace.SDFGState + inp_name: str + inp: data.Data + in_subset: dace.subsets.Range + out_name: str + out: data.Data + out_subset: dace.subsets.Range + in_shape_collapsed: List[symbolic.SymExpr] + out_shape_collapsed: List[symbolic.SymExpr] + + +def _is_cross_cpu_gpu(src_storage: dtypes.StorageType, dst_storage: dtypes.StorageType, copy_node: "CopyLibraryNode", + parent_state: dace.SDFGState) -> bool: + """True if src/dst cross the CPU/GPU boundary. ``Register`` follows scope: GPU scope -> GPU, + else CPU.""" + in_gpu = is_devicelevel_gpu(parent_state.sdfg, parent_state, copy_node) + + src_gpu = (src_storage in GPU_RESIDENT_STORAGES) or (src_storage == dtypes.StorageType.Register and in_gpu) + dst_gpu = (dst_storage in GPU_RESIDENT_STORAGES) or (dst_storage == dtypes.StorageType.Register and in_gpu) + + src_cpu = (src_storage in CPU_RESIDENT_STORAGES) or (src_storage == dtypes.StorageType.Register and not in_gpu) + dst_cpu = (dst_storage in CPU_RESIDENT_STORAGES) or (dst_storage == dtypes.StorageType.Register and not in_gpu) + + return (src_cpu and dst_gpu) or (src_gpu and dst_cpu) + + +def _copy_waives_volume_check(copy_node: "CopyLibraryNode", parent_state: dace.SDFGState) -> bool: + """True if a wired memlet carries ``allow_oob``, i.e. the author waived the src/dst volume check. + + ``SDFGState.validate`` (``validation.py``) skips its own equal-volume check on such an edge, and + plain copy-edge codegen sizes the transfer from the source subset. A lifted ``CopyLibraryNode`` + must honour the same waiver, else an SDFG that is legal as a direct copy edge stops expanding. + """ + return any(e.data.allow_oob for e in parent_state.all_edges(copy_node) if not e.data.is_empty()) + + +def _both_packed_same_layout(inp: data.Data, out: data.Data) -> bool: + """True if both descriptors share packed major order (both C or both Fortran).""" + return ((inp.is_packed_c_strides() and out.is_packed_c_strides()) + or (inp.is_packed_fortran_strides() and out.is_packed_fortran_strides())) + + +def _delinearized_index(b_i: symbolic.symbol, shape: List[symbolic.SymExpr], layout: str) -> List[symbolic.SymExpr]: + """Multi-dim index for a 1-D walker into a packed-layout array. Only C (row-major) and + F (column-major) layouts are supported. + + :param b_i: the 1-D map symbol. + :param shape: per-dim extents in descriptor order. + :param layout: ``'C'`` (stride-1 is the last dim) or ``'F'`` (stride-1 is the first dim). + :returns: list of per-dim symbolic index expressions, in descriptor order. + """ + cum_strides = [] + cum = 1 + iter_shape = reversed(shape) if layout == 'C' else iter(shape) + for s in iter_shape: + cum_strides.append(cum) + cum *= s + if layout == 'C': + cum_strides.reverse() + return [symbolic.int_floor(b_i, cum_strides[d]) % shape[d] for d in range(len(shape))] + + +def cuda2d_pitch_params( + copy_shape: List[symbolic.SymExpr], src_strides: List[symbolic.SymExpr], dst_strides: List[symbolic.SymExpr] +) -> Optional[Tuple[symbolic.SymExpr, symbolic.SymExpr, symbolic.SymExpr, symbolic.SymExpr]]: + """Element-count ``cudaMemcpy2DAsync`` pitch params ``(dpitch, spitch, width, height)`` for a + 2D (or ``(N, 1)``-promoted) copy, or ``None`` if not a single ``cudaMemcpy2DAsync``. Single + source of truth for the ``MemcpyCUDA2D`` selector gate and the expander, so the two can't + drift: selector treats a non-``None`` result as "applies"; expander formats the same + components into the emitted call. Values are in elements; caller multiplies pitch/width by + ``sizeof(dtype)``. + + :param copy_shape: Two-element collapsed copy shape ``(rows, columns)``. + :param src_strides: Two-element source strides aligned with ``copy_shape``. + :param dst_strides: Two-element destination strides aligned with ``copy_shape``. + :returns: ``(dpitch, spitch, width, height)`` in elements, or ``None`` if not a single + ``cudaMemcpy2DAsync``. + """ + if src_strides[1] == 1 and dst_strides[1] == 1: + return dst_strides[0], src_strides[0], copy_shape[1], copy_shape[0] + if src_strides[0] == 1 and dst_strides[0] == 1: + return dst_strides[1], src_strides[1], copy_shape[0], copy_shape[1] + try: + if (not symbolic.inequal_symbols(src_strides[0] / src_strides[1], copy_shape[1]) + and not symbolic.inequal_symbols(dst_strides[0] / dst_strides[1], copy_shape[1])): + return dst_strides[1], src_strides[1], 1, copy_shape[0] * copy_shape[1] + except (TypeError, ZeroDivisionError): + return None + return None + + +def _make_expansion_sdfg(node: "CopyLibraryNode", + parent_state: dace.SDFGState, + allow_cross_storage: bool = False) -> CopyExpansion: + """Shared validation + wrapper-SDFG skeleton for expansions. + + :param node: the :class:`CopyLibraryNode` being expanded. + :param parent_state: state containing ``node``. + :param allow_cross_storage: permit differing src/dst storages. + :returns: a :class:`CopyExpansion` with the skeleton SDFG and collapsed shape/stride state. + """ + inp_name, inp, in_subset, out_name, out, out_subset = node.validate(parent_state.sdfg, + parent_state, + allow_cross_storage=allow_cross_storage) + + in_shape_collapsed, in_strides_collapsed = collapse_shape_and_strides(in_subset, inp.strides) + out_shape_collapsed, out_strides_collapsed = collapse_shape_and_strides(out_subset, out.strides) + + # The label is built from data names, and a struct member carries a '.' -- the SDFG name reaches + # C++ as an identifier (``dace/sdfg/sdfg.py`` sanitizes data names the same way). + label = node.label.replace('.', '_') + sdfg = dace.SDFG(f"{label}_sdfg") + sdfg.add_array(inp_name, in_shape_collapsed, inp.dtype, inp.storage, strides=in_strides_collapsed) + sdfg.add_array(out_name, out_shape_collapsed, out.dtype, out.storage, strides=out_strides_collapsed) + # Match the ambient stream connector if the experimental GPU codegen wired one in. + if CURRENT_STREAM_NAME in node.in_connectors: + sdfg.add_scalar(CURRENT_STREAM_NAME, dtypes.gpuStream_t, transient=False) + + state = sdfg.add_state(f"{label}_state", is_start_block=True) + + return CopyExpansion(sdfg=sdfg, + state=state, + inp_name=inp_name, + inp=inp, + in_subset=in_subset, + out_name=out_name, + out=out, + out_subset=out_subset, + in_shape_collapsed=in_shape_collapsed, + out_shape_collapsed=out_shape_collapsed) + + +def _make_mapped_tasklet_expansion(node: "CopyLibraryNode", + parent_state: dace.SDFGState, + allow_cross_storage: bool = False) -> dace.SDFG: + """Element-wise mapped tasklet expansion. Raises if the copy crosses the CPU/GPU boundary. + + Schedule from storages: ``Sequential`` for thread-level (Register/Register, + Register<->GPU_Shared) or any in-kernel copy; ``GPU_Device`` if any side is GPU storage + at host level; else ``Default``. + + :param node: the :class:`CopyLibraryNode` being expanded. + :param parent_state: state containing ``node``. + :param allow_cross_storage: permit differing src/dst storages. + :returns: the wrapper SDFG holding the mapped tasklet. + :raises ValueError: the copy crosses the CPU/GPU boundary. + """ + ctx = _make_expansion_sdfg(node, parent_state, allow_cross_storage=allow_cross_storage) + inp, out = ctx.inp, ctx.out + + if _is_cross_cpu_gpu(inp.storage, out.storage, node, parent_state): + raise ValueError("MappedTasklet expansion cannot cross the CPU/GPU boundary " + f"(got {inp.storage} -> {out.storage}). Use a MemcpyCUDA1D variant.") + + is_register = lambda s: s == dtypes.StorageType.Register + is_thread_local = (is_register(inp.storage) and is_register(out.storage)) or ( + (is_register(inp.storage) and out.storage == dtypes.StorageType.GPU_Shared) or + (is_register(out.storage) and inp.storage == dtypes.StorageType.GPU_Shared)) + in_kernel = is_devicelevel_gpu(parent_state.sdfg, parent_state, node) + if is_thread_local or in_kernel: + schedule = dtypes.ScheduleType.Sequential + elif inp.storage in GPU_RESIDENT_STORAGES or out.storage in GPU_RESIDENT_STORAGES: + schedule = dtypes.ScheduleType.GPU_Device + else: + # Default, not Sequential, inside an enclosing map: SCOPEDEFAULT_SCHEDULE already resolves a + # map nested in a CPU_Multicore one to single-threaded. + schedule = dtypes.ScheduleType.Default + + ctx.sdfg.schedule = dtypes.ScheduleType.Default + + # Must not collide with the wrapper SDFG's parameter arrays (named after outer connectors). + inner_in, inner_out = "_in", "_out" + in_shape, out_shape = ctx.in_shape_collapsed, ctx.out_shape_collapsed + + # A copy of zero elements moves nothing; the map below has an empty range, so the per-dim shape + # agreement the transfer would otherwise need is vacuous. + # ``is_length=False``: the default assumes every operand is positive, which 0 contradicts. + zero_size = any(symbolic.equal(s, 0, is_length=False) is True for s in in_shape + out_shape) + + if len(in_shape) == len(out_shape): + # Same-rank: per-dim map params, shared access expr on both sides. Shapes must match + # (per-dim permutations are transposes, not reshapes), unless the author waived the check. + # Refuse only on a PROVEN mismatch -- `equal` answers None when it cannot tell, and a + # symbolic pair such as ``ceiling(N/2)`` vs ``floor(N/2)`` is not grounds to reject. + if not zero_size and not _copy_waives_volume_check(node, parent_state) and any( + symbolic.equal(a, b) is False for a, b in zip(in_shape, out_shape)): + raise ValueError(f"MappedTasklet same-rank copy requires matching per-dim shapes; got src " + f"{tuple(in_shape)} vs dst {tuple(out_shape)}. Per-dim permutations are not " + f"supported -- use a Transpose libnode. Reshapes must change rank.") + map_params = [f"__i{i}" for i in range(len(in_shape))] + # Symbolic bounds, never a rendered string: ``sym2cpp`` spells a symbolic extent as C++ + # (``dace::math::ipow(R, K)``), and the range parser splits on ':', so the qualified name + # comes back as four bogus tokens. + map_rng = {i: (0, s - 1, 1) for i, s in zip(map_params, in_shape)} + access_expr = ','.join(map_params) + inputs = {inner_in: dace.memlet.Memlet(f"{ctx.inp_name}[{access_expr}]")} + outputs = {inner_out: dace.memlet.Memlet(f"{ctx.out_name}[{access_expr}]")} + else: + # Rank-mismatch reshape: 1-D walker over the total element count. Needs both endpoints + # packed same major order -- mixed layouts have no shared flat order. + if not _both_packed_same_layout(inp, out): + raise ValueError( + f"MappedTasklet rank-mismatched copy ({tuple(in_shape)} -> {tuple(out_shape)}) requires " + f"both endpoints to be packed in the same major order (both C-contiguous or both " + f"Fortran-contiguous). Got src '{ctx.inp_name}' strides {tuple(inp.strides)} on shape " + f"{tuple(inp.shape)} and dst '{ctx.out_name}' strides {tuple(out.strides)} on shape " + f"{tuple(out.shape)}. Mixed layouts are transposes -- use a same-rank Tasklet copy instead.") + layout = 'C' if inp.is_packed_c_strides() else 'F' + if layout == 'F': + # Under Fortran order the walker visits the FIRST collapsed dim fastest, but a strided or + # tiled dim collapses to (step count, tile) with the tile innermost -- only the C walk + # order then matches subset iteration order, so Fortran still needs flat subsets. + in_contig = ctx.in_subset.is_contiguous_subset(inp) + out_contig = ctx.out_subset.is_contiguous_subset(out) + if not (in_contig and out_contig): + raise ValueError( + f"MappedTasklet rank-mismatched copy ({tuple(in_shape)} -> {tuple(out_shape)}) requires " + f"contiguous subsets on both endpoints in Fortran layout (the 1-D walker treats the data " + f"as a flat sequence). Got src subset {ctx.in_subset} (contiguous: {in_contig}) on shape " + f"{tuple(inp.shape)} and dst subset {ctx.out_subset} (contiguous: {out_contig}) on shape " + f"{tuple(out.shape)}.") + + # Product of the collapsed shape, not ``num_elements_exact`` -- the latter is a BOUNDING BOX + # (``subsets.py``), which overcounts a strided or tiled subset. + total = functools.reduce(operator.mul, in_shape, 1) + b_i_name = "__b_i" + b_i = symbolic.symbol(b_i_name) + map_rng = {b_i_name: (0, total - 1, 1)} + + def _side_access(arr_name, shape): + idx = [b_i] if len(shape) == 1 else _delinearized_index(b_i, shape, layout) + return dace.memlet.Memlet(data=arr_name, subset=subsets.Range([(e, e, 1) for e in idx])) + + inputs = {inner_in: _side_access(ctx.inp_name, in_shape)} + outputs = {inner_out: _side_access(ctx.out_name, out_shape)} + + _, map_entry, _ = ctx.state.add_mapped_tasklet(f"{node.label}_tasklet", + map_rng, + inputs, + f"{inner_out} = {inner_in}", + outputs, + schedule=schedule, + external_edges=True) + + return ctx.sdfg + + +def _memcpy_kind(inp: data.Data, out: data.Data) -> str: + """``cudaMemcpyTo`` from endpoint storages.""" + src_loc = "Device" if inp.storage == dace.dtypes.StorageType.GPU_Global else "Host" + dst_loc = "Device" if out.storage == dace.dtypes.StorageType.GPU_Global else "Host" + backend = get_gpu_backend() + return f"{backend}Memcpy{src_loc}To{dst_loc}" + + +def _make_memcpy_tasklet(node: "CopyLibraryNode", parent_state: dace.SDFGState, *, cuda: bool) -> nodes.Tasklet: + """Build a Tasklet emitting one contiguous-block copy. Raises ``ValueError`` on a + non-contiguous subset (the single-call form would overrun the region; use ``MappedTasklet``). + + Emits ``cudaMemcpyAsync`` when ``cuda`` is set -- cross-CPU/GPU allowed, direction + (HostToDevice/DeviceToHost/DeviceToDevice/HostToHost) inferred from endpoint storages -- + else a same-storage ``std::memcpy``. + + :param node: the :class:`CopyLibraryNode` being expanded. + :param parent_state: state containing ``node`` (owning SDFG is ``parent_state.sdfg``). + :param cuda: emit ``cudaMemcpyAsync`` (else ``memcpy``). + :returns: a :class:`~dace.sdfg.nodes.Tasklet` issuing the copy. + :raises ValueError: a subset is non-contiguous. + """ + label = "MemcpyCUDA1D" if cuda else "MemcpyCPU" + inp_name, inp, in_subset, out_name, out, out_subset = node.validate(parent_state.sdfg, + parent_state, + allow_cross_storage=cuda) + single_elt = (in_subset.num_elements_exact() == 1 and out_subset.num_elements_exact() == 1) + if single_elt: + pass + elif not (in_subset.is_contiguous_subset(inp) and out_subset.is_contiguous_subset(out)): + raise ValueError(f"{label} requires contiguous subsets; got src '{inp_name}' subset {in_subset} " + f"(shape {inp.shape} strides {inp.strides}) and dst '{out_name}' subset {out_subset} " + f"(shape {out.shape} strides {out.strides}). Use MappedTasklet for strided subsets.") + + in_conn = INPUT_CONNECTOR_NAME + out_conn = OUTPUT_CONNECTOR_NAME + nbytes = f"{sym2cpp(in_subset.num_elements_exact())} * sizeof({inp.dtype.ctype})" + if cuda: + backend = get_gpu_backend() + code = f"{backend}MemcpyAsync({out_conn}, {in_conn}, {nbytes}, {_memcpy_kind(inp, out)}, {CURRENT_STREAM_NAME});" + else: + code = f"memcpy({out_conn}, {in_conn}, {nbytes});" + + return nodes.Tasklet(node.name, + inputs={in_conn: dace.dtypes.pointer(inp.dtype)}, + outputs={out_conn: dace.dtypes.pointer(out.dtype)}, + code=code, + language=dace.Language.CPP) + + +def _build_shmem_collective_copy_code(node: "CopyLibraryNode", parent_state: dace.SDFGState, inp: data.Data, + in_subset: dace.subsets.Range, out: data.Data, + out_subset: dace.subsets.Range) -> str: + """Build the C++ code for ``ExpandSharedMemoryCollective``. + + A static 1-D transfer inside a kernel uses DaCe's block-collective runtime helpers + (``dace::GlobalToShared1D`` / ``dace::SharedToGlobal1D``), which split the elements across the + thread block -- the same call plain copy-edge codegen emits. Everything else falls back to a + ``dace::CopyND<...>::Copy(...)`` plus ``__syncthreads()``: the most-specific static template + (``CopyNDDynamic`` for symbolic shapes), refined by ``ConstDst``/``ConstSrc``/``Dynamic`` on + whichever stride set is constexpr, with the rest passed as runtime args. + + :param node: the :class:`CopyLibraryNode` being expanded. + :param parent_state: state containing ``node`` (supplies the enclosing thread-block size). + :param inp: source descriptor (provides ``ctype`` and ``strides``). + :param in_subset: source memlet subset. + :param out: destination descriptor (provides ``strides``). + :param out_subset: destination memlet subset. + :returns: the tasklet body. + """ + copy_shape, src_strides = collapse_shape_and_strides(in_subset, inp.strides) + _, dst_strides = collapse_shape_and_strides(out_subset, out.strides) + ndims = len(copy_shape) + + in_conn = INPUT_CONNECTOR_NAME + out_conn = OUTPUT_CONNECTOR_NAME + block_dims = devicelevel_block_size(parent_state.sdfg, parent_state, node) + if ndims == 1 and block_dims is not None and not any( + symbolic.issymbolic(s) for s in (copy_shape[0], src_strides[0], dst_strides[0])): + bdims = ', '.join(sym2cpp(b) for b in block_dims) + args = f"{inp.dtype.ctype}, {bdims}, {sym2cpp(copy_shape[0])}" + if out.storage == dtypes.StorageType.GPU_Shared: + return (f"dace::GlobalToShared1D<{args}, {sym2cpp(dst_strides[0])}, false>" + f"({in_conn}, {sym2cpp(src_strides[0])}, {out_conn});") + return (f"dace::SharedToGlobal1D<{args}, false>::Copy" + f"({in_conn}, {sym2cpp(src_strides[0])}, {out_conn}, {sym2cpp(dst_strides[0])});") + + shape_strs = [sym2cpp(s) for s in copy_shape] + src_stride_strs = [sym2cpp(s) for s in src_strides] + dst_stride_strs = [sym2cpp(s) for s in dst_strides] + + dims_static = not any(symbolic.issymbolic(s) for s in copy_shape) + src_static = not any(symbolic.issymbolic(s) for s in src_strides) + dst_static = not any(symbolic.issymbolic(s) for s in dst_strides) + + ctype = inp.dtype.ctype + if dims_static: + copy_tmpl = f"dace::CopyND<{ctype}, 1, false, {', '.join(shape_strs)}>" + else: + copy_tmpl = f"dace::CopyNDDynamic<{ctype}, 1, false, {ndims}>" + + # Prefer ConstDst, else ConstSrc, else Dynamic; the rest go as runtime args, per-dim order. + if dst_static: + shape_tmpl = f"template ConstDst<{', '.join(dst_stride_strs)}>" + elif src_static: + shape_tmpl = f"template ConstSrc<{', '.join(src_stride_strs)}>" + else: + shape_tmpl = "Dynamic" + + stride_args = [] + for d in range(ndims): + if not dims_static: + stride_args.append(shape_strs[d]) + if not src_static or dst_static: + stride_args.append(src_stride_strs[d]) + if not dst_static: + stride_args.append(dst_stride_strs[d]) + + all_args = [in_conn, out_conn] + stride_args + # ``CopyND`` makes every thread copy the whole region, so a shared SOURCE is read across thread + # boundaries and needs the block's writes to have landed first. + preamble = "__syncthreads();\n" if inp.storage == dtypes.StorageType.GPU_Shared else "" + return f"{preamble}{copy_tmpl}::{shape_tmpl}::Copy({', '.join(all_args)});\n__syncthreads();" diff --git a/dace/libraries/standard/nodes/copy/expansions/__init__.py b/dace/libraries/standard/nodes/copy/expansions/__init__.py new file mode 100644 index 0000000000..5d6813b048 --- /dev/null +++ b/dace/libraries/standard/nodes/copy/expansions/__init__.py @@ -0,0 +1,10 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Every ``CopyLibraryNode`` expansion. Imported here so registration runs on package import.""" +from dace.libraries.standard.nodes.copy.expansions.auto import ExpandAuto +from dace.libraries.standard.nodes.copy.expansions.mapped_tasklet import ExpandMappedTasklet +from dace.libraries.standard.nodes.copy.expansions.memcpy_cpu import ExpandMemcpyCPU +from dace.libraries.standard.nodes.copy.expansions.memcpy_cuda1d import ExpandMemcpyCUDA1D +from dace.libraries.standard.nodes.copy.expansions.memcpy_cuda2d import ExpandMemcpyCUDA2D +from dace.libraries.standard.nodes.copy.expansions.memcpy_cuda_nd import ExpandMemcpyCUDANDStrided +from dace.libraries.standard.nodes.copy.expansions.shmem_collective import ExpandSharedMemoryCollective +from dace.libraries.standard.nodes.copy.expansions.tasklet import ExpandTasklet diff --git a/dace/libraries/standard/nodes/copy/expansions/auto.py b/dace/libraries/standard/nodes/copy/expansions/auto.py new file mode 100644 index 0000000000..8c24026ef0 --- /dev/null +++ b/dace/libraries/standard/nodes/copy/expansions/auto.py @@ -0,0 +1,24 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Default expansion: dispatches via :func:`select_copy_implementation`. +""" +from typing import TYPE_CHECKING + +from dace import library +from dace.libraries.standard.helper import auto_dispatch +from dace.transformation.transformation import ExpandTransformation +from dace.libraries.standard.nodes.copy.select import select_copy_implementation + +if TYPE_CHECKING: + pass + + +@library.expansion +class ExpandAuto(ExpandTransformation): + """Default expansion: dispatches to the implementation chosen by + :func:`select_copy_implementation` from endpoint storages, subset shapes, and scope.""" + environments = [] + + @staticmethod + def expansion(node, parent_state, parent_sdfg): + from dace.libraries.standard.nodes.copy.node import CopyLibraryNode # Avoid import loop + return auto_dispatch(node, parent_state, select_copy_implementation, CopyLibraryNode) diff --git a/dace/libraries/standard/nodes/copy/expansions/mapped_tasklet.py b/dace/libraries/standard/nodes/copy/expansions/mapped_tasklet.py new file mode 100644 index 0000000000..dba69c991b --- /dev/null +++ b/dace/libraries/standard/nodes/copy/expansions/mapped_tasklet.py @@ -0,0 +1,23 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Element-wise mapped copy: the general form every other expansion falls back to. +""" +from typing import TYPE_CHECKING + +from dace import library +from dace.transformation.transformation import ExpandTransformation +from dace.libraries.standard.nodes.copy.common import (_make_mapped_tasklet_expansion) + +if TYPE_CHECKING: + pass + + +@library.expansion +class ExpandMappedTasklet(ExpandTransformation): + """Mapped element-wise tasklet ``_cpy_out = _cpy_in`` over the collapsed copy shape. Schedule + from endpoint storages: ``Sequential`` for Register/Register<->GPU_Shared (thread-level), + ``GPU_Device`` if any side is GPU, else ``Default``. Raises across the CPU/GPU boundary.""" + environments = [] + + @staticmethod + def expansion(node, parent_state, parent_sdfg): + return _make_mapped_tasklet_expansion(node, parent_state, allow_cross_storage=True) diff --git a/dace/libraries/standard/nodes/copy/expansions/memcpy_cpu.py b/dace/libraries/standard/nodes/copy/expansions/memcpy_cpu.py new file mode 100644 index 0000000000..5bcf43dacc --- /dev/null +++ b/dace/libraries/standard/nodes/copy/expansions/memcpy_cpu.py @@ -0,0 +1,22 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Single ``memcpy`` on the host. +""" +from typing import TYPE_CHECKING + +from dace import library +from dace.libraries.standard import environments +from dace.transformation.transformation import ExpandTransformation +from dace.libraries.standard.nodes.copy.common import (_make_memcpy_tasklet) + +if TYPE_CHECKING: + pass + + +@library.expansion +class ExpandMemcpyCPU(ExpandTransformation): + """One ``std::memcpy`` for a contiguous CPU<->CPU copy.""" + environments = [environments.CPU] + + @staticmethod + def expansion(node, parent_state, parent_sdfg): + return _make_memcpy_tasklet(node, parent_state, cuda=False) diff --git a/dace/libraries/standard/nodes/copy/expansions/memcpy_cuda1d.py b/dace/libraries/standard/nodes/copy/expansions/memcpy_cuda1d.py new file mode 100644 index 0000000000..5e11dca972 --- /dev/null +++ b/dace/libraries/standard/nodes/copy/expansions/memcpy_cuda1d.py @@ -0,0 +1,23 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Contiguous device copy through ``cudaMemcpyAsync``. +""" +from typing import TYPE_CHECKING + +from dace import library +from dace.libraries.standard import environments +from dace.transformation.transformation import ExpandTransformation +from dace.libraries.standard.nodes.copy.common import (_make_memcpy_tasklet) + +if TYPE_CHECKING: + pass + + +@library.expansion +class ExpandMemcpyCUDA1D(ExpandTransformation): + """One ``cudaMemcpyAsync`` for a contiguous copy; direction (H2D/D2H/D2D/H2H) inferred from + endpoint storages.""" + environments = [environments.CUDA] + + @staticmethod + def expansion(node, parent_state, parent_sdfg): + return _make_memcpy_tasklet(node, parent_state, cuda=True) diff --git a/dace/libraries/standard/nodes/copy/expansions/memcpy_cuda2d.py b/dace/libraries/standard/nodes/copy/expansions/memcpy_cuda2d.py new file mode 100644 index 0000000000..c7e28b9087 --- /dev/null +++ b/dace/libraries/standard/nodes/copy/expansions/memcpy_cuda2d.py @@ -0,0 +1,77 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Pitched 2-D device copy through ``cudaMemcpy2DAsync``. +""" +from typing import TYPE_CHECKING + +import dace +from dace import library, nodes +from dace.codegen.common import sym2cpp, get_gpu_backend +from dace.libraries.standard import environments +from dace.libraries.standard.helper import (CURRENT_STREAM_NAME, collapse_shape_and_strides) +from dace.transformation.transformation import ExpandTransformation +from dace.libraries.standard.nodes.copy.common import (_memcpy_kind, cuda2d_pitch_params, INPUT_CONNECTOR_NAME, + OUTPUT_CONNECTOR_NAME) + +if TYPE_CHECKING: + pass + + +@library.expansion +class ExpandMemcpyCUDA2D(ExpandTransformation): + """2D strided copy via ``cudaMemcpy2DAsync`` between any GPU_Global/host storage combination: + row-major contiguous rows, column-major contiguous columns, or outer stride a multiple of + inner (degenerate).""" + environments = [environments.CUDA] + + @staticmethod + def expansion(node, parent_state, parent_sdfg): + inp_name, inp, in_subset, out_name, out, out_subset = node.validate(parent_sdfg, + parent_state, + allow_cross_storage=True) + + in_shape_collapsed, in_strides_collapsed = collapse_shape_and_strides(in_subset, inp.strides) + out_shape_collapsed, out_strides_collapsed = collapse_shape_and_strides(out_subset, out.strides) + + # 1D-collapsed shapes promote to (N, 1) so one cudaMemcpy2D call covers strided 1D patterns. + if len(in_shape_collapsed) == 1 and len(out_shape_collapsed) == 1: + in_shape_2d = [in_shape_collapsed[0], 1] + out_shape_2d = [out_shape_collapsed[0], 1] + in_strides_2d = [in_strides_collapsed[0], 1] + out_strides_2d = [out_strides_collapsed[0], 1] + elif len(in_shape_collapsed) == 2 and len(out_shape_collapsed) == 2: + in_shape_2d = in_shape_collapsed + out_shape_2d = out_shape_collapsed + in_strides_2d = in_strides_collapsed + out_strides_2d = out_strides_collapsed + else: + raise ValueError("MemcpyCUDA2D requires 1D or 2D collapsed shapes, got " + f"{in_shape_collapsed} (src) / {out_shape_collapsed} (dst).") + + kind = _memcpy_kind(inp, out) + + copy_shape = in_shape_2d + src_strides = in_strides_2d + dst_strides = out_strides_2d + ctype = inp.dtype.ctype + backend = get_gpu_backend() + + pitch = cuda2d_pitch_params(copy_shape, src_strides, dst_strides) + if pitch is None: + raise NotImplementedError(f"Unsupported 2D memory copy: shape={copy_shape}, " + f"src_strides={src_strides}, dst_strides={dst_strides}.") + dpitch_elems, spitch_elems, width_elems, height_elems = pitch + dpitch = f"{sym2cpp(dpitch_elems)} * sizeof({ctype})" + spitch = f"{sym2cpp(spitch_elems)} * sizeof({ctype})" + width = f"{sym2cpp(width_elems)} * sizeof({ctype})" + height = sym2cpp(height_elems) + + code = (f"{backend}Memcpy2DAsync({OUTPUT_CONNECTOR_NAME}, {dpitch}, {INPUT_CONNECTOR_NAME}, {spitch}, " + f"{width}, {height}, {kind}, {CURRENT_STREAM_NAME});") + + in_conns = {INPUT_CONNECTOR_NAME: dace.dtypes.pointer(inp.dtype)} + tasklet = nodes.Tasklet(node.name, + inputs=in_conns, + outputs={OUTPUT_CONNECTOR_NAME: dace.dtypes.pointer(out.dtype)}, + code=code, + language=dace.Language.CPP) + return tasklet diff --git a/dace/libraries/standard/nodes/copy/expansions/memcpy_cuda_nd.py b/dace/libraries/standard/nodes/copy/expansions/memcpy_cuda_nd.py new file mode 100644 index 0000000000..25caa8532b --- /dev/null +++ b/dace/libraries/standard/nodes/copy/expansions/memcpy_cuda_nd.py @@ -0,0 +1,105 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Strided N-D device copy, issued as a loop of contiguous ``cudaMemcpyAsync`` calls. +""" +from typing import TYPE_CHECKING + +import dace +from dace import library, nodes, subsets, symbolic +from dace.codegen.common import sym2cpp, get_gpu_backend +from dace.libraries.standard import environments +from dace.libraries.standard.helper import (CURRENT_STREAM_NAME, collapse_shape_and_strides) +from dace.transformation.transformation import ExpandTransformation +from dace.libraries.standard.nodes.copy.common import (_make_expansion_sdfg, _memcpy_kind, INPUT_CONNECTOR_NAME, + OUTPUT_CONNECTOR_NAME) + +if TYPE_CHECKING: + pass + + +@library.expansion +class ExpandMemcpyCUDANDStrided(ExpandTransformation): + """Fallback for >=3D-strided cross-boundary copies that can't collapse to one + ``cudaMemcpyAsync`` / ``cudaMemcpy2DAsync``: a Sequential map issuing one ``cudaMemcpyAsync`` + per row over every collapsed dim except the chunk axis (``stride == 1`` both sides).""" + environments = [environments.CUDA] + + @staticmethod + def expansion(node, parent_state, parent_sdfg): + inp_name, inp, in_subset, out_name, out, out_subset = node.validate(parent_sdfg, + parent_state, + allow_cross_storage=True) + in_shape_collapsed, in_strides_collapsed = collapse_shape_and_strides(in_subset, inp.strides) + out_shape_collapsed, out_strides_collapsed = collapse_shape_and_strides(out_subset, out.strides) + + if len(in_shape_collapsed) != len(out_shape_collapsed): + raise NotImplementedError("ExpandCUDANDStrided requires src and dst to share the collapsed rank " + f"(got {in_shape_collapsed} vs {out_shape_collapsed}).") + ndims = len(in_shape_collapsed) + if ndims < 1: + raise NotImplementedError("ExpandCUDANDStrided requires at least one collapsed dimension.") + + # Chunk axis: innermost dim with stride 1 on both sides. + chunk_dim = None + for d in reversed(range(ndims)): + if in_strides_collapsed[d] == 1 and out_strides_collapsed[d] == 1: + chunk_dim = d + break + if chunk_dim is None: + raise NotImplementedError("ExpandCUDANDStrided requires at least one common stride-1 axis on both sides " + f"(got src_strides={in_strides_collapsed}, dst_strides={out_strides_collapsed}).") + + ctype = inp.dtype.ctype + chunk = sym2cpp(in_shape_collapsed[chunk_dim]) + kind = _memcpy_kind(inp, out) + backend = get_gpu_backend() + + if ndims == 1: + code = (f"DACE_GPU_CHECK({backend}MemcpyAsync({OUTPUT_CONNECTOR_NAME}, {INPUT_CONNECTOR_NAME}, " + f"{chunk} * sizeof({ctype}), {kind}, {CURRENT_STREAM_NAME}));") + in_conns = {INPUT_CONNECTOR_NAME: dace.dtypes.pointer(inp.dtype)} + return nodes.Tasklet(node.name, + inputs=in_conns, + outputs={OUTPUT_CONNECTOR_NAME: dace.dtypes.pointer(out.dtype)}, + code=code, + language=dace.Language.CPP) + + ctx = _make_expansion_sdfg(node, parent_state, allow_cross_storage=True) + map_axes = [d for d in range(ndims) if d != chunk_dim] + map_params = [f"__cpy_i{d}" for d in map_axes] + # Symbolic bounds, never a rendered string: ``sym2cpp`` spells a symbolic extent as C++ + # (``dace::math::ipow(R, K)``), and the range parser splits on ':', so the qualified name + # comes back as four bogus tokens. + map_ranges = {p: (0, ctx.in_shape_collapsed[d] - 1, 1) for d, p in zip(map_axes, map_params)} + + def _row_subset(shape): + parts = [] + map_pi = 0 + for d in range(ndims): + if d == chunk_dim: + parts.append((0, shape[d] - 1, 1)) + else: + p = symbolic.symbol(map_params[map_pi]) + parts.append((p, p, 1)) + map_pi += 1 + return subsets.Range(parts) + + in_memlet = dace.memlet.Memlet(data=ctx.inp_name, subset=_row_subset(ctx.in_shape_collapsed)) + out_memlet = dace.memlet.Memlet(data=ctx.out_name, subset=_row_subset(ctx.out_shape_collapsed)) + inner_in, inner_out = "_in", "_out" + backend = get_gpu_backend() + code = (f"DACE_GPU_CHECK({backend}MemcpyAsync({inner_out}, {inner_in}, " + f"{chunk} * sizeof({ctype}), {kind}, {CURRENT_STREAM_NAME}));") + + inner_tasklet, map_entry, _map_exit = ctx.state.add_mapped_tasklet(name=f"{node.label}_tasklet", + map_ranges=map_ranges, + inputs={inner_in: in_memlet}, + code=code, + outputs={inner_out: out_memlet}, + schedule=dace.dtypes.ScheduleType.Sequential, + language=dace.Language.CPP, + external_edges=True) + # Force pointer connectors so codegen types them T*, matching cudaMemcpyAsync's signature. + inner_tasklet.in_connectors[inner_in] = dace.dtypes.pointer(inp.dtype) + inner_tasklet.out_connectors[inner_out] = dace.dtypes.pointer(out.dtype) + + return ctx.sdfg diff --git a/dace/libraries/standard/nodes/copy/expansions/shmem_collective.py b/dace/libraries/standard/nodes/copy/expansions/shmem_collective.py new file mode 100644 index 0000000000..2d3230854f --- /dev/null +++ b/dace/libraries/standard/nodes/copy/expansions/shmem_collective.py @@ -0,0 +1,52 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Block-collective copy into shared memory. +""" +from typing import TYPE_CHECKING + +import dace +from dace import library, nodes, dtypes +from dace.sdfg.scope import is_in_scope +from dace.transformation.transformation import ExpandTransformation +from dace.libraries.standard.nodes.copy.common import (_build_shmem_collective_copy_code, INPUT_CONNECTOR_NAME, + OUTPUT_CONNECTOR_NAME) + +if TYPE_CHECKING: + pass + + +@library.expansion +class ExpandSharedMemoryCollective(ExpandTransformation): + """Block-collective Shared <-> Shared/Global copy: a single Tasklet emitting + ``dace::CopyND<...>::Copy + __syncthreads()``, with ``_in``/``_out`` connectors matching the + libnode's directly (no NSDFG wrapper -- the parent kernel's ``__shared__`` array binds + straight in, no scope-id name mangling). + + Caller must place this outside any enclosing ``GPU_ThreadBlock`` map -- this expansion *is* + the thread-block-level operation. Shared <-> Register goes through ``MappedTasklet`` instead + (auto selector routes it there).""" + environments = [] + + @staticmethod + def expansion(node, parent_state, parent_sdfg): + inp_name, inp, in_subset, out_name, out, out_subset = node.validate(parent_sdfg, + parent_state, + allow_cross_storage=True) + + valid_storages = {dtypes.StorageType.GPU_Shared, dtypes.StorageType.GPU_Global} + if inp.storage not in valid_storages or out.storage not in valid_storages: + raise ValueError(f"SharedMemoryCollective requires GPU_Shared / GPU_Global storages " + f"(got {inp.storage} -> {out.storage}). Use MappedTasklet for " + "Shared <-> Register thread-level copies.") + if inp.storage != dtypes.StorageType.GPU_Shared and out.storage != dtypes.StorageType.GPU_Shared: + raise ValueError("SharedMemoryCollective requires at least one side to be GPU_Shared.") + + if is_in_scope(parent_sdfg, parent_state, node, [dtypes.ScheduleType.GPU_ThreadBlock]): + raise ValueError("SharedMemoryCollective IS the thread-block-level operation " + "and must not be nested inside a GPU_ThreadBlock map.") + + return nodes.Tasklet(node.name, + inputs={INPUT_CONNECTOR_NAME: dace.dtypes.pointer(inp.dtype)}, + outputs={OUTPUT_CONNECTOR_NAME: dace.dtypes.pointer(out.dtype)}, + code=_build_shmem_collective_copy_code(node, parent_state, inp, in_subset, out, + out_subset), + language=dace.Language.CPP) diff --git a/dace/libraries/standard/nodes/copy/expansions/tasklet.py b/dace/libraries/standard/nodes/copy/expansions/tasklet.py new file mode 100644 index 0000000000..c88b14bca0 --- /dev/null +++ b/dace/libraries/standard/nodes/copy/expansions/tasklet.py @@ -0,0 +1,44 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Single-element same-side scalar assignment. +""" +from typing import TYPE_CHECKING + +import dace +from dace import library, nodes +from dace.sdfg.scope import is_devicelevel_gpu +from dace.transformation.transformation import ExpandTransformation +from dace.libraries.standard.nodes.copy.common import (_is_cross_cpu_gpu, INPUT_CONNECTOR_NAME, OUTPUT_CONNECTOR_NAME) + +if TYPE_CHECKING: + pass + + +@library.expansion +class ExpandTasklet(ExpandTransformation): + """Single-element same-side scalar copy: ``_cpy_out = _cpy_in`` as a Python tasklet""" + environments = [] + + @staticmethod + def expansion(node, parent_state, parent_sdfg): + inp_name, inp, in_subset, out_name, out, out_subset = node.validate(parent_sdfg, + parent_state, + allow_cross_storage=True) + in_volume = in_subset.num_elements_exact() + out_volume = out_subset.num_elements_exact() + if in_volume != 1 or out_volume != 1: + raise ValueError(f"Tasklet expansion requires single-element subsets " + f"(got input volume {in_volume}, output volume {out_volume}). " + f"Use MappedTasklet for multi-element copies.") + # Single-element Shared involvement is a valid thread-level assignment, routed here. + # In-kernel the boundary does not exist for a single element: the host side arrives as a + # by-value kernel argument (or is pinned, hence device-addressable), so assignment is right. + if not is_devicelevel_gpu(parent_sdfg, parent_state, node) and _is_cross_cpu_gpu( + inp.storage, out.storage, node, parent_state): + raise ValueError(f"Tasklet expansion: storage types must match (no CPU/GPU boundary); " + f"got {inp.storage} -> {out.storage}. Use a MemcpyCUDA1D variant instead.") + + return nodes.Tasklet(node.name, + inputs={INPUT_CONNECTOR_NAME: inp.dtype}, + outputs={OUTPUT_CONNECTOR_NAME: out.dtype}, + code=f"{OUTPUT_CONNECTOR_NAME} = {INPUT_CONNECTOR_NAME}", + language=dace.Language.Python) diff --git a/dace/libraries/standard/nodes/copy/node.py b/dace/libraries/standard/nodes/copy/node.py new file mode 100644 index 0000000000..87eacb6820 --- /dev/null +++ b/dace/libraries/standard/nodes/copy/node.py @@ -0,0 +1,125 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""``CopyLibraryNode``: an explicit copy between two data containers. +""" +from typing import TYPE_CHECKING + +from dace import library, nodes, dtypes +from dace.libraries.standard.helper import (CURRENT_STREAM_NAME, CPU_RESIDENT_STORAGES) +from dace.libraries.standard.nodes.copy.common import INPUT_CONNECTOR_NAME, OUTPUT_CONNECTOR_NAME +from dace.libraries.standard.nodes.copy.expansions import (ExpandAuto, ExpandMappedTasklet, ExpandMemcpyCPU, + ExpandMemcpyCUDA1D, ExpandMemcpyCUDA2D, + ExpandMemcpyCUDANDStrided, ExpandSharedMemoryCollective, + ExpandTasklet) + +if TYPE_CHECKING: + from dace.libraries.standard.nodes.copy.node import CopyLibraryNode + + +@library.node +class CopyLibraryNode(nodes.LibraryNode): + """Library node representing a data copy between two access nodes. Implementations: + ``MappedTasklet`` (element-wise tasklet, also the rank-mismatch/reshape and large-CPU-copy + path), ``Tasklet`` (bare assignment, no map), ``MemcpyCPU`` (single ``std::memcpy``), + ``MemcpyCUDA1D``/``2D`` (``cudaMemcpyAsync``/``cudaMemcpy2DAsync``), ``MemcpyCUDANDStrided`` + (Sequential map of ``cudaMemcpyAsync``), ``SharedMemoryCollective`` (``dace::CopyND`` + + ``__syncthreads()``). + + Does NOT accept dynamic (Scalar) input connectors -- subset expressions must use symbols + already in scope at construction time, so the auto selector reasons purely from static + memlet subsets. + """ + + implementations = { + "Auto": ExpandAuto, + "MappedTasklet": ExpandMappedTasklet, + "Tasklet": ExpandTasklet, + "MemcpyCPU": ExpandMemcpyCPU, + "MemcpyCUDA1D": ExpandMemcpyCUDA1D, + "MemcpyCUDA2D": ExpandMemcpyCUDA2D, + "MemcpyCUDANDStrided": ExpandMemcpyCUDANDStrided, + "SharedMemoryCollective": ExpandSharedMemoryCollective, + } + default_implementation = 'Auto' + + INPUT_CONNECTOR_NAME = "_cpy_in" + OUTPUT_CONNECTOR_NAME = "_cpy_out" + + def __init__(self, name, *args, **kwargs): + super().__init__(name, *args, inputs={INPUT_CONNECTOR_NAME}, outputs={OUTPUT_CONNECTOR_NAME}, **kwargs) + + def src_storage(self, state) -> dtypes.StorageType: + """Storage of the array feeding ``_cpy_in``, or ``Default`` if unwired. + + :param state: state containing this libnode (owning SDFG is ``state.sdfg``). + :returns: the source :class:`~dace.dtypes.StorageType`. + """ + in_edges = [e for e in state.in_edges(self) if e.dst_conn == INPUT_CONNECTOR_NAME] + if not in_edges: + return dtypes.StorageType.Default + outer = state.memlet_path(in_edges[0])[0].src + if not isinstance(outer, nodes.AccessNode): + return dtypes.StorageType.Default + return state.sdfg.arrays[outer.data].storage + + def dst_storage(self, state) -> dtypes.StorageType: + """Storage of the array fed by ``_cpy_out``, or ``Default`` if unwired. + + :param state: state containing this libnode (owning SDFG is ``state.sdfg``). + :returns: the destination :class:`~dace.dtypes.StorageType`. + """ + out_edges = [e for e in state.out_edges(self) if e.src_conn == OUTPUT_CONNECTOR_NAME] + if not out_edges: + return dtypes.StorageType.Default + outer = state.memlet_path(out_edges[0])[-1].dst + if not isinstance(outer, nodes.AccessNode): + return dtypes.StorageType.Default + return state.sdfg.arrays[outer.data].storage + + def validate(self, sdfg, state, allow_cross_storage=True): + """Resolve in/out edges, names, and subsets: ``(inp_name, inp, in_subset, out_name, out, + out_subset)``. Raises ``ValueError`` if not wired with exactly one input and one output + data edge, dtypes mismatch, an extraneous non-reserved input connector wired, or (when + ``allow_cross_storage`` is False) the two storages differ. + + :param sdfg: SDFG containing ``state``. + :param state: state containing this libnode. + :param allow_cross_storage: when False, require matching src/dst storages. + :returns: ``(inp_name, inp, in_subset, out_name, out, out_subset)``. + :raises ValueError: see above. + """ + out_edges = [oe for oe in state.out_edges(self) if oe.src_conn == OUTPUT_CONNECTOR_NAME] + if len(out_edges) != 1: + raise ValueError(f"{type(self).__name__} expects exactly one " + f"``{OUTPUT_CONNECTOR_NAME}`` output edge.") + oe = out_edges[0] + out = sdfg.arrays[oe.data.data] + out_subset = oe.data.subset + out_name = oe.src_conn + + reserved = {INPUT_CONNECTOR_NAME, CURRENT_STREAM_NAME} + extra = [ie.dst_conn for ie in state.in_edges(self) if ie.dst_conn not in reserved and not ie.data.is_empty()] + if extra: + raise ValueError(f"{type(self).__name__} does not accept dynamic input connectors; got {extra}. " + f"Subset expressions must use symbols already in scope.") + + in_edges = [ie for ie in state.in_edges(self) if ie.dst_conn == INPUT_CONNECTOR_NAME] + if len(in_edges) != 1: + raise ValueError(f"{type(self).__name__} expects exactly one data input edge " + f"connected to the ``{INPUT_CONNECTOR_NAME}`` connector.") + ie = in_edges[0] + inp = sdfg.arrays[ie.data.data] + in_subset = ie.data.subset + inp_name = ie.dst_conn + + if inp.dtype != out.dtype: + raise ValueError(f"Input and output data types must match (got {inp.dtype} vs {out.dtype}).") + + # Two host storages differ only in the allocator, so a plain memcpy between them is correct; + # only a CPU/GPU (or other target-specific) pairing genuinely needs a different expansion. + host_pair = {inp.storage, out.storage} <= (CPU_RESIDENT_STORAGES | {dtypes.StorageType.Default}) + if not allow_cross_storage and inp.storage != out.storage and not host_pair: + raise ValueError(f"Input and output storage types must match for this expansion " + f"(got {inp.storage} vs {out.storage}). Use a cross-storage " + f"expansion or the pure fallback.") + + return inp_name, inp, in_subset, out_name, out, out_subset diff --git a/dace/libraries/standard/nodes/copy/select.py b/dace/libraries/standard/nodes/copy/select.py new file mode 100644 index 0000000000..b3bba07fbb --- /dev/null +++ b/dace/libraries/standard/nodes/copy/select.py @@ -0,0 +1,139 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Implementation selection for ``CopyLibraryNode``. +""" +from typing import Optional, TYPE_CHECKING + +import dace +from dace import dtypes, symbolic +from dace.libraries.standard.helper import (CPU_RESIDENT_STORAGES, collapse_shape_and_strides, is_in_parallel_scope, + is_parallel_cpu_transfer_size) +from dace.sdfg.scope import is_devicelevel_gpu, is_in_scope +from dace.libraries.standard.nodes.copy.common import (cuda2d_pitch_params, _both_packed_same_layout, _is_cross_cpu_gpu) + +if TYPE_CHECKING: + from dace.libraries.standard.nodes.copy.node import CopyLibraryNode + + +def select_copy_implementation(node: "CopyLibraryNode", parent_state: dace.SDFGState) -> str: + """Resolve ``CopyLibraryNode.implementation`` when set to ``'Auto'`` (the default); never + returns ``'Auto'`` itself. + + :param node: the :class:`CopyLibraryNode` being expanded. + :param parent_state: state containing ``node``. + :returns: a concrete implementation name from ``CopyLibraryNode.implementations``. + """ + inp_name, inp, in_subset, out_name, out, out_subset = node.validate(parent_state.sdfg, + parent_state, + allow_cross_storage=True) + + # A 0-D map crashes memlet propagation, so single-element copies use Tasklet/MemcpyCUDA1D. + single_elt = (in_subset.num_elements_exact() == 1 and out_subset.num_elements_exact() == 1) + + # GPU_Shared: SharedMemoryCollective, unless thread-level (Register endpoint or in a map). + # TODO: replace dace::CopyND with a vectorized 128-bit collective load. + if inp.storage == dtypes.StorageType.GPU_Shared or out.storage == dtypes.StorageType.GPU_Shared: + thread_level = (inp.storage == dtypes.StorageType.Register or out.storage == dtypes.StorageType.Register + or is_in_scope(parent_state.sdfg, parent_state, node, [dtypes.ScheduleType.GPU_ThreadBlock])) + if thread_level: + return 'Tasklet' if single_elt else 'MappedTasklet' + return 'SharedMemoryCollective' + + # Single-element non-Shared: MemcpyCUDA1D crossing CPU/GPU or GPU<->GPU from host; else Tasklet. + if single_elt: + # Device code cannot issue cudaMemcpyAsync at all, so an in-kernel single-element transfer + # is a plain assignment whichever storages it spans -- a host scalar reaches the kernel as + # a by-value argument, and CPU_Pinned is directly device-addressable. + inside_kernel = is_devicelevel_gpu(parent_state.sdfg, parent_state, node) + if inside_kernel: + return 'Tasklet' + if _is_cross_cpu_gpu(inp.storage, out.storage, node, parent_state): + return 'MemcpyCUDA1D' + both_gpu_global = (inp.storage == dtypes.StorageType.GPU_Global + and out.storage == dtypes.StorageType.GPU_Global) + if both_gpu_global: + return 'MemcpyCUDA1D' + return 'Tasklet' + + # cudaMemcpyAsync can't issue from device code, so in-kernel multi-element copies map instead. + if is_devicelevel_gpu(parent_state.sdfg, parent_state, node): + return 'MappedTasklet' + + # Host CPU-resident: same-shape/contiguous/same-layout below the parallel-transfer threshold + # is one MemcpyCPU; otherwise falls through to a parallel MappedTasklet. The threshold only + # applies where the copy runs once: inside a parallel map the mapped form is sequentialized to + # an element loop, which is strictly worse than the single call at any size. + host_storages = CPU_RESIDENT_STORAGES | {dtypes.StorageType.Default} + same_shape = (len(inp.shape) == len(out.shape) + and not any(symbolic.inequal_symbols(a, b) for a, b in zip(in_subset.size(), out_subset.size()))) + if ({inp.storage, out.storage} <= host_storages and same_shape and in_subset.is_contiguous_subset(inp) + and out_subset.is_contiguous_subset(out) and _both_packed_same_layout(inp, out) + and not (is_parallel_cpu_transfer_size(in_subset.num_elements()) + and not is_in_parallel_scope(node, parent_state))): + return 'MemcpyCPU' + + 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) + + if impl == 'MemcpyCUDA1D': + refined = _refine_cuda_impl_for_subsets(node, parent_state) + if refined is not None: + impl = refined + + # Rank-mismatched copies (e.g. (2,3,4) -> (8,3)) fall through to the MappedTasklet 1-D walker. + return impl or 'MappedTasklet' + + +def _refine_cuda_impl_for_subsets(node: "CopyLibraryNode", parent_state: dace.SDFGState) -> Optional[str]: + """Upgrade ``MemcpyCUDA1D`` to a more specific impl for non-contiguous subsets. + + both subsets contiguous -> ``None`` (keep CUDA1D) + collapsed rank 2, 2D pitched layout matches -> ``MemcpyCUDA2D`` + collapsed rank 1, both sides equal length -> ``MemcpyCUDA2D`` (degenerate ``(1, N)``) + same-side (no CPU/GPU boundary) -> ``MappedTasklet`` (per-element loop nest) + cross CPU/GPU, same rank, common stride-1 axis -> ``MemcpyCUDANDStrided`` (seq cudaMemcpyAsync/chunk) + cross CPU/GPU, no common stride-1 axis -> raise (no ``cudaMemcpy*`` lowering exists; + host can't issue cudaMemcpyAsync for + non-contiguous regions, device code can't + issue it at all) + + :param node: the :class:`CopyLibraryNode` being expanded. + :param parent_state: state containing ``node``. + :returns: the refined impl name, or ``None`` when both subsets are contiguous (keeps + ``MemcpyCUDA1D``). + :raises ValueError: cross-CPU/GPU strided pattern with no common stride-1 axis. + """ + _, inp, in_subset, _, out, out_subset = node.validate(parent_state.sdfg, parent_state, allow_cross_storage=True) + + if in_subset.is_contiguous_subset(inp) and out_subset.is_contiguous_subset(out): + return None + + in_shape_collapsed, in_strides_collapsed = collapse_shape_and_strides(in_subset, inp.strides) + out_shape_collapsed, out_strides_collapsed = collapse_shape_and_strides(out_subset, out.strides) + + src_rank, dst_rank = len(in_shape_collapsed), len(out_shape_collapsed) + if src_rank == 2 and dst_rank == 2: + # Shared with the expander so selector and expander cannot disagree. + if cuda2d_pitch_params(in_shape_collapsed, in_strides_collapsed, out_strides_collapsed) is not None: + return 'MemcpyCUDA2D' + + elif src_rank == 1 and dst_rank == 1: + # Degenerate (1, N) case: neither side needs stride-1, e.g. `a[:, 1] = b[4, :]` (C order). + if not symbolic.inequal_symbols(in_shape_collapsed[0], out_shape_collapsed[0]): + return 'MemcpyCUDA2D' + + if not _is_cross_cpu_gpu(inp.storage, out.storage, node, parent_state): + return 'MappedTasklet' + + if (len(in_shape_collapsed) == len(out_shape_collapsed) and len(in_shape_collapsed) >= 1 + and any(in_strides_collapsed[d] == 1 and out_strides_collapsed[d] == 1 + for d in range(len(in_shape_collapsed)))): + return 'MemcpyCUDANDStrided' + + raise ValueError(f"CopyLibraryNode '{node.name}' has a strided cross-CPU/GPU copy pattern that " + f"cannot be lowered to a single cudaMemcpy or cudaMemcpy2DAsync and has no " + f"common stride-1 axis for chunked memcpy " + f"(src_shape={in_shape_collapsed}, src_strides={in_strides_collapsed}, " + f"dst_shape={out_shape_collapsed}, dst_strides={out_strides_collapsed}); " + f"pick an explicit implementation manually.") diff --git a/dace/libraries/standard/nodes/fill/__init__.py b/dace/libraries/standard/nodes/fill/__init__.py new file mode 100644 index 0000000000..529fed1cef --- /dev/null +++ b/dace/libraries/standard/nodes/fill/__init__.py @@ -0,0 +1,5 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +from dace.libraries.standard.nodes.fill.common import byte_pattern, cpp_literal, python_literal +from dace.libraries.standard.nodes.fill.node import FillLibraryNode +from dace.libraries.standard.nodes.fill.select import select_fill_implementation +from dace.libraries.standard.nodes.fill.expansions import (ExpandAuto, ExpandCPU, ExpandCUDA, ExpandPure, ExpandTasklet) diff --git a/dace/libraries/standard/nodes/fill/common.py b/dace/libraries/standard/nodes/fill/common.py new file mode 100644 index 0000000000..8e16d10083 --- /dev/null +++ b/dace/libraries/standard/nodes/fill/common.py @@ -0,0 +1,91 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Shared pieces of the ``FillLibraryNode`` expansions. + +Imported by both the node and its expansions, so it must not import either. +""" +from typing import List, Optional, Tuple, TYPE_CHECKING + +import numpy as np + +import dace +from dace.libraries.standard.helper import collapse_shape_and_strides + +if TYPE_CHECKING: + from dace.libraries.standard.nodes.fill.node import FillLibraryNode + +OUTPUT_CONNECTOR_NAME = "_fill_out" + + +def numpy_scalar(value, dtype: dace.dtypes.typeclass) -> np.generic: + """The fill value as a numpy scalar of the destination type. + + :param value: The Python constant held by the node. + :param dtype: Destination element type. + :returns: ``value`` narrowed to ``dtype``. + """ + return np.array(value, dtype=dtype.as_numpy_dtype())[()] + + +def byte_pattern(value, dtype: dace.dtypes.typeclass) -> Optional[int]: + """The single byte a ``memset`` would need, or ``None`` if the value is not byte-splat. + + ``memset`` writes one byte over the whole range, so it can only express values whose object + representation repeats that byte -- every zero, ``-1`` in two's complement, any 1-byte type. + ``1.0f`` (``0000803f``) is not one of them. + + :param value: The Python constant held by the node. + :param dtype: Destination element type. + :returns: The byte to pass to ``memset``, or ``None``. + """ + raw = numpy_scalar(value, dtype).tobytes() + return raw[0] if len(set(raw)) == 1 else None + + +def cpp_literal(value, dtype: dace.dtypes.typeclass) -> str: + """Render the fill value as a C++ literal of ``dtype``. + + :param value: The Python constant held by the node. + :param dtype: Destination element type. + :returns: A C++ expression of type ``dtype.ctype``. + """ + narrowed = numpy_scalar(value, dtype).item() + ctype = dtype.ctype + if isinstance(narrowed, bool): + return 'true' if narrowed else 'false' + if isinstance(narrowed, complex): + return f"{ctype}({narrowed.real!r}, {narrowed.imag!r})" + if isinstance(narrowed, float): + return f"{narrowed!r}f" if ctype == 'float' else repr(narrowed) + return repr(narrowed) + + +def python_literal(value, dtype: dace.dtypes.typeclass) -> str: + """Render the fill value for a Python-language tasklet body. + + :param value: The Python constant held by the node. + :param dtype: Destination element type. + :returns: A Python expression. + """ + return repr(numpy_scalar(value, dtype).item()) + + +def make_fill_skeleton(node: "FillLibraryNode", + parent_state: dace.SDFGState) -> Tuple[dace.SDFG, dace.SDFGState, str, dace.data.Data, List]: + """Build the shared SDFG skeleton for the mapped (``ExpandPure``) fill expansion. + + :param node: The fill library node being expanded. + :param parent_state: The state containing ``node`` (owning SDFG is ``parent_state.sdfg``). + :returns: ``(sdfg, state, out_name, out, map_lengths)``. + """ + out_name, out, out_subset = node.validate(parent_state.sdfg, parent_state) + out_shape_collapsed, out_strides_collapsed = collapse_shape_and_strides(out_subset, out.strides) + + sdfg = dace.SDFG(f"{node.label}_sdfg") + sdfg.add_array(out_name, out_shape_collapsed, out.dtype, out.storage, strides=out_strides_collapsed) + sdfg.schedule = dace.dtypes.ScheduleType.Sequential + + state = sdfg.add_state(f"{node.label}_state") + # Reuse the array descriptor's collapsed shape as map bounds so extents can't diverge. + map_lengths = out_shape_collapsed + + return sdfg, state, out_name, out, map_lengths diff --git a/dace/libraries/standard/nodes/fill/expansions/__init__.py b/dace/libraries/standard/nodes/fill/expansions/__init__.py new file mode 100644 index 0000000000..c039a724ea --- /dev/null +++ b/dace/libraries/standard/nodes/fill/expansions/__init__.py @@ -0,0 +1,7 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Every ``FillLibraryNode`` expansion. Imported here so registration runs on package import.""" +from dace.libraries.standard.nodes.fill.expansions.auto import ExpandAuto +from dace.libraries.standard.nodes.fill.expansions.cpu import ExpandCPU +from dace.libraries.standard.nodes.fill.expansions.cuda import ExpandCUDA +from dace.libraries.standard.nodes.fill.expansions.mapped_tasklet import ExpandPure +from dace.libraries.standard.nodes.fill.expansions.tasklet import ExpandTasklet diff --git a/dace/libraries/standard/nodes/fill/expansions/auto.py b/dace/libraries/standard/nodes/fill/expansions/auto.py new file mode 100644 index 0000000000..cb40285d30 --- /dev/null +++ b/dace/libraries/standard/nodes/fill/expansions/auto.py @@ -0,0 +1,22 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Default expansion: dispatches via :func:`select_fill_implementation`.""" +from typing import TYPE_CHECKING + +import dace +from dace import library +from dace.libraries.standard.helper import auto_dispatch +from dace.libraries.standard.nodes.fill.select import select_fill_implementation +from dace.transformation.transformation import ExpandTransformation + +if TYPE_CHECKING: + from dace.libraries.standard.nodes.fill.node import FillLibraryNode + + +@library.expansion +class ExpandAuto(ExpandTransformation): + environments = [] + + @staticmethod + def expansion(node: "FillLibraryNode", parent_state: dace.SDFGState, parent_sdfg: dace.SDFG): + from dace.libraries.standard.nodes.fill.node import FillLibraryNode # Avoid import loop + return auto_dispatch(node, parent_state, select_fill_implementation, FillLibraryNode) diff --git a/dace/libraries/standard/nodes/fill/expansions/cpu.py b/dace/libraries/standard/nodes/fill/expansions/cpu.py new file mode 100644 index 0000000000..0fa8498783 --- /dev/null +++ b/dace/libraries/standard/nodes/fill/expansions/cpu.py @@ -0,0 +1,39 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Single-call host fill through ``std::fill_n``.""" +from typing import TYPE_CHECKING + +import dace +from dace import library, nodes +from dace.codegen.common import sym2cpp +from dace.libraries.standard import environments +from dace.libraries.standard.nodes.fill.common import OUTPUT_CONNECTOR_NAME, cpp_literal +from dace.transformation.transformation import ExpandTransformation + +if TYPE_CHECKING: + from dace.libraries.standard.nodes.fill.node import FillLibraryNode + + +@library.expansion +class ExpandCPU(ExpandTransformation): + environments = [environments.CPU] + + @staticmethod + def expansion(node: "FillLibraryNode", parent_state: dace.SDFGState, parent_sdfg: dace.SDFG) -> nodes.Tasklet: + out_name, out, out_subset = node.validate(parent_state.sdfg, parent_state) + if not out_subset.is_contiguous_subset(out): + raise ValueError(f"FillLibraryNode CPU expansion requires a contiguous subset; got '{out_name}' " + f"subset {out_subset} on shape {tuple(out.shape)} strides {tuple(out.strides)}. " + f"Use the 'pure' expansion (mapped tasklet) for non-contiguous regions.") + + # Both gcc and clang turn this into a memset at -O2 and above whenever the value's object + # representation allows it, including through the one-byte fp8 wrappers, so spelling the + # memset out here would only repeat what the build already does at the Release level dace + # always compiles with. + count = sym2cpp(out_subset.num_elements_exact()) + code = f"std::fill_n({OUTPUT_CONNECTOR_NAME}, {count}, {cpp_literal(node.value, out.dtype)});" + + return nodes.Tasklet(node.name, + inputs={}, + outputs={OUTPUT_CONNECTOR_NAME: dace.dtypes.pointer(out.dtype)}, + code=code, + language=dace.Language.CPP) diff --git a/dace/libraries/standard/nodes/fill/expansions/cuda.py b/dace/libraries/standard/nodes/fill/expansions/cuda.py new file mode 100644 index 0000000000..8979a23adc --- /dev/null +++ b/dace/libraries/standard/nodes/fill/expansions/cuda.py @@ -0,0 +1,44 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Host-issued ``MemsetAsync`` over GPU memory. Byte-splat values only.""" +from typing import TYPE_CHECKING + +import dace +from dace import library, nodes +from dace.codegen.common import sym2cpp, get_gpu_backend +from dace.libraries.standard import environments +from dace.libraries.standard.helper import CURRENT_STREAM_NAME +from dace.libraries.standard.nodes.fill.common import OUTPUT_CONNECTOR_NAME, byte_pattern +from dace.transformation.transformation import ExpandTransformation + +if TYPE_CHECKING: + from dace.libraries.standard.nodes.fill.node import FillLibraryNode + + +@library.expansion +class ExpandCUDA(ExpandTransformation): + environments = [environments.CUDA] + + @staticmethod + def expansion(node: "FillLibraryNode", parent_state: dace.SDFGState, parent_sdfg: dace.SDFG) -> nodes.Tasklet: + out_name, out, out_subset = node.validate(parent_state.sdfg, parent_state) + if not out_subset.is_contiguous_subset(out): + raise ValueError(f"FillLibraryNode CUDA expansion requires a contiguous subset; got '{out_name}' " + f"subset {out_subset} on shape {tuple(out.shape)} strides {tuple(out.strides)}. " + f"Use the 'pure' expansion (mapped tasklet) for non-contiguous regions.") + + pattern = byte_pattern(node.value, out.dtype) + if pattern is None: + raise ValueError(f"FillLibraryNode CUDA expansion requires a byte-splat value; {node.value!r} as " + f"{out.dtype} is not one. Use the 'pure' expansion (mapped tasklet).") + + nbytes = f"{sym2cpp(out_subset.num_elements_exact())} * sizeof({out.dtype.ctype})" + # The API name follows the configured backend (cuda/hip), like every sibling copy expansion. + backend = get_gpu_backend() + code = (f"DACE_GPU_CHECK({backend}MemsetAsync({OUTPUT_CONNECTOR_NAME}, {pattern}, {nbytes}, " + f"{CURRENT_STREAM_NAME}));") + + return nodes.Tasklet(node.name, + inputs={}, + outputs={OUTPUT_CONNECTOR_NAME: dace.dtypes.pointer(out.dtype)}, + code=code, + language=dace.Language.CPP) diff --git a/dace/libraries/standard/nodes/fill/expansions/mapped_tasklet.py b/dace/libraries/standard/nodes/fill/expansions/mapped_tasklet.py new file mode 100644 index 0000000000..a4b0c19a1d --- /dev/null +++ b/dace/libraries/standard/nodes/fill/expansions/mapped_tasklet.py @@ -0,0 +1,40 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Element-wise mapped fill: the device-neutral parallel form, and the fallback for everything a +single call cannot express (non-contiguous subsets, non-byte-splat GPU values, device scope).""" +from typing import TYPE_CHECKING + +import dace +from dace import library +from dace.libraries.standard.nodes.fill.common import make_fill_skeleton, python_literal +from dace.transformation.transformation import ExpandTransformation + +if TYPE_CHECKING: + from dace.libraries.standard.nodes.fill.node import FillLibraryNode + + +@library.expansion +class ExpandPure(ExpandTransformation): + environments = [] + + @staticmethod + def expansion(node: "FillLibraryNode", parent_state: dace.SDFGState, parent_sdfg: dace.SDFG) -> dace.SDFG: + sdfg, state, out_name, out, map_lengths = make_fill_skeleton(node, parent_state) + + # Must not collide with the wrapper SDFG's parameter array (named after outer connector). + inner_out = "_out" + map_params = [f"__i{i}" for i in range(len(map_lengths))] + # Symbolic bounds, never a rendered string: the range parser splits on ':', so any extent + # whose spelling carries one comes back as bogus tokens. + map_rng = {i: (0, s - 1, 1) for i, s in zip(map_params, map_lengths)} + outputs = {inner_out: dace.memlet.Memlet(f"{out_name}[{','.join(map_params)}]")} + schedule = (dace.dtypes.ScheduleType.GPU_Device + if out.storage == dace.dtypes.StorageType.GPU_Global else dace.dtypes.ScheduleType.Default) + state.add_mapped_tasklet(f"{node.label}_tasklet", + map_rng, + dict(), + f"{inner_out} = {python_literal(node.value, out.dtype)}", + outputs, + schedule=schedule, + external_edges=True) + + return sdfg diff --git a/dace/libraries/standard/nodes/fill/expansions/tasklet.py b/dace/libraries/standard/nodes/fill/expansions/tasklet.py new file mode 100644 index 0000000000..7a8ea0586b --- /dev/null +++ b/dace/libraries/standard/nodes/fill/expansions/tasklet.py @@ -0,0 +1,37 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Single-element same-side scalar assignment.""" +from typing import TYPE_CHECKING + +import dace +from dace import library, nodes +from dace.libraries.standard.helper import GPU_RESIDENT_STORAGES +from dace.libraries.standard.nodes.fill.common import OUTPUT_CONNECTOR_NAME, python_literal +from dace.sdfg.scope import is_devicelevel_gpu +from dace.transformation.transformation import ExpandTransformation + +if TYPE_CHECKING: + from dace.libraries.standard.nodes.fill.node import FillLibraryNode + + +@library.expansion +class ExpandTasklet(ExpandTransformation): + environments = [] + + @staticmethod + def expansion(node: "FillLibraryNode", parent_state: dace.SDFGState, parent_sdfg: dace.SDFG) -> nodes.Tasklet: + out_name, out, out_subset = node.validate(parent_sdfg, parent_state) + out_volume = out_subset.num_elements_exact() + if out_volume != 1: + raise ValueError(f"Tasklet expansion requires single-element subsets " + f"(got output volume {out_volume}). Use 'pure' for multi-element fills.") + + # Host scope can't write device memory directly; route that case to 'CUDA' instead. + if (not is_devicelevel_gpu(parent_state.sdfg, parent_state, node) and out.storage in GPU_RESIDENT_STORAGES): + raise ValueError(f"Tasklet expansion cannot fill GPU-resident storage ({out.storage}) for " + f"'{out_name}' from host scope; use the 'CUDA' Fill expansion instead.") + + return nodes.Tasklet(node.name, + inputs={}, + outputs={OUTPUT_CONNECTOR_NAME: out.dtype}, + code=f"{OUTPUT_CONNECTOR_NAME} = {python_literal(node.value, out.dtype)}", + language=dace.Language.Python) diff --git a/dace/libraries/standard/nodes/fill/node.py b/dace/libraries/standard/nodes/fill/node.py new file mode 100644 index 0000000000..c51e5285bd --- /dev/null +++ b/dace/libraries/standard/nodes/fill/node.py @@ -0,0 +1,65 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""``FillLibraryNode``: write one constant over a contiguous output subset.""" +from typing import Tuple + +import dace +from dace import library, nodes, properties +from dace.libraries.standard.helper import CURRENT_STREAM_NAME +from dace.libraries.standard.nodes.fill.common import OUTPUT_CONNECTOR_NAME +from dace.libraries.standard.nodes.fill.expansions import (ExpandAuto, ExpandCPU, ExpandCUDA, ExpandPure, ExpandTasklet) + + +@library.node +class FillLibraryNode(nodes.LibraryNode): + """Library node writing a constant over a contiguous output subset. + + Does NOT accept dynamic (Scalar) input connectors: subset expressions must use symbols + already in scope, so the auto selector reasons purely from the static memlet subset. + """ + + implementations = { + "Auto": ExpandAuto, + "pure": ExpandPure, + "CUDA": ExpandCUDA, + "CPU": ExpandCPU, + "tasklet": ExpandTasklet + } + default_implementation = 'Auto' + + OUTPUT_CONNECTOR_NAME = OUTPUT_CONNECTOR_NAME + + # dtype=None takes any Python constant; numpy scalars are normalized by Property.__set__. + value = properties.Property(dtype=None, default=0, desc='The constant written over the subset.') + + def __init__(self, name: str, *args, value=0, **kwargs): + # Dotted structure-member data names reach here through the callers that build the label; + # the label names the wrapper SDFG, i.e. a C++ function. See CopyLibraryNode.__init__. + super().__init__(name.replace('.', '_'), *args, outputs={OUTPUT_CONNECTOR_NAME}, **kwargs) + self.value = value + + def validate(self, sdfg: dace.SDFG, state: dace.SDFGState) -> Tuple[str, dace.data.Data, dace.subsets.Range]: + """Validate wiring and resolve the output edge. + + :param sdfg: The SDFG owning the data descriptors. + :param state: The state containing this node. + :returns: ``(out_name, out, out_subset)``. + :raises ValueError: If the node lacks exactly one output edge, or has a non-empty + non-reserved input connector wired. + """ + data_oes = [oe for oe in state.out_edges(self) if oe.src_conn == OUTPUT_CONNECTOR_NAME] + if len(data_oes) != 1: + raise ValueError(f"{type(self).__name__} expects exactly one " + f"``{OUTPUT_CONNECTOR_NAME}`` output edge.") + + reserved = {CURRENT_STREAM_NAME} + extra = [ie.dst_conn for ie in state.in_edges(self) if ie.dst_conn not in reserved and not ie.data.is_empty()] + if extra: + raise ValueError(f"{type(self).__name__} does not accept dynamic input connectors; got {extra}. " + f"Subset expressions must use symbols already in scope.") + + oe = data_oes[0] + out = sdfg.arrays[oe.data.data] + out_subset = oe.data.subset + out_name = oe.src_conn + + return out_name, out, out_subset diff --git a/dace/libraries/standard/nodes/fill/select.py b/dace/libraries/standard/nodes/fill/select.py new file mode 100644 index 0000000000..3b7e4f79f5 --- /dev/null +++ b/dace/libraries/standard/nodes/fill/select.py @@ -0,0 +1,53 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Implementation selection for ``FillLibraryNode``.""" +from typing import TYPE_CHECKING + +import dace +from dace.libraries.standard.helper import (CPU_RESIDENT_STORAGES, is_in_parallel_scope, is_parallel_cpu_transfer_size) +from dace.libraries.standard.nodes.fill.common import byte_pattern +from dace.sdfg.scope import is_devicelevel_gpu + +if TYPE_CHECKING: + from dace.libraries.standard.nodes.fill.node import FillLibraryNode + + +def select_fill_implementation(node: "FillLibraryNode", parent_state: dace.SDFGState) -> str: + """Resolve an ``'Auto'`` ``FillLibraryNode`` implementation to a concrete one. + + ``'pure'``: device scope (no ``cudaMemsetAsync`` from a kernel), non-contiguous subsets, a GPU + destination whose value is not byte-splat, or a contiguous CPU fill that is not provably + sub-threshold. ``'CUDA'``: host-issued byte-splat fill of GPU memory. ``'CPU'``: a single + ``memset``/``std::fill_n``. ``'tasklet'``: a single element. + + :param node: The fill library node being expanded. + :param parent_state: The state containing ``node`` (owning SDFG is ``parent_state.sdfg``). + :returns: One of ``'pure'``, ``'CUDA'``, ``'CPU'``, or ``'tasklet'``. + """ + _out_name, out, out_subset = node.validate(parent_state.sdfg, parent_state) + + if is_devicelevel_gpu(parent_state.sdfg, parent_state, node): + if out_subset.num_elements_exact() == 1: + return 'tasklet' + return 'pure' + + if out_subset.num_elements_exact() == 1 and (out.storage in CPU_RESIDENT_STORAGES + or out.storage == dace.dtypes.StorageType.Register): + return 'tasklet' + + if not out_subset.is_contiguous_subset(out): + return 'pure' + + if out.storage == dace.dtypes.StorageType.GPU_Global: + # cudaMemsetAsync writes ONE byte over the range; dace links only the CUDA runtime API, and + # the 32-bit setter (cuMemsetD32Async) is driver-API. A non-byte-splat value fills by kernel. + return 'CUDA' if byte_pattern(node.value, out.dtype) is not None else 'pure' + + # CPU main-memory fill: the element map ('pure') is the DEFAULT, taken unless the count is + # PROVABLY below parallel_transfer_min_elements, which keeps the single call ('CPU'). A symbolic + # count is assumed big. Inside a parallel map the element map is sequentialized anyway, so the + # single call wins there at any size. + allowed = CPU_RESIDENT_STORAGES | {dace.dtypes.StorageType.Default} + if (out.storage in allowed and is_parallel_cpu_transfer_size(out_subset.num_elements()) + and not is_in_parallel_scope(node, parent_state)): + return 'pure' + return 'CPU' diff --git a/dace/sdfg/infer_types.py b/dace/sdfg/infer_types.py index 4666e7af31..07490b503a 100644 --- a/dace/sdfg/infer_types.py +++ b/dace/sdfg/infer_types.py @@ -261,6 +261,12 @@ def _determine_schedule_from_storage(state: SDFGState, node: nodes.Node) -> Opti continue constraints.add(sched) + # Copy/Memset library nodes legitimately bridge storages; schedule on the GPU if involved. + from dace.libraries.standard.nodes.copy import CopyLibraryNode + from dace.libraries.standard.nodes.fill import FillLibraryNode + if isinstance(node, (CopyLibraryNode, FillLibraryNode)) and dtypes.ScheduleType.GPU_Device in constraints: + return dtypes.ScheduleType.GPU_Device + if not constraints: # No constraints found child_schedule = None elif len(constraints) > 1: diff --git a/dace/sdfg/nodes.py b/dace/sdfg/nodes.py index 086aa3b2c0..addd00e150 100644 --- a/dace/sdfg/nodes.py +++ b/dace/sdfg/nodes.py @@ -1352,6 +1352,11 @@ class LibraryNode(CodeNode): "the node upon expansion, if expanded to a nested SDFG.", default=dtypes.ScheduleType.Default) debuginfo = DebugInfoProperty(allow_none=True) + # Codegen dispatches ``on_node_begin``/``on_node_end`` for a library node like any other code + # node, and expansion carries this onto whatever the node expands into. + instrument = EnumProperty(dtype=dtypes.InstrumentationType, + desc="Measure execution statistics with given method", + default=dtypes.InstrumentationType.No_Instrumentation) def __init__(self, name, *args, schedule=None, **kwargs): super().__init__(*args, **kwargs) diff --git a/dace/sdfg/state.py b/dace/sdfg/state.py index aae219c20a..d4f9b82afd 100644 --- a/dace/sdfg/state.py +++ b/dace/sdfg/state.py @@ -921,17 +921,18 @@ def unordered_arglist(self, } if top_source_edge.src.data not in descs else {}) elif isinstance(edge.dst, nd.ExitNode) and isinstance(edge.src, (nd.AccessNode, nd.CodeNode)): - # Same case as above, but for outgoing Memlets. Every edge on the matching connector - # is inspected, since the data can go to more than one destination, and each is - # followed to where it lands: one hop still names the inner transient whenever the - # write leaves through more than one exit, as it does in a tiled map. + # Outgoing counterpart of the above. A source-relative Memlet's .data names the + # inner transient, not the written array, so resolve the real destination via the + # memlet-tree root -- else its shape/stride symbols drop from the kernel signature. additional_descs = {} connector_to_look = "OUT_" + edge.dst_conn[3:] for oedge in self.graph.out_edges_by_connector(edge.dst, connector_to_look): - outermost = self.graph.memlet_path(oedge)[-1].data - if ((not outermost.is_empty()) and (outermost.data not in descs) - and (outermost.data not in additional_descs)): - additional_descs[outermost.data] = sdfg.arrays[outermost.data] + if oedge.data.is_empty(): + continue + root_dst = self.graph.memlet_tree(oedge).root().edge.dst + dst_name = root_dst.data if isinstance(root_dst, nd.AccessNode) else oedge.data.data + if dst_name not in descs and dst_name not in additional_descs: + additional_descs[dst_name] = sdfg.arrays[dst_name] else: # Case is ignored. diff --git a/dace/subsets.py b/dace/subsets.py index b93c955b4b..57ab0aa00e 100644 --- a/dace/subsets.py +++ b/dace/subsets.py @@ -974,9 +974,9 @@ def is_contiguous_subset(self, array: 'dace.data.Array') -> bool: array: array descriptor to check against Returns: - True if the subset is contiguous, False otherwise - Returns False on all arrays that are not have a packed layout, - meaning that the complete array is contiguously stored in 1D memory. + True if the subset addresses one uninterrupted run of memory: the whole array has a + packed layout, or -- even on a non-packed (padded) descriptor -- the subset is a 1D + slice (at most one dimension has size > 1, and that dimension has stride 1). """ # Any step size != 1 -> not contiguous if any(s != 1 for (_, _, s) in self): diff --git a/dace/transformation/passes/insert_explicit_copies.py b/dace/transformation/passes/insert_explicit_copies.py new file mode 100644 index 0000000000..74890d7ee8 --- /dev/null +++ b/dace/transformation/passes/insert_explicit_copies.py @@ -0,0 +1,314 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Pass replacing implicit copy patterns with explicit ``CopyLibraryNode`` instances.""" +import copy +from typing import Any, Dict, Optional + +from dace import data, dtypes, nodes, properties, subsets, symbolic +from dace.memlet import Memlet +from dace.sdfg import SDFG +from dace.sdfg import utils as sdutils +from dace.sdfg.state import SDFGState +from dace.transformation import pass_pipeline as ppl, transformation +from dace.libraries.standard.helper import CPU_RESIDENT_STORAGES, GPU_RESIDENT_STORAGES +from dace.libraries.standard.nodes.copy import CopyLibraryNode + + +def _derive_matching_dst_subset(src_subset: subsets.Range, dst_desc: data.Data) -> subsets.Range: + """Derive the absent side of a copy memlet: the full array when volumes are not + provably unequal, else ``src_subset``. + + :param src_subset: the known (source) side of the copy. + :param dst_desc: descriptor whose subset is being derived. + :returns: the destination :class:`~dace.subsets.Range`. + """ + dst_range = subsets.Range.from_array(dst_desc) + # Equalize first: two instances of the same symbol name make equal() answer None on identical counts. + src_count, dst_count = symbolic.equalize_symbols(src_subset.num_elements(), dst_range.num_elements()) + if symbolic.equal(src_count, dst_count) is not False: + return dst_range + return src_subset + + +def _competing_writer(state: SDFGState, target: nodes.Node, edge, name: str, subset: subsets.Subset) -> bool: + """True if another edge into ``target`` writes a region of ``name`` that may overlap ``subset``. + + Nothing in the graph orders two writes to the same region that reach a node on separate edges: + plain copy-edge codegen emits a copy when its SOURCE access node is visited, so the copy lands + before every other consumer of that node. Lifting the copy to a node of its own re-sorts it + against the competing write, which silently swaps which value survives (measured on npbench + ``vadv``: a dead ``dcol`` write moved after the tasklet that supersedes it). Where the order + cannot be shown to be irrelevant, leave the copy implicit. + + :param state: the state holding ``target``. + :param target: the node the copy writes through (access node, or map exit when staging out). + :param edge: the copy edge itself, excluded from the scan. + :param name: data name the copy writes. + :param subset: region the copy writes. + :returns: ``True`` when a possibly-overlapping competing write exists. + """ + for other in state.in_edges(target): + if other is edge or other.data.is_empty() or other.data.data != name: + continue + other_subset = other.data.get_dst_subset(other, state) or other.data.subset + if subsets.intersects(other_subset, subset) is not False: + return True + return False + + +def _carry_write_ordering(state: SDFGState, written: nodes.AccessNode, libnode: nodes.Node) -> None: + """Repeat onto ``libnode`` the ordering edges that sequenced writes to ``written``. + + An empty memlet is a happens-before edge, and the write it constrained is no longer the access + node's own -- it is the libnode's. Left behind, it orders a node that no longer writes anything, + and the libnode is free to be scheduled ahead of the write it was supposed to follow. + + :param state: the state both nodes live in. + :param written: the access node the libnode now writes. + :param libnode: the inserted copy node. + """ + for edge in state.in_edges(written): + if not edge.data.is_empty() or edge.src is libnode: + continue + if any(existing.src is edge.src for existing in state.in_edges(libnode)): + continue + state.add_edge(edge.src, None, libnode, None, Memlet()) + + +@properties.make_properties +@transformation.explicit_cf_compatible +class InsertExplicitCopies(ppl.Pass): + """Replaces implicit copy patterns with ``CopyLibraryNode`` instances: direct + ``AccessNode -> AccessNode`` edges (View endpoints included), and stage-in/stage-out through + chained ``MapEntry``/``MapExit`` (libnode placed inside the map scope).""" + + # Other storages (TensorCore_*, FPGA_*, Snitch_*) use their own codegen ``copy_memory`` hook. + _STANDARD_STORAGES = (CPU_RESIDENT_STORAGES | GPU_RESIDENT_STORAGES + | {dtypes.StorageType.Default, dtypes.StorageType.Register}) + + def modifies(self) -> ppl.Modifies: + return ppl.Modifies.States | ppl.Modifies.Nodes | ppl.Modifies.Edges + + def should_reapply(self, modified: ppl.Modifies) -> bool: + return False + + def depends_on(self): + return set() + + def apply_pass(self, sdfg: SDFG, pipeline_results: Dict[str, Any]) -> Optional[int]: + """Lift every implicit copy in ``sdfg`` (and nested SDFGs) to a ``CopyLibraryNode``. + + :param sdfg: The SDFG to transform, recursively including nested SDFGs. + :param pipeline_results: Results of previously applied passes (unused). + :returns: The number of copy nodes inserted, or ``None`` if none. + """ + count = 0 + for nsdfg in sdfg.all_sdfgs_recursive(): + for state in nsdfg.states(): + count += self._replace_direct_copies(state) + count += self._replace_map_staging_copies(state) + return count if count > 0 else None + + def _replace_direct_copies(self, state: SDFGState) -> int: + """Replace direct ``AccessNode -> AccessNode`` edges with ``CopyLibraryNode`` instances. + + :param state: The state to scan for direct copy edges (owning SDFG is ``state.sdfg``). + :returns: The number of copy nodes inserted in ``state``. + """ + sdfg = state.sdfg + edges = list(state.edges()) + count = 0 + for edge in edges: + if not (isinstance(edge.src, nodes.AccessNode) and isinstance(edge.dst, nodes.AccessNode)): + continue + + src_node: nodes.AccessNode = edge.src + dst_node: nodes.AccessNode = edge.dst + memlet: Memlet = edge.data + + if memlet.is_empty(): + continue + + # WCR edges aren't copies. + if memlet.wcr is not None: + continue + + # A reference-set edge assigns a POINTER; rewriting it would drop the ``set`` connector. + if edge.dst_conn == 'set': + continue + + src_desc = sdfg.arrays[src_node.data] + dst_desc = sdfg.arrays[dst_node.data] + + # A view's alias (view-defining) edge references the underlying buffer, not data -- skip. + if any( + isinstance(sdfg.arrays[an.data], data.View) and sdutils.get_view_edge(state, an) is edge + for an in (src_node, dst_node)): + continue + + if not isinstance(src_desc, (data.Array, data.Scalar)) \ + or not isinstance(dst_desc, (data.Array, data.Scalar)): + continue + + # Custom-target storages are handled by their own codegen, not CopyLibraryNode. + if (src_desc.storage not in self._STANDARD_STORAGES or dst_desc.storage not in self._STANDARD_STORAGES): + continue + + # A dtype-converting copy is a cast, not a byte move: CopyLibraryNode (memcpy) + # cannot express it, so leave it for tasklet lowering (mirrors _lift_staging_edge). + if src_desc.dtype != dst_desc.dtype: + continue + + src_name = src_node.data + dst_name = dst_node.data + + # Self-copy: subset is the dst side; otherwise the memlet path maps ``data`` to an endpoint. + if src_name == dst_name: + src_subset, dst_subset = memlet.other_subset, memlet.subset + else: + src_subset = memlet.get_src_subset(edge, state) + dst_subset = memlet.get_dst_subset(edge, state) + + # Derive any side the memlet omitted from the array shape (same-volume, different-shape). + if src_subset is None: + src_subset = _derive_matching_dst_subset(dst_subset, src_desc) + if dst_subset is None: + dst_subset = _derive_matching_dst_subset(src_subset, dst_desc) + + # A copy of zero elements moves nothing, and plain copy-edge codegen emits nothing for + # it. Lifting it would put a node in the state that has no work to do. + if symbolic.equal(src_subset.num_elements(), 0, is_length=False) is True: + continue + + if _competing_writer(state, dst_node, edge, dst_name, dst_subset): + continue + + in_memlet = Memlet(data=src_name, subset=copy.deepcopy(src_subset)) + in_memlet.dynamic = memlet.dynamic + out_memlet = Memlet(data=dst_name, subset=copy.deepcopy(dst_subset)) + out_memlet.dynamic = memlet.dynamic + # ``allow_oob`` is the author's waiver of the src/dst volume check (``validation.py`` + # honours it the same way); dropping it here turns a legal copy into an expansion error. + in_memlet.allow_oob = memlet.allow_oob + out_memlet.allow_oob = memlet.allow_oob + + label = f"copy_{src_name}_to_{dst_name}" + libnode = CopyLibraryNode(name=label) + # Instrumentation providers decide a copy edge's instrumentation from the state it is in + # (``on_copy_begin``); as a node the copy needs its own setting to stay measured. + libnode.instrument = state.instrument + + state.remove_edge(edge) + state.add_node(libnode) + state.add_edge(src_node, None, libnode, CopyLibraryNode.INPUT_CONNECTOR_NAME, in_memlet) + state.add_edge(libnode, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, dst_node, None, out_memlet) + _carry_write_ordering(state, dst_node, libnode) + count += 1 + + return count + + def _replace_map_staging_copies(self, state: SDFGState) -> int: + """Lift stage-in / stage-out copies through ``MapEntry`` / ``MapExit`` to ``CopyLibraryNode``. + + The libnode sits inside the map scope; chained MapEntries / MapExits are followed via + ``memlet_path``. + + :param state: The state to scan (owning SDFG is ``state.sdfg``). + :returns: Number of libnodes inserted. + """ + count = 0 + for node in state.nodes(): + if isinstance(node, nodes.MapEntry): + for edge in list(state.out_edges(node)): + if self._lift_staging_edge(state, edge, stage_in=True): + count += 1 + elif isinstance(node, nodes.MapExit): + for edge in list(state.in_edges(node)): + if self._lift_staging_edge(state, edge, stage_in=False): + count += 1 + return count + + def _lift_staging_edge(self, state: SDFGState, edge, stage_in: bool) -> bool: + """Lift one stage-in (``stage_in=True``) or stage-out copy edge to a libnode. + + :returns: True iff the edge was lifted. + """ + sdfg = state.sdfg + inner_node = edge.dst if stage_in else edge.src + if not isinstance(inner_node, nodes.AccessNode) or edge.data.is_empty(): + return False + # A WCR edge isn't a copy -- it's a reduction (e.g. AccumulateTransient's tile merge back + # into the real output). CopyLibraryNode's expansions (ExpandMemcpyCPU et al.) always emit + # an unconditional store; lifting a WCR edge here would silently turn the accumulate into + # an overwrite. Mirrors the same guard in ``_replace_direct_copies``. + if edge.data.wcr is not None: + return False + # A reference-set edge binds a POINTER rather than moving data; lifting it would drop the + # ``set`` connector and leave the Reference unbound. + if edge.dst_conn == 'set': + return False + inner_desc = sdfg.arrays[inner_node.data] + if isinstance(inner_desc, data.View): + return False + find_outer = sdutils.find_input_arraynode if stage_in else sdutils.find_output_arraynode + try: + outer = find_outer(state, edge) + except RuntimeError: + return False + outer_desc = sdfg.arrays[outer.data] + if (outer_desc.storage not in self._STANDARD_STORAGES or inner_desc.storage not in self._STANDARD_STORAGES + or outer_desc.dtype != inner_desc.dtype): + return False + + outer_memlet = edge.data + # May be dst-relative (subset in ``other_subset``); resolve via ``get_src/dst_subset``. + if stage_in: + outer_subset = outer_memlet.get_src_subset(edge, state) or outer_memlet.subset + else: + outer_subset = outer_memlet.get_dst_subset(edge, state) or outer_memlet.subset + if _competing_writer(state, edge.dst, edge, outer.data, outer_subset): + return False + outer_side_memlet = Memlet(data=outer.data, subset=copy.deepcopy(outer_subset)) + outer_side_memlet.dynamic = outer_memlet.dynamic + outer_side_memlet.wcr = outer_memlet.wcr + # When the memlet already names both sides that mapping IS the copy -- deriving from the + # outer subset instead silently retargets the write. + if stage_in: + inner_subset = outer_memlet.get_dst_subset(edge, state) + else: + inner_subset = outer_memlet.get_src_subset(edge, state) + if inner_subset is None or outer_memlet.other_subset is None: + inner_subset = _derive_matching_dst_subset(outer_subset, inner_desc) + else: + inner_subset = copy.deepcopy(inner_subset) + if stage_in and _competing_writer(state, inner_node, edge, inner_node.data, inner_subset): + return False + inner_memlet = Memlet(data=inner_node.data, subset=inner_subset) + label = (f"copy_{outer.data}_to_{inner_node.data}" if stage_in else f"copy_{inner_node.data}_to_{outer.data}") + libnode = CopyLibraryNode(name=label) + libnode.instrument = state.instrument + state.add_node(libnode) + if stage_in: + map_node = edge.src + state.add_edge(map_node, edge.src_conn, libnode, CopyLibraryNode.INPUT_CONNECTOR_NAME, outer_side_memlet) + state.add_edge(libnode, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, inner_node, None, inner_memlet) + _carry_write_ordering(state, inner_node, libnode) + boundary_conn = 'IN_' + edge.src_conn[len('OUT_'):] + boundary_edges = list(state.in_edges_by_connector(map_node, boundary_conn)) + else: + map_node = edge.dst + state.add_edge(inner_node, None, libnode, CopyLibraryNode.INPUT_CONNECTOR_NAME, inner_memlet) + state.add_edge(libnode, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, map_node, edge.dst_conn, outer_side_memlet) + boundary_conn = 'OUT_' + edge.dst_conn[len('IN_'):] + boundary_edges = list(state.out_edges_by_connector(map_node, boundary_conn)) + state.remove_edge(edge) + + # The scope-boundary edge on this connector may still carry a memlet whose ``.data`` + # names the inner array, relying on memlet_path continuing through the scope entry/exit + # straight to inner_node for validation (validation.py resolves src/dst from the full + # path, not the edge's immediate neighbours). That continuation broke: the libnode now + # sits between the scope node and inner_node, so the path ends at a non-AccessNode and + # the boundary edge needs its own outer-relative memlet instead. + for bedge in boundary_edges: + if bedge.data.data != outer.data: + bedge.data = Memlet(data=outer.data, subset=copy.deepcopy(outer_subset)) + return True diff --git a/dace/transformation/transformation.py b/dace/transformation/transformation.py index 73222ad1f6..a1c6bece0a 100644 --- a/dace/transformation/transformation.py +++ b/dace/transformation/transformation.py @@ -727,6 +727,9 @@ def apply(self, state, sdfg, *args, **kwargs): else: raise TypeError("Node expansion must be a CodeNode or an SDFG") + # The node the expansion replaces is the one the user asked to measure. + expansion.instrument = node.instrument + expansion.environments = copy.copy(set(map(lambda a: a.full_class_path(), type(self).environments))) sdutil.change_edge_dest(state, node, expansion) sdutil.change_edge_src(state, node, expansion) diff --git a/tests/codegen/multicopy_test.py b/tests/codegen/multicopy_test.py index e6646c3374..82b48ae17b 100644 --- a/tests/codegen/multicopy_test.py +++ b/tests/codegen/multicopy_test.py @@ -16,8 +16,11 @@ def test_multicopy(): state.add_nedge(a, b, dace.Memlet('A[0]')) state.add_nedge(a, c, dace.Memlet('C[0]')) - # Check generated code - assert sdfg.generate_code()[0].clean_code.count('CopyND') == 2 + # Check generated code. The regression under test is DUPLICATED copy code, so what matters is + # the count, not which lowering produced it: with explicit copy nodes a single-element copy + # becomes a scalar-assignment tasklet instead of the dace::CopyND runtime template. + code = sdfg.generate_code()[0].clean_code + assert code.count('CopyND') + code.count('_cpy_out = _cpy_in;') == 2 # Check outputs A = np.random.rand(1) diff --git a/tests/library/copy_node_test.py b/tests/library/copy_node_test.py new file mode 100644 index 0000000000..4dac05d7e0 --- /dev/null +++ b/tests/library/copy_node_test.py @@ -0,0 +1,2004 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Tests for ``CopyLibraryNode`` and its pure, CPU, CUDA, cross-storage, register, and shared-memory expansions.""" +import contextlib +from dataclasses import dataclass +from typing import Optional, Sequence, Tuple + +import dace +from dace import symbolic +from dace.sdfg.graph import SubgraphView +from dace.transformation.subgraph import GPUPersistentKernel +from dace.libraries.standard.helper import collapse_shape_and_strides +from dace.libraries.standard.nodes.copy import CopyLibraryNode, select_copy_implementation +from dace.libraries.standard.nodes.fill import FillLibraryNode +from dace.libraries.standard.nodes.copy.common import _make_expansion_sdfg, cuda2d_pitch_params + +import pytest +import numpy as np + + +@dataclass +class _ArraySpec: + """Per-side array spec for :func:`_make_copy_sdfg`. + + :param shape: array shape. + :param storage: storage type. + :param strides: explicit strides; ``None`` keeps DaCe's packed-C default. + :param total_size: explicit buffer total size; only consulted when ``strides`` is set + (defaults to ``prod(shape)``). + :param transient: transient-array flag. + :param subset: memlet subset string; defaults to the full per-dim range. + :param name: SDFG-visible array name; defaults to ``src`` / ``dst`` from position. + :param dtype: element type; ``None`` defers to the helper's ``dtype`` argument. + """ + shape: Sequence[int] + storage: dace.dtypes.StorageType + strides: Optional[Sequence[int]] = None + total_size: Optional[int] = None + transient: bool = False + subset: Optional[str] = None + name: Optional[str] = None + dtype: Optional[dace.dtypes.typeclass] = None + + +def _make_copy_sdfg(src: _ArraySpec, + dst: _ArraySpec, + *, + implementation: Optional[str] = None, + name: str = "copy_sdfg", + libnode_name: str = "cp", + dtype: dace.dtypes.typeclass = dace.float64) -> Tuple[dace.SDFG, CopyLibraryNode]: + """One-state SDFG copying ``src`` -> ``dst`` via a single ``CopyLibraryNode``. + + :param src: source-side array spec. + :param dst: destination-side array spec. + :param implementation: pinned ``CopyLibraryNode.implementation`` (``None`` keeps ``'Auto'``). + :param name: SDFG name. + :param libnode_name: libnode label. + :param dtype: fallback dtype when a spec leaves ``dtype=None``. + :returns: ``(sdfg, libnode)``. + """ + sdfg, src_name, dst_name, src_acc, dst_acc, src_subset, dst_subset = _make_copy_skeleton(src, dst, name, dtype) + libnode = CopyLibraryNode(name=libnode_name) + if implementation is not None: + libnode.implementation = implementation + state = sdfg.start_state + state.add_edge(src_acc, None, libnode, CopyLibraryNode.INPUT_CONNECTOR_NAME, + dace.memlet.Memlet(f"{src_name}[{src_subset}]")) + state.add_edge(libnode, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, dst_acc, None, + dace.memlet.Memlet(f"{dst_name}[{dst_subset}]")) + return sdfg, libnode + + +def _make_copy_skeleton(src: _ArraySpec, dst: _ArraySpec, name: str, dtype: dace.dtypes.typeclass): + """Shared scaffolding for :func:`_make_copy_sdfg` and :func:`_make_legacy_copy_sdfg`: builds the arrays + AccessNodes and returns the subsets.""" + sdfg = dace.SDFG(name) + src_name = src.name or "src" + dst_name = dst.name or "dst" + for arr_name, spec in ((src_name, src), (dst_name, dst)): + kwargs = {"transient": spec.transient} + if spec.strides is not None: + # Sympify each stride so string entries (``"src_stride"``) become + # SDFG symbols. ``Array.validate`` rejects raw strings, and + # ``add_array`` only sympifies the shape, not the strides. + kwargs["strides"] = [dace.symbolic.pystr_to_symbolic(s) for s in spec.strides] + kwargs["total_size"] = spec.total_size if spec.total_size is not None else int(np.prod(spec.shape)) + sdfg.add_array(arr_name, spec.shape, spec.dtype or dtype, storage=spec.storage, **kwargs) + state = sdfg.add_state("main") + src_acc = state.add_access(src_name) + dst_acc = state.add_access(dst_name) + src_subset = src.subset if src.subset is not None else ", ".join(f"0:{s}" for s in src.shape) + dst_subset = dst.subset if dst.subset is not None else ", ".join(f"0:{s}" for s in dst.shape) + return sdfg, src_name, dst_name, src_acc, dst_acc, src_subset, dst_subset + + +def _make_legacy_copy_sdfg(src: _ArraySpec, + dst: _ArraySpec, + *, + name: str = "copy_legacy", + dtype: dace.dtypes.typeclass = dace.float64) -> dace.SDFG: + """One-state SDFG copying ``src`` -> ``dst`` via a canonical direct AN -> AN edge. + + Legacy DaCe memlet convention (``data=dst``, ``subset``=dst write region, + ``other_subset``=src read region) -- the standard copy lowering's output and + the baseline for comparing against the :class:`CopyLibraryNode` path. + """ + sdfg, src_name, dst_name, src_acc, dst_acc, src_subset, dst_subset = _make_copy_skeleton(src, dst, name, dtype) + sdfg.start_state.add_edge(src_acc, None, dst_acc, None, + dace.memlet.Memlet(data=dst_name, subset=dst_subset, other_subset=src_subset)) + return sdfg + + +def _fortran_strides(shape): + """Column-major Fortran-packed strides, via the same helper ``Array.is_packed_fortran_strides`` checks against.""" + return dace.data.Array(dace.float64, shape=shape)._get_packed_fortran_strides() + + +def _compile_no_copynd(sdfg: dace.SDFG): + """Assert the generated C++ contains no ``dace::CopyND`` template, then compile. + + The libnodes displace the runtime CopyND fallback entirely. The only intentional + ``CopyND`` user is ``ExpandSharedMemoryCollective``; tests exercising that expansion + inspect tasklet bodies directly and don't run codegen, so this assertion is safe here. + """ + for obj in sdfg.generate_code(): + assert 'CopyND<' not in obj.code, f"unexpected dace::CopyND in generated code object {obj.name}" + return sdfg.compile() + + +def test_copy_pure_cpu(): + """Pure (mapped tasklet) expansion on CPU_Heap -> CPU_Heap.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[200], storage=dace.dtypes.StorageType.CPU_Heap, subset="150:200", name="A"), + _ArraySpec(shape=[200], storage=dace.dtypes.StorageType.CPU_Heap, subset="50:100", name="B"), + implementation="MappedTasklet", + name="copy_pure_cpu", + ) + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = _compile_no_copynd(sdfg) + + A = np.ones(200, dtype=np.float64) + B = np.zeros(200, dtype=np.float64) + exe(A=A, B=B) + + np.testing.assert_array_equal(B[50:100], A[150:200]) + assert np.all(B[:50] == 0) + assert np.all(B[100:] == 0) + + +def test_copy_cpu_memcpy(): + """CPU expansion (std::memcpy) on CPU_Heap -> CPU_Heap.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[200], storage=dace.dtypes.StorageType.CPU_Heap, subset="150:200", name="A"), + _ArraySpec(shape=[200], storage=dace.dtypes.StorageType.CPU_Heap, subset="50:100", name="B"), + implementation="MemcpyCPU", + name="copy_cpu_memcpy", + ) + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = _compile_no_copynd(sdfg) + + A = np.arange(200, dtype=np.float64) + B = np.zeros(200, dtype=np.float64) + exe(A=A, B=B) + + np.testing.assert_array_equal(B[50:100], A[150:200]) + + +def test_copy_fortran_packed_same_rank(): + """Same-rank Fortran-packed (column-major) full copy is contiguous and same-layout, so the + Auto path routes it to the serial ``std::memcpy`` (``MemcpyCPU``); a flat byte copy is exact + for two Fortran-packed operands.""" + sdfg, libnode = _make_copy_sdfg( + _ArraySpec(shape=(4, 5, 6), storage=dace.dtypes.StorageType.CPU_Heap, strides=(1, 4, 20)), + _ArraySpec(shape=(4, 5, 6), storage=dace.dtypes.StorageType.CPU_Heap, strides=(1, 4, 20)), + name="copy_fortran_packed_same_rank", + ) + sdfg.validate() + sdfg.expand_library_nodes() + assert libnode.implementation == 'MemcpyCPU' + + src_data = np.arange(120, dtype=np.float64).reshape(4, 5, 6, order='F').copy(order='F') + dst_data = np.zeros((4, 5, 6), dtype=np.float64, order='F') + sdfg(src=src_data, dst=dst_data) + assert np.array_equal(dst_data, src_data) + + +def test_copy_fortran_packed_strided_slice(): + """Same-rank Fortran-packed strided-slice copy via the Auto-routed MappedTasklet.""" + sdfg, libnode = _make_copy_sdfg( + _ArraySpec(shape=(8, 10, 12), + storage=dace.dtypes.StorageType.CPU_Heap, + strides=(1, 8, 80), + subset="2:6, 3:7, 4:8"), + _ArraySpec(shape=(8, 10, 12), + storage=dace.dtypes.StorageType.CPU_Heap, + strides=(1, 8, 80), + subset="2:6, 3:7, 4:8"), + name="copy_fortran_packed_strided_slice", + ) + sdfg.validate() + sdfg.expand_library_nodes() + assert libnode.implementation == 'MappedTasklet' + + src_data = np.arange(960, dtype=np.float64).reshape(8, 10, 12, order='F').copy(order='F') + dst_data = np.zeros((8, 10, 12), dtype=np.float64, order='F') + sdfg(src=src_data, dst=dst_data) + assert np.array_equal(dst_data[2:6, 3:7, 4:8], src_data[2:6, 3:7, 4:8]) + untouched = dst_data.copy() + untouched[2:6, 3:7, 4:8] = 0 + assert np.all(untouched == 0) + + +def test_copy_mixed_c_fortran_via_mapped_tasklet(): + """Mixed C-packed -> Fortran-packed same-rank copy lowers via MappedTasklet.""" + sdfg, libnode = _make_copy_sdfg( + _ArraySpec(shape=(6, 7), storage=dace.dtypes.StorageType.CPU_Heap, strides=(7, 1)), + _ArraySpec(shape=(6, 7), storage=dace.dtypes.StorageType.CPU_Heap, strides=(1, 6)), + name="copy_mixed_c_fortran", + ) + sdfg.validate() + sdfg.expand_library_nodes() + assert libnode.implementation == 'MappedTasklet' + + src_data = np.arange(42, dtype=np.float64).reshape(6, 7).copy(order='C') + dst_data = np.zeros((6, 7), dtype=np.float64, order='F') + sdfg(src=src_data, dst=dst_data) + assert np.array_equal(dst_data, src_data) + + +def test_copy_rank_mismatch_mixed_layouts_raises(): + """Rank-mismatch with mixed C/F packed layouts is rejected (1-D walker has no shared layout).""" + # src is C-packed (3, 8) -- strides (8, 1); dst is Fortran-packed (2, 3, 4) + # -- strides (1, 2, 6). Same volume = 24. + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=(3, 8), storage=dace.dtypes.StorageType.CPU_Heap), + _ArraySpec(shape=(2, 3, 4), storage=dace.dtypes.StorageType.CPU_Heap, strides=(1, 2, 6)), + name="copy_rank_mismatch_mixed_raises", + ) + sdfg.validate() + with pytest.raises(ValueError, match="same major order"): + sdfg.expand_library_nodes() + + +def test_copy_rank_mismatch_padded_src_raises(): + """Rank-mismatch with padded (neither C- nor F-packed) strides is rejected.""" + # src padded (row stride 8 instead of 6), dst flat (120,). + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=(4, 5, 6), + storage=dace.dtypes.StorageType.CPU_Heap, + strides=(5 * 8, 8, 1), + total_size=4 * 5 * 8), + _ArraySpec(shape=(120, ), storage=dace.dtypes.StorageType.CPU_Heap), + name="copy_rank_mismatch_padded_raises", + ) + sdfg.validate() + with pytest.raises(ValueError, match="same major order"): + sdfg.expand_library_nodes() + + +def test_copy_rank_mismatch_strided_src_subset(): + """Rank-mismatch from a non-contiguous C-layout src subset walks the collapsed strides.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=(8, 10), storage=dace.dtypes.StorageType.CPU_Heap, subset="0:8, 2:6"), + _ArraySpec(shape=(32, ), storage=dace.dtypes.StorageType.CPU_Heap), + name="copy_rank_mismatch_strided_subset", + ) + sdfg.validate() + sdfg.expand_library_nodes() + exe = _compile_no_copynd(sdfg) + A = np.arange(80, dtype=np.float64).reshape(8, 10).copy() + B = np.zeros(32, dtype=np.float64) + exe(src=A, dst=B) + np.testing.assert_array_equal(B, A[:, 2:6].reshape(32)) + + +def test_copy_rank_mismatch_strided_dst_subset(): + """Symmetric to the src-side variant: non-contiguous C-layout subset on the dst side.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=(32, ), storage=dace.dtypes.StorageType.CPU_Heap), + _ArraySpec(shape=(8, 10), storage=dace.dtypes.StorageType.CPU_Heap, subset="0:8, 2:6"), + name="copy_rank_mismatch_strided_dst_subset", + ) + sdfg.validate() + sdfg.expand_library_nodes() + exe = _compile_no_copynd(sdfg) + A = np.arange(32, dtype=np.float64) + B = np.zeros((8, 10), dtype=np.float64) + exe(src=A, dst=B) + np.testing.assert_array_equal(B[:, 2:6], A.reshape(8, 4)) + + +def test_copy_same_subset_different_array_shapes(): + """A ``0:N`` slice copies between arrays of different total shape as long as the per-dim subset sizes match.""" + N = 10 + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=(200, ), storage=dace.dtypes.StorageType.CPU_Heap, subset=f"0:{N}", name="A"), + _ArraySpec(shape=(500, ), storage=dace.dtypes.StorageType.CPU_Heap, subset=f"0:{N}", name="B"), + name="copy_same_subset_diff_shape", + ) + sdfg.expand_library_nodes() + sdfg.validate() + exe = _compile_no_copynd(sdfg) + A = np.arange(200, dtype=np.float64) + B = np.zeros(500, dtype=np.float64) + exe(A=A, B=B) + np.testing.assert_array_equal(B[:N], A[:N]) + + +def test_copy_1d_slice_from_2d_source(): + """A row-slice ``[i, 0:N]`` of a 2D array copies into a 1D array (singleton dims collapse to same rank).""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=(5, 10), storage=dace.dtypes.StorageType.CPU_Heap, subset="2, 0:10", name="A"), + _ArraySpec(shape=(10, ), storage=dace.dtypes.StorageType.CPU_Heap, subset="0:10", name="B"), + name="copy_1d_slice_from_2d", + ) + sdfg.expand_library_nodes() + sdfg.validate() + exe = _compile_no_copynd(sdfg) + A = np.arange(50, dtype=np.float64).reshape(5, 10).copy() + B = np.zeros(10, dtype=np.float64) + exe(A=A, B=B) + np.testing.assert_array_equal(B, A[2]) + + +def test_copy_transpose_pattern_rejected(): + """Same-rank copy with per-dim shapes swapped (transpose) is rejected upfront.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=(3, 4), storage=dace.dtypes.StorageType.CPU_Heap), + _ArraySpec(shape=(4, 3), storage=dace.dtypes.StorageType.CPU_Heap), + name="copy_transpose_pattern", + ) + sdfg.validate() + with pytest.raises(ValueError, match="matching per-dim shapes"): + sdfg.expand_library_nodes() + + +def test_copy_4d_to_1d_flatten_c_packed(): + """4D -> 1D flatten via MappedTasklet rank-mismatch (extends beyond the 3D->1D coverage).""" + sdfg, libnode = _make_copy_sdfg( + _ArraySpec(shape=(2, 3, 4, 5), storage=dace.dtypes.StorageType.CPU_Heap), + _ArraySpec(shape=(120, ), storage=dace.dtypes.StorageType.CPU_Heap), + name="copy_4d_to_1d_c", + ) + sdfg.validate() + sdfg.expand_library_nodes() + assert libnode.implementation == 'MappedTasklet' + + src = np.arange(120, dtype=np.float64).reshape(2, 3, 4, 5).copy(order='C') + dst = np.zeros(120, dtype=np.float64) + sdfg(src=src, dst=dst) + assert np.array_equal(dst, src.ravel(order='C')) + + +def test_copy_1d_to_4d_inflate_c_packed(): + """1D -> 4D inflate (higher-rank destination); inverse direction of the flatten path.""" + sdfg, libnode = _make_copy_sdfg( + _ArraySpec(shape=(24, ), storage=dace.dtypes.StorageType.CPU_Heap), + _ArraySpec(shape=(2, 3, 4), storage=dace.dtypes.StorageType.CPU_Heap), + name="copy_1d_to_3d_c", + ) + sdfg.validate() + sdfg.expand_library_nodes() + assert libnode.implementation == 'MappedTasklet' + + src = np.arange(24, dtype=np.float64) + dst = np.zeros((2, 3, 4), dtype=np.float64) + sdfg(src=src, dst=dst) + assert np.array_equal(dst, src.reshape(2, 3, 4)) + + +def test_copy_3d_to_2d_collapse_first_two_dims(): + """3D -> 2D collapse of the first two dims (C-order) via MappedTasklet rank-mismatch.""" + sdfg, libnode = _make_copy_sdfg( + _ArraySpec(shape=(2, 3, 4), storage=dace.dtypes.StorageType.CPU_Heap), + _ArraySpec(shape=(6, 4), storage=dace.dtypes.StorageType.CPU_Heap), + name="copy_3d_to_2d_collapse", + ) + sdfg.validate() + sdfg.expand_library_nodes() + assert libnode.implementation == 'MappedTasklet' + + src = np.arange(24, dtype=np.float64).reshape(2, 3, 4).copy(order='C') + dst = np.zeros((6, 4), dtype=np.float64) + sdfg(src=src, dst=dst) + assert np.array_equal(dst, src.reshape(6, 4)) + + +def test_copy_4d_to_2d_collapse_pair_dims_fortran(): + """4D -> 2D Fortran-packed reshape: walk both sides in column-major order.""" + sdfg, libnode = _make_copy_sdfg( + _ArraySpec(shape=(2, 3, 4, 5), storage=dace.dtypes.StorageType.CPU_Heap, strides=(1, 2, 6, 24)), + _ArraySpec(shape=(6, 20), storage=dace.dtypes.StorageType.CPU_Heap, strides=(1, 6)), + name="copy_4d_to_2d_f", + ) + sdfg.validate() + sdfg.expand_library_nodes() + assert libnode.implementation == 'MappedTasklet' + + src = np.arange(120, dtype=np.float64).reshape(2, 3, 4, 5, order='F').copy(order='F') + dst = np.zeros((6, 20), dtype=np.float64, order='F') + sdfg(src=src, dst=dst) + assert np.array_equal(dst, src.reshape(6, 20, order='F')) + + +def test_copy_strided_step_2_cpu_same_rank(): + """Same-rank 1D copy with subset step=2 (every other element).""" + sdfg, libnode = _make_copy_sdfg( + _ArraySpec(shape=(10, ), storage=dace.dtypes.StorageType.CPU_Heap, subset="0:10:2"), + _ArraySpec(shape=(5, ), storage=dace.dtypes.StorageType.CPU_Heap, subset="0:5"), + name="copy_step2_cpu", + ) + sdfg.validate() + sdfg.expand_library_nodes() + assert libnode.implementation == 'MappedTasklet' + + src = np.arange(10, dtype=np.float64) + dst = np.zeros(5, dtype=np.float64) + sdfg(src=src, dst=dst) + assert np.array_equal(dst, src[0:10:2]) + + +@pytest.mark.gpu +def test_copy_pure_gpu(): + """Pure (mapped tasklet) expansion on GPU_Global -> GPU_Global.""" + import cupy as cp + + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[200], storage=dace.dtypes.StorageType.GPU_Global, subset="150:200", name="gpu_A"), + _ArraySpec(shape=[200], storage=dace.dtypes.StorageType.GPU_Global, subset="50:100", name="gpu_B"), + implementation="MappedTasklet", + name="copy_pure_gpu", + ) + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = _compile_no_copynd(sdfg) + + A = cp.ones(200, dtype=cp.float64) + B = cp.zeros(200, dtype=cp.float64) + exe(gpu_A=A, gpu_B=B) + + cp.testing.assert_array_equal(B[50:100], A[150:200]) + assert cp.all(B[:50] == 0) + assert cp.all(B[100:] == 0) + + +@pytest.mark.gpu +def test_copy_cuda_d2d(): + """CUDA expansion (cudaMemcpyDeviceToDevice) on GPU_Global -> GPU_Global.""" + import cupy as cp + + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[200], storage=dace.dtypes.StorageType.GPU_Global, subset="150:200", name="gpu_A"), + _ArraySpec(shape=[200], storage=dace.dtypes.StorageType.GPU_Global, subset="50:100", name="gpu_B"), + implementation="MemcpyCUDA1D", + name="copy_cuda_d2d", + ) + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = _compile_no_copynd(sdfg) + + A = cp.arange(200, dtype=cp.float64) + B = cp.zeros(200, dtype=cp.float64) + exe(gpu_A=A, gpu_B=B) + + cp.testing.assert_array_equal(B[50:100], A[150:200]) + + +@pytest.mark.gpu +def test_copy_cuda_1d_single_element(): + """CUDA expansion (cudaMemcpyDeviceToDevice) on GPU_Global -> GPU_Global for a single element.""" + import cupy as cp + + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[200], + strides=["src_stride"], + storage=dace.dtypes.StorageType.GPU_Global, + subset="130", + name="gpu_A"), + _ArraySpec(shape=[200], + strides=["dst_stride"], + storage=dace.dtypes.StorageType.GPU_Global, + subset="15", + name="gpu_B"), + implementation="MemcpyCUDA1D", + name="copy_cuda_1d_single_element", + ) + # ``add_array`` registered the stride symbols at default int. + sdfg.validate() + + sdfg.expand_library_nodes() + sdfg.validate() + + exe = _compile_no_copynd(sdfg) + + A = cp.arange(200, dtype=cp.float64) + B = cp.zeros(200, dtype=cp.float64) + + ref = B.copy() + ref[15] = A[130] + + exe(gpu_A=A, gpu_B=B, src_stride=1, dst_stride=1) + + cp.testing.assert_array_equal(ref, B) + + +def test_copy_pure_host_to_device_rejected(): + """Pure expansion must reject CPU_Heap -> GPU_Global (needs cudaMemcpy).""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[128], storage=dace.dtypes.StorageType.CPU_Heap), + _ArraySpec(shape=[128], storage=dace.dtypes.StorageType.GPU_Global), + implementation="MappedTasklet", + name="copy_pure_h2d_reject", + ) + sdfg.validate() + with pytest.raises(Exception, match="CPU/GPU boundary"): + sdfg.expand_library_nodes() + + +def test_copy_pure_device_to_host_rejected(): + """Pure expansion must reject GPU_Global -> CPU_Heap (needs cudaMemcpy).""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[128], storage=dace.dtypes.StorageType.GPU_Global), + _ArraySpec(shape=[128], storage=dace.dtypes.StorageType.CPU_Heap), + implementation="MappedTasklet", + name="copy_pure_d2h_reject", + ) + sdfg.validate() + with pytest.raises(Exception, match="CPU/GPU boundary"): + sdfg.expand_library_nodes() + + +@pytest.mark.gpu +def test_copy_cuda_host_to_device(): + """CUDAHostToDevice expansion for CPU_Heap -> GPU_Global.""" + import cupy as cp + + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[128], storage=dace.dtypes.StorageType.CPU_Heap), + _ArraySpec(shape=[128], storage=dace.dtypes.StorageType.GPU_Global), + implementation="MemcpyCUDA1D", + name="copy_cuda_h2d", + ) + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = _compile_no_copynd(sdfg) + + src = np.arange(128, dtype=np.float64) + dst = cp.zeros(128, dtype=cp.float64) + exe(src=src, dst=dst) + + cp.testing.assert_array_equal(dst, cp.asarray(src)) + + +@pytest.mark.gpu +def test_copy_cuda_device_to_host(): + """CUDADeviceToHost expansion for GPU_Global -> CPU_Heap.""" + import cupy as cp + + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[128], storage=dace.dtypes.StorageType.GPU_Global), + _ArraySpec(shape=[128], storage=dace.dtypes.StorageType.CPU_Heap), + implementation="MemcpyCUDA1D", + name="copy_cuda_d2h", + ) + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = _compile_no_copynd(sdfg) + + src = cp.arange(128, dtype=cp.float64) + dst = np.zeros(128, dtype=np.float64) + exe(src=src, dst=dst) + + np.testing.assert_array_equal(dst, cp.asnumpy(src)) + + +@pytest.mark.gpu +def test_copy_cuda_4d_strided_host_to_device(): + """A 4D strided CPU_Heap -> GPU_Global slice copy via ``MemcpyCUDANDStrided`` produces correct output.""" + import cupy as cp + + # Slice into a larger array so the outer dims are strided, exercising the + # per-row strided CUDA path rather than a single contiguous memcpy. + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=(7, 8, 9, 10), + storage=dace.dtypes.StorageType.CPU_Heap, + subset="1:6, 1:7, 1:8, 1:9", + name="A_full"), + _ArraySpec(shape=(5, 6, 7, 8), storage=dace.dtypes.StorageType.GPU_Global, name="B_dst"), + implementation="MemcpyCUDANDStrided", + name="copy_cuda_4d_strided_h2d", + libnode_name="cp_4d_strided", + ) + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = _compile_no_copynd(sdfg) + + # ``reshape`` returns a numpy view; DaCe rejects views by default + # (``compiler.allow_view_arguments``). Build directly as a fresh array. + A = np.empty((7, 8, 9, 10), dtype=np.float64) + A[:] = np.arange(7 * 8 * 9 * 10).reshape(7, 8, 9, 10) + B = cp.zeros((5, 6, 7, 8), dtype=cp.float64) + exe(A_full=A, B_dst=B) + + expected = A[1:6, 1:7, 1:8, 1:9] + cp.testing.assert_array_equal(B, cp.asarray(expected)) + + +def test_copy_fortran_packed_cpu_default_pure(): + """A same-side CPU copy of a Fortran-packed array expands and produces correct output.""" + shape = (4, 5, 6) + f_strides = _fortran_strides(shape) + total = int(np.prod(shape)) + + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=shape, storage=dace.dtypes.StorageType.CPU_Heap, strides=f_strides, total_size=total), + _ArraySpec(shape=shape, storage=dace.dtypes.StorageType.CPU_Heap, strides=f_strides, total_size=total), + name="copy_fortran_cpu", + libnode_name="cp_fortran_cpu", + ) + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = _compile_no_copynd(sdfg) + + A = np.arange(total, dtype=np.float64).reshape(shape, order='F').copy(order='F') + B = np.zeros(shape, dtype=np.float64, order='F') + exe(src=A, dst=B) + np.testing.assert_array_equal(B, A) + + +@pytest.mark.gpu +def test_copy_fortran_packed_gpu_falls_back_to_pure(): + """A same-side GPU copy of a Fortran-packed array expands and produces correct output.""" + import cupy as cp + + shape = (4, 5, 6) + f_strides = _fortran_strides(shape) + total = int(np.prod(shape)) + + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=shape, storage=dace.dtypes.StorageType.GPU_Global, strides=f_strides, total_size=total), + _ArraySpec(shape=shape, storage=dace.dtypes.StorageType.GPU_Global, strides=f_strides, total_size=total), + implementation="MemcpyCUDA1D", + name="copy_fortran_gpu", + libnode_name="cp_fortran_gpu", + ) + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = _compile_no_copynd(sdfg) + + host = np.arange(total, dtype=np.float64).reshape(shape, order='F').copy(order='F') + A = cp.asfortranarray(cp.asarray(host)) + B = cp.asfortranarray(cp.zeros(shape, dtype=cp.float64)) + exe(src=A, dst=B) + cp.testing.assert_array_equal(B, A) + + +@pytest.mark.gpu +def test_copy_fortran_packed_cpu_to_gpu_uses_outermost_chunk(): + """A cross-CPU/GPU copy of a Fortran-packed array expands and produces correct output.""" + import cupy as cp + + shape = (4, 5, 6) + f_strides = _fortran_strides(shape) + total = int(np.prod(shape)) + + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=shape, storage=dace.dtypes.StorageType.CPU_Heap, strides=f_strides, total_size=total), + _ArraySpec(shape=shape, storage=dace.dtypes.StorageType.GPU_Global, strides=f_strides, total_size=total), + implementation="MemcpyCUDA1D", + name="copy_fortran_h2d", + libnode_name="cp_fortran_h2d", + ) + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = _compile_no_copynd(sdfg) + + host = np.arange(total, dtype=np.float64).reshape(shape, order='F').copy(order='F') + dev = cp.asfortranarray(cp.zeros(shape, dtype=cp.float64)) + exe(src=host, dst=dev) + cp.testing.assert_array_equal(dev, cp.asarray(host)) + + +def test_copy_no_common_stride1_axis_raises(): + """Cross-CPU/GPU copy with no shared stride-1 axis is rejected.""" + # src C-packed (stride-1 innermost), dst Fortran-packed (stride-1 + # outermost): after the partial slice the two have no shared stride-1 axis. + shape = (4, 5, 6) + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=shape, + storage=dace.dtypes.StorageType.CPU_Heap, + strides=(30, 6, 1), + total_size=120, + subset="0:4, 0:4, 0:5"), + _ArraySpec(shape=shape, + storage=dace.dtypes.StorageType.GPU_Global, + strides=(1, 4, 20), + total_size=120, + subset="0:4, 0:4, 0:5"), + implementation="Auto", # exercise the refine-time strided-pattern check + name="copy_no_common_stride1", + libnode_name="cp_no_common", + ) + sdfg.validate() + with pytest.raises(ValueError, match="cross-CPU/GPU"): + sdfg.expand_library_nodes() + + +def test_copy_node_storage_from_edges(): + """``src_storage`` / ``dst_storage`` resolve live from the node's ``_in`` / ``_out`` edges.""" + sdfg, node = _make_copy_sdfg( + _ArraySpec(shape=[10], storage=dace.dtypes.StorageType.CPU_Heap, name="A"), + _ArraySpec(shape=[10], storage=dace.dtypes.StorageType.GPU_Global, name="B"), + name="storage_from_edges", + libnode_name="edges_to_storage", + ) + state = sdfg.start_state + assert node.src_storage(state) == dace.dtypes.StorageType.CPU_Heap + assert node.dst_storage(state) == dace.dtypes.StorageType.GPU_Global + + +def test_copy_node_storage_defaults_when_unattached(): + """Without edges, the storage methods fall back to ``StorageType.Default``.""" + sdfg = dace.SDFG("storage_unattached") + state = sdfg.add_state("main") + node = CopyLibraryNode(name="unattached") + state.add_node(node) + + assert node.src_storage(state) == dace.dtypes.StorageType.Default + assert node.dst_storage(state) == dace.dtypes.StorageType.Default + + +def test_copy_cross_storage_validation_rejects_without_flag(): + """The ``MemcpyCPU`` expansion rejects a CPU<->GPU storage mismatch at expansion time.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[64], storage=dace.dtypes.StorageType.CPU_Heap), + _ArraySpec(shape=[64], storage=dace.dtypes.StorageType.GPU_Global), + implementation="MemcpyCPU", + name="copy_cross_reject", + ) + sdfg.validate() # the SDFG is valid; only the expansion rejects the mismatch + with pytest.raises(Exception): + sdfg.expand_library_nodes() + + +def test_copy_dtype_mismatch_rejected(): + """CopyLibraryNode must reject mismatched dtypes.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[10], storage=dace.dtypes.StorageType.CPU_Heap, dtype=dace.float32, name="A"), + _ArraySpec(shape=[10], storage=dace.dtypes.StorageType.CPU_Heap, dtype=dace.float64, name="B"), + name="dtype_mismatch", + libnode_name="cp_bad", + ) + with pytest.raises(ValueError, match="data types must match"): + sdfg.expand_library_nodes() + + +def test_cpu_memcpy_rejects_non_contiguous_subset(): + """CPU (memcpy) expansion must reject a non-contiguous 2D slice.""" + # Partial dim 0 over a smaller dim 1 makes the source slice non-contiguous. + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[10, 20], storage=dace.dtypes.StorageType.CPU_Heap, subset="2:6, 0:10", name="A"), + _ArraySpec(shape=[4, 20], storage=dace.dtypes.StorageType.CPU_Heap, subset="0:4, 0:10", name="B"), + implementation="MemcpyCPU", + name="cpu_noncontig", + libnode_name="cp_nc", + ) + with pytest.raises(Exception, match="contiguous"): + sdfg.expand_library_nodes() + + +def test_strided_expansions_accept_non_contiguous(): + """The ``MappedTasklet`` expansion accepts a non-contiguous subset.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[10, 20], storage=dace.dtypes.StorageType.CPU_Heap, subset="2:6, 0:10", name="A"), + _ArraySpec(shape=[4, 20], storage=dace.dtypes.StorageType.CPU_Heap, subset="0:4, 0:10", name="B"), + implementation="MappedTasklet", + name="noncontig_MappedTasklet", + ) + sdfg.expand_library_nodes() + + +# A (1, N) array whose unit leading dim carries a padded stride (here 64) is a +# non-packed descriptor, yet the accessed row ``[0, 0:N]`` is one physical run of +# N contiguous elements: the pad sits on an extent-1 axis that is never stepped. +# ``is_contiguous_subset`` therefore reports True (the 1D-slice special case), and +# the copy safely lowers to a single flat block. A fresh contiguous (1, N) array +# backs it with no view (``total_size`` only needs to cover the accessed run). +_PADDED_N = 60 +_PADDED_STRIDE = 64 +_PADDED_ROWS = 3 + + +def _padded_unit_spec(storage, name): + """``_ArraySpec`` for a (1, ``_PADDED_N``) array with a padded (non-packed) leading stride.""" + return _ArraySpec(shape=(1, _PADDED_N), + storage=storage, + strides=(_PADDED_STRIDE, 1), + total_size=_PADDED_N, + name=name) + + +def _padded_multirow_spec(storage, name): + """``_ArraySpec`` for a (``_PADDED_ROWS``, ``_PADDED_N``) array whose rows carry the padded stride. + + Unlike the unit-row spec, the leading dim has extent > 1, so the inter-row pitch gap + (``_PADDED_STRIDE - _PADDED_N`` unused elements per row) is actually stepped over: the full + ``[0:ROWS, 0:N]`` copy is genuinely non-contiguous. + """ + return _ArraySpec(shape=(_PADDED_ROWS, _PADDED_N), + storage=storage, + strides=(_PADDED_STRIDE, 1), + total_size=_PADDED_STRIDE * _PADDED_ROWS, + name=name) + + +def test_copy_padded_unit_dim_same_storage_cpu(): + """Same-storage CPU copy of a padded (1, N) array: contiguous run, CPU<->CPU map fallback, exact result.""" + sdfg, node = _make_copy_sdfg( + _padded_unit_spec(dace.dtypes.StorageType.CPU_Heap, "A"), + _padded_unit_spec(dace.dtypes.StorageType.CPU_Heap, "B"), + name="copy_padded_unit_cpu", + libnode_name="cp_padded_cpu", + ) + state = sdfg.start_state + _, inp, in_sub, _, out, out_sub = node.validate(state.sdfg, state, allow_cross_storage=True) + # The accessed row is a single contiguous run (1D-slice special case). + assert in_sub.is_contiguous_subset(inp) + assert out_sub.is_contiguous_subset(out) + # CPU<->CPU multi-element copies never route to a memcpy libnode; they fall back to a map. + assert select_copy_implementation(node, state) == "MappedTasklet" + + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = _compile_no_copynd(sdfg) + + A = np.zeros((1, _PADDED_N), dtype=np.float64) # fresh + contiguous: A.base is None, so no view rejection + B = np.zeros((1, _PADDED_N), dtype=np.float64) + A[0, :] = np.arange(1, _PADDED_N + 1, dtype=np.float64) + exe(A=A, B=B) + np.testing.assert_array_equal(B, A) + + +def test_copy_padded_unit_dim_cross_storage_selection(): + """Cross CPU/GPU copy of a padded (1, N) array is a single contiguous row: flat ``cudaMemcpy``, not pitched. + + With only one row the pitch gap is never crossed, so the row is one contiguous run on both sides and + ``MemcpyCUDA1D`` is exact (a pitched ``cudaMemcpy2D`` would be equivalent but needlessly 2D).""" + for src_storage, dst_storage in ( + (dace.dtypes.StorageType.CPU_Heap, dace.dtypes.StorageType.GPU_Global), + (dace.dtypes.StorageType.GPU_Global, dace.dtypes.StorageType.CPU_Heap), + ): + sdfg, node = _make_copy_sdfg( + _padded_unit_spec(src_storage, "A"), + _padded_unit_spec(dst_storage, "B"), + name="copy_padded_unit_cross", + libnode_name="cp_padded_cross", + ) + state = sdfg.start_state + _, inp, in_sub, _, out, out_sub = node.validate(state.sdfg, state, allow_cross_storage=True) + assert in_sub.is_contiguous_subset(inp) + assert out_sub.is_contiguous_subset(out) + assert select_copy_implementation(node, state) == "MemcpyCUDA1D" + + +def test_copy_padded_multirow_cross_storage_uses_pitched(): + """Cross CPU/GPU copy of a padded multi-row (ROWS, N) array must route to the pitched ``cudaMemcpy2D``. + + With more than one row the inter-row pitch gap is stepped over, so the region is genuinely + non-contiguous and a flat ``MemcpyCUDA1D`` would drag the padding bytes between rows into the copy. + This pins the dangerous direction: were ``is_contiguous_subset`` to ever wrongly report this subset + contiguous, ``_refine_cuda_impl_for_subsets`` would keep the flat copy and silently corrupt the data -- + this test would catch it before the numerical damage. + """ + for src_storage, dst_storage in ( + (dace.dtypes.StorageType.CPU_Heap, dace.dtypes.StorageType.GPU_Global), + (dace.dtypes.StorageType.GPU_Global, dace.dtypes.StorageType.CPU_Heap), + ): + sdfg, node = _make_copy_sdfg( + _padded_multirow_spec(src_storage, "A"), + _padded_multirow_spec(dst_storage, "B"), + name="copy_padded_multirow_cross", + libnode_name="cp_padded_multirow", + ) + state = sdfg.start_state + _, inp, in_sub, _, out, out_sub = node.validate(state.sdfg, state, allow_cross_storage=True) + assert not in_sub.is_contiguous_subset(inp) + assert not out_sub.is_contiguous_subset(out) + assert select_copy_implementation(node, state) == "MemcpyCUDA2D" + + +@pytest.mark.gpu +def test_copy_padded_unit_dim_cross_storage_gpu(): + """Cross CPU->GPU copy of a padded (1, N) array expands to a pitched copy and is numerically exact.""" + import cupy as cp + + sdfg, _ = _make_copy_sdfg( + _padded_unit_spec(dace.dtypes.StorageType.CPU_Heap, "A"), + _padded_unit_spec(dace.dtypes.StorageType.GPU_Global, "B"), + name="copy_padded_unit_h2d", + libnode_name="cp_padded_h2d", + ) + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = _compile_no_copynd(sdfg) + + A = np.zeros((1, _PADDED_N), dtype=np.float64) + A[0, :] = np.arange(1, _PADDED_N + 1, dtype=np.float64) + B = cp.zeros((1, _PADDED_N), dtype=cp.float64) + exe(A=A, B=B) + cp.testing.assert_array_equal(B, cp.asarray(A)) + + +def test_register_copy_expands_with_register_storage(): + """A Register -> Register ``MappedTasklet`` copy expands to a Sequential (thread-level) map.""" + reg = dace.dtypes.StorageType.Register + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[8], storage=reg, transient=True, name="R_in"), + _ArraySpec(shape=[8], storage=reg, transient=True, name="R_out"), + implementation="MappedTasklet", + name="reg_copy_ok", + libnode_name="regcpy", + ) + sdfg.expand_library_nodes() + + found_sequential = False + for n, _ in sdfg.all_nodes_recursive(): + if isinstance(n, dace.sdfg.nodes.MapEntry): + if n.schedule == dace.dtypes.ScheduleType.Sequential: + found_sequential = True + break + assert found_sequential, "RegisterCopy expansion should contain a Sequential map." + + +def test_direct_assignment_cpu_same_storage(): + """``Tasklet`` impl on CPU_Heap -> CPU_Heap (single element) compiles and runs.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[4], storage=dace.dtypes.StorageType.CPU_Heap, subset="2:3", name="A"), + _ArraySpec(shape=[4], storage=dace.dtypes.StorageType.CPU_Heap, subset="1:2", name="B"), + implementation="Tasklet", + name="direct_assign_cpu", + ) + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = _compile_no_copynd(sdfg) + + A = np.arange(4, dtype=np.float64) + B = np.zeros(4, dtype=np.float64) + exe(A=A, B=B) + assert B[1] == A[2] + + +def test_direct_assignment_register_to_register(): + """A size-1 Register -> Register ``Tasklet`` copy expands to a Python tasklet with no map.""" + reg = dace.dtypes.StorageType.Register + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[1], storage=reg, transient=True, subset="0", name="R_in"), + _ArraySpec(shape=[1], storage=reg, transient=True, subset="0", name="R_out"), + implementation="Tasklet", + name="direct_assign_reg", + libnode_name="da", + ) + sdfg.expand_library_nodes() + + found_tasklet = False + found_map = False + for n, _ in sdfg.all_nodes_recursive(): + if (isinstance(n, dace.sdfg.nodes.Tasklet) and n.language == dace.Language.Python + and "_cpy_out = _cpy_in" in n.code.as_string): + found_tasklet = True + if isinstance(n, dace.sdfg.nodes.MapEntry): + found_map = True + assert found_tasklet, "Tasklet impl should produce a Python tasklet with ``_cpy_out = _cpy_in``." + assert not found_map, "Tasklet impl should NOT produce a map." + + +def test_direct_assignment_rejects_multi_element(): + """``Tasklet`` is size-1 only; rejects multi-element copies.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[32], storage=dace.dtypes.StorageType.GPU_Global, transient=True, name="G_in"), + _ArraySpec(shape=[32], storage=dace.dtypes.StorageType.GPU_Shared, transient=True, name="S_out"), + implementation="Tasklet", + name="da_multi_bad", + libnode_name="da_multi_bad", + ) + with pytest.raises(Exception, match="single-element subsets"): + sdfg.expand_library_nodes() + + +def test_direct_assignment_rejects_cross_boundary(): + """``Tasklet`` rejects CPU<->GPU pairings via the same-storage check.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[1], storage=dace.dtypes.StorageType.CPU_Heap, subset="0", name="C_in"), + _ArraySpec(shape=[1], storage=dace.dtypes.StorageType.GPU_Global, transient=True, subset="0", name="G_out"), + implementation="Tasklet", + name="da_cross_bad", + libnode_name="da_cross", + ) + sdfg.validate() + with pytest.raises(Exception, match="storage types must match"): + sdfg.expand_library_nodes() + + +def test_shared_memory_copy_global_to_shared_is_collective(): + """Global -> Shared collective copy emits a CPP tasklet with __syncthreads() and no GPU_ThreadBlock map.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[64], storage=dace.dtypes.StorageType.GPU_Global, transient=True, name="G_in"), + _ArraySpec(shape=[64], storage=dace.dtypes.StorageType.GPU_Shared, transient=True, name="S_out"), + implementation="SharedMemoryCollective", + name="shmcpy_collective", + libnode_name="shmcpy", + ) + sdfg.expand_library_nodes() + + found_syncthreads = False + for n, _ in sdfg.all_nodes_recursive(): + if isinstance(n, dace.sdfg.nodes.Tasklet): + if n.language == dace.Language.CPP and "__syncthreads" in n.code.as_string: + found_syncthreads = True + break + assert found_syncthreads, ("SharedMemoryCopy (Global->Shared) should generate a CPP tasklet " + "containing __syncthreads().") + + # No GPU_ThreadBlock map: the collective tasklet is itself the block-level op. + for n, _ in sdfg.all_nodes_recursive(): + if isinstance(n, dace.sdfg.nodes.MapEntry): + assert n.schedule != dace.dtypes.ScheduleType.GPU_ThreadBlock, ( + "SharedMemoryCopy (Global->Shared) should not generate a " + "GPU_ThreadBlock map.") + + +def _libnode_in_tblock_scope(src_storage, dst_storage, src_subset, dst_subset, src_shape=None, dst_shape=None): + """Build an SDFG with a ``CopyLibraryNode`` nested inside a ``GPU_ThreadBlock`` + map; returns ``(sdfg, libnode, state)`` for scope-aware dispatcher tests.""" + src_shape = src_shape or [16] + dst_shape = dst_shape or [16] + sdfg = dace.SDFG(f"in_tblock_{src_storage.name}_{dst_storage.name}") + sdfg.add_array("src", + src_shape, + dace.float64, + storage=src_storage, + transient=(src_storage != dace.dtypes.StorageType.CPU_Heap)) + sdfg.add_array("dst", + dst_shape, + dace.float64, + storage=dst_storage, + transient=(dst_storage != dace.dtypes.StorageType.CPU_Heap)) + state = sdfg.add_state("main") + src_acc = state.add_access("src") + dst_acc = state.add_access("dst") + ome, omx = state.add_map("device_map", {"bi": "0:1"}, schedule=dace.dtypes.ScheduleType.GPU_Device) + ime, imx = state.add_map("tblock_map", {"ti": "0:16"}, schedule=dace.dtypes.ScheduleType.GPU_ThreadBlock) + libnode = CopyLibraryNode(name="cp") + state.add_memlet_path(src_acc, + ome, + ime, + libnode, + dst_conn=CopyLibraryNode.INPUT_CONNECTOR_NAME, + memlet=dace.memlet.Memlet(f"src[{src_subset}]")) + state.add_memlet_path(libnode, + imx, + omx, + dst_acc, + src_conn=CopyLibraryNode.OUTPUT_CONNECTOR_NAME, + memlet=dace.memlet.Memlet(f"dst[{dst_subset}]")) + return sdfg, libnode, state + + +# Auto-dispatch unit tests for Shared-involved copies. One exact-impl +# assertion per unique routing rule (symmetric directions share the rule); +# end-to-end correctness lives in the ``test_copy_*_roundtrip`` tests. +# The "no single-element -> MappedTasklet" invariant is exhaustively +# covered by ``test_auto_dispatch_single_element_never_mapped_tasklet``. + + +def test_auto_dispatch_multi_element_shared_register_routes_to_mapped_tasklet(): + """Rule 2 (multi): Shared <-> Register multi-element -> ``MappedTasklet``.""" + sdfg, node = _make_copy_sdfg( + _ArraySpec(shape=[8], storage=dace.dtypes.StorageType.GPU_Shared, transient=True, name="S_in"), + _ArraySpec(shape=[8], storage=dace.dtypes.StorageType.Register, transient=True, name="R_out"), + name="auto_shm_to_reg", + libnode_name="cp_shm_reg", + ) + assert select_copy_implementation(node, sdfg.start_state) == "MappedTasklet" + + +def test_auto_dispatch_single_element_shared_register_routes_to_tasklet(): + """Rule 2 (single): Shared <-> Register single-element -> ``Tasklet``.""" + sdfg, node = _make_copy_sdfg( + _ArraySpec(shape=[8], storage=dace.dtypes.StorageType.GPU_Shared, transient=True, subset="3", name="S_in"), + _ArraySpec(shape=[1], storage=dace.dtypes.StorageType.Register, transient=True, subset="0", name="R_out"), + name="auto_shm_reg_single", + libnode_name="cp_shm_reg_single", + ) + assert select_copy_implementation(node, sdfg.start_state) == "Tasklet" + + +def test_auto_dispatch_global_shared_outside_tblock_routes_to_collective(): + """Rule 3 (multi): Global <-> Shared outside a ThreadBlock map -> ``SharedMemoryCollective``.""" + sdfg, node = _make_copy_sdfg( + _ArraySpec(shape=[64], storage=dace.dtypes.StorageType.GPU_Global, transient=True, name="G_in"), + _ArraySpec(shape=[64], storage=dace.dtypes.StorageType.GPU_Shared, transient=True, name="S_out"), + name="auto_global_to_shm", + libnode_name="cp_global_shm", + ) + assert select_copy_implementation(node, sdfg.start_state) == "SharedMemoryCollective" + + +def test_auto_dispatch_single_element_global_shared_outside_tblock_still_collective(): + """Rule 3 (single): Global <-> Shared single-element outside ThreadBlock -> ``SharedMemoryCollective``.""" + sdfg, node = _make_copy_sdfg( + _ArraySpec(shape=[64], storage=dace.dtypes.StorageType.GPU_Global, transient=True, subset="5", name="G_in"), + _ArraySpec(shape=[8], storage=dace.dtypes.StorageType.GPU_Shared, transient=True, subset="3", name="S_out"), + name="auto_global_shm_single", + libnode_name="cp_global_shm_single", + ) + assert select_copy_implementation(node, sdfg.start_state) == "SharedMemoryCollective" + + +def test_auto_dispatch_shared_shared_outside_tblock_routes_to_collective(): + """Rule 3 (Shared<->Shared): outside ThreadBlock -> ``SharedMemoryCollective``.""" + sdfg, node = _make_copy_sdfg( + _ArraySpec(shape=[32], storage=dace.dtypes.StorageType.GPU_Shared, transient=True, name="S_a"), + _ArraySpec(shape=[32], storage=dace.dtypes.StorageType.GPU_Shared, transient=True, name="S_b"), + name="auto_shm_to_shm", + libnode_name="cp_shm_shm", + ) + assert select_copy_implementation(node, sdfg.start_state) == "SharedMemoryCollective" + + +def test_auto_dispatch_global_shared_inside_tblock_routes_to_mapped_tasklet(): + """Rule 4 (multi): Global -> Shared *inside* a ThreadBlock map is per-thread -> ``MappedTasklet``.""" + sdfg, node, state = _libnode_in_tblock_scope(dace.dtypes.StorageType.GPU_Global, + dace.dtypes.StorageType.GPU_Shared, + src_subset="0:4", + dst_subset="0:4") + assert select_copy_implementation(node, state) == "MappedTasklet" + + +def test_auto_dispatch_global_shared_inside_tblock_single_element_routes_to_tasklet(): + """Rule 4 (single): Global -> Shared single-element *inside* a ThreadBlock map -> ``Tasklet``.""" + sdfg, node, state = _libnode_in_tblock_scope(dace.dtypes.StorageType.GPU_Global, + dace.dtypes.StorageType.GPU_Shared, + src_subset="ti", + dst_subset="ti") + assert select_copy_implementation(node, state) == "Tasklet" + + +def test_shared_memory_collective_single_element_emits_syncthreads(): + """Single-element collective Global -> Shared must emit ``__syncthreads()`` (the barrier is volume-independent).""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[64], storage=dace.dtypes.StorageType.GPU_Global, transient=True, subset="5", name="G_in"), + _ArraySpec(shape=[8], storage=dace.dtypes.StorageType.GPU_Shared, transient=True, subset="3", name="S_out"), + name="auto_global_shm_single_e2e", + libnode_name="cp_global_shm_single_e2e", + ) + sdfg.expand_library_nodes() + assert any(isinstance(n, dace.sdfg.nodes.Tasklet) and n.language == dace.Language.CPP + and "__syncthreads" in n.code.as_string + for n, _ in sdfg.all_nodes_recursive()), \ + "Single-element collective Global->Shared must still emit __syncthreads()." + + +_SINGLE_ELT_STORAGES = [ + dace.dtypes.StorageType.CPU_Heap, + dace.dtypes.StorageType.GPU_Global, + dace.dtypes.StorageType.GPU_Shared, + dace.dtypes.StorageType.Register, +] + + +@pytest.mark.parametrize("src_storage", _SINGLE_ELT_STORAGES) +@pytest.mark.parametrize("dst_storage", _SINGLE_ELT_STORAGES) +def test_auto_dispatch_single_element_never_mapped_tasklet(src_storage, dst_storage): + """Invariant: no single-element copy is ever routed to ``MappedTasklet`` (a 0-D map crashes in propagation), over every storage pair.""" + src_kwargs = {"transient": True} if src_storage != dace.dtypes.StorageType.CPU_Heap else {} + dst_kwargs = {"transient": True} if dst_storage != dace.dtypes.StorageType.CPU_Heap else {} + sdfg, node = _make_copy_sdfg( + _ArraySpec(shape=[8], storage=src_storage, subset="3", name="src", **src_kwargs), + _ArraySpec(shape=[8], storage=dst_storage, subset="5", name="dst", **dst_kwargs), + name=f"auto_single_{src_storage.name}_{dst_storage.name}", + libnode_name=f"cp_single_{src_storage.name}_{dst_storage.name}", + ) + state = sdfg.start_state + impl = select_copy_implementation(node, state) + assert impl != "MappedTasklet", ( + f"Single-element {src_storage.name} -> {dst_storage.name} routed to MappedTasklet; " + "single-element copies must use Tasklet / MemcpyCUDA1D / SharedMemoryCollective.") + + +def test_shared_memory_copy_rejects_no_shared(): + """SharedMemoryCopy expansion rejects if neither side is GPU_Shared.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[32], storage=dace.dtypes.StorageType.GPU_Global, transient=True, name="G_in"), + _ArraySpec(shape=[32], storage=dace.dtypes.StorageType.Register, transient=True, name="R_out"), + implementation="SharedMemoryCollective", + name="shmcpy_bad", + libnode_name="shmcpy_bad", + ) + with pytest.raises(Exception, match="GPU_Shared / GPU_Global storages"): + sdfg.expand_library_nodes() + + +def test_shared_memory_copy_rejects_cpu(): + """SharedMemoryCopy expansion rejects CPU_Heap storage.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[32], storage=dace.dtypes.StorageType.CPU_Heap, name="C_in"), + _ArraySpec(shape=[32], storage=dace.dtypes.StorageType.GPU_Shared, transient=True, name="S_out"), + implementation="SharedMemoryCollective", + name="shmcpy_cpu", + libnode_name="shmcpy_cpu", + ) + with pytest.raises(Exception, match="GPU_Shared / GPU_Global storages"): + sdfg.expand_library_nodes() + + +def test_shared_memory_copy_rejects_inside_tblock_map(): + """A collective ``SharedMemoryCollective`` copy nested in a GPU_ThreadBlock map raises at expansion.""" + sdfg = dace.SDFG("shmcpy_in_tblock") + sdfg.add_array("A", [256], dace.float64, dace.dtypes.StorageType.GPU_Global) + sdfg.add_array("B", [256], dace.float64, dace.dtypes.StorageType.GPU_Global) + sdfg.add_array("shmem", [32], dace.float64, dace.dtypes.StorageType.GPU_Shared, transient=True) + + state = sdfg.add_state("main") + a = state.add_access("A") + shm = state.add_access("shmem") + + ome, omx = state.add_map("device_map", {"bi": "0:256:32"}, schedule=dace.dtypes.ScheduleType.GPU_Device) + # ThreadBlock map is an invalid parent for a collective copy. + ime, imx = state.add_map("tblock_map", {"ti": "0:32"}, schedule=dace.dtypes.ScheduleType.GPU_ThreadBlock) + + libnode = CopyLibraryNode(name="shmcpy_bad") + libnode.implementation = "SharedMemoryCollective" + + state.add_memlet_path(a, + ome, + ime, + libnode, + dst_conn=CopyLibraryNode.INPUT_CONNECTOR_NAME, + memlet=dace.Memlet("A[bi:bi+32]")) + state.add_memlet_path(libnode, + imx, + omx, + shm, + src_conn=CopyLibraryNode.OUTPUT_CONNECTOR_NAME, + memlet=dace.Memlet("shmem[0:32]")) + + with pytest.raises(Exception, match="GPU_ThreadBlock"): + sdfg.expand_library_nodes() + + +@pytest.mark.gpu +def test_copy_roundtrip_variant_a_cooperative_load(): + """Variant A: collective load OUTSIDE the tblock_map (block-cooperative ``dace::CopyND`` + ``__syncthreads()``), per-thread writeback inside it, round-tripping through Global ``B``.""" + import cupy as cp + + N = 256 + TILE = 32 + sdfg = dace.SDFG("roundtrip_variant_a") + sdfg.add_array("A", [N], dace.float64, dace.dtypes.StorageType.GPU_Global) + sdfg.add_array("B", [N], dace.float64, dace.dtypes.StorageType.GPU_Global) + sdfg.add_array("tile", [TILE], dace.float64, dace.dtypes.StorageType.GPU_Shared, transient=True) + + state = sdfg.add_state("main") + a = state.add_access("A") + tile = state.add_access("tile") + b = state.add_access("B") + + ome, omx = state.add_map("device_map", {"bi": f"0:{N}:{TILE}"}, schedule=dace.dtypes.ScheduleType.GPU_Device) + + # Cooperative load: libnode sits OUTSIDE the tblock map (between ome and ime). + load = CopyLibraryNode(name="load_a_to_tile") + state.add_memlet_path(a, + ome, + load, + dst_conn=CopyLibraryNode.INPUT_CONNECTOR_NAME, + memlet=dace.Memlet(f"A[bi:bi+{TILE}]")) + state.add_edge(load, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, tile, None, dace.Memlet(f"tile[0:{TILE}]")) + + ime, imx = state.add_map("tblock_map", {"ti": f"0:{TILE}"}, schedule=dace.dtypes.ScheduleType.GPU_ThreadBlock) + t = state.add_tasklet("writeback", {"v"}, {"o"}, "o = v") + state.add_memlet_path(tile, ime, t, dst_conn="v", memlet=dace.Memlet("tile[ti]")) + state.add_memlet_path(t, imx, omx, b, src_conn="o", memlet=dace.Memlet("B[bi+ti]")) + + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + + A = cp.arange(N, dtype=cp.float64) * 3.0 + 0.5 + B = cp.zeros(N, dtype=cp.float64) + sdfg(A=A, B=B) + cp.testing.assert_array_equal(B, A) + + +@pytest.mark.gpu +def test_copy_roundtrip_variant_b_per_thread_load(): + """Variant B: per-thread load INSIDE the tblock_map -- each thread copies ``A[bi+ti] -> tile[ti] -> B[bi+ti]`` via its own ``Tasklet`` (no block-collective).""" + import cupy as cp + + N = 256 + TILE = 32 + sdfg = dace.SDFG("roundtrip_variant_b") + sdfg.add_array("A", [N], dace.float64, dace.dtypes.StorageType.GPU_Global) + sdfg.add_array("B", [N], dace.float64, dace.dtypes.StorageType.GPU_Global) + sdfg.add_array("tile", [TILE], dace.float64, dace.dtypes.StorageType.GPU_Shared, transient=True) + + state = sdfg.add_state("main") + a = state.add_access("A") + tile = state.add_access("tile") + b = state.add_access("B") + + ome, omx = state.add_map("device_map", {"bi": f"0:{N}:{TILE}"}, schedule=dace.dtypes.ScheduleType.GPU_Device) + ime, imx = state.add_map("tblock_map", {"ti": f"0:{TILE}"}, schedule=dace.dtypes.ScheduleType.GPU_ThreadBlock) + + # Per-thread load: libnode INSIDE the tblock map -- each thread copies one cell. + load = CopyLibraryNode(name="load_a_to_tile_per_thread") + state.add_memlet_path(a, + ome, + ime, + load, + dst_conn=CopyLibraryNode.INPUT_CONNECTOR_NAME, + memlet=dace.Memlet("A[bi+ti]")) + state.add_edge(load, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, tile, None, dace.Memlet("tile[ti]")) + + # Per-thread store: libnode INSIDE the tblock map -- each thread writes its cell. + store = CopyLibraryNode(name="store_tile_to_b_per_thread") + state.add_edge(tile, None, store, CopyLibraryNode.INPUT_CONNECTOR_NAME, dace.Memlet("tile[ti]")) + state.add_memlet_path(store, + imx, + omx, + b, + src_conn=CopyLibraryNode.OUTPUT_CONNECTOR_NAME, + memlet=dace.Memlet("B[bi+ti]")) + + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + + A = cp.arange(N, dtype=cp.float64) * 5.0 - 2.0 + B = cp.zeros(N, dtype=cp.float64) + sdfg(A=A, B=B) + cp.testing.assert_array_equal(B, A) + + +@pytest.mark.gpu +def test_copy_full_pipeline_roundtrip(): + """Pipeline: Global -> Shared (collective) -> per-thread (Register -> Register -> Shared) -> Global; exercises auto-dispatched Shared<->Register libnodes alongside the block-cooperative load.""" + import cupy as cp + + N = 256 + TILE = 32 + sdfg = dace.SDFG("full_pipeline_roundtrip") + sdfg.add_array("A", [N], dace.float64, dace.dtypes.StorageType.GPU_Global) + sdfg.add_array("B", [N], dace.float64, dace.dtypes.StorageType.GPU_Global) + sdfg.add_array("shm_in", [TILE], dace.float64, dace.dtypes.StorageType.GPU_Shared, transient=True) + sdfg.add_array("shm_out", [TILE], dace.float64, dace.dtypes.StorageType.GPU_Shared, transient=True) + sdfg.add_array("reg_a", [1], dace.float64, dace.dtypes.StorageType.Register, transient=True) + sdfg.add_array("reg_b", [1], dace.float64, dace.dtypes.StorageType.Register, transient=True) + + state = sdfg.add_state("main") + a = state.add_access("A") + shm_in = state.add_access("shm_in") + shm_out = state.add_access("shm_out") + b = state.add_access("B") + + ome, omx = state.add_map("device_map", {"bi": f"0:{N}:{TILE}"}, schedule=dace.dtypes.ScheduleType.GPU_Device) + + # Global -> Shared (collective load). + load = CopyLibraryNode(name="load_a_to_shm") + state.add_memlet_path(a, + ome, + load, + dst_conn=CopyLibraryNode.INPUT_CONNECTOR_NAME, + memlet=dace.Memlet(f"A[bi:bi+{TILE}]")) + state.add_edge(load, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, shm_in, None, dace.Memlet(f"shm_in[0:{TILE}]")) + + # Single GPU_ThreadBlock map carries: + # Shared(shm_in) -> Register(reg_a) -> Register(reg_b) -> Shared(shm_out) + # -> Global(B) (per-thread tasklet for the last leg) + ime, imx = state.add_map("tblock_map", {"ti": f"0:{TILE}"}, schedule=dace.dtypes.ScheduleType.GPU_ThreadBlock) + s2r = CopyLibraryNode(name="shm_to_reg_a") + r2r = CopyLibraryNode(name="reg_a_to_reg_b") + r2s = CopyLibraryNode(name="reg_b_to_shm") + reg_a = state.add_access("reg_a") + reg_b = state.add_access("reg_b") + + state.add_memlet_path(shm_in, + ime, + s2r, + dst_conn=CopyLibraryNode.INPUT_CONNECTOR_NAME, + memlet=dace.Memlet("shm_in[ti]")) + state.add_edge(s2r, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, reg_a, None, dace.Memlet("reg_a[0]")) + state.add_edge(reg_a, None, r2r, CopyLibraryNode.INPUT_CONNECTOR_NAME, dace.Memlet("reg_a[0]")) + state.add_edge(r2r, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, reg_b, None, dace.Memlet("reg_b[0]")) + state.add_edge(reg_b, None, r2s, CopyLibraryNode.INPUT_CONNECTOR_NAME, dace.Memlet("reg_b[0]")) + state.add_memlet_path(r2s, + imx, + shm_out, + src_conn=CopyLibraryNode.OUTPUT_CONNECTOR_NAME, + memlet=dace.Memlet("shm_out[ti]")) + + # Per-thread Shared -> Global writeback via a tasklet -- avoids a + # second block-collective copy in the same kernel. + ime2, imx2 = state.add_map("writeback_map", {"tj": f"0:{TILE}"}, schedule=dace.dtypes.ScheduleType.GPU_ThreadBlock) + tw = state.add_tasklet("writeback", {"v"}, {"o"}, "o = v") + state.add_memlet_path(shm_out, ime2, tw, dst_conn="v", memlet=dace.Memlet("shm_out[tj]")) + state.add_memlet_path(tw, imx2, omx, b, src_conn="o", memlet=dace.Memlet("B[bi+tj]")) + + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + + A = cp.arange(N, dtype=cp.float64) * 2.0 + 1.0 + B = cp.zeros(N, dtype=cp.float64) + sdfg(A=A, B=B) + cp.testing.assert_array_equal(B, A) + + +def test_copy_pure_cpu_2d(): + """Pure expansion on a 2D slice copy, CPU_Heap.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[10, 20], storage=dace.dtypes.StorageType.CPU_Heap, subset="2:8, 5:15", name="A"), + _ArraySpec(shape=[10, 20], storage=dace.dtypes.StorageType.CPU_Heap, subset="0:6, 0:10", name="B"), + implementation="MappedTasklet", + name="copy_2d_cpu", + libnode_name="cp2d", + ) + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = _compile_no_copynd(sdfg) + + A = np.arange(200, dtype=np.float64).reshape(10, 20).copy() + B = np.zeros((10, 20), dtype=np.float64) + exe(A=A, B=B) + + np.testing.assert_array_equal(B[0:6, 0:10], A[2:8, 5:15]) + + +@pytest.mark.gpu +def test_copy_single_element_h2d(): + """Single-element host -> GPU copy compiles and round-trips.""" + pytest.importorskip('cupy') + import cupy as cp + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[1], storage=dace.dtypes.StorageType.CPU_Heap, name="host"), + _ArraySpec(shape=[1], storage=dace.dtypes.StorageType.GPU_Global, name="dev"), + name="single_elem_h2d", + libnode_name="copy_h2d", + ) + + host = np.array([3.14159], dtype=np.float64) + dev = cp.zeros(1, dtype=cp.float64) + + _compile_no_copynd(sdfg)(host=host, dev=dev) + np.testing.assert_allclose(cp.asnumpy(dev), host) + + +@pytest.mark.gpu +def test_copy_two_element_h2d(): + """A 2-element host -> GPU copy compiles and round-trips (pointer-typed connectors, unlike single element).""" + pytest.importorskip('cupy') + import cupy as cp + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[2], storage=dace.dtypes.StorageType.CPU_Heap, name="host"), + _ArraySpec(shape=[2], storage=dace.dtypes.StorageType.GPU_Global, name="dev"), + name="two_elem_h2d", + libnode_name="copy_h2d_2", + ) + + host = np.array([1.0, 2.0], dtype=np.float64) + dev = cp.zeros(2, dtype=cp.float64) + _compile_no_copynd(sdfg)(host=host, dev=dev) + np.testing.assert_allclose(cp.asnumpy(dev), host) + + +@pytest.mark.gpu +def test_copy_single_element_d2h(): + """Single-element GPU -> host copy compiles and round-trips.""" + pytest.importorskip('cupy') + import cupy as cp + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=[1], storage=dace.dtypes.StorageType.GPU_Global, name="dev"), + _ArraySpec(shape=[1], storage=dace.dtypes.StorageType.CPU_Heap, name="host"), + name="single_elem_d2h", + libnode_name="copy_d2h", + ) + + dev = cp.array([2.71828], dtype=cp.float64) + host = np.zeros(1, dtype=np.float64) + + _compile_no_copynd(sdfg)(host=host, dev=dev) + np.testing.assert_allclose(host, cp.asnumpy(dev)) + + +# Legacy direct-edge miscompile regression pins: each test builds the SDFG twice +# -- with a CopyLibraryNode and with the canonical direct AN -> AN edge -- and checks +# both against a NumPy for-loop. The libnode's advantage is rank-mismatch reshapes +# with per-side layout strides, which the legacy memcpy path miscompiles or fails to +# compile. The legacy-fails assertions are informational: if legacy ever produces +# correct output, the test fails and should be deleted (the advantage is gone). + + +def _legacy_fails(sdfg_leg: dace.SDFG, expected: np.ndarray, run) -> bool: + """``True`` if compiling/running the legacy SDFG raises OR produces output diverging from ``expected``. + + :param sdfg_leg: SDFG with libnodes already replaced by direct edges. + :param expected: NumPy ground truth. + :param run: a callable ``run(exe) -> np.ndarray`` that runs the compiled SDFG and returns the dst array. + """ + # ``compiler.cpu.explicit_copy`` defaults to on, which lifts this very direct edge to a + # CopyLibraryNode -- comparing against it would compare the libnode path with itself. + try: + with dace.config.set_temporary('compiler', 'cpu', 'explicit_copy', value=False): + exe = sdfg_leg.compile() + return not np.array_equal(run(exe), expected) + except Exception: + return True + + +def test_legacy_silently_miscompiles_rank_mismatch_fortran_collapse(): + """Pin: legacy direct-edge miscompiles a 4D->2D Fortran-packed reshape.""" + src = _ArraySpec(shape=(2, 3, 4, 5), + storage=dace.dtypes.StorageType.CPU_Heap, + strides=(1, 2, 6, 24), + total_size=120) + dst = _ArraySpec(shape=(6, 20), storage=dace.dtypes.StorageType.CPU_Heap, strides=(1, 6), total_size=120) + sdfg_lib, _ = _make_copy_sdfg(src, dst, name="legacy_fortran_collapse_lib") + sdfg_leg = _make_legacy_copy_sdfg(src, dst, name="legacy_fortran_collapse_leg") + + A = np.arange(120, dtype=np.float64).reshape(2, 3, 4, 5, order='F').copy(order='F') + expected = np.zeros((6, 20), dtype=np.float64, order='F') + # Fortran-order flat walk: src index (i,j,k,l) -> flat n = i + j*2 + k*6 + l*24 + # dst index (p, q) -> flat n = p + q*6 + flat = np.empty(120, dtype=np.float64) + for l in range(5): + for k in range(4): + for j in range(3): + for i in range(2): + flat[i + j * 2 + k * 6 + l * 24] = A[i, j, k, l] + for q in range(20): + for p in range(6): + expected[p, q] = flat[p + q * 6] + + B_lib = np.zeros((6, 20), dtype=np.float64, order='F') + sdfg_lib.expand_library_nodes() + _compile_no_copynd(sdfg_lib)(src=A, dst=B_lib) + np.testing.assert_array_equal(B_lib, expected) + + def run(exe): + out = np.zeros((6, 20), dtype=np.float64, order='F') + exe(src=A, dst=out) + return out + + assert _legacy_fails(sdfg_leg, expected, run), ("Legacy direct-edge no longer fails on 4D->2D Fortran reshape; " + "remove this test, the libnode advantage is gone.") + + +def test_single_element_in_kernel_register_to_gpu_global_routes_to_tasklet(): + """Single-element in-kernel Register -> GPU_Global routes to a direct Tasklet, not MappedTasklet.""" + sdfg = dace.SDFG('reg_to_gpuglobal_in_kernel') + sdfg.add_array('R', [1, 1, 1], dace.float64, dace.StorageType.Register, transient=True) + sdfg.add_array('G', [4, 4, 4], dace.float64, dace.StorageType.GPU_Global, transient=True) + state = sdfg.add_state('s') + + # Wrap the copy inside a GPU_Device map so ``is_devicelevel_gpu`` returns True. + me, mx = state.add_map('kernel', dict(i='0:1'), schedule=dace.dtypes.ScheduleType.GPU_Device) + r = state.add_access('R') + g = state.add_access('G') + libnode = CopyLibraryNode(name='reg_to_g') + state.add_node(libnode) + state.add_memlet_path(me, r, memlet=dace.Memlet()) + state.add_edge(r, None, libnode, CopyLibraryNode.INPUT_CONNECTOR_NAME, dace.Memlet('R[0, 0, 0]')) + state.add_edge(libnode, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, g, None, dace.Memlet('G[0, 0, 0]')) + state.add_memlet_path(g, mx, memlet=dace.Memlet()) + + sdfg.expand_library_nodes() + + nsdfg_count = sum(1 for n, _ in sdfg.all_nodes_recursive() if isinstance(n, dace.nodes.NestedSDFG)) + assert nsdfg_count == 0, (f"Single-element in-kernel copy should expand to a direct Tasklet, " + f"not a NestedSDFG; got {nsdfg_count} NestedSDFG(s).") + assignments = [ + n for n, _ in sdfg.all_nodes_recursive() + if isinstance(n, dace.nodes.Tasklet) and '_cpy_out = _cpy_in' in n.code.as_string + ] + assert assignments, "Expected at least one ``_cpy_out = _cpy_in`` Tasklet from the expansion." + + +def test_register_location_detection(): + """Register location detection distinguishes in-kernel from host-side copies.""" + sdfg = dace.SDFG('register_location_detection') + sdfg.add_array('R', [1], dace.float64, dace.StorageType.Register, transient=True) + sdfg.add_array('G', [1], dace.float64, dace.StorageType.GPU_Global, transient=True) + state = sdfg.add_state('s') + + r = state.add_access('R') + g = state.add_access('G') + libnode = CopyLibraryNode(name='reg_to_g') + state.add_node(libnode) + state.add_edge(r, None, libnode, CopyLibraryNode.INPUT_CONNECTOR_NAME, dace.Memlet('R[0]')) + state.add_edge(libnode, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, g, None, dace.Memlet('G[0]')) + + sdfg.expand_library_nodes() + + nsdfg_count = sum(1 for n, _ in sdfg.all_nodes_recursive() if isinstance(n, dace.nodes.NestedSDFG)) + assert nsdfg_count == 0, (f"Single-element in-kernel copy should expand to a direct Memcpy (cross-boundary), " + f"not a NestedSDFG; got {nsdfg_count} NestedSDFG(s).") + assignments = [ + n for n, _ in sdfg.all_nodes_recursive() + if isinstance(n, dace.nodes.Tasklet) and 'cudaMemcpy' in n.code.as_string + ] + assert assignments, "Expected at least one ``cudaMemcpy`` Tasklet from the expansion." + + +def test_cuda2d_pitch_params_branches(): + """``cuda2d_pitch_params`` returns element-count ``(dpitch, spitch, width, height)`` for each + supported 2D stride pattern and ``None`` otherwise. It is the single source of truth shared by + the ``MemcpyCUDA2D`` selector gate and the expander, so selector and expander cannot drift.""" + # Contiguous rows (inner stride 1): pitch = outer stride, width = columns, height = rows. + assert cuda2d_pitch_params([4, 3], [3, 1], [10, 1]) == (10, 3, 3, 4) + # Contiguous columns (outer stride 1): the roles of the two axes swap. + assert cuda2d_pitch_params([4, 3], [1, 4], [1, 8]) == (8, 4, 4, 3) + # Neither axis unit-strided, but outer/inner ratio equals the inner width -> one strided run. + assert cuda2d_pitch_params([4, 2], [4, 2], [6, 3]) == (3, 2, 1, 8) + # No single cudaMemcpy2DAsync expresses this pattern. + assert cuda2d_pitch_params([4, 3], [5, 2], [5, 2]) is None + + +@contextlib.contextmanager +def _pinned_transfer_threshold(value): + """Pin ``compiler.cpu.parallel_transfer_min_elements`` so Auto selection is deterministic.""" + orig = dace.config.Config.get("compiler", "cpu", "parallel_transfer_min_elements") + dace.config.Config.set("compiler", "cpu", "parallel_transfer_min_elements", value=value) + try: + yield + finally: + dace.config.Config.set("compiler", "cpu", "parallel_transfer_min_elements", value=orig) + + +def _cpu_copy_sdfg(extent, name): + """Single-state CPU_Heap -> CPU_Heap ``CopyLibraryNode`` copy over ``0:extent``.""" + sdfg = dace.SDFG(name) + sdfg.add_array("src", [extent], dace.float64, dace.dtypes.StorageType.CPU_Heap) + sdfg.add_array("dst", [extent], dace.float64, dace.dtypes.StorageType.CPU_Heap) + state = sdfg.add_state("s") + libnode = CopyLibraryNode(name="cp") + state.add_edge(state.add_access("src"), None, libnode, CopyLibraryNode.INPUT_CONNECTOR_NAME, + dace.Memlet(f"src[0:{extent}]")) + state.add_edge(libnode, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, state.add_access("dst"), None, + dace.Memlet(f"dst[0:{extent}]")) + sdfg.validate() + return sdfg, libnode + + +def _generated_code(sdfg): + return "\n".join(obj.code for obj in sdfg.generate_code()) + + +def test_copy_below_threshold_emits_memcpy(): + """A constant-size CPU copy below the threshold lowers to a single ``memcpy``, not an OpenMP loop.""" + with _pinned_transfer_threshold(1024): + sdfg, libnode = _cpu_copy_sdfg(100, "copy_below_threshold") + sdfg.expand_library_nodes(recursive=True) + assert libnode.implementation == 'MemcpyCPU' + code = _generated_code(sdfg) + assert 'memcpy(' in code + assert '#pragma omp parallel for' not in code + + +def test_copy_at_threshold_emits_omp_parallel_for(): + """A constant-size CPU copy at/above the threshold lowers to an OpenMP element map, not ``memcpy``.""" + with _pinned_transfer_threshold(1024): + sdfg, libnode = _cpu_copy_sdfg(4096, "copy_at_threshold") + sdfg.expand_library_nodes(recursive=True) + assert libnode.implementation == 'MappedTasklet' + code = _generated_code(sdfg) + assert '#pragma omp parallel for' in code + assert 'memcpy(' not in code + + +def test_copy_symbolic_size_emits_omp_parallel_for(): + """A symbolic (compile-time-unknown) CPU copy size is assumed large, so it takes the same + OpenMP-parallel path as a large constant, never the single-call ``memcpy``.""" + with _pinned_transfer_threshold(1024): + n = dace.symbol('N_copy_symbolic') + sdfg, libnode = _cpu_copy_sdfg(n, "copy_symbolic_size") + sdfg.expand_library_nodes(recursive=True) + assert libnode.implementation == 'MappedTasklet' + code = _generated_code(sdfg) + assert '#pragma omp parallel for' in code + assert 'memcpy(' not in code + + +# --- Regression pins for the codegen bugs fixed alongside the explicit-copy lowering --- + + +def test_collapse_expands_tiled_dimension(): + """A tiled range addresses ``tile`` contiguous elements per step; collapsing must yield two dims.""" + subset = dace.subsets.Range.from_string('1, 0:10:8:2, 3') + shape, strides = collapse_shape_and_strides(subset, (64, 4, 1)) + assert [int(s) for s in shape] == [2, 2] + assert [int(s) for s in strides] == [32, 4] + + +def test_copy_tiled_subset_rank_mismatch_numbers(): + """The 1-D walker over a tiled C-layout source visits tile-innermost, matching subset order.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=(2, 16, 4), storage=dace.dtypes.StorageType.CPU_Heap, subset="1, 0:10:8:2, 3", name="A"), + _ArraySpec(shape=(4, ), storage=dace.dtypes.StorageType.CPU_Heap, name="B"), + name="copy_tiled_subset", + ) + sdfg.validate() + sdfg.expand_library_nodes() + exe = _compile_no_copynd(sdfg) + A = np.arange(128, dtype=np.float64).reshape(2, 16, 4).copy() + B = np.zeros(4, dtype=np.float64) + exe(A=A, B=B) + np.testing.assert_array_equal(B, np.concatenate([A[1, 0:2, 3], A[1, 8:10, 3]])) + + +def test_copy_symbolic_shapes_that_cannot_be_compared_are_not_refused(): + """``ceiling(N/2)`` vs ``floor(N/2)`` is "cannot tell", not "different" -- expansion must proceed.""" + n = dace.symbol('N_sym_cmp') + sdfg = dace.SDFG('copy_symbolic_shape_cmp') + sdfg.add_array('A', [n], dace.float64) + sdfg.add_array('B', [n], dace.float64) + state = sdfg.add_state('main') + libnode = CopyLibraryNode(name='cp') + state.add_edge(state.add_access('A'), None, libnode, CopyLibraryNode.INPUT_CONNECTOR_NAME, + dace.memlet.Memlet('A[0:int_ceil(N_sym_cmp, 2)]')) + state.add_edge(libnode, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, state.add_access('B'), None, + dace.memlet.Memlet('B[0:int_floor(N_sym_cmp, 2)]')) + sdfg.expand_library_nodes() + assert libnode.implementation == 'MappedTasklet' + + +def test_copy_zero_element_expands_to_an_empty_map(): + """A zero-element copy transfers nothing, so mismatched per-dim shapes are not grounds to refuse.""" + sdfg, _ = _make_copy_sdfg( + _ArraySpec(shape=(20, 20), storage=dace.dtypes.StorageType.CPU_Heap, subset="2:17, 2:2", name="A"), + _ArraySpec(shape=(20, 20), storage=dace.dtypes.StorageType.CPU_Heap, subset="2:18, 3:3", name="B"), + name="copy_zero_size", + ) + sdfg.expand_library_nodes() + sdfg.validate() + A = np.arange(400, dtype=np.float64).reshape(20, 20).copy() + B = np.zeros((20, 20), dtype=np.float64) + _compile_no_copynd(sdfg)(A=A, B=B) + np.testing.assert_array_equal(B, np.zeros((20, 20))) + + +def test_copy_between_two_cpu_storages_is_a_memcpy(): + """CPU_ThreadLocal and CPU_Heap differ only in the allocator; a plain memcpy between them is correct.""" + sdfg, libnode = _make_copy_sdfg( + _ArraySpec(shape=(16, ), storage=dace.dtypes.StorageType.CPU_ThreadLocal, transient=True, name="A"), + _ArraySpec(shape=(16, ), storage=dace.dtypes.StorageType.CPU_Heap, transient=True, name="B"), + implementation="MemcpyCPU", + name="copy_threadlocal_to_heap", + ) + sdfg.expand_library_nodes() + assert 'memcpy(' in _generated_code(sdfg) + + +def test_copy_struct_member_name_is_a_valid_identifier(): + """A struct member name carries a '.', which cannot appear in the emitted C++ SDFG/function name.""" + sdfg = dace.SDFG('copy_struct_member_name') + sdfg.add_datadesc('S', dace.data.Structure({'f': dace.data.Array(dace.float64, (16, ))}, name='S')) + sdfg.add_array('B', [16], dace.float64) + state = sdfg.add_state('main') + libnode = CopyLibraryNode(name='copy_S.f_to_B') + state.add_edge(state.add_access('S.f'), None, libnode, CopyLibraryNode.INPUT_CONNECTOR_NAME, + dace.memlet.Memlet('S.f[0:16]')) + state.add_edge(libnode, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, state.add_access('B'), None, + dace.memlet.Memlet('B[0:16]')) + ctx = _make_expansion_sdfg(libnode, state, allow_cross_storage=True) + assert '.' not in ctx.sdfg.name + assert '.' not in ctx.state.label + + +def test_single_element_in_kernel_cross_boundary_is_a_tasklet(): + """cudaMemcpyAsync cannot be issued from device code, so an in-kernel single-element copy assigns.""" + sdfg = dace.SDFG('single_elt_in_kernel_cross') + sdfg.add_array('H', [1], dace.float64, dace.dtypes.StorageType.CPU_Heap) + sdfg.add_array('G', [4], dace.float64, dace.dtypes.StorageType.GPU_Global, transient=True) + state = sdfg.add_state('s') + me, mx = state.add_map('kernel', dict(i='0:1'), schedule=dace.dtypes.ScheduleType.GPU_Device) + h = state.add_access('H') + g = state.add_access('G') + libnode = CopyLibraryNode(name='h_to_g') + state.add_node(libnode) + state.add_memlet_path(me, h, memlet=dace.Memlet()) + state.add_edge(h, None, libnode, CopyLibraryNode.INPUT_CONNECTOR_NAME, dace.Memlet('H[0]')) + state.add_edge(libnode, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, g, None, dace.Memlet('G[0]')) + state.add_memlet_path(g, mx, memlet=dace.Memlet('G[0]')) + assert select_copy_implementation(libnode, state) == 'Tasklet' + + +def test_host_scalar_endpoint_memcpy_takes_its_address(): + """A pointer connector bound to a scalar-defined variable must bind its ADDRESS, not its value.""" + sdfg = dace.SDFG('scalar_endpoint_memcpy') + sdfg.add_scalar('s', dace.float32) + sdfg.add_scalar('gs', dace.float32, dace.dtypes.StorageType.GPU_Global, transient=True) + sdfg.add_array('out', [1], dace.float32, dace.dtypes.StorageType.GPU_Global, transient=True) + state = sdfg.add_state('s0') + state.add_nedge(state.add_read('s'), state.add_write('gs'), dace.Memlet('s')) + state.add_nedge(state.add_read('gs'), state.add_write('out'), dace.Memlet('gs')) + code = _generated_code(sdfg) + assert 'MemcpyAsync(_cpy_out, _cpy_in' in code + assert '_cpy_in = &s;' in code, code + + +def test_opaque_handle_endpoint_is_not_addressed(): + """The other half of the scalar-endpoint rule: an opaque handle (``MPI_Comm`` and friends) IS + already the pointer its connector names, so it passes through by value instead of gaining an + ``&`` -- ``&handle`` is one indirection too many and does not compile.""" + handle = dace.dtypes.opaque('MPI_Comm') + sdfg = dace.SDFG('opaque_handle_endpoint') + sdfg.add_scalar('h', handle, transient=True) + sdfg.add_array('out', [1], dace.int32) + state = sdfg.add_state('s0') + tasklet = state.add_tasklet('use_handle', {'_h'}, {'_o'}, 'take_handle(_h);\n_o = 0;', language=dace.Language.CPP) + tasklet.in_connectors['_h'] = dace.dtypes.pointer(handle) + state.add_edge(state.add_access('h'), None, tasklet, '_h', dace.memlet.Memlet('h')) + state.add_edge(tasklet, '_o', state.add_write('out'), None, dace.memlet.Memlet('out[0]')) + + code = _generated_code(sdfg) + assert 'MPI_Comm _h = h;' in code, code + assert '&h' not in code, code + + +def test_host_tasklet_writing_gpu_memory_gets_the_stream_in_scope(): + """``__dace_current_stream`` is declared for a host tasklet that only WRITES GPU memory.""" + sdfg = dace.SDFG('h2d_stream_scope') + sdfg.add_array('A', [64], dace.float64) + sdfg.add_array('gA', [64], dace.float64, dace.dtypes.StorageType.GPU_Global, transient=True) + sdfg.add_array('B', [64], dace.float64, dace.dtypes.StorageType.GPU_Global, transient=True) + state = sdfg.add_state('s0') + state.add_nedge(state.add_read('A'), state.add_write('gA'), dace.Memlet('A[0:64]')) + state.add_nedge(state.add_read('gA'), state.add_write('B'), dace.Memlet('gA[0:64]')) + with dace.config.set_temporary('compiler', 'cuda', 'max_concurrent_streams', value=-1): + code = _generated_code(sdfg) + stream_decl = code.find('__dace_current_stream = ') + memcpy_call = code.find('MemcpyAsync(_cpy_out, _cpy_in') + assert stream_decl != -1 and memcpy_call != -1, code + assert stream_decl < memcpy_call + + +def test_in_kernel_copy_does_not_emit_a_grid_barrier(): + """A grid barrier releases only once EVERY thread reaches it. The copy expansion runs inside a + single-thread component of a persistent kernel, so a barrier at its own state boundary is + reached by one thread of one block and hangs the grid. Ordering comes from the enclosing + state's barrier, which sits outside that guard.""" + N = dace.symbol('N_gbar', dtype=dace.int64) + + @dace.program(auto_optimize=False, device=dace.dtypes.DeviceType.GPU) + def gbar_prog(A: dace.float64[N], B: dace.float64[N]): + a = 10.2 + for t in range(1, 10): + if t < N: + A[:] = (A + B + a) / 2 + a += 1 + + sdfg = gbar_prog.to_sdfg() + sdfg.apply_gpu_transformations() + content_nodes = set(sdfg.nodes()) - {sdfg.start_state, sdfg.sink_nodes()[0]} + transform = GPUPersistentKernel() + transform.setup_match(SubgraphView(sdfg, content_nodes)) + transform.kernel_prefix = 'stuff' + transform.apply(sdfg) + + cuda = next(obj.clean_code for obj in sdfg.generate_code() if obj.language == 'cu') + marker = 'DACE_DFI void copy_' + assert marker in cuda, cuda + start = cuda.index(marker) + end = cuda.index('DACE_DFI', start + len(marker)) + assert '__gbar.Sync();' not in cuda[start:end], cuda[start:end] + # The kernel's own state machine still gets its barriers. + assert '__gbar.Sync();' in cuda[end:] + + +def test_a_multi_state_nested_sdfg_below_the_kernel_keeps_its_barriers(): + """A nested SDFG with several states below the kernel map is a state machine, and its states + still need grid barriers between them -- the lone-state narrowing must not eat those.""" + inner = dace.SDFG('inner_states') + inner.add_array('X', [32], dace.float64, storage=dace.dtypes.StorageType.GPU_Global) + s1 = inner.add_state('one', is_start_block=True) + s2 = inner.add_state('two') + inner.add_edge(s1, s2, dace.InterstateEdge()) + t1 = s1.add_tasklet('w1', {}, {'o'}, 'o = 1.0') + s1.add_edge(t1, 'o', s1.add_write('X'), None, dace.Memlet('X[0]')) + t2 = s2.add_tasklet('w2', {}, {'o'}, 'o = 2.0') + s2.add_edge(t2, 'o', s2.add_write('X'), None, dace.Memlet('X[1]')) + + sdfg = dace.SDFG('persistent_nested_state_machine') + sdfg.add_array('X', [32], dace.float64, storage=dace.dtypes.StorageType.GPU_Global) + state = sdfg.add_state('launch') + entry, exit_ = state.add_map('kernel_launch_map', + dict(ignore='0'), + schedule=dace.dtypes.ScheduleType.GPU_Persistent) + nsdfg = state.add_nested_sdfg(inner, [], ['X']) + state.add_nedge(entry, nsdfg, dace.Memlet()) + state.add_edge_pair(exit_, + nsdfg, + state.add_write('X'), + internal_connector='X', + internal_memlet=dace.Memlet.from_array('X', sdfg.arrays['X'])) + sdfg.validate() + + cuda = next(obj.clean_code for obj in sdfg.generate_code() if obj.language == 'cu') + # One barrier per inner state: the transition between them is a sync point. + assert cuda.count('__gbar.Sync();') >= 2, cuda + + +def test_shared_to_global_uses_the_block_collective_helper(): + """Inside a kernel, a 1-D Shared -> Global copy splits across the thread block instead of + every thread copying the whole region (which races on shared memory written by other threads).""" + sdfg = dace.SDFG('shared_to_global_collective') + sdfg.add_array('G', [32], dace.float64, dace.dtypes.StorageType.GPU_Global) + sdfg.add_array('S', [32], dace.float64, dace.dtypes.StorageType.GPU_Shared, transient=True) + state = sdfg.add_state('s') + me, mx = state.add_map('kernel', dict(i='0:1'), schedule=dace.dtypes.ScheduleType.GPU_Device) + s_acc = state.add_access('S') + g_acc = state.add_access('G') + libnode = CopyLibraryNode(name='s_to_g') + state.add_node(libnode) + state.add_memlet_path(me, s_acc, memlet=dace.Memlet()) + state.add_edge(s_acc, None, libnode, CopyLibraryNode.INPUT_CONNECTOR_NAME, dace.Memlet('S[0:32]')) + state.add_edge(libnode, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, g_acc, None, dace.Memlet('G[0:32]')) + state.add_memlet_path(g_acc, mx, memlet=dace.Memlet('G[0:32]')) + libnode.implementation = 'SharedMemoryCollective' + sdfg.expand_library_nodes() + bodies = [n.code.as_string for n, _ in sdfg.all_nodes_recursive() if isinstance(n, dace.nodes.Tasklet)] + assert any('dace::SharedToGlobal1D<' in b for b in bodies), bodies + + +def test_symbolic_extent_expansions_keep_their_ranges_symbolic(): + """A copy / fill whose extent is a symbolic POWER must expand. + + ``sym2cpp(R ** K)`` is ``dace::math::ipow(R, K)``; rendering an extent through it and handing the + text back to the range parser splits the qualified name on ':' and raises + ``SyntaxError: Invalid range``. Both expansions must therefore build their map ranges and memlet + subsets from the symbolic expression, never from a rendered string. Reproduces the npbench + ``stockham_fft`` expansion failure. + """ + R, K = dace.symbol('R'), dace.symbol('K') + extent = R**K + + sdfg = dace.SDFG('symbolic_extent_copy') + sdfg.add_array('src', [extent], dace.float64) + sdfg.add_array('dst', [extent], dace.float64) + state = sdfg.add_state('main') + libnode = CopyLibraryNode('cpy') + libnode.implementation = 'MappedTasklet' + state.add_node(libnode) + state.add_edge(state.add_access('src'), None, libnode, CopyLibraryNode.INPUT_CONNECTOR_NAME, + dace.Memlet(f'src[0:{extent}]')) + state.add_edge(libnode, CopyLibraryNode.OUTPUT_CONNECTOR_NAME, state.add_access('dst'), None, + dace.Memlet(f'dst[0:{extent}]')) + sdfg.expand_library_nodes() + sdfg.validate() + + fill_sdfg = dace.SDFG('symbolic_extent_fill') + fill_sdfg.add_array('out', [extent], dace.float64) + fill_state = fill_sdfg.add_state('main') + fill = FillLibraryNode('zero') + fill.implementation = 'pure' + fill_state.add_node(fill) + fill_state.add_edge(fill, FillLibraryNode.OUTPUT_CONNECTOR_NAME, fill_state.add_access('out'), None, + dace.Memlet(f'out[0:{extent}]')) + fill_sdfg.expand_library_nodes() + fill_sdfg.validate() + + for expanded in (sdfg, fill_sdfg): + entries = [n for n, _ in expanded.all_nodes_recursive() if isinstance(n, dace.nodes.MapEntry)] + assert entries, f'{expanded.name}: no map emitted' + for entry in entries: + for _, end, _ in entry.map.range: + # The extent survives as a symbolic expression over R and K (possibly wrapped in a + # ceiling by the element count), not as a C++ rendering of it. + assert '::' not in str(end), f'{expanded.name}: C++ spelling leaked into a map range: {end}' + assert {str(s) for s in symbolic.pystr_to_symbolic(str(end)).free_symbols} == {'R', 'K'}, \ + f'{expanded.name}: extent lost its symbols: {end}' + for st in expanded.all_states(): + for e in st.edges(): + if e.data is not None and not e.data.is_empty() and e.data.subset is not None: + assert '::' not in str(e.data.subset), \ + f'{expanded.name}: C++ spelling leaked into a memlet subset: {e.data.subset}' + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/tests/library/fill_node_test.py b/tests/library/fill_node_test.py new file mode 100644 index 0000000000..8dce0d49ac --- /dev/null +++ b/tests/library/fill_node_test.py @@ -0,0 +1,441 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Tests for :class:`FillLibraryNode` and its pure / CPU / CUDA / tasklet expansions.""" +import contextlib +from typing import Optional, Sequence + +import dace +from dace.libraries.standard.nodes.fill import FillLibraryNode, byte_pattern, select_fill_implementation + +import pytest +import numpy as np + + +def make_fill_sdfg(implementation: Optional[str], + shape: Sequence[int], + subset: str, + gpu: bool = True, + name: str = "fill_sdfg", + dtype: dace.dtypes.typeclass = dace.dtypes.float64, + value=0) -> dace.SDFG: + """Build an SDFG that fills a sub-region of a single array. + + :param implementation: ``FillLibraryNode.implementation`` (``None`` keeps ``'Auto'``). + :param shape: array shape (sequence of dim extents). + :param subset: memlet subset string for the fill's output edge. + :param gpu: True for ``GPU_Global`` storage, False for ``CPU_Heap``. + :param name: SDFG name. + :param dtype: element type of the filled array. + :param value: the constant the node writes. + :returns: the constructed SDFG. + """ + sdfg = dace.SDFG(name) + arr_name = "gpuB" if gpu else "B" + storage = dace.dtypes.StorageType.GPU_Global if gpu else dace.dtypes.StorageType.CPU_Heap + sdfg.add_array(name=arr_name, shape=list(shape), dtype=dtype, storage=storage, transient=False) + + state = sdfg.add_state("main") + out = state.add_access(arr_name) + libnode = FillLibraryNode(name="fill_libnode", value=value) + if implementation is not None: + libnode.implementation = implementation + state.add_edge(libnode, FillLibraryNode.OUTPUT_CONNECTOR_NAME, out, None, + dace.memlet.Memlet(f"{arr_name}[{subset}]")) + return sdfg + + +def _get_sdfg(implementation: Optional[str], gpu: bool = True) -> dace.SDFG: + """1-D slice fill.""" + return make_fill_sdfg(implementation, (200, ), "50:100", gpu=gpu, name="fill_sdfg") + + +def _get_multi_dim_sdfg(implementation: Optional[str], gpu: bool = True) -> dace.SDFG: + """3-D sub-block fill.""" + return make_fill_sdfg(implementation, (50, 2, 2), "40:50, 0:2, 0:2", gpu=gpu, name="fill_sdfg2") + + +def test_fill_pure_1d_cpu(): + """``pure`` zeros the 1D CPU slice, leaving the rest unchanged.""" + sdfg = _get_sdfg("pure", gpu=False) + sdfg.name += "_pure_cpu" + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = sdfg.compile() + + B = np.ones((200, ), dtype=np.float64) + exe(B=B) + + assert np.all(B[:50] == 1) + assert np.all(B[100:] == 1) + assert np.all(B[50:100] == 0) + + +def test_fill_pure_3d_cpu(): + """``pure`` zeros the 3D CPU sub-block, leaving the rest unchanged.""" + sdfg = _get_multi_dim_sdfg("pure", gpu=False) + sdfg.name += "_pure_cpu_multi_dim" + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = sdfg.compile() + + B = np.ones((50, 2, 2), dtype=np.float64) + exe(B=B) + + assert np.all(B[0:40, :, :] == 1) + assert np.all(B[40:50, :, :] == 0) + + +@pytest.mark.gpu +def test_fill_pure_1d_gpu(): + """``pure`` zeros the 1D GPU slice, leaving the rest unchanged.""" + import cupy as cp + + sdfg = _get_sdfg("pure", gpu=True) + sdfg.name += "_pure_gpu" + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = sdfg.compile() + + B = cp.ones((200, ), dtype=cp.float64) + exe(gpuB=B) + + assert cp.all(B[:50] == 1) + assert cp.all(B[100:] == 1) + assert cp.all(B[50:100] == 0) + + +@pytest.mark.gpu +def test_fill_pure_3d_gpu(): + """``pure`` zeros the 3D GPU sub-block, leaving the rest unchanged.""" + import cupy as cp + + sdfg = _get_multi_dim_sdfg("pure", gpu=True) + sdfg.name += "_pure_gpu_multi_dim" + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = sdfg.compile() + + B = cp.ones((50, 2, 2), dtype=np.float64) + exe(gpuB=B) + + assert cp.all(B[0:40, :, :] == 1) + assert cp.all(B[40:50, :, :] == 0) + + +@pytest.mark.gpu +def test_fill_cuda_1d_gpu(): + """``CUDA`` zeros the 1D GPU slice, leaving the rest unchanged.""" + import cupy as cp + + sdfg = _get_sdfg("CUDA", gpu=True) + sdfg.name += "_cuda_gpu" + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = sdfg.compile() + + B = cp.ones((200, ), dtype=cp.float64) + exe(gpuB=B) + + assert cp.all(B[:50] == 1) + assert cp.all(B[100:] == 1) + assert cp.all(B[50:100] == 0) + + +@pytest.mark.gpu +def test_fill_cuda_3d_gpu(): + """``CUDA`` zeros the 3D GPU sub-block, leaving the rest unchanged.""" + import cupy as cp + + sdfg = _get_multi_dim_sdfg("CUDA", gpu=True) + sdfg.name += "_cuda_gpu_multi_dim" + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = sdfg.compile() + + B = cp.ones((50, 2, 2), dtype=np.float64) + exe(gpuB=B) + + assert cp.all(B[0:40, :, :] == 1) + assert cp.all(B[40:50, :, :] == 0) + + +@pytest.mark.gpu +def test_fill_cuda_rejects_cpu_storage(): + """``CUDA`` targeting a CPU array is rejected.""" + sdfg = _get_sdfg("CUDA", gpu=False) + sdfg.name += "_cuda_cpu" + sdfg.validate() + sdfg.expand_library_nodes() + with pytest.raises(Exception): + sdfg.validate() + sdfg.compile() + + +def test_fill_auto_routes_non_contiguous_to_pure_cpu(): + """Auto routes a non-contiguous CPU subset to ``pure`` (one call would write outside the region).""" + sdfg = make_fill_sdfg(None, (10, 20), "2:8, 5:15", gpu=False, name="fill_noncontig_cpu_auto") + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() + exe = sdfg.compile() + + B = np.ones((10, 20), dtype=np.float64) + exe(B=B) + # The 6x10 sub-block is zeroed; everything else stays 1. + expected = np.ones((10, 20), dtype=np.float64) + for i in range(2, 8): + for j in range(5, 15): + expected[i, j] = 0 + np.testing.assert_array_equal(B, expected) + + +def test_fill_cpu_rejects_non_contiguous_subset(): + """Explicit ``CPU`` expansion rejects a non-contiguous subset (one call would overrun the region).""" + sdfg = make_fill_sdfg("CPU", (10, 20), "2:8, 5:15", gpu=False, name="fill_noncontig_cpu_explicit") + sdfg.validate() + with pytest.raises(ValueError, match="contiguous"): + sdfg.expand_library_nodes() + + +@pytest.mark.gpu +def test_fill_cuda_rejects_non_contiguous_subset(): + """Explicit ``CUDA`` expansion rejects a non-contiguous subset (one ``cudaMemsetAsync`` would overrun).""" + sdfg = make_fill_sdfg("CUDA", (10, 20), "2:8, 5:15", gpu=True, name="fill_noncontig_cuda_explicit") + sdfg.validate() + with pytest.raises(ValueError, match="contiguous"): + sdfg.expand_library_nodes() + + +def test_fill_register_outside_kernel_routes_to_cpu_tasklet(): + """A Fill on a Register outside a GPU kernel scope lowers to a direct host-side Tasklet.""" + sdfg = dace.SDFG('fill_reg_outside_kernel') + sdfg.add_array('R', [1], dace.float64, dace.StorageType.Register, transient=True) + state = sdfg.add_state('s') + + r = state.add_access('R') + fill_node = FillLibraryNode(name='fill_r') + state.add_node(fill_node) + state.add_edge(fill_node, FillLibraryNode.OUTPUT_CONNECTOR_NAME, r, None, dace.Memlet('R[0]')) + + sdfg.expand_library_nodes() + + # Verify no complex structures or CUDA launch strings are generated on the host for raw registers + nsdfg_count = sum(1 for n, _ in sdfg.all_nodes_recursive() if isinstance(n, dace.nodes.NestedSDFG)) + assert nsdfg_count == 0, "Host register fill should expand to a direct Tasklet, not a NestedSDFG." + + assignments = [ + n for n, _ in sdfg.all_nodes_recursive() if isinstance(n, dace.nodes.Tasklet) and '= 0' in n.code.as_string + ] + assert assignments, "Expected a basic literal assignment tasklet on the host." + + +def test_fill_register_inside_kernel_routes_to_sequential(): + """A multi-element Fill targeting a Register array inside a GPU kernel maps to sequential in-kernel logic.""" + sdfg = dace.SDFG('fill_reg_inside_kernel') + sdfg.add_array('R', [4], dace.float64, dace.StorageType.Register, transient=True) + state = sdfg.add_state('s') + + me, mx = state.add_map('kernel', dict(i='0:1'), schedule=dace.dtypes.ScheduleType.GPU_Device) + r = state.add_access('R') + fill_node = FillLibraryNode(name='fill_r') + state.add_node(fill_node) + + state.add_memlet_path(me, fill_node, memlet=dace.Memlet()) + state.add_edge(fill_node, FillLibraryNode.OUTPUT_CONNECTOR_NAME, r, None, dace.Memlet('R[0:4]')) + state.add_memlet_path(r, mx, memlet=dace.Memlet()) + + sdfg.expand_library_nodes() + + # Ensure it did not lower to a host-side or invalid device-side memset call. The expansion + # names the API after the configured backend, so match both spellings, not just cuda's. + memsets = [ + n for n, _ in sdfg.all_nodes_recursive() + if isinstance(n, dace.nodes.Tasklet) and ('cudaMemset' in n.code.as_string or 'hipMemset' in n.code.as_string) + ] + assert len(memsets) == 0, "Cannot issue a device memset on local GPU registers." + + # It should fall back to an internal loop/unrolled tasklet chain inside the device state + assert any(isinstance(n, dace.nodes.Tasklet) for n, _ in sdfg.all_nodes_recursive()) + + +def test_fill_single_gpu_shared_inside_kernel_expands_clean(): + """A single-element fill targeting GPU-resident storage *inside* a GPU kernel is valid device + code (a device-side ``_out = 0``) and must expand cleanly. Regression: the ``tasklet`` guard fired + on exactly this valid case, and its error path dereferenced the output *name* (a ``str``) as + ``inp.storage`` -> ``AttributeError``.""" + sdfg = dace.SDFG('fill_shared_inside_kernel') + sdfg.add_array('s', [1], dace.float64, dace.StorageType.GPU_Shared, transient=True) + state = sdfg.add_state('s') + + me, mx = state.add_map('kernel', dict(i='0:1'), schedule=dace.dtypes.ScheduleType.GPU_Device) + s_acc = state.add_access('s') + fill_node = FillLibraryNode(name='fill_s') + state.add_node(fill_node) + state.add_memlet_path(me, fill_node, memlet=dace.Memlet()) + state.add_edge(fill_node, FillLibraryNode.OUTPUT_CONNECTOR_NAME, s_acc, None, dace.Memlet('s[0]')) + state.add_memlet_path(s_acc, mx, memlet=dace.Memlet()) + + sdfg.expand_library_nodes() # must not raise + + assert any(isinstance(n, dace.nodes.Tasklet) and '= 0' in n.code.as_string + for n, _ in sdfg.all_nodes_recursive()), "Expected a scalar zero-assignment tasklet." + + +def test_fill_tasklet_rejects_gpu_storage_from_host_scope(): + """The single-element ``tasklet`` expansion emits ``_out = 0`` in its own scope; from host scope it + cannot target GPU-resident storage (a scalar assignment cannot write device memory), so it must + raise a clean ``ValueError``. Regression: the guard tested the wrong side, letting this host->GPU + case through instead of rejecting it.""" + sdfg = make_fill_sdfg("tasklet", (1, ), "0:1", gpu=True, name="fill_tasklet_host_gpu") + sdfg.validate() + with pytest.raises(ValueError): + sdfg.expand_library_nodes() + + +def test_fill_pure_strided_map_matches_array(): + """A ``pure`` fill over a strided subset must give the mapped tasklet the same collapsed + extent as the wrapper array descriptor. Regression: ``map_lengths`` was recomputed from + ``out_subset.size()`` instead of the collapsed shape used for the array, so the map bounds + could diverge from the array rank/extent.""" + sdfg = make_fill_sdfg(None, (9, ), "0:9:3", gpu=False, name="fill_strided_cpu") + sdfg.validate() + sdfg.expand_library_nodes() + sdfg.validate() # a diverged map/array would fail validation here + exe = sdfg.compile() + + B = np.ones((9, ), dtype=np.float64) + exe(B=B) + + expected = np.ones((9, ), dtype=np.float64) + expected[0:9:3] = 0 # indices 0, 3, 6 + np.testing.assert_array_equal(B, expected) + + +@contextlib.contextmanager +def _pinned_transfer_threshold(value): + """Pin ``compiler.cpu.parallel_transfer_min_elements`` so Auto selection is deterministic.""" + orig = dace.config.Config.get("compiler", "cpu", "parallel_transfer_min_elements") + dace.config.Config.set("compiler", "cpu", "parallel_transfer_min_elements", value=value) + try: + yield + finally: + dace.config.Config.set("compiler", "cpu", "parallel_transfer_min_elements", value=orig) + + +def cpu_fill_sdfg(extent, name): + """Single-state CPU_Heap ``FillLibraryNode`` zeroing ``0:extent``.""" + sdfg = dace.SDFG(name) + sdfg.add_array("dst", [extent], dace.float64, dace.dtypes.StorageType.CPU_Heap) + state = sdfg.add_state("s") + libnode = FillLibraryNode(name="ms") + state.add_edge(libnode, FillLibraryNode.OUTPUT_CONNECTOR_NAME, state.add_access("dst"), None, + dace.Memlet(f"dst[0:{extent}]")) + sdfg.validate() + return sdfg, libnode + + +def _generated_code(sdfg): + return "\n".join(obj.code for obj in sdfg.generate_code()) + + +def test_fill_below_threshold_emits_a_single_call(): + """A constant-size CPU fill below the threshold lowers to one ``std::fill_n``, not an OpenMP loop. + gcc and clang reduce that call to a memset at the Release level dace builds with when the + value's object representation allows it.""" + with _pinned_transfer_threshold(1024): + sdfg, libnode = cpu_fill_sdfg(100, "fill_below_threshold") + sdfg.expand_library_nodes(recursive=True) + assert libnode.implementation == 'CPU' + code = _generated_code(sdfg) + assert 'std::fill_n' in code + assert '#pragma omp parallel for' not in code + + +def test_fill_at_threshold_emits_omp_parallel_for(): + """A constant-size CPU fill at/above the threshold lowers to an OpenMP element map, not one call.""" + with _pinned_transfer_threshold(1024): + sdfg, libnode = cpu_fill_sdfg(4096, "fill_at_threshold") + sdfg.expand_library_nodes(recursive=True) + assert libnode.implementation == 'pure' + code = _generated_code(sdfg) + assert '#pragma omp parallel for' in code + assert 'std::fill_n' not in code + + +def test_fill_symbolic_size_emits_omp_parallel_for(): + """A symbolic (compile-time-unknown) CPU fill size is assumed large, so it takes the same + OpenMP-parallel path as a large constant, never the single ``std::fill_n``.""" + with _pinned_transfer_threshold(1024): + n = dace.symbol('N_fill_symbolic') + sdfg, libnode = cpu_fill_sdfg(n, "fill_symbolic_size") + sdfg.expand_library_nodes(recursive=True) + assert libnode.implementation == 'pure' + code = _generated_code(sdfg) + assert '#pragma omp parallel for' in code + assert 'std::fill_n' not in code + + +if __name__ == "__main__": + pytest.main([__file__]) + +NARROW_FLOATS = [ + pytest.param(dace.float16, 'float16'), + pytest.param(dace.bfloat16, 'bfloat16'), + pytest.param(dace.float8_e4m3fn, 'float8_e4m3fn'), + pytest.param(dace.float8_e5m2, 'float8_e5m2'), +] + + +@pytest.mark.parametrize('dtype,label', NARROW_FLOATS) +@pytest.mark.parametrize('value', [0.0, 1.0, -1.0, 0.5, 2.0]) +def test_fill_narrow_float_writes_the_value(dtype, label, value): + """A reduced-precision fill must write the value the destination type can hold, whichever + lowering the object representation selects. One byte (fp8) makes every value byte-splat, so + those always memset; two bytes (fp16, bfloat16) only do for zero.""" + n = 64 + sdfg = make_fill_sdfg("CPU", (n, ), + f"0:{n}", + gpu=False, + name=f"fill_{label}_{str(value).replace('.', '_').replace('-', 'neg')}", + dtype=dtype, + value=value) + npdt = dtype.as_numpy_dtype() + buf = np.full(n, 7, dtype=npdt) + sdfg(B=buf) + assert np.array_equal(buf, np.full(n, value, dtype=npdt)), f"{label} fill of {value}" + + +@pytest.mark.parametrize('dtype,label', NARROW_FLOATS) +@pytest.mark.parametrize('value', [0.0, 1.0]) +def test_fill_narrow_float_host_lowering_is_one_call(dtype, label, value): + """The host fill is a single ``std::fill_n`` whatever the value's object representation is: + gcc and clang both reduce it to a memset when the representation allows, and dace always + builds Release. What must hold here is that the fill stays one call and never degrades to an + element map.""" + n = 64 + tag = str(value).replace('.', '_') + sdfg = make_fill_sdfg("CPU", (n, ), f"0:{n}", gpu=False, dtype=dtype, value=value, name=f"lowering_{label}_{tag}") + code = sdfg.generate_code()[0].clean_code + assert 'std::fill_n' in code + assert '#pragma omp parallel for' not in code + + +@pytest.mark.parametrize('dtype,label', NARROW_FLOATS) +def test_fill_narrow_float_gpu_routing_follows_the_byte_pattern(dtype, label): + """``cudaMemsetAsync`` writes one byte, and no optimizer can widen it, so the GPU choice must + still be made from the object representation. A one-byte type (fp8) can memset any value; a + two-byte type (fp16, bfloat16) only zero, and everything else has to reach the kernel.""" + for value in (0.0, 1.0): + pattern = byte_pattern(value, dtype) + expected_memset = (dtype.bytes == 1) or value == 0.0 + assert (pattern is not None) == expected_memset, f"{label} {value}" + + sdfg = make_fill_sdfg(None, (64, ), "0:64", gpu=True, dtype=dtype, value=1.0, name=f"gpu_route_{label}") + node = next(n for n in sdfg.start_state.nodes() if isinstance(n, FillLibraryNode)) + chosen = select_fill_implementation(node, sdfg.start_state) + assert chosen == ('CUDA' if dtype.bytes == 1 else 'pure'), f"{label} routed to {chosen}" diff --git a/tests/passes/insert_explicit_copies_test.py b/tests/passes/insert_explicit_copies_test.py new file mode 100644 index 0000000000..4322a83a04 --- /dev/null +++ b/tests/passes/insert_explicit_copies_test.py @@ -0,0 +1,1218 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Tests for the ``InsertExplicitCopies`` pass.""" +import copy as _copy +import importlib.util +import os +import sys + +import dace +import numpy as np +import pytest +from dace import nodes +from dace.memlet import Memlet +from dace.sdfg import utils as sdutils +from dace.libraries.standard.nodes.copy import CopyLibraryNode +from dace.transformation.passes.insert_explicit_copies import InsertExplicitCopies + +import tests.polybench + +# The polybench programs import their ``polybench`` harness as a top-level module, which only +# resolves when run as scripts (own directory on sys.path); importing them as a package needs it too. +sys.path.append(os.path.dirname(tests.polybench.__file__)) + +from tests.polybench.correlation import correlation, init_array as _correlation_init_array +from tests.polybench.covariance import covariance, init_array as _covariance_init_array + +# fdtd-2d.py's hyphenated filename is not a valid module identifier. Load it from +# its path under a clean module name so the SDFG name (derived from the module +# path) is valid -- without importing or mutating the canonical hyphenated module. +_fdtd2d_path = os.path.join(os.path.dirname(tests.polybench.__file__), "fdtd-2d.py") +_fdtd2d_spec = importlib.util.spec_from_file_location("polybench_fdtd_2d", _fdtd2d_path) +_fdtd2d_module = importlib.util.module_from_spec(_fdtd2d_spec) +_fdtd2d_spec.loader.exec_module(_fdtd2d_module) +fdtd2d = _fdtd2d_module.fdtd2d +_fdtd2d_init_array = _fdtd2d_module.init_array + + +def _wcr_edges(sdfg): + """``(state, edge)`` for every edge still carrying a WCR.""" + return [(st, e) for st in sdfg.all_states() for e in st.edges() if e.data.wcr is not None] + + +def _count_copy_nodes(sdfg): + """Count CopyLibraryNode instances across all states (recursive).""" + return sum(1 for n, _ in sdfg.all_nodes_recursive() if isinstance(n, CopyLibraryNode)) + + +def _count_direct_copy_edges(sdfg): + """Count AccessNode -> AccessNode non-empty edges (recursive).""" + count = 0 + for nsdfg in sdfg.all_sdfgs_recursive(): + for state in nsdfg.states(): + for e in state.edges(): + if (isinstance(e.src, nodes.AccessNode) and isinstance(e.dst, nodes.AccessNode) + and not e.data.is_empty()): + count += 1 + return count + + +def _assert_no_other_subset(sdfg: dace.SDFG) -> None: + """Assert no data-movement memlet still carries an ``other_subset`` after copy-node insertion. + + View-defining (alias) edges are excluded: they reference the underlying buffer rather than + moving data, so ``InsertExplicitCopies`` correctly leaves them direct with ``other_subset`` + intact (mirrors the pass's own view-edge skip); a copy would change the SDFG's semantics. + """ + for nsdfg in sdfg.all_sdfgs_recursive(): + for state in nsdfg.states(): + for edge in state.edges(): + memlet = edge.data + if memlet.is_empty(): + continue + if any( + isinstance(an, nodes.AccessNode) and isinstance(nsdfg.arrays[an.data], dace.data.View) + and sdutils.get_view_edge(state, an) is edge for an in (edge.src, edge.dst)): + continue + assert memlet.other_subset is None, ( + f"Memlet on edge {edge.src}->{edge.dst} in SDFG '{nsdfg.name}' still " + f"has other_subset={memlet.other_subset}; expected None after copy insertion.") + + +def _assert_no_copynd(sdfg: dace.SDFG) -> None: + """Assert ``generate_code`` emits no ``dace::CopyND`` template instantiations.""" + sdfg.expand_library_nodes() + for obj in sdfg.generate_code(): + code = obj.code if isinstance(obj.code, str) else getattr(obj.code, 'code', str(obj.code)) + assert 'CopyND<' not in code, f"unexpected CopyND in code object {obj.title}" + + +def _build_copy_sdfg(name, arrays, edge_memlet): + """Build an SDFG with two AccessNodes wired by a single edge.""" + sdfg = dace.SDFG(name) + for arr_name, shape, storage in arrays: + sdfg.add_array(arr_name, shape, dace.float64, storage) + st = sdfg.add_state("s") + src = st.add_access(arrays[0][0]) + dst = st.add_access(arrays[1][0]) + st.add_edge(src, None, dst, None, edge_memlet) + return sdfg, st, src, dst + + +def _assert_copy_storages(sdfg, src_storage, dst_storage): + """Assert that every CopyLibraryNode in ``sdfg`` has the given storages.""" + found = False + for n, parent in sdfg.all_nodes_recursive(): + if isinstance(n, CopyLibraryNode): + assert n.src_storage(parent) == src_storage + assert n.dst_storage(parent) == dst_storage + found = True + assert found, "No CopyLibraryNode found in SDFG" + + +def _compile_and_run(sdfg, inputs): + sdfg.expand_library_nodes() + exe = sdfg.compile() + exe(**inputs) + + +def test_insert_cpu_to_cpu_1d(): + """CPU_Heap -> CPU_Heap 1D copy.""" + cpu = dace.StorageType.CPU_Heap + sdfg, _, _, _ = _build_copy_sdfg("insert_cpu_cpu_1d", [("A", [100], cpu), ("B", [100], cpu)], + Memlet("A[10:60]", other_subset="20:70")) + + assert _count_direct_copy_edges(sdfg) == 1 + InsertExplicitCopies().apply_pass(sdfg, {}) + _assert_no_other_subset(sdfg) + assert _count_direct_copy_edges(sdfg) == 0 + assert _count_copy_nodes(sdfg) == 1 + _assert_copy_storages(sdfg, cpu, cpu) + + A = np.arange(100, dtype=np.float64) + B = np.zeros(100, dtype=np.float64) + _compile_and_run(sdfg, dict(A=A, B=B)) + np.testing.assert_array_equal(B[20:70], A[10:60]) + assert np.all(B[:20] == 0) and np.all(B[70:] == 0) + + +def test_insert_cpu_to_cpu_2d_slice(): + """CPU 2D slice copy with explicit other_subset.""" + cpu = dace.StorageType.CPU_Heap + sdfg, _, _, _ = _build_copy_sdfg("insert_cpu_2d", [("A", [10, 20], cpu), ("B", [10, 20], cpu)], + Memlet(data="A", subset="2:8, 5:15", other_subset="0:6, 0:10")) + + InsertExplicitCopies().apply_pass(sdfg, {}) + _assert_no_other_subset(sdfg) + assert _count_direct_copy_edges(sdfg) == 0 + assert _count_copy_nodes(sdfg) == 1 + + A = np.arange(200, dtype=np.float64).reshape(10, 20).copy() + B = np.zeros((10, 20), dtype=np.float64) + _compile_and_run(sdfg, dict(A=A, B=B)) + np.testing.assert_array_equal(B[0:6, 0:10], A[2:8, 5:15]) + + +@pytest.mark.parametrize("sdfg_name,memlet", [ + ("insert_other_dst", Memlet(data="B", subset="0:8", other_subset="2:10")), + ("insert_other_src", Memlet(data="A", subset="2:10", other_subset="0:8")), +], + ids=["data_is_dst", "data_is_src"]) +def test_insert_other_subset_data_convention(sdfg_name, memlet): + """Both memlet conventions yield the same copy ``_in=A[2:10]``, ``_out=B[0:8]`` with no ``other_subset``.""" + cpu = dace.StorageType.CPU_Heap + sdfg, st, _, _ = _build_copy_sdfg(sdfg_name, [("A", [20], cpu), ("B", [20], cpu)], memlet) + + InsertExplicitCopies().apply_pass(sdfg, {}) + _assert_no_other_subset(sdfg) + assert _count_copy_nodes(sdfg) == 1 + + for n in st.nodes(): + if isinstance(n, CopyLibraryNode): + in_m = list(st.in_edges(n))[0].data + out_m = list(st.out_edges(n))[0].data + assert in_m.data == "A" and str(in_m.subset) == "2:10" + assert in_m.other_subset is None + assert out_m.data == "B" and str(out_m.subset) == "0:8" + assert out_m.other_subset is None + break + + A = np.arange(20, dtype=np.float64) + B = np.full(20, -1.0, dtype=np.float64) + _compile_and_run(sdfg, dict(A=A, B=B)) + np.testing.assert_array_equal(B[0:8], A[2:10]) + assert np.all(B[8:] == -1.0) + + +def test_insert_cpu_to_cpu_full_array(): + """Full array copy.""" + cpu = dace.StorageType.CPU_Heap + sdfg, _, _, _ = _build_copy_sdfg("insert_full", [("A", [64], cpu), ("B", [64], cpu)], Memlet("A[0:64]")) + + InsertExplicitCopies().apply_pass(sdfg, {}) + _assert_no_other_subset(sdfg) + A = np.arange(64, dtype=np.float64) + B = np.zeros(64, dtype=np.float64) + _compile_and_run(sdfg, dict(A=A, B=B)) + np.testing.assert_array_equal(B, A) + + +def test_insert_multiple_copies_same_state(): + """Two copies in the same state: A->B and A->C.""" + sdfg = dace.SDFG("insert_multi") + for name in ("A", "B", "C"): + sdfg.add_array(name, [32], dace.float64, dace.StorageType.CPU_Heap) + st = sdfg.add_state("s") + a = st.add_access("A") + b = st.add_access("B") + c = st.add_access("C") + st.add_edge(a, None, b, None, Memlet("A[0:32]")) + st.add_edge(a, None, c, None, Memlet("A[0:32]")) + + result = InsertExplicitCopies().apply_pass(sdfg, {}) + _assert_no_other_subset(sdfg) + assert result == 2 + assert _count_copy_nodes(sdfg) == 2 + + A = np.arange(32, dtype=np.float64) + B = np.zeros(32, dtype=np.float64) + C = np.zeros(32, dtype=np.float64) + _compile_and_run(sdfg, dict(A=A, B=B, C=C)) + np.testing.assert_array_equal(B, A) + np.testing.assert_array_equal(C, A) + + +def test_insert_empty_memlet_skipped(): + """Empty memlets (control edges) are not replaced.""" + cpu = dace.StorageType.CPU_Heap + sdfg, _, _, _ = _build_copy_sdfg("insert_empty", [("A", [10], cpu), ("B", [10], cpu)], Memlet()) + + result = InsertExplicitCopies().apply_pass(sdfg, {}) + _assert_no_other_subset(sdfg) + assert result is None + assert _count_copy_nodes(sdfg) == 0 + + +def test_insert_no_copies_returns_none(): + """If there are no copy edges, return None.""" + sdfg = dace.SDFG("no_copies") + sdfg.add_array("A", [10], dace.float64, dace.StorageType.CPU_Heap) + st = sdfg.add_state("s") + a = st.add_access("A") + t = st.add_tasklet("noop", {"_in"}, {"_out"}, "_out = _in + 1") + a2 = st.add_access("A") + st.add_edge(a, None, t, "_in", Memlet("A[0]")) + st.add_edge(t, "_out", a2, None, Memlet("A[0]")) + + result = InsertExplicitCopies().apply_pass(sdfg, {}) + _assert_no_other_subset(sdfg) + assert result is None + + +def test_insert_nested_sdfg(): + """Copy inside a nested SDFG is also replaced.""" + inner = dace.SDFG("inner") + inner.add_array("X", [20], dace.float64, dace.StorageType.CPU_Heap) + inner.add_array("Y", [20], dace.float64, dace.StorageType.CPU_Heap) + ist = inner.add_state("is") + x = ist.add_access("X") + y = ist.add_access("Y") + ist.add_edge(x, None, y, None, Memlet("X[0:20]")) + + outer = dace.SDFG("outer") + outer.add_array("A", [20], dace.float64, dace.StorageType.CPU_Heap) + outer.add_array("B", [20], dace.float64, dace.StorageType.CPU_Heap) + ost = outer.add_state("os") + nsdfg = ost.add_nested_sdfg(inner, {"X"}, {"Y"}) + a = ost.add_access("A") + b = ost.add_access("B") + ost.add_edge(a, None, nsdfg, "X", Memlet("A[0:20]")) + ost.add_edge(nsdfg, "Y", b, None, Memlet("B[0:20]")) + + result = InsertExplicitCopies().apply_pass(outer, {}) + _assert_no_other_subset(outer) + assert result == 1 + assert _count_copy_nodes(outer) == 1 + + +def _count_nested_sdfgs(sdfg): + """Count NestedSDFGs in ``sdfg`` (top level only -- not recursive into them).""" + return sum(1 for n, _ in sdfg.all_nodes_recursive() if isinstance(n, nodes.NestedSDFG)) + + +def test_single_element_copies_expand_to_tasklets_no_nested_sdfg(): + """Single-element copies expand to direct ``_cpy_out = _cpy_in`` Tasklets, never a NestedSDFG. + + The ``MappedTasklet`` path would build a 0-D map for these and crash + propagation, so routing must short-circuit to the ``Tasklet`` impl. + """ + cpu = dace.StorageType.CPU_Heap + pinned = dace.StorageType.CPU_Pinned + register = dace.StorageType.Register + gpu = dace.StorageType.GPU_Global + + sdfg = dace.SDFG("scalar_copies") + sdfg.add_array("c_in", [1], dace.float64, cpu) + sdfg.add_array("c_out", [1], dace.float64, pinned) + sdfg.add_array("r_in", [1], dace.float64, register, transient=True) + sdfg.add_array("r_out", [1], dace.float64, register, transient=True) + + st = sdfg.add_state("s") + c_in = st.add_access("c_in") + c_out = st.add_access("c_out") + r_in = st.add_access("r_in") + r_out = st.add_access("r_out") + st.add_edge(c_in, None, c_out, None, Memlet("c_in[0]")) + st.add_edge(r_in, None, r_out, None, Memlet("r_in[0]")) + + InsertExplicitCopies().apply_pass(sdfg, {}) + assert _count_copy_nodes(sdfg) == 2 + + sdfg.expand_library_nodes() + + assert _count_nested_sdfgs(sdfg) == 0, ( + "Single-element copies should expand to a direct Tasklet, not a NestedSDFG. " + f"Found {_count_nested_sdfgs(sdfg)} NestedSDFG(s) after expansion.") + + tasklets = [n for n, _ in sdfg.all_nodes_recursive() if isinstance(n, nodes.Tasklet)] + assert any( + "_cpy_out = _cpy_in" in t.code.as_string + for t in tasklets), (f"Expected at least one ``_cpy_out = _cpy_in`` Tasklet from CopyLibraryNode expansion; " + f"got tasklets with code: {[t.code.as_string for t in tasklets]}") + + +def test_insert_validates_after_pass(): + """SDFG passes validation after InsertExplicitCopies.""" + cpu = dace.StorageType.CPU_Heap + sdfg, _, _, _ = _build_copy_sdfg("validate_after", [("A", [100], cpu), ("B", [100], cpu)], Memlet("A[0:100]")) + + InsertExplicitCopies().apply_pass(sdfg, {}) + _assert_no_other_subset(sdfg) + sdfg.validate() + + +def _make_view_round_trip_sdfg(name, *, dst_side=False): + """Build a round-trip through ``A_view``, a 5x6 view of the 4x5x6 array ``A``. + + Source-side (default) flows ``A[1] -> A_view -> other``; dst-side flows + ``other -> A_view -> A[1]`` (view aliases the write target). + + :returns: ``(sdfg, state, a, view, other)`` -- ``a`` is the 4x5x6 array, ``other`` the 5x6 one. + """ + cpu = dace.StorageType.CPU_Heap + sdfg = dace.SDFG(name) + sdfg.add_array("A", [4, 5, 6], dace.float64, storage=cpu) + sdfg.add_view("A_view", [5, 6], dace.float64, storage=cpu) + sdfg.add_array("other", [5, 6], dace.float64, storage=cpu) + st = sdfg.add_state("s") + a, v, o = st.add_access("A"), st.add_access("A_view"), st.add_access("other") + if dst_side: + st.add_edge(o, None, v, None, Memlet("other[0:5, 0:6]")) + st.add_edge(v, None, a, None, Memlet("A[1, 0:5, 0:6]")) + else: + st.add_edge(a, None, v, None, Memlet("A[1, 0:5, 0:6]")) + st.add_edge(v, None, o, None, Memlet("A_view[0:5, 0:6]")) + return sdfg, st, a, v, o + + +def test_insert_view_src_round_trip_lifts_movement_edge(): + """``A -> A_view -> sink``: alias edge kept, movement edge lifted to ``A -> A_view -> Copy -> sink``.""" + sdfg, st, a, v, out = _make_view_round_trip_sdfg("view_src_movement") + InsertExplicitCopies().apply_pass(sdfg, {}) + sdfg.validate() + + assert v in st.nodes(), "the view must be preserved as a copy endpoint" + assert _count_copy_nodes(sdfg) == 1 + a_out = list(st.out_edges(a)) + assert len(a_out) == 1 and a_out[0].dst is v, "alias edge A -> A_view must be untouched" + v_out = list(st.out_edges(v)) + assert len(v_out) == 1 and isinstance(v_out[0].dst, CopyLibraryNode) + assert isinstance(list(st.in_edges(out))[0].src, CopyLibraryNode) + + +def test_insert_view_src_round_trip_numerical(): + """The copy lifted onto a source-side view reads the viewed slice correctly end to end.""" + sdfg, st, a, v, out = _make_view_round_trip_sdfg("view_src_numerical") + InsertExplicitCopies().apply_pass(sdfg, {}) + _assert_no_copynd(sdfg) + + A = np.arange(4 * 5 * 6, dtype=np.float64).reshape(4, 5, 6).copy() + other = np.zeros((5, 6), dtype=np.float64) + sdfg(A=A, other=other) + np.testing.assert_array_equal(other, A[1]) + + +def test_insert_view_dst_round_trip_numerical(): + """``other -> A_view -> A``: the view aliases the write target, is preserved, and data lands in ``A[1]``.""" + sdfg, st, a, v, o = _make_view_round_trip_sdfg("view_dst_numerical", dst_side=True) + InsertExplicitCopies().apply_pass(sdfg, {}) + sdfg.validate() + assert v in st.nodes(), "the view must be preserved as a copy endpoint" + assert _count_copy_nodes(sdfg) == 1 + + _assert_no_copynd(sdfg) + other = np.arange(5 * 6, dtype=np.float64).reshape(5, 6).copy() + A = np.zeros((4, 5, 6), dtype=np.float64) + sdfg(A=A, other=other) + np.testing.assert_array_equal(A[1], other) + assert np.all(A[0] == 0) and np.all(A[2:] == 0) + + +def test_insert_self_copy_subset_is_dst_side(): + """Self-copy ``p -> p``: ``subset`` maps to ``_out`` (dst), ``other_subset`` to ``_in`` (src); reversing them + would silently produce a backwards copy.""" + sdfg = dace.SDFG("self_copy_subset_dst") + sdfg.add_array("p", [4, 5], dace.float64) + + st = sdfg.add_state("s") + a = st.add_access("p") + b = st.add_access("p") + st.add_edge(a, None, b, None, Memlet(data="p", subset="0:4, 4", other_subset="0:4, 3")) + + InsertExplicitCopies().apply_pass(sdfg, {}) + sdfg.validate() + + copies = [n for n in st.nodes() if isinstance(n, CopyLibraryNode)] + assert len(copies) == 1 + cn = copies[0] + in_e = [e for e in st.in_edges(cn) if e.dst_conn == CopyLibraryNode.INPUT_CONNECTOR_NAME][0] + out_e = [e for e in st.out_edges(cn) if e.src_conn == CopyLibraryNode.OUTPUT_CONNECTOR_NAME][0] + + assert str(in_e.data.subset) == "0:4, 3", (f"src side should read column 3 (other_subset); got {in_e.data.subset}") + assert str(out_e.data.subset) == "0:4, 4", (f"dst side should write column 4 (subset); got {out_e.data.subset}") + + +def _check_reshape_copy(sdfg, dst_name, dst_shape): + """Assert the SDFG validates and the single lifted ``CopyLibraryNode`` output memlet spans the full ``dst_shape``.""" + sdfg.validate() + copies = [n for n, _ in sdfg.all_nodes_recursive() if isinstance(n, CopyLibraryNode)] + assert len(copies) == 1, f"expected exactly one CopyLibraryNode, got {len(copies)}" + cn = copies[0] + parent = next(p for n, p in sdfg.all_nodes_recursive() if n is cn) + out_e = [e for e in parent.out_edges(cn) if e.src_conn == CopyLibraryNode.OUTPUT_CONNECTOR_NAME][0] + assert out_e.data.data == dst_name + assert str(out_e.data.subset) == ', '.join( + f"0:{s}" for s in dst_shape), (f"dst memlet subset should span full {dst_shape}, got {out_e.data.subset}") + + +def _run_reshape_copy_test(prefix, src_shape, dst_shape): + """Build ``A[full] -> B`` (no other_subset), lift, and assert the derived destination range spans all of ``B``.""" + cpu = dace.StorageType.CPU_Heap + sdfg, _, _, _ = _build_copy_sdfg(f"{prefix}_{len(src_shape)}_to_{len(dst_shape)}", [("A", src_shape, cpu), + ("B", dst_shape, cpu)], + Memlet(data="A", subset=', '.join(f"0:{s}" for s in src_shape))) + InsertExplicitCopies().apply_pass(sdfg, {}) + _check_reshape_copy(sdfg, "B", dst_shape) + + +@pytest.mark.parametrize( + "src_shape,dst_shape", + [ + ([8, 12, 5, 3], [96, 5, 3]), # collapse leading two: einsum_blas test_4x4 pattern + ([8, 10, 12], [80, 12]), # collapse leading two: einsum_blas test_3x2 pattern + ([8, 12, 5, 3], [8, 60, 3]), # collapse middle two + ([2, 3, 4, 5], [6, 20]), # double collapse: dims 0-1 and dims 2-3 + ([8, 12, 5, 3], [1440]), # full flatten + ]) +def test_insert_consecutive_collapse_reshape(src_shape, dst_shape): + """Collapsing contiguous source dims derives a full-destination subset, not the rank-mismatched ``src_subset``.""" + _run_reshape_copy_test("reshape_collapse", src_shape, dst_shape) + + +@pytest.mark.parametrize( + "src_shape,dst_shape", + [ + ([80, 12], [8, 10, 12]), # split leading dim + ([96, 5, 3], [8, 12, 5, 3]), # split leading dim + ([1440], [8, 12, 5, 3]), # full unflatten + ([6, 20], [2, 3, 4, 5]), # double split + ]) +def test_insert_consecutive_split_reshape(src_shape, dst_shape): + """Inverse split case: a higher-rank destination by splitting source dims, via the same symmetric code path.""" + _run_reshape_copy_test("reshape_split", src_shape, dst_shape) + + +@pytest.mark.parametrize( + "src_shape,dst_shape", + [ + ([8, 1, 12], [8, 12]), # squeeze a length-1 dim + ([8, 12, 1, 5], [96, 5]), # squeeze + collapse + ([1, 96, 5, 3], [8, 12, 5, 3]), # leading 1 + split + ]) +def test_insert_reshape_with_squeezed_ones(src_shape, dst_shape): + """Unit-length dims on either side are ignored when matching a consecutive collapse or split.""" + _run_reshape_copy_test("reshape_squeeze", src_shape, dst_shape) + + +def test_insert_view_rewrite_is_idempotent_under_repeated_apply(): + """Repeated ``apply_pass`` calls do not accumulate extra ``CopyLibraryNode``s; the only remaining ``AN -> AN`` + edge is the view's alias edge, so later runs are no-ops.""" + sdfg, st, _, _, _ = _make_view_round_trip_sdfg("view_rewrite_idempotent") + p = InsertExplicitCopies() + p.apply_pass(sdfg, {}) + n_after_first = _count_copy_nodes(sdfg) + assert n_after_first == 1 + + for _ in range(5): + p.apply_pass(sdfg, {}) + + assert _count_copy_nodes(sdfg) == n_after_first + sdfg.validate() + + +@pytest.mark.gpu +@pytest.mark.parametrize("sdfg_name,src_name,src_storage,dst_name,dst_storage,size", [ + ("insert_cpu_gpu", "H", dace.StorageType.CPU_Heap, "G", dace.StorageType.GPU_Global, 64), + ("insert_gpu_cpu", "G", dace.StorageType.GPU_Global, "H", dace.StorageType.CPU_Heap, 64), + ("insert_gpu_gpu", "A", dace.StorageType.GPU_Global, "B", dace.StorageType.GPU_Global, 128), +], + ids=["cpu_to_gpu", "gpu_to_cpu", "gpu_to_gpu"]) +def test_insert_cross_storage_transfer(sdfg_name, src_name, src_storage, dst_name, dst_storage, size): + """Structural check for cross-storage (CPU<->GPU, GPU<->GPU) transfers.""" + sdfg, _, _, _ = _build_copy_sdfg(sdfg_name, [(src_name, [size], src_storage), (dst_name, [size], dst_storage)], + Memlet(f"{src_name}[0:{size}]")) + + InsertExplicitCopies().apply_pass(sdfg, {}) + _assert_no_other_subset(sdfg) + assert _count_copy_nodes(sdfg) == 1 + assert _count_direct_copy_edges(sdfg) == 0 + _assert_copy_storages(sdfg, src_storage, dst_storage) + + +_N = dace.symbol('_N') + + +def test_iec_skips_array_to_view_edge(): + """An AccessNode -> View edge is left direct (no ``CopyLibraryNode`` inserted).""" + sdfg = dace.SDFG('skip_array_to_view') + sdfg.add_array('A', [4, 5, 6], dace.float64) + sdfg.add_view('Av', [5, 6], dace.float64) + state = sdfg.add_state() + a = state.add_access('A') + v = state.add_access('Av') + state.add_edge(a, None, v, None, Memlet('A[1, 0:5, 0:6]')) + InsertExplicitCopies().apply_pass(sdfg, {}) + assert _count_copy_nodes(sdfg) == 0 + in_e = list(state.in_edges(v)) + assert len(in_e) == 1 and in_e[0].src is a + + +def test_iec_round_trip_view_lifts_one_copy(): + """An A -> View -> sink round-trip lifts one ``CopyLibraryNode``, keeps the View, and stays correct.""" + sdfg, state, _, v, _ = _make_view_round_trip_sdfg("round_trip_view") + InsertExplicitCopies().apply_pass(sdfg, {}) + assert _count_copy_nodes(sdfg) == 1 + assert v in state.nodes() + sdfg.validate() + A = np.copy(np.arange(120, dtype=np.float64).reshape(4, 5, 6)) + other = np.zeros((5, 6), dtype=np.float64) + sdfg(A=A, other=other) + assert np.array_equal(other, A[1]) + + +def test_iec_view_multiple_consumers_each_lifted(): + """Each movement edge off a multiply-consumed View is lifted; the View is kept.""" + sdfg, state, _, v, _ = _make_view_round_trip_sdfg("view_multiple_consumers") + sdfg.add_array("also_reads", [5, 6], dace.float64, storage=dace.StorageType.CPU_Heap) + state.add_edge(v, None, state.add_access("also_reads"), None, Memlet("A_view[0:5, 0:6]")) + InsertExplicitCopies().apply_pass(sdfg, {}) + assert v in state.nodes() + assert _count_copy_nodes(sdfg) == 2 + sdfg.validate() + + +def test_iec_skips_reshape_view_edge(): + """A reshape (rank-changing) AccessNode -> View edge is left direct with no ``CopyLibraryNode``.""" + sdfg = dace.SDFG('skip_reshape_view') + sdfg.add_array('A', [2, 3, 4], dace.float64) + sdfg.add_view('Av', [8, 3], dace.float64) + state = sdfg.add_state() + a = state.add_access('A') + v = state.add_access('Av') + state.add_edge(a, None, v, None, Memlet(data='A', subset='0:2, 0:3, 0:4', other_subset='0:8, 0:3')) + InsertExplicitCopies().apply_pass(sdfg, {}) + assert _count_copy_nodes(sdfg) == 0 + + +@pytest.mark.parametrize( + "name,src_shape,dst_shape,subset,other_subset,expected", + [ + # constant-index dims collapse to matching rank... + ("const_first", [5, 4, 3], [4, 3], "2, 0:4, 0:3", "0:4, 0:3", lambda s: s[2]), + ("const_middle", [4, 5, 3], [4, 3], "0:4, 2, 0:3", "0:4, 0:3", lambda s: s[:, 2, :]), + # ...and volume-equal reshapes take the MappedTasklet rank-mismatch path. + ("rank_change", [2, 3, 4], [8, 3], "0:2, 0:3, 0:4", "0:8, 0:3", lambda s: s.reshape(8, 3)), + ("flatten", [4, 3], [12], "0:4, 0:3", "0:12", lambda s: s.reshape(12)), + ]) +def test_iec_array_to_array_rank_mismatch(name, src_shape, dst_shape, subset, other_subset, expected): + """Rank-mismatched copies (constant-index collapse or volume-equal reshape) copy correctly.""" + default = dace.StorageType.Default + sdfg, _, _, _ = _build_copy_sdfg(f"a2a_{name}", [("src", src_shape, default), ("dst", dst_shape, default)], + Memlet(data="src", subset=subset, other_subset=other_subset)) + InsertExplicitCopies().apply_pass(sdfg, {}) + sdfg.validate() + src = np.copy(np.arange(int(np.prod(src_shape)), dtype=np.float64).reshape(src_shape)) + dst = np.zeros(dst_shape, dtype=np.float64) + sdfg(src=src, dst=dst) + assert np.array_equal(dst, expected(src)) + + +@dace.program +def _iec_pin_reshape_rank_change(A: dace.float64[2, 3, 4], B: dace.float64[8, 3]): + C = np.reshape(A, [8, 3]) + B[:] += C + + +def test_iec_reshape_does_not_lift_view(): + """The pass does not lift a reshape view in a real program; output stays numerically correct.""" + sdfg = _iec_pin_reshape_rank_change.to_sdfg(simplify=True) + InsertExplicitCopies().apply_pass(sdfg, {}) + sdfg.validate() + A = np.random.rand(2, 3, 4) + B = np.random.rand(8, 3) + expected = np.reshape(A, [8, 3]) + B + sdfg(A=A, B=B) + assert np.allclose(B, expected) + + +@dace.program +def _iec_pin_reinterpret_dtype(A: dace.int32[_N]): + C = A.view(dace.int16) + C[:] += 1 + + +def test_iec_reinterpret_does_not_lift_view(): + """The pass does not lift a dtype-reinterpret view; output stays numerically correct.""" + sdfg = _iec_pin_reinterpret_dtype.to_sdfg(simplify=True) + InsertExplicitCopies().apply_pass(sdfg, {}) + sdfg.validate() + A = np.random.randint(0, 262144, size=[10], dtype=np.int32) + expected = np.copy(A) + expected.view(np.int16)[:] += 1 + sdfg(A=A, _N=10) + assert np.array_equal(A, expected) + + +# Map-staging lift: AN -> MapEntry/MapExit -> AN copies put a CopyLibraryNode INSIDE the map +# scope, wired directly to the map connector; outer-side Views stay in place; chained +# MapEntries/MapExits are followed via memlet_path; generated code emits no CopyND. + +_CPU = dace.dtypes.StorageType.CPU_Heap +_N_STAGE = 128 +_TILE = 32 + + +def _build_stage_in_sdfg(name: str, with_view: bool = False) -> dace.SDFG: + """Build ``A -> MapEntry -> local -> inner work -> B``, optionally with a View aliasing ``A``.""" + sdfg = dace.SDFG(name) + sdfg.add_array("A", [_N_STAGE], dace.float64, storage=_CPU) + sdfg.add_array("B", [_N_STAGE], dace.float64, storage=_CPU) + sdfg.add_array("local", [_TILE], dace.float64, storage=_CPU, transient=True) + if with_view: + sdfg.add_view("Av", [_N_STAGE], dace.float64, storage=_CPU) + + state = sdfg.add_state("s") + a = state.add_access("A") + b = state.add_access("B") + local = state.add_access("local") + me, mx = state.add_map("tile", {"bi": f"0:{_N_STAGE}:{_TILE}"}) + + if with_view: + av = state.add_access("Av") + state.add_edge(a, None, av, None, Memlet(f"A[0:{_N_STAGE}]")) + state.add_memlet_path(av, me, local, memlet=Memlet(f"Av[bi:bi+{_TILE}]")) + else: + state.add_memlet_path(a, me, local, memlet=Memlet(f"A[bi:bi+{_TILE}]")) + + ime, imx = state.add_map("inner", {"ti": f"0:{_TILE}"}) + t = state.add_tasklet("incr", {"_in"}, {"_out"}, "_out = _in + 1.0") + state.add_memlet_path(local, ime, t, dst_conn="_in", memlet=Memlet("local[ti]")) + state.add_memlet_path(t, imx, mx, b, src_conn="_out", memlet=Memlet("B[bi+ti]")) + return sdfg + + +def _build_stage_out_sdfg(name: str, with_view: bool = False) -> dace.SDFG: + """Build ``A -> inner work -> local -> MapExit -> B``, optionally with a View aliasing ``B``.""" + sdfg = dace.SDFG(name) + sdfg.add_array("A", [_N_STAGE], dace.float64, storage=_CPU) + sdfg.add_array("B", [_N_STAGE], dace.float64, storage=_CPU) + sdfg.add_array("local", [_TILE], dace.float64, storage=_CPU, transient=True) + if with_view: + sdfg.add_view("Bv", [_N_STAGE], dace.float64, storage=_CPU) + + state = sdfg.add_state("s") + a = state.add_access("A") + b = state.add_access("B") + local = state.add_access("local") + me, mx = state.add_map("tile", {"bi": f"0:{_N_STAGE}:{_TILE}"}) + + ime, imx = state.add_map("inner", {"ti": f"0:{_TILE}"}) + t = state.add_tasklet("incr", {"_in"}, {"_out"}, "_out = _in + 1.0") + state.add_memlet_path(a, me, ime, t, dst_conn="_in", memlet=Memlet("A[bi+ti]")) + state.add_memlet_path(t, imx, local, src_conn="_out", memlet=Memlet("local[ti]")) + + if with_view: + bv = state.add_access("Bv") + state.add_memlet_path(local, mx, bv, memlet=Memlet(f"Bv[bi:bi+{_TILE}]")) + state.add_edge(bv, None, b, None, Memlet(f"B[0:{_N_STAGE}]")) + else: + state.add_memlet_path(local, mx, b, memlet=Memlet(f"B[bi:bi+{_TILE}]")) + return sdfg + + +def _find_libnode_and_scope(state): + libnodes = [n for n in state.nodes() if isinstance(n, CopyLibraryNode)] + assert len(libnodes) == 1, f"expected exactly one CopyLibraryNode, got {len(libnodes)}" + cn = libnodes[0] + return cn, state.entry_node(cn) + + +def _assert_lifted_libnode(state, side: str, expected_scope=None): + """Assert exactly one libnode in ``state`` is inside a map scope and wired directly to it. + + :param side: ``'in'`` for stage-in (libnode input edge from MapEntry) or + ``'out'`` for stage-out (libnode output edge to MapExit). + :param expected_scope: optional MapEntry to require for the libnode's + enclosing scope; when ``None``, any MapEntry passes. + :returns: ``(libnode, enclosing_map_entry)``. + """ + cn, parent = _find_libnode_and_scope(state) + assert isinstance(parent, nodes.MapEntry), f"libnode parent scope is {type(parent).__name__}, expected MapEntry" + if expected_scope is not None: + assert parent is expected_scope, "libnode must sit in the expected (innermost) map scope" + if side == "in": + in_edges = [e for e in state.in_edges(cn) if e.dst_conn == CopyLibraryNode.INPUT_CONNECTOR_NAME] + assert len(in_edges) == 1 and in_edges[0].src is parent, \ + "libnode's input must wire directly to the MapEntry connector" + else: + out_edges = [e for e in state.out_edges(cn) if e.src_conn == CopyLibraryNode.OUTPUT_CONNECTOR_NAME] + assert len(out_edges) == 1 and isinstance(out_edges[0].dst, nodes.MapExit), \ + "libnode's output must wire directly to the MapExit connector" + return cn, parent + + +def _run_and_check(sdfg: dace.SDFG, expected_b): + A = np.arange(_N_STAGE, dtype=np.float64) + B = np.zeros(_N_STAGE, dtype=np.float64) + sdfg(A=A, B=B) + np.testing.assert_array_equal(B, expected_b(A)) + + +def test_lift_stage_in_copy(): + """``A -> MapEntry -> local`` lifts to a libnode INSIDE the map scope, wired directly to MapEntry.""" + sdfg = _build_stage_in_sdfg("stage_in") + InsertExplicitCopies().apply_pass(sdfg, {}) + + _assert_lifted_libnode(sdfg.start_state, side="in") + _assert_no_copynd(sdfg) + _run_and_check(sdfg, lambda A: A + 1.0) + + +def test_lift_stage_out_copy(): + """``local -> MapExit -> B`` lifts to a libnode INSIDE the map scope, wired directly to MapExit.""" + sdfg = _build_stage_out_sdfg("stage_out") + InsertExplicitCopies().apply_pass(sdfg, {}) + + _assert_lifted_libnode(sdfg.start_state, side="out") + _assert_no_copynd(sdfg) + _run_and_check(sdfg, lambda A: A + 1.0) + + +def _view_an_names(sdfg, state): + return [ + n.data for n in state.nodes() + if isinstance(n, nodes.AccessNode) and isinstance(sdfg.arrays[n.data], dace.data.View) + ] + + +def test_lift_stage_in_copy_through_view(): + """``A -> A_view -> MapEntry -> local``: View stays in place; libnode placed between MapEntry and inner AN.""" + sdfg = _build_stage_in_sdfg("stage_in_view", with_view=True) + InsertExplicitCopies().apply_pass(sdfg, {}) + + _assert_lifted_libnode(sdfg.start_state, side="in") + assert _view_an_names(sdfg, sdfg.start_state) == ["Av"] + _assert_no_copynd(sdfg) + _run_and_check(sdfg, lambda A: A + 1.0) + + +def test_lift_stage_out_copy_through_view(): + """``local -> MapExit -> B_view -> B``: View stays in place; libnode placed between local and MapExit.""" + sdfg = _build_stage_out_sdfg("stage_out_view", with_view=True) + InsertExplicitCopies().apply_pass(sdfg, {}) + + _assert_lifted_libnode(sdfg.start_state, side="out") + assert _view_an_names(sdfg, sdfg.start_state) == ["Bv"] + _assert_no_copynd(sdfg) + _run_and_check(sdfg, lambda A: A + 1.0) + + +def _build_chained_stage_sdfg(name, *, stage_in): + """2-level tiled map nest with a chained stage-in (``A -> ME1 -> ME2 -> local``) or + stage-out (``local -> MX2 -> MX1 -> B``) copy through the inner-block scope. + + :returns: ``(sdfg, state, inner_block_entry)`` -- the inner-block map (ME2), where the + lifted libnode is expected to land. + """ + N, TILE, INNER = 64, 16, 4 + sdfg = dace.SDFG(name) + sdfg.add_array("A", [N], dace.float64, storage=_CPU) + sdfg.add_array("B", [N], dace.float64, storage=_CPU) + sdfg.add_array("local", [INNER], dace.float64, storage=_CPU, transient=True) + state = sdfg.add_state("s") + a, b, local = state.add_access("A"), state.add_access("B"), state.add_access("local") + me1, mx1 = state.add_map("outer", {"bi": f"0:{N}:{TILE}"}) + me2, mx2 = state.add_map("inner_block", {"si": f"0:{TILE}:{INNER}"}) + ime, imx = state.add_map("inner", {"ti": f"0:{INNER}"}) + t = state.add_tasklet("incr", {"_in"}, {"_out"}, "_out = _in + 1.0") + if stage_in: + state.add_memlet_path(a, me1, me2, local, memlet=Memlet(f"A[bi+si:bi+si+{INNER}]")) + state.add_memlet_path(local, ime, t, dst_conn="_in", memlet=Memlet("local[ti]")) + state.add_memlet_path(t, imx, mx2, mx1, b, src_conn="_out", memlet=Memlet("B[bi+si+ti]")) + else: + state.add_memlet_path(a, me1, me2, ime, t, dst_conn="_in", memlet=Memlet("A[bi+si+ti]")) + state.add_memlet_path(t, imx, local, src_conn="_out", memlet=Memlet("local[ti]")) + state.add_memlet_path(local, mx2, mx1, b, memlet=Memlet(f"B[bi+si:bi+si+{INNER}]")) + return sdfg, state, me2 + + +def test_lift_stage_in_copy_chained_map_entries(): + """``A -> ME1 -> ME2 -> local``: lift through nested MapEntries; libnode at innermost scope.""" + sdfg, state, me2 = _build_chained_stage_sdfg("stage_in_nested", stage_in=True) + InsertExplicitCopies().apply_pass(sdfg, {}) + _assert_lifted_libnode(state, side="in", expected_scope=me2) + _assert_no_copynd(sdfg) + A = np.arange(64, dtype=np.float64) + B = np.zeros(64, dtype=np.float64) + sdfg(A=A, B=B) + np.testing.assert_array_equal(B, A + 1.0) + + +def test_lift_stage_out_copy_chained_map_exits(): + """Symmetric: ``local -> MX2 -> MX1 -> B`` -- libnode at innermost scope, wired directly to MX2.""" + sdfg, state, me2 = _build_chained_stage_sdfg("stage_out_nested", stage_in=False) + InsertExplicitCopies().apply_pass(sdfg, {}) + _assert_lifted_libnode(state, side="out", expected_scope=me2) + _assert_no_copynd(sdfg) + A = np.arange(64, dtype=np.float64) + B = np.zeros(64, dtype=np.float64) + sdfg(A=A, B=B) + np.testing.assert_array_equal(B, A + 1.0) + + +def _make_inner_nested_sdfg(body_name: str, inout_name: str, size: int, op: str) -> dace.SDFG: + """Tiny NestedSDFG: ``inout[i] = op(inout[i])`` over ``i = 0:size``.""" + nsdfg = dace.SDFG(body_name) + nsdfg.add_array(inout_name, [size], dace.float64) + st = nsdfg.add_state("body") + a = st.add_access(inout_name) + b = st.add_access(inout_name) + me, mx = st.add_map("inner", {"ti": f"0:{size}"}) + t = st.add_tasklet("op", {"_in"}, {"_out"}, f"_out = {op}") + st.add_memlet_path(a, me, t, dst_conn="_in", memlet=Memlet(f"{inout_name}[ti]")) + st.add_memlet_path(t, mx, b, src_conn="_out", memlet=Memlet(f"{inout_name}[ti]")) + return nsdfg + + +def test_lift_stage_in_copy_with_nested_sdfg_consumer(): + """``A -> MapEntry -> local`` where ``local`` feeds a NestedSDFG inside the map: lift unaffected.""" + sdfg = dace.SDFG("stage_in_nsdfg") + sdfg.add_array("A", [_N_STAGE], dace.float64, storage=_CPU) + sdfg.add_array("B", [_N_STAGE], dace.float64, storage=_CPU) + sdfg.add_array("local", [_TILE], dace.float64, storage=_CPU, transient=True) + state = sdfg.add_state("s") + a = state.add_access("A") + b = state.add_access("B") + local = state.add_access("local") + me, mx = state.add_map("tile", {"bi": f"0:{_N_STAGE}:{_TILE}"}) + state.add_memlet_path(a, me, local, memlet=Memlet(f"A[bi:bi+{_TILE}]")) + + nsdfg = _make_inner_nested_sdfg("inner_body", "buf", _TILE, "_in + 1.0") + nnode = state.add_nested_sdfg(nsdfg, {"buf"}, {"buf"}) + state.add_edge(local, None, nnode, "buf", Memlet(f"local[0:{_TILE}]")) + out_local = state.add_access("local") + state.add_edge(nnode, "buf", out_local, None, Memlet(f"local[0:{_TILE}]")) + state.add_memlet_path(out_local, mx, b, memlet=Memlet(f"B[bi:bi+{_TILE}]")) + + InsertExplicitCopies().apply_pass(sdfg, {}) + state = sdfg.start_state + # Both the stage-in and stage-out edges lift. + libnodes = [n for n in state.nodes() if isinstance(n, CopyLibraryNode)] + assert len(libnodes) == 2 + for cn in libnodes: + assert isinstance(state.entry_node(cn), nodes.MapEntry) + + _assert_no_copynd(sdfg) + A = np.arange(_N_STAGE, dtype=np.float64) + B = np.zeros(_N_STAGE, dtype=np.float64) + sdfg(A=A, B=B) + np.testing.assert_array_equal(B, A + 1.0) + + +# Polybench-derived tests: the pass must preserve numerical output on real programs. +# Init wrappers allocate the arrays and delegate to the canonical programs' ``init_array``. + + +def _run_and_compare(program, init_fn, check_arrays, sizes, name): + """Run a program before and after InsertExplicitCopies, asserting numerical correctness.""" + sdfg_ref = program.to_sdfg(simplify=True) + ref_exe = sdfg_ref.compile() + ref_arrays = init_fn(**sizes) + ref_exe(**{k: v for k, v in ref_arrays.items()}, **sizes) + ref_values = {k: ref_arrays[k].copy() for k in check_arrays} + + sdfg_pass = _copy.deepcopy(sdfg_ref) + InsertExplicitCopies().apply_pass(sdfg_pass, {}) + _assert_no_other_subset(sdfg_pass) + sdfg_pass.expand_library_nodes() + pass_exe = sdfg_pass.compile() + pass_arrays = init_fn(**sizes) + pass_exe(**{k: v for k, v in pass_arrays.items()}, **sizes) + + for arr_name in check_arrays: + np.testing.assert_allclose(pass_arrays[arr_name], + ref_values[arr_name], + rtol=1e-10, + atol=1e-12, + err_msg=f"{name}: array '{arr_name}' mismatch after pass") + + +def _init_fdtd2d(NX, NY, TMAX): + ex = np.zeros((NX, NY), dtype=np.float64) + ey = np.zeros((NX, NY), dtype=np.float64) + hz = np.zeros((NX, NY), dtype=np.float64) + fict = np.zeros(TMAX, dtype=np.float64) + _fdtd2d_init_array(ex, ey, hz, fict, NX, NY, TMAX) + return {"ex": ex, "ey": ey, "hz": hz, "_fict_": fict} + + +def _init_correlation(N, M): + data = np.zeros((N, M), dtype=np.float64) + corr = np.zeros((M, M), dtype=np.float64) + mean = np.zeros(M, dtype=np.float64) + stddev = np.zeros(M, dtype=np.float64) + _correlation_init_array(data, corr, mean, stddev, N, M) + return {"data": data, "corr": corr, "mean": mean, "stddev": stddev} + + +def _init_covariance(N, M): + data = np.zeros((N, M), dtype=np.float64) + cov = np.zeros((M, M), dtype=np.float64) + mean = np.zeros(M, dtype=np.float64) + _covariance_init_array(data, cov, mean, N, M) + return {"data": data, "cov": cov, "mean": mean} + + +def test_polybench_fdtd2d(): + """``InsertExplicitCopies`` preserves fdtd2d output versus the untransformed reference.""" + _run_and_compare(fdtd2d, _init_fdtd2d, ["ex", "ey", "hz"], {"NX": 20, "NY": 30, "TMAX": 10}, "fdtd2d") + + +def test_polybench_correlation(): + """``InsertExplicitCopies`` preserves correlation output versus the untransformed reference.""" + _run_and_compare(correlation, _init_correlation, ["corr"], {"N": 32, "M": 28}, "correlation") + + +def test_polybench_covariance(): + """``InsertExplicitCopies`` preserves covariance output versus the untransformed reference.""" + _run_and_compare(covariance, _init_covariance, ["cov"], {"N": 32, "M": 28}, "covariance") + + +def test_iec_skips_dtype_converting_copy(): + """A direct copy between different dtypes is a cast, not a byte move: the pass must leave it + for tasklet lowering rather than insert a ``CopyLibraryNode`` (memcpy), which cannot convert. + Regression: the direct-copy path lacked the dtype guard its staging path already has.""" + cpu = dace.StorageType.CPU_Heap + sdfg = dace.SDFG("iec_dtype_convert") + sdfg.add_array("A", [64], dace.float32, cpu) + sdfg.add_array("B", [64], dace.float64, cpu) + st = sdfg.add_state("s") + a = st.add_access("A") + b = st.add_access("B") + st.add_edge(a, None, b, None, Memlet("A[0:64]")) + + InsertExplicitCopies().apply_pass(sdfg, {}) + + assert _count_copy_nodes(sdfg) == 0, "a dtype-converting copy must not be lowered to CopyLibraryNode" + assert _count_direct_copy_edges(sdfg) == 1, "the dtype-converting edge must be left in place" + + +def test_iec_skips_reference_set_edge(): + """A ``set`` edge binds a POINTER, it does not move data. Rewriting it into a ``CopyLibraryNode`` + drops the ``set`` connector, so the Reference is never bound and validation rejects the SDFG. + Regression: both the direct-copy and the map-staging path lifted reference-set edges.""" + sdfg = dace.SDFG("iec_reference_set") + sdfg.add_array("A", [20], dace.float64) + sdfg.add_reference("ref", [20], dace.float64) + setstate = sdfg.add_state("set_ref") + setstate.add_edge(setstate.add_read("A"), None, setstate.add_write("ref"), "set", Memlet("A[0:20]")) + usestate = sdfg.add_state_after(setstate, "use_ref") + t = usestate.add_tasklet("one", {}, {"o"}, "o = 1.0") + usestate.add_edge(t, "o", usestate.add_write("ref"), None, Memlet("ref[3]")) + + InsertExplicitCopies().apply_pass(sdfg, {}) + + assert _count_copy_nodes(sdfg) == 0, "a reference-set edge must not be lowered to CopyLibraryNode" + assert [e for e in setstate.edges() if e.dst_conn == "set"], "the reference-set edge must survive the pass" + sdfg.validate() + A = np.zeros(20, dtype=np.float64) + sdfg(A=A) + assert A[3] == 1.0, "the write through the Reference must land in the aliased array" + + +def test_iec_skips_reference_set_edge_through_map(): + """Same guard on the map-staging path: ``A -> MapEntry -> ref[set]`` must not be lifted.""" + sdfg = dace.SDFG("iec_reference_set_staged") + sdfg.add_array("A", [20], dace.float64) + sdfg.add_reference("ref", [4], dace.float64) + state = sdfg.add_state("s") + me, mx = state.add_map("m", {"i": "0:20:4"}, schedule=dace.dtypes.ScheduleType.Sequential) + ref = state.add_access("ref") + state.add_memlet_path(state.add_read("A"), me, ref, dst_conn="set", memlet=Memlet("A[i:i+4]")) + t = state.add_tasklet("one", {}, {"o"}, "o = 1.0") + state.add_edge(ref, None, t, None, Memlet()) + state.add_memlet_path(t, mx, state.add_write("A"), src_conn="o", memlet=Memlet("A[i]")) + + InsertExplicitCopies().apply_pass(sdfg, {}) + + assert _count_copy_nodes(sdfg) == 0, "a reference-set edge must not be lifted through a MapEntry" + assert [e for e in state.edges() if e.dst_conn == "set"], "the reference-set edge must survive the pass" + sdfg.validate() + + +def test_iec_staging_keeps_memlet_named_inner_subset(): + """When a staging memlet names BOTH sides, that mapping IS the copy. Deriving the inner subset + from the outer one instead silently retargets the access: ``B[1, i] -> [i + 2, 3]`` copied + ``A[1, i]``, a wrong-address copy that still validates and still runs.""" + sdfg = dace.SDFG("iec_staging_two_sided") + sdfg.add_array("B", [4, 4], dace.float64, storage=_CPU) + sdfg.add_array("A", [4, 4], dace.float64, storage=_CPU, transient=True) + state = sdfg.add_state("s") + a = state.add_access("A") + me, mx = state.add_map("outer", {"i": "0:2"}, schedule=dace.dtypes.ScheduleType.Sequential) + ime, imx = state.add_map("fill", {"r": "0:4", "c": "0:4"}, schedule=dace.dtypes.ScheduleType.Sequential) + t = state.add_tasklet("fill", {}, {"_out"}, "_out = 10.0 * r + c") + state.add_edge(me, None, ime, None, Memlet()) + state.add_edge(ime, None, t, None, Memlet()) + state.add_memlet_path(t, imx, a, src_conn="_out", memlet=Memlet("A[r, c]")) + mx.add_in_connector("IN_b") + mx.add_out_connector("OUT_b") + state.add_edge(a, None, mx, "IN_b", Memlet(data="B", subset="1, i", other_subset="i + 2, 3")) + state.add_edge(mx, "OUT_b", state.add_write("B"), None, Memlet(data="B", subset="1, i")) + + assert InsertExplicitCopies().apply_pass(sdfg, {}) == 1 + + cn, _ = _find_libnode_and_scope(state) + in_edges = [e for e in state.in_edges(cn) if e.dst_conn == CopyLibraryNode.INPUT_CONNECTOR_NAME] + assert str(in_edges[0].data.subset) == "i + 2, 3", \ + f"inner subset must keep the mapping the memlet named, got {in_edges[0].data.subset}" + + # A[r, c] == 10 * r + c, so the named mapping yields A[2, 3] and A[3, 3]; the derived one gave A[1, 0:2]. + B = np.zeros((4, 4), dtype=np.float64) + sdfg(B=B) + np.testing.assert_array_equal(B[1], np.array([23.0, 33.0, 0.0, 0.0])) + + +def test_iec_symbolic_reshape_targets_the_whole_destination(): + """``A[1:N-1, 0:M]`` into a transient shaped ``[N-2, M]`` moves the whole destination, so the + derived side must be the destination's full range. The element counts are equal but can come + from two symbol instances of the same name, so a cancel-based comparison may answer None; + equalizing first keeps the comparison conclusive instead of relying on the ``is not False`` + gate to paper over an unresolved symbol identity.""" + N = dace.symbol("N", dtype=dace.int64) + M = dace.symbol("M", dtype=dace.int64) + sdfg = dace.SDFG("iec_symbolic_reshape") + sdfg.add_array("A", [N, M], dace.float32) + sdfg.add_transient("At", [N - 2, M], dace.float32) + st = sdfg.add_state("s") + st.add_edge(st.add_access("A"), None, st.add_access("At"), None, Memlet("A[1:N-1, 0:M]")) + + InsertExplicitCopies().apply_pass(sdfg, {}) + + assert _count_copy_nodes(sdfg) == 1 + out_edges = [e for state in sdfg.states() for e in state.edges() if e.data.data == "At"] + assert len(out_edges) == 1 + assert str(out_edges[0].data.subset) == "0:N - 2, 0:M" + + +def test_iec_skips_wcr_staging_edge(): + """A WCR edge is a reduction, not a copy. ``CopyLibraryNode``'s expansions emit an unconditional + store, so lifting the tile-merge edge AccumulateTransient produces turns ``out[i] += tile[i]`` + into ``out[i] = tile[i]`` -- silently, with a valid SDFG and a wrong answer.""" + sdfg = dace.SDFG("iec_wcr_staging") + sdfg.add_array("A", [8], dace.float64) + sdfg.add_array("out", [1], dace.float64) + sdfg.add_transient("tile", [1], dace.float64) + state = sdfg.add_state("s") + + me, mx = state.add_map("m", {"i": "0:8"}, schedule=dace.dtypes.ScheduleType.Sequential) + t = state.add_tasklet("copy", {"inp"}, {"o"}, "o = inp") + tile = state.add_access("tile") + state.add_memlet_path(state.add_read("A"), me, t, dst_conn="inp", memlet=Memlet("A[i]")) + state.add_edge(t, "o", tile, None, Memlet("tile[0]")) + # The merge back out of the map scope accumulates -- this is the edge that must not be lifted. + state.add_memlet_path(tile, mx, state.add_write("out"), memlet=Memlet("out[0]", wcr="lambda a, b: a + b")) + + InsertExplicitCopies().apply_pass(sdfg, {}) + + assert _count_copy_nodes(sdfg) == 0, "a WCR edge must not be lowered to CopyLibraryNode" + wcr_edges = [e for _, e in _wcr_edges(sdfg)] + assert wcr_edges, "the accumulate must survive the pass" + sdfg.validate() + + A = np.arange(8, dtype=np.float64) + out = np.zeros(1, dtype=np.float64) + sdfg(A=A, out=out) + assert out[0] == A.sum(), f"expected the accumulate {A.sum()}, got {out[0]} (an overwrite keeps only the last)" + + +def test_iec_keeps_the_ordering_edge_on_the_node_that_writes(): + """An empty memlet is a happens-before edge, and lifting a copy moves the write it constrained. + + The map reads ``A[i]`` into ``tmp_A`` and writes ``A[(i+1)%2]`` from ``tmp_B``, with an ordering + edge saying the read happens after that write. Left on the access node, the constraint no longer + reaches the node that performs the read, and the copy is free to be scheduled ahead of it -- a + silently wrong answer, not an error. + """ + sdfg = dace.SDFG("iec_ordering_edge") + sdfg.add_array("A", [2], dace.int32) + sdfg.add_array("B", [2], dace.int32) + sdfg.add_transient("tmp_A", [1], dace.int32) + sdfg.add_transient("tmp_B", [1], dace.int32) + state = sdfg.add_state("s") + + me, mx = state.add_map("m", {"i": "0:2"}, schedule=dace.dtypes.ScheduleType.Sequential) + for conn in ("IN_A", "IN_B"): + me.add_in_connector(conn) + for conn in ("OUT_A", "OUT_B"): + me.add_out_connector(conn) + mx.add_in_connector("IN_A") + mx.add_out_connector("OUT_A") + + a_write, a_ordered = state.add_write("A"), state.add_write("A") + tmp_a, tmp_b = state.add_write("tmp_A"), state.add_write("tmp_B") + state.add_edge(state.add_read("A"), None, me, "IN_A", Memlet("A[0:2]")) + state.add_edge(state.add_read("B"), None, me, "IN_B", Memlet("B[0:2]")) + state.add_edge(me, "OUT_A", tmp_a, None, Memlet("A[i]")) + state.add_edge(me, "OUT_B", tmp_b, None, Memlet("B[i]")) + state.add_edge(tmp_a, None, a_write, None, Memlet("tmp_A[0] -> [((i+1)%2)]")) + state.add_edge(a_write, None, mx, "IN_A", Memlet("A[0:2]")) + state.add_edge(tmp_b, None, a_ordered, None, Memlet("tmp_B[0] -> [((i+1)%2)]")) + state.add_edge(a_ordered, None, tmp_a, None, Memlet()) # the ordering edge + state.add_edge(a_ordered, None, mx, "IN_A", Memlet("A[0:2]")) + state.add_edge(mx, "OUT_A", state.add_write("A"), None, Memlet("A[0:2]")) + sdfg.validate() + + InsertExplicitCopies().apply_pass(sdfg, {}) + sdfg.validate() + + state = sdfg.states()[0] + writers = [e.src for e in state.in_edges(tmp_a) if isinstance(e.src, CopyLibraryNode)] + assert len(writers) == 1, "the stage-in copy of tmp_A was not lifted" + ordered_after = [e.src for e in state.in_edges(writers[0]) if e.data.is_empty()] + assert a_ordered in ordered_after, "the copy that now writes tmp_A is not ordered after the write it followed" + + a = np.array([7, 3], dtype=np.int32) + sdfg(A=a, B=np.array([11, 13], dtype=np.int32)) + assert a[0] == a[1], f"the ordering edge was not honoured: got {a}" + + +def test_copy_is_left_implicit_when_another_edge_writes_the_same_region(): + """Nothing orders two writes to one region that reach a node on separate edges. Plain copy-edge + codegen emits the copy when its SOURCE access node is visited, so it lands before the tasklet + that supersedes it; lifting it to a node would re-sort it after and flip which write survives.""" + sdfg = dace.SDFG("competing_writer") + sdfg.add_array("A", [4], dace.float64) + sdfg.add_array("B", [4], dace.float64) + state = sdfg.add_state("main", is_start_block=True) + a = state.add_access("A") + b = state.add_access("B") + tasklet = state.add_tasklet("supersede", {"i"}, {"o"}, "o = i + 1.0") + state.add_nedge(a, b, Memlet("A[0:4] -> [0:4]")) + state.add_edge(a, None, tasklet, "i", Memlet("A[0]")) + state.add_edge(tasklet, "o", b, None, Memlet("B[0]")) + sdfg.validate() + + InsertExplicitCopies().apply_pass(sdfg, {}) + + lifted = [n for n in state.nodes() if isinstance(n, CopyLibraryNode)] + assert not lifted, "a copy competing with another write to B was lifted" + + +def test_zero_element_copy_is_not_lifted(): + """A copy that moves nothing needs no node: plain copy-edge codegen emits nothing for it.""" + sdfg = dace.SDFG("zero_element_copy") + for name in "AB": + sdfg.add_array(name, [20, 20], dace.float64, transient=True) + state = sdfg.add_state("main", is_start_block=True) + state.add_nedge(state.add_access("A"), state.add_access("B"), Memlet("A[2:17, 2:2] -> [2:18, 3:3]")) + + InsertExplicitCopies().apply_pass(sdfg, {}) + + assert state.number_of_nodes() == 2 + assert not [n for n in state.nodes() if isinstance(n, CopyLibraryNode)] + + +def test_lifted_copy_inherits_the_state_instrumentation(): + """Providers read a copy edge's instrumentation off the state it lives in (``on_copy_begin``); + as a node the copy carries its own setting, or it drops out of the report.""" + sdfg = dace.SDFG("instrumented_copy") + sdfg.add_array("A", [4], dace.float64) + sdfg.add_array("B", [4], dace.float64) + state = sdfg.add_state("main", is_start_block=True) + state.instrument = dace.InstrumentationType.GPU_TX_MARKERS + state.add_nedge(state.add_access("A"), state.add_access("B"), Memlet("A[0:4] -> [0:4]")) + + InsertExplicitCopies().apply_pass(sdfg, {}) + + lifted = [n for n in state.nodes() if isinstance(n, CopyLibraryNode)] + assert len(lifted) == 1 + assert lifted[0].instrument == dace.InstrumentationType.GPU_TX_MARKERS + + sdfg.expand_library_nodes() + expanded = [n for n in state.nodes() if isinstance(n, (nodes.Tasklet, nodes.NestedSDFG))] + assert len(expanded) == 1 + assert expanded[0].instrument == dace.InstrumentationType.GPU_TX_MARKERS + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/tests/specialize_test.py b/tests/specialize_test.py index 5bc6db0065..402032ad2c 100644 --- a/tests/specialize_test.py +++ b/tests/specialize_test.py @@ -43,12 +43,18 @@ def test_constant_specialization(): code_nonspec = spec_sdfg.generate_code() - assert 'Dynamic' in code_nonspec[0].code + # Was ``'Dynamic' in code``: the implicit copy used to lower to the ``dace::CopyNDDynamic`` + # runtime template, whose name doubled as the "shape is still symbolic" marker. Explicit copy + # nodes lower it to a mapped tasklet, so assert specialization itself -- N and M are runtime + # arguments before it and compile-time constants after. + assert 'constexpr int64_t N' not in code_nonspec[0].code + assert 'constexpr int64_t M' not in code_nonspec[0].code spec_sdfg.specialize(dict(N=n, M=m)) code_spec = spec_sdfg.generate_code() - assert 'Dynamic' not in code_spec[0].code + assert f'constexpr int64_t N = {n};' in code_spec[0].code + assert f'constexpr int64_t M = {m};' in code_spec[0].code func = spec_sdfg.compile() func(A=input, B=output, N=n, M=m)