Missed jump threading for predecessor-constant switch selectors

[tdavidovicNV]

Motivation and Desired Source

We generate material-specific data and heterogeneous BSDF storage, but want the material algorithm itself to remain ordinary handwritten Slang. Assuming the total selection weight is already available as selectionTotal, ideally that algorithm would be able to write:

const float target = nextFloat(sg) * selectionTotal;
float cdf = 0.0;
int index = TCount - 1;
for (int i = 0; i < TCount - 1; ++i)
{
    cdf += selection[i];
    if (target <= cdf)
    {
        index = i;
        break;
    }
}

result = weights[index] * data.evalBsdf(index, /* ... */);

Here, the loop performs the real algorithmic work of selecting a lobe by comparing a random target against the selection CDF. Generated data hides the mechanical heterogeneous dispatch behind evalBsdf().

Producing machine code comparable to the manually specialized, force-unrolled workaround shown below requires at least two independent capabilities:

• scalar replacement of weights[index], tracked separately by #12787;
• preserving the correlation between the index chosen by the CDF traversal and the concrete operation exposed by generated heterogeneous dispatch, which is the subject of this issue.

A current source workaround is to follow that selection traversal with a second, force-unrolled traversal that exposes a constant index to generated data explicitly:

const float target = nextFloat(sg) * selectionTotal;
float cdf = 0.0;
int index = TCount - 1;
for (int i = 0; i < TCount - 1; ++i)
{
    cdf += selection[i];
    if (target <= cdf)
    {
        index = i;
        break;
    }
}

[ForceUnroll]
for (int i = 0; i < TCount; ++i)
{
    if (i == index)
        result = weights[i] * data.evalBsdf(i, /* ... */);
}

The first loop performs the weighted choice. The second loop does no new algorithmic work: it only redispatches the already selected index. The second traversal is a functional source-level workaround: after forced unrolling, each i is constant, allowing generated data to resolve the concrete heterogeneous operation and scalarize fixed storage. However, it represents one logical selection-and-dispatch operation as two separate control flows. We do not consider this the desired source form. The compiler should eliminate the redundant dispatch traversal--either by fusing it with the CDF-selection traversal at loop IR or through equivalent jump threading after CFG lowering--so handwritten code can select the index once and use it directly.

Reduced Optimization Problem

Assuming generated heterogeneous dispatch or an unrolled visitor has already exposed concrete per-arm operations, and the CDF-selection loop is visible in the caller, the resulting CFG can contain two serial decisions:

1. A runtime selection chain assigns an edge-known literal arm ID.
2. A later switch uses the merged ID to invoke the corresponding concrete operation.

The selector is not globally constant, so ordinary SCCP cannot fold the switch. It is nevertheless known along every selecting path: the path that assigns N can only enter switch case N. Performing the concrete operation directly in that arm avoids the second dispatch.

The motivating split implementation explicitly computed the index and then invoked generated dynamic-index concrete dispatch. The reproducer models that same two-stage control structure with an explicit switch, while removing unrelated array, visitor, and heterogeneous-storage mechanics. The switch is not code we want users to write, and the issue does not claim that the ideal one-line source already lowers to this shape; reaching it still depends on the earlier transformations described above.

The desired transformation can be understood as loop fusion at source or loop IR level, or as predecessor-constant jump threading after unrolling and CFG lowering. This issue does not prescribe which compiler representation or pass should implement it. It does not require Slang to infer heterogeneous storage from an opaque runtime index. The reproducer deliberately removes arrays, loops, dynamic indexing, dead values, and application-specific types so that only the final CFG form remains.

Minimal Reproducer

Save the following as constant-phi-switch-repro.slang:

// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

struct SampleResult
{
    float3 value;
    float pdf;
    uint tag;
}

interface ILobe
{
    [ForceInline]
    SampleResult sample(float3 wi, float2 u);
}

struct Lobe<let ID : int> : ILobe
{
    float3 tint;
    float roughness;

