AMDGPU: Fix implicitarg.ptr handling special inputs

llvm-svn: 310002

5 files changed, 111 insertions, 11 deletions

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.cpp b/llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.cpp
index cc2f99d08d7..3533ec30505 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.cpp
@@ -65,6 +65,7 @@ void AMDGPUArgumentUsageInfo::print(raw_ostream &OS, const Module *M) const {
        << "  PrivateSegmentWaveByteOffset: "
           << FI.second.PrivateSegmentWaveByteOffset
        << "  ImplicitBufferPtr: " << FI.second.ImplicitBufferPtr
+       << "  ImplicitArgPtr: " << FI.second.ImplicitArgPtr
        << "  WorkItemIDX " << FI.second.WorkItemIDX
        << "  WorkItemIDY " << FI.second.WorkItemIDY
        << "  WorkItemIDZ " << FI.second.WorkItemIDZ
@@ -101,6 +102,9 @@ AMDGPUFunctionArgInfo::getPreloadedValue(
   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);
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.h b/llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.h
index 0379d83f558..9373b80e87b 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.h
@@ -90,12 +90,13 @@ struct AMDGPUFunctionArgInfo {
     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.
@@ -120,6 +121,10 @@ struct AMDGPUFunctionArgInfo {
   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;
 
diff --git a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
index 5a53d7914c0..da7d04bad25 100644
--- a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
+++ b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
@@ -1195,6 +1195,9 @@ static void allocateSpecialInputSGPRs(CCState &CCInfo,
 
   if (Info.hasWorkGroupIDZ())
     ArgInfo.WorkGroupIDZ = allocateSGPR32Input(CCInfo);
+
+  if (Info.hasImplicitArgPtr())
+    ArgInfo.ImplicitArgPtr = allocateSGPR64Input(CCInfo);
 }
 
 // Allocate special inputs passed in user SGPRs.
@@ -1914,7 +1917,8 @@ void SITargetLowering::passSpecialInputs(
     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) {
@@ -1933,7 +1937,17 @@ void SITargetLowering::passSpecialInputs(
 
     // 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 {
@@ -3662,7 +3676,8 @@ SDValue SITargetLowering::LowerINTRINSIC_WO_CHAIN(SDValue Op,
   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,
diff --git a/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp b/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp
index c4405309e12..baf26039a5f 100644
--- a/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp
@@ -60,7 +60,8 @@ SIMachineFunctionInfo::SIMachineFunctionInfo(const MachineFunction &MF)
     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);
diff --git a/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.implicitarg.ptr.ll b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.implicitarg.ptr.ll
index 7edb8533672..71088844146 100644
--- a/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.implicitarg.ptr.ll
+++ b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.implicitarg.ptr.ll
@@ -1,6 +1,6 @@
-; 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
@@ -32,7 +32,82 @@ define amdgpu_kernel void @kernel_implicitarg_ptr([112 x i8]) #0 {
   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 }