AMDGPU/GlobalISel: Add types to special inputs

[arsenm]

Summary

When passing special ABI inputs, we have no existing context for the
type to use.

[arsenm]

f25d020c2ec7cb1971fa56b99381d416799d8145

llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.h

 //==- AMDGPUArgumentrUsageInfo.h - Function Arg Usage Info -------*- C++ -*-==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_TARGET_AMDGPU_AMDGPUARGUMENTUSAGEINFO_H
 #define LLVM_LIB_TARGET_AMDGPU_AMDGPUARGUMENTUSAGEINFO_H
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/CodeGen/Register.h"
+#include "llvm/Support/LowLevelTypeImpl.h"
 #include "llvm/Pass.h"
 
 namespace llvm {
 
 class Function;
 class raw_ostream;
 class TargetRegisterClass;
 class TargetRegisterInfo;
struct AMDGPUFunctionArgInfo {
   // Input registers for non-HSA ABI
   ArgDescriptor ImplicitBufferPtr;
 
   // VGPRs inputs. These are always v0, v1 and v2 for entry functions.
   ArgDescriptor WorkItemIDX;
   ArgDescriptor WorkItemIDY;
   ArgDescriptor WorkItemIDZ;
 
-  std::pair<const ArgDescriptor *, const TargetRegisterClass *>
+  std::tuple<const ArgDescriptor *, const TargetRegisterClass *, LLT>
   getPreloadedValue(PreloadedValue Value) const;
 
   static constexpr AMDGPUFunctionArgInfo fixedABILayout();
 };
 
 class AMDGPUArgumentUsageInfo : public ImmutablePass {
 private:
   DenseMap<const Function *, AMDGPUFunctionArgInfo> ArgInfoMap;

llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.cpp

OS << "Arguments for " << FI.first->getName() << '\n'
        << "  ImplicitArgPtr: " << FI.second.ImplicitArgPtr
        << "  WorkItemIDX " << FI.second.WorkItemIDX
        << "  WorkItemIDY " << FI.second.WorkItemIDY
        << "  WorkItemIDZ " << FI.second.WorkItemIDZ
        << '\n';
   }
 }
 
