AMDGPU: Fix implicitarg.ptr handling special inputs

[rampitec]

LGTM

[arsenm]

r310002

lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.h

enum PreloadedValue {
     KERNARG_SEGMENT_PTR =  3,
     DISPATCH_ID         =  4,
     FLAT_SCRATCH_INIT   =  5,
     WORKGROUP_ID_X      = 10,
     WORKGROUP_ID_Y      = 11,
     WORKGROUP_ID_Z      = 12,
     PRIVATE_SEGMENT_WAVE_BYTE_OFFSET = 14,
     IMPLICIT_BUFFER_PTR = 15,
+    IMPLICIT_ARG_PTR = 16,
 
     // VGPRS:
-    FIRST_VGPR_VALUE    = 16,
-    WORKITEM_ID_X       = FIRST_VGPR_VALUE,
-    WORKITEM_ID_Y       = 17,
-    WORKITEM_ID_Z       = 18
+    WORKITEM_ID_X       = 17,
+    WORKITEM_ID_Y       = 18,
+    WORKITEM_ID_Z       = 19,
+    FIRST_VGPR_VALUE    = WORKITEM_ID_X
   };
 
   // Kernel input registers setup for the HSA ABI in allocation order.
 
   // User SGPRs in kernels
   // XXX - Can these require argument spills?
   ArgDescriptor PrivateSegmentBuffer;
   ArgDescriptor DispatchPtr;
   ArgDescriptor QueuePtr;
   ArgDescriptor KernargSegmentPtr;
   ArgDescriptor DispatchID;
   ArgDescriptor FlatScratchInit;
   ArgDescriptor PrivateSegmentSize;
   ArgDescriptor GridWorkGroupCountX;
   ArgDescriptor GridWorkGroupCountY;
   ArgDescriptor GridWorkGroupCountZ;
 
   // System SGPRs in kernels.
   ArgDescriptor WorkGroupIDX;
   ArgDescriptor WorkGroupIDY;
   ArgDescriptor WorkGroupIDZ;
   ArgDescriptor WorkGroupInfo;
   ArgDescriptor PrivateSegmentWaveByteOffset;
 
+  // Pointer with offset from kernargsegmentptr to where special ABI arguments
+  // are passed to callable functions.
+  ArgDescriptor ImplicitArgPtr;
+
   // Input registers for non-HSA ABI
   ArgDescriptor ImplicitBufferPtr = 0;
 
   // VGPRs inputs. These are always v0, v1 and v2 for entry functions.
   ArgDescriptor WorkItemIDX;
   ArgDescriptor WorkItemIDY;
   ArgDescriptor WorkItemIDZ;
 

lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.cpp

OS << "Arguments for " << FI.first->getName() << '\n'
        << "  GridWorkgroupCountZ: " << FI.second.GridWorkGroupCountZ
        << "  WorkGroupIDX: " << FI.second.WorkGroupIDX
        << "  WorkGroupIDY: " << FI.second.WorkGroupIDY
        << "  WorkGroupIDZ: " << FI.second.WorkGroupIDZ
        << "  WorkGroupInfo: " << FI.second.WorkGroupInfo
        << "  PrivateSegmentWaveByteOffset: "
           << FI.second.PrivateSegmentWaveByteOffset
        << "  ImplicitBufferPtr: " << FI.second.ImplicitBufferPtr
+       << "  ImplicitArgPtr: " << FI.second.ImplicitArgPtr
        << "  WorkItemIDX " << FI.second.WorkItemIDX
        << "  WorkItemIDY " << FI.second.WorkItemIDY
        << "  WorkItemIDZ " << FI.second.WorkItemIDZ
        << '\n';
   }
 }
 
 std::pair<const ArgDescriptor *, const TargetRegisterClass *>
return std::make_pair(WorkGroupIDZ ? &WorkGroupIDZ : nullptr,
                           &AMDGPU::SGPR_32RegClass);
   case AMDGPUFunctionArgInfo::PRIVATE_SEGMENT_WAVE_BYTE_OFFSET:
     return std::make_pair(
       PrivateSegmentWaveByteOffset ? &PrivateSegmentWaveByteOffset : nullptr,
       &AMDGPU::SGPR_32RegClass);
   case AMDGPUFunctionArgInfo::KERNARG_SEGMENT_PTR:
     return std::make_pair(KernargSegmentPtr ? &KernargSegmentPtr : nullptr,
                           &AMDGPU::SGPR_64RegClass);
+  case AMDGPUFunctionArgInfo::IMPLICIT_ARG_PTR:
+    return std::make_pair(ImplicitArgPtr ? &ImplicitArgPtr : nullptr,
+                          &AMDGPU::SGPR_64RegClass);
   case AMDGPUFunctionArgInfo::DISPATCH_ID:
     return std::make_pair(DispatchID ? &DispatchID : nullptr,
                           &AMDGPU::SGPR_64RegClass);
   case AMDGPUFunctionArgInfo::FLAT_SCRATCH_INIT:
     return std::make_pair(FlatScratchInit ? &FlatScratchInit : nullptr,
                           &AMDGPU::SGPR_64RegClass);
   case AMDGPUFunctionArgInfo::DISPATCH_PTR:
     return std::make_pair(DispatchPtr ? &DispatchPtr : nullptr,

lib/Target/AMDGPU/SIISelLowering.cpp

static void allocateSpecialInputSGPRs(CCState &CCInfo,
   if (Info.hasWorkGroupIDX())
     ArgInfo.WorkGroupIDX = allocateSGPR32Input(CCInfo);
 
   if (Info.hasWorkGroupIDY())
     ArgInfo.WorkGroupIDY = allocateSGPR32Input(CCInfo);
 
   if (Info.hasWorkGroupIDZ())
     ArgInfo.WorkGroupIDZ = allocateSGPR32Input(CCInfo);
+
+  if (Info.hasImplicitArgPtr())
+    ArgInfo.ImplicitArgPtr = allocateSGPR64Input(CCInfo);
 }
 
 // Allocate special inputs passed in user SGPRs.
 static void allocateHSAUserSGPRs(CCState &CCInfo,
                                  MachineFunction &MF,
                                  const SIRegisterInfo &TRI,
                                  SIMachineFunctionInfo &Info) {
   if (Info.hasImplicitBufferPtr()) {
AMDGPUFunctionArgInfo::PreloadedValue InputRegs[] = {
     AMDGPUFunctionArgInfo::QUEUE_PTR,
     AMDGPUFunctionArgInfo::KERNARG_SEGMENT_PTR,
     AMDGPUFunctionArgInfo::DISPATCH_ID,
     AMDGPUFunctionArgInfo::WORKGROUP_ID_X,
     AMDGPUFunctionArgInfo::WORKGROUP_ID_Y,
     AMDGPUFunctionArgInfo::WORKGROUP_ID_Z,
     AMDGPUFunctionArgInfo::WORKITEM_ID_X,
     AMDGPUFunctionArgInfo::WORKITEM_ID_Y,
-    AMDGPUFunctionArgInfo::WORKITEM_ID_Z
+    AMDGPUFunctionArgInfo::WORKITEM_ID_Z,
+    AMDGPUFunctionArgInfo::IMPLICIT_ARG_PTR
   };
 
   for (auto InputID : InputRegs) {
     const ArgDescriptor *OutgoingArg;
     const TargetRegisterClass *ArgRC;
 
     std::tie(OutgoingArg, ArgRC) = CalleeArgInfo.getPreloadedValue(InputID);
     if (!OutgoingArg)
       continue;
 
     const ArgDescriptor *IncomingArg;
     const TargetRegisterClass *IncomingArgRC;
     std::tie(IncomingArg, IncomingArgRC)
       = 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 = loadInputValue(DAG, ArgRC, ArgVT, DL, *IncomingArg);
+    SDValue InputReg;
+
+    if (IncomingArg) {
+      InputReg = loadInputValue(DAG, ArgRC, ArgVT, DL, *IncomingArg);
+    } else {
+      // The implicit arg ptr is special because it doesn't have a corresponding
+      // input for kernels, and is computed from the kernarg segment pointer.
+      assert(InputID == AMDGPUFunctionArgInfo::IMPLICIT_ARG_PTR);
+      InputReg = getImplicitArgPtr(DAG, DL);
+    }
+
     if (OutgoingArg->isRegister()) {
       RegsToPass.emplace_back(OutgoingArg->getRegister(), InputReg);
     } else {
       SDValue ArgStore = storeStackInputValue(DAG, DL, Chain, StackPtr,
                                               InputReg,
                                               OutgoingArg->getStackOffset());
       MemOpChains.push_back(ArgStore);
     }
case Intrinsic::amdgcn_queue_ptr: {
 
     auto RegID = IntrinsicID == Intrinsic::amdgcn_dispatch_ptr ?
       AMDGPUFunctionArgInfo::DISPATCH_PTR : AMDGPUFunctionArgInfo::QUEUE_PTR;
     return getPreloadedValue(DAG, *MFI, VT, RegID);
   }
   case Intrinsic::amdgcn_implicitarg_ptr: {
     if (MFI->isEntryFunction())
       return getImplicitArgPtr(DAG, DL);
-    report_fatal_error("amdgcn.implicitarg.ptr not implemented for functions");
+    return getPreloadedValue(DAG, *MFI, VT,
+                             AMDGPUFunctionArgInfo::IMPLICIT_ARG_PTR);
   }
   case Intrinsic::amdgcn_kernarg_segment_ptr: {
     return getPreloadedValue(DAG, *MFI, VT,
                              AMDGPUFunctionArgInfo::KERNARG_SEGMENT_PTR);
   }
   case Intrinsic::amdgcn_dispatch_id: {
     return getPreloadedValue(DAG, *MFI, VT, AMDGPUFunctionArgInfo::DISPATCH_ID);
   }

lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

: AMDGPUMachineFunction(MF),
     WorkGroupIDX(false),
     WorkGroupIDY(false),
     WorkGroupIDZ(false),
     WorkGroupInfo(false),
     PrivateSegmentWaveByteOffset(false),
     WorkItemIDX(false),
     WorkItemIDY(false),
     WorkItemIDZ(false),
-    ImplicitBufferPtr(false) {
+    ImplicitBufferPtr(false),
+    ImplicitArgPtr(false) {
   const SISubtarget &ST = MF.getSubtarget<SISubtarget>();
   const Function *F = MF.getFunction();
   FlatWorkGroupSizes = ST.getFlatWorkGroupSizes(*F);
   WavesPerEU = ST.getWavesPerEU(*F);
 
   if (!isEntryFunction()) {
     // Non-entry functions have no special inputs for now, other registers
     // required for scratch access.

test/CodeGen/AMDGPU/llvm.amdgcn.implicitarg.ptr.ll

-; RUN: llc -mtriple=amdgcn-amd-amdhsa -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,HSA,HSA-NOENV %s
-; RUN: llc -mtriple=amdgcn-amd-amdhsa-opencl -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,HSA,HSA-OPENCL %s
-; RUN: llc -mtriple=amdgcn-mesa-mesa3d -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,MESA %s
+; RUN: llc -mtriple=amdgcn-amd-amdhsa -amdgpu-function-calls -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,HSA,HSA-NOENV %s
+; RUN: llc -mtriple=amdgcn-amd-amdhsa-opencl -amdgpu-function-calls -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,HSA,HSA-OPENCL %s
+; RUN: llc -mtriple=amdgcn-mesa-mesa3d -amdgpu-function-calls -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,MESA %s
 
 ; GCN-LABEL: {{^}}kernel_implicitarg_ptr_empty:
 ; GCN: enable_sgpr_kernarg_segment_ptr = 1
 
 ; HSA-NOENV: kernarg_segment_byte_size = 0
 ; HSA-OPENCL: kernarg_segment_byte_size = 32
 ; MESA: kernarg_segment_byte_size = 16
 

 ; HSA: s_load_dword s0, s[4:5], 0x1c
 define amdgpu_kernel void @kernel_implicitarg_ptr([112 x i8]) #0 {
   %implicitarg.ptr = call i8 addrspace(2)* @llvm.amdgcn.implicitarg.ptr()
   %cast = bitcast i8 addrspace(2)* %implicitarg.ptr to i32 addrspace(2)*
   %load = load volatile i32, i32 addrspace(2)* %cast
   ret void
 }
 
+; GCN-LABEL: {{^}}func_implicitarg_ptr:
+; GCN: s_waitcnt
+; GCN-NEXT: s_load_dword s{{[0-9]+}}, s[6:7], 0x0{{$}}
+; GCN-NEXT: s_waitcnt
+; GCN-NEXT: s_setpc_b64
+define void @func_implicitarg_ptr() #1 {
+  %implicitarg.ptr = call i8 addrspace(2)* @llvm.amdgcn.implicitarg.ptr()
+  %cast = bitcast i8 addrspace(2)* %implicitarg.ptr to i32 addrspace(2)*
+  %load = load volatile i32, i32 addrspace(2)* %cast
+  ret void
+}
+
+; GCN-LABEL: {{^}}kernel_call_implicitarg_ptr_func_empty:
+; GCN: enable_sgpr_kernarg_segment_ptr = 1
+; HSA-NOENV: kernarg_segment_byte_size = 0
+; HSA-OPENCL: kernarg_segment_byte_size = 32
+; MESA: kernarg_segment_byte_size = 16
+; GCN: s_mov_b64 s[6:7], s[4:5]
+; GCN: s_swappc_b64
+define amdgpu_kernel void @kernel_call_implicitarg_ptr_func_empty() #0 {
+  call void @func_implicitarg_ptr()
+  ret void
+}
+
+; GCN-LABEL: {{^}}kernel_call_implicitarg_ptr_func:
+; GCN: enable_sgpr_kernarg_segment_ptr = 1
+; HSA-OPENCL: kernarg_segment_byte_size = 144
+; HSA-NOENV: kernarg_segment_byte_size = 112
+; MESA: kernarg_segment_byte_size = 464
+
+; HSA: s_add_u32 s6, s4, 0x70
+; MESA: s_add_u32 s6, s4, 0x1c0
+
+; GCN: s_addc_u32 s7, s5, 0{{$}}
+; GCN: s_swappc_b64
+define amdgpu_kernel void @kernel_call_implicitarg_ptr_func([112 x i8]) #0 {
+  call void @func_implicitarg_ptr()
+  ret void
+}
+
+; GCN-LABEL: {{^}}func_call_implicitarg_ptr_func:
+; GCN-NOT: s6
+; GCN-NOT: s7
+; GCN-NOT: s[6:7]
+define void @func_call_implicitarg_ptr_func() #1 {
+  call void @func_implicitarg_ptr()
+  ret void
+}
+
+; GCN-LABEL: {{^}}func_kernarg_implicitarg_ptr:
+; GCN: s_waitcnt
+; GCN: s_load_dword s{{[0-9]+}}, s[6:7], 0x0{{$}}
+; GCN: s_load_dword s{{[0-9]+}}, s[8:9], 0x0{{$}}
+define void @func_kernarg_implicitarg_ptr() #1 {
+  %kernarg.segment.ptr = call i8 addrspace(2)* @llvm.amdgcn.kernarg.segment.ptr()
+  %implicitarg.ptr = call i8 addrspace(2)* @llvm.amdgcn.implicitarg.ptr()
+  %cast.kernarg.segment.ptr = bitcast i8 addrspace(2)* %kernarg.segment.ptr to i32 addrspace(2)*
+  %cast.implicitarg = bitcast i8 addrspace(2)* %implicitarg.ptr to i32 addrspace(2)*
+  %load0 = load volatile i32, i32 addrspace(2)* %cast.kernarg.segment.ptr
+  %load1 = load volatile i32, i32 addrspace(2)* %cast.implicitarg
+  ret void
+}
+
+; GCN-LABEL: {{^}}kernel_call_kernarg_implicitarg_ptr_func:
+; GCN: s_mov_b64 s[6:7], s[4:5]
+; HSA: s_add_u32 s8, s6, 0x70
+; MESA: s_add_u32 s8, s6, 0x1c0
+; GCN: s_addc_u32 s9, s7, 0
+; GCN: s_swappc_b64
+define amdgpu_kernel void @kernel_call_kernarg_implicitarg_ptr_func([112 x i8]) #0 {
+  call void @func_kernarg_implicitarg_ptr()
+  ret void
+}
+
 declare i8 addrspace(2)* @llvm.amdgcn.implicitarg.ptr() #2
+declare i8 addrspace(2)* @llvm.amdgcn.kernarg.segment.ptr() #2
 
 attributes #0 = { nounwind noinline }
 attributes #1 = { nounwind noinline }
 attributes #2 = { nounwind readnone speculatable }