[llvm-branch-commits] [llvm] [AMDGPU] Count WMMA scale copies against both operands (PR #234104)

[Kazu Hirata via llvm-branch-commits]

================
@@ -210,34 +281,119 @@ bool SIMaterializeWMMAScale::run(MachineFunction &MF,
   MLI = &Loops;
   ExecLoopCache.clear();
 
-  // Insertion order keeps copy placement deterministic. Subregister lanes of
-  // one tuple are separate scales.
-  MapVector<std::pair<Register, unsigned>, SmallVector<ScaleUse, 4>> ScaleUses;
+  // Insertion order keeps copy placement deterministic. Each lane is its own
+  // scale. A physical SGPR, or a subregister with no legal VGPR copy, still
+  // occupies va_ssrc and is not a key.
+  MapVector<ScaleKey, SmallVector<ScaleUse, 4>> ScaleUses;
+  DenseMap<ScaleKey, unsigned> KeyIndex;
+  SmallVector<const MachineBasicBlock *, 8> KeyDefBB;
+  SmallVector<ScaleCandidate, 8> Candidates;
+
+  auto canCopy = [&](const MachineOperand &MO, const MachineBasicBlock *UseBB,
+                     const MachineBasicBlock *&DefBB) {
+    if (!MO.isReg() || !MO.getReg().isVirtual())
+      return false;
+    MachineInstr *Def = MRI->getVRegDef(MO.getReg());
+    if (!Def || Def->isTerminator())
+      return false;
+    unsigned SubReg = MO.getSubReg();
+    unsigned BitWidth =
+        SubReg ? TRI->getSubRegIdxSize(SubReg)
+               : TRI->getRegSizeInBits(*MRI->getRegClass(MO.getReg()))
+                     .getFixedValue();
+    if (!TRI->getAlignedLo256VGPRClassForBitWidth(BitWidth))
+      return false;
+    DefBB = Def->getParent();
+    return !crossesUnsafeDivergentLoop(DefBB, UseBB);
+  };
+
   for (MachineBasicBlock &MBB : MF) {
     for (MachineInstr &MI : MBB) {
       if (!SIInstrInfo::isWMMA(MI) ||
           !AMDGPU::hasNamedOperand(MI.getOpcode(), AMDGPU::OpName::scale_src0))
         continue;
 
+      bool FixedSGPR = false;
+      SmallVector<unsigned, 2> Keys;
       for (AMDGPU::OpName Name :
            {AMDGPU::OpName::scale_src0, AMDGPU::OpName::scale_src1}) {
         int Idx = AMDGPU::getNamedOperandIdx(MI.getOpcode(), Name);
         if (Idx < 0)
           continue;
         MachineOperand &MO = MI.getOperand(Idx);
-        if (!MO.isReg() || !MO.getReg().isVirtual())
+        if (!MO.isReg() || !MO.getReg() || !TRI->isSGPRReg(*MRI, MO.getReg()))
           continue;
-        if (!TRI->isSGPRReg(*MRI, MO.getReg()))
+        const MachineBasicBlock *DefBB = nullptr;
+        if (!canCopy(MO, MI.getParent(), DefBB)) {
+          FixedSGPR = true;
           continue;
-        ScaleUses[{MO.getReg(), MO.getSubReg()}].push_back(
-            {&MI, static_cast<unsigned>(Idx)});
+        }
+        ScaleKey Key{MO.getReg(), MO.getSubReg()};
+        auto [It, Inserted] = KeyIndex.try_emplace(Key, KeyIndex.size());
+        if (Inserted) {
+          ScaleUses[Key] = {};
+          KeyDefBB.push_back(DefBB);
+        }
----------------
kazutakahirata wrote:

May I suggest removing `ScaleUses[Key] = {};` since `ScaleUses[Key].push_back(...)` on line 337 already default-constructs the entry on first access?

```suggestion
        if (Inserted)
          KeyDefBB.push_back(DefBB);
```

https://github.com/llvm/llvm-project/pull/234104