OffloadToAcceleratorPass - #2553
Open
FinnHeckman wants to merge 88 commits into
Open
Conversation
…rained information and to integrate with dtypes better
…tcases; added some tests too
wrote OtA.get_data_locations(sdfg) which finds all accessed array and assigns them to gpu_set, cpu_set or both the function currently passes all tests
…o that it now passes the test
… detection & input-output checks for all maps in case arrays aren't used within
…l exploratory but promising
…py instaed of just one 2) clean code, insert IR logic into OtA class
- add test suite & populate with tests - get OtA to work with line graphs and loop regions -> bug fixes - improve printing for IR nodes Rn copies are inserted eagerly. This leads to too many copies and doesn't work for branches at all. Next: refactor & make insertion lazy
…s late as possible" - the structure simply wasn't designed for this additional complexity. Yakup suggested doing "as late as possible" and then going over it with an optimization pass. After spending a full night in the weeds, I agree. I'll be rolling back and starting over. Yay!
… from "as late as possible" - the structure simply wasn't designed for this additional complexity. Yakup suggested doing "as late as possible" and then going over it with an optimization pass. After spending a full night in the weeds, I agree. I'll be rolling back and starting over. Yay!" This reverts commit df0bea8. role back & start over
…n individually 2) get branches to work somewhat - still working on a non-determinism bug
…y using join nodes for both loops and brnaches, looks good
…nnot rename the arrays within the nested sdfg
…f skipping container states, might be in for a larger refactor, so this is my return point
…s 100 lines shorter, more powerful and more readable (both code and output) than before. One gotcha: at the moment, everything assumes linegraphs and non-conditional branching as well as edge conditions aren't handled atm
…h later pass?) and the eval function looks extremely ugly rn, but I think I refactor it with a pair of fresh eyes tommorrow morning
…x issues. Biggest issue is my handling of library nodes atm
- I changed the naming scheme to respect dace's __return convention - I added a catch to the state analysis, checking whether gpu & cpu sets overlap - if so, it first tries to set all LibraryNodes to CPU - if that doesn't work, it sets all maps to CPU too (could this actually be one step?) - it then raises an error and continues to treat this state as fully sequential (further design necessary, see whether something can be salvaged here) Almost all large polybench graphs have sections that can't be offloaded due to mixed map/libnode & tasklet use -> ask Yakup how to proceed
…dded runtime comparison & plotting: polts look terrible, offloaded graphs take way longer :O
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 spcl#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.
…ad spclstorage 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.
ThrudPrimrose
self-requested a review
September 7, 2026 10:16
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.
The final close node put every non-transient back on the host, including ones the SDFG never writes. That copy-back writes the container, which validation refuses as soon as the SDFG is nested and the container reaches it through an input connector only -- the graph GPUTransformMap builds around a map that reads one array and writes another, and the laplace transformation sweep's two failures.
Restoring a container's home is what the previous commit removed from the IR, which left a read-only array claiming to end on the device. Put the close node back at the starting position and refuse the copy itself: a container never moves, a twin is staged beside it, and its home copy goes stale only once something writes the twin. Filling the twin still stands in both directions, so an array whose home is the device keeps the host copy that makes it readable.
…ller's A clone carries the caller's descriptor class, so a reduce reading an ArrayView -- lenet's second one reads a view of the maxpool input -- gave its nested SDFG an "_in" that views nothing and is read by several edges, which validation refuses. It is now built the way "_out" already was.
Placement decides for the container and treats an access through a view as an access to it, but the alias was left naming the other side: a descriptor carries one storage, so a view read from a kernel in one state and from host code in another cannot serve both, and the dispatcher answers the losing side with "Illegal copy!". Staging C now stages C_0 with it, the twin viewing the twin.
A view owns no storage, so what its descriptor declares is right only where some other pass has already matched it to the container it aliases. Read directly, a view of a GPU_Global array still declaring Default sent a reduction that belongs on the device down the Pure fallback, leaving it on the wrong side of the machine rather than failing outright.
The corpus skipped the five kernels whose own test_gpu upstream disabled, but those marks record a runtime verdict and this file only claims a kernel lowers. Following them cost five kernels and hid two bugs, both now fixed; running them turned up two more in the view staging itself, where a chain reaches a name registered later or a node that is not an access node at all.
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.
The OffloadToAcceleratorPass is a new GPU-offloading pass to replace GPUTransformSDFG. It yields more performant benchmarks, extends previous capabilities and improves maintainability. It was tested on the NPBench testsuite. The pass has been developed as the main contribution of the Bachelor Thesis "Principled Accelerator Offloading in DaCe" by Finn R. Heckman, supervised by Yakup K. Budandaz.