Zero-cost GPU bounds elision with Lean 4 and Zig

Published April 10, 2026

WebGPU has to be conservative around runtime-sized buffers. If the compiler cannot prove an index is safe, it keeps clamp machinery in the shader. Doe asks a narrower question: when a proof covers the index family, can the runtime move that safety check to dispatch time and delete the per-access guard?

The split is simple: Lean proves the whole bounds family; Zig refuses to build unless the checked-out artifact still matches that proof. If the implementation can match a shader access to a proof-covered pattern in the artifact, the WebGPU implementation can apply a simpler rewrite. Doe is a concrete WebGPU implementation that shows that path end to end.

The evidence is in the repo, not in a diagram. These are the exact files that carry the proof from theorem to build gate to shader rewrite:

• ComputeBounds.lean: the Lean theorem family itself, including the affine loop case this article uses.
• Extract.lean: the Lean extraction step that turns proved conditions into a machine-readable artifact.
• proven-conditions.json: the extracted proof manifest Zig actually consumes at build time.
• build.zig: where the build reads the Lean artifact, provenance, and proof options into build_options.
• lean-proof-patterns.json: the map from stable runtime pattern names to exact theorem names and callsites.
• lean_proof.zig: the comptime gate that rejects missing theorems or provenance drift with @compileError.
• affine_loop_storage.wgsl: the exact motivating affine-loop shader used in the proof metrics pass.
• dispatch_proof_match.zig: the matcher that recognizes affine WGSL index shapes and emits one dispatch precondition.
• ir_transform_robustness.zig: the rewrite pass where the fallback clamp stays or disappears.
• dispatch_preconditions.zig: the byte-count check that validates a matched precondition before launch.
• doe_compute_preconditions_native.zig: the host-side launch check that replaces per-access shader guards.

What The Default Path Keeps Alive

Consider this shader:

@group(0) @binding(0) var<storage, read_write> data: array<u32>;

@compute @workgroup_size(8)
fn main(@builtin(global_invocation_id) gid: vec3u) {
    for (var i: u32 = 1u; i <= 4u; i = i + 2u) {
        data[gid.x * 2u + i * 3u + 1u] = 1u;
    }
}

In the default robust path, the compiler cannot assume that

gid.x * 2u + i * 3u + 1u

is in bounds for every legal dispatch. So it lowers the access through a runtime clamp:

if (config.elide_proven_bounds) {
    if (dispatch_proof_match.try_elide_storage_index(module, function, function_id, i, index_data)) |precondition| {
        module.dispatch_preconditions.append(module.allocator, precondition) catch return error.OutOfMemory;
        continue;
    }
}

try clamp_runtime_sized(allocator, module, function, i);

In Zig, the fallback builds the clamp path explicitly:

const al_args = try function.append_expr_args(allocator, &.{base_ref});
const array_length_id = try function.append_expr(allocator, .{
    .ty = u32_ty,
    .category = .value,
    .data = .{ .call = .{
        .name = try ir.dup_string(allocator, "arrayLength"),
        .kind = .builtin,
        .args = al_args,
    } },
});

const one_id = try function.append_expr(allocator, .{
    .ty = u32_ty,
    .category = .value,
    .data = .{ .int_lit = 1 },
});

const sub_id = try function.append_expr(allocator, .{
    .ty = u32_ty,
    .category = .value,
    .data = .{ .binary = .{
        .op = .sub,
        .lhs = array_length_id,
        .rhs = one_id,
    } },
});

const min_args = try function.append_expr_args(allocator, &.{ original_index, sub_id });
const min_id = try function.append_expr(allocator, .{
    .ty = u32_ty,
    .category = .value,
    .data = .{ .call = .{
        .name = try ir.dup_string(allocator, "min"),
        .kind = .builtin,
        .args = min_args,
    } },
});

That is the baseline. min(...) stays in the kernel. Runtime size plumbing stays alive with it.

What Lean Proves

Lean proves the affine family with one inequality. In ComputeBounds.lean, the affine loop case is stated for arbitrary naturals:

