GVN excludes `CoopMma` wholesale, so identical cooperative-matrix loads are never CSE'd — measured ~2% end-to-end cost on Vulkan (ptxas hides it on CUDA)

[laurigates]

Summary

cubecl-opt's GVN refuses to value-number any Operation::CoopMma(_). In crates/cubecl-opt/src/gvn/numbering.rs (same on 0.10.0 and current main):

Operation::Branch(_)
| Operation::Synchronization(_)
| Operation::CoopMma(_)
| Operation::NonSemantic(_)
| Operation::Barrier(_)
| Operation::Tma(_)
| ...
| Operation::Marker(_) => Err(None),

That is clearly the right default for Execute/Store and the manual-register ops, but it also covers pure CoopMma::Loads: two loads of the same buffer, offset, stride, and layout with no intervening store or barrier can never be deduplicated. The optimizer does GVN the scalar address bindings feeding them, so the redundancy is visible directly in the optimizer IR — repeated matrix_loads referencing the same offset binding:

matrix(96)  = matrix_load(input(1), stride: u32(64), offset: binding(93))
matrix(102) = matrix_load(input(1), stride: u32(64), offset: binding(100))
…
matrix(183) = matrix_load(input(1), stride: u32(64), offset: binding(93))   // duplicate of matrix(96)
matrix(189) = matrix_load(input(1), stride: u32(64), offset: binding(100))  // duplicate of matrix(102)

Minimal shape of the repro

Any kernel where a B-fragment load sits inside an #[unroll] loop it doesn't depend on:

#[cube(launch_unchecked)]
fn repro(k: &Array<f16>, out: &mut Array<f32>, #[comptime] d: usize) {
    // 2 accumulator row-strips × shared B fragment: the B load doesn't
    // depend on `mb`, but is emitted (and kept) once per iteration.
    #[unroll]
    for mb in 0..2usize {
        let acc = cmma::Matrix::<f32>::from_value(cmma::MatrixIdent::Accumulator,
            16usize, 16usize, 16usize, cmma::MatrixLayout::Undefined, f32::new(0.0));
        let b = cmma::Matrix::<f16>::from_slice(cmma::MatrixIdent::B,
            16usize, 16usize, 16usize, cmma::MatrixLayout::ColMajor,
            &k.slice(0, k.len()), comptime!(d as u32));   // ← identical both iterations
        // … execute with a per-mb A fragment, store …
    }
}

In our FlashAttention kernel this pattern (K and V fragment loads inside a 2-way mb unroll) emits 24 cooperative-matrix loads of 12 unique addresses per key block; we counted the identical duplication in the cubecl optimizer IR, in the emitted SPIR-V (OpCooperativeMatrixLoadKHR count), and in the generated CUDA C++ (wmma::load_matrix_sync count) — 80 loads total vs 68 after deduplicating at the source level.

Measured consequence (RTX 4090, f16 m16n16k16, both backends)

• CUDA: end-to-end timing was insensitive to the duplication itself — a source-level hoist of the duplicated loads into Sequence-held fragments changed CUDA time only through register pressure (-Xptxas -v: 168 → 254 regs/lane for a whole-loop K hoist, dropping occupancy 12 → 8 blocks/SM at 32 threads/block, +4–6% slower). So NVRTC/ptxas evidently already eliminates the redundant loads downstream (inferred from timing + register data; we did not diff SASS).
• Vulkan (wgpu<spirv>, NVIDIA 580.126.18): the driver does not appear to absorb them. The same source-level K hoist measured reproducibly ~2% faster end-to-end at large shapes — two independently dispatched runs, same direction and magnitude, under a harness with clock-steadying convergence rungs and a 2×-iterations proportionality check on every quoted number — despite paying the same occupancy drop, so the pure cost of the duplicated OpCooperativeMatrixLoadKHRs is larger than 2% of this kernel. A likely aggravating factor (backend context, possibly a separable cubecl-spirv item rather than part of the GVN question): cubecl-spirv materializes every fragment as a Function-storage OpVariable with OpStore/OpLoad round-trips around each OpCooperativeMatrixMulAddKHR (94 such variables in this kernel), so the driver would have to promote coopmat-typed variables to values (the classic mem2reg pass) and CSE the loads to recover what nvcc gets from wmma::fragment locals in SSA form.

We've worked around it at the source level (hoisted, comptime-indexed Sequence fragments), so nothing is blocking — filing because the IR-level cause is cleanly isolated and the Vulkan cost is measured, in case coopmat loads are ever made numberable.

Suggestion

Treat CoopMma::Load (and possibly Fill/Cast) as numberable expressions keyed on (buffer, offset value-number, stride, layout, fragment shape/ident), with the same invalidation discipline GVN already applies to scalar Index reads from the same memory space (any store to that space, and barriers, kill the number). Execute/Store and the manual-register ops stay unnumberable. The aliasing/invalidations questions are the real work here and you know those trade-offs better than we do — happy to test a branch, provide the raw .spv/IR/CUDA-C++ dumps, or attempt a PR if the direction sounds right.

Environment / links

• cubecl 0.10.0 (crates.io) — duplication + measurements; numbering.rs exclusion confirmed unchanged on main at filing time.
• Full kernel, dispatched-run IDs, register table, and dump artifacts (public): https://github.com/laurigates/custom-attention-engine-framework/issues/25 ("Lever 2 findings" comment) and https://github.com/laurigates/custom-attention-engine-framework/pull/29.