Skip to content

fix(mir): keep small runtime-indexed arrays in SSA#398

Draft
niklebedenko wants to merge 15 commits into
NVlabs:mainfrom
niklebedenko:upstreaming/fixed-array-ssa-final
Draft

fix(mir): keep small runtime-indexed arrays in SSA#398
niklebedenko wants to merge 15 commits into
NVlabs:mainfrom
niklebedenko:upstreaming/fixed-array-ssa-final

Conversation

@niklebedenko

Copy link
Copy Markdown
Contributor

Fixes #397.

Summary

  • lower runtime extraction from fixed arrays of up to 16 elements as constant extractvalue plus value selects
  • lower accesses to compiler-owned small local arrays as constant-address selections, including nested loads and stores
  • retain ordinary dynamic GEP lowering for larger arrays and borrowed/external pointers
  • add a bounded post-inline full-unroll pass (maximum trip count 16, threshold 512, no partial/runtime unrolling) before final O2 cleanup
  • extend the GPU array_for_loop regression with nested aggregate runtime reads/writes and PTX local-memory assertions

The local-pointer restriction is intentional: speculatively touching all candidate elements is safe for compiler-owned local storage, but would change observable access behavior for shared or externally borrowed arrays.

Validation

  • cargo test -p mir-lower -p mir-importer -p cuda-oxide-codegen
  • cargo clippy -p mir-lower -p mir-importer -p cuda-oxide-codegen -- -D warnings
  • cargo oxide run array_for_loop on RTX 3080
  • downstream Impulse cross-ghost prism/tet GPU regression: runtime checks pass and strict PTX audit reports zero local-memory/stack warnings across five kernels

@niklebedenko
niklebedenko marked this pull request as draft July 14, 2026 06:25
@nihalpasham
nihalpasham force-pushed the upstreaming/fixed-array-ssa-final branch from 58a2632 to 971e77d Compare July 22, 2026 11:30
@nihalpasham nihalpasham added perf Performance of generated code or of the compiler itself bug Something isn't working enhancement New feature or request codegen Device code-generation pipeline (Rust MIR to IR to PTX) IR-lowering Lowering between dialects (dialect-mir to LLVM dialect) build-related Build scripts, toolchain discovery, or compilation setup labels Jul 22, 2026
niklebedenko and others added 14 commits July 22, 2026 21:51
Signed-off-by: nik <nlebedenko@hotmail.com>
Signed-off-by: nik <nlebedenko@hotmail.com>
Signed-off-by: nik <nlebedenko@hotmail.com>
Signed-off-by: nik <nlebedenko@hotmail.com>
…he session's

The collector's dead-edge pruning evaluated RuntimeChecks switch
discriminants as check.value(tcx.sess) while the importer's reachability
walk and operand translation hardcode the device value (false). Only
release-like cargo-oxide profiles force -Cdebug-assertions=off, so a
dev-profile device build keeps session checks on: the collector would
prune a different SwitchInt edge than the importer translates, leaving
calls to functions that were never collected.

Introduce one shared constant, mir_importer::DEVICE_RUNTIME_CHECKS, and
make every consumer read it: the operand translator's constant-false
lowering, the importer reachability fold, the drop-glue no-op proof, and
the collector's edge pruning (including the block-local lookback, which
previously did not fold a RuntimeChecks assignment source at all while
the importer's did).

The other folding class stays aligned by construction: rustc_public's
BodyBuilder eagerly evaluates every MIR constant into Const::Val
(reported as ConstantKind::Allocated), so the collector's
instantiate+try_eval_bits folds exactly the constants the importer's
Allocated-scalar read folds; both fail together when const-eval cannot
produce a scalar. This is now documented at the collector fold site.

Regression tests: a unit test locks the importer fold to the shared
constant, and array_for_loop gains the issue-384 associated-const shape
(generic switch on an associated const with a dead panic arm, live arm
carrying a ub-checks RuntimeChecks edge).

