[llvm] [AMDGPU] Set WMMA source-operand reuse bits in SIPreEmitPeephole (PR #204050)

[Shilei Tian via llvm-commits]

================
@@ -778,6 +795,78 @@ llvm::SIPreEmitPeepholePass::run(MachineFunction &MF,
   return PreservedAnalyses::all();
 }
 
+bool SIPreEmitPeephole::setWMMAReuseBits(MachineBasicBlock &MBB) const {
+  // Returns true if WMMA MI reads Reg as its A (src0) or B (src1) source.
+  auto WMMAReadsAB = [&](MachineInstr &MI, Register Reg) {
+    for (AMDGPU::OpName Name : {AMDGPU::OpName::src0, AMDGPU::OpName::src1}) {
+      const MachineOperand *MO = TII->getNamedOperand(MI, Name);
+      if (MO && MO->isReg() && MO->getReg() == Reg)
+        return true;
+    }
+    return false;
+  };
+
+  // The A operand corresponds to src0 and the B operand to src1.
+  struct ReuseCand {
+    AMDGPU::OpName BitName;
+    AMDGPU::OpName SrcName;
+    const char *DebugName;
+  };
+  static constexpr ReuseCand Cands[] = {
+      {AMDGPU::OpName::matrix_a_reuse, AMDGPU::OpName::src0, "matrix_a_reuse"},
+      {AMDGPU::OpName::matrix_b_reuse, AMDGPU::OpName::src1, "matrix_b_reuse"}};
+
+  bool Changed = false;
+  for (MachineInstr &W1 : MBB) {
+    if (!SIInstrInfo::isWMMA(W1))
+      continue;
+    unsigned Opc = W1.getOpcode();
+    if (!AMDGPU::hasNamedOperand(Opc, AMDGPU::OpName::matrix_a_reuse) &&
+        !AMDGPU::hasNamedOperand(Opc, AMDGPU::OpName::matrix_b_reuse))
+      continue;
+
+    for (const ReuseCand &C : Cands) {
+      MachineOperand *Bit = TII->getNamedOperand(W1, C.BitName);
+      // Never clear a bit the frontend already requested; only set it.
+      if (!Bit || Bit->getImm() != 0)
+        continue;
+      MachineOperand *Src = TII->getNamedOperand(W1, C.SrcName);
+      if (!Src || !Src->isReg())
+        continue;
+      Register Reg = Src->getReg();
+      if (!Reg.isValid())
+        continue;
+
+      // Scan forward for a WMMA that reuses Reg as an A or B source. The
+      // source-operand cache is only a few entries deep and a high-temporality
+      // entry survives only while few other entries are sticky.
+      unsigned WMMASeen = 0;
+      for (MachineBasicBlock::iterator
+               It = std::next(MachineBasicBlock::iterator(W1)),
+               End = MBB.end();
+           It != End; ++It) {
+        MachineInstr &MI = *It;
+        if (MI.isDebugInstr())
+          continue;
+        bool IsWMMA = SIInstrInfo::isWMMA(MI);
+        if (IsWMMA && WMMAReadsAB(MI, Reg) && !MI.modifiesRegister(Reg, TRI)) {
+          LLVM_DEBUG(dbgs() << "Setting " << C.DebugName << " on " << W1
+                            << "  reused by " << MI);
+          Bit->setImm(1);
+          Changed = true;
+          break;
+        }
+        if (MI.modifiesRegister(Reg, TRI))
----------------
shiltian wrote:

If my question https://github.com/llvm/llvm-project/pull/204050/changes#r3424725190 holds, I think this can be moved above.

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