[AMDGPU] Support alloca in AS0

shiltian · shiltian · commit 01e69911c063 · 2025-04-23T16:51:31.000-04:00
This PR lowers an alloca in AS0 to an alloca in AS5 followed by an addrspacecast
back to AS0.
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp b/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp
@@ -749,6 +749,16 @@ struct AAAMDAttributesFunction : public AAAMDAttributes {
                                    UsedAssumedInformation))
       return true;
 
+    // Check all alloca instructions. FlatScratchInit is needed if there is an
+    // alloca in AS0 since that implies a cast from PRIVATE_ADDRESS.
+    auto AllocaNotInAS0 = [](Instruction &I) {
+      return cast<AllocaInst>(I).getAddressSpace() != AMDGPUAS::FLAT_ADDRESS;
+    };
+
+    if (!A.checkForAllInstructions(AllocaNotInAS0, *this, {Instruction::Alloca},
+                                   UsedAssumedInformation))
+      return true;
+
     // Check for addrSpaceCast from PRIVATE_ADDRESS in constant expressions
     auto &InfoCache = static_cast<AMDGPUInformationCache &>(A.getInfoCache());
 
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp b/llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp
@@ -385,9 +385,12 @@ AMDGPUTargetLowering::AMDGPUTargetLowering(const TargetMachine &TM,
 
   setOperationAction({ISD::BR_JT, ISD::BRIND}, MVT::Other, Expand);
 
+  setOperationAction(ISD::FrameIndex, MVT::i64, Custom);
+
   // For R600, this is totally unsupported, just custom lower to produce an
   // error.
   setOperationAction(ISD::DYNAMIC_STACKALLOC, MVT::i32, Custom);
+  setOperationAction(ISD::DYNAMIC_STACKALLOC, MVT::i64, Custom);
 
   // Library functions.  These default to Expand, but we have instructions
   // for them.
diff --git a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
@@ -912,12 +912,15 @@ AMDGPULegalizerInfo::AMDGPULegalizerInfo(const GCNSubtarget &ST_,
       .widenScalarToNextPow2(0, 32)
       .clampMaxNumElements(0, S32, 16);
 
-  getActionDefinitionsBuilder(G_FRAME_INDEX).legalFor({PrivatePtr});
+  getActionDefinitionsBuilder(G_FRAME_INDEX)
+      .legalFor({PrivatePtr})
+      .customFor({FlatPtr});
 
   // If the amount is divergent, we have to do a wave reduction to get the
   // maximum value, so this is expanded during RegBankSelect.
   getActionDefinitionsBuilder(G_DYN_STACKALLOC)
-    .legalFor({{PrivatePtr, S32}});
+      .legalFor({{PrivatePtr, S32}})
+      .customFor({FlatPtr, S32});
 
   getActionDefinitionsBuilder(G_STACKSAVE)
     .customFor({PrivatePtr});
@@ -2221,6 +2224,10 @@ bool AMDGPULegalizerInfo::legalizeCustom(
     return legalizeTrap(MI, MRI, B);
   case TargetOpcode::G_DEBUGTRAP:
     return legalizeDebugTrap(MI, MRI, B);
+  case TargetOpcode::G_FRAME_INDEX:
+    return legalizeFrameIndex(MI, MRI, B);
+  case TargetOpcode::G_DYN_STACKALLOC:
+    return legalizeDynStackAlloc(MI, MRI, B);
   default:
     return false;
   }
@@ -7668,3 +7675,25 @@ bool AMDGPULegalizerInfo::legalizeIntrinsic(LegalizerHelper &Helper,
 
   return true;
 }
+
+bool AMDGPULegalizerInfo::legalizeFrameIndex(MachineInstr &MI,
+                                             MachineRegisterInfo &MRI,
+                                             MachineIRBuilder &B) const {
+  MachineInstrBuilder FI = B.buildFrameIndex(
+      LLT::pointer(AMDGPUAS::PRIVATE_ADDRESS, 32), MI.getOperand(1).getIndex());
+  B.buildAddrSpaceCast(MI.getOperand(0).getReg(), FI);
+  MI.eraseFromParent();
+  return true;
+}
+
+bool AMDGPULegalizerInfo::legalizeDynStackAlloc(MachineInstr &MI,
+                                                MachineRegisterInfo &MRI,
+                                                MachineIRBuilder &B) const {
+  MachineInstrBuilder Size = B.buildTrunc(S32, MI.getOperand(1));
+  Align Alignment(MI.getOperand(2).getImm());
+  MachineInstrBuilder DynStackAlloc = B.buildDynStackAlloc(
+      LLT::pointer(AMDGPUAS::PRIVATE_ADDRESS, 32), Size, Alignment);
+  B.buildAddrSpaceCast(MI.getOperand(0).getReg(), DynStackAlloc);
+  MI.eraseFromParent();
+  return true;
+}
diff --git a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h
@@ -246,6 +246,11 @@ class AMDGPULegalizerInfo final : public LegalizerInfo {
 
   bool legalizeIntrinsic(LegalizerHelper &Helper,
                          MachineInstr &MI) const override;
+
+  bool legalizeFrameIndex(MachineInstr &MI, MachineRegisterInfo &MRI,
+                          MachineIRBuilder &B) const;
+  bool legalizeDynStackAlloc(MachineInstr &MI, MachineRegisterInfo &MRI,
+                             MachineIRBuilder &B) const;
 };
 } // End llvm namespace.
 #endif
diff --git a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
@@ -4117,6 +4117,17 @@ SDValue SITargetLowering::LowerCall(CallLoweringInfo &CLI,
                          InVals, /*IsThisReturn=*/false, SDValue());
 }
 