Signed-off-by: nihalp <nihalp@nvidia.com>
The rebased array::map support exposed a chain of dead-path shapes main's
enum verifiers (from NVlabs#349) rightly refuse as real operations:
array::try_from_fn keeps the residual ControlFlow arms in MIR even though
NeverShortCircuitResidual is uninhabited, so the importer emitted
mir.enum_payload on an uninhabited variant, mir.construct_enum of an
uninhabited variant, and mir.get_discriminant on a fully uninhabited
enum -- each failing verification.

No value can inhabit those paths, so the reads and constructions are
dynamically dead. Translate each to a typed mir.undef instead (the same
treatment [T; 0] extraction and uninit union constants already get),
keeping the dead path representable while the dialect verifiers stay the
loud stop against real uninhabited-variant operations. New dialect tests
lock in the fail-loud half for construct_enum and get_discriminant.

The undef placeholders also exposed a textual-export ordering bug: undef
results were only named at op-emit time, so a use in an earlier-printed
block (a PHI incoming from a later-printed defining block) failed with
'missing exported name'. Pre-register undef results in the naming
pre-pass like constants and addressof, which are the other two virtual
ops.

The map_over_aligned_array kernel in array_for_loop exercises the whole
chain end to end on GPU.

Signed-off-by: nihalp <nihalp@nvidia.com>
The cl::opt applied alongside -passes is process-global: it changed the
unrolling cost model inside BOTH default<O2> runs for every kernel, not
just the bounded post-inline loop-unroll it was meant to feed. LLVM 21
and 22 expose no threshold= per-pass parameter for loop-unroll
(opt -print-passes lists only O0-O3, full-unroll-max=N, and the
partial/peeling/runtime/upperbound toggles), so the override cannot be
scoped to that single pass. Remove it entirely and let the default
threshold apply; the bounded pass keeps its per-pass O3 cost model and
full-unroll-max=16 cap.

The array_for_loop and array_index GPU runs, including their PTX
no-.local assertions, pass unchanged without the override.

Signed-off-by: nihalp <nihalp@nvidia.com>
Evaluated replacing the second default<O2> with a targeted cleanup
pipeline (sroa,early-cse,instcombine,simplifycfg) after the bounded
loop-unroll. Findings on the heavyweight examples (compile-only, .ll
piped through opt-22 three times each):

  middle-end (opt) time      single O2   O2+unroll+O2   targeted
  gemm_sol   (7616-line .ll)   0.08 s       0.13 s      LLVM ERROR
  wgmma       (223-line .ll)   0.01 s       0.01 s        0.01 s
  tcgen05_matmul (735 lines)   0.01 s       0.01 s        0.01 s

The targeted set hard-fails on gemm_sol: standalone instcombine in a
-passes string verifies its fixpoint after one iteration and raises
'did not reach a fixpoint' unless <no-verify-fixpoint> suppresses the
check. Suppressing an LLVM verifier to save ~0.05 s (<0.5% of the
~16 s end-to-end compile) is a bad trade, so the full second O2 stays.

PTX-size deltas vs a single default<O2> come from the bounded unroll
itself, not the duplicated O2: gemm_sol 5243 -> 8473 lines after the
unroll stage, 8593 after the second O2; tcgen05_matmul 414 -> 784 and
unchanged by the second O2; wgmma identical throughout. The array
examples keep their no-.local guarantees under both pipelines.

Signed-off-by: nihalp <nihalp@nvidia.com>
…te fan-out

is_readonly_aggregate_array_address gated only on 'GEP whose source
element type is a struct' plus a shared pointer, never on where the base
lives. For an aggregate genuinely resident in global memory (a
&BigStruct kernel param), one dynamic load fanned out into up to
MAX_SSA_ARRAY_ELEMENTS speculative global loads plus selects that SROA
can never delete -- a fan-out nvcc never takes. Require the GEP base to
be alloca-rooted so the rewrite only fires for function-local
aggregates, where the candidate loads collapse back into registers.

Tests: the alloca-rooted shared-reference shape still scalarizes; a
block-argument-rooted (external-memory) shared reference now keeps the
single dynamic GEP, matching the mutable case. The
runtime_borrowed_array_field kernel keeps its no-.local PTX guarantee on
GPU because rustc MIR-inlines the #[inline(always)] helper, leaving the
field GEP rooted at the caller's alloca.

Signed-off-by: nihalp <nihalp@nvidia.com>
…t elements

MAX_SSA_ARRAY_ELEMENTS counts elements only: [BigStruct; 16] flowed
sixteen whole-aggregate extractvalues plus fifteen aggregate selects
through registers. Add MAX_SSA_ARRAY_BYTES = 128 (the budget the element
cap already implies for word-sized elements, 16 x u64) next to the
element cap and enforce it on both the by-value extract path and the
constant-address selection path. Unsizeable layouts conservatively keep
the memory lowering. Tests cover an 8 x 32-byte array that satisfies
the element cap but overflows the byte budget on both paths.

Also document the select-chain clamp at both default sites: an
out-of-bounds UNCHECKED index now clamps to the last element instead of
touching wild memory -- UB either way, but the loud fault becomes a
silent stale read, which the code cannot express on its own.

Signed-off-by: nihalp <nihalp@nvidia.com>
clippy's assertions_on_constants rejects the runtime assert on
DEVICE_RUNTIME_CHECKS; a const assertion is also the stronger pin (the
invariant fails the build, not the test run). The aggregate byte-cap
tests passed `&mut ctx` where `IntegerType::get` takes `&Context`
(unnecessary_mut_passed).

Signed-off-by: nihalp <nihalp@nvidia.com>
…rapper

The importer's drop_glue_is_noop now applies the array::Drain
element-type reduction before the body proof. The collector consulted
the instance-level predicate directly, so for Drain-of-trivially-
droppable-element types it would collect and translate a Drain::drop
shim the importer never calls. Route DropGlue(_, Some(ty)) through the
same type-level wrapper so collection and emission stay in lockstep.

Signed-off-by: nihalp <nihalp@nvidia.com>
@nihalpasham
nihalpasham force-pushed the upstreaming/fixed-array-ssa-final branch from 971e77d to bd85728 Compare July 22, 2026 16:32
CI's fmt gate checks every example crate individually (they are
standalone workspaces the root cargo fmt never visits); the extended
no-local assert needed rustfmt's single-line form.

