Repository navigation
Avoid dynamic dispatch in back-end IR code - #1071
Merged
Merged
Conversation
The arguments were collected in a `Vector{Any}`, so every use of their
fields in the back-ends (Metal's argument metadata and parameter
attributes, `irgen`'s by-value rewriting) dispatched dynamically, and
so did the LLVM.jl and Base methods called with them. Many of those
are only inferred in an upstream package image, without native code,
and a back-end's own image can't add native code to them
(JuliaLang/julia#58465), so they were compiled in every session.
Splatting a collection of LLVM values or types into a vector literal dispatches on the run-time types of the elements, compiling a method for every combination of instruction and type kinds encountered. Build these vectors with `collect`, `pushfirst!` or by concatenating concretely typed collections instead. Also assert the types of the kernel state getter call and of the caller instruction `backtrace` moves to, which kept every property access in its loop dynamic.
The comprehension collecting the atomics widened its element type at
run time, so it was inferred as an untyped vector and every helper
called on its elements was dispatched dynamically. Collect into a
`Vector{LLVM.Instruction}`, don't specialize the helpers on the kind of
instruction or pointer (which compiled each of them again for every
kind), and give the other vector literals built from IR objects an
element type.
As for PTX, collect the atomics to lower into a typed vector, don't specialize the helpers on the kind of instruction or pointer, and build the argument and parameter-type vectors of the AIR atomic intrinsics by appending to typed literals rather than by splatting into untyped ones. Give the table of kernel intrinsics that become arguments concrete key and value types, which made all of `add_input_arguments!` dynamic. When emitting argument metadata, assert the types of sizes and address spaces computed from untyped values, so that the `Int32` and `ConstantInt` calls on them are resolved statically.
Like for PTX and Metal, the untyped comprehension made every call on the collected instructions dispatch dynamically.
Their callers only know they have an `LLVM.Instruction` or `LLVM.Value`, so every call dispatched dynamically and compiled the helper again for each concrete kind of instruction or pointer it met. With `@nospecialize`, a single method instance serves all of them and the calls resolve statically.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1071 +/- ##
==========================================
- Coverage 92.49% 92.46% -0.03%
==========================================
Files 31 31
Lines 7313 7315 +2
==========================================
Hits 6764 6764
- Misses 549 551 +2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
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
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.
Compiling the first kernel of a session JIT-compiles a lot of LLVM.jl, GPUCompiler and Base host code, much of it because GPUCompiler's back-end code calls it through dynamic dispatch: values come out of untyped containers (
classify_arguments'Vector{Any}, Metal's untypedkernel_intrinsicsdict, comprehensions whose element type widens at run time), IR objects get splatted into vector literals, and the atomics helpers are specialized on every instruction kind they see. Many of the callees are only inferred in an upstream package image, without native code, and a back-end's own image can't add native code to those (JuliaLang/julia#58465), so this happens in every session.These commits make those call sites type-stable, or stop the specialization where it doesn't buy anything: a typed vector from
classify_arguments, vectors of IR objects built withcollect/pushfirst!/concatenation instead of splatting, typed atomics worklists for PTX, Metal and SPIR-V,@nospecializeon the local-atomics and shared atomics helpers, and a few type assertions in Metal's argument metadata. Behavior is unchanged.Together with #1070, on Julia 1.13 the first new kernel takes 939 instead of 1470 ms with CUDA.jl on an RTX 5080 (JIT-compiled methods 156 → 124), and 233 instead of 270 ms with Metal.jl on an M1 (56 → 47).
What remains is mostly in LLVM.jl (accessors specialized on the concrete instruction type,
first_usedispatched insideiterate(::ValueUseSet)) and in the back-ends' own packages.