AMDGPU: Fix UBSan error caused by r335942

[tstellar]

Summary

Fixes PR38071.

[arsenm]

lib/Target/AMDGPU/AMDGPUISelLowering.cpp

3942–3944

Why can't this keep getting the offset from the subtarget?

lib/Target/AMDGPU/AMDGPUISelLowering.h

297–298

Weird formatting

[tstellar]

lib/Target/AMDGPU/AMDGPUISelLowering.cpp

3942–3944

That would require moving another function (isAMDCodeObjectV2) into the AMDGPUCommonSubtarget class. I've been trying not to add more things to this class, but I can make this change.

[arsenm]

lib/Target/AMDGPU/AMDGPUISelLowering.cpp

3942–3944

This ABI stuff probably does need to move there, but my main concern is how this broke so subtly

[tstellar]

Use subtarget query to get kernel arg offset.

[tstellar]

lib/Target/AMDGPU/AMDGPUISelLowering.cpp

3942–3944

AMDGPUSubtarget is now GCN specific, so I just forgot to rename AMDGPUSubtarget->AMDGPUCommonSubtarget.

I have a few cleanups to do after splitting the targets and making the subtarget names make sense is one of them.

[dstenb]

I don't really know anything about the AMDGPU target, so I'll resign as reviewer to make that clear, but thanks for notifying me about the fix!

[arsenm]

LGTM

lib/Target/AMDGPU/AMDGPUISelLowering.h

enum ImplicitParameter {
     FIRST_IMPLICIT,
     GRID_DIM = FIRST_IMPLICIT,
     GRID_OFFSET,
   };
 
   /// Helper function that returns the byte offset of the given
   /// type of implicit parameter.
   uint32_t getImplicitParameterOffset(const MachineFunction &MF,
-                                      const ImplicitParameter Param) const;
+                                      const ImplicitParameter Param,
+              unsigned ExplicitArgOffset) const;
 
   AMDGPUAS getAMDGPUAS() const {
     return AMDGPUASI;
   }
 
   MVT getFenceOperandTy(const DataLayout &DL) const override {
     return MVT::i32;
   }

lib/Target/AMDGPU/AMDGPUISelLowering.cpp

SDValue AMDGPUTargetLowering::loadInputValue(SelectionDAG &DAG,
   assert(Arg && "Attempting to load missing argument");
 
   if (Arg.isRegister())
     return CreateLiveInRegister(DAG, RC, Arg.getRegister(), VT, SL);
   return loadStackInputValue(DAG, VT, SL, Arg.getStackOffset());
 }
 
 uint32_t AMDGPUTargetLowering::getImplicitParameterOffset(
-    const MachineFunction &MF, const ImplicitParameter Param) const {
+    const MachineFunction &MF, const ImplicitParameter Param,
+    unsigned ExplicitArgOffset) const {
   const AMDGPUMachineFunction *MFI = MF.getInfo<AMDGPUMachineFunction>();
-  const AMDGPUSubtarget &ST = MF.getSubtarget<AMDGPUSubtarget>();
-  unsigned ExplicitArgOffset = ST.getExplicitKernelArgOffset(MF.getFunction());
+  const AMDGPUCommonSubtarget &ST =
+      AMDGPUCommonSubtarget::get(getTargetMachine(), MF.getFunction());
   unsigned Alignment = ST.getAlignmentForImplicitArgPtr();
   uint64_t ArgOffset = alignTo(MFI->getExplicitKernArgSize(), Alignment) +
                        ExplicitArgOffset;
   switch (Param) {
   case GRID_DIM:
     return ArgOffset;
   case GRID_OFFSET:
     return ArgOffset + 4;

lib/Target/AMDGPU/R600ISelLowering.cpp

case Intrinsic::r600_dot4: {
       DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, MVT::f32, Op.getOperand(2),
           DAG.getConstant(3, DL, MVT::i32))
       };
       return DAG.getNode(AMDGPUISD::DOT4, DL, MVT::f32, Args);
     }
 
     case Intrinsic::r600_implicitarg_ptr: {
       MVT PtrVT = getPointerTy(DAG.getDataLayout(), AMDGPUASI.PARAM_I_ADDRESS);
-      uint32_t ByteOffset = getImplicitParameterOffset(MF, FIRST_IMPLICIT);
+      uint32_t ByteOffset = getImplicitParameterOffset(MF, FIRST_IMPLICIT,
+                                     Subtarget->getExplicitKernelArgOffset(MF));
       return DAG.getConstant(ByteOffset, DL, PtrVT);
     }
     case Intrinsic::r600_read_ngroups_x:
       return LowerImplicitParameter(DAG, VT, DL, 0);
     case Intrinsic::r600_read_ngroups_y:
       return LowerImplicitParameter(DAG, VT, DL, 1);
     case Intrinsic::r600_read_ngroups_z:
       return LowerImplicitParameter(DAG, VT, DL, 2);

lib/Target/AMDGPU/SIISelLowering.cpp

SDValue SITargetLowering::lowerKernArgParameterPtr(SelectionDAG &DAG,
   SDValue BasePtr = DAG.getCopyFromReg(Chain, SL,
     MRI.getLiveInVirtReg(InputPtrReg->getRegister()), PtrVT);
 
   return DAG.getObjectPtrOffset(SL, BasePtr, Offset);
 }
 
 SDValue SITargetLowering::getImplicitArgPtr(SelectionDAG &DAG,
                                             const SDLoc &SL) const {
+  MachineFunction &MF = DAG.getMachineFunction();
+  const Function &F = MF.getFunction();
+  const SISubtarget &ST = MF.getSubtarget<SISubtarget>();
   uint64_t Offset = getImplicitParameterOffset(DAG.getMachineFunction(),
-                                               FIRST_IMPLICIT);
+                                               FIRST_IMPLICIT,
+                                               ST.getExplicitKernelArgOffset(F));
   return lowerKernArgParameterPtr(DAG, SL, DAG.getEntryNode(), Offset);
 }
 
 SDValue SITargetLowering::convertArgType(SelectionDAG &DAG, EVT VT, EVT MemVT,
                                          const SDLoc &SL, SDValue Val,
                                          bool Signed,
                                          const ISD::InputArg *Arg) const {
   if (Arg && (Arg->Flags.isSExt() || Arg->Flags.isZExt()) &&