Signed-off-by: nihalp <nihalp@nvidia.com>
@nihalpasham nihalpasham added the on-hold Paused: valid but deferred until a prerequisite decision or dependency resolves label Jul 22, 2026
@nihalpasham

Copy link
Copy Markdown
Collaborator

Thanks for the PR. I think there's merit in examining this further. But as it stands, after looking at the full picture, I'm putting it on hold rather than merging.

The blocker is where the optimization lives, not whether it works:

  • This bakes a code-growing transformation into the default compile pipeline for every kernel: the select-chain lowering is always-on, and the added loop-unroll + second O2 stage runs on all device code unconditionally.
  • The side effects are global: kernels that never touch a runtime-indexed array pay for it — gemm_sol's PTX grows ~64%, and generated code becomes less predictable (a loop crossing the trip-count boundary silently flips its whole shape).

The general recommendation or guidelines for something like this is the three tiers for optimizations:

  1. compiler defaults, which require broad applicability and have to clear a high bar (FMA-fusion class);
  2. custom opts that live in a user's own project;
  3. an oxide-owned opt registry; upstream, but off by default, where users pick what fits their workload.

This PR would qualify as a tier-3 opt. It's workload-specific (small runtime-indexed local arrays), and most importantly it doesn't clear the default bar: kernels that never touch the pattern pay for it. The implementation shape should follow the registry design too: a pliron pass (Pass/OpPass with the pass-manager and AnalysisManager), not lowering-embedded logic plus an unconditional change to the opt pipeline every kernel runs through.

Before landing anything, I think examining prior art like how nvcc/nvc++ handle small runtime-indexed local arrays today (their register-promotion and unroll policies) would be a good idea to see if it makes sense for us.

An additional note: the enabling fixes in the PR (dead-branch pruning, union over-alignment, uninhabited enums, array::Drain drop glue) are valuable independent of the pipeline question, and we'll likely split those out. The branch stays as repaired, labeled on-hold.

This addresses #397, but the issue stays open until the in-house version lands.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working build-related Build scripts, toolchain discovery, or compilation setup codegen Device code-generation pipeline (Rust MIR to IR to PTX) enhancement New feature or request IR-lowering Lowering between dialects (dialect-mir to LLVM dialect) on-hold Paused: valid but deferred until a prerequisite decision or dependency resolves perf Performance of generated code or of the compiler itself

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Small runtime-indexed local arrays spill to PTX local memory

2 participants