Slang Autodiff Kernel Opt

Autodiff Kernel Optimization

Methodology for investigating and optimizing Slang autodiff backward kernel performance.

Investigation workflow

1. Isolation

Large projects have complex build systems (Bazel, CMake) with many dependencies. For kernel investigation, extract the target function into a self-contained .slang file that compiles directly with slangc. This cuts the edit-compile cycle from minutes to seconds.

Guidelines:

• Identify the target kernel — the one with the worst backward performance or highest register count.
• Create a new .slang file containing only the code needed for that kernel. Inline small dependencies; #include larger shared ones.
• If the original code uses a module system (import), avoid it — standalone compilation via #include is more reliable and avoids module deserialization issues.
• Use #ifdef guards for A/B variant testing from the same source file:

#ifndef USE_AUTODIFF
[BackwardDerivativeOf(my_func)]
void my_func_bwd(...) { /* custom backward */ }
#endif

Compile with/without -DUSE_AUTODIFF to compare.

• For fully portable benchmarks, inline all dependencies so the file compiles with zero -I flags. This makes the experiment shareable as a standalone package.

CUDA example:

Under WSL with a Windows-hosted Slang build, invoke the matching .exe binary through $SLANGC (for example slangc.exe) and stop if it is unavailable. Do not silently run a WSL-native slangc from another build.

# slangc generates a self-contained .cu (includes Slang CUDA prelude)
"$SLANGC" my_shader.slang -O3 -target cuda -o /tmp/out.cu

# nvcc compiles with register/spill reporting
nvcc /tmp/out.cu -arch=sm_89 --use_fast_math \
  -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ \
  -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ \
  --ptxas-options=-v -cubin -o /tmp/out.cubin 2>&1 | grep -A2 "_bwd_diff$"

The -target cuda output is self-contained — TensorView/DiffTensorView become plain C++ structs. No PyTorch or slangtorch needed.

2. Register/spill analysis

Compile the isolated kernel to the target backend and check resource usage for the _bwd_diff entry points.

3. Checkpoint analysis

Use -report-checkpoint-intermediates with any target to see what autodiff stores in the forward pass for use in the backward pass:

"$SLANGC" my_shader.slang -O3 -target cuda -report-checkpoint-intermediates -o /dev/null

Example. Given this kernel:

uint2 get_index<let N : int>(no_diff uint index) {
    return uint2(index / N, index - (index / N) * N);
}

[CUDAKernel] [Differentiable] [AutoPyBindCUDA]
void loop_kernel(no_diff uint count, DiffTensorView output, DiffTensorView input) {
    uint tid = cudaBlockIdx().x * cudaBlockDim().x + cudaThreadIdx().x;
    if (tid >= count) return;
    [ForceUnroll]
    for (int i = 0; i < 8; ++i) {
        uint2 index = get_index<8>(i * count + tid);
        float x = input.loadOnce(index);
        output.storeOnce(index, square(x));
    }
}

The report shows:

note: checkpointing context of 96 bytes associated with: 'loop_kernel'
  --> example.slang:6:6
   |
 6 | void loop_kernel(...)
   |      ^^^^^^^^^^^
...
13 |         uint2 index = get_index<8>(i * count + tid);
   |                                   - 8 instances of 8 bytes (Vector<uint32_t, 2>)
14 |         float x = input.loadOnce(index);
   |                                 - 8 instances of 4 bytes (float)
---'

Two variables are stored per unrolled iteration: index (uint2, 8 bytes) and x (float, 4 bytes), for 8 iterations = 96 bytes total. Note that index is computed purely from no_diff values and shouldn't need checkpointing — this is a candidate for [PreferRecompute] (see Strategy 1 below).

For [MaxIters(N)] loops, storage appears as FixedArray<T, N+2> instead of N instances of.

What to look for:

• Large per-iteration storage: many instances of a variable in an unrolled loop. The total bytes = instances x size_per_instance. This is the main source of register pressure.
• Redundant storage: values recomputable from already-stored values.
• Non-diff values being stored: computations derived entirely from no_diff inputs that shouldn't need checkpointing. These are candidates for [PreferRecompute].
• Tasks with zero checkpoint: functions with [BackwardDerivativeOf] show no checkpoint. When comparing, this tells you the cost that autodiff would add if the custom backward were removed.

4. Benchmarking

Register counts and checkpoint analysis predict performance but don't prove it. Always measure actual runtime to confirm improvements.

Guidelines:

