[llvm] [AMDGPU] Fix illegal AGPR reclassification in RewriteMFMAFormStage (PR #200972)

Dmitry Sidorov via llvm-commits llvm-commits at lists.llvm.org
Fri Jun 12 06:27:26 PDT 2026


https://github.com/MrSidims updated https://github.com/llvm/llvm-project/pull/200972

>From 938a973818b8bac8356b3a8aa7abce230d5f2aff Mon Sep 17 00:00:00 2001
From: Dmitry Sidorov <Dmitry.Sidorov at amd.com>
Date: Mon, 1 Jun 2026 02:44:51 +0200
Subject: [PATCH 1/5] [AMDGPU] Fix illegal AGPR reclassification in
 RewriteMFMAFormStage

If src2 escapes rewrite group then bridge copy AGPR -> VGPR must be inserted.
---
 llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp   |  36 +++-
 .../AMDGPU/sched_mfma_rewrite_copies.mir      | 196 ++++++++++++++++++
 2 files changed, 229 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
index 11b783a1024da..5f1e640d4aad6 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
@@ -2251,10 +2251,11 @@ void GCNSchedStage::modifyRegionSchedule(unsigned RegionIdx,
 /// Returns true when \p RD will already be in AGPR-form after the rewrite, so
 /// no bridge copy is needed at this reaching definition.
 static bool isReachingDefAGPRForm(MachineInstr *RD,
+                                  const SmallPtrSetImpl<MachineInstr *> &RewriteSet,
                                   const DenseSet<Register> &CandSrc2Regs,
                                   const SIInstrInfo &TII) {
   if (TII.isMAI(*RD))
-    return true;
+    return RewriteSet.contains(RD);
   if (RD->getOpcode() == AMDGPU::AV_MOV_B32_IMM_PSEUDO ||
       RD->getOpcode() == AMDGPU::AV_MOV_B64_IMM_PSEUDO)
     return true;
@@ -2263,6 +2264,23 @@ static bool isReachingDefAGPRForm(MachineInstr *RD,
   return false;
 }
 
+/// Returns true when \p Reg has a use that requires it to stay in VGPR form
+/// after the rewrite, i.e. any use other than a group MFMA.
+static bool
+src2RegEscapesRewriteGroup(Register Reg, const MachineRegisterInfo &MRI,
+                           const SmallPtrSetImpl<MachineInstr *> &RewriteSet,
+                           const SIInstrInfo &TII) {
+  for (const MachineOperand &Use : MRI.use_nodbg_operands(Reg)) {
+    const MachineInstr *UseMI = Use.getParent();
+    if (UseMI->isCopy())
+      continue;
+    if (TII.isMAI(*UseMI) && RewriteSet.contains(UseMI))
+      continue;
+    return true;
+  }
+  return false;
+}
+
 bool RewriteMFMAFormStage::isRewriteCandidate(MachineInstr *MI) const {
   if (!static_cast<const SIInstrInfo *>(DAG.TII)->isMAI(*MI))
     return false;
@@ -2315,9 +2333,15 @@ bool RewriteMFMAFormStage::initHeuristics(
         SmallVector<SlotIndex, 8> Src2ReachingDefs;
         findReachingDefs(*Src2, DAG.LIS, Src2ReachingDefs);
 
+        // If src2 has a use that must remain VGPR, it cannot be reclassified to
+        // AGPR.
+        bool Src2Escapes = src2RegEscapesRewriteGroup(Src2->getReg(), DAG.MRI,
+                                                      RewriteSet, *TII);
+
         for (SlotIndex RDIdx : Src2ReachingDefs) {
           MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIdx);
-          if (isReachingDefAGPRForm(RD, CandSrc2Regs, *TII))
+          if (!Src2Escapes &&
+              isReachingDefAGPRForm(RD, RewriteSet, CandSrc2Regs, *TII))
             continue;
           CopyForDef.insert(RD);
         }
@@ -2592,9 +2616,15 @@ bool RewriteMFMAFormStage::rewrite(
       findReachingDefs(*Src2, DAG.LIS, Src2ReachingDefs);
       SmallSetVector<MachineInstr *, 8> Src2DefsReplace;
 
+      // If src2 has a use that must remain VGPR, it cannot be reclassified to
+      // AGPR.
+      bool Src2Escapes =
+          src2RegEscapesRewriteGroup(Src2Reg, DAG.MRI, RewriteCandsSet, *TII);
+
       for (SlotIndex RDIndex : Src2ReachingDefs) {
         MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIndex);
-        if (isReachingDefAGPRForm(RD, RewriteSrc2Regs, *TII))
+        if (!Src2Escapes &&
+            isReachingDefAGPRForm(RD, RewriteCandsSet, RewriteSrc2Regs, *TII))
           continue;
 
         Src2DefsReplace.insert(RD);
diff --git a/llvm/test/CodeGen/AMDGPU/sched_mfma_rewrite_copies.mir b/llvm/test/CodeGen/AMDGPU/sched_mfma_rewrite_copies.mir
index 2238914b24e7d..b2f6536e06b7f 100644
--- a/llvm/test/CodeGen/AMDGPU/sched_mfma_rewrite_copies.mir
+++ b/llvm/test/CodeGen/AMDGPU/sched_mfma_rewrite_copies.mir
@@ -197,6 +197,16 @@
     unreachable
   }
 
+  define void @noncandidate_mfma_def_feeds_candidate_src2() #0 {
+  entry:
+    unreachable
+  }
+
+  define void @av_mov_imm_src2_def_escapes_to_vgpr_user() #0 {
+  entry:
+    unreachable
+  }
+
   attributes #0 = { "amdgpu-waves-per-eu"="1,1" "amdgpu-flat-work-group-size"="64,64"}
 ...
 
@@ -5423,3 +5433,189 @@ body:             |
     S_NOP 0, implicit %50, implicit %51
     S_ENDPGM 0
 ...
+
+---
+name:            noncandidate_mfma_def_feeds_candidate_src2
+tracksRegLiveness: true
+machineFunctionInfo:
+  isEntryFunction: true
+  scratchRSrcReg:  '$sgpr96_sgpr97_sgpr98_sgpr99'
+  stackPtrOffsetReg: '$sgpr32'
+  argumentInfo:
+    privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+    kernargSegmentPtr: { reg: '$sgpr4_sgpr5' }
+    workGroupIDX:    { reg: '$sgpr6' }
+    privateSegmentWaveByteOffset: { reg: '$sgpr7' }
+    workItemIDX:     { reg: '$vgpr0' }
+  sgprForEXECCopy: '$sgpr100_sgpr101'
+body:             |
+  ; CHECK-LABEL: name: noncandidate_mfma_def_feeds_candidate_src2
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.1(0x80000000)
+  ; CHECK-NEXT:   liveins: $vgpr0, $sgpr4_sgpr5
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[DEF:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF1:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF2:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF3:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF4:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF5:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF6:%[0-9]+]]:vreg_512 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF7:%[0-9]+]]:vreg_64 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF8:%[0-9]+]]:vgpr_32 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF9:%[0-9]+]]:vreg_128 = IMPLICIT_DEF
+  ; CHECK-NEXT:   S_NOP 0, implicit-def %12
+  ; CHECK-NEXT:   S_NOP 0, implicit-def %13
+  ; CHECK-NEXT:   [[DEF10:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   SCHED_BARRIER 0
+  ; CHECK-NEXT:   [[DEF11:%[0-9]+]]:av_128_align2 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF12:%[0-9]+]]:av_128_align2 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF13:%[0-9]+]]:vreg_128_align2 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF14:%[0-9]+]]:vreg_64_align2 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF15:%[0-9]+]]:vgpr_32 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF16:%[0-9]+]]:vreg_128_align2 = IMPLICIT_DEF
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   successors: %bb.2(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64_:%[0-9]+]]:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 [[DEF11]], [[DEF12]], [[DEF16]], 4, 4, [[DEF14]].sub0, [[DEF15]], 0, 0, implicit $mode, implicit $exec
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:areg_128_align2 = COPY [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64_]]
+  ; CHECK-NEXT:   [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_e64_:%[0-9]+]]:areg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_e64 [[DEF11]], [[DEF12]], [[COPY]], 4, 4, [[DEF14]].sub0, [[DEF15]], 0, 0, implicit $mode, implicit $exec
+  ; CHECK-NEXT:   [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_e64_1:%[0-9]+]]:areg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_e64 [[DEF11]], [[DEF12]], [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_e64_]], 4, 4, [[DEF14]].sub0, [[DEF15]], 0, 0, implicit $mode, implicit $exec
+  ; CHECK-NEXT:   [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_e64_2:%[0-9]+]]:areg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_e64 [[DEF11]], [[DEF12]], [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_e64_1]], 4, 4, [[DEF14]].sub0, [[DEF15]], 0, 0, implicit $mode, implicit $exec
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   undef [[V_ADD_U32_e32_:%[0-9]+]].sub0:vreg_128_align2 = V_ADD_U32_e32 [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64_]].sub0, [[DEF15]], implicit $exec
+  ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:vreg_128_align2 = COPY [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_e64_2]]
+  ; CHECK-NEXT:   [[COPY2:%[0-9]+]]:vreg_128_align2 = COPY [[COPY1]]
+  ; CHECK-NEXT:   SCHED_BARRIER 0
+  ; CHECK-NEXT:   [[DEF17:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   KILL [[DEF17]], [[DEF]], [[DEF1]], [[DEF2]], [[DEF3]], [[DEF4]], [[DEF5]], [[DEF6]], [[DEF7]], [[DEF8]], [[DEF9]], [[DEF10]], [[DEF13]], [[DEF16]], [[V_ADD_U32_e32_]], [[COPY2]]
+  ; CHECK-NEXT:   S_NOP 0, implicit %12, implicit %13
+  ; CHECK-NEXT:   S_ENDPGM 0
+  bb.0:
+   liveins: $vgpr0, $sgpr4_sgpr5
+    %1:vreg_1024 = IMPLICIT_DEF
+    %2:vreg_1024 = IMPLICIT_DEF
+    %3:vreg_1024 = IMPLICIT_DEF
+    %4:vreg_1024 = IMPLICIT_DEF
+    %5:vreg_1024 = IMPLICIT_DEF
+    %6:vreg_1024 = IMPLICIT_DEF
+    %7:vreg_1024 = IMPLICIT_DEF
+    %8:vreg_512 = IMPLICIT_DEF
+    %10:vreg_64 = IMPLICIT_DEF
+    %11:vgpr_32 = IMPLICIT_DEF
+    %12:vreg_128 = IMPLICIT_DEF
+    %13:vreg_1024 = IMPLICIT_DEF
+    S_NOP 0, implicit-def %50:av_512
+    S_NOP 0, implicit-def %51:av_512
+    SCHED_BARRIER 0
+    %60:av_128_align2 = IMPLICIT_DEF
+    %61:av_128_align2 = IMPLICIT_DEF
+    %62:vreg_128_align2 = IMPLICIT_DEF
+    %63:vreg_64_align2 = IMPLICIT_DEF
+    %64:vgpr_32 = IMPLICIT_DEF
+    %72:vreg_128_align2 = IMPLICIT_DEF
+
+  bb.1:
+    %85:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 %60, %61, %72, 4, 4, %63.sub0, %64, 0, 0, implicit $mode, implicit $exec
+    %86:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 %60, %61, %85, 4, 4, %63.sub0, %64, 0, 0, implicit $mode, implicit $exec
+    %87:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 %60, %61, %86, 4, 4, %63.sub0, %64, 0, 0, implicit $mode, implicit $exec
+    %88:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 %60, %61, %87, 4, 4, %63.sub0, %64, 0, 0, implicit $mode, implicit $exec
+
+  bb.2:
+    undef %94.sub0:vreg_128_align2 = V_ADD_U32_e32 %85.sub0, %64, implicit $exec
+    %95:vreg_128_align2 = COPY %88
+    SCHED_BARRIER 0
+    KILL %1, %2, %3, %4, %5, %6, %7, %8, %10, %11, %12, %13, %62, %72, %94, %95
+    S_NOP 0, implicit %50, implicit %51
+    S_ENDPGM 0
+...
+
+---
+name:            av_mov_imm_src2_def_escapes_to_vgpr_user
+tracksRegLiveness: true
+machineFunctionInfo:
+  isEntryFunction: true
+  scratchRSrcReg:  '$sgpr96_sgpr97_sgpr98_sgpr99'
+  stackPtrOffsetReg: '$sgpr32'
+  argumentInfo:
+    privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+    kernargSegmentPtr: { reg: '$sgpr4_sgpr5' }
+    workGroupIDX:    { reg: '$sgpr6' }
+    privateSegmentWaveByteOffset: { reg: '$sgpr7' }
+    workItemIDX:     { reg: '$vgpr0' }
+  sgprForEXECCopy: '$sgpr100_sgpr101'
+body:             |
+  ; CHECK-LABEL: name: av_mov_imm_src2_def_escapes_to_vgpr_user
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.1(0x80000000)
+  ; CHECK-NEXT:   liveins: $vgpr0, $sgpr4_sgpr5
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[DEF:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF1:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF2:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF3:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF4:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF5:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF6:%[0-9]+]]:vreg_512 = IMPLICIT_DEF
+  ; CHECK-NEXT:   S_NOP 0, implicit-def %9
+  ; CHECK-NEXT:   S_NOP 0, implicit-def %10
+  ; CHECK-NEXT:   [[DEF7:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   SCHED_BARRIER 0
+  ; CHECK-NEXT:   [[DEF8:%[0-9]+]]:vreg_64_align2 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF9:%[0-9]+]]:vreg_64_align2 = IMPLICIT_DEF
+  ; CHECK-NEXT:   %acc:vreg_64_align2 = AV_MOV_B64_IMM_PSEUDO 0, implicit $exec
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:areg_64_align2 = COPY %acc
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   successors: %bb.2(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[V_MFMA_F64_4X4X4F64_e64_:%[0-9]+]]:areg_64_align2 = contract nofpexcept V_MFMA_F64_4X4X4F64_e64 [[DEF8]], [[DEF9]], [[COPY]], 0, 0, 0, implicit $mode, implicit $exec
+  ; CHECK-NEXT:   [[V_MFMA_F64_4X4X4F64_e64_1:%[0-9]+]]:areg_64_align2 = contract nofpexcept V_MFMA_F64_4X4X4F64_e64 [[DEF8]], [[DEF9]], [[V_MFMA_F64_4X4X4F64_e64_]], 0, 0, 0, implicit $mode, implicit $exec
+  ; CHECK-NEXT:   [[V_MFMA_F64_4X4X4F64_e64_2:%[0-9]+]]:areg_64_align2 = contract nofpexcept V_MFMA_F64_4X4X4F64_e64 [[DEF8]], [[DEF9]], [[V_MFMA_F64_4X4X4F64_e64_1]], 0, 0, 0, implicit $mode, implicit $exec
+  ; CHECK-NEXT:   [[V_MFMA_F64_4X4X4F64_e64_3:%[0-9]+]]:areg_64_align2 = contract nofpexcept V_MFMA_F64_4X4X4F64_e64 [[DEF8]], [[DEF9]], [[V_MFMA_F64_4X4X4F64_e64_2]], 0, 0, 0, implicit $mode, implicit $exec
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   [[DEF10:%[0-9]+]]:vgpr_32 = IMPLICIT_DEF
+  ; CHECK-NEXT:   undef [[V_ADD_U32_e32_:%[0-9]+]].sub0:vreg_64_align2 = V_ADD_U32_e32 %acc.sub0, [[DEF10]], implicit $exec
+  ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:vreg_64_align2 = COPY [[V_MFMA_F64_4X4X4F64_e64_3]]
+  ; CHECK-NEXT:   [[COPY2:%[0-9]+]]:vreg_64_align2 = COPY [[COPY1]]
+  ; CHECK-NEXT:   SCHED_BARRIER 0
+  ; CHECK-NEXT:   [[DEF11:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   KILL [[DEF11]], [[DEF]], [[DEF1]], [[DEF2]], [[DEF3]], [[DEF4]], [[DEF5]], [[DEF6]], [[DEF7]], [[V_ADD_U32_e32_]], [[COPY2]]
+  ; CHECK-NEXT:   S_NOP 0, implicit %9, implicit %10
+  ; CHECK-NEXT:   S_ENDPGM 0
+  bb.0:
+   liveins: $vgpr0, $sgpr4_sgpr5
+    %1:vreg_1024 = IMPLICIT_DEF
+    %2:vreg_1024 = IMPLICIT_DEF
+    %3:vreg_1024 = IMPLICIT_DEF
+    %4:vreg_1024 = IMPLICIT_DEF
+    %5:vreg_1024 = IMPLICIT_DEF
+    %6:vreg_1024 = IMPLICIT_DEF
+    %7:vreg_1024 = IMPLICIT_DEF
+    %8:vreg_512 = IMPLICIT_DEF
+    %13:vreg_1024 = IMPLICIT_DEF
+    S_NOP 0, implicit-def %50:av_512
+    S_NOP 0, implicit-def %51:av_512
+    SCHED_BARRIER 0
+    %60:vreg_64_align2 = IMPLICIT_DEF
+    %61:vreg_64_align2 = IMPLICIT_DEF
+    %64:vgpr_32 = IMPLICIT_DEF
+    %acc:vreg_64_align2 = AV_MOV_B64_IMM_PSEUDO 0, implicit $exec
+
+  bb.1:
+    %85:vreg_64_align2 = contract nofpexcept V_MFMA_F64_4X4X4F64_vgprcd_e64 %60, %61, %acc, 0, 0, 0, implicit $mode, implicit $exec
+    %86:vreg_64_align2 = contract nofpexcept V_MFMA_F64_4X4X4F64_vgprcd_e64 %60, %61, %85, 0, 0, 0, implicit $mode, implicit $exec
+    %87:vreg_64_align2 = contract nofpexcept V_MFMA_F64_4X4X4F64_vgprcd_e64 %60, %61, %86, 0, 0, 0, implicit $mode, implicit $exec
+    %88:vreg_64_align2 = contract nofpexcept V_MFMA_F64_4X4X4F64_vgprcd_e64 %60, %61, %87, 0, 0, 0, implicit $mode, implicit $exec
+
+  bb.2:
+    undef %94.sub0:vreg_64_align2 = V_ADD_U32_e32 %acc.sub0, %64, implicit $exec
+    %95:vreg_64_align2 = COPY %88
+    SCHED_BARRIER 0
+    KILL %1, %2, %3, %4, %5, %6, %7, %8, %13, %94, %95
+    S_NOP 0, implicit %50, implicit %51
+    S_ENDPGM 0
+...

>From 24266e8dad215c49f73843b1c16df9a0bd113a1e Mon Sep 17 00:00:00 2001
From: Dmitry Sidorov <Dmitry.Sidorov at amd.com>
Date: Tue, 2 Jun 2026 01:50:17 +0200
Subject: [PATCH 2/5] format

---
 llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
index 5f1e640d4aad6..71a5a7cc3c094 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
@@ -2250,10 +2250,9 @@ void GCNSchedStage::modifyRegionSchedule(unsigned RegionIdx,
 
 /// Returns true when \p RD will already be in AGPR-form after the rewrite, so
 /// no bridge copy is needed at this reaching definition.
-static bool isReachingDefAGPRForm(MachineInstr *RD,
-                                  const SmallPtrSetImpl<MachineInstr *> &RewriteSet,
-                                  const DenseSet<Register> &CandSrc2Regs,
-                                  const SIInstrInfo &TII) {
+static bool isReachingDefAGPRForm(
+    MachineInstr *RD, const SmallPtrSetImpl<MachineInstr *> &RewriteSet,
+    const DenseSet<Register> &CandSrc2Regs, const SIInstrInfo &TII) {
   if (TII.isMAI(*RD))
     return RewriteSet.contains(RD);
   if (RD->getOpcode() == AMDGPU::AV_MOV_B32_IMM_PSEUDO ||

>From 9c39d4e3253e7e123aee7fadd46d8dcf3fa261bb Mon Sep 17 00:00:00 2001
From: Dmitry Sidorov <Dmitry.Sidorov at amd.com>
Date: Tue, 2 Jun 2026 14:44:59 +0200
Subject: [PATCH 3/5] apply comments

---
 llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp   | 40 ++++-----
 llvm/lib/Target/AMDGPU/GCNSchedStrategy.h     |  6 ++
 .../AMDGPU/sched_mfma_rewrite_copies.mir      | 84 +++++++++----------
 3 files changed, 69 insertions(+), 61 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
index 71a5a7cc3c094..5c509e45b1596 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
@@ -2248,8 +2248,10 @@ void GCNSchedStage::modifyRegionSchedule(unsigned RegionIdx,
   DAG.Regions[RegionIdx].first = MIOrder.front();
 }
 
-/// Returns true when \p RD will already be in AGPR-form after the rewrite, so
-/// no bridge copy is needed at this reaching definition.
+/// Returns true if reaching def \p RD will be in AGPR form after the rewrite
+/// and so needs no bridge copy: a candidate MFMA in \p RewriteSet, an
+/// AV_MOV_*_IMM_PSEUDO, or a copy from a candidate src2 reg in \p CandSrc2Regs.
+/// A non-candidate MFMA stays in VGPR form and still needs a bridge.
 static bool isReachingDefAGPRForm(
     MachineInstr *RD, const SmallPtrSetImpl<MachineInstr *> &RewriteSet,
     const DenseSet<Register> &CandSrc2Regs, const SIInstrInfo &TII) {
@@ -2263,19 +2265,21 @@ static bool isReachingDefAGPRForm(
   return false;
 }
 
-/// Returns true when \p Reg has a use that requires it to stay in VGPR form
-/// after the rewrite, i.e. any use other than a group MFMA.
-static bool
-src2RegEscapesRewriteGroup(Register Reg, const MachineRegisterInfo &MRI,
-                           const SmallPtrSetImpl<MachineInstr *> &RewriteSet,
-                           const SIInstrInfo &TII) {
-  for (const MachineOperand &Use : MRI.use_nodbg_operands(Reg)) {
-    const MachineInstr *UseMI = Use.getParent();
-    if (UseMI->isCopy())
-      continue;
-    if (TII.isMAI(*UseMI) && RewriteSet.contains(UseMI))
-      continue;
-    return true;
+bool RewriteMFMAFormStage::hasEscapingUse(
+    ArrayRef<SlotIndex> Src2ReachingDefs,
+    const SmallPtrSetImpl<MachineInstr *> &RewriteSet) {
+  for (SlotIndex RDIdx : Src2ReachingDefs) {
+    MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIdx);
+    SmallVector<MachineOperand *, 8> ReachingUses;
+    findReachingUses(RD, DAG.LIS, ReachingUses);
+    for (MachineOperand *UseMO : ReachingUses) {
+      MachineInstr *UseMI = UseMO->getParent();
+      if (UseMI->isCopy())
+        continue;
+      if (TII->isMAI(*UseMI) && RewriteSet.contains(UseMI))
+        continue;
+      return true;
+    }
   }
   return false;
 }
@@ -2334,8 +2338,7 @@ bool RewriteMFMAFormStage::initHeuristics(
 
         // If src2 has a use that must remain VGPR, it cannot be reclassified to
         // AGPR.
-        bool Src2Escapes = src2RegEscapesRewriteGroup(Src2->getReg(), DAG.MRI,
-                                                      RewriteSet, *TII);
+        bool Src2Escapes = hasEscapingUse(Src2ReachingDefs, RewriteSet);
 
         for (SlotIndex RDIdx : Src2ReachingDefs) {
           MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIdx);
@@ -2617,8 +2620,7 @@ bool RewriteMFMAFormStage::rewrite(
 
       // If src2 has a use that must remain VGPR, it cannot be reclassified to
       // AGPR.
-      bool Src2Escapes =
-          src2RegEscapesRewriteGroup(Src2Reg, DAG.MRI, RewriteCandsSet, *TII);
+      bool Src2Escapes = hasEscapingUse(Src2ReachingDefs, RewriteCandsSet);
 
       for (SlotIndex RDIndex : Src2ReachingDefs) {
         MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIndex);
diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
index 3682f39c8259b..87945afa376e4 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
@@ -479,6 +479,12 @@ class RewriteMFMAFormStage : public GCNSchedStage {
   void findReachingUses(MachineInstr *DefMI, LiveIntervals *LIS,
                         SmallVectorImpl<MachineOperand *> &ReachingUses);
 
+  /// Returns true if the src2 register with reaching defs \p Src2ReachingDefs
+  /// has a use other than a group MFMA (in \p RewriteSet) or a copy, which
+  /// would keep it in VGPR form rather than let it be reclassified to AGPR.
+  bool hasEscapingUse(ArrayRef<SlotIndex> Src2ReachingDefs,
+                      const SmallPtrSetImpl<MachineInstr *> &RewriteSet);
+
 public:
   bool initGCNSchedStage() override;
 
diff --git a/llvm/test/CodeGen/AMDGPU/sched_mfma_rewrite_copies.mir b/llvm/test/CodeGen/AMDGPU/sched_mfma_rewrite_copies.mir
index b2f6536e06b7f..71838881aee99 100644
--- a/llvm/test/CodeGen/AMDGPU/sched_mfma_rewrite_copies.mir
+++ b/llvm/test/CodeGen/AMDGPU/sched_mfma_rewrite_copies.mir
@@ -5495,40 +5495,40 @@ body:             |
   ; CHECK-NEXT:   S_ENDPGM 0
   bb.0:
    liveins: $vgpr0, $sgpr4_sgpr5
+    %0:vreg_1024 = IMPLICIT_DEF
     %1:vreg_1024 = IMPLICIT_DEF
     %2:vreg_1024 = IMPLICIT_DEF
     %3:vreg_1024 = IMPLICIT_DEF
     %4:vreg_1024 = IMPLICIT_DEF
     %5:vreg_1024 = IMPLICIT_DEF
     %6:vreg_1024 = IMPLICIT_DEF
-    %7:vreg_1024 = IMPLICIT_DEF
-    %8:vreg_512 = IMPLICIT_DEF
-    %10:vreg_64 = IMPLICIT_DEF
-    %11:vgpr_32 = IMPLICIT_DEF
-    %12:vreg_128 = IMPLICIT_DEF
-    %13:vreg_1024 = IMPLICIT_DEF
-    S_NOP 0, implicit-def %50:av_512
-    S_NOP 0, implicit-def %51:av_512
+    %7:vreg_512 = IMPLICIT_DEF
+    %8:vreg_64 = IMPLICIT_DEF
+    %9:vgpr_32 = IMPLICIT_DEF
+    %10:vreg_128 = IMPLICIT_DEF
+    %11:vreg_1024 = IMPLICIT_DEF
+    S_NOP 0, implicit-def %12:av_512
+    S_NOP 0, implicit-def %13:av_512
     SCHED_BARRIER 0
-    %60:av_128_align2 = IMPLICIT_DEF
-    %61:av_128_align2 = IMPLICIT_DEF
-    %62:vreg_128_align2 = IMPLICIT_DEF
-    %63:vreg_64_align2 = IMPLICIT_DEF
-    %64:vgpr_32 = IMPLICIT_DEF
-    %72:vreg_128_align2 = IMPLICIT_DEF
+    %14:av_128_align2 = IMPLICIT_DEF
+    %15:av_128_align2 = IMPLICIT_DEF
+    %16:vreg_128_align2 = IMPLICIT_DEF
+    %17:vreg_64_align2 = IMPLICIT_DEF
+    %18:vgpr_32 = IMPLICIT_DEF
+    %19:vreg_128_align2 = IMPLICIT_DEF
 
   bb.1:
-    %85:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 %60, %61, %72, 4, 4, %63.sub0, %64, 0, 0, implicit $mode, implicit $exec
-    %86:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 %60, %61, %85, 4, 4, %63.sub0, %64, 0, 0, implicit $mode, implicit $exec
-    %87:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 %60, %61, %86, 4, 4, %63.sub0, %64, 0, 0, implicit $mode, implicit $exec
-    %88:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 %60, %61, %87, 4, 4, %63.sub0, %64, 0, 0, implicit $mode, implicit $exec
+    %20:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 %14, %15, %19, 4, 4, %17.sub0, %18, 0, 0, implicit $mode, implicit $exec
+    %21:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 %14, %15, %20, 4, 4, %17.sub0, %18, 0, 0, implicit $mode, implicit $exec
+    %22:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 %14, %15, %21, 4, 4, %17.sub0, %18, 0, 0, implicit $mode, implicit $exec
+    %23:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 %14, %15, %22, 4, 4, %17.sub0, %18, 0, 0, implicit $mode, implicit $exec
 
   bb.2:
-    undef %94.sub0:vreg_128_align2 = V_ADD_U32_e32 %85.sub0, %64, implicit $exec
-    %95:vreg_128_align2 = COPY %88
+    undef %24.sub0:vreg_128_align2 = V_ADD_U32_e32 %20.sub0, %18, implicit $exec
+    %25:vreg_128_align2 = COPY %23
     SCHED_BARRIER 0
-    KILL %1, %2, %3, %4, %5, %6, %7, %8, %10, %11, %12, %13, %62, %72, %94, %95
-    S_NOP 0, implicit %50, implicit %51
+    KILL %0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %16, %19, %24, %25
+    S_NOP 0, implicit %12, implicit %13
     S_ENDPGM 0
 ...
 
@@ -5565,8 +5565,8 @@ body:             |
   ; CHECK-NEXT:   SCHED_BARRIER 0
   ; CHECK-NEXT:   [[DEF8:%[0-9]+]]:vreg_64_align2 = IMPLICIT_DEF
   ; CHECK-NEXT:   [[DEF9:%[0-9]+]]:vreg_64_align2 = IMPLICIT_DEF
-  ; CHECK-NEXT:   %acc:vreg_64_align2 = AV_MOV_B64_IMM_PSEUDO 0, implicit $exec
-  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:areg_64_align2 = COPY %acc
+  ; CHECK-NEXT:   [[AV_MOV_:%[0-9]+]]:vreg_64_align2 = AV_MOV_B64_IMM_PSEUDO 0, implicit $exec
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:areg_64_align2 = COPY [[AV_MOV_]]
   ; CHECK-NEXT: {{  $}}
   ; CHECK-NEXT: bb.1:
   ; CHECK-NEXT:   successors: %bb.2(0x80000000)
@@ -5578,7 +5578,7 @@ body:             |
   ; CHECK-NEXT: {{  $}}
   ; CHECK-NEXT: bb.2:
   ; CHECK-NEXT:   [[DEF10:%[0-9]+]]:vgpr_32 = IMPLICIT_DEF
-  ; CHECK-NEXT:   undef [[V_ADD_U32_e32_:%[0-9]+]].sub0:vreg_64_align2 = V_ADD_U32_e32 %acc.sub0, [[DEF10]], implicit $exec
+  ; CHECK-NEXT:   undef [[V_ADD_U32_e32_:%[0-9]+]].sub0:vreg_64_align2 = V_ADD_U32_e32 [[AV_MOV_]].sub0, [[DEF10]], implicit $exec
   ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:vreg_64_align2 = COPY [[V_MFMA_F64_4X4X4F64_e64_3]]
   ; CHECK-NEXT:   [[COPY2:%[0-9]+]]:vreg_64_align2 = COPY [[COPY1]]
   ; CHECK-NEXT:   SCHED_BARRIER 0
@@ -5588,34 +5588,34 @@ body:             |
   ; CHECK-NEXT:   S_ENDPGM 0
   bb.0:
    liveins: $vgpr0, $sgpr4_sgpr5
+    %0:vreg_1024 = IMPLICIT_DEF
     %1:vreg_1024 = IMPLICIT_DEF
     %2:vreg_1024 = IMPLICIT_DEF
     %3:vreg_1024 = IMPLICIT_DEF
     %4:vreg_1024 = IMPLICIT_DEF
     %5:vreg_1024 = IMPLICIT_DEF
     %6:vreg_1024 = IMPLICIT_DEF
-    %7:vreg_1024 = IMPLICIT_DEF
-    %8:vreg_512 = IMPLICIT_DEF
-    %13:vreg_1024 = IMPLICIT_DEF
-    S_NOP 0, implicit-def %50:av_512
-    S_NOP 0, implicit-def %51:av_512
+    %7:vreg_512 = IMPLICIT_DEF
+    %8:vreg_1024 = IMPLICIT_DEF
+    S_NOP 0, implicit-def %9:av_512
+    S_NOP 0, implicit-def %10:av_512
     SCHED_BARRIER 0
-    %60:vreg_64_align2 = IMPLICIT_DEF
-    %61:vreg_64_align2 = IMPLICIT_DEF
-    %64:vgpr_32 = IMPLICIT_DEF
-    %acc:vreg_64_align2 = AV_MOV_B64_IMM_PSEUDO 0, implicit $exec
+    %11:vreg_64_align2 = IMPLICIT_DEF
+    %12:vreg_64_align2 = IMPLICIT_DEF
+    %13:vgpr_32 = IMPLICIT_DEF
+    %14:vreg_64_align2 = AV_MOV_B64_IMM_PSEUDO 0, implicit $exec
 
   bb.1:
-    %85:vreg_64_align2 = contract nofpexcept V_MFMA_F64_4X4X4F64_vgprcd_e64 %60, %61, %acc, 0, 0, 0, implicit $mode, implicit $exec
-    %86:vreg_64_align2 = contract nofpexcept V_MFMA_F64_4X4X4F64_vgprcd_e64 %60, %61, %85, 0, 0, 0, implicit $mode, implicit $exec
-    %87:vreg_64_align2 = contract nofpexcept V_MFMA_F64_4X4X4F64_vgprcd_e64 %60, %61, %86, 0, 0, 0, implicit $mode, implicit $exec
-    %88:vreg_64_align2 = contract nofpexcept V_MFMA_F64_4X4X4F64_vgprcd_e64 %60, %61, %87, 0, 0, 0, implicit $mode, implicit $exec
+    %15:vreg_64_align2 = contract nofpexcept V_MFMA_F64_4X4X4F64_vgprcd_e64 %11, %12, %14, 0, 0, 0, implicit $mode, implicit $exec
+    %16:vreg_64_align2 = contract nofpexcept V_MFMA_F64_4X4X4F64_vgprcd_e64 %11, %12, %15, 0, 0, 0, implicit $mode, implicit $exec
+    %17:vreg_64_align2 = contract nofpexcept V_MFMA_F64_4X4X4F64_vgprcd_e64 %11, %12, %16, 0, 0, 0, implicit $mode, implicit $exec
+    %18:vreg_64_align2 = contract nofpexcept V_MFMA_F64_4X4X4F64_vgprcd_e64 %11, %12, %17, 0, 0, 0, implicit $mode, implicit $exec
 
   bb.2:
-    undef %94.sub0:vreg_64_align2 = V_ADD_U32_e32 %acc.sub0, %64, implicit $exec
-    %95:vreg_64_align2 = COPY %88
+    undef %19.sub0:vreg_64_align2 = V_ADD_U32_e32 %14.sub0, %13, implicit $exec
+    %20:vreg_64_align2 = COPY %18
     SCHED_BARRIER 0
-    KILL %1, %2, %3, %4, %5, %6, %7, %8, %13, %94, %95
-    S_NOP 0, implicit %50, implicit %51
+    KILL %0, %1, %2, %3, %4, %5, %6, %7, %8, %19, %20
+    S_NOP 0, implicit %9, implicit %10
     S_ENDPGM 0
 ...

>From 4dc96f67ee617f17b5b204d1244409ca4faa38b3 Mon Sep 17 00:00:00 2001
From: Dmitry Sidorov <Dmitry.Sidorov at amd.com>
Date: Sun, 7 Jun 2026 02:22:05 +0200
Subject: [PATCH 4/5] Apply review comments

---
 llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp | 19 ++++++++++---------
 llvm/lib/Target/AMDGPU/GCNSchedStrategy.h   |  6 +++---
 2 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
index d596b4b6e85cf..2dccd595240ba 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
@@ -1318,7 +1318,7 @@ void RewriteMFMAFormStage::findReachingDefs(
 }
 
 void RewriteMFMAFormStage::findReachingUses(
-    MachineInstr *DefMI, LiveIntervals *LIS,
+    const MachineInstr *DefMI, LiveIntervals *LIS,
     SmallVectorImpl<MachineOperand *> &ReachingUses) {
   SlotIndex DefIdx = LIS->getInstructionIndex(*DefMI);
   for (MachineOperand &UseMO :
@@ -2265,15 +2265,15 @@ static bool isReachingDefAGPRForm(
   return false;
 }
 
-bool RewriteMFMAFormStage::hasEscapingUse(
+bool RewriteMFMAFormStage::hasUseRequiringVGPR(
     ArrayRef<SlotIndex> Src2ReachingDefs,
     const SmallPtrSetImpl<MachineInstr *> &RewriteSet) {
   for (SlotIndex RDIdx : Src2ReachingDefs) {
-    MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIdx);
+    const MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIdx);
     SmallVector<MachineOperand *, 8> ReachingUses;
     findReachingUses(RD, DAG.LIS, ReachingUses);
-    for (MachineOperand *UseMO : ReachingUses) {
-      MachineInstr *UseMI = UseMO->getParent();
+    for (const MachineOperand *UseMO : ReachingUses) {
+      const MachineInstr *UseMI = UseMO->getParent();
       if (UseMI->isCopy())
         continue;
       if (TII->isMAI(*UseMI) && RewriteSet.contains(UseMI))
@@ -2361,11 +2361,11 @@ bool RewriteMFMAFormStage::initHeuristics(
 
         // If src2 has a use that must remain VGPR, it cannot be reclassified to
         // AGPR.
-        bool Src2Escapes = hasEscapingUse(Src2ReachingDefs, RewriteSet);
+        bool Src2NeedsVGPR = hasUseRequiringVGPR(Src2ReachingDefs, RewriteSet);
 
         for (SlotIndex RDIdx : Src2ReachingDefs) {
           MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIdx);
-          if (!Src2Escapes &&
+          if (!Src2NeedsVGPR &&
               isReachingDefAGPRForm(RD, RewriteSet, CandSrc2Regs, *TII))
             continue;
           CopyForDef.insert(RD);
@@ -2626,11 +2626,12 @@ bool RewriteMFMAFormStage::rewrite(
 
       // If src2 has a use that must remain VGPR, it cannot be reclassified to
       // AGPR.
-      bool Src2Escapes = hasEscapingUse(Src2ReachingDefs, RewriteCandsSet);
+      bool Src2NeedsVGPR =
+          hasUseRequiringVGPR(Src2ReachingDefs, RewriteCandsSet);
 
       for (SlotIndex RDIndex : Src2ReachingDefs) {
         MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIndex);
-        if (!Src2Escapes &&
+        if (!Src2NeedsVGPR &&
             isReachingDefAGPRForm(RD, RewriteCandsSet, RewriteSrc2Regs, *TII))
           continue;
 
diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
index ff080fa0adb28..0b1ce4797f075 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
@@ -479,14 +479,14 @@ class RewriteMFMAFormStage : public GCNSchedStage {
 
   /// Finds all the reaching uses of \p DefMI and stores the use operands in \p
   /// ReachingUses
-  void findReachingUses(MachineInstr *DefMI, LiveIntervals *LIS,
+  void findReachingUses(const MachineInstr *DefMI, LiveIntervals *LIS,
                         SmallVectorImpl<MachineOperand *> &ReachingUses);
 
   /// Returns true if the src2 register with reaching defs \p Src2ReachingDefs
   /// has a use other than a group MFMA (in \p RewriteSet) or a copy, which
   /// would keep it in VGPR form rather than let it be reclassified to AGPR.
-  bool hasEscapingUse(ArrayRef<SlotIndex> Src2ReachingDefs,
-                      const SmallPtrSetImpl<MachineInstr *> &RewriteSet);
+  bool hasUseRequiringVGPR(ArrayRef<SlotIndex> Src2ReachingDefs,
+                           const SmallPtrSetImpl<MachineInstr *> &RewriteSet);
 
 public:
   bool initGCNSchedStage() override;

>From 2750e6b7239a708e4da79d8ac262a96bb4287fbb Mon Sep 17 00:00:00 2001
From: Dmitry Sidorov <Dmitry.Sidorov at amd.com>
Date: Fri, 12 Jun 2026 14:19:57 +0200
Subject: [PATCH 5/5] cache

---
 llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp | 4 ++--
 llvm/lib/Target/AMDGPU/GCNSchedStrategy.h   | 4 ++++
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
index 2dccd595240ba..dc5628d0fceb5 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
@@ -2362,6 +2362,7 @@ bool RewriteMFMAFormStage::initHeuristics(
         // If src2 has a use that must remain VGPR, it cannot be reclassified to
         // AGPR.
         bool Src2NeedsVGPR = hasUseRequiringVGPR(Src2ReachingDefs, RewriteSet);
+        Src2NeedsVGPRCache[&MI] = Src2NeedsVGPR;
 
         for (SlotIndex RDIdx : Src2ReachingDefs) {
           MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIdx);
@@ -2626,8 +2627,7 @@ bool RewriteMFMAFormStage::rewrite(
 
       // If src2 has a use that must remain VGPR, it cannot be reclassified to
       // AGPR.
-      bool Src2NeedsVGPR =
-          hasUseRequiringVGPR(Src2ReachingDefs, RewriteCandsSet);
+      bool Src2NeedsVGPR = Src2NeedsVGPRCache.lookup(MI);
 
       for (SlotIndex RDIndex : Src2ReachingDefs) {
         MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIndex);
diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
index 0b1ce4797f075..2059f4e6479ff 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
@@ -443,6 +443,10 @@ class RewriteMFMAFormStage : public GCNSchedStage {
   const SIInstrInfo *TII;
   const SIRegisterInfo *SRI;
 
+  /// Per-candidate cache of the src2 "needs VGPR" decision, computed once
+  /// and reused on-demand.
+  DenseMap<const MachineInstr *, bool> Src2NeedsVGPRCache;
+
   /// Do a speculative rewrite and collect copy locations. The speculative
   /// rewrite allows us to calculate the RP of the code after the rewrite, and
   /// the copy locations allow us to calculate the total cost of copies required



More information about the llvm-commits mailing list