[AMDGPU] Avoid offset register in MUBUF for direct stack object accesses

[cdevadas]

Summary

We use an absolute address for stack objects and
it would be necessary to have a constant 0 for soffset field.

[arsenm]

Should make corresponding globalisel change. I think there are some missing assert changes too that were checking the soffset

[cdevadas]

At global-isel, it has already been taken care of.
The asserts - the cases where it is expecting a register is specifically for sp relative accesses and won't cause a problem after this patch.

[arsenm]

In D89234#2327303, @cdevadas wrote:

At global-isel, it has already been taken care of.
The asserts - the cases where it is expecting a register is specifically for sp relative accesses and won't cause a problem after this patch.

It's not:

[=](MachineInstrBuilder &MIB) { // soffset
  // If we don't know this private access is a local stack object, it
  // needs to be relative to the entry point's scratch wave offset.
  // TODO: Should split large offsets that don't fit like above.
  // TODO: Don't use scratch wave offset just because the offset
  // didn't fit.
  if (!Info->isEntryFunction() && FI.hasValue())
    MIB.addReg(Info->getStackPtrOffsetReg());
  else
    MIB.addImm(0);
},

[arsenm]

llvm/test/CodeGen/AMDGPU/stack-pointer-offset-relative-frameindex.ll

43–46 ↗

(On Diff #297559)

This is worse. Why did this break the local stack offset optimization?

[arsenm]

This should also remove my recent change to SIRegisterInfo::resolveFrameIndex and switch it to an assert

[cdevadas]

Used the FI directly in MI if available during ISel. Fixed an unhandled case during eliminateFrameIndex.
Also fixed the assertions to make sure the soffset field is relevant.

[cdevadas]

Ping

llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp

 
 static bool isStackPtrRelative(const MachinePointerInfo &PtrInfo) {
   auto PSV = PtrInfo.V.dyn_cast<const PseudoSourceValue *>();
   return PSV && PSV->isStack();
 }
 
 std::pair<SDValue, SDValue> AMDGPUDAGToDAGISel::foldFrameIndex(SDValue N) const {
   SDLoc DL(N);
-  const MachineFunction &MF = CurDAG->getMachineFunction();
-  const SIMachineFunctionInfo *Info = MF.getInfo<SIMachineFunctionInfo>();
-
-  if (auto FI = dyn_cast<FrameIndexSDNode>(N)) {
-    SDValue TFI = CurDAG->getTargetFrameIndex(FI->getIndex(),
-                                              FI->getValueType(0));
-
-    // If we can resolve this to a frame index access, this will be relative to
-    // either the stack or frame pointer SGPR.
-    return std::make_pair(
-        TFI, CurDAG->getRegister(Info->getStackPtrOffsetReg(), MVT::i32));
-  }
 
-  // If we don't know this private access is a local stack object, it needs to
-  // be relative to the entry point's scratch wave offset.
-  return std::make_pair(N, CurDAG->getTargetConstant(0, DL, MVT::i32));
+  auto *FI = dyn_cast<FrameIndexSDNode>(N);
+  SDValue TFI =
+      FI ? CurDAG->getTargetFrameIndex(FI->getIndex(), FI->getValueType(0)) : N;
+
+  // We rebase the base address into an absolute stack address and hence
+  // use constant 0 for soffset.
+  return std::make_pair(TFI, CurDAG->getTargetConstant(0, DL, MVT::i32));
 }
 
 bool AMDGPUDAGToDAGISel::SelectMUBUFScratchOffen(SDNode *Parent,
                                                  SDValue Addr, SDValue &Rsrc,
                                                  SDValue &VAddr, SDValue &SOffset,
                                                  SDValue &ImmOffset) const {
 
   SDLoc DL(Addr);

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

 #ifndef NDEBUG
   MachineBasicBlock *MBB = MI.getParent();
   MachineFunction *MF = MBB->getParent();
 #endif
   assert(FIOp && FIOp->isFI() && "frame index must be address operand");
   assert(TII->isMUBUF(MI));
 
   MachineOperand *SOffset = TII->getNamedOperand(MI, AMDGPU::OpName::soffset);
-  assert(SOffset->getReg() ==
-             MF->getInfo<SIMachineFunctionInfo>()->getStackPtrOffsetReg() &&
-         "should only be seeing stack pointer offset relative FrameIndex");
+  assert((SOffset->isReg() &&
+          SOffset->getReg() ==
+              MF->getInfo<SIMachineFunctionInfo>()->getStackPtrOffsetReg()) ||
+         (SOffset->isImm() && SOffset->getImm() == 0));
 
   MachineOperand *OffsetOp = TII->getNamedOperand(MI, AMDGPU::OpName::offset);
   int64_t NewOffset = OffsetOp->getImm() + Offset;
   assert(SIInstrInfo::isLegalMUBUFImmOffset(NewOffset) &&
          "offset should be legal");
 
   FIOp->ChangeToRegister(BaseReg, false);
   OffsetOp->setImm(NewOffset);
-
-  // The move materializing the base address will be an absolute stack address,
-  // so clear the base offset.
-  SOffset->ChangeToImmediate(0);
 }
 
 bool SIRegisterInfo::isFrameOffsetLegal(const MachineInstr *MI,
                                         Register BaseReg,
                                         int64_t Offset) const {
   if (!SIInstrInfo::isMUBUF(*MI))
     return false;
 
default: {
                 SOffset.getReg() == MFI->getStackPtrOffsetReg()) ||
                (SOffset.isImm() && SOffset.getImm() == 0));
         if (SOffset.isReg()) {
           if (FrameReg == AMDGPU::NoRegister) {
             SOffset.ChangeToImmediate(0);
           } else {
             SOffset.setReg(FrameReg);
           }
+        } else if (SOffset.isImm() && FrameReg != AMDGPU::NoRegister) {
+          SOffset.ChangeToRegister(FrameReg, false);
         }
 
         int64_t Offset = FrameInfo.getObjectOffset(Index);
         int64_t OldImm
           = TII->getNamedOperand(*MI, AMDGPU::OpName::offset)->getImm();
         int64_t NewOffset = OldImm + Offset;
 
         if (SIInstrInfo::isLegalMUBUFImmOffset(NewOffset) &&

llvm/test/CodeGen/AMDGPU/call-argument-types.ll

 define amdgpu_kernel void @stack_passed_arg_alignment_v32i32_f64(<32 x i32> %val, double %tmp) #0 {
 entry:
   call void @stack_passed_f64_arg(<32 x i32> %val, double %tmp)
   ret void
 }
 
 ; GCN-LABEL: {{^}}tail_call_byval_align16:
 ; GCN-NOT: s32
-; GCN: buffer_load_dword v32, off, s[0:3], s32 offset:12
-; GCN: buffer_load_dword v33, off, s[0:3], s32 offset:8
+; GCN: buffer_load_dword [[VREG1:v[0-9]+]], off, s[0:3], s32 offset:8
+; GCN: buffer_load_dword [[VREG2:v[0-9]+]], off, s[0:3], s32 offset:12
 
 ; GCN: s_getpc_b64
 
-; GCN: buffer_store_dword v32, off, s[0:3], s32 offset:4
-; GCN: buffer_store_dword v33, off, s[0:3], s32{{$}}
+; GCN: buffer_store_dword [[VREG2]], off, s[0:3], s32 offset:4
+; GCN: buffer_store_dword [[VREG1]], off, s[0:3], s32{{$}}
 ; GCN-NOT: s32
 ; GCN: s_setpc_b64
 define void @tail_call_byval_align16(<32 x i32> %val, double %tmp) #0 {
 entry:
   %alloca = alloca double, align 8, addrspace(5)
   tail call void @byval_align16_f64_arg(<32 x i32> %val, double addrspace(5)* byval align 16 %alloca)
   ret void
 }