theorem gid_affine_plus_scaled_loop_index_inbounds_when_dispatch_fits
    (workgroup_id_x local_id_x workgroup_size_x num_workgroups_x array_length : Nat)
    (gid_stride limit i loop_stride offset : Nat)
    (h_gid_stride : 0 < gid_stride)
    (h_loop_stride : 0 < loop_stride)
    (h_i : i < limit)
    (h_wid : workgroup_id_x < num_workgroups_x)
    (h_lid : local_id_x < workgroup_size_x)
    (h_fit : workgroup_size_x * num_workgroups_x * gid_stride + limit * loop_stride + offset ≤ array_length) :
    globalInvocationId workgroup_id_x local_id_x workgroup_size_x * gid_stride + i * loop_stride + offset < array_length := by
  let gid := globalInvocationId workgroup_id_x local_id_x workgroup_size_x
  let total := workgroup_size_x * num_workgroups_x
  have h_gid_lt_total : gid < total := by
    dsimp [gid, total]
    exact gid_component_lt_total
      workgroup_id_x local_id_x workgroup_size_x num_workgroups_x
      (workgroup_size_x * num_workgroups_x)
      h_wid h_lid (Nat.le_refl _)
  have h_gid_scaled_lt :
      gid * gid_stride < total * gid_stride := by
    exact Nat.mul_lt_mul_of_pos_right h_gid_lt_total h_gid_stride
  have h_loop_scaled_lt :
      i * loop_stride < limit * loop_stride := by
    exact Nat.mul_lt_mul_of_pos_right h_i h_loop_stride
  have h_sum_lt :
      gid * gid_stride + i * loop_stride <
      total * gid_stride + limit * loop_stride := by
    exact Nat.add_lt_add h_gid_scaled_lt h_loop_scaled_lt
  have h_offset_lt :
      gid * gid_stride + i * loop_stride + offset <
      (total * gid_stride + limit * loop_stride) + offset := by
    exact Nat.add_lt_add_right h_sum_lt offset
  calc
    gid * gid_stride + i * loop_stride + offset <
        (total * gid_stride + limit * loop_stride) + offset := h_offset_lt
    _ ≤ array_length := h_fit

Zig comptime can check one concrete tuple, but not the whole quantified family. Lean proves the family. Zig makes it operational.

What Zig Matches

On the implementation side, the affine-loop matcher path is real code, not a schematic:

if (match_gid_component_loop_affine_plus_offset(function, index_data.index, .global_invocation_id)) |match| {
    if (loop_match.find_bounded_loop_limit(function, expr_id, match.local_idx)) |loop_limit| {
        return .{
            .kind = .gid_component,
            .gid_axis = match.axis,
            .storage_binding = binding,
            .element_multiplier = match.gid_multiplier,
            .loop_limit = loop_limit,
            .loop_limit_multiplier = match.loop_multiplier,
            .element_stride_bytes = element_stride_bytes,
            .element_offset = match.offset,
        };
    }
}

For the motivating shader, those matched fields resolve to axis 0, binding (0, 0), gid multiplier 2, loop limit 5, loop multiplier 3, 4-byte elements, and offset 1. The 5 is the exclusive upper bound after normalizing i <= 4 to < 5. Once that match succeeds, the shader access stays in its original form and the guard machinery disappears. There is no per-access min(...), no arrayLength(&data), and no _doe_sizes buffer kept alive just to answer a runtime query inside the kernel. The safety check moves to dispatch time instead. The byte-count path is also straight from the runtime:

.gid_component => blk: {
    const axis = precondition.gid_axis;
    if (axis >= dispatch_workgroups.len or axis >= workgroup_size.len) {
        return error.DispatchPreconditionFailed;
    }
    const total_invocations = try invocation_extent(dispatch_workgroups[axis], workgroup_size[axis]);
    const scaled_invocations = try std.math.mul(u64, total_invocations, precondition.element_multiplier);
    const loop_elements = try std.math.mul(u64, precondition.loop_limit, precondition.loop_limit_multiplier);
    const affine_elements = try std.math.add(u64, scaled_invocations, loop_elements);
    const total_elements = try std.math.add(u64, affine_elements, precondition.element_offset);
    break :blk try std.math.mul(u64, total_elements, precondition.element_stride_bytes);
},

The exchange is:

1. GPU-side clamp removed
2. size-buffer machinery removed with it
3. one host-side bound check added before launch

Compile p50 drops by about a quarter. Steady-state dispatch gets a smaller but measurable lift. GPU timestamps move only slightly.

Why This Matters

Performance is a benefit, but the longer-term appeal is simpler invariants, fewer ad hoc guards, and a tighter proof-to-build contract. The proofs are not treated like a loose JSON hint. The extraction step writes a fixed schema with provenance and theorem inventory. The Zig build embeds the artifact and the expected hashes. lean_proof.zig recomputes the same provenance from the current checkout and rejects drift with @compileError. The artifact is trusted only when the local tree can re-derive it.

The point is not that Lean makes GPU code magically faster. The point is that a proof can become an enforceable build artifact, then a matcher, then a specific IR rewrite. The clamp disappears only when that chain still holds.

Reproduction note: Build with -Dlean-verified=false and -Dlean-verified=true. See proof_metrics_summary.json.