Internalize non-kernel deferred codegen entrypoints - #928
Conversation
Deferred codegen entrypoints were kept externally visible across the InternalizePass, so that linking could resolve them. Once linked (and, for Enzyme's wrappers, alwaysinlined) they are dead, but their external linkage kept GlobalDCEPass from dropping them. Back-ends like SPIR-V then still had to translate functions they cannot express. Only deferred jobs that are kernels themselves (e.g. child kernels for dynamic parallelism) are entrypoints in their own right. Classify the entrypoints once after deferred codegen, and use that set both for internalization and for the per-entrypoint `finish_ir!` loop. Other deferred functions are internalized like any other function, and dropped by the existing clean-up passes once inlined. The mock Enzyme test helper gains an opt-in `always_inline` flag that marks the generated function `alwaysinline`, like Enzyme's wrappers, so the removal can be tested on the native, PTX and SPIR-V targets. Assisted-by: Claude Code (Fable 5.1)
|
Ugh, this is becoming one hell of an ugly interface. |
| # wrappers Enzyme generates) are only called from within this module, so they should | ||
| # be internalized like any other function and dropped once inlined, rather than kept | ||
| # around for back-ends that cannot express their signatures (e.g. SPIR-V). | ||
| entrypoints = filter(((job′, _),) -> job′ === job || job′.config.kernel, jobs) |
There was a problem hiding this comment.
jobs is being used for too many things here. It was only intended to be for deduplicating compilation, but now it's also for preserving externally visible symbols, and scheduling finalization... That makes this code really hard to read.
I'm fine with the bugfix, but we should really replace this with separate mechanisms at some point.
|
Not disagreeing, and I am in a catch-22 of my own making. KA 0.10 should keep Enzyme support... I would like to re-engineer the GPU integration that Enzyme currently has. Maybe Claude has more follow-through than I do on #582 |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #928 +/- ##
==========================================
- Coverage 85.76% 81.42% -4.34%
==========================================
Files 29 29
Lines 5598 5895 +297
==========================================
- Hits 4801 4800 -1
- Misses 797 1095 +298 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
2 participants
Motivated by JuliaGPU/KernelAbstractions.jl#736.
Problem
Deferred codegen entrypoints (notably the wrappers Enzyme generates) were held externally live across GPUCompiler's
InternalizePassso that linking could resolve them. Once linked andalwaysinlined they are dead, but external linkage keptGlobalDCEPassfrom dropping them, so e.g. the SPIR-V back-end still had to translate functions it cannot express.Change
Only deferred jobs that are kernels themselves (e.g. child kernels for CUDA dynamic parallelism, which inherit the parent's config) are entrypoints in their own right. The driver now classifies the entrypoints once after deferred codegen, and uses that set both for internalization and for the per-entrypoint
finish_ir!loop. Other deferred functions (Enzyme registers its jobs withkernel=false) are internalized like any other function, and dropped by the existing clean-upGlobalDCEPassonce inlined. No back-end hooks are touched, so thefinish_ir!override in the KernelAbstractions PR should no longer be necessary.Tests
always_inlineflag onEnzyme.deferred_codegen, marking the generated functionalwaysinlinelike Enzyme's wrappers.define internal, and that it is gone entirely whenalwaysinline. The PTX test previously compiled for the native target; it now compiles an actual PTX kernel.alwaysinlinechild is absent from the optimized module, and the translated SPIR-V contains exactly oneOpFunction. Without the driver change this fails on both back-ends, as the external definition is left behind.Note that the reflection
code_nativesetsonly_entryunlessdump_module=true, which silently skips deferred codegen; the SPIR-V test passesdump_module=truefor that reason.Tested locally with Julia 1.10 and 1.12 (
native,ptx,spirv,utils,gcn,metal).🤖 Generated with Claude Code
https://claude.ai/code/session_01Jap1wwekUHPKo3iET2dKhn