llvm/test/CodeGen/AMDGPU/callee-special-input-vgprs.ll

 ; var abi stack layout:
 ; sp[0] = byval
 ; sp[1] = ??
 ; sp[2] = stack passed workitem ID x
 
 ; GCN-LABEL: {{^}}kern_call_too_many_args_use_workitem_id_x_byval:
 ; VARABI: enable_vgpr_workitem_id = 0
 ; VARABI: v_mov_b32_e32 [[K:v[0-9]+]], 0x3e7{{$}}
-; VARABI: buffer_store_dword [[K]], off, s[0:3], 0 offset:4
-; VARABI: buffer_load_dword [[RELOAD_BYVAL:v[0-9]+]], off, s[0:3], 0 offset:4
 ; VARABI: s_movk_i32 s32, 0x400{{$}}
-
-; VARABI-NOT: s32
+; VARABI: buffer_store_dword [[K]], off, s[0:3], 0 offset:4
 ; VARABI: buffer_store_dword v0, off, s[0:3], s32 offset:4
+; VARABI: buffer_load_dword [[RELOAD_BYVAL:v[0-9]+]], off, s[0:3], 0 offset:4
 
 ; VARABI: buffer_store_dword [[RELOAD_BYVAL]], off, s[0:3], s32{{$}}
 ; VARABI: v_mov_b32_e32 [[RELOAD_BYVAL]],
 ; VARABI: s_swappc_b64
 
 
 ; FIXEDABI: v_mov_b32_e32 [[K0:v[0-9]+]], 0x3e7
 ; FIXEDABI: buffer_store_dword [[K0]], off, s[0:3], 0 offset:4{{$}}
call void @too_many_args_use_workitem_id_x_byval(
     i32 290, i32 300, i32 310, i32 320,
     i32 addrspace(5)* %alloca)
   ret void
 }
 
 ; GCN-LABEL: {{^}}func_call_too_many_args_use_workitem_id_x_byval:
 ; VARABI: v_mov_b32_e32 [[K:v[0-9]+]], 0x3e7{{$}}
 ; VARABI: buffer_store_dword [[K]], off, s[0:3], s33{{$}}
-; VARABI: buffer_load_dword [[RELOAD_BYVAL:v[0-9]+]], off, s[0:3], s33{{$}}
 ; VARABI: buffer_store_dword v0, off, s[0:3], s32 offset:4
+; VARABI: buffer_load_dword [[RELOAD_BYVAL:v[0-9]+]], off, s[0:3], s33{{$}}
 ; VARABI: buffer_store_dword [[RELOAD_BYVAL]], off, s[0:3], s32{{$}}
 ; VARABI: v_mov_b32_e32 [[RELOAD_BYVAL]],
 ; VARABI: s_swappc_b64
 
 
 ; FIXED-ABI-NOT: v31
 ; FIXEDABI: v_mov_b32_e32 [[K0:v[0-9]+]], 0x3e7{{$}}
 ; FIXEDABI: buffer_store_dword [[K0]], off, s[0:3], s33{{$}}