[stinkytofu] Add WmmaVgprReorder analysis pass with explicit pool tagging

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

@@ -0,0 +1,151 @@
+# 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.
+
+```cpp
+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.
+
+```cpp
+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.
+
+```cpp
+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:
+
+```cpp
+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.
+
+```cpp
+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:
+
+```cpp
+auto pass = createStinkyWmmaVgprReorderPass(
+    std::make_unique<MyLivenessBackend>(),
+    std::make_unique<MyReorderAlgorithm>());
+```
+
+Passing `nullptr` for either uses the built-in default.

shared/stinkytofu/include/stinkytofu/ir/asm/StinkyModifiers.hpp

@@ -293,6 +293,7 @@ struct Modifier {
         COMMENT,
         MATRIX_FMT,
         MEM_TOKEN,
+        WMMA_POOL_INDEX,
     };
 
     Modifier(Type type) : type(type) {}
@@ -1034,4 +1035,15 @@ struct MemTokenData : public TypedModifier<MemTokenData> {
         : TypedModifier<MemTokenData>(), tokens(tokens) {}
 };
 
+/// Buffer pool index for WMMA instructions in double/triple/N-buffered GEMM kernels.
+/// Set by TensileLite during rocisa → StinkyTofu conversion. Consumed by
+/// StinkyWmmaVgprReorderPass to group wmma instructions into pools without heuristics.
+struct WmmaPoolData : public TypedModifier<WmmaPoolData> {
+    static constexpr Modifier::Type Type = Modifier::Type::WMMA_POOL_INDEX;
+
+    uint32_t poolIndex = 0;
+
+    explicit WmmaPoolData(uint32_t idx) : TypedModifier<WmmaPoolData>(), poolIndex(idx) {}
+};
+
 }  // namespace stinkytofu

shared/stinkytofu/include/stinkytofu/transforms/asm/StinkyWmmaVgprReorderPass.hpp

@@ -0,0 +1,185 @@
+/* ************************************************************************
+ * Copyright (C) 2025-2026 Advanced Micro Devices, Inc.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ *
+ * ************************************************************************ */
+
+#pragma once
+
+#include <map>
+#include <memory>
+#include <vector>
+
+#include "stinkytofu/Export.hpp"
+#include "stinkytofu/ir/asm/StinkyModifiers.hpp"
+#include "stinkytofu/ir/asm/StinkyRegister.hpp"
+
+namespace stinkytofu {
+
+class BasicBlock;
+class Pass;
+struct StinkyInstruction;
+
+// ─────────────────────────────────────────────────────────────────────────────
+// Shared data types
+// ─────────────────────────────────────────────────────────────────────────────
+
+/// A contiguous VGPR group identified by its base index and size.
+struct RegGroup {
+    uint32_t base;
+    uint16_t size;
+    bool operator==(const RegGroup& o) const {
+        return base == o.base && size == o.size;
+    }
+    bool operator<(const RegGroup& o) const {
+        return base < o.base || (base == o.base && size < o.size);
+    }
+};
+
+/// Live interval expressed as wmma-sequence positions (or full-instruction
+/// positions for richer liveness backends). Consumers treat values as opaque
+/// ordinals and only compare them for overlap.
+struct RegInterval {
+    unsigned first = 0;
+    unsigned last = 0;
+    bool overlaps(const RegInterval& o) const {
+        return first <= o.last && o.first <= last;
+    }
+};
+
+/// A single per-operand rewrite entry. The downstream register-renaming pass
+/// applies this list mechanically — no further analysis is required.
+struct RegReplacement {
+    StinkyInstruction* inst;  ///< instruction whose operand must change
+    unsigned operandIdx;      ///< index into srcRegs (isSrc=true) or destRegs
+    bool isSrc;
+    StinkyRegister oldReg;
+    StinkyRegister newReg;
+};
+
+/// Stable output contract between this pass and all downstream consumers.
+/// Never changes regardless of which liveness backend or algorithm is active.
+struct WmmaReorderAnalysisResult {
+    /// True when the optimization applies to the analyzed basic block.
+    bool applicable = false;
+
+    /// Desired wmma ordering: a permutation of the original wmma instruction
+    /// pointers. A downstream reorder pass applies this directly to the BB.
+    std::vector<StinkyInstruction*> desiredWmmaOrder;
+
+    /// Flat per-operand replacement map. Each entry describes exactly one
+    /// operand that must be rewritten to eliminate the aliasable register range.
+    std::vector<RegReplacement> replacements;
+
+    /// Total VGPRs freed when all replacements are applied.
+    unsigned totalVgprSaved = 0;
+};
+
+// ─────────────────────────────────────────────────────────────────────────────
+// ABI 1 — Liveness backend (swappable)
+// ─────────────────────────────────────────────────────────────────────────────
+
+/// Internal wmma descriptor used by both ABIs.
+struct WmmaNode {
+    StinkyInstruction* inst;
+    RegGroup aGroup;  ///< src0 (A matrix tile)
+    RegGroup bGroup;  ///< src1 (B matrix tile)
+    RegGroup cGroup;  ///< dest (C accumulator tile)
+};
+
+/// Compute live intervals for A/B register groups referenced by wmma.
+/// Implementations may use wmma-sequence positions (fast) or full-instruction
+/// positions (precise). Callers treat interval values as opaque ordinals.
+class STINKYTOFU_EXPORT IRegLivenessAnalysis {
+   public:
+    virtual ~IRegLivenessAnalysis() = default;
+
+    /// Return the live interval for every RegGroup that appears as an A or B
+    /// source in @p wmmaSeq. The @p bb argument is provided so full-instruction
+    /// backends can inspect all instruction types; wmma-only backends may ignore it.
+    virtual std::map<RegGroup, RegInterval> computeLiveness(
+        const BasicBlock& bb, const std::vector<WmmaNode>& wmmaSeq) const = 0;
+};
+
+/// Fast wmma-only liveness: interval = [first wmma index, last wmma index]
+/// that reads the group. Sufficient for the double-buffered GEMM pattern.
+class STINKYTOFU_EXPORT WmmaIntervalLiveness : public IRegLivenessAnalysis {
+   public:
+    std::map<RegGroup, RegInterval> computeLiveness(
+        const BasicBlock& bb, const std::vector<WmmaNode>& wmmaSeq) const override;
+};
+
+// ─────────────────────────────────────────────────────────────────────────────
+// ABI 2 — Reorder algorithm (swappable)
+// ─────────────────────────────────────────────────────────────────────────────
+
+struct AliasCandidate {
+    RegGroup canonical;  ///< register group to keep (e.g. A_X0[i])
+    RegGroup aliasable;  ///< register group to remap onto canonical (e.g. A_X1[i])
+    unsigned vgprSaved;  ///< == canonical.size
+};
+
+/// Given the wmma sequence and precomputed live intervals, decide the optimal
+/// wmma ordering and which register groups can be aliased. The algorithm does
+/// not know how liveness was computed.
+class STINKYTOFU_EXPORT IWmmaReorderAlgorithm {
+   public:
+    virtual ~IWmmaReorderAlgorithm() = default;
+
+    struct Result {
+        std::vector<WmmaNode> desiredOrder;  ///< permutation of wmmaSeq
+        std::vector<AliasCandidate> aliases;
+    };
+
+    virtual Result solve(const std::vector<std::vector<WmmaNode>>& pools,
+                         const std::map<RegGroup, RegInterval>& liveness) const = 0;
+};
+
+/// Reorders each pool so its pool-varying operand (A after detectABIndices relabeling)
+/// is contiguous, tightening its liveness and enabling cross-pool register aliasing.
+/// O(n_wmma). Covers standard N-buffered GEMM kernels regardless of which hardware
+/// src (A or B) carries the pool-varying registers.
+class STINKYTOFU_EXPORT PoolVaryingReorderAlgorithm : public IWmmaReorderAlgorithm {
+   public:
+    Result solve(const std::vector<std::vector<WmmaNode>>& pools,
+                 const std::map<RegGroup, RegInterval>& liveness) const override;
+};
+
+// ─────────────────────────────────────────────────────────────────────────────
+// Pass factory
+// ─────────────────────────────────────────────────────────────────────────────
+
+/// Analysis-only pass that determines the optimal wmma reordering to reduce
+/// VGPR usage by aliasing A_X1 register groups onto A_X0.
+///
+/// Runs before StinkyDAGSchedulerPass. Does not mutate any instruction or
+/// register operand — all output is in WmmaReorderAnalysisResult.
+///
+/// The liveness backend and reorder algorithm are injected at construction so
+/// either axis can be swapped independently without touching consumers.
+STINKYTOFU_EXPORT std::unique_ptr<Pass> createStinkyWmmaVgprReorderPass(
+    std::unique_ptr<IRegLivenessAnalysis> liveness = nullptr,
+    std::unique_ptr<IWmmaReorderAlgorithm> algorithm = nullptr);
+
+/// Retrieve the analysis result produced for @p bb by the most recent pass run.
+/// Returns nullptr if the pass has not run or the BB was not processed.
+STINKYTOFU_EXPORT const WmmaReorderAnalysisResult* getWmmaReorderResult(const BasicBlock& bb);
+
+}  // namespace stinkytofu

shared/stinkytofu/src/transforms/asm/CMakeLists.txt

@@ -32,4 +32,5 @@ target_sources(stinkytofu_objs PRIVATE
     InsertWaitAluPass.cpp
     RederiveExpertScopePass.cpp
     EstimateAsmCyclesPass.cpp
+    StinkyWmmaVgprReorderPass.cpp
 )