Explicit copy memset nodes - #2380
Open
ThrudPrimrose wants to merge 113 commits into
Open
Conversation
Introduce explicit copy / memset library nodes and a pass that lifts
every implicit AccessNode -> AccessNode (or scope-staging) edge into
a CopyLibraryNode, so the dataflow shape post-simplification is
self-describing and the legacy "copy edge gets lowered by ad-hoc
codegen" path can be deprecated.
dace/libraries/standard/nodes/copy_node.py (CopyLibraryNode + 8 expansions)
dace/libraries/standard/nodes/memset_node.py (MemsetLibraryNode + 3 expansions)
dace/libraries/standard/helper.py (shared expansion helpers)
dace/libraries/standard/environments/cpu.py (CPU environment used by ExpandMemcpyCPU)
dace/sdfg/construction_utils.py (small utilities used by copy_node)
dace/transformation/passes/insert_explicit_copies.py
(the pass)
Pass wiring: appended `InsertExplicitCopies` to `SIMPLIFY_PASSES` so
``SDFG.simplify()`` lifts implicit copies as part of standard cleanup
(idempotent — once lifted, no further implicit edges remain to match).
Tests:
tests/library/copy_node_test.py
tests/library/memset_node_test.py
tests/passes/insert_explicit_copies_test.py
…ryNode
Three small fixes that unblock the Copy/Memset libnodes' schedule
inference:
* ``infer_out_connector_type`` / ``infer_connector_types``: wrap the
``e.data.subset and e.data.subset.num_elements() == 1`` expression
in ``bool()``. A single-element ``Range`` returns False from
``__bool__`` and the bare ``and`` chain leaks the Range object
instead of a bool, which later trips the ``scalar |= …`` operator
with TypeError when the libnode's collapsed subset is empty.
* ``_determine_schedule_from_storage``: when the node under inspection
is a ``CopyLibraryNode`` / ``MemsetLibraryNode`` AND any neighbouring
memlet imposes a ``GPU_Device`` constraint, return ``GPU_Device``
directly. Without this, an H2D copy has both ``CPU_Multicore``
(from the CPU source) and ``GPU_Device`` (from the GPU sink) in its
constraint set and the existing ``len(constraints) > 1`` branch
raises ``InvalidSDFGNodeError: Cannot determine default schedule
for node copy_A_to_gpu_A``. The libnode is exactly the node class
designed to bridge storages; routing it to GPU_Device when GPU is
involved is the intended resolution.
…simplify InsertExplicitCopies materialises implicit AccessNode->AccessNode edges into CopyLibraryNodes — a shape-changing lowering step, not a shape-preserving simplification. Including it in SIMPLIFY_PASSES broke 22 tests (numpy reshape/flatten/view, redundant_copy count assertions, range_indirection / reinterpret validators) that rely on simplify preserving the implicit-edge form. Pass remains available as a standalone Pass for consumers that want it (e.g. the GPU codegen lowering pipeline, which calls it explicitly via InsertExplicitGPUGlobalMemoryCopies).
Three xfail-strict pins: - AccessNode<->View edges must not be lifted (policy: views are aliases). - Rank-changing reshape lift produces mismatched-rank memlets that trip codegen IndexError in cpp_offset_expr. - Dtype-reinterpret View lift produces CopyLibraryNode with mismatched element types, failing sdfg.validate().
…atch InsertExplicitCopies changes: - Drop the AN<->View edge lift (was inserting a Copy + intermediate buffer for every AN<->View edge, including pure aliasing edges that don't need one). - Add a round-trip collapse: AN_src -> View -> AN_dst becomes a single AN_src -> CopyLibraryNode -> AN_dst direct edge with the composed memlet (src side from the view-underlying subset, dst side from the access-side subset). The View AccessNode is removed when it has no other consumers. CopyLibraryNode expansion: - select_copy_implementation now routes rank-mismatched volume-equal copies (different rank after collapse_shape_and_strides) to CopyNDTemplate when both sides are same-storage C-packed contiguous. MappedTasklet reuses a single access expression for both endpoints, which produces a rank-mismatched memlet on the smaller side and crashes codegen. - ExpandCopyNDTemplate flattens to a 1D pointer walk when the in/out collapsed shapes have different ranks (both sides packed contiguous, so the linearization is sound). Tests: - Update the 4 view-lift tests to assert the new policy (round-trip collapses to a single Copy; View is removed if it has no other consumers; repeat applies are no-ops). - Add 10 new test_iec_* pins covering both patterns: AN<->View edges kept direct, AN->View->AN round-trip collapse, AN->View round-trip with another consumer keeping the View, AN<->AN copies with rank differing because of a constant-index dim, AN<->AN rank-mismatched volume-equal copies routed through CopyND, plus @dace.program reshape and reinterpret cases.
…h legacy and experimental codegen)
… explicit-copy-memset-nodes
…onstruction_utils Use SDFG.parent (O(1)) instead of the recursive _get_parent_state scan to find the state containing a nested-SDFG node; verified equivalent. copy_node.py imports the helper from dace.transformation.helpers (top-level, no import cycle). dace/sdfg/construction_utils.py removed (it only held these two helpers on this branch).
…nearization; selector bypasses CopyND
…se feature-regression tests
…trides in memset, Range.num_elements over reduce
…D only for shared memory)
…y2D 1D+2D branches
…converting copies Removes latent divergences in the copy/memset lowering, each with a regression test: - copy_node: extract cuda2d_pitch_params as the single source of truth for the MemcpyCUDA2D pitch/width/height. The selector gate and the expander previously encoded the same 2D stride test twice; drift meant the selector could pick MemcpyCUDA2D while the expander raised NotImplementedError. The selector now asks the same function (non-None => applies). - memset: derive the ExpandPure map bounds from the collapsed shape already used for the array descriptor instead of recomputing them from out_subset.size(), so the map rank/extent cannot diverge from the array. - insert_explicit_copies: skip dtype-converting direct copies (a cast, not a byte move) in the direct-copy path, mirroring the guard the staging path already has; a CopyLibraryNode (memcpy) cannot express a conversion. Tests: cuda2d_pitch_params branch coverage, strided pure-memset map/array agreement, and the dtype-conversion skip. Full CPU suites for all three files pass (126).
… explicit-copy-memset-nodes
select_copy_implementation had no host branch, so every CPU<->CPU multi-element copy fell through to MappedTasklet even when a single std::memcpy is exact. Add a host-resident step: when both endpoints are CPU-resident (or Default), same rank with matching per-dimension subset sizes, contiguous, and same packed layout, pick MemcpyCPU -- matching ExpandMemcpyCPU's preconditions so it always expands cleanly. Non-contiguous, mixed-layout, transpose (mismatched per-dim), and rank-changing reshapes still fall through to MappedTasklet. Fixes test_copy_fortran_packed_same_rank; transpose/reshape rejections intact.
…ubset assert _assert_no_other_subset flagged view-defining (alias) edges, which legitimately keep other_subset -- InsertExplicitCopies correctly skips them (a view references the underlying buffer, not a data move; converting one would change the SDFG). The covariance kernel's data[:,i] @ data[:,i:M] matmul-view pattern produces such edges. Mirror the pass's own view-edge skip in the assertion. Fixes test_polybench_covariance; the basic copy asserts and test_iec_skips_reshape_view_edge (which have no / expect view edges) are unaffected.
Adopt the improved Auto picker (a contiguous same-layout CPU copy lowers to a single std::memcpy instead of an element-wise MappedTasklet loop).
…iew guard to 120 cols
A contiguous CPU copy or zero whose element count is a compile-time constant >= the configurable threshold compiler.cpu.parallel_transfer_min_elements (default 1024) expands to the element map (MappedTasklet / pure), which DaCe schedules across OpenMP threads at top level and sequentially when nested. A small or symbolic (unknown at compile time) size keeps the single std::memcpy / std::memset: we do not fork an OpenMP region for a size that may be tiny at runtime.
Comment-only cleanup of copy_node.py / memset_node.py / helper.py: compress over-written prose to a lower comment ratio while keeping every sphinx :param/:returns/:raises field and every load-bearing rationale (the memcpy-vs-map threshold, the CUDA pitch anti-drift notes, the shared-memory-collective thread-block invariant, the NestedSDFG stream-connector gotcha). No code change.
Remove local secrets/tokens from .gitignore
…py insertion Two defects in InsertExplicitCopies: A reference-set edge was rewritten into a CopyLibraryNode at both the direct-copy and the map-staging site. A set binds a pointer rather than moving data, so the 'set' connector was dropped and the Reference was never bound -- validation then rejects it as used before set. _lift_staging_edge always derived the inner subset from the outer one, discarding the mapping the memlet already named: B[1, i] -> [i + 2, 3] emitted A[1, i], a silent wrong-address copy that still validates and still runs.
The branch had moved 'import polybench' out of module scope in 25 programs. Revert that so the subfolder matches main exactly. Those programs resolve their harness as a top-level module, which works when they are run as scripts because the script's own directory lands on sys.path. The pass test imports them as package modules instead, so it now puts that directory on sys.path itself.
Ranks of one job derive the same build folder, so ranks that each compile build on top of each other and can load a library another rank is still writing. Eight processes running one GPU test out of one folder failed six times; with a folder each, none. The new cache_distaware config entry names the build cache root after the rank the launcher (MPI, Flux, Slurm) advertises. It is off by default, because sharing one build is also a valid setup: distributed_compile has rank 0 build and every other rank load its folder. That path now pins the broadcast folder on the ranks that hold the SDFG, the others being free to pass None.
Raise parallel_transfer_min_elements default to 262144 (~2 MiB) and make symbolic (unknown-at-compile-time) sizes take the OpenMP element map instead of the single-call memcpy/memset path. copy_node.py already routed through the shared helper.py selector, so it inherits the new threshold without changes. Add structural tests asserting memcpy/memset vs "#pragma omp parallel for" in generated code on both sides of the threshold, and for symbolic size, for both copy and memset; supersedes the deleted non-structural selection test.
All workers/ranks see every GPU and pile CUDA contexts onto device 0, which flakes as invalid device ordinal (101) under -n 32 on cscs CI.
The 'unique' build folder hashed only the pid, so a later process with a recycled pid and the same SDFG name reused a stale build folder. Hash a per-process token (pid plus import timestamp) instead.
InsertExplicitCopies built fresh memlets that dropped allow_oob, so the author's waiver of the src/dst volume check was lost and the MappedTasklet same-rank guard rejected copies that validation and plain copy-edge codegen accept. Propagate the flag and skip the guard when it is set.
A WCR edge is a reduction, not a copy: it is what AccumulateTransient leaves behind to merge a per-tile local back into the real output. CopyLibraryNode's expansions always emit an unconditional store, so lifting one turned out[i] += tile[i] into out[i] = tile[i] -- a valid SDFG with a wrong answer. _replace_direct_copies already refused WCR edges; _lift_staging_edge now does too. Lifting a staging edge also leaves the scope-boundary edge on that connector naming the inner array. Validation resolves an edge's endpoints through the full memlet path, which used to run through the scope node to the inner AccessNode; the libnode now sits in between, so the path ends on a non-AccessNode and the boundary memlet has to be outer-relative. The test pins the first half where it is silent: without the guard the accumulate becomes an overwrite and only the last tile survives.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The pass implements 0-memset and copy nodes to make copies explicit as a part of pre-codegen passes.
The auto-lowering tries to find the most suitable implementations. It also extends the patterns DaCe can correctly lower.