[WIP] New CPU codegen, new GPU codegen and GPU offloading pass (combined) - #2620
Draft
ThrudPrimrose wants to merge 1017 commits into
Draft
ThrudPrimrose wants to merge 1017 commits into
ThrudPrimrose wants to merge 1017 commits into
Conversation
…L_models group The exported graph needs a pure And expansion to lower; eager attention keeps the SDPA mask guards out of the graph.
…k factory transformers builds the mask from inputs_embeds.shape[1], which the TorchScript tracer turns into a 0-dim tensor. A 4D mask is returned as-is, so the tracer never reaches that code.
Deleting dace/transformation/interstate/gpu_transform_sdfg.py left about a hundred references to it behind, so the branch did not import: SDFG.apply_gpu_ transformations still imported the deleted module (through a `dace.dace.` path that would not have resolved either), and GPUTransformMap, the sparse csrmm / csrmv expansions, eight test modules, a codegen sample and three doc pages named it directly. apply_gpu_transformations now runs OffloadToAccelerator and honours its own simplify / validate contract. sequential_innermaps, register_transients, permissive, host_maps and host_data were that transformation's parameters and go with it -- the pass derives what stays on the host rather than taking it as an input. Callers that passed them are updated. tests/host_map_host_data_test.py tested nothing but the two removed host parameters and goes with them. gpu_transform_test.py loses the two cases that asserted the old transformation's own output shape rather than a general property: its toplevel_trans lifetimes, and its copy elision for a fully-written array. The doc entry moves off the interstate page, which no longer holds that module, and onto the passes page under its real module path.
Five separate reasons the pass could not have run as merged. The package had no __init__.py, at either level, so neither the module nor a wheel built from it was importable. Phase 1 and phase 6 each built an f-string with a nested same-quote expression, which is PEP 701 and therefore Python 3.12 and up; the project declares >= 3.10, where both are syntax errors. Phase 3 imported the length-one array <-> scalar conversion from a private copy of it under passes/vectorization/ -- a directory that was also not a package -- rather than the version merged to main in #2532; the copy is deleted and the merged module, a superset of it, is used instead. Every set in the package is now an OrderedSet. The one the review asked about drives which partitions phase 4 wraps, and the rest reach node insertion the same way, so a plain set makes codegen depend on hash order. The tests were named *_unittests.py, which pytest.ini's python_files does not match, so CI collected none of the fifteen. They also carried gpu_offload and current markers that were never registered -- pytest_configure only takes effect in a conftest -- and five of them called sdfg.view(), which opens a browser. They are renamed, marked gpu, and no longer draw graphs. MAX_ITERATIONS and VERBOSE were mutable class attributes; they are now max_iterations and verbose Properties on the Pass, with the bound raised well past what any SDFG needs and a docstring for why the loop terminates. Also dropped: an unused GPU_CPU_Unified enum pair, a committed npbench.db, a stray whitespace-only edit to validation.py, the debug edits to auto_optimize.py that had disabled expand_library_nodes, and a dead match arm that named a CLOSE node "open".
Renaming the test file so pytest collects it turned up two failures and one crash, none of which had ever run in CI. Both failures were one bug in the single-element copy optimization. Moving an access node through a map entry leaves ``other_subset`` describing a node that is no longer on the edge: validation then reads it against the surviving node's descriptor and reports a dimension mismatch, or -- when that node is a tasklet, which has no descriptor -- crashes reaching for ``.data`` on it. The subset describes the second container of a container-to-container copy, so it is dropped wherever the rewire puts a scope node or a tasklet on one end. The crash was a nested SDFG node in a state. The copy analysis dispatched on ControlFlowRegion and had no case for one, so it raised by name. A nested SDFG is a graph of its own and is now classified by its contents: one holding no GPU-scheduled node runs on the host whole, which is what a sequential scan in a loop region is. Classifying it by the state around it made the state hybrid and handed its body to the size-1 wrapper, which is one kernel launch per scan step. The size-1 wrapper also carried three defects fixed on the extended branch. An empty memlet orders and carries no data, so a connector on it is invalid and infer_connector_types raised KeyError; empty boundary edges now cross the wrapper connector-free, and the connector numbering follows a list rather than a set. Ordering edges for dangling roots cannot be the else-branch of the rewiring, since rewiring one boundary edge does not make the other roots any less dangling. And a partition is a dataflow component while a map scope spans one, so a partition holding half a scope is closed under its scopes, or left alone where closing would reach across a kernel. Three gaps behind the schedule assignment go with them: a Stream has no subset to collect and is invisible to the analysis rather than an error, an expansion now declares whether it runs inside a kernel, and a transient live only inside one kernel becomes a Register instead of a host allocation the dispatcher would have to answer with an illegal copy. Schedule assignment is now a single recursive walk: GPU_Device at a host level, Sequential below one, and a nested SDFG at a host level is a host level of its own.
The offloading makes every top-level map a kernel. That is wrong for a map whose purpose is to LAUNCH work rather than do it -- ICON's shape, an nblks map over nproma/nlev bodies -- and it cannot work at all for a map holding a callback. host_maps takes None to name nothing, True to derive them structurally, or a list of the maps themselves, each given as a label or as the MapEntry node. It is a constructor argument rather than a Property: a MapEntry cannot round-trip through JSON, and a caller passing node objects is driving the pass in process anyway. A named map is final except where the lowering could not be emitted -- an inner extent naming the outer map's own parameter reaches the launch configuration, where that parameter is not in scope. Automatic detection is structural only: a scope that launches rather than computes. Deciding a map by comparing how many threads each lowering would launch is a heuristic and is deliberately left out. Callbacks are the case that needs no naming at all. Neither kind can be offloaded -- a Python callback needs the interpreter, and a GPU callback is itself a launch, so a kernel cannot issue one -- so a map around one stays on the host whatever host_maps says, and a callback bounds a partition rather than being swept into a size-1 map. This derives what host_data=['__pystate'] and exclude_tasklets were pinning by hand against the transformation this pass replaced, which is why neither parameter comes back. The same rule already covered a map around a device-wide library node: a cuBLAS call is issued by host code, so that answer does not wait for automatic detection either. Nine tests, on the ICON zekinh gather kernel the shape is named for. The control asserts the outer map IS offloaded when nothing is named, so the pinning tests cannot pass by asserting the default; the negative control asserts a parent map that does its own work is not auto-detected, so detection cannot pass by calling every outer map a host map. A frontend-generated callback and a hand-built map around one cover both routes, and the scan case pins that a loop-carried accumulation gains no map at all.
… with dropout torch.onnx.export restores the module's original training flag when it finishes. A freshly constructed wrapper defaults to training mode, so the restore switched the wrapped BertModel back on and the reference outputs, computed after the export, were sampled with dropout active.
…entions
Three things, all in the offloading package.
A map enclosing a device-wide library node was kept on the host unconditionally,
on the grounds that a vendor call is issued by host code. That is an opinion
about scheduling, and applying it whatever the caller asked for moved the kernel
one level inwards wherever a library node sat under a map: npbench spmv came out
with its outer map sequential and the gather map inside a nested SDFG promoted to
a kernel, which then read host memory ('Illegal copy! (from x to indirection)'),
and trmm failed to compile. Both pass again with the rule gone. host_maps is now
what it says: a list the caller gives, automatic structural detection, or
nothing. The ExpandTransformation.runs_inside_kernel flag existed only to serve
that rule and goes with it, which leaves dace/transformation/transformation.py
and dace/libraries/standard/nodes/reduce.py identical to main again.
A program returning several values names them __return_0, __return_1, ... The
device twin tested for equality with "__return", so only the single-return case
was moved out of the reserved namespace; the rest became __return_0_gpu, which
still starts with __return and the runtime refuses as a transient. The host twin
next to it already used startswith -- the asymmetry was the bug. bicg,
gramschmidt and ludcmp pass again.
The rest is convention. The phase modules were named offl_phaseN_<what>, an
abbreviation and a numbering with gaps (no 5, a 7b) in a directory called
offload_to_accelerator_phases, where every sibling package under passes/ uses
plain descriptive module names; they are now phases/<what>.py, and the ordering
stays where it is load-bearing, in apply_pass. Every public function carries
parameter and return annotations. A map parameter becomes a SYMBOL, so
get_new_map_identifiers now checks the name against every namespace one can live
in -- symbol tables up the parent chain, descriptors, and every map in the SDFG
-- rather than one state's parameters, since reusing a name that is already a
symbol elsewhere gives one string two sets of assumptions. Extent comparison
reads names through symbolic rather than str() on sympy objects. Also gone: two
plain sets whose iteration order reaches codegen, a hasattr, two direct function
imports, and a duplicated copyright header.
tests/npbench runs these kernels end to end on the device, but every one of them is marked gpu, so on a machine without one the offloading gets no corpus coverage at all -- which is how two defects reached the branch and were only found by reading a failure on hardware. This runs the same pipeline those tests run, auto_optimize for the GPU, and stops before the compiler. validate and generate_code need neither a device nor a toolchain, take about a second a kernel, and still catch the whole placement family: a descriptor moved to the device that host code still reads fails validate with the container named, and a name minted into a namespace the runtime reserves fails at code generation. Running the whole pipeline rather than the pass alone is the point. auto_optimize picks library implementations before it offloads, and a raw parsed SDFG has none: offloading one directly leaves a host LAPACKE call under a GPU schedule with its status scalar in device memory, which is a graph the pipeline never produces. Asserting on it reported a library bug that cannot happen, and nearly bought a fix for it. Going through auto_optimize also means the cuBLAS and cuSolverDn lowerings are the ones under test, rather than the reference expansions. A kernel whose own test_gpu is skipped is skipped here too, with upstream's reason: mandelbrot2 is issue #1139 and lenet raises std::runtime_error, and reporting either here would be reporting somebody else's open bug as this pass's. Reading the marker off the module keeps the two in step when one is re-enabled. 33 kernels, 5 skipped.
A view owns no storage: it is an alias, and it belongs wherever the container it aliases ends up. The pass did not treat it that way, and the two halves of the mistake hid each other. ArrayView, ContainerView and ContainerArray all derive from Array, so is_array answered True for every one of them and the view branch sitting below it was unreachable. An alias was therefore placed as though it owned a buffer, which is how a view came to be renamed into a device twin of itself -- C_0 aliasing C_0_host aliasing C_0, a cycle the code generator follows until it runs out of stack. Meanwhile a view that was NOT placed kept whatever storage it was declared with, so npbench correlation handed cuBLAS three Default views of GPU_Global arrays and the call was emitted into the host translation unit, where the stream it wants does not exist. Both are fixed together, because fixing either alone breaks the other: is_array now recognises only real buffers, and a view takes its origin's storage as the LAST step of copy insertion -- only then does every container carry the storage the placement gave it, and doing it earlier copies a host storage that has not been updated yet. A view inside a kernel keeps Register, which the code generator requires, and a view chain that reaches no access node is left alone instead of recursing on None and asking the state for the edges of a node it does not hold. Structures and container arrays have no single buffer to place -- one is a record of descriptors, the other an array of them -- so they are skipped by name like a Stream rather than falling through to the raise. A length-one view is excluded from the scalar conversion for the same reason ConvertLengthOneArraysToScalars exempts one: a Scalar cannot carry the views alias edge.
…pclstorage host
OffloadToAccelerator takes host_maps -- the maps that keep a host schedule, so that the maps under them become the kernels -- but apply_gpu_transformations did not forward it, and that method is how everything outside this package offloads. A caller could not reach the feature at all, which is the same gap the old transformation's own host_maps parameter used to cover. The parameter defaults to naming nothing, so nothing changes for a caller that does not pass it. Two tests: one asserts the method and the pass reach the same schedules for the same named map, so the plumbing cannot rot without saying so, and one asserts the default still offloads the outer map, so the first cannot pass by asserting the behaviour the pass produces anyway.
host_maps defaulted to None, which read as "unset" rather than as an answer. The default is now False -- run no host-map detection at all -- and None and [] mean the same thing, so a caller that computes the list and finds it empty gets "name none" instead of falling into a different branch. True is what asks for the built-in heuristics, and a list names the maps outright. Nothing changes for a caller that passes nothing: all three spellings already produced no host maps, and the default is the same answer under a name that says so. A test pins the four spellings against each other, so the day one of them drifts into meaning something else it says so. A map holding a callback is still kept on the host whatever this says: a Python callback needs the interpreter and a GPU callback is itself a launch, so a kernel cannot issue either, and offloading one produces code that cannot run rather than code that runs slower.
Two ways an SDFG stopped surviving a serialization round trip, both found by running the GPU pipeline over the npbench kernels with `testing.serialization` on. A memlet parsed from a string mints its symbols with the default dtype and no assumptions, while a map parameter carries the frontend's -- int64 and nonnegative. SymPy compares assumptions, so the two `k` atoms stay distinct and offsetting one subset by the other leaves `k - k` in the range instead of `0`. A symbol in DaCe is its name, so both denote one value; `Range.offset` and `offset_new` now collapse such a pair, guarded by a name count so the substitution only runs when a name really is duplicated (vadv). The GPUAuto reduce expansion reshaped its `_in` descriptor by assigning `shape` alone. `offset` is rank-dependent, so it stayed at the source array's rank and the descriptor failed its own `validate` with "Offset must be the same size as shape"; deserialization then answered with an unregistered placeholder instead of an Array (softmax). It goes through `set_shape`, which recomputes what the rank decides.
`compiler.cuda.max_concurrent_streams` went back to 0. Changing the shipped default to -1
puts every SDFG on the default stream, offloaded by this pass or not, which is not this
pass's call to make: it dropped the cross-stream event synchronization and the host wait
before a device-to-host copy's reader, taking four codegen tests with it.
Copy placement only propagated forwards, so an array first used on the device late in the
program was copied there between two device states -- a host state in the middle of a run
of kernels, which is what GPUPersistentKernel then had to swallow and could not
("MappedTasklet expansion cannot cross the CPU/GPU boundary"). A state that never touches
the array cannot care which side it is on, so the copy now moves above it, and only when
every successor wants the array on the device, which leaves a copy inside the one branch
that needs it.
`OffloadToAccelerator(verbose=True)` raised a TypeError instead of setting the property;
both Properties are now named parameters, as the other passes spell it.
The offset corrections in vertical map fusion bring together two mints of one name -- a memlet parsed from a string carries the default dtype and no assumptions, a map parameter carries the frontend's -- and SymPy compares assumptions, so the subtraction leaves `k - k` standing where it means `0`. That reaches code generation and stops the SDFG from surviving a serialization round trip, because only one of the two atoms is written with its dtype. Folding it in `Range.offset` fixed it for every caller at once, but that is a change to the subset arithmetic every transformation shares, and it does not belong to this PR. The fold now runs where the two mints actually meet. The general question -- how a name comes to carry two different symbols at all -- is left to its own change.
… copy The fold now lives in `Range.offset` and `Range.compose`, where the two mints of one name actually meet, so map fusion needs no copy of it and `compose` is covered too -- the other place a composed bound kept `i + (M - i - 1)`. This is the same change as the bugfix PR against main; it sits here so the corpus test is green before that lands.
…te names The walk moves into PowerRelaxer, the sign facts of an SDFG become a frozen SignFacts that is passed down instead of saved and restored on the pass, and the provers are public functions that the tests call directly. Behaviour is unchanged.
# Conflicts: # dace/transformation/passes/gpu_specialization/split_state_by_gpu_class.py # tests/gpu_specialization/auto_single_stream_test.py # tests/gpu_specialization/gpu_stream_test.py # tests/gpu_specialization/monolithic_single_stream_test.py # tests/gpu_specialization/npbench_gpu_correctness_test.py
A size held in a scalar (nt = Nt + 1; np.empty(nt)) is a data descriptor, but an extent has to be a symbol. The frontend now reads each such scalar into its own __sym_ symbol on an interstate edge and substitutes it into the shape, leaving the descriptor in place so the program can still read or reassign it. A fresh symbol per shape keeps two arrays sized from one reassigned name from collapsing onto a single value, and keeps a later index on the same name from rebinding the extent. A size-1 array is read through a subscript, since assigning it by name would take its pointer. np.zeros, np.ones and np.full go through the same promotion as np.empty and dace.define_local, so they no longer reject a computed size.
A transient sized by a symbol that an interstate edge assigns is allocated at the dominator of its accesses. When that dominator only dominates them and its only outgoing edge carries the assignment and enters a control flow region, the allocation was emitted at the start of the dominator, before the symbol was assigned, and the array got an undefined size. The block the edge leads to now allocates the data, and a control flow region allocates on entry what it is the scope of. The deallocation at a postdominating region was never emitted, so the array leaked. A region now deallocates on exit what it is the scope of.
…claration determine_allocation_lifetime only saw access nodes, so a transient named as a free name in a tasklet's code, with no connector, memlet or access node, did not count as a use. Its declaration then landed in the brace scope of the one state with a real write, and the sibling state reading it from code did not compile. Such uses now count, through a stand-in access node that is never made for a view, whose allocation reads the edges of its node.
…mbined branch The merge of the slimmed GPU codegen branch dropped both, and the offloaded cholesky2 kernel then wrote device-resident _info from the host.
…perator new, not an alignment attribute
new double DACE_ALIGN(64)[N] names an over-aligned element type, which GCC 16 rejects for a constant bound
("alignment of array elements is greater than element size"). Defer to the base heap_alloc_stmt/heap_free_stmt,
which emit new (std::align_val_t(n)) T[count] at the descriptor's alignment and the matching aligned delete, and
fuse constant-extent allocations like runtime ones.
(cherry picked from commit 89266363be90b413e79e86067c40fc77e1e748d6)
… and leave every other power to std::pow
The tile library nodes of dace.libraries.tileops, their per-ISA lowerings and runtime headers, and the dace.tile.add and dace.tile.masked_copy calls of a @dace.program, with the tests of the nodes and the calls. The vectorization passes that emit the nodes are not part of this commit. The nodes need the ONE broadcast marker and shapes_equal of dace.symbolic, DACE_UNROLL and c_mod of the runtime headers, and the library registration of a node that subclasses a node of another library.
…he way to its accesses
# Conflicts: # dace/libraries/standard/nodes/reduce.py
…nt' into ngc-2620 # Conflicts: # dace/codegen/targets/framecode.py
# Conflicts: # dace/codegen/codegen.py
# Conflicts: # dace/codegen/targets/cpp.py # dace/codegen/targets/framecode.py # dace/symbolic.py
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
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.
Combines the readable CPU code generator, the new GPU code generator (#2259) and the GPU offloading pass on top of main, so all three run through CI together.
It also replaces the const-init marking with a small pass that promotes literal-only transients to SDFG constants.
🤖 Generated with Claude Code