+SDValue SITargetLowering::lowerFrameIndex(SDValue Op, SelectionDAG &DAG) const {
+  // Since address space information is lost here, we assume that an i64 frame
+  // index comes from an alloca in AS0.
+  SDLoc DL(Op);
+  auto *FI = cast<FrameIndexSDNode>(Op);
+  SDValue TFI = DAG.getFrameIndex(FI->getIndex(), MVT::i32);
+  return DAG.getAddrSpaceCast(DL, Op.getValueType(), TFI,
+                              AMDGPUAS::PRIVATE_ADDRESS,
+                              AMDGPUAS::FLAT_ADDRESS);
+}
+
 // This is similar to the default implementation in ExpandDYNAMIC_STACKALLOC,
 // except for:
 // 1. Stack growth direction(default: downwards, AMDGPU: upwards), and
@@ -4129,13 +4140,27 @@ SDValue SITargetLowering::LowerDYNAMIC_STACKALLOC(SDValue Op,
   SDLoc dl(Op);
   EVT VT = Op.getValueType();
   SDValue Chain = Op.getOperand(0);
+  SDValue Size = Op.getOperand(1);
+
+  // Since address space information is lost here, we assume that an i64 dynamic
+  // alloca comes from an alloca in AS0.
+  if (VT == MVT::i64) {
+    SDValue Align = Op.getOperand(2);
+    Size = DAG.getZExtOrTrunc(Size, dl, MVT::i32);
+    SDValue Ops[] = {Chain, Size, Align};
+    SDValue DynAlloc =
+        DAG.getNode(ISD::DYNAMIC_STACKALLOC, dl, {MVT::i32, MVT::Other}, Ops);
+    SDValue Cast = DAG.getAddrSpaceCast(
+        dl, VT, DynAlloc, AMDGPUAS::PRIVATE_ADDRESS, AMDGPUAS::FLAT_ADDRESS);
+    return DAG.getMergeValues({Cast, DynAlloc.getValue(1)}, dl);
+  }
+
   Register SPReg = Info->getStackPtrOffsetReg();
 
   // Chain the dynamic stack allocation so that it doesn't modify the stack
   // pointer when other instructions are using the stack.
   Chain = DAG.getCALLSEQ_START(Chain, 0, 0, dl);
 
-  SDValue Size = Op.getOperand(1);
   SDValue BaseAddr = DAG.getCopyFromReg(Chain, dl, SPReg, VT);
   Align Alignment = cast<ConstantSDNode>(Op.getOperand(2))->getAlignValue();
 
@@ -6087,6 +6112,8 @@ SDValue SITargetLowering::LowerOperation(SDValue Op, SelectionDAG &DAG) const {
   case ISD::SMUL_LOHI:
   case ISD::UMUL_LOHI:
     return lowerXMUL_LOHI(Op, DAG);
+  case ISD::FrameIndex:
+    return lowerFrameIndex(Op, DAG);
   case ISD::DYNAMIC_STACKALLOC:
     return LowerDYNAMIC_STACKALLOC(Op, DAG);
   case ISD::STACKSAVE:
diff --git a/llvm/lib/Target/AMDGPU/SIISelLowering.h b/llvm/lib/Target/AMDGPU/SIISelLowering.h
@@ -428,6 +428,7 @@ class SITargetLowering final : public AMDGPUTargetLowering {
   SDValue LowerCall(CallLoweringInfo &CLI,
                     SmallVectorImpl<SDValue> &InVals) const override;
 
+  SDValue lowerFrameIndex(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerDYNAMIC_STACKALLOC(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerSTACKSAVE(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerGET_ROUNDING(SDValue Op, SelectionDAG &DAG) const;
diff --git a/llvm/test/CodeGen/AMDGPU/alloca-as0.ll b/llvm/test/CodeGen/AMDGPU/alloca-as0.ll
diff --git a/llvm/test/CodeGen/AMDGPU/assert-wrong-alloca-addrspace.ll b/llvm/test/CodeGen/AMDGPU/assert-wrong-alloca-addrspace.ll
diff --git a/llvm/test/CodeGen/AMDGPU/attributor-flatscratchinit.ll b/llvm/test/CodeGen/AMDGPU/attributor-flatscratchinit.ll
