Update [AMDGPU] Dynamic VGPR support for llvm.amdgcn.cs.chain llvm#130094

Author: qiaojbao

llvm/lib/Target/AMDGPU/AMDGPUCallLowering.cpp

@@ -1246,15 +1246,15 @@ bool AMDGPUCallLowering::lowerTailCall(
       // On GFX12, we can only change the VGPR allocation for wave32.
       if (!ST.isWave32()) {
         F.getContext().diagnose(DiagnosticInfoUnsupported(
-            F, "Dynamic VGPR mode is only supported for wave32\n"));
+            F, "dynamic VGPR mode is only supported for wave32"));
         return false;
       }
 
       ArgInfo FallbackExecArg = Info.OrigArgs[ChainCallArgIdx::FallbackExec];
       assert(FallbackExecArg.Regs.size() == 1 &&
              "Expected single register for fallback EXEC");
       if (!FallbackExecArg.Ty->isIntegerTy(ST.getWavefrontSize())) {
-        LLVM_DEBUG(dbgs() << "Bad type for fallback EXEC");
+        LLVM_DEBUG(dbgs() << "Bad type for fallback EXEC\n");
         return false;
       }
     }

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

@@ -5475,6 +5475,7 @@ const char* AMDGPUTargetLowering::getTargetNodeName(unsigned Opcode) const {
   NODE_NAME_CASE(TC_RETURN)
   NODE_NAME_CASE(TC_RETURN_GFX)
   NODE_NAME_CASE(TC_RETURN_CHAIN)
+  NODE_NAME_CASE(TC_RETURN_CHAIN_DVGPR)
   NODE_NAME_CASE(TRAP)
   NODE_NAME_CASE(RET_GLUE)
   NODE_NAME_CASE(WAVE_ADDRESS)

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.h

@@ -402,6 +402,7 @@ enum NodeType : unsigned {
   TC_RETURN,
   TC_RETURN_GFX,
   TC_RETURN_CHAIN,
+  TC_RETURN_CHAIN_DVGPR,
   TRAP,
 
   // Masked control flow nodes.

llvm/lib/Target/AMDGPU/AMDGPUInstrInfo.td

@@ -99,6 +99,12 @@ def AMDGPUtc_return_chain: SDNode<"AMDGPUISD::TC_RETURN_CHAIN",
   [SDNPHasChain, SDNPOptInGlue, SDNPVariadic]
 >;
 
+// With dynamic VGPRs.
+def AMDGPUtc_return_chain_dvgpr: SDNode<"AMDGPUISD::TC_RETURN_CHAIN_DVGPR",
+  SDTypeProfile<0, -1, [SDTCisPtrTy<0>]>,
+  [SDNPHasChain, SDNPOptInGlue, SDNPVariadic]
+>;
+
 def AMDGPUtrap : SDNode<"AMDGPUISD::TRAP",
   SDTypeProfile<0, 1, [SDTCisVT<0, i16>]>,
     [SDNPHasChain, SDNPVariadic, SDNPSideEffect, SDNPOptInGlue]

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

@@ -3710,6 +3710,7 @@ SDValue SITargetLowering::LowerCall(CallLoweringInfo &CLI,
   SDValue Callee = CLI.Callee;
 
   llvm::SmallVector<SDValue, 6> ChainCallSpecialArgs;
+  bool UsesDynamicVGPRs = false;
   if (IsChainCallConv) {
     // The last arguments should be the value that we need to put in EXEC,
     // followed by the flags and any other arguments with special meanings.
@@ -3758,6 +3759,7 @@ SDValue SITargetLowering::LowerCall(CallLoweringInfo &CLI,
         return lowerUnhandledCall(CLI, InVals, "Expected 3 additional args");
       }
 
+      UsesDynamicVGPRs = true;
       std::for_each(CLI.Args.begin() + ChainCallArgIdx::NumVGPRs,
                     CLI.Args.end(), PushNodeOrTargetConstant);
     }
@@ -4091,7 +4093,8 @@ SDValue SITargetLowering::LowerCall(CallLoweringInfo &CLI,
       break;
     case CallingConv::AMDGPU_CS_Chain:
     case CallingConv::AMDGPU_CS_ChainPreserve:
-      OPC = AMDGPUISD::TC_RETURN_CHAIN;
+      OPC = UsesDynamicVGPRs ? AMDGPUISD::TC_RETURN_CHAIN_DVGPR
+                             : AMDGPUISD::TC_RETURN_CHAIN;
       break;
     }
 

llvm/lib/Target/AMDGPU/SIInstructions.td

@@ -704,16 +704,16 @@ multiclass SI_CS_CHAIN_TC<
       SchedRW = [WriteBranch],
       isConvergent = 1,
       WaveSizePredicate = wavesizepred in {
-  // This is essentially a tail call, but it also takes a mask to put in EXEC
-  // right before jumping to the callee.
-  def NAME: SPseudoInstSI <(outs),
-      (ins CCR_SGPR_64:$src0, unknown:$callee, i32imm:$fpdiff, execrc:$exec)>;
-
-  // Same as above, but it will first try to reallocate the VGPRs, and choose an
-  // EXEC mask and a callee depending on the success of the reallocation attempt.
-  def _DVGPR : SPseudoInstSI <(outs),
-      (ins CCR_SGPR_64:$src0, i64imm:$callee, i32imm:$fpdiff, execrc:$exec,
-           SSrc_b32:$numvgprs, execrc:$fbexec, CCR_SGPR_64:$fbcallee)>;
+    // This is essentially a tail call, but it also takes a mask to put in EXEC
+    // right before jumping to the callee.
+    def NAME: SPseudoInstSI <(outs),
+        (ins CCR_SGPR_64:$src0, unknown:$callee, i32imm:$fpdiff, execrc:$exec)>;
+
+    // Same as above, but it will first try to reallocate the VGPRs, and choose an
+    // EXEC mask and a callee depending on the success of the reallocation attempt.
+    def _DVGPR : SPseudoInstSI <(outs),
+        (ins CCR_SGPR_64:$src0, i64imm:$callee, i32imm:$fpdiff, execrc:$exec,
+             SSrc_b32:$numvgprs, execrc:$fbexec, CCR_SGPR_64:$fbcallee)>;
   } // End FixedSize = 0 etc
 }
 
@@ -747,16 +747,15 @@ defm : si_cs_chain_tc_patterns<i64>;
 multiclass si_cs_chain_tc_dvgpr_patterns<
   ValueType execvt, RegisterOperand execrc = getSOPSrcForVT<execvt>.ret,
   Instruction tc = SI_CS_CHAIN_TC_W32_DVGPR> {
-  let AddedComplexity = 90 in {
   foreach callee = [ (i64 0), (tglobaladdr) ] in {
   def : GCNPat<
-    (AMDGPUtc_return_chain i64:$src0, callee, (i32 timm:$fpdiff), execvt:$exec,
-                           i32:$numvgprs, execvt:$fbexec, i64:$fbcallee),
+    (AMDGPUtc_return_chain_dvgpr i64:$src0, callee, (i32 timm:$fpdiff),
+                                 execvt:$exec, i32:$numvgprs,
+                                 execvt:$fbexec, i64:$fbcallee),
     (tc CCR_SGPR_64:$src0, (i64 0), i32imm:$fpdiff, execrc:$exec,
         SSrc_b32:$numvgprs, execrc:$fbexec, CCR_SGPR_64:$fbcallee)
   >;
   }
-  } // AddedComplexity
 }
 
 defm : si_cs_chain_tc_dvgpr_patterns<i32>; // On GFX12, dVGPR mode is wave32-only.

llvm/lib/Target/AMDGPU/SILateBranchLowering.cpp

@@ -117,11 +117,12 @@ static void splitBlock(MachineBasicBlock &MBB, MachineInstr &MI,
   MDT->applyUpdates(DTUpdates);
 }
 
-static void addRegOrCopyOp(MachineInstrBuilder &MIB, MachineOperand &Op) {
+static void copyOpWithoutRegFlags(MachineInstrBuilder &MIB,
+                                  MachineOperand &Op) {
   if (Op.isReg())
     MIB.addReg(Op.getReg());
   else
-    MIB->addOperand(Op);
+    MIB.add(Op);
 }
 
 void SILateBranchLowering::expandChainCall(MachineInstr &MI,
@@ -138,30 +139,35 @@ void SILateBranchLowering::expandChainCall(MachineInstr &MI,
     // * Try to change the VGPR allocation
     // * Select the callee based on the result of the reallocation attempt
     // * Select the EXEC mask based on the result of the reallocation attempt
+    // If any of the register operands of the chain pseudo is used in more than
+    // one of these instructions, we need to make sure that the kill flags
+    // aren't copied along.
     auto AllocMI =
         BuildMI(*MI.getParent(), MI, DL, TII->get(AMDGPU::S_ALLOC_VGPR));
-    addRegOrCopyOp(AllocMI,
-                   *TII->getNamedOperand(MI, AMDGPU::OpName::numvgprs));
+    copyOpWithoutRegFlags(AllocMI,
+                          *TII->getNamedOperand(MI, AMDGPU::OpName::numvgprs));
 
     auto SelectCallee =
         BuildMI(*MI.getParent(), MI, DL, TII->get(AMDGPU::S_CSELECT_B64))
             .addDef(TII->getNamedOperand(MI, AMDGPU::OpName::src0)->getReg());
-    addRegOrCopyOp(SelectCallee,
-                   *TII->getNamedOperand(MI, AMDGPU::OpName::src0));
-    addRegOrCopyOp(SelectCallee,
-                   *TII->getNamedOperand(MI, AMDGPU::OpName::fbcallee));
+    copyOpWithoutRegFlags(SelectCallee,
+                          *TII->getNamedOperand(MI, AMDGPU::OpName::src0));
+    copyOpWithoutRegFlags(SelectCallee,
+                          *TII->getNamedOperand(MI, AMDGPU::OpName::fbcallee));
 
     auto SelectExec = BuildMI(*MI.getParent(), MI, DL,
                               TII->get(ST.isWave32() ? AMDGPU::S_CSELECT_B32
                                                      : AMDGPU::S_CSELECT_B64))
                           .addDef(ExecReg);
 
-    addRegOrCopyOp(SelectExec, *TII->getNamedOperand(MI, AMDGPU::OpName::exec));
-    addRegOrCopyOp(SelectExec,
-                   *TII->getNamedOperand(MI, AMDGPU::OpName::fbexec));
+    copyOpWithoutRegFlags(SelectExec,
+                          *TII->getNamedOperand(MI, AMDGPU::OpName::exec));
+    copyOpWithoutRegFlags(SelectExec,
+                          *TII->getNamedOperand(MI, AMDGPU::OpName::fbexec));
   } else {
     auto SetExec = BuildMI(*MI.getParent(), MI, DL, TII->get(MovOpc), ExecReg);
-    addRegOrCopyOp(SetExec, *TII->getNamedOperand(MI, AMDGPU::OpName::exec));
+    copyOpWithoutRegFlags(SetExec,
+                          *TII->getNamedOperand(MI, AMDGPU::OpName::exec));
   }
 
   for (unsigned OpIdx = MI.getNumExplicitOperands() - 1; OpIdx >= ExecIdx;
@@ -215,12 +221,12 @@ bool SILateBranchLowering::runOnMachineFunction(MachineFunction &MF) {
 
       case AMDGPU::SI_CS_CHAIN_TC_W32:
       case AMDGPU::SI_CS_CHAIN_TC_W64:
-        expandChainCall(MI, ST, /*DynamicVGPR=*/ false);
+        expandChainCall(MI, ST, /*DynamicVGPR=*/false);
         MadeChange = true;
         break;
       case AMDGPU::SI_CS_CHAIN_TC_W32_DVGPR:
       case AMDGPU::SI_CS_CHAIN_TC_W64_DVGPR:
-        expandChainCall(MI, ST, /*DynamicVGPR=*/ true);
+        expandChainCall(MI, ST, /*DynamicVGPR=*/true);
         MadeChange = true;
         break;