[tlopex]

This pr fixes InjectSoftwarePipeline to rewrite opaque buffer offsets for legacy PTX MMA intrinsics, including ptx_ldmatrix_legacy, ptx_mma_legacy, mma_store_legacy, and mma_fill_legacy.

Failure

test_async_nested_pipeline_mma_gemm_ideal_annotation failed with a real numerical mismatch during the final GEMM result check:

Mismatched elements: 6744298 / 16777216 (40.2%)
Max absolute difference: 6.706421
Max relative difference: 0.00653978

A diagnostic run with a fixed seed reproduced the same class of error:

nested_full bad 6826830 / 16777216
max_abs 7.0405273
mean_abs 0.9874781

The adjacent simple pipeline test passed, and disabling async copy did not change the mismatch, so this was not a cp.async lowering issue, tolerance noise, or GPU flakiness.

Root Cause

InjectSoftwarePipeline may add a leading version dimension to pipeline buffers. Normal buffer loads/stores and newer opaque PTX intrinsics already had their offsets rewritten to include the pipeline version slot.

However, the legacy MMA intrinsics were not covered. In this test, the warp buffers became multi-versioned, but the legacy ldmatrix/mma offsets still pointed to the original slot. As a result, the second ldmatrix stage overwrote the first warp fragment, and both MMA stages read the same fragment. That skipped one K fragment and duplicated another, producing the numerical mismatch above.

[gemini-code-assist]

Code Review

This pull request adds support for legacy operations (such as legacy PTX ldmatrix and MMA ops) in the software pipeline transformation by introducing a helper function IsOp and updating the buffer access rewriter. The feedback suggests declaring the IsOp helper function as static to ensure internal linkage and prevent potential linker collisions.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.