TheRock/rocm-libraries/shared/stinkytofu/docs/developer/wmma-vgpr-reorder-pass.md

WMMA VGPR Reorder Pass

Overview

StinkyWmmaVgprReorderPass is a read-only analysis pass that detects VGPR savings in N-buffered GEMM kernels by identifying when wmma instructions can be reordered to allow cross-pool register aliasing.

The pass never mutates any instruction or register operand. All output is in WmmaReorderAnalysisResult for downstream passes to act on.

---

Problem

In an N-buffered GEMM loop body the wmma instructions are split into N pools, each using its own set of registers for the pool-varying operand (either A or B depending on the kernel). When that operand is the outer loop dimension its registers stay live across all iterations of the inner dimension, preventing cross-pool aliasing and wasting VGPRs.

Making the pool-varying operand the inner dimension (grouping its instructions contiguously within each pool) tightens its liveness interval so that pool N's registers die before pool 0's next use. The two pools can then be aliased to the same physical registers.

---

Three-Layer ABI

The pass is structured as three independent layers. Each can be swapped without touching the others.

Layer 1 — IRegLivenessAnalysis

Computes live intervals for the A and B register groups referenced by the wmma sequence.

class IRegLivenessAnalysis {
public:
    virtual std::map<RegGroup, RegInterval> computeLiveness(
        const BasicBlock&            bb,
        const std::vector<WmmaNode>& wmmaSeq) const = 0;
};

Built-in implementation: WmmaIntervalLiveness

Fast wmma-only backend. The interval for a register group is [first wmma index that reads it, last wmma index that reads it]. Sufficient for standard N-buffered GEMM patterns.

A full-instruction liveness backend can be substituted here without changing anything else.

---

Layer 2 — IWmmaReorderAlgorithm

Given the pre-grouped pools and precomputed live intervals, decides the optimal reordering and which register groups can be aliased.

class IWmmaReorderAlgorithm {
public:
    struct Result {
        std::vector<WmmaNode>       desiredOrder;
        std::vector<AliasCandidate> aliases;
    };

    virtual Result solve(
        const std::vector<std::vector<WmmaNode>>& pools,
        const std::map<RegGroup, RegInterval>&    liveness) const = 0;
};

pools arrives already grouped and relabeled by detectABIndices so that WmmaNode::aGroup is always the pool-varying operand and WmmaNode::bGroup is always pool-shared.

Built-in implementation: PoolVaryingReorderAlgorithm

Checks whether the pool-varying operand (A) has inflated liveness (interval width > number of distinct B groups). If so, reorders each pool via reorderContiguously to make A groups contiguous, then pairs up corresponding A groups across pools as alias candidates.

---

Layer 3 — WmmaReorderAnalysisResult

The stable output contract consumed by all downstream passes.

struct WmmaReorderAnalysisResult {
    bool                          applicable;
    std::vector<StinkyInstruction*> desiredWmmaOrder;  // permutation of wmma pointers
    std::vector<RegReplacement>   replacements;         // per-operand rewrite map
    unsigned                      totalVgprSaved;
};

Retrieve after the pass runs:

const WmmaReorderAnalysisResult* result = getWmmaReorderResult(bb);

---

Pool Tagging

TensileLite stamps each wmma instruction with a WmmaPoolData modifier at generation time so the pass can group instructions by pool without heuristics.

struct WmmaPoolData : public TypedModifier<WmmaPoolData> {
    uint32_t poolIndex = 0;
};

If any wmma instruction is missing this modifier the pass bails out for the entire basic block. Partial tagging indicates a misconfigured pipeline and the pass must not proceed on incomplete information.

---

A/B Detection

detectABIndices determines which src operand index is pool-varying (A) vs pool-shared (B) by checking register group intersection across pools. If src0's groups appear in multiple pools it is pool-shared and src1 is A; otherwise src0 is A.

This means the pass handles both the standard case (hardware A is pool-varying) and the symmetric case (hardware B is pool-varying) without any special-casing.

---

Injecting Custom Implementations

Both the liveness backend and the algorithm are injected at construction:

auto pass = createStinkyWmmaVgprReorderPass(
    std::make_unique<MyLivenessBackend>(),
    std::make_unique<MyReorderAlgorithm>());

Passing nullptr for either uses the built-in default.