    [ForceInline]
    SampleResult sample(float3 wi, float2 u)
    {
        const float id = float(ID);
        const float3 axis = normalize(float3(0.2 + 0.03 * id, 0.3 + 0.02 * id, 0.9 - 0.04 * id));
        const float cosine = abs(dot(normalize(wi), axis));
        float response = cosine * (0.8 + 0.03 * id) + roughness * (0.02 + 0.004 * id);
        response = (response * response + u.x * 0.03125) / (1.0 + roughness * (1.0 - cosine));

        SampleResult result;
        result.value = tint * response * (0.75 + 0.25 * u.y);
        result.pdf = saturate(cosine * (0.5 + 0.008 * id) + roughness * 0.125);
        result.tag = ID;
        return result;
    }
}

struct MaterialData
{
    Lobe<0> lobe0; float selection0;
    Lobe<1> lobe1; float selection1;
    Lobe<2> lobe2; float selection2;
    Lobe<3> lobe3; float selection3;
    Lobe<4> lobe4; float selection4;
    Lobe<5> lobe5; float selection5;
    Lobe<6> lobe6; float selection6;
    Lobe<7> lobe7; float selection7;
}

struct InputData
{
    float3 wi;
    float selector;
    float2 u;
    uint materialIndex;
}

[ForceInline]
SampleResult sampleConcrete<T : ILobe>(T lobe, float3 wi, float2 u)
{
    return lobe.sample(wi, u);
}

[ForceInline]
SampleResult finishSample(SampleResult result)
{
    result.value = result.value * 0.875 + result.pdf * float3(0.01, 0.03, 0.04);
    result.pdf = saturate(dot(result.value, float3(0.25, 0.5, 0.25)) + 0.01);
    return result;
}

// Reduced split form: selection and concrete dispatch are separate.
[ForceInline]
SampleResult sampleSplit(MaterialData data, float3 wi, float selector, float2 u)
{
    const float c0 = data.selection0;
    const float c1 = c0 + data.selection1;
    const float c2 = c1 + data.selection2;
    const float c3 = c2 + data.selection3;
    const float c4 = c3 + data.selection4;
    const float c5 = c4 + data.selection5;
    const float c6 = c5 + data.selection6;
    const float c7 = c6 + data.selection7;
    const float target = saturate(selector) * c7;

    int selected;
    if (target < c0) selected = 0;
    else if (target < c1) selected = 1;
    else if (target < c2) selected = 2;
    else if (target < c3) selected = 3;
    else if (target < c4) selected = 4;
    else if (target < c5) selected = 5;
    else if (target < c6) selected = 6;
    else selected = 7;

    SampleResult result = {};
    switch (selected)
    {
    case 0: result = sampleConcrete(data.lobe0, wi, u); break;
    case 1: result = sampleConcrete(data.lobe1, wi, u); break;
    case 2: result = sampleConcrete(data.lobe2, wi, u); break;
    case 3: result = sampleConcrete(data.lobe3, wi, u); break;
    case 4: result = sampleConcrete(data.lobe4, wi, u); break;
    case 5: result = sampleConcrete(data.lobe5, wi, u); break;
    case 6: result = sampleConcrete(data.lobe6, wi, u); break;
    case 7: result = sampleConcrete(data.lobe7, wi, u); break;
    }
    return finishSample(result);
}

// Control form: preserve the selected type until its concrete call.
[ForceInline]
SampleResult sampleFused(MaterialData data, float3 wi, float selector, float2 u)
{
    const float c0 = data.selection0;
    const float c1 = c0 + data.selection1;
    const float c2 = c1 + data.selection2;
    const float c3 = c2 + data.selection3;
    const float c4 = c3 + data.selection4;
    const float c5 = c4 + data.selection5;
    const float c6 = c5 + data.selection6;
    const float c7 = c6 + data.selection7;
    const float target = saturate(selector) * c7;

    SampleResult result;
    if (target < c0) result = sampleConcrete(data.lobe0, wi, u);
    else if (target < c1) result = sampleConcrete(data.lobe1, wi, u);
    else if (target < c2) result = sampleConcrete(data.lobe2, wi, u);
    else if (target < c3) result = sampleConcrete(data.lobe3, wi, u);
    else if (target < c4) result = sampleConcrete(data.lobe4, wi, u);
    else if (target < c5) result = sampleConcrete(data.lobe5, wi, u);
    else if (target < c6) result = sampleConcrete(data.lobe6, wi, u);
    else result = sampleConcrete(data.lobe7, wi, u);
    return finishSample(result);
}

StructuredBuffer<MaterialData> materials;
StructuredBuffer<InputData> inputs;
RWStructuredBuffer<SampleResult> results;

