Defer GPU shared/heap allocation fusing to the offload site#9209
Open
abadams wants to merge 3 commits into
Open
Defer GPU shared/heap allocation fusing to the offload site#9209abadams wants to merge 3 commits into
abadams wants to merge 3 commits into
Conversation
fuse_gpu_thread_loops still hoists shared, heap, and register allocations and normalizes GPU loop structure, and still computes the lifetime/offset fusion analysis where the barrier structure is intact. But instead of erasing the per-Func allocations and rewriting their loads and stores onto the fused buffer, it now emits a backing allocation plus one aliasing Allocate per Func, whose new_expr is a new offset_pointer(backing, offset) intrinsic. Loads and stores keep their original per-Func names. A new flatten pass in inject_gpu_offload folds each aliasing allocation's offset back into its loads and stores and drops the aliasing nodes, reproducing the flat representation the device backends expect. It runs after the conceptual stmt is captured, so the profiler and the conceptual stmt see per-Func allocation names and stores can be attributed to a specific Func. Supporting fixes: - offset_pointer is a non-pure Intrinsic so CSE/LICM won't lift it out of new_expr and escape the backing allocation's scope. - RemoveDeadAllocations traverses new_expr, so a backing allocation referenced only through aliasing new_exprs isn't dropped as dead. - PartitionLoops no longer moves a let inside an Allocate when the let is used in the allocation's new_expr or condition (evaluated in the outer scope), which previously left a dangling reference for dynamically-sized allocations. - The flatten pass re-derives load/store alignment from the offset, rather than keeping the original alignment across a possibly misaligned shift. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
LowerWarpShuffles treated any non-shared allocation above a lane loop as per-lane register storage, striping its size by the warp size and moving the Allocate inside the lane loop. With GPU allocation fusing now emitting per-Func aliasing allocations, a heap allocation used inside a gpu_lanes loop (e.g. apps/iir_blur) reaches this pass as a real Allocate node and was wrongly striped, which also separated it from its Free and left a dangling reference. Heap is global memory and is never warp-level storage, so skip it alongside GPUShared. On main no heap Allocate nodes reach this pass, so this is a no-op there. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #9209 +/- ##
=======================================
Coverage ? 69.81%
=======================================
Files ? 255
Lines ? 78593
Branches ? 18802
=======================================
Hits ? 54872
Misses ? 18136
Partials ? 5585 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
|
This is great. I think it's totally fine now to drop the complicated name of that alloc group. AFAICT, there will be only one. Just allocgroup0 for example would be fine. |
Member
Author
|
Oh good point, I don't need the monster name now. Will change. |
The per-Func names now live on the aliasing allocations, so the backing
allocation no longer needs a name concatenating every fused Func. Name
it unique_name("shared_alloc") or unique_name("global_alloc") instead of
the "allocgroup__f1__f2__..." monster name. The single-shared-allocation
case still names the backing after its Func, since there are no aliasing
wrappers there.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
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.
fuse_gpu_loops fuses and renames allocations halfway through lowering, making a confusing mess like this (from camera_pipe):
There is a single fused allocation for all the block-level funcs, and all stores and loads are rewritten to offset accesses to this single allocation. This is a problem for humans to read, but it also makes it hard for later lowering passes (the profiler change I'm working on) to attribute a store to a particular Func. This PR produces this instead:
Note the stores are now to f1.0, so we know which Func is being computed. The original distinct per-Func Allocate nodes remain, but they're just aliases of slices of the shared allocation. This IMO keeps things more understandable and will solve my profiler problem. In offload_gpu_loops the aliases are resolved and it's rewritten into the first form, so there's no impact on backends.