[SYCL] Separate host instantiation from HostKernel (#18534)

The HostKernel class is currently used for two different things:
1) To extend the lifetime of the kernel function object/lambda; and
2) To force the kernel to be instantiated on the host.

Requiring a single solution to solve both problems prevents
optimization. Providing separate solutions will enable us to use the
fastest solution for each problem.

---

A few notes that might help with review:

- I would like us to eventually reach a point where we can remove the
`HostKernel` class entirely and avoid extending the lifetime of the
kernel function object. We can't do that if it's being relied upon for
debugging.

- I had to change the `runKernelWithoutArg` signature to remove SFINAE
because MSVC complained about creating a pointer to an overloaded
function. I would have expected that the compiler could figure out the
function pointer _after_ SFINAE, but it failed.

- Removing the old mechanism of instantiating on the host is an ABI
break because it removes a function from the vtable.

- With the new approach, we can avoid instantiating (and compiling!) the
kernel on the host under `NDEBUG`. I've verified that `-DNDEBUG
-fpreview-breaking-changes` removes the kernel symbols from the binary.
`-DNDEBUG` alone isn't currently sufficient, because the old
instantiation mechanism is still triggered. I wasn't sure how much of
the old implementation to leave in tact.

---------

Signed-off-by: John Pennycook <[email protected]>

Author: Pennycook

sycl/include/sycl/detail/cg_types.hpp

@@ -124,32 +124,25 @@ struct KernelLambdaHasKernelHandlerArgT {
 
 // Helpers for running kernel lambda on the host device
 
-template <typename KernelType>
-typename std::enable_if_t<KernelLambdaHasKernelHandlerArgT<KernelType>::value>
-runKernelWithoutArg(KernelType KernelName) {
-  kernel_handler KH;
-  KernelName(KH);
-}
-
-template <typename KernelType>
-typename std::enable_if_t<!KernelLambdaHasKernelHandlerArgT<KernelType>::value>
-runKernelWithoutArg(KernelType KernelName) {
-  KernelName();
-}
-
-template <typename ArgType, typename KernelType>
-typename std::enable_if_t<
-    KernelLambdaHasKernelHandlerArgT<KernelType, ArgType>::value>
-runKernelWithArg(KernelType KernelName, ArgType Arg) {
-  kernel_handler KH;
-  KernelName(Arg, KH);
+template <typename KernelType, bool HasKernelHandlerArg>
+void runKernelWithoutArg(KernelType KernelName,
+                         const std::bool_constant<HasKernelHandlerArg> &) {
+  if constexpr (HasKernelHandlerArg) {
+    kernel_handler KH;
+    KernelName(KH);
+  } else {
+    KernelName();
+  }
 }
-
-template <typename ArgType, typename KernelType>
-typename std::enable_if_t<
-    !KernelLambdaHasKernelHandlerArgT<KernelType, ArgType>::value>
-runKernelWithArg(KernelType KernelName, ArgType Arg) {
-  KernelName(Arg);
+template <typename ArgType, typename KernelType, bool HasKernelHandlerArg>
+void runKernelWithArg(KernelType KernelName, ArgType Arg,
+                      const std::bool_constant<HasKernelHandlerArg> &) {
+  if constexpr (HasKernelHandlerArg) {
+    kernel_handler KH;
+    KernelName(Arg, KH);
+  } else {
+    KernelName(Arg);
+  }
 }
 
 // The pure virtual class aimed to store lambda/functors of any type.
@@ -159,8 +152,10 @@ class HostKernelBase {
   // Used to extract captured variables.
   virtual char *getPtr() = 0;
   virtual ~HostKernelBase() = default;
+#ifndef __INTEL_PREVIEW_BREAKING_CHANGES
   // NOTE: InstatiateKernelOnHost() should not be called.
   virtual void InstantiateKernelOnHost() = 0;
+#endif
 };
 
 // Class which stores specific lambda object.
@@ -176,17 +171,21 @@ class HostKernel : public HostKernelBase {
 
   ~HostKernel() = default;
 
+#ifndef __INTEL_PREVIEW_BREAKING_CHANGES
   // This function is needed for host-side compilation to keep kernels
   // instantitated. This is important for debuggers to be able to associate
   // kernel code instructions with source code lines.
   // NOTE: InstatiateKernelOnHost() should not be called.
   void InstantiateKernelOnHost() override {
     using IDBuilder = sycl::detail::Builder;
+    constexpr bool HasKernelHandlerArg =
+        KernelLambdaHasKernelHandlerArgT<KernelType, KernelArgType>::value;
     if constexpr (std::is_same_v<KernelArgType, void>) {
-      runKernelWithoutArg(MKernel);
+      runKernelWithoutArg(MKernel, std::bool_constant<HasKernelHandlerArg>());
     } else if constexpr (std::is_same_v<KernelArgType, sycl::id<Dims>>) {
       sycl::id ID = InitializedVal<Dims, id>::template get<0>();
-      runKernelWithArg<const KernelArgType &>(MKernel, ID);
+      runKernelWithArg<const KernelArgType &>(
+          MKernel, ID, std::bool_constant<HasKernelHandlerArg>());
     } else if constexpr (std::is_same_v<KernelArgType, item<Dims, true>> ||
                          std::is_same_v<KernelArgType, item<Dims, false>>) {
       constexpr bool HasOffset =
@@ -195,13 +194,15 @@ class HostKernel : public HostKernelBase {
         KernelArgType Item = IDBuilder::createItem<Dims, HasOffset>(
             InitializedVal<Dims, range>::template get<1>(),
             InitializedVal<Dims, id>::template get<0>());
-        runKernelWithArg<KernelArgType>(MKernel, Item);
+        runKernelWithArg<KernelArgType>(
+            MKernel, Item, std::bool_constant<HasKernelHandlerArg>());
       } else {
         KernelArgType Item = IDBuilder::createItem<Dims, HasOffset>(
             InitializedVal<Dims, range>::template get<1>(),
             InitializedVal<Dims, id>::template get<0>(),
             InitializedVal<Dims, id>::template get<0>());
-        runKernelWithArg<KernelArgType>(MKernel, Item);
+        runKernelWithArg<KernelArgType>(
+            MKernel, Item, std::bool_constant<HasKernelHandlerArg>());
       }
     } else if constexpr (std::is_same_v<KernelArgType, nd_item<Dims>>) {
       sycl::range<Dims> Range = InitializedVal<Dims, range>::template get<1>();
@@ -214,22 +215,44 @@ class HostKernel : public HostKernelBase {
           IDBuilder::createItem<Dims, false>(Range, ID);
       KernelArgType NDItem =
           IDBuilder::createNDItem<Dims>(GlobalItem, LocalItem, Group);
-      runKernelWithArg<const KernelArgType>(MKernel, NDItem);
+      runKernelWithArg<const KernelArgType>(
+          MKernel, NDItem, std::bool_constant<HasKernelHandlerArg>());
     } else if constexpr (std::is_same_v<KernelArgType, sycl::group<Dims>>) {
       sycl::range<Dims> Range = InitializedVal<Dims, range>::template get<1>();
       sycl::id<Dims> ID = InitializedVal<Dims, id>::template get<0>();
       KernelArgType Group =
           IDBuilder::createGroup<Dims>(Range, Range, Range, ID);
-      runKernelWithArg<KernelArgType>(MKernel, Group);
+      runKernelWithArg<KernelArgType>(
+          MKernel, Group, std::bool_constant<HasKernelHandlerArg>());
     } else {
       // Assume that anything else can be default-constructed. If not, this
       // should fail to compile and the implementor should implement a generic
       // case for the new argument type.
-      runKernelWithArg<KernelArgType>(MKernel, KernelArgType{});
+      runKernelWithArg<KernelArgType>(
+          MKernel, KernelArgType{}, std::bool_constant<HasKernelHandlerArg>());
     }
   }
+#endif
 };
 
+// This function is needed for host-side compilation to keep kernels
+// instantitated. This is important for debuggers to be able to associate
+// kernel code instructions with source code lines.
+template <class KernelType, class KernelArgType, int Dims>
+constexpr void *GetInstantiateKernelOnHostPtr() {
+  if constexpr (std::is_same_v<KernelArgType, void>) {
+    constexpr bool HasKernelHandlerArg =
+        KernelLambdaHasKernelHandlerArgT<KernelType>::value;
+    return reinterpret_cast<void *>(
+        &runKernelWithoutArg<KernelType, HasKernelHandlerArg>);
+  } else {
+    constexpr bool HasKernelHandlerArg =
+        KernelLambdaHasKernelHandlerArgT<KernelType, KernelArgType>::value;
+    return reinterpret_cast<void *>(
+        &runKernelWithArg<KernelArgType, KernelType, HasKernelHandlerArg>);
+  }
+}
+
 } // namespace detail
 } // namespace _V1
 } // namespace sycl

sycl/include/sycl/handler.hpp

@@ -821,6 +821,14 @@ class __SYCL_EXPORT handler {
     MHostKernel.reset(new detail::HostKernel<KernelType, LambdaArgType, Dims>(
         std::forward<KernelTypeUniversalRef>(KernelFunc)));
 
+    // Instantiating the kernel on the host improves debugging.
+    // Passing this pointer to another translation unit prevents optimization.
+#ifndef NDEBUG
+    instantiateKernelOnHost(
+        detail::GetInstantiateKernelOnHostPtr<KernelType, LambdaArgType,
+                                              Dims>());
+#endif
+
     constexpr bool KernelHasName =
         detail::getKernelName<KernelName>() != nullptr &&
         detail::getKernelName<KernelName>()[0] != '\0';
@@ -3834,6 +3842,8 @@ class __SYCL_EXPORT handler {
                      detail::kernel_param_desc_t (*KernelParamDescGetter)(int),
                      bool KernelIsESIMD, bool KernelHasSpecialCaptures);
 
+  void instantiateKernelOnHost(void *InstantiateKernelOnHostPtr);
+
   friend class detail::HandlerAccess;
 
 #ifdef __INTEL_PREVIEW_BREAKING_CHANGES

sycl/source/handler.cpp

@@ -2386,6 +2386,12 @@ void handler::setKernelInfo(
   impl->MKernelHasSpecialCaptures = KernelHasSpecialCaptures;
 }
 
+void handler::instantiateKernelOnHost(void *InstantiateKernelOnHostPtr) {
+  // Passing the pointer to the runtime is enough to prevent optimization.
+  // We don't need to use the pointer for anything.
+  (void)InstantiateKernelOnHostPtr;
+}
+
 void handler::saveCodeLoc(detail::code_location CodeLoc, bool IsTopCodeLoc) {
   MCodeLoc = CodeLoc;
   impl->MIsTopCodeLoc = IsTopCodeLoc;

sycl/test/abi/sycl_symbols_linux.dump

@@ -3575,6 +3575,7 @@ _ZN4sycl3_V17handler22setHandlerKernelBundleENS0_6kernelE
 _ZN4sycl3_V17handler22setHandlerKernelBundleERKSt10shared_ptrINS0_6detail18kernel_bundle_implEE
 _ZN4sycl3_V17handler22setKernelClusterLaunchENS0_5rangeILi3EEEi
 _ZN4sycl3_V17handler22setKernelIsCooperativeEb
+_ZN4sycl3_V17handler23instantiateKernelOnHostEPv
 _ZN4sycl3_V17handler24GetRangeRoundingSettingsERmS2_S2_
 _ZN4sycl3_V17handler24ext_intel_read_host_pipeENS0_6detail11string_viewEPvmb
 _ZN4sycl3_V17handler24ext_oneapi_memcpy2d_implEPvmPKvmmm

sycl/test/abi/sycl_symbols_windows.dump

@@ -4230,6 +4230,7 @@
 ?has_kernel_bundle_impl@detail@_V1@sycl@@YA_NAEBVcontext@23@AEBV?$vector@Vdevice@_V1@sycl@@V?$allocator@Vdevice@_V1@sycl@@@std@@@std@@W4bundle_state@23@@Z
 ?has_specialization_constant_impl@kernel_bundle_plain@detail@_V1@sycl@@IEBA_NPEBD@Z
 ?increase_threshold_to@memory_pool@experimental@oneapi@ext@_V1@sycl@@QEAAX_K@Z
+?instantiateKernelOnHost@handler@_V1@sycl@@AEAAXPEAX@Z
 ?internalProfilingTagImpl@handler@_V1@sycl@@AEAAXXZ
 ?isBackendSupportedFillSize@handler@_V1@sycl@@CA_N_K@Z
 ?isConstOrGlobal@handler@_V1@sycl@@CA_NW4target@access@23@@Z