[AMDGPU] Emit metadata for hidden arguments for kernel enqueue

[yaxunl]

Summary

Identifies kernels which performs device side kernel enqueues and emit
metadata for the associated hidden kernel arguments. Such kernels are
marked with "calls-enqueue-kernel" function attribute by
AMDGPUOpenCLEnqueueKernelLowering pass and later on
hidden kernel arguments metadata HiddenDefaultQueue and
HiddenCompletionAction are emitted for them.

[t-tye]

If the kernel does not use the hidden arguments should they be listed in the metadata since the runtime is not required to pass them in? If not needed, the we still want to have a kernarg location at the known location, it will just not be initialized by the runtime.

[yaxunl]

In D39255#905627, @t-tye wrote:

If the kernel does not use the hidden arguments should they be listed in the metadata since the runtime is not required to pass them in? If not needed, the we still want to have a kernarg location at the known location, it will just not be initialized by the runtime.

In that case we just need runtime to fill in 0 values for the missing argument so that each hidden argument is in the fixed position. e.g. if runtime does not find printf metadata, it will still insert a nullptr as the PrintfBuffer argument. I think only library code uses these hidden arguments. Then library code can assume fixed position for each hidden argument.

[t-tye]

In D39255#905658, @yaxunl wrote:

In D39255#905627, @t-tye wrote:

If the kernel does not use the hidden arguments should they be listed in the metadata since the runtime is not required to pass them in? If not needed, the we still want to have a kernarg location at the known location, it will just not be initialized by the runtime.

In that case we just need runtime to fill in 0 values for the missing argument so that each hidden argument is in the fixed position. e.g. if runtime does not find printf metadata, it will still insert a nullptr as the PrintfBuffer argument. I think only library code uses these hidden arguments. Then library code can assume fixed position for each hidden argument.

So will the metadata only be generated for the hidden arguments that actually need to be initialized by the runtime? I believe the metadata format contains the position of the argument so it is possible to set up non-contiguous arguments (namely, can skip the ones that will not be used).

[yaxunl]

Revised to emit hidden kernel argument metadata only when necessary.

[yaxunl]

In D39255#905873, @t-tye wrote:

In D39255#905658, @yaxunl wrote:

In D39255#905627, @t-tye wrote:

If the kernel does not use the hidden arguments should they be listed in the metadata since the runtime is not required to pass them in? If not needed, the we still want to have a kernarg location at the known location, it will just not be initialized by the runtime.

In that case we just need runtime to fill in 0 values for the missing argument so that each hidden argument is in the fixed position. e.g. if runtime does not find printf metadata, it will still insert a nullptr as the PrintfBuffer argument. I think only library code uses these hidden arguments. Then library code can assume fixed position for each hidden argument.

So will the metadata only be generated for the hidden arguments that actually need to be initialized by the runtime? I believe the metadata format contains the position of the argument so it is possible to set up non-contiguous arguments (namely, can skip the ones that will not be used).

I have updated the patch so that the metadata are generated only when they are needed. The metadata format does not contain the position of the argument. The runtime needs to always emit the Printf hidden kernel argument so that library code does not need to know if it is emitted or not.

[t-tye]

In D39255#906591, @yaxunl wrote:

In D39255#905873, @t-tye wrote:

In D39255#905658, @yaxunl wrote:

In D39255#905627, @t-tye wrote:

If the kernel does not use the hidden arguments should they be listed in the metadata since the runtime is not required to pass them in? If not needed, the we still want to have a kernarg location at the known location, it will just not be initialized by the runtime.

In that case we just need runtime to fill in 0 values for the missing argument so that each hidden argument is in the fixed position. e.g. if runtime does not find printf metadata, it will still insert a nullptr as the PrintfBuffer argument. I think only library code uses these hidden arguments. Then library code can assume fixed position for each hidden argument.

So will the metadata only be generated for the hidden arguments that actually need to be initialized by the runtime? I believe the metadata format contains the position of the argument so it is possible to set up non-contiguous arguments (namely, can skip the ones that will not be used).

I have updated the patch so that the metadata are generated only when they are needed. The metadata format does not contain the position of the argument. The runtime needs to always emit the Printf hidden kernel argument so that library code does not need to know if it is emitted or not.

I think the metadata has a HIDDENNONE kind which needs to be used for kernel arguments that need no set up.

[yaxunl]

In D39255#906718, @t-tye wrote:

In D39255#906591, @yaxunl wrote:

In D39255#905873, @t-tye wrote:

In D39255#905658, @yaxunl wrote:

In D39255#905627, @t-tye wrote:

If the kernel does not use the hidden arguments should they be listed in the metadata since the runtime is not required to pass them in? If not needed, the we still want to have a kernarg location at the known location, it will just not be initialized by the runtime.

In that case we just need runtime to fill in 0 values for the missing argument so that each hidden argument is in the fixed position. e.g. if runtime does not find printf metadata, it will still insert a nullptr as the PrintfBuffer argument. I think only library code uses these hidden arguments. Then library code can assume fixed position for each hidden argument.

So will the metadata only be generated for the hidden arguments that actually need to be initialized by the runtime? I believe the metadata format contains the position of the argument so it is possible to set up non-contiguous arguments (namely, can skip the ones that will not be used).

I have updated the patch so that the metadata are generated only when they are needed. The metadata format does not contain the position of the argument. The runtime needs to always emit the Printf hidden kernel argument so that library code does not need to know if it is emitted or not.

I think the metadata has a HIDDENNONE kind which needs to be used for kernel arguments that need no set up.

Good idea. Will do.

[yaxunl]

Emit dummy hidden argument when necessary.

[t-tye]

lib/Target/AMDGPU/MCTargetDesc/AMDGPUHSAMetadataStreamer.cpp

265–267

Should these only be requested if there are uses of the corresponding get intrinsics along the lines of only generating the following ones if they are actually used? If not generated then the HIDDENNONE needs to be generated if necessary.

[yaxunl]

lib/Target/AMDGPU/MCTargetDesc/AMDGPUHSAMetadataStreamer.cpp

265–267

Can we do that in a separate patch? This patch is intended for metadata for enqueue_kernel.

[kzhuravl]

Can we do that in a separate patch? This patch is intended for metadata for enqueue_kernel.

I think separate patch for this is better.

[kzhuravl]

lib/Target/AMDGPU/MCTargetDesc/AMDGPUHSAMetadataStreamer.cpp

265–267

I have this and other minor cleanups partially completed.

docs/AMDGPUUsage.rst

                                                 "HiddenDefaultQueue"
                                                   A global address space pointer
                                                   to the OpenCL device enqueue
                                                   queue that should be used by
                                                   the kernel by default is
                                                   passed in the kernarg.
 
                                                 "HiddenCompletionAction"
-                                                  *TBD*
-
-                                                  .. TODO
-                                                     Add description.
+                                                  A global address space pointer
+                                                  to help link enqueued kernels into
+                                                  the ancestor tree for determining
+                                                  when the parent kernel has finished.
 
      "ValueType"       string         Required  Kernel argument value type. Only
                                                 present if "ValueKind" is
                                                 "ByValue". For vector data
                                                 types, the value is for the
                                                 element type. Values include:
 
                                                 - "Struct"

lib/Target/AMDGPU/AMDGPUOpenCLEnqueuedBlockLowering.cpp

 // literal is actually runtime handle and loads the kernel address from it
 // and put it into AQL packet for dispatching.
 //
 // This cannot be done in FE since FE cannot create a unique global variable
 // with external linkage across LLVM modules. The global variable with internal
 // linkage does not work since optimization passes will try to replace loads
 // of the global variable with its initialization value.
 //
+// It also identifies the kernels directly or indirectly enqueues kernels
+// and adds "calls-enqueue-kernel" function attribute to them, which will
+// be used to determine whether to emit runtime metadata for the kernel
+// enqueue related hidden kernel arguments.
+//
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/User.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "amdgpu-lower-enqueued-block"
 
 using namespace llvm;
 

 
 INITIALIZE_PASS(AMDGPUOpenCLEnqueuedBlockLowering, DEBUG_TYPE,
                 "Lower OpenCL enqueued blocks", false, false)
 
 ModulePass* llvm::createAMDGPUOpenCLEnqueuedBlockLoweringPass() {
   return new AMDGPUOpenCLEnqueuedBlockLowering();
 }
 
+/// Collect direct or indrect callers of \p F and save them
+/// to \p Callers.
+static void collectCallers(Function *F, DenseSet<Function *> &Callers) {
+  for (auto U : F->users()) {
+    if (auto *CI = dyn_cast<CallInst>(&*U)) {
+      auto *Caller = CI->getParent()->getParent();
+      if (Callers.count(Caller))
+        continue;
+      Callers.insert(Caller);
+      collectCallers(Caller, Callers);
+    }
+  }
+}
+
 bool AMDGPUOpenCLEnqueuedBlockLowering::runOnModule(Module &M) {
+  DenseSet<Function *> Callers;
   auto &C = M.getContext();
   auto AS = AMDGPU::getAMDGPUAS(M);
   bool Changed = false;
   for (auto &F : M.functions()) {
     if (F.hasFnAttribute("enqueued-block")) {
       if (!F.hasOneUse() || !F.user_begin()->hasOneUse() ||
           !isa<ConstantExpr>(*F.user_begin()) ||
           !isa<ConstantExpr>(*F.user_begin()->user_begin())) {
         continue;
       }
       auto *BitCast = cast<ConstantExpr>(*F.user_begin());
       auto *AddrCast = cast<ConstantExpr>(*BitCast->user_begin());
       auto RuntimeHandle = (F.getName() + "_runtime_handle").str();
       auto *GV = new GlobalVariable(
           M, Type::getInt8Ty(C)->getPointerTo(AS.GLOBAL_ADDRESS),
           /*IsConstant=*/true, GlobalValue::ExternalLinkage,
           /*Initializer=*/nullptr, RuntimeHandle, /*InsertBefore=*/nullptr,
           GlobalValue::NotThreadLocal, AS.GLOBAL_ADDRESS,
           /*IsExternallyInitialized=*/true);
       DEBUG(dbgs() << "runtime handle created: " << *GV << '\n');
       auto *NewPtr = ConstantExpr::getPointerCast(GV, AddrCast->getType());
       AddrCast->replaceAllUsesWith(NewPtr);
       F.addFnAttr("runtime-handle", RuntimeHandle);
       F.setLinkage(GlobalValue::ExternalLinkage);
+
+      // Collect direct or indirect callers of enqueue_kernel.
+      for (auto U : NewPtr->users()) {
+        if (auto *I = dyn_cast<Instruction>(&*U)) {
+          auto *F = I->getParent()->getParent();
+          Callers.insert(F);
+          collectCallers(F, Callers);
+        }
+      }
       Changed = true;
     }
   }
+
+  for (auto F : Callers) {
+    if (F->getCallingConv() != CallingConv::AMDGPU_KERNEL)
+      continue;
+    F->addFnAttr("calls-enqueue-kernel");
+  }
   return Changed;
 }

lib/Target/AMDGPU/MCTargetDesc/AMDGPUHSAMetadataStreamer.cpp

void MetadataStreamer::emitKernelArgs(const Function &Func) {
 
   // TODO: What about other languages?
   if (!Func.getParent()->getNamedMetadata("opencl.ocl.version"))
     return;
 
   auto &DL = Func.getParent()->getDataLayout();
   auto Int64Ty = Type::getInt64Ty(Func.getContext());
 
   emitKernelArg(DL, Int64Ty, ValueKind::HiddenGlobalOffsetX);
   emitKernelArg(DL, Int64Ty, ValueKind::HiddenGlobalOffsetY);
   emitKernelArg(DL, Int64Ty, ValueKind::HiddenGlobalOffsetZ);
 
-  if (!Func.getParent()->getNamedMetadata("llvm.printf.fmts"))
-    return;
-
   auto Int8PtrTy = Type::getInt8PtrTy(Func.getContext(),
                                       AMDGPUASI.GLOBAL_ADDRESS);
+  if (Func.getParent()->getNamedMetadata("llvm.printf.fmts"))
     emitKernelArg(DL, Int8PtrTy, ValueKind::HiddenPrintfBuffer);
+  if (Func.hasFnAttribute("calls-enqueue-kernel")) {
+    emitKernelArg(DL, Int8PtrTy, ValueKind::HiddenDefaultQueue);
+    emitKernelArg(DL, Int8PtrTy, ValueKind::HiddenCompletionAction);
+  }
 }
 
 void MetadataStreamer::emitKernelArg(const Argument &Arg) {
   auto Func = Arg.getParent();
   auto ArgNo = Arg.getArgNo();
   const MDNode *Node;
 
   StringRef Name;

test/CodeGen/AMDGPU/enqueue-kernel.ll

 ; RUN: opt -amdgpu-lower-enqueued-block -S < %s | FileCheck %s
 
 ; CHECK: @__test_block_invoke_kernel_runtime_handle = external addrspace(1) externally_initialized constant i8 addrspace(1)*
 ; CHECK: @__test_block_invoke_2_kernel_runtime_handle = external addrspace(1) externally_initialized constant i8 addrspace(1)*
 
 target datalayout = "e-p:32:32-p1:64:64-p2:64:64-p3:32:32-p4:64:64-p5:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64"
 target triple = "amdgcn-amdhsa-amd-opencl"
 
 %struct.ndrange_t = type { i32 }
 %opencl.queue_t = type opaque
 
-define amdgpu_kernel void @test(i8 addrspace(1)* %a, i8 %b, i64 addrspace(1)* %c, i64 %d) local_unnamed_addr
+; CHECK: define amdgpu_kernel void @non_caller(i8 addrspace(1)* %a, i8 %b, i64 addrspace(1)* %c, i64 %d) local_unnamed_addr !kernel_arg_addr_space
+define amdgpu_kernel void @non_caller(i8 addrspace(1)* %a, i8 %b, i64 addrspace(1)* %c, i64 %d) local_unnamed_addr
+  !kernel_arg_addr_space !3 !kernel_arg_access_qual !4 !kernel_arg_type !5 !kernel_arg_base_type !5 !kernel_arg_type_qual !6 {
+  ret void
+}
+
+; CHECK: define amdgpu_kernel void @caller_indirect(i8 addrspace(1)* %a, i8 %b, i64 addrspace(1)* %c, i64 %d) local_unnamed_addr #[[AT_CALLER:[0-9]+]]
+define amdgpu_kernel void @caller_indirect(i8 addrspace(1)* %a, i8 %b, i64 addrspace(1)* %c, i64 %d) local_unnamed_addr
+  !kernel_arg_addr_space !3 !kernel_arg_access_qual !4 !kernel_arg_type !5 !kernel_arg_base_type !5 !kernel_arg_type_qual !6 {
+  call void @caller(i8 addrspace(1)* %a, i8 %b, i64 addrspace(1)* %c, i64 %d)
+  ret void
+}
+
+; CHECK: define amdgpu_kernel void @caller(i8 addrspace(1)* %a, i8 %b, i64 addrspace(1)* %c, i64 %d) local_unnamed_addr #[[AT_CALLER]]
+define amdgpu_kernel void @caller(i8 addrspace(1)* %a, i8 %b, i64 addrspace(1)* %c, i64 %d) local_unnamed_addr
   !kernel_arg_addr_space !3 !kernel_arg_access_qual !4 !kernel_arg_type !5 !kernel_arg_base_type !5 !kernel_arg_type_qual !6 {
 entry:
   %block = alloca <{ i32, i32, i8 addrspace(4)*, i8 addrspace(1)*, i8 }>, align 8
   %tmp = alloca %struct.ndrange_t, align 4
   %block2 = alloca <{ i32, i32, i8 addrspace(4)*, i8 addrspace(1)*, i64 addrspace(1)*, i64, i8 }>, align 8
   %tmp3 = alloca %struct.ndrange_t, align 4
   %block.size = getelementptr inbounds <{ i32, i32, i8 addrspace(4)*, i8 addrspace(1)*, i8 }>, <{ i32, i32, i8 addrspace(4)*, i8 addrspace(1)*, i8 }>* %block, i32 0, i32 0
   store i32 25, i32* %block.size, align 8
entry:
   %.fca.4.extract = extractvalue <{ i32, i32, i8 addrspace(4)*, i8 addrspace(1)*, i64 addrspace(1)*, i64, i8 }> %arg, 4
   %.fca.5.extract = extractvalue <{ i32, i32, i8 addrspace(4)*, i8 addrspace(1)*, i64 addrspace(1)*, i64, i8 }> %arg, 5
   %.fca.6.extract = extractvalue <{ i32, i32, i8 addrspace(4)*, i8 addrspace(1)*, i64 addrspace(1)*, i64, i8 }> %arg, 6
   store i8 %.fca.6.extract, i8 addrspace(1)* %.fca.3.extract, align 1
   store i64 %.fca.5.extract, i64 addrspace(1)* %.fca.4.extract, align 8
   ret void
 }
 
+; CHECK: attributes #[[AT_CALLER]] = { "calls-enqueue-kernel" }
 ; CHECK: attributes #[[AT1]] = {{.*}}"runtime-handle"="__test_block_invoke_kernel_runtime_handle"
 ; CHECK: attributes #[[AT2]] = {{.*}}"runtime-handle"="__test_block_invoke_2_kernel_runtime_handle"
 
 attributes #0 = { "enqueued-block" }
 
 !3 = !{i32 1, i32 0, i32 1, i32 0}
 !4 = !{!"none", !"none", !"none", !"none"}
 !5 = !{!"char*", !"char", !"long*", !"long"}
 !6 = !{!"", !"", !"", !""}
 !14 = !{i32 0}
 !15 = !{!"none"}
 !16 = !{!"__block_literal"}
 !17 = !{!""}

test/CodeGen/AMDGPU/hsa-metadata-from-llvm-ir-full.ll

 ; CHECK-NEXT:       Align:         8
 ; CHECK-NEXT:       ValueKind:     HiddenGlobalOffsetZ
 ; CHECK-NEXT:       ValueType:     I64
 ; CHECK-NEXT:     - Size:          8
 ; CHECK-NEXT:       Align:         8
 ; CHECK-NEXT:       ValueKind:     HiddenPrintfBuffer
 ; CHECK-NEXT:       ValueType:     I8
 ; CHECK-NEXT:       AddrSpaceQual: Global
+; CHECK-NOT:        ValueKind:     HiddenDefaultQueue
+; CHECK-NOT:        ValueKind:     HiddenCompletionAction
 define amdgpu_kernel void @test_char(i8 %a)
     !kernel_arg_addr_space !1 !kernel_arg_access_qual !2 !kernel_arg_type !9
     !kernel_arg_base_type !9 !kernel_arg_type_qual !4 {
   ret void
 }
 
 ; CHECK:      - Name:            test_ushort2
 ; CHECK-NEXT:   SymbolName:      'test_ushort2@kd'

 ; CHECK-NEXT:       AddrSpaceQual: Global
 define amdgpu_kernel void @__test_block_invoke_kernel(
     <{ i32, i32, i8 addrspace(4)*, i8 addrspace(1)*, i8 }> %arg) #0
     !kernel_arg_addr_space !1 !kernel_arg_access_qual !2 !kernel_arg_type !110
     !kernel_arg_base_type !110 !kernel_arg_type_qual !4 {
   ret void
 }
 
+; CHECK:      - Name:            test_enqueue_kernel_caller
+; CHECK-NEXT:   SymbolName:      'test_enqueue_kernel_caller@kd'
+; CHECK-NEXT:   Language:        OpenCL C
+; CHECK-NEXT:   LanguageVersion: [ 2, 0 ]
+; CHECK-NEXT:   Args:
+; CHECK-NEXT:     - TypeName:      char
+; CHECK-NEXT:       Size:          1
+; CHECK-NEXT:       Align:         1
+; CHECK-NEXT:       ValueKind:     ByValue
+; CHECK-NEXT:       ValueType:     I8
+; CHECK-NEXT:       AccQual:       Default
+; CHECK-NEXT:     - Size:          8
+; CHECK-NEXT:       Align:         8
+; CHECK-NEXT:       ValueKind:     HiddenGlobalOffsetX
+; CHECK-NEXT:       ValueType:     I64
+; CHECK-NEXT:     - Size:          8
+; CHECK-NEXT:       Align:         8
+; CHECK-NEXT:       ValueKind:     HiddenGlobalOffsetY
+; CHECK-NEXT:       ValueType:     I64
+; CHECK-NEXT:     - Size:          8
+; CHECK-NEXT:       Align:         8
+; CHECK-NEXT:       ValueKind:     HiddenGlobalOffsetZ
+; CHECK-NEXT:       ValueType:     I64
+; CHECK-NEXT:     - Size:          8
+; CHECK-NEXT:       Align:         8
+; CHECK-NEXT:       ValueKind:     HiddenPrintfBuffer
+; CHECK-NEXT:       ValueType:     I8
+; CHECK-NEXT:       AddrSpaceQual: Global
+; CHECK-NEXT:     - Size:          8
+; CHECK-NEXT:       Align:         8
+; CHECK-NEXT:       ValueKind:     HiddenDefaultQueue
+; CHECK-NEXT:       ValueType:     I8
+; CHECK-NEXT:       AddrSpaceQual: Global
+; CHECK-NEXT:     - Size:          8
+; CHECK-NEXT:       Align:         8
+; CHECK-NEXT:       ValueKind:     HiddenCompletionAction
+; CHECK-NEXT:       ValueType:     I8
+; CHECK-NEXT:       AddrSpaceQual: Global
+define amdgpu_kernel void @test_enqueue_kernel_caller(i8 %a) #1
+    !kernel_arg_addr_space !1 !kernel_arg_access_qual !2 !kernel_arg_type !9
+    !kernel_arg_base_type !9 !kernel_arg_type_qual !4 {
+  ret void
+}
+
 attributes #0 = { "runtime-handle"="__test_block_invoke_kernel_runtime_handle" }
+attributes #1 = { "calls-enqueue-kernel" }
 
 !llvm.printf.fmts = !{!100, !101}
 
 !1 = !{i32 0}
 !2 = !{!"none"}
 !3 = !{!"int"}
 !4 = !{!""}
 !5 = !{i32 undef, i32 1}