AMDGPU: Fix UBSan error caused by r335942

[tstellar]

Summary

Fixes PR38071.

[arsenm]

lib/Target/AMDGPU/AMDGPUISelLowering.cpp

3941–3945

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

lib/Target/AMDGPU/AMDGPUISelLowering.h

297–298 ↗

(On Diff #154263)

Weird formatting

[tstellar]

lib/Target/AMDGPU/AMDGPUISelLowering.cpp

3941–3945

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

3941–3945

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

3941–3945

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.cpp

SDValue AMDGPUTargetLowering::loadInputValue(SelectionDAG &DAG,
 
   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 AMDGPUMachineFunction *MFI = MF.getInfo<AMDGPUMachineFunction>();
-  const AMDGPUSubtarget &ST = MF.getSubtarget<AMDGPUSubtarget>();
+  const AMDGPUCommonSubtarget &ST =
+      AMDGPUCommonSubtarget::get(getTargetMachine(), MF.getFunction());
   unsigned ExplicitArgOffset = ST.getExplicitKernelArgOffset(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/AMDGPUSubtarget.h

public:
   bool isAmdHsaOS() const {
     return TargetTriple.getOS() == Triple::AMDHSA;
   }
 
   bool isAmdPalOS() const {
     return TargetTriple.getOS() == Triple::AMDPAL;
   }
 
+  bool isMesa3DOS() const {
+    return TargetTriple.getOS() == Triple::Mesa3D;
+  }
+
+  bool isMesaKernel(const Function &F) const {
+    return isMesa3DOS() && !AMDGPU::isShader(F.getCallingConv());
+  }
+
+  bool isAmdCodeObjectV2(const Function &F) const {
+    return isAmdHsaOS() || isMesaKernel(F);
+  }
+
   bool has16BitInsts() const {
     return Has16BitInsts;
   }
 
   bool hasMadMixInsts() const {
     return HasMadMixInsts;
   }
 
public:
   int getLocalMemorySize() const {
     return LocalMemorySize;
   }
 
   unsigned getAlignmentForImplicitArgPtr() const {
     return isAmdHsaOS() ? 8 : 4;
   }
 
+  /// Returns the offset in bytes from the start of the input buffer
+  ///        of the first explicit kernel argument.
+  unsigned getExplicitKernelArgOffset(const Function &F) const {
+    return isAmdCodeObjectV2(F) ? 0 : 36;
+  }
+
   /// \returns Maximum number of work groups per compute unit supported by the
   /// subtarget and limited by given \p FlatWorkGroupSize.
   unsigned getMaxWorkGroupsPerCU(unsigned FlatWorkGroupSize) const {
     return AMDGPU::IsaInfo::getMaxWorkGroupsPerCU(SubtargetFeatureBits,
                                                   FlatWorkGroupSize);
   }
 
   /// \returns Minimum flat work group size supported by the subtarget.
public:
 
   // Nothing implemented, just prevent crashes on use.
   const SelectionDAGTargetInfo *getSelectionDAGInfo() const override {
     return &TSInfo;
   }
 
   void ParseSubtargetFeatures(StringRef CPU, StringRef FS);
 
-  bool isMesa3DOS() const {
-    return TargetTriple.getOS() == Triple::Mesa3D;
-  }
-
   Generation getGeneration() const {
     return (Generation)Gen;
   }
 
   unsigned getWavefrontSizeLog2() const {
     return Log2_32(WavefrontSize);
   }
 
public:
   bool hasAddNoCarry() const {
     return AddNoCarryInsts;
   }
 
   bool hasUnpackedD16VMem() const {
     return HasUnpackedD16VMem;
   }
 
-  bool isMesaKernel(const Function &F) const {
-    return isMesa3DOS() && !AMDGPU::isShader(F.getCallingConv());
-  }
-
   // Covers VS/PS/CS graphics shaders
   bool isMesaGfxShader(const Function &F) const {
     return isMesa3DOS() && AMDGPU::isShader(F.getCallingConv());
   }
 
-  bool isAmdCodeObjectV2(const Function &F) const {
-    return isAmdHsaOS() || isMesaKernel(F);
-  }
-
   bool hasMad64_32() const {
     return getGeneration() >= SEA_ISLANDS;
   }
 
   bool hasSDWAOmod() const {
     return HasSDWAOmod;
   }
 
public:
   bool hasDLInsts() const {
     return HasDLInsts;
   }
 
   bool d16PreservesUnusedBits() const {
     return D16PreservesUnusedBits;
   }
 
-  /// Returns the offset in bytes from the start of the input buffer
-  ///        of the first explicit kernel argument.
-  unsigned getExplicitKernelArgOffset(const Function &F) const {
-    return isAmdCodeObjectV2(F) ? 0 : 36;
-  }
-
   /// \returns Number of bytes of arguments that are passed to a shader or
   /// kernel in addition to the explicit ones declared for the function.
   unsigned getImplicitArgNumBytes(const Function &F) const {
     if (isMesaKernel(F))
       return 16;
     return AMDGPU::getIntegerAttribute(F, "amdgpu-implicitarg-num-bytes", 0);
   }
 
public:
   }
 
   bool hasFFBH() const {
     return (getGeneration() >= EVERGREEN);
   }
 
   bool hasFMA() const { return FMA; }
 
-  unsigned getExplicitKernelArgOffset(const MachineFunction &MF) const {
-    return 36;
-  }
-
   bool hasCFAluBug() const { return CFALUBug; }
 
   bool hasVertexCache() const { return HasVertexCache; }
 
   short getTexVTXClauseSize() const { return TexVTXClauseSize; }
 
   AMDGPUAS getAMDGPUAS() const { return AS; }
 

lib/Target/AMDGPU/R600ISelLowering.cpp

for (unsigned i = 0, e = Ins.size(); i < e; ++i) {
     }
 
     // Compute the offset from the value.
     // XXX - I think PartOffset should give you this, but it seems to give the
     // size of the register which isn't useful.
 
     unsigned ValBase = ArgLocs[In.getOrigArgIndex()].getLocMemOffset();
     unsigned PartOffset = VA.getLocMemOffset();
-    unsigned Offset = Subtarget->getExplicitKernelArgOffset(MF) +
+    unsigned Offset = Subtarget->getExplicitKernelArgOffset(MF.getFunction()) +
                       VA.getLocMemOffset();
 
     MachinePointerInfo PtrInfo(UndefValue::get(PtrTy), PartOffset - ValBase);
     SDValue Arg = DAG.getLoad(
         ISD::UNINDEXED, Ext, VT, DL, Chain,
         DAG.getConstant(Offset, DL, MVT::i32), DAG.getUNDEF(MVT::i32), PtrInfo,
         MemVT, /* Alignment = */ 4, MachineMemOperand::MONonTemporal |
                                         MachineMemOperand::MODereferenceable |