Repository navigation
[New GPU Codegen Prep] Add MoveArrayOutOfKernel - #2535
Open
ThrudPrimrose wants to merge 86 commits into
Open
ThrudPrimrose wants to merge 86 commits into
ThrudPrimrose wants to merge 86 commits into
Conversation
DaCe's symbolic C++ printer emits dace::math::pow (libm, double) for a non-constant exponent, which is illegal where a C++ integer is required -- an array size, a subscript, or a loop bound. For example the radix decomposition R**(K-i-1) used as a stride in stockham_fft fails to compile. dace::math::ipow (repeated multiply, exact integer) is the correct lowering iff the exponent is a non-negative integer (its C++ exponent is unsigned, so a negative one is catastrophic). - Add an ipow SymPy Function (dace/symbolic.py) lowering to dace::math::ipow; registered so it round-trips through parsing and serialization. - Add the RelaxIntegerPowers pass rewriting Pow -> ipow for non-negative integer exponents: constants, integer-valued float literals, and symbolic exponents proven >= 0 by interval analysis over the enclosing loop / map iterator ranges. A scoped recursive descent binds each iterator (and each SDFG's declared symbol signs) as it enters loops, maps and nested SDFGs and drops them on exit; the proof reuses equalize_symbol and the assuming/ask idiom. Wired last in SIMPLIFY_PASSES. - Teach infer_symbols_from_datadescriptor to read ipow as Pow when solving a shape for its symbols (ipow is opaque to the solver otherwise). - Unskip the stockham_fft CPU test; add tests/passes/relax_integer_powers_test.py.
_SerializedSymbolicParser._functions (the serialize_symbolic/
deserialize_symbolic path used for SDFG properties like an Array's
total_size) lacked ipow, so a serialized 64*ipow(P, 2) deserialized to an
opaque Function('ipow') that never folds; int(evaluate(total_size, {P: n}))
then raised "Cannot convert symbols to int" (mpi subarrays_test BlockGather,
whose return shape is 8*P x 8*P). Register the real class so it round-trips
and folds.
…ove pass pre-codegen Root cause of the stockham_fft heap corruption: ``dace::math::ipow`` seeded its accumulator at ``a`` and looped ``b - 1`` times, so ``ipow(a, 0)`` returned ``a`` instead of ``1``. Every zero-exponent size/stride/bound was corrupted (e.g. ``ipow(R, 0)`` -> ``R``), doubling the innermost tile-transpose loop and writing one element past the output array. Seed at ``1`` and loop ``b`` times so the base case is correct; pass the scalar exponent (and base) by value. Also close the gaps that left an integer-context power as a ``double`` ``dace::math::pow``: - RelaxIntegerPowers now visits loop bounds/conditions, interstate-edge conditions + assignments, and nested-SDFG symbol mappings -- not just descriptors and memlet subsets. The scope analysis already proved these exponents non-negative; they were simply never offered to it. - cppunparse qualifies an ``ipow`` call as ``dace::math::ipow`` (the symbolic printer already does for descriptor expressions; loop bounds/interstate edges unparse through the AST path instead). - Move RelaxIntegerPowers out of SIMPLIFY_PASSES to a pre-codegen pass. A size is just a more complex expression that should not perturb simplification; keeping powers as ``Pow`` through simplify also lets SymPy fold ``R**i * R**(K-i-1) -> R**(K-1)`` before the opaque ``ipow`` freezes the form.
The single ipow template seeded the accumulator with T(1), but vector value types (half4/half8, and the CPU vector types) have no scalar constructor, so T(1) fails to compile -- regressing the CUDA half-vector gelu path that raises a value to a constant integer power (a**3 -> ipow(a, 3)). Split into two SFINAE overloads keyed on is_constructible<T, int>: scalar types seed at T(1) (keeping the a**0 == 1 base case), vector types seed at a and loop from 1. Vector ipow is only ever emitted with a compile-time exponent >= 1 (the constant-power path writes a literal 1 for exponent 0), so the base case only matters for scalars.
…-case The cppunparse/interstate path (loop bounds, interstate-edge assignments) emitted ipow through a bespoke t.func.id == 'ipow' branch that hand-qualified it to dace::math::ipow. Every other sympy/DaCe builtin on that path (min, max, int_ceil, Min, Max, Abs) instead resolves by being defined at global scope in the runtime headers, so the default _Call emits the bare name. Give ipow the same treatment: a global-scope ipow wrapper in math.h forwarding to dace::math::ipow, and drop the special-case branch.
… reject negative exponents
…s not crash the pass
…onent
dace::math::pow had integer overloads for `int` and `unsigned int` only, so int64 symbols
fell to the generic std::pow and came back `double`. stockham_fft's `R ** K` then reached
codegen as a floating-point value in two places gcc rejects outright: an OpenMP loop bound
("invalid controlling predicate") and a pointer offset in the MatMul expansion's WCR
("invalid operands of types 'dace::complex128*' and 'double'"). A constrained template now
covers every integral width, and the generic overload excludes both-integral so the two
cannot be ambiguous.
Addresses the review on the ipow overloads. The global wrapper takes an arbitrary `U` and
forwards it to an `unsigned int` parameter, so a signed -1 arrives as 4294967295 and the
loop spins instead of returning; it now answers 0, which is what pow(int, int) has always
done for a negative exponent. The guard is limited to arithmetic T because pi.h defines
ipow for the typeless_pi types, which have no T(0).
The 3 sve and 1 argument_signature failures in tests/codegen are pre-existing on this
branch: they fail identically with this commit's math.h reverted.
Hoists a transient GPU_Global array declared inside a GPU_Device map out of the kernel, giving it one slice per map iteration, and demotes one small enough to per-thread registers instead. GPUs cannot allocate device-global memory inside a kernel, so this exists for backward compatibility. The lift builds its memlet paths with add_memlet_path and finds a free name with utils.find_new_name rather than open-coding either, and the node-to-state lookup is built once for the whole pass.
ThrudPrimrose
changed the title
…r redesign The pass now walks each owning SDFG's transient once, computes the GPU levels its accesses sit under, and prefixes every memlet from that one fact. Ported from extended, each with a reproducer: - index a lifted slice from the map's first iteration, not the raw param; - order the prepended indices like the prepended dimensions (outermost first); - carry one slice per iteration out through the map exits, not the whole buffer; - translate a lifted descriptor through the nested symbol_mapping (simultaneous); - reshape and rewrite only the defining SDFG and its ancestors, never the nests below it that receive the buffer through a connector; - rewrite subscripts that interstate edges, loops and branches read; - bind the kernel params the new index names into every nest on the way down, and never into a nest that assigns that name itself (refused instead). Also drops the node-to-state cache, a class-level mutable default shared across instances, and refuses a transient shared by two kernels or read by control flow under per-thread maps instead of lifting it wrongly.
…ls, control flow reading a per-thread buffer
…ansient in place when a nest would misread its index
… where the accumulator lives
…im redundant tests enclosing_maps walks helpers.get_parent_map. prefix_control_flow hands its subscript prefixer to rewrite_code_slots, which also rewrites NestedSDFG symbol_mapping values, so a nest bound to an element of a lifted buffer now reads it with the kernel index (new test). Merge duplicate tile_extent, layout and non-transient accumulator tests.
…not when unparsing reformats
ThrudPrimrose
marked this pull request as ready for review
September 28, 2026 16:33
…lyses Both passes spent most of their time outside their dataflow analyses. Output is unchanged (identical SDFG JSON on all benchmarks): - ConstantPropagation: 17.3s -> 4.1s on a test SDFG with ~900 nested SDFGs and ~10k symbols, 53s -> 4.1s on a 4000-block CFG. - SymbolPropagation: 11.7s -> 0.5s on the same test SDFG, 57s -> 6.1s on a 4000-block CFG, and time no longer grows with loop nesting depth. Shared: - Add cfg.blockorder_reverse_postorder, a linear block order for fixed-point analyses, used by both passes instead of blockorder_topological_sort, whose dominator and branch analyses are quadratic in the length of a region. - Add Node.new_symbol_names so used_symbols no longer infers the types of map parameters only to discard them. - replace_properties_dict substitutes into symbolic properties only when a replaced symbol occurs, and _internal_replace filters replacements by name without invoking the SymPy printer. ConstantPropagation: - Process nested SDFGs before removing unused symbols from their parent and memoize their used symbols (nodes.memoize_nested_used_symbols), so each nested SDFG is traversed once instead of once per enclosing SDFG. - Skip empty descriptor replacements and memoize free-symbol and parsing queries on assignment values. SymbolPropagation: - Dead assignment elimination skips SDFGs without assignments, sweeps each SDFG to its own fixed point, and assembles the symbols read from blocks, meta code and edges instead of recomputing free symbols per region. - Revisit a region after its contents so that its successors see final values in the same sweep. - _update_syms computes each free-symbol set once and skips regions whose meta code and outgoing edges read no candidate. - Memoize free-symbol and parsing queries on values, and precompute the view and opaque scalar names per SDFG. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
On a test SDFG with ~4.4k states, ~8.5k control flow regions and ~4.7k data containers, code generation ran out of memory (over 24 GB) in the reachability analysis; past that, it was on track to take ~50 minutes, dominated by per-array and per-tasklet rescans of the whole SDFG. It now takes 54 s with a peak of 1.4 GB, and loading takes 13 s instead of 25 s. Generated code and serialized SDFGs are unchanged on all tested SDFGs. - ControlFlowBlockReachability / StateReachability: return read-only set views backed by per-region bitset transitive closures instead of materialized reachable sets (quadratic memory in deeply nested control flow). - Allocation lifetime: compute per-SDFG control-flow symbols and access nodes once instead of per array. - InterstateEdge.new_symbols / LoopRegion.new_symbols: skip edges without assignments and only look up the data containers named in the expressions instead of building a dictionary of all of them. - Code generation resolves the symbols of each state once (SymbolResolver) and the struct types once per SDFG; tasklets without struct initializers skip StructInitializer. - ControlFlowRegion.cfg_id: cache the position in the CFG list. - Validation: check connector names against symbols/constants/arrays without enumerating all nested data names. - Deserialization: reset the CFG list once per SDFG instead of per added region; cache the version comparison of the symbolic deserializer. - Serialization: cache the default assumptions of symbols per type. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…nt list - The final clean-up of the frame code appended every kept line to one string (and truncated it to drop a goto followed by its label), which copies the code generated so far for each line: quadratic in the code size. Lines are now collected in a list and joined once, labels are looked up in the goto counter instead of the list of all gotos, and the regular expressions are compiled once. - The top-level interstate variable declarations reuse the argument list computed on construction (the one the function signature and the targets use) instead of traversing the whole SDFG again with ``SDFG.arglist()``. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- ``read_set |= rs.keys()`` does not update the set in place (``set.__ior__`` does not accept dictionary views, so ``dict_keys.__ror__`` builds a new set), which made collecting the data containers of all states quadratic in their number. The sets are now updated in place. - The state-level read/write analysis deep-copied its result, although ``SDFG.read_and_write_sets`` and ``SDFGState.read_and_write_sets`` only use the container names. Both now use the uncopied ``_read_and_write_subsets``; ``_read_and_write_sets`` still returns a copy. Both are used by ``SDFG.used_symbols``, which code generation calls on the whole SDFG. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- ``blockorder_topological_sort``: in a region where no block has more than one outgoing edge, the traversal only follows single edges and never consults the parent tree, branch merges or loop exits, so their computation (immediate dominators, transitive closure of the dominator tree, dominance frontiers) is skipped. - ``block_parent_tree``: branch merges are only computed when loops are respected (the only place they are used), and reuse the dominators computed there instead of computing them a second time. - ``ControlFlowRaising``: regions with a single sink and no branches have no conditionals to lift, so the dominator analysis is skipped for them (the start block is still queried, as the analysis would). Code generation calls these once per control flow region. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
networkx 3.6+ omits the start block from immediate_dominators, so block_parent_tree and blockorder_topological_sort raised KeyError on the entry block during codegen. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…scriptor and the slice index
It is get_parent_map repeated to the outermost scope, so it belongs next to it rather than in the pass.
Remove the yapf pre-commit hook, .style.yapf and the yapf pin in the linting group.
Add the ruff-format hook next to ruff-check. line-length and quote-style carry over the yapf settings; notebooks stay unformatted as before.
ruff honours fmt: off/on at statement level only, so the parametrize decorator uses a trailing fmt: skip.
# Conflicts: # dace/autodiff/library/library.py # dace/codegen/codeobject.py # dace/codegen/targets/cpu.py # dace/data/core.py # dace/data/distributed.py # dace/frontend/python/replacements/array_manipulation.py # dace/frontend/python/replacements/filtering.py # dace/libraries/blas/blas_helpers.py # dace/libraries/blas/nodes/batched_matmul.py # dace/libraries/blas/nodes/dot.py # dace/libraries/blas/nodes/einsum.py # dace/libraries/blas/nodes/gemm.py # dace/libraries/blas/nodes/gemv.py # dace/libraries/blas/nodes/ger.py # dace/libraries/blas/nodes/matmul.py # dace/libraries/fft/nodes/fft.py # dace/libraries/linalg/nodes/tensordot.py # dace/libraries/mpi/nodes/allreduce.py # dace/libraries/mpi/nodes/redistribute.py # dace/libraries/mpi/nodes/reduce.py # dace/libraries/onnx/nodes/onnx_op.py # dace/libraries/onnx/nodes/onnx_op_registry.py # dace/libraries/onnx/schema.py # dace/libraries/pblas/nodes/pgemv.py # dace/libraries/sparse/nodes/csrmm.py # dace/libraries/sparse/nodes/csrmv.py # dace/libraries/standard/helper.py # dace/libraries/standard/nodes/copy/node.py # dace/libraries/standard/nodes/fill/node.py # dace/libraries/standard/nodes/reduce.py # dace/libraries/stencil/stencil.py # dace/memlet.py # dace/properties.py # dace/sdfg/nodes.py # dace/sdfg/sdfg.py # dace/sdfg/state.py # dace/transformation/dataflow/copy_to_device.py # dace/transformation/dataflow/copy_to_map.py # dace/transformation/dataflow/local_storage.py # dace/transformation/dataflow/map_expansion.py # dace/transformation/dataflow/mapreduce.py # dace/transformation/dataflow/reduce_expansion.py # dace/transformation/dataflow/stream_transient.py # dace/transformation/dataflow/streaming_memory.py # dace/transformation/dataflow/strip_mining.py # dace/transformation/dataflow/tasklet_fusion.py # dace/transformation/dataflow/tiling.py # dace/transformation/dataflow/vectorization.py # dace/transformation/dataflow/warp_tiling.py # dace/transformation/estimator/enumeration/enumerator.py # dace/transformation/estimator/enumeration/greedy_enumerator.py # dace/transformation/estimator/enumeration/scoring_enumerator.py # dace/transformation/helpers.py # dace/transformation/interstate/gpu_transform_sdfg.py # dace/transformation/interstate/loop_peeling.py # dace/transformation/interstate/loop_unroll.py # dace/transformation/interstate/sdfg_nesting.py # dace/transformation/pass_pipeline.py # dace/transformation/passes/analysis/analysis.py # dace/transformation/passes/constant_propagation.py # dace/transformation/passes/dead_dataflow_elimination.py # dace/transformation/passes/full_map_fusion.py # dace/transformation/passes/fusion_inline.py # dace/transformation/passes/length_one_array_scalar_conversion.py # dace/transformation/passes/loop_local_memory_reduction.py # dace/transformation/passes/pattern_matching.py # dace/transformation/passes/prune_symbols.py # dace/transformation/passes/scalar_promotion.py # dace/transformation/passes/scalar_to_symbol.py # dace/transformation/passes/simplification/control_flow_raising.py # dace/transformation/passes/simplify.py # dace/transformation/subgraph/composite.py # dace/transformation/subgraph/expansion.py # dace/transformation/subgraph/stencil_tiling.py # dace/transformation/subgraph/subgraph_fusion.py # dace/transformation/transformation.py # tests/codegen/aligned_allocation_test.py # tests/codegen/allocation_after_size_assignment_test.py # tests/codegen/allocation_lifetime_test.py # tests/codegen/cpp_test.py # tests/codegen/external_memory_test.py # tests/library/copy_node_test.py # tests/numpy/searching_test.py # tests/passes/symbol_propagation_test.py # tests/transformations/nest_subgraph_test.py
…t_parent_map per level
ThrudPrimrose
added a commit
that referenced
this pull request
Oct 13, 2026
…NestedGPUDeviceMapLowering (#2534) MoveArrayOutOfKernel keeps constant-size per-thread arrays of up to register_limit (16) elements in registers and moves larger and symbolically-sized ones out of the kernel; the GPU codegen preprocess pipeline keeps calling both passes. Branch code that iterated states() as a list or called all_states() follows main's generator states().
ThrudPrimrose
added a commit
that referenced
this pull request
Oct 13, 2026
…NestedGPUDeviceMapLowering (#2534) MoveArrayOutOfKernel keeps constant-size per-thread arrays of up to register_limit (16) elements in registers and moves larger and symbolically-sized ones out of the kernel; the GPU codegen preprocess pipeline keeps calling both passes.
… up to register_limit and moves the rest Binds the kernel parameters before reshaping the nests below the owner.
ThrudPrimrose
force-pushed
the
gpu-prep/hoist-kernel-transients
branch
from
October 13, 2026 16:37
747ace5 to
568468f
Compare
This branch has not been deployed
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
2 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.
Hoists a transient
GPU_Globalarray declared inside aGPU_Devicemap out of the kernel, giving it one slice per map iteration, and demotes one small enough to per-thread registers instead; GPUs cannot allocate device-global memory inside a kernel, so this exists for backward compatibility. The lift builds its memlet paths withadd_memlet_pathand finds a free name withutils.find_new_namerather than open-coding either.