CodeGen: Let frame index value type match alloca addr space

[yaxunl]

Summary

Recently alloca address space has been added to data layout. Due to this change, pointer returned by alloca may have different size as pointer in address space 0.

However, currently the value type of frame index is assumed to be of the same size as pointer in address space 0.

This patch fixes that.

Most targets assume alloca returning pointer in address space 0, which is the default alloca address space. Therefore it is NFC for them.

AMDGCN target with amdgiz environment requires this change since it assumes alloca returning pointer to addr space 5 and its size is 32, which is different from the size of pointer in addr space 0 which is 64.

[arsenm]

Can you add a test using the new mapping for this

lib/CodeGen/SelectionDAG/SelectionDAG.cpp

1829 ↗

(On Diff #95134)

Add a new TLI->getAllocaPointerTy?

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

5418 ↗

(On Diff #95134)

This one hasn't been updated with mangling yet, so might as well leave it for now

[arsenm]

lib/CodeGen/SelectionDAG/SelectionDAG.cpp

1829 ↗

(On Diff #95134)

getFrameIndexTy would be better

[yaxunl]

lib/CodeGen/SelectionDAG/SelectionDAG.cpp

1829 ↗

(On Diff #95134)

will do.

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

5418 ↗

(On Diff #95134)

will do.

[yaxunl]

Revised as Matt suggested.

[t-tye]

LGTM

[yaxunl]

Ping. Any further comments? Thanks.

[arsenm]

LGTM with test cleanups

test/CodeGen/AMDGPU/frame-index-amdgiz.ll

14 ↗

(On Diff #95580)

Drop metadata

54 ↗

(On Diff #95580)

Strip unnecessary string attributes

58 ↗

(On Diff #95580)

You can drop this

llvm/trunk/include/llvm/Target/TargetLowering.h

public:
 
   /// Return the pointer type for the given address space, defaults to
   /// the pointer type from the data layout.
   /// FIXME: The default needs to be removed once all the code is updated.
   MVT getPointerTy(const DataLayout &DL, uint32_t AS = 0) const {
     return MVT::getIntegerVT(DL.getPointerSizeInBits(AS));
   }
 
+  /// Return the type for frame index, which is determined by
+  /// the alloca address space specified through the data layout.
+  MVT getFrameIndexTy(const DataLayout &DL) const {
+    return getPointerTy(DL, DL.getAllocaAddrSpace());
+  }
+
   /// EVT is not used in-tree, but is used by out-of-tree target.
   /// A documentation for this function would be nice...
   virtual MVT getScalarShiftAmountTy(const DataLayout &, EVT) const;
 
   EVT getShiftAmountTy(EVT LHSTy, const DataLayout &DL) const;
 
   /// Returns the type to be used for the index operand of:
   /// ISD::INSERT_VECTOR_ELT, ISD::EXTRACT_VECTOR_ELT,

llvm/trunk/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

 SDValue SelectionDAG::CreateStackTemporary(EVT VT, unsigned minAlign) {
   MachineFrameInfo &MFI = getMachineFunction().getFrameInfo();
   unsigned ByteSize = VT.getStoreSize();
   Type *Ty = VT.getTypeForEVT(*getContext());
   unsigned StackAlign =
       std::max((unsigned)getDataLayout().getPrefTypeAlignment(Ty), minAlign);
 
   int FrameIdx = MFI.CreateStackObject(ByteSize, StackAlign, false);
-  return getFrameIndex(FrameIdx, TLI->getPointerTy(getDataLayout()));
+  return getFrameIndex(FrameIdx, TLI->getFrameIndexTy(getDataLayout()));
 }
 
 SDValue SelectionDAG::CreateStackTemporary(EVT VT1, EVT VT2) {
   unsigned Bytes = std::max(VT1.getStoreSize(), VT2.getStoreSize());
   Type *Ty1 = VT1.getTypeForEVT(*getContext());
   Type *Ty2 = VT2.getTypeForEVT(*getContext());
   const DataLayout &DL = getDataLayout();
   unsigned Align =
       std::max(DL.getPrefTypeAlignment(Ty1), DL.getPrefTypeAlignment(Ty2));
 
   MachineFrameInfo &MFI = getMachineFunction().getFrameInfo();
   int FrameIdx = MFI.CreateStackObject(Bytes, Align, false);
-  return getFrameIndex(FrameIdx, TLI->getPointerTy(getDataLayout()));
+  return getFrameIndex(FrameIdx, TLI->getFrameIndexTy(getDataLayout()));
 }
 
 SDValue SelectionDAG::FoldSetCC(EVT VT, SDValue N1, SDValue N2,
                                 ISD::CondCode Cond, const SDLoc &dl) {
   // These setcc operations always fold.
   switch (Cond) {
   default: break;
   case ISD::SETFALSE:

llvm/trunk/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

SDValue SelectionDAGBuilder::getValueImpl(const Value *V) {
 
   // If this is a static alloca, generate it as the frameindex instead of
   // computation.
   if (const AllocaInst *AI = dyn_cast<AllocaInst>(V)) {
     DenseMap<const AllocaInst*, int>::iterator SI =
       FuncInfo.StaticAllocaMap.find(AI);
     if (SI != FuncInfo.StaticAllocaMap.end())
       return DAG.getFrameIndex(SI->second,
-                               TLI.getPointerTy(DAG.getDataLayout()));
+                               TLI.getFrameIndexTy(DAG.getDataLayout()));
   }
 
   // If this is an instruction which fast-isel has deferred, select it now.
   if (const Instruction *Inst = dyn_cast<Instruction>(V)) {
     unsigned InReg = FuncInfo.InitializeRegForValue(Inst);
     RegsForValue RFV(*DAG.getContext(), TLI, DAG.getDataLayout(), InReg,
                      Inst->getType());
     SDValue Chain = DAG.getEntryNode();
for (SmallVectorImpl<Value*>::iterator Object = Allocas.begin(),
       if (SI == FuncInfo.StaticAllocaMap.end())
         return nullptr;
 
       int FI = SI->second;
 
       SDValue Ops[2];
       Ops[0] = getRoot();
       Ops[1] =
-          DAG.getFrameIndex(FI, TLI.getPointerTy(DAG.getDataLayout()), true);
+          DAG.getFrameIndex(FI, TLI.getFrameIndexTy(DAG.getDataLayout()), true);
       unsigned Opcode = (IsStart ? ISD::LIFETIME_START : ISD::LIFETIME_END);
 
       Res = DAG.getNode(Opcode, sdl, MVT::Other, Ops);
       DAG.setRoot(Res);
     }
     return nullptr;
   }
   case Intrinsic::invariant_start:
static SDValue getAddressForMemoryInput(SDValue Chain, const SDLoc &Location,
 
   // Otherwise, create a stack slot and emit a store to it before the asm.
   Type *Ty = OpVal->getType();
   auto &DL = DAG.getDataLayout();
   uint64_t TySize = DL.getTypeAllocSize(Ty);
   unsigned Align = DL.getPrefTypeAlignment(Ty);
   MachineFunction &MF = DAG.getMachineFunction();
   int SSFI = MF.getFrameInfo().CreateStackObject(TySize, Align, false);
-  SDValue StackSlot = DAG.getFrameIndex(SSFI, TLI.getPointerTy(DL));
+  SDValue StackSlot = DAG.getFrameIndex(SSFI, TLI.getFrameIndexTy(DL));
   Chain = DAG.getStore(Chain, Location, OpInfo.CallOperand, StackSlot,
                        MachinePointerInfo::getFixedStack(MF, SSFI));
   OpInfo.CallOperand = StackSlot;
 
   return Chain;
 }
 
 /// GetRegistersForValue - Assign registers (virtual or physical) for the
for (unsigned i = StartIdx, e = CS.arg_size(); i != e; ++i) {
     if (ConstantSDNode *C = dyn_cast<ConstantSDNode>(OpVal)) {
       Ops.push_back(
         Builder.DAG.getTargetConstant(StackMaps::ConstantOp, DL, MVT::i64));
       Ops.push_back(
         Builder.DAG.getTargetConstant(C->getSExtValue(), DL, MVT::i64));
     } else if (FrameIndexSDNode *FI = dyn_cast<FrameIndexSDNode>(OpVal)) {
       const TargetLowering &TLI = Builder.DAG.getTargetLoweringInfo();
       Ops.push_back(Builder.DAG.getTargetFrameIndex(
-          FI->getIndex(), TLI.getPointerTy(Builder.DAG.getDataLayout())));
+          FI->getIndex(), TLI.getFrameIndexTy(Builder.DAG.getDataLayout())));
     } else
       Ops.push_back(OpVal);
   }
 }
 
 /// \brief Lower llvm.experimental.stackmap directly to its target opcode.
 void SelectionDAGBuilder::visitStackmap(const CallInst &CI) {
   // void @llvm.experimental.stackmap(i32 <id>, i32 <numShadowBytes>,
if (!CanLowerReturn) {
     // assert(!CS.hasInAllocaArgument() &&
     //        "sret demotion is incompatible with inalloca");
     uint64_t TySize = DL.getTypeAllocSize(CLI.RetTy);
     unsigned Align = DL.getPrefTypeAlignment(CLI.RetTy);
     MachineFunction &MF = CLI.DAG.getMachineFunction();
     DemoteStackIdx = MF.getFrameInfo().CreateStackObject(TySize, Align, false);
     Type *StackSlotPtrType = PointerType::getUnqual(CLI.RetTy);
 
-    DemoteStackSlot = CLI.DAG.getFrameIndex(DemoteStackIdx, getPointerTy(DL));
+    DemoteStackSlot = CLI.DAG.getFrameIndex(DemoteStackIdx, getFrameIndexTy(DL));
     ArgListEntry Entry;
     Entry.Node = DemoteStackSlot;
     Entry.Ty = StackSlotPtrType;
     Entry.IsSExt = false;
     Entry.IsZExt = false;
     Entry.IsInReg = false;
     Entry.IsSRet = true;
     Entry.IsNest = false;

llvm/trunk/test/CodeGen/AMDGPU/frame-index-amdgiz.ll

+; RUN: llc -verify-machineinstrs < %s | FileCheck %s
+;
+; The original OpenCL kernel:
+; kernel void f(global int *a, int i,  int j) {
+;  int x[100];
+;  x[i] = 7;
+;  a[0] = x[j];
+; }
+; clang -cc1 -triple amdgcn---amdgizcl -emit-llvm -o - 
+
+target datalayout = "e-p:64:64-p1:64:64-p2:64:64-p3:32:32-p4:32:32-p5:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-A5"
+target triple = "amdgcn---amdgiz"
+
+define amdgpu_kernel void @f(i32 addrspace(1)* nocapture %a, i32 %i, i32 %j) local_unnamed_addr #0 {
+entry:
+; CHECK: s_load_dword s2, s[0:1], 0xb
+; CHECK: s_load_dwordx2 s[4:5], s[0:1], 0x9
+; CHECK: s_load_dword s0, s[0:1], 0xc
+; CHECK: s_mov_b32 s8, SCRATCH_RSRC_DWORD0
+; CHECK: s_mov_b32 s9, SCRATCH_RSRC_DWORD1
+; CHECK: s_mov_b32 s10, -1
+; CHECK: s_waitcnt lgkmcnt(0)
+; CHECK: s_lshl_b32 s1, s2, 2
+; CHECK: v_mov_b32_e32 v0, 4
+; CHECK: s_mov_b32 s11, 0xe8f000
+; CHECK: v_add_i32_e32 v1, vcc, s1, v0
+; CHECK: v_mov_b32_e32 v2, 7
+; CHECK: s_lshl_b32 s0, s0, 2
+; CHECK: buffer_store_dword v2, v1, s[8:11], s3 offen
+; CHECK: v_add_i32_e32 v0, vcc, s0, v0
+; CHECK: buffer_load_dword v0, v0, s[8:11], s3 offen
+; CHECK: s_mov_b32 s7, 0xf000
+; CHECK: s_mov_b32 s6, -1
+; CHECK: s_waitcnt vmcnt(0)
+; CHECK: buffer_store_dword v0, off, s[4:7], 0
+; CHECK: s_endpgm
+
+  %x = alloca [100 x i32], align 4, addrspace(5)
+  %0 = bitcast [100 x i32] addrspace(5)* %x to i8 addrspace(5)*
+  call void @llvm.lifetime.start.p5i8(i64 400, i8 addrspace(5)* nonnull %0) #0
+  %arrayidx = getelementptr inbounds [100 x i32], [100 x i32] addrspace(5)* %x, i32 0, i32 %i
+  store i32 7, i32 addrspace(5)* %arrayidx, align 4
+  %arrayidx2 = getelementptr inbounds [100 x i32], [100 x i32] addrspace(5)* %x, i32 0, i32 %j
+  %1 = load i32, i32 addrspace(5)* %arrayidx2, align 4
+  store i32 %1, i32 addrspace(1)* %a, align 4
+  call void @llvm.lifetime.end.p5i8(i64 400, i8 addrspace(5)* nonnull %0) #0
+  ret void
+}
+
+declare void @llvm.lifetime.start.p5i8(i64, i8 addrspace(5)* nocapture) #1
+
+declare void @llvm.lifetime.end.p5i8(i64, i8 addrspace(5)* nocapture) #1
+
+attributes #0 = { nounwind }
+attributes #1 = { argmemonly nounwind }