• Compile both the original and optimized variants from the same source (using #ifdef toggles).
• Use the target platform's profiling/timing mechanism.
• Run warmup iterations before measurement to stabilize caches and clocks.
• Run multiple times and average to reduce noise.

CUDA example: Write a host launcher that #includes the slangc-generated .cu, allocates GPU tensors, and times kernel launches with cudaEvent_t.

Optimization strategies

Strategy 1: [PreferRecompute] on pure/cheap functions

If a function's output is checkpointed per iteration but can be recomputed cheaply, mark it [PreferRecompute]. Two proven patterns:

Pattern A: Pure functions of no_diff inputs. If a function's inputs are all no_diff (loop counters, thread IDs, constants), its output is free to recompute. Mark it [PreferRecompute] to avoid checkpointing the return value per iteration.

// Before: autodiff stores the returned uint2 per iteration
uint2 computeCoord(no_diff uint flatIndex) { ... }

// After: backward recomputes it instead of storing
[PreferRecompute]
uint2 computeCoord(no_diff uint flatIndex) { ... }

If the function has side effects (e.g., tensor loads via loadOnce), use [PreferRecompute(SideEffectBehavior::Allow)] to suppress the warning that recomputing may re-execute side effects. This is safe when the side effect is an idempotent read.

Pattern B: Cheap-to-redo work. If a value can be reproduced by re-reading memory or re-running a trivial computation, wrap that work in a [PreferRecompute] helper. This trades re-execution in the backward pass for fewer bytes of checkpoint storage. Good when the re-execution cost (e.g., one global memory load) is cheaper than the register pressure from storing the value across many iterations.

// Before: autodiff stores val (4 bytes) per iteration
for (int i = 0; i < N; ++i) {
    float val = input.load(index);
    float newVal = operationsOn(val);
    output.store(index, newVal);
}

// After: wrap load+operations so backward re-reads from memory
[PreferRecompute]
[Differentiable]
float loadAndCompute(DiffTensorView input, uint2 index) {
    return operationsOn(input.load(index));
}

for (int i = 0; i < N; ++i) {
    output.store(index, loadAndCompute(input, index));
}

Strategy 2: [MaxIters(N)] instead of [ForceUnroll]

When a [ForceUnroll] loop stores checkpoint context per iteration, switching to [MaxIters(N)] makes autodiff use loop replay with compact FixedArray storage instead of N separate SSA values. Best for large N (e.g., 48-iteration copy loops). Not beneficial when per-iteration state is small enough to fit in registers without spilling.

Per-call threshold pattern. If a shared differentiable helper is not easy to switch globally to [MaxIters], add a generic threshold parameter and choose the loop annotation from compile-time constants. This lets each task or kernel specialization opt into loop replay only where profiling proves it helps.

[Differentiable]
void visitComponents<let ComponentCount : int, let ReplayThreshold : int>(
    no_diff uint elementId)
{
    if (ComponentCount > ReplayThreshold) {
        [MaxIters(ComponentCount)]
        for (int componentId = 0; componentId < ComponentCount; ++componentId) {
            // Call the target differentiable function here.
        }
    }
    else {
        [ForceUnroll]
        for (int componentId = 0; componentId < ComponentCount; ++componentId) {
            // Call the target differentiable function here.
        }
    }
}

// Keep small or sensitive call sites unrolled.
visitComponents<20, 1024>(elementId);

// Replay only the specialization where 45 iterations reduced register pressure.
visitComponents<45, 44>(elementId);

Use this pattern when:

• A global [MaxIters] change improves one kernel but regresses another.
• The same helper is instantiated by multiple tasks with different register pressure, occupancy, or checkpoint behavior.
• You need to preserve manual-backward parity for some call sites while improving a large-loop specialization.

Benchmark each threshold candidate. In practice, a threshold just below the problematic loop trip count, such as 44 for a 45-element copy, can isolate the benefit while leaving other specializations on the original unrolled path.

Strategy 3: Eliminate redundant loop-carried variables

If a loop-carried variable is only used in a non-differentiable context (e.g., convergence checks, early termination), it doesn't need gradients. Use detach() to prevent autodiff from tracking it:

// Before: prev_point is a loop phi — stored per iteration for autodiff
float2 prev_point = float2(0, 0);
for (...) {
    float2 pixel_diff = result.point - prev_point;  // used for convergence check
    if (length(pixel_diff) < threshold) break;
    prev_point = result.point;
}

// After: detach() tells autodiff this value doesn't need gradients
for (...) {
    float2 prev = detach(result.point);
    float2 pixel_diff = detach(result.point) - prev;
    if (length(pixel_diff) < threshold) break;
}

Only use detach() when the variable genuinely doesn't need derivatives. If prev_point participates in differentiable computation (not just control flow), detaching it will produce incorrect gradients.

Strategy 4: Split phases (convergence loop + single diff step)

For iterative convergence loops. Extract the loop into a non-[Differentiable] solver, call via no_diff, then do one differentiable step:

ConvergenceResult _solve(...) { /* plain loop, not [Differentiable] */ }

[Differentiable]
Result my_func(...) {
    ConvergenceResult cr = no_diff _solve(...);
    return differentiable_step(cr.converged_params, ...);
}

Eliminates ALL loop checkpoint storage. Produces code equivalent to a hand-written [BackwardDerivativeOf] without writing the backward manually.

Strategy 5: Custom derivatives with [BackwardDerivativeOf]

As a last resort, write a manual backward when autodiff fundamentally cannot derive the correct result:

• Different math: the backward requires a mathematically different formulation (e.g., implicit differentiation, adjoint methods) that cannot be expressed by mechanically differentiating the forward code.
• Synchronization / memory access patterns: the backward needs warp-level reductions (WaveActiveSum), shared memory atomics, or other coordination that autodiff doesn't generate.

[BackwardDerivativeOf(my_func)]
void my_func_bwd(...) { /* hand-written backward */ }

Use this minimally. Custom derivatives duplicate logic (forward and backward must stay in sync), increase code size, and bypass the compiler's checkpoint optimizations. Prefer Strategies 1-4 first — they let autodiff generate correct derivatives while controlling performance via annotations. Only fall back to [BackwardDerivativeOf] when the backward truly requires different logic that autodiff cannot produce.