[shader("compute")]
[numthreads(64, 1, 1)]
void splitMain(uint3 tid : SV_DispatchThreadID)
{
    const InputData input = inputs[tid.x];
    results[tid.x] = sampleSplit(materials[input.materialIndex], input.wi, input.selector, input.u);
}

[shader("compute")]
[numthreads(64, 1, 1)]
void fusedMain(uint3 tid : SV_DispatchThreadID)
{
    const InputData input = inputs[tid.x];
    results[tid.x] = sampleFused(materials[input.materialIndex], input.wi, input.selector, input.u);
}

Compile the two entry points independently:

slangc constant-phi-switch-repro.slang -target cuda -entry splitMain -O3 -o split.cu
slangc constant-phi-switch-repro.slang -target ptx -entry splitMain -O3 -o split.ptx
slangc constant-phi-switch-repro.slang -target spirv-asm -entry splitMain -O3 -o split.spvasm

slangc constant-phi-switch-repro.slang -target cuda -entry fusedMain -O3 -o fused.cu
slangc constant-phi-switch-repro.slang -target ptx -entry fusedMain -O3 -o fused.ptx
slangc constant-phi-switch-repro.slang -target spirv-asm -entry fusedMain -O3 -o fused.spvasm

sampleSplit() represents the lowered selection-then-dispatch shape. sampleFused() is the control: it preserves the selected type until the concrete call. Both execute the same lobe calculation and common continuation.

Expected Behavior

After the generic calls are specialized and inlined, a bounded CFG transform should be able to redirect each selecting path to its known switch case:

selector arm N -> selected = N -> switch(selected) -> case N

becomes:

selector arm N -> case N

The concrete cases can still merge before finishSample(). This does not require cloning the common continuation.

The resulting machine-code shape should be comparable to fusedMain: one runtime selection chain, no subsequent tag dispatch.

Actual Behavior

With unmodified Slang 2026.16-26-g1e74b400f, the split form retains the second dispatch in all inspected targets:

Metric	splitMain	fusedMain
CUDA entry bytes	11,887	10,365
CUDA switches / cases	1 / 8	0 / 0
PTX bytes	25,508	23,996
PTX branches	25	21
PTX predicate definitions	18	7
SPIR-V OpSwitch	1	0
ptxas registers (sm_75)	26	26
Stack / spills	0 / 0	0 / 0

The byte and instruction-count rows are structural diagnostics, not runtime measurements. Text bytes are counted after removing Slang's trailing NUL. CUDA counts are restricted to the emitted entry function; PTX branch and predicate counts are bra* and setp.* instruction lines. The listed -target ptx commands use Slang's configured NVRTC downstream path and produced sm_75 PTX with CUDA Toolkit 13.3.

SPIR-V makes the missed correlation particularly direct:

%211 = OpPhi %int ... %int_3 ...
%212 = OpPhi %int %211 ... %int_2 ...
%213 = OpPhi %int %212 ... %int_1 ...
%214 = OpPhi %int %213 ... %int_0 ...
OpSwitch %214 ... 0 ... 1 ... 2 ... 3 ... 4 ... 5 ... 6 ... 7 ...

The CUDA backend also emits an eight-case switch. NVRTC receives that expanded selection chain and separate switch, and the final PTX retains four additional branches and eleven additional predicate definitions. On sm_75, the two reduced forms have identical register and spill counts, so the structural difference is not accompanied by a resource-count difference on that target.

Real-World Impact

The motivating renderer workload has eleven live heterogeneous operations. Two successive one-process RTX 5090 compute_120 experiments used the same all-active material and timing harness. Each dispatch made 4,147,200 material calls (1920 x 1080 x 2), followed 200 warmup dispatches, and was measured by an SGL GPU profiler zone over 50 dispatches:

Experiment	Material form, mean +/- standard deviation	Same-run V10 control
Fused weighted choice and concrete call	5.2292 +/- 0.1610 ms	5.4087 +/- 0.1423 ms
Split weighted choice then concrete dispatch	6.3505 +/- 0.1414 ms	5.4150 +/- 0.1566 ms

The split form is 21.4% slower than the fused form, while the V10 controls differ by 0.1% between runs. The source change moved the cumulative weighted-choice loop into handwritten code and followed it with generated concrete dispatch; the material inventory and surrounding sample algorithm were retained.

The slower form added 10,232 PTX bytes, 13 branches, and 14 predicate comparisons. It used fewer registers (202 versus 212), and both forms had a 32-byte frame and zero spills. This argues against a simple explanation that the regression was caused by allocating more registers or spilling, while the control-flow delta matches the reduction.

Why Downstream Optimization Is Not Sufficient

For CUDA, Slang emits the already-expanded selection chain followed by a separate switch and large case bodies. This requires NVRTC to reconstruct the original edge/value correlation across both constructs; the retained PTX difference shows that it does not do so in this case.

SPIR-V independently retains the phi-to-switch shape before any NVIDIA downstream compiler participates, so this is not CUDA-specific.

Potential Initial Scope

A conservative implementation could target switches where:

• the selector is a phi, or a promotable local represented by constant incoming assignments;
• each relevant predecessor path, possibly through a bounded chain of trivial merge blocks and phis, proves one switch destination;
• threading can redirect edges without cloning a nontrivial shared continuation;
• case count and CFG growth are explicitly bounded;
• structured-control-flow, derivative, and convergence constraints permit the rewrite.

This appears most useful after specialization, inlining, and scalar promotion have exposed the concrete operations, but before target-specific structured-CFG legalization obscures the correlation.

Source inspection suggests that Slang's SCCP folds a switch when the selector becomes globally constant, while CFG simplification handles trivial switches and local block fusion. Neither currently performs this predecessor-specific threading in the reproduced shape.

LLVM's bounded JumpThreading pass is a standard precedent for exploiting values known along specific predecessor paths. The requested behavior need not recognize interfaces, materials, BSDFs, or this source spelling.

An earlier rewrite attempted this before ordinary CFG/value compaction and cloned the common continuation, substantially increasing output size and compile cost. That failure motivates the join-preserving and cost-bounded constraints above rather than early continuation cloning.

Broader Relevance

The same source pattern can arise from tagged-union dispatch, enum/state-machine lowering, visitor implementations, specialized interface witnesses, and generated heterogeneous containers. The optimization is therefore useful independently of the motivating renderer.

Environment

• Slang: 2026.16-26-g1e74b400f (1e74b400fc9cd3312290706670f7806548f4f411)
• Relevant SCCP and CFG-simplification files were unchanged through inspected upstream commit 400ef93be5bc771df174a87fabb39e3fabb6117a
• Windows 11
• CUDA Toolkit 13.3
• ptxas -arch=sm_75 for register/stack/spill comparison
• RTX 5090, CUDA compute_120 for the motivating runtime comparison

[nv-slang-bot]

Triage — missed-optimization, routed to a fixer.

• Status: Triaged → handed to a Slang fixer; fix incoming. Classified Feature / medium / P2, component IR (new optimization pass). Set Issue Type = Feature (this is a missed optimization, output is correct — not a defect).
• Confirmed on top-of-tree (c1cffad25), GPU-free: compiling your reducer with a Debug slangc, splitMain -target spirv-asm -O3 emits 1 OpSwitch (8 cases) whose selector is a phi chain merging int_0…int_7 per predecessor edge, while fusedMain emits 0 — reproducing the phi→OpSwitch shape from your report at the IR/SPIR-V level, independent of any downstream compiler. (The RTX-5090 +21.4% figure is GPU-only and was not re-measured here.)
• Root gap: SCCP folds a switch only when the selector is a globally constant value; it models one lattice value per instruction, so a phi that carries a different constant along each incoming edge collapses to "unknown" and the switch survives. simplifyCFG does trivial-switch elision and block fusion but performs no value-based, predecessor-constant threading. So neither existing pass covers the reduced shape — consistent with your analysis.
• Recommended direction: a dedicated, bounded, join-preserving switch-threading pass that redirects each predecessor edge whose incoming selector constant proves exactly one case, leaving the cases to merge at the existing continuation (no continuation cloning — matching your rejected-experiment finding). Slotted after specialization/inlining/scalar-promotion and before target-specific structured-CFG legalization, exactly as you scoped it. Structured-control-flow / convergence / derivative constraints will be respected by gating on same-region edges.
• Next action: a fixer will evaluate the approach and, if it holds up, open a PR (Fixes #12792); this issue will be updated when a resolution lands. No further input needed from you right now — thank you for the exceptionally complete reducer and metrics.

_{🤖 Generated by an automated Slang coworker — may be inaccurate. A human maintainer should verify.}

[tdavidovicNV]

@nv-slang-bot please re-evaluate this issue using the current description and reproducer.

The motivation and desired source have been substantially revised since the initial triage. They now show the actual CDF-selection traversal followed by the force-unrolled heterogeneous-dispatch traversal. The requested result is to eliminate that redundant second control flow, potentially through loop fusion at loop IR or equivalent predecessor-constant jump threading after unrolling and CFG lowering. Please base any updated analysis and fix direction on the current issue body rather than its original revision.

[nv-slang-bot]

Re-triaged against the revised body — thanks for the added source context. Direct answer to your fusion-vs-threading question:

Which layer: I recommend keeping the fix at the predecessor-constant jump-threading layer (after CFG lowering), not loop fusion. Concretely, in the current architecture loop fusion is not a practical path:

• Slang has no loop-fusion pass of any kind (the loop passes are unroll, inversion, trivial-single-iteration removal, multi-level-break elimination, and autodiff loop analysis — none merge two serial loops).
• [ForceUnroll] is unrolled very early — inside the specialization fixpoint (unrollLoopsInModule, slang-ir-specialize.cpp:1734), before the main optimization pipeline's force-inlining and simplification. So the [ForceUnroll] dispatch loop is already peeled into the phi→OpSwitch CFG before any late pass could fuse it, and a hypothetical fuse-before-unroll pass would have to merge two structurally dissimilar loops (a data-dependent early-break selection loop and a fixed-count unrolled loop) — a strictly harder, speculative transform.
• Your reducer itself contains no loops (you removed them so "only the final CFG form remains"), so it exercises exactly this jump-threading layer — which I confirmed reproduces on top-of-tree.

Two refinements the revised body adds to the fix direction, which I've passed to the fixer:

1. The threading must walk a bounded chain of trivial merge blocks and phis, not a single-block phi. The if-else-if ladder lowers to a phi cascade — in the SPIR-V I see %214 = OpPhi %213 … int_0, %213 = OpPhi %212 … int_1, … — so the pass has to follow that chain back to the per-predecessor-edge constant that proves one case. (Matches your "possibly through a bounded chain of trivial merge blocks and phis" scope note.)
2. This issue stays scoped to the index↔concrete-operation correlation. The ideal single-loop source additionally needs scalar replacement of the dynamic-indexed weights[index], which is Promote small non-escaping fixed local arrays after specialization and loop unrolling #12787 and composes separately — I'm not conflating the two.

The join-preserving / no-continuation-cloning / cost-bounded constraints you reaffirmed match the recommended approach exactly. Classification is unchanged: Feature / medium / P2 / IR (new optimization pass). Handoff to the fixer has been re-synced with this direction; the issue will be updated when a resolution lands.

_{🤖 Generated by an automated Slang coworker — may be inaccurate. A human maintainer should verify.}

[nv-slang-bot]

Fix in draft PR #12795 (held pending review).

• Status: Fixed — draft PR Add jump-threading for predecessor-constant switch selectors (#12792) #12795 open against master, pr: non-breaking, CI dispatched.
• Root cause: SCCP folds a switch only when the selector is globally constant; here the selector is a phi carrying a distinct constant per incoming edge (constant along every path, not globally), so neither SCCP nor simplifyCFG removes the redundant second dispatch.
• Fix: new bounded, join-preserving IR pass threadSwitchOnConstantPhi — when every reaching edge proves one case, it re-routes each selecting arm to its case body and removes the now-dead switch, reusing the existing phi chain (no continuation cloning). Matches the fusedMain control shape.
• Verified: the reducer's splitMain now emits 0 OpSwitch (was 1); tests/optimization/ 42/42, tests/spirv/ 565/565, tests/autodiff/ 901/901 pass; SPIR-V validation passes. Decline-path tests confirm the switch survives for runtime selectors, side-effecting headers, and loop-boundary shapes.
• Next: peer + CI review on Add jump-threading for predecessor-constant switch selectors (#12792) #12795; a human maintainer decides on merge.

_{🤖 Generated by an automated Slang coworker — may be inaccurate. A human maintainer should verify.}