[AMDGPU] Propagate AGPR RC from PHI to its PHI operands

[rampitec]

Summary

We can fix register class of PHI based on its all AGPR uses.
That leaves behind all PHIs which were already processed
earlier. Propagate RC back to PHI operands of a PHI.

[arsenm]

llvm/lib/Target/AMDGPU/SIFixSGPRCopies.cpp

854

s/themself/themselves

llvm/test/CodeGen/AMDGPU/mfma-loop.ll

534

Should also have a test where a phi directly refers to itself

[rampitec]

llvm/test/CodeGen/AMDGPU/mfma-loop.ll

534

Is it a real case?! This is SSA, value should be defined earlier.

[rampitec]

Fixed comment.

[rtaylor]

LGTM

[rampitec]

llvm/test/CodeGen/AMDGPU/mfma-loop.ll

534

I cannot create a relevant test. If PHI refers itself it does not refer an MFMA inside the loop which is needed to start propagation. I.e. such PHI never triggers this code.

llvm/lib/Target/AMDGPU/SIFixSGPRCopies.cpp

bool SIFixSGPRCopies::runOnMachineFunction(MachineFunction &MF) {
   return true;
 }
 
 void SIFixSGPRCopies::processPHINode(MachineInstr &MI) {
   unsigned numVGPRUses = 0;
   bool AllAGPRUses = true;
   SetVector<const MachineInstr *> worklist;
   SmallSet<const MachineInstr *, 4> Visited;
+  SetVector<MachineInstr *> PHIOperands;
   worklist.insert(&MI);
   Visited.insert(&MI);
   while (!worklist.empty()) {
     const MachineInstr *Instr = worklist.pop_back_val();
     unsigned Reg = Instr->getOperand(0).getReg();
     for (const auto &Use : MRI->use_operands(Reg)) {
       const MachineInstr *UseMI = Use.getParent();
       AllAGPRUses &= (UseMI->isCopy() &&
while (!worklist.empty()) {
     }
   }
 
   Register PHIRes = MI.getOperand(0).getReg();
   const TargetRegisterClass *RC0 = MRI->getRegClass(PHIRes);
   if (AllAGPRUses && numVGPRUses && !TRI->hasAGPRs(RC0)) {
     LLVM_DEBUG(dbgs() << "Moving PHI to AGPR: " << MI);
     MRI->setRegClass(PHIRes, TRI->getEquivalentAGPRClass(RC0));
+    for (unsigned I = 1, N = MI.getNumOperands(); I != N; I += 2) {
+      MachineInstr *DefMI = MRI->getVRegDef(MI.getOperand(I).getReg());
+      if (DefMI && DefMI->isPHI())
+        PHIOperands.insert(DefMI);
+    }
   }
 
   bool hasVGPRInput = false;
   for (unsigned i = 1; i < MI.getNumOperands(); i += 2) {
     unsigned InputReg = MI.getOperand(i).getReg();
     MachineInstr *Def = MRI->getVRegDef(InputReg);
     if (TRI->isVectorRegister(*MRI, InputReg)) {
       if (Def->isCopy()) {
if ((!TRI->isVectorRegister(*MRI, PHIRes) &&
     LLVM_DEBUG(dbgs() << "Fixing PHI: " << MI);
     TII->moveToVALU(MI);
   }
   else {
     LLVM_DEBUG(dbgs() << "Legalizing PHI: " << MI);
     TII->legalizeOperands(MI, MDT);
   }
 
+  // Propagate register class back to PHI operands which are PHI themselves.
+  while (!PHIOperands.empty()) {
+    processPHINode(*PHIOperands.pop_back_val());
+  }
 }

llvm/test/CodeGen/AMDGPU/mfma-loop.ll

for.cond.preheader:
   %cc = icmp eq i32 %inc, 16
   br i1 %cc, label %exit, label %for.cond.preheader
 
 exit:
   store <32 x float> %mai.1, <32 x float> addrspace(1)* %arg
   ret void
 }
 
+; GCN-COUNT-32: v_accvgpr_write_b32 a{{[0-9]+}}, 0{{$}}
+
+; Check that we do not copy agprs to vgprs and back in an outer loop.
+
+; GCN: [[OUTER_LOOP:BB[0-9_]+]]:
+; GCN-NOT: v_accvgpr
+; GCN: [[INNER_LOOP:BB[0-9_]+]]:
+; GCN-NOT: v_accvgpr
+; GCN: v_mfma_f32_32x32x1f32
+; GCN-NOT: v_accvgpr
+; GCN: s_cbranch_scc1 [[INNER_LOOP]]
+; GCN-NOT: v_accvgpr
+; GCN: s_cbranch_scc1 [[OUTER_LOOP]]
+
+; Final result should be read only once after the loop.
+
+; GCN-COUNT-32: v_accvgpr_read_b32
+
+define amdgpu_kernel void @test_mfma_nested_loop_zeroinit(<32 x float> addrspace(1)* %arg) {
+entry:
+  br label %for.cond.preheader
+
+for.cond.preheader:
+  %phi.0 = phi <32 x float> [ zeroinitializer, %entry ], [ %mai.1, %inner.exit ]
+  %c.0 = phi i32 [ 0, %entry ], [ %inc.0, %inner.exit ]
+  br label %inner.for.cond.preheader
+
+inner.for.cond.preheader:
+  %phi = phi <32 x float> [ %phi.0, %for.cond.preheader ], [ %mai.1, %inner.for.cond.preheader ]
+  %c = phi i32 [ 0, %for.cond.preheader ], [ %inc, %inner.for.cond.preheader ]
+  %mai.1 = tail call <32 x float> @llvm.amdgcn.mfma.f32.32x32x1f32(float 1.0, float 2.0, <32 x float> %phi, i32 0, i32 0, i32 0)
+  %inc = add nuw nsw i32 %c, 1
+  %cc = icmp eq i32 %inc, 16
+  br i1 %cc, label %inner.exit, label %inner.for.cond.preheader
+
+inner.exit:
+  %inc.0 = add nuw nsw i32 %c.0, 1
+  %cc.0 = icmp eq i32 %inc.0, 16
+  br i1 %cc.0, label %exit, label %for.cond.preheader
+
+exit:
+  store <32 x float> %mai.1, <32 x float> addrspace(1)* %arg
+  ret void
+}
+
 declare <32 x float> @llvm.amdgcn.mfma.f32.32x32x1f32(float, float, <32 x float>, i32, i32, i32)
 declare i32 @llvm.amdgcn.workitem.id.x()