Defer arg extraction until handler::finalize

Rather than extracting arguments from the lambda when the kernel is enqueued,
store a pointer to the lambda alongside relevant information from the
integration header or compiler builtins.

Storing this information will allow us to defer the extraction of arguments
until we reach handler::finalize(), at which point it may be possible to
set the kernel arguments directly without populating MArgs.

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

Author: Pennycook

sycl/include/sycl/handler.hpp

@@ -749,8 +749,8 @@ class __SYCL_EXPORT handler {
 
   /// Stores lambda to the template-free object
   ///
-  /// Also initializes kernel name, list of arguments and requirements using
-  /// information from the integration header/built-ins.
+  /// Also initializes the kernel name and prepares for arguments to
+  /// be extracted from the lambda in handler::finalize().
   ///
   /// \param KernelFunc is a SYCL kernel function
   /// \param ParamDescs is the vector of kernel parameter descriptors.
@@ -796,11 +796,11 @@ class __SYCL_EXPORT handler {
     if constexpr (KernelHasName) {
       // TODO support ESIMD in no-integration-header case too.
 
-      clearArgs();
-      extractArgsAndReqsFromLambda(MHostKernel->getPtr(),
-                                   &(detail::getKernelParamDesc<KernelName>),
-                                   detail::getKernelNumParams<KernelName>(),
-                                   detail::isKernelESIMD<KernelName>());
+      setKernelInfo((void *)MHostKernel->getPtr(),
+                    detail::getKernelNumParams<KernelName>(),
+                    &(detail::getKernelParamDesc<KernelName>),
+                    detail::isKernelESIMD<KernelName>());
+
       MKernelName = detail::getKernelName<KernelName>();
     } else {
       // In case w/o the integration header it is necessary to process
@@ -3761,6 +3761,10 @@ class __SYCL_EXPORT handler {
                                   sycl::range<3> LocalSize, sycl::id<3> Offset,
                                   int Dims);
 
+  void setKernelInfo(void *KernelFuncPtr, int KernelNumArgs,
+                     detail::kernel_param_desc_t (*KernelParamDescGetter)(int),
+                     bool KernelIsESIMD);
+
   friend class detail::HandlerAccess;
 
 #ifdef __INTEL_PREVIEW_BREAKING_CHANGES

sycl/source/detail/handler_impl.hpp

@@ -198,6 +198,12 @@ class handler_impl {
 
   // Allocation ptr to be freed asynchronously.
   void *MFreePtr = nullptr;
+
+  // Store information about the kernel arguments.
+  void *MKernelFuncPtr = nullptr;
+  int MKernelNumArgs = 0;
+  detail::kernel_param_desc_t (*MKernelParamDescGetter)(int) = nullptr;
+  bool MKernelIsESIMD = false;
 };
 
 } // namespace detail

sycl/source/handler.cpp

@@ -411,6 +411,16 @@ event handler::finalize() {
     return MLastEvent;
   MIsFinalized = true;
 
+  // Extract arguments from the kernel lambda.
+  // TODO: Skip this in simple cases.
+  const auto &type = getType();
+  if (type == detail::CGType::Kernel && impl->MKernelFuncPtr) {
+    clearArgs();
+    extractArgsAndReqsFromLambda((char *)impl->MKernelFuncPtr,
+                                 impl->MKernelParamDescGetter,
+                                 impl->MKernelNumArgs, impl->MKernelIsESIMD);
+  }
+
   // According to 4.7.6.9 of SYCL2020 spec, if a placeholder accessor is passed
   // to a command without being bound to a command group, an exception should
   // be thrown.
@@ -448,7 +458,6 @@ event handler::finalize() {
     }
   }
 
-  const auto &type = getType();
   if (type == detail::CGType::Kernel) {
     // If there were uses of set_specialization_constant build the kernel_bundle
     std::shared_ptr<detail::kernel_bundle_impl> KernelBundleImpPtr =
@@ -2254,6 +2263,16 @@ void handler::setNDRangeDescriptorPadded(sycl::range<3> NumWorkItems,
   impl->MNDRDesc = NDRDescT{NumWorkItems, LocalSize, Offset, Dims};
 }
 
+void handler::setKernelInfo(
+    void *KernelFuncPtr, int KernelNumArgs,
+    detail::kernel_param_desc_t (*KernelParamDescGetter)(int),
+    bool KernelIsESIMD) {
+  impl->MKernelFuncPtr = KernelFuncPtr;
+  impl->MKernelNumArgs = KernelNumArgs;
+  impl->MKernelParamDescGetter = KernelParamDescGetter;
+  impl->MKernelIsESIMD = KernelIsESIMD;
+}
+
 void handler::saveCodeLoc(detail::code_location CodeLoc, bool IsTopCodeLoc) {
   MCodeLoc = CodeLoc;
   impl->MIsTopCodeLoc = IsTopCodeLoc;

sycl/test/abi/sycl_symbols_linux.dump

@@ -3523,6 +3523,7 @@ _ZN4sycl3_V17handler12addReductionERKSt10shared_ptrIKvE
 _ZN4sycl3_V17handler12setArgHelperEiRNS0_3ext6oneapi12experimental6detail30dynamic_work_group_memory_baseE
 _ZN4sycl3_V17handler12setArgHelperEiRNS0_6detail22work_group_memory_implE
 _ZN4sycl3_V17handler13getKernelNameEv
+_ZN4sycl3_V17handler13setKernelInfoEPviPFNS0_6detail19kernel_param_desc_tEiEb
 _ZN4sycl3_V17handler14addAccessorReqESt10shared_ptrINS0_6detail16AccessorImplHostEE
 _ZN4sycl3_V17handler14setNDRangeUsedEb
 _ZN4sycl3_V17handler15ext_oneapi_copyENS0_3ext6oneapi12experimental16image_mem_handleENS0_5rangeILi3EEERKNS4_16image_descriptorEPvS7_S7_S7_

sycl/test/abi/sycl_symbols_windows.dump

@@ -4377,6 +4377,7 @@
 ?setHandlerKernelBundle@handler@_V1@sycl@@AEAAXVkernel@23@@Z
 ?setKernelCacheConfig@handler@_V1@sycl@@AEAAXW4StableKernelCacheConfig@123@@Z
 ?setKernelClusterLaunch@handler@_V1@sycl@@AEAAXV?$range@$02@23@H@Z
+?setKernelInfo@handler@_V1@sycl@@AEAAXPEAXHP6A?AUkernel_param_desc_t@detail@23@H@Z_N@Z
 ?setKernelIsCooperative@handler@_V1@sycl@@AEAAX_N@Z
 ?setKernelWorkGroupMem@handler@_V1@sycl@@AEAAX_K@Z
 ?setLocalAccessorArgHelper@handler@_V1@sycl@@AEAAXHAEAVLocalAccessorBaseHost@detail@23@@Z