[SYCL][NATIVECPU] Emit subhandler as LLVM IR (#10282)

This PR allows the emit the Native CPU subhandler as LLVM-IR, instead of
emitting its declaration as C++ code in the Native CPU integration
header. This is a first step towards removing the Native CPU integration
header, using clang-offload-wrapper and also supporting multiple SYCL
targets alongside SYCL Native CPU, which we will implement in following
PRs.

---------

Co-authored-by: Uwe Dolinsky <[email protected]>
Co-authored-by: Alexey Sachkov <[email protected]>
Co-authored-by: Steffen Larsen <[email protected]>

Author: 4 people

clang/lib/CodeGen/CodeGenModule.cpp

@@ -2115,10 +2115,6 @@ void CodeGenModule::GenKernelArgMetadata(llvm::Function *Fn,
       Fn->setMetadata("kernel_arg_exclusive_ptr",
                       llvm::MDNode::get(VMContext, argSYCLAccessorPtrs));
     }
-    if (LangOpts.SYCLIsNativeCPU) {
-      Fn->setMetadata("kernel_arg_type",
-                      llvm::MDNode::get(VMContext, argTypeNames));
-    }
   } else {
     if (getLangOpts().OpenCL || getLangOpts().SYCLIsDevice) {
       Fn->setMetadata("kernel_arg_addr_space",

clang/test/CodeGenSYCL/native_cpu_basic.cpp

@@ -1,6 +1,5 @@
 // This test checks for some basic Front End features for Native CPU:
 // * Kernel name mangling
-// * kernel_arg_type metadata node
 // * is-native-cpu module flag
 // RUN: %clang_cc1 -fsycl-is-device -S -emit-llvm -internal-isystem %S/Inputs -fsycl-is-native-cpu -o %t.ll %s 
 // RUN: FileCheck -input-file=%t.ll %s 
@@ -50,11 +49,9 @@ void gen() {
   test<float>(q);
 }
 
-// Check name mangling and kernel_arg_type metadata
-// CHECK-DAG: @_ZTS6init_aIiE_NativeCPUKernel({{.*}}){{.*}}!kernel_arg_type ![[TYPE1:[0-9]*]]
-// CHECK-DAG: @_ZTS6init_aIfE_NativeCPUKernel({{.*}}){{.*}}!kernel_arg_type ![[TYPE3:[0-9]*]]
-// CHECK-DAG: ![[TYPE1]] = !{!"int*", !"sycl::range<1>", !"sycl::range<1>", !"sycl::id<1>", !"int"}
-// CHECK-DAG: ![[TYPE3]] = !{!"float*", !"sycl::range<1>", !"sycl::range<1>", !"sycl::id<1>", !"float"}
+// Check name mangling 
+// CHECK-DAG: @_ZTS6init_aIiE_NativeCPUKernel_NativeCPUKernel({{.*}})
+// CHECK-DAG: @_ZTS6init_aIfE_NativeCPUKernel_NativeCPUKernel({{.*}})
 
 // Check Native CPU module flag
 // CHECK-DAG: !{{[0-9]*}} = !{i32 1, !"is-native-cpu", i32 1}

llvm/lib/SYCLLowerIR/EmitSYCLNativeCPUHeader.cpp

@@ -33,133 +33,12 @@
 using namespace llvm;
 
 namespace {
-SmallVector<bool> getArgMask(const Function *F) {
-  SmallVector<bool> Res;
-  auto *UsedNode = F->getMetadata("sycl_kernel_omit_args");
-  if (!UsedNode) {
-    // the metadata node is not available if -fenable-sycl-dae
-    // was not set; set everything to true in the mask.
-    for (unsigned I = 0; I < F->getFunctionType()->getNumParams(); I++) {
-      Res.push_back(true);
-    }
-    return Res;
-  }
-  auto NumOperands = UsedNode->getNumOperands();
-  for (unsigned I = 0; I < NumOperands; I++) {
-    auto &Op = UsedNode->getOperand(I);
-    if (auto *CAM = dyn_cast<ConstantAsMetadata>(Op.get())) {
-      if (auto *Const = dyn_cast<ConstantInt>(CAM->getValue())) {
-        auto Val = Const->getValue();
-        Res.push_back(!Val.getBoolValue());
-      } else {
-        report_fatal_error("Unable to retrieve constant int from "
-                           "sycl_kernel_omit_args metadata node");
-      }
-    } else {
-      report_fatal_error(
-          "Error while processing sycl_kernel_omit_args metadata node");
-    }
-  }
-  return Res;
-}
 
-SmallVector<StringRef> getArgTypeNames(const Function *F) {
-  SmallVector<StringRef> Res;
-  auto *TNNode = F->getMetadata("kernel_arg_type");
-  assert(TNNode &&
-         "kernel_arg_type metadata node is required for sycl native CPU");
-  auto NumOperands = TNNode->getNumOperands();
-  for (unsigned I = 0; I < NumOperands; I++) {
-    auto &Op = TNNode->getOperand(I);
-    auto *MDS = dyn_cast<MDString>(Op.get());
-    if (!MDS)
-      report_fatal_error("error while processing kernel_arg_types metadata");
-    Res.push_back(MDS->getString());
-  }
-  return Res;
-}
-
-void emitKernelDecl(const Function *F, const SmallVector<bool> &ArgMask,
-                    const SmallVector<StringRef> &ArgTypeNames,
-                    raw_ostream &O) {
-  auto EmitArgDecl = [&](const Argument *Arg, unsigned Index) {
-    Type *ArgTy = Arg->getType();
-    if (isa<PointerType>(ArgTy))
-      return "void *";
-    return ArgTypeNames[Index].data();
-  };
-
-  auto NumParams = F->getFunctionType()->getNumParams();
-  O << "extern \"C\" void " << F->getName() << "(";
-
-  unsigned I = 0, UsedI = 0;
-  for (; I + 1 < ArgMask.size() && UsedI + 1 < NumParams; I++) {
-    if (!ArgMask[I])
-      continue;
-    O << EmitArgDecl(F->getArg(UsedI), I) << ", ";
-    UsedI++;
-  }
-
-  // parameters may have been removed.
-  bool NoUsedArgs = true;
-  for (auto &Entry : ArgMask) {
-    NoUsedArgs &= !Entry;
-  }
-  if (NoUsedArgs) {
-    O << ");\n";
-    return;
-  }
-  // find the index of the last used arg
-  while (!ArgMask[I] && I + 1 < ArgMask.size())
-    I++;
-  O << EmitArgDecl(F->getArg(UsedI), I) << ", __nativecpu_state *);\n";
-}
-
-void emitSubKernelHandler(const Function *F, const SmallVector<bool> &ArgMask,
-                          const SmallVector<StringRef> &ArgTypeNames,
-                          raw_ostream &O) {
-  SmallVector<unsigned> UsedArgIdx;
-  auto EmitParamCast = [&](Argument *Arg, unsigned Index) {
-    std::string Res;
-    llvm::raw_string_ostream OS(Res);
-    UsedArgIdx.push_back(Index);
-    if (isa<PointerType>(Arg->getType())) {
-      OS << "  void* arg" << Index << " = ";
-      OS << "MArgs[" << Index << "].getPtr();\n";
-      return OS.str();
-    }
-    auto TN = ArgTypeNames[Index].str();
-    OS << "  " << TN << " arg" << Index << " = ";
-    OS << "*(" << TN << "*)"
-       << "MArgs[" << Index << "].getPtr();\n";
-    return OS.str();
-  };
-
-  O << "\ninline static void " << F->getName() << "subhandler(";
+void emitSubKernelHandler(const Function *F, raw_ostream &O) {
+  O << "\nextern \"C\" void " << F->getName() << "subhandler(";
   O << "const sycl::detail::NativeCPUArgDesc *MArgs, "
-       "__nativecpu_state *state) {\n";
-  // Retrieve only the args that are used
-  for (unsigned I = 0, UsedI = 0;
-       I < ArgMask.size() && UsedI < F->getFunctionType()->getNumParams();
-       I++) {
-    if (ArgMask[I]) {
-      O << EmitParamCast(F->getArg(UsedI), I);
-      UsedI++;
-    }
-  }
-  // Emit the actual kernel call
-  O << "  " << F->getName() << "(";
-  if (UsedArgIdx.size() == 0) {
-    O << ");\n";
-  } else {
-    for (unsigned I = 0; I < UsedArgIdx.size() - 1; I++) {
-      O << "arg" << UsedArgIdx[I] << ", ";
-    }
-    if (UsedArgIdx.size() >= 1)
-      O << "arg" << UsedArgIdx.back();
-    O << ", state);\n";
-  }
-  O << "};\n\n";
+       "__nativecpu_state *state);\n";
+  return;
 }
 
 // Todo: maybe we could use clang-offload-wrapper for this,
@@ -254,10 +133,7 @@ PreservedAnalyses EmitSYCLNativeCPUHeaderPass::run(Module &M,
   O << "extern \"C\" void __sycl_register_lib(pi_device_binaries desc);\n";
 
   for (auto *F : Kernels) {
-    auto ArgMask = getArgMask(F);
-    auto ArgTypeNames = getArgTypeNames(F);
-    emitKernelDecl(F, ArgMask, ArgTypeNames, O);
-    emitSubKernelHandler(F, ArgMask, ArgTypeNames, O);
+    emitSubKernelHandler(F, O);
     emitSYCLRegisterLib(F, O);
   }
 

llvm/lib/SYCLLowerIR/PrepareSYCLNativeCPU.cpp

@@ -23,6 +23,7 @@
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/Value.h"
 #include "llvm/InitializePasses.h"
@@ -62,7 +63,88 @@ void fixCallingConv(Function *F) {
     F->setLinkage(GlobalValue::LinkageTypes::WeakAnyLinkage);
 }
 
-// Clone the function and returns a new function with a new argument on type T
+// returns the indexes of the used arguments
+SmallVector<unsigned> getUsedIndexes(const Function *F) {
+  SmallVector<unsigned> res;
+  auto UsedNode = F->getMetadata("sycl_kernel_omit_args");
+  if (!UsedNode) {
+    // the metadata node is not available if -fenable-sycl-dae
+    // was not set; set everything to true
+    // Exclude one arg because we already added the state ptr
+    for (unsigned I = 0; I + 1 < F->getFunctionType()->getNumParams(); I++) {
+      res.push_back(I);
+    }
+    return res;
+  }
+  auto NumOperands = UsedNode->getNumOperands();
+  for (unsigned I = 0; I < NumOperands; I++) {
+    auto &Op = UsedNode->getOperand(I);
+    if (auto CAM = dyn_cast<ConstantAsMetadata>(Op.get())) {
+      if (auto Const = dyn_cast<ConstantInt>(CAM->getValue())) {
+        auto Val = Const->getValue();
+        if (!Val.getBoolValue()) {
+          res.push_back(I);
+        }
+      } else {
+        report_fatal_error("Unable to retrieve constant int from "
+                           "sycl_kernel_omit_args metadata node");
+      }
+    } else {
+      report_fatal_error(
+          "Error while processing sycl_kernel_omit_args metadata node");
+    }
+  }
+  return res;
+}
+
+void emitSubkernelForKernel(Function *F, Type *NativeCPUArgDescType,
+                            Type *StatePtrType) {
+  LLVMContext &Ctx = F->getContext();
+  Type *NativeCPUArgDescPtrType = PointerType::getUnqual(NativeCPUArgDescType);
+
+  // Create function signature
+  const std::string SubHandlerName = F->getName().str() + "subhandler";
+  FunctionType *FTy = FunctionType::get(
+      Type::getVoidTy(Ctx), {NativeCPUArgDescPtrType, StatePtrType}, false);
+  auto SubhFCallee = F->getParent()->getOrInsertFunction(SubHandlerName, FTy);
+  Function *SubhF = cast<Function>(SubhFCallee.getCallee());
+
+  // Emit function body, unpack kernel args
+  auto UsedIndexes = getUsedIndexes(F);
+  auto *KernelTy = F->getFunctionType();
+  // assert(UsedIndexes.size() + 1 == KernelTy->getNumParams() && "mismatch
+  // between number of params and used args");
+  IRBuilder<> Builder(Ctx);
+  BasicBlock *Block = BasicBlock::Create(Ctx, "entry", SubhF);
+  Builder.SetInsertPoint(Block);
+  unsigned NumArgs = UsedIndexes.size();
+  auto *BaseNativeCPUArg = SubhF->getArg(0);
+  SmallVector<Value *, 5> KernelArgs;
+  for (unsigned I = 0; I < NumArgs; I++) {
+    auto *Arg = F->getArg(I);
+    auto UsedI = UsedIndexes[I];
+    // Load the correct NativeCPUDesc and load the pointer from it
+    auto *Addr = Builder.CreateGEP(NativeCPUArgDescType, BaseNativeCPUArg,
+                                   {Builder.getInt64(UsedI)});
+    auto *Load = Builder.CreateLoad(PointerType::getUnqual(Ctx), Addr);
+    if (Arg->getType()->isPointerTy()) {
+      // If the arg is a pointer, just use it
+      KernelArgs.push_back(Load);
+    } else {
+      // Otherwise, load the scalar value and use that
+      auto *Scalar = Builder.CreateLoad(Arg->getType(), Load);
+      KernelArgs.push_back(Scalar);
+    }
+  }
+
+  // Call the kernel
+  // Add the nativecpu state as arg
+  KernelArgs.push_back(SubhF->getArg(1));
+  Builder.CreateCall(KernelTy, F, KernelArgs);
+  Builder.CreateRetVoid();
+}
+
+// Clones the function and returns a new function with a new argument on type T
 // added as last argument
 Function *cloneFunctionAndAddParam(Function *OldF, Type *T) {
   auto *OldT = OldF->getFunctionType();
@@ -166,6 +248,14 @@ PreservedAnalyses PrepareSYCLNativeCPUPass::run(Module &M,
     ModuleChanged |= true;
   }
 
+  StructType *NativeCPUArgDescType =
+      StructType::create({PointerType::getUnqual(M.getContext())});
+  for (auto &NewK : NewKernels) {
+    emitSubkernelForKernel(NewK, NativeCPUArgDescType, StatePtrType);
+    std::string NewName = NewK->getName().str() + "_NativeCPUKernel";
+    NewK->setName(NewName);
+  }
+
   // Then we iterate over all the supported builtins, find their uses and
   // replace them with calls to our Native CPU functions.
   for (auto &Entry : BuiltinNamesMap) {

sycl/doc/design/SYCLNativeCPU.md

@@ -122,26 +122,35 @@ entry:
   ret void
 }
 ```
-This pass will also set the correct calling convention for the target, and handle calling convention-related function attributes, allowing to call the kernel from the runtime.
+This pass will also set the correct calling convention for the target, and handle calling convention-related function attributes, allowing to call the kernel from the runtime. \\
+Additionally, this pass emits the definition for a `subhandler` function, which unpacks the vector of kernel arguments coming from the SYCL runtime, and forwards only the used arguments to the kernel. For our example the `subhandler` IR is:
+
+```llvm
+define weak void @_Z6Samplesubhandler(ptr %0, ptr %1) #4 {
+entry:
+  %2 = getelementptr %0, ptr %0, i64 0
+  %3 = load ptr, ptr %2, align 8
+  %4 = getelementptr %0, ptr %0, i64 3
+  %5 = load ptr, ptr %4, align 8
+  %6 = getelementptr %0, ptr %0, i64 4
+  %7 = load ptr, ptr %6, align 8
+  %8 = getelementptr %0, ptr %0, i64 7
+  %9 = load ptr, ptr %8, align 8
+  call void @_ZTS10SimpleVaddIiE_NativeCPUKernel(ptr %3, ptr %5, ptr %7, ptr %9, ptr %1)
+  ret void
+}
+```
 
 ## EmitSYCLNativeCPUHeader pass
 
 This pass emits an additional integration header, that will be compiled by the host compiler during the host compilation step. This header is included by the main integration footer and does not need to be managed manually. Its main purpose is to enable the SYCL runtime to register kernels and to call kernels that had unused parameters removed by the optimizer. The header contains, for each kernel:
-* The kernel declaration as a C++ function, all pointer arguments are emitted as `void *`, the scalar arguments maintain their type.
-* A `subhandler` definition, which unpacks the vector of kernel arguments coming from the SYCL runtime, and forwards only the used arguments to the kernel.
+* The subhandler declaration as a C++ function.
 * The definition of `_pi_offload_entry_struct`, `pi_device_binary_struct` and `pi_device_binaries_struct` variables, and a call to `__sycl_register_lib`, which allows to register the kernel to the sycl runtime (the call to `__sycl_register_lib` is performed at program startup via the constructor of a global). The Native CPU integration header is always named `<main-sycl-int-header>.hc`.
 
 The Native CPU integration header for our example is:
 
 ```c++
-extern "C" void _Z6Sample(void *, void *, void *, nativecpu_state *);
-
-inline static void _Z6Samplesubhandler(const sycl::detail::NativeCPUArgDesc *MArgs, nativecpu_state *state) {
-  void* arg0 = MArgs[0].getPtr();
-  void* arg1 = MArgs[1].getPtr();
-  void* arg2 = MArgs[2].getPtr();
-  _Z6Sample(arg0, arg1, arg2, state);
-};
+extern "C" void _Z6Samplesubhandler(const sycl::detail::NativeCPUArgDesc *MArgs, nativecpu_state *state);
 
 static _pi_offload_entry_struct _pi_offload_entry_struct_Z6Sample{(void*)&_Z6Samplesubhandler, const_cast<char*>("_Z6Sample"), 1, 0, 0 };
 static pi_device_binary_struct pi_device_binary_struct_Z6Sample{0, 4, 0, __SYCL_PI_DEVICE_BINARY_TARGET_UNKNOWN, nullptr, nullptr, nullptr, nullptr, (unsigned char*)&_Z6Samplesubhandler, (unsigned char*)&_Z6Samplesubhandler + 1, &_pi_offload_entry_struct_Z6Sample, &_pi_offload_entry_struct_Z6Sample+1, nullptr, nullptr };

sycl/include/sycl/detail/native_cpu.hpp

@@ -18,7 +18,6 @@ namespace detail {
 struct NativeCPUArgDesc {
   void *MPtr;
 
-  void *getPtr() const { return MPtr; }
   NativeCPUArgDesc(void *Ptr) : MPtr(Ptr){};
 };
 

sycl/plugins/native_cpu/pi_native_cpu.cpp

@@ -1385,6 +1385,17 @@ pi_result piextEnqueueCommandBuffer(pi_ext_command_buffer, pi_queue, pi_uint32,
   DIE_NO_IMPLEMENTATION;
 }
 
+pi_result piextEnablePeerAccess(pi_device, pi_device) { DIE_NO_IMPLEMENTATION; }
+
+pi_result piextDisablePeerAccess(pi_device, pi_device) {
+  DIE_NO_IMPLEMENTATION;
+}
+
+pi_result piextPeerAccessGetInfo(pi_device, pi_device, pi_peer_attr, size_t,
+                                 void *, size_t *) {
+  DIE_NO_IMPLEMENTATION;
+}
+
 pi_result piTearDown(void *) {
   // Todo: is it fine as a no-op?
   return PI_SUCCESS;