-std::pair<const ArgDescriptor *, const TargetRegisterClass *>
+std::tuple<const ArgDescriptor *, const TargetRegisterClass *, LLT>
 AMDGPUFunctionArgInfo::getPreloadedValue(
   AMDGPUFunctionArgInfo::PreloadedValue Value) const {
   switch (Value) {
   case AMDGPUFunctionArgInfo::PRIVATE_SEGMENT_BUFFER: {
-    return std::make_pair(
+    return std::make_tuple(
       PrivateSegmentBuffer ? &PrivateSegmentBuffer : nullptr,
-      &AMDGPU::SGPR_128RegClass);
+      &AMDGPU::SGPR_128RegClass,
+      LLT::vector(4, 32));
   }
   case AMDGPUFunctionArgInfo::IMPLICIT_BUFFER_PTR:
-    return std::make_pair(ImplicitBufferPtr ? &ImplicitBufferPtr : nullptr,
-                          &AMDGPU::SGPR_64RegClass);
+    return std::make_tuple(ImplicitBufferPtr ? &ImplicitBufferPtr : nullptr,
+                          &AMDGPU::SGPR_64RegClass,
+                          LLT::pointer(AMDGPUAS::CONSTANT_ADDRESS, 64));
   case AMDGPUFunctionArgInfo::WORKGROUP_ID_X:
-    return std::make_pair(WorkGroupIDX ? &WorkGroupIDX : nullptr,
-                          &AMDGPU::SGPR_32RegClass);
-
+    return std::make_tuple(WorkGroupIDX ? &WorkGroupIDX : nullptr,
+                           &AMDGPU::SGPR_32RegClass, LLT::scalar(32));
   case AMDGPUFunctionArgInfo::WORKGROUP_ID_Y:
-    return std::make_pair(WorkGroupIDY ? &WorkGroupIDY : nullptr,
-                          &AMDGPU::SGPR_32RegClass);
+    return std::make_tuple(WorkGroupIDY ? &WorkGroupIDY : nullptr,
+                           &AMDGPU::SGPR_32RegClass, LLT::scalar(32));
   case AMDGPUFunctionArgInfo::WORKGROUP_ID_Z:
-    return std::make_pair(WorkGroupIDZ ? &WorkGroupIDZ : nullptr,
-                          &AMDGPU::SGPR_32RegClass);
+    return std::make_tuple(WorkGroupIDZ ? &WorkGroupIDZ : nullptr,
+                           &AMDGPU::SGPR_32RegClass, LLT::scalar(32));
   case AMDGPUFunctionArgInfo::PRIVATE_SEGMENT_WAVE_BYTE_OFFSET:
-    return std::make_pair(
+    return std::make_tuple(
       PrivateSegmentWaveByteOffset ? &PrivateSegmentWaveByteOffset : nullptr,
-      &AMDGPU::SGPR_32RegClass);
+      &AMDGPU::SGPR_32RegClass, LLT::scalar(32));
   case AMDGPUFunctionArgInfo::KERNARG_SEGMENT_PTR:
-    return std::make_pair(KernargSegmentPtr ? &KernargSegmentPtr : nullptr,
-                          &AMDGPU::SGPR_64RegClass);
+    return std::make_tuple(KernargSegmentPtr ? &KernargSegmentPtr : nullptr,
+                           &AMDGPU::SGPR_64RegClass,
+                           LLT::pointer(AMDGPUAS::CONSTANT_ADDRESS, 64));
   case AMDGPUFunctionArgInfo::IMPLICIT_ARG_PTR:
-    return std::make_pair(ImplicitArgPtr ? &ImplicitArgPtr : nullptr,
-                          &AMDGPU::SGPR_64RegClass);
+    return std::make_tuple(ImplicitArgPtr ? &ImplicitArgPtr : nullptr,
+                           &AMDGPU::SGPR_64RegClass,
+                           LLT::pointer(AMDGPUAS::CONSTANT_ADDRESS, 64));
   case AMDGPUFunctionArgInfo::DISPATCH_ID:
-    return std::make_pair(DispatchID ? &DispatchID : nullptr,
-                          &AMDGPU::SGPR_64RegClass);
+    return std::make_tuple(DispatchID ? &DispatchID : nullptr,
+                           &AMDGPU::SGPR_64RegClass, LLT::scalar(64));
   case AMDGPUFunctionArgInfo::FLAT_SCRATCH_INIT:
-    return std::make_pair(FlatScratchInit ? &FlatScratchInit : nullptr,
-                          &AMDGPU::SGPR_64RegClass);
+    return std::make_tuple(FlatScratchInit ? &FlatScratchInit : nullptr,
+                           &AMDGPU::SGPR_64RegClass, LLT::scalar(64));
   case AMDGPUFunctionArgInfo::DISPATCH_PTR:
-    return std::make_pair(DispatchPtr ? &DispatchPtr : nullptr,
-                          &AMDGPU::SGPR_64RegClass);
+    return std::make_tuple(DispatchPtr ? &DispatchPtr : nullptr,
+                           &AMDGPU::SGPR_64RegClass,
+                           LLT::pointer(AMDGPUAS::CONSTANT_ADDRESS, 64));
   case AMDGPUFunctionArgInfo::QUEUE_PTR:
-    return std::make_pair(QueuePtr ? &QueuePtr : nullptr,
-                          &AMDGPU::SGPR_64RegClass);
+    return std::make_tuple(QueuePtr ? &QueuePtr : nullptr,
+                           &AMDGPU::SGPR_64RegClass,
+                           LLT::pointer(AMDGPUAS::CONSTANT_ADDRESS, 64));
   case AMDGPUFunctionArgInfo::WORKITEM_ID_X:
-    return std::make_pair(WorkItemIDX ? &WorkItemIDX : nullptr,
-                          &AMDGPU::VGPR_32RegClass);
+    return std::make_tuple(WorkItemIDX ? &WorkItemIDX : nullptr,
+                           &AMDGPU::VGPR_32RegClass, LLT::scalar(32));
   case AMDGPUFunctionArgInfo::WORKITEM_ID_Y:
-    return std::make_pair(WorkItemIDY ? &WorkItemIDY : nullptr,
-                          &AMDGPU::VGPR_32RegClass);
+    return std::make_tuple(WorkItemIDY ? &WorkItemIDY : nullptr,
+                           &AMDGPU::VGPR_32RegClass, LLT::scalar(32));
   case AMDGPUFunctionArgInfo::WORKITEM_ID_Z:
-    return std::make_pair(WorkItemIDZ ? &WorkItemIDZ : nullptr,
-                          &AMDGPU::VGPR_32RegClass);
+    return std::make_tuple(WorkItemIDZ ? &WorkItemIDZ : nullptr,
+                           &AMDGPU::VGPR_32RegClass, LLT::scalar(32));
   }
   llvm_unreachable("unexpected preloaded value type");
 }
 
 constexpr AMDGPUFunctionArgInfo AMDGPUFunctionArgInfo::fixedABILayout() {
   AMDGPUFunctionArgInfo AI;
   AI.PrivateSegmentBuffer
     = ArgDescriptor::createRegister(AMDGPU::SGPR0_SGPR1_SGPR2_SGPR3);

llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

Register AMDGPULegalizerInfo::getLiveInRegister(MachineIRBuilder &B,
   return insertLiveInCopy(B, MRI, LiveIn, PhyReg);
 }
 
 const ArgDescriptor *AMDGPULegalizerInfo::getArgDescriptor(
     MachineIRBuilder &B, AMDGPUFunctionArgInfo::PreloadedValue ArgType) const {
   const SIMachineFunctionInfo *MFI = B.getMF().getInfo<SIMachineFunctionInfo>();
   const ArgDescriptor *Arg;
   const TargetRegisterClass *RC;
-  std::tie(Arg, RC) = MFI->getPreloadedValue(ArgType);
+  LLT ArgTy;
+  std::tie(Arg, RC, ArgTy) = MFI->getPreloadedValue(ArgType);
   if (!Arg) {
     LLVM_DEBUG(dbgs() << "Required arg register missing\n");
     return nullptr;
   }
   return Arg;
 }
 
 bool AMDGPULegalizerInfo::loadInputValue(Register DstReg, MachineIRBuilder &B,
uint64_t Offset =
     ST.getTargetLowering()->getImplicitParameterOffset(
       B.getMF(), AMDGPUTargetLowering::FIRST_IMPLICIT);
   Register DstReg = MI.getOperand(0).getReg();
   LLT DstTy = MRI.getType(DstReg);
   LLT IdxTy = LLT::scalar(DstTy.getSizeInBits());
 
   const ArgDescriptor *Arg;
   const TargetRegisterClass *RC;
-  std::tie(Arg, RC)
+  LLT ArgTy;
+  std::tie(Arg, RC, ArgTy)
     = MFI->getPreloadedValue(AMDGPUFunctionArgInfo::KERNARG_SEGMENT_PTR);
   if (!Arg)
     return false;
 
   Register KernargPtrReg = MRI.createGenericVirtualRegister(DstTy);
   if (!loadInputValue(KernargPtrReg, B, Arg))
     return false;
 

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

SDValue SITargetLowering::lowerKernArgParameterPtr(SelectionDAG &DAG,
                                                    SDValue Chain,
                                                    uint64_t Offset) const {
   const DataLayout &DL = DAG.getDataLayout();
   MachineFunction &MF = DAG.getMachineFunction();
   const SIMachineFunctionInfo *Info = MF.getInfo<SIMachineFunctionInfo>();
 
   const ArgDescriptor *InputPtrReg;
   const TargetRegisterClass *RC;
+  LLT ArgTy;
 
-  std::tie(InputPtrReg, RC)
+  std::tie(InputPtrReg, RC, ArgTy)
     = Info->getPreloadedValue(AMDGPUFunctionArgInfo::KERNARG_SEGMENT_PTR);
 
   MachineRegisterInfo &MRI = DAG.getMachineFunction().getRegInfo();
   MVT PtrVT = getPointerTy(DL, AMDGPUAS::CONSTANT_ADDRESS);
   SDValue BasePtr = DAG.getCopyFromReg(Chain, SL,
     MRI.getLiveInVirtReg(InputPtrReg->getRegister()), PtrVT);
 
   return DAG.getObjectPtrOffset(SL, BasePtr, Offset);

 }
 
 SDValue SITargetLowering::getPreloadedValue(SelectionDAG &DAG,
   const SIMachineFunctionInfo &MFI,
   EVT VT,
   AMDGPUFunctionArgInfo::PreloadedValue PVID) const {
   const ArgDescriptor *Reg;
   const TargetRegisterClass *RC;
+  LLT Ty;
 
-  std::tie(Reg, RC) = MFI.getPreloadedValue(PVID);
+  std::tie(Reg, RC, Ty) = MFI.getPreloadedValue(PVID);
   return CreateLiveInRegister(DAG, RC, Reg->getRegister(), VT);
 }
 
 static void processShaderInputArgs(SmallVectorImpl<ISD::InputArg> &Splits,
                                    CallingConv::ID CallConv,
                                    ArrayRef<ISD::InputArg> Ins,
                                    BitVector &Skipped,
                                    FunctionType *FType,
AMDGPUFunctionArgInfo::PreloadedValue InputRegs[] = {
     AMDGPUFunctionArgInfo::WORKGROUP_ID_X,
     AMDGPUFunctionArgInfo::WORKGROUP_ID_Y,
     AMDGPUFunctionArgInfo::WORKGROUP_ID_Z
   };
 
   for (auto InputID : InputRegs) {
     const ArgDescriptor *OutgoingArg;
     const TargetRegisterClass *ArgRC;
+    LLT ArgTy;
 
-    std::tie(OutgoingArg, ArgRC) = CalleeArgInfo->getPreloadedValue(InputID);
+    std::tie(OutgoingArg, ArgRC, ArgTy)
+      = CalleeArgInfo->getPreloadedValue(InputID);
     if (!OutgoingArg)
       continue;
 
     const ArgDescriptor *IncomingArg;
     const TargetRegisterClass *IncomingArgRC;
-    std::tie(IncomingArg, IncomingArgRC)
+    LLT Ty;
+    std::tie(IncomingArg, IncomingArgRC, Ty)
       = CallerArgInfo.getPreloadedValue(InputID);
     assert(IncomingArgRC == ArgRC);
 
     // All special arguments are ints for now.
     EVT ArgVT = TRI->getSpillSize(*ArgRC) == 8 ? MVT::i64 : MVT::i32;
     SDValue InputReg;
 
     if (IncomingArg) {
if (OutgoingArg->isRegister()) {
       MemOpChains.push_back(ArgStore);
     }
   }
 
   // Pack workitem IDs into a single register or pass it as is if already
   // packed.
   const ArgDescriptor *OutgoingArg;
   const TargetRegisterClass *ArgRC;
+  LLT Ty;
 
-  std::tie(OutgoingArg, ArgRC) =
+  std::tie(OutgoingArg, ArgRC, Ty) =
     CalleeArgInfo->getPreloadedValue(AMDGPUFunctionArgInfo::WORKITEM_ID_X);
   if (!OutgoingArg)
-    std::tie(OutgoingArg, ArgRC) =
+    std::tie(OutgoingArg, ArgRC, Ty) =
       CalleeArgInfo->getPreloadedValue(AMDGPUFunctionArgInfo::WORKITEM_ID_Y);
   if (!OutgoingArg)
-    std::tie(OutgoingArg, ArgRC) =
+    std::tie(OutgoingArg, ArgRC, Ty) =
       CalleeArgInfo->getPreloadedValue(AMDGPUFunctionArgInfo::WORKITEM_ID_Z);
   if (!OutgoingArg)
     return;
 
   const ArgDescriptor *IncomingArgX
-    = CallerArgInfo.getPreloadedValue(AMDGPUFunctionArgInfo::WORKITEM_ID_X).first;
+    = std::get<0>(CallerArgInfo.getPreloadedValue(AMDGPUFunctionArgInfo::WORKITEM_ID_X));
   const ArgDescriptor *IncomingArgY
-    = CallerArgInfo.getPreloadedValue(AMDGPUFunctionArgInfo::WORKITEM_ID_Y).first;
+    = std::get<0>(CallerArgInfo.getPreloadedValue(AMDGPUFunctionArgInfo::WORKITEM_ID_Y));
   const ArgDescriptor *IncomingArgZ
-    = CallerArgInfo.getPreloadedValue(AMDGPUFunctionArgInfo::WORKITEM_ID_Z).first;
+    = std::get<0>(CallerArgInfo.getPreloadedValue(AMDGPUFunctionArgInfo::WORKITEM_ID_Z));
 
   SDValue InputReg;
   SDLoc SL;
 
   // If incoming ids are not packed we need to pack them.
   if (IncomingArgX && !IncomingArgX->isMasked() && CalleeArgInfo->WorkItemIDX)
     InputReg = loadInputValue(DAG, ArgRC, MVT::i32, DL, *IncomingArgX);
 

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

public:
   AMDGPUFunctionArgInfo &getArgInfo() {
     return ArgInfo;
   }
 
   const AMDGPUFunctionArgInfo &getArgInfo() const {
     return ArgInfo;
   }
 
-  std::pair<const ArgDescriptor *, const TargetRegisterClass *>
+  std::tuple<const ArgDescriptor *, const TargetRegisterClass *, LLT>
   getPreloadedValue(AMDGPUFunctionArgInfo::PreloadedValue Value) const {
     return ArgInfo.getPreloadedValue(Value);
   }
 
   Register getPreloadedReg(AMDGPUFunctionArgInfo::PreloadedValue Value) const {
-    auto Arg = ArgInfo.getPreloadedValue(Value).first;
+    auto Arg = std::get<0>(ArgInfo.getPreloadedValue(Value));
     return Arg ? Arg->getRegister() : Register();
   }
 
   unsigned getGITPtrHigh() const {
     return GITPtrHigh;
   }
 
   Register getGITPtrLoReg(const MachineFunction &MF) const;