[SYCL] Implement updates to the root group extension (#22552)

Implement the updates made to the root group extension as described
here: #21838

Author: lbushi25

sycl/include/sycl/ext/oneapi/get_kernel_info.hpp

@@ -63,9 +63,103 @@ get_kernel_info(const queue &Q) {
   return get_kernel_info<KernelName, Param>(Q.get_context(), Q.get_device());
 }
 
-// For free functions.
 namespace experimental {
 
+// Param must be info::kernel::max_num_work_groups_sync.
+template <typename KernelName, typename Param,
+          typename LaunchProperties = empty_properties_t>
+typename std::enable_if_t<
+    std::is_same_v<Param, sycl::info::kernel::max_num_work_groups_sync>,
+    typename Param::return_type>
+get_kernel_info(const context &ctxt, const device &dev, sycl::range<1> r,
+                LaunchProperties props = {}, size_t bytes = 0) {
+  auto bundle =
+      sycl::get_kernel_bundle<KernelName, sycl::bundle_state::executable>(ctxt);
+  sycl::kernel k = bundle.template get_kernel<KernelName>();
+  return k.ext_oneapi_get_info<Param>(dev, r, props, bytes);
+}
+
+// Param must be info::kernel::max_num_work_groups_sync.
+template <typename KernelName, typename Param,
+          typename LaunchProperties = empty_properties_t>
+typename std::enable_if_t<
+    std::is_same_v<Param, sycl::info::kernel::max_num_work_groups_sync>,
+    typename Param::return_type>
+get_kernel_info(const context &ctxt, const device &dev, sycl::range<2> r,
+                LaunchProperties props = {}, size_t bytes = 0) {
+  auto bundle =
+      sycl::get_kernel_bundle<KernelName, sycl::bundle_state::executable>(ctxt);
+  sycl::kernel k = bundle.template get_kernel<KernelName>();
+  return k.ext_oneapi_get_info<Param>(dev, r, props, bytes);
+}
+
+// Param must be info::kernel::max_num_work_groups_sync.
+template <typename KernelName, typename Param,
+          typename LaunchProperties = empty_properties_t>
+typename std::enable_if_t<
+    std::is_same_v<Param, sycl::info::kernel::max_num_work_groups_sync>,
+    typename Param::return_type>
+get_kernel_info(const context &ctxt, const device &dev, sycl::range<3> r,
+                LaunchProperties props = {}, size_t bytes = 0) {
+  auto bundle =
+      sycl::get_kernel_bundle<KernelName, sycl::bundle_state::executable>(ctxt);
+  sycl::kernel k = bundle.template get_kernel<KernelName>();
+  return k.ext_oneapi_get_info<Param>(dev, r, props, bytes);
+}
+
+// Param must be info::kernel::max_num_work_groups_sync.
+template <typename KernelName, typename Param,
+          typename LaunchProperties = empty_properties_t>
+typename std::enable_if_t<
+    std::is_same_v<Param, sycl::info::kernel::max_num_work_groups_sync>,
+    typename Param::return_type>
+
+get_kernel_info(const queue &q, sycl::range<1> r,
+
+                LaunchProperties props = {}, size_t bytes = 0) {
+  sycl::context ctxt = q.get_context();
+  sycl::device dev = q.get_device();
+  auto bundle =
+      sycl::get_kernel_bundle<KernelName, sycl::bundle_state::executable>(ctxt);
+  sycl::kernel k = bundle.template get_kernel<KernelName>();
+  return k.ext_oneapi_get_info<Param>(dev, r, props, bytes);
+}
+
+// Param must be info::kernel::max_num_work_groups_sync.
+template <typename KernelName, typename Param,
+          typename LaunchProperties = empty_properties_t>
+typename std::enable_if_t<
+    std::is_same_v<Param, sycl::info::kernel::max_num_work_groups_sync>,
+    typename Param::return_type>
+
+get_kernel_info(const queue &q, sycl::range<2> r, LaunchProperties props = {},
+                size_t bytes = 0) {
+  sycl::context ctxt = q.get_context();
+  sycl::device dev = q.get_device();
+  auto bundle =
+      sycl::get_kernel_bundle<KernelName, sycl::bundle_state::executable>(ctxt);
+  sycl::kernel k = bundle.template get_kernel<KernelName>();
+  return k.ext_oneapi_get_info<Param>(dev, r, props, bytes);
+}
+
+// Param must be info::kernel::max_num_work_groups_sync.
+template <typename KernelName, typename Param,
+          typename LaunchProperties = empty_properties_t>
+typename std::enable_if_t<
+    std::is_same_v<Param, sycl::info::kernel::max_num_work_groups_sync>,
+    typename Param::return_type>
+
+get_kernel_info(const queue &q, sycl::range<3> r, LaunchProperties props = {},
+                size_t bytes = 0) {
+  sycl::context ctxt = q.get_context();
+  sycl::device dev = q.get_device();
+  auto bundle =
+      sycl::get_kernel_bundle<KernelName, sycl::bundle_state::executable>(ctxt);
+  sycl::kernel k = bundle.template get_kernel<KernelName>();
+  return k.ext_oneapi_get_info<Param>(dev, r, props, bytes);
+}
+
+// For free functions.
 template <auto *Func, typename Param>
 std::enable_if_t<ext::oneapi::experimental::is_kernel_v<Func>,
                  typename sycl::detail::is_kernel_info_desc<Param>::return_type>
@@ -106,6 +200,104 @@ get_kernel_info(const context &, const device &) {
   return sycl::detail::FreeFunctionInfoData<Func>::getNumParams();
 }
 
+// Param must be equal to max_num_work_groups_sync.
+template <auto *Func, typename Param,
+          typename LaunchProperties = empty_properties_t>
+std::enable_if_t<
+    sycl::ext::oneapi::experimental::is_kernel_v<Func> &&
+        std::is_same_v<Param, sycl::info::kernel::max_num_work_groups_sync>,
+    typename Param::return_type>
+get_kernel_info(const context &ctxt, const device &dev, sycl::range<1> r,
+                LaunchProperties props = {}, size_t bytes = 0) {
+  auto bundle = sycl::ext::oneapi::experimental::get_kernel_bundle<
+      Func, sycl::bundle_state::executable>(ctxt);
+  sycl::kernel k = bundle.template ext_oneapi_get_kernel<Func>();
+  return k.ext_oneapi_get_info<Param>(dev, r, props, bytes);
+}
+
+// Param must be equal to max_num_work_groups_sync.
+template <auto *Func, typename Param,
+          typename LaunchProperties = empty_properties_t>
+std::enable_if_t<
+    sycl::ext::oneapi::experimental::is_kernel_v<Func> &&
+        std::is_same_v<Param, sycl::info::kernel::max_num_work_groups_sync>,
+    typename Param::return_type>
+
+get_kernel_info(const context &ctxt, const device &dev, sycl::range<2> r,
+                LaunchProperties props = {}, size_t bytes = 0) {
+  auto bundle = sycl::ext::oneapi::experimental::get_kernel_bundle<
+      Func, sycl::bundle_state::executable>(ctxt);
+  sycl::kernel k = bundle.template ext_oneapi_get_kernel<Func>();
+  return k.ext_oneapi_get_info<Param>(dev, r, props, bytes);
+}
+
+// Param must be equal to max_num_work_groups_sync.
+template <auto *Func, typename Param,
+          typename LaunchProperties = empty_properties_t>
+std::enable_if_t<
+    sycl::ext::oneapi::experimental::is_kernel_v<Func> &&
+        std::is_same_v<Param, sycl::info::kernel::max_num_work_groups_sync>,
+    typename Param::return_type>
+
+get_kernel_info(const context &ctxt, const device &dev, sycl::range<3> r,
+                LaunchProperties props = {}, size_t bytes = 0) {
+  auto bundle = sycl::ext::oneapi::experimental::get_kernel_bundle<
+      Func, sycl::bundle_state::executable>(ctxt);
+  sycl::kernel k = bundle.template ext_oneapi_get_kernel<Func>();
+  return k.ext_oneapi_get_info<Param>(dev, r, props, bytes);
+}
+
+// Param must be equal to max_num_work_groups_sync.
+template <auto *Func, typename Param,
+          typename LaunchProperties = empty_properties_t>
+std::enable_if_t<
+    sycl::ext::oneapi::experimental::is_kernel_v<Func> &&
+        std::is_same_v<Param, sycl::info::kernel::max_num_work_groups_sync>,
+    typename Param::return_type>
+get_kernel_info(const queue &q, sycl::range<1> r, LaunchProperties props = {},
+                size_t bytes = 0) {
+  sycl::context ctxt = q.get_context();
+  sycl::device dev = q.get_device();
+  auto bundle = sycl::ext::oneapi::experimental::get_kernel_bundle<
+      Func, sycl::bundle_state::executable>(ctxt);
+  sycl::kernel k = bundle.template ext_oneapi_get_kernel<Func>();
+  return k.ext_oneapi_get_info<Param>(dev, r, props, bytes);
+}
+
+// Param must be equal to max_num_work_groups_sync.
+template <auto *Func, typename Param,
+          typename LaunchProperties = empty_properties_t>
+std::enable_if_t<
+    sycl::ext::oneapi::experimental::is_kernel_v<Func> &&
+        std::is_same_v<Param, sycl::info::kernel::max_num_work_groups_sync>,
+    typename Param::return_type>
+get_kernel_info(const queue &q, sycl::range<2> r, LaunchProperties props = {},
+                size_t bytes = 0) {
+  sycl::context ctxt = q.get_context();
+  sycl::device dev = q.get_device();
+  auto bundle = sycl::ext::oneapi::experimental::get_kernel_bundle<
+      Func, sycl::bundle_state::executable>(ctxt);
+  sycl::kernel k = bundle.template ext_oneapi_get_kernel<Func>();
+  return k.ext_oneapi_get_info<Param>(dev, r, props, bytes);
+}
+
+// Param must be equal to max_num_work_groups_sync.
+template <auto *Func, typename Param,
+          typename LaunchProperties = empty_properties_t>
+std::enable_if_t<
+    sycl::ext::oneapi::experimental::is_kernel_v<Func> &&
+        std::is_same_v<Param, sycl::info::kernel::max_num_work_groups_sync>,
+    typename Param::return_type>
+get_kernel_info(const queue &q, sycl::range<3> r, LaunchProperties props = {},
+                size_t bytes = 0) {
+  sycl::context ctxt = q.get_context();
+  sycl::device dev = q.get_device();
+  auto bundle = sycl::ext::oneapi::experimental::get_kernel_bundle<
+      Func, sycl::bundle_state::executable>(ctxt);
+  sycl::kernel k = bundle.template ext_oneapi_get_kernel<Func>();
+  return k.ext_oneapi_get_info<Param>(dev, r, props, bytes);
+}
+
 } // namespace experimental
 } // namespace ext::oneapi
 } // namespace _V1

sycl/include/sycl/info/kernel.hpp

@@ -23,8 +23,12 @@ class context;
 
 namespace info {
 
-// A.5 Kernel information desctiptors
+// A.5 Kernel information descriptors
 namespace kernel {
+struct max_num_work_groups_sync {
+  using return_type = size_t;
+};
+
 template <ur_kernel_info_t UrCode>
 using kernel_traits =
     sycl::detail::ur_traits_base<sycl::detail::info_class::kernel, UrCode>;

sycl/include/sycl/kernel.hpp

@@ -13,6 +13,9 @@
 #include <sycl/detail/export.hpp>             // for __SYCL_EXPORT
 #include <sycl/detail/owner_less_base.hpp>    // for OwnerLessBase
 #include <sycl/detail/util.hpp>
+#include <sycl/ext/intel/experimental/kernel_execution_properties.hpp> // for cache_config
+#include <sycl/ext/oneapi/properties.hpp> // for empty_properties_t
+#include <sycl/ext/oneapi/work_group_scratch_memory.hpp> // for work_group_scratch_size
 #include <sycl/info/kernel.hpp>         // for is_kernel_device_specif...
 #include <sycl/kernel_bundle_enums.hpp> // for bundle_state
 #include <unified-runtime/ur_api.h>     // for ur_native_handle_t
@@ -253,6 +256,75 @@ class __SYCL_EXPORT kernel : public detail::OwnerLessBase<kernel> {
   typename detail::is_kernel_queue_specific_info_desc<Param>::return_type
   ext_oneapi_get_info(queue Queue, const range<1> &WG) const;
 
+  /// Query the maximum number of work-groups that are allowed when the kernel
+  /// uses root-group synchronization, assuming the kernel is launched with a
+  /// work-group size of r, kernel launch properties props, and bytes bytes of
+  /// dynamic work-group local memory. If the kernel can be submitted to the
+  /// device without an error, the minimum value returned by this query is 1,
+  /// otherwise it is 0. Param must be equal to max_num_work_groups_sync.
+  /// \param Dev is a valid SYCL device
+  /// \param Range is the execution range of the kernel
+  /// \param Props are the launch properties of the kernel
+  /// \param Bytes is the amount of dynamic work group local memory allocated by
+  /// the kernel
+  template <
+      typename Param,
+      typename LaunchProperties = ext::oneapi::experimental::empty_properties_t,
+      typename = std::enable_if_t<
+          std::is_same_v<Param, info::kernel::max_num_work_groups_sync>, void>>
+  typename Param::return_type
+  ext_oneapi_get_info(const device &Dev, sycl::range<1> Range,
+                      LaunchProperties Props = {}, size_t Bytes = 0) const {
+    return queryMaxNumWorkGroups<1>(Dev, Range, Bytes,
+                                    ProcessLaunchProperties(Props));
+  }
+
+  /// Query the maximum number of work-groups that are allowed when the kernel
+  /// uses root-group synchronization, assuming the kernel is launched with a
+  /// work-group size of r, kernel launch properties props, and bytes bytes of
+  /// dynamic work-group local memory. If the kernel can be submitted to the
+  /// device without an error, the minimum value returned by this query is 1,
+  /// otherwise it is 0. Param must be equal to max_num_work_groups_sync.
+  /// \param Dev is a valid SYCL device
+  /// \param Range is the execution range of the kernel
+  /// \param Props are the launch properties of the kernel
+  /// \param Bytes is the amount of dynamic work group local memory allocated by
+  /// the kernel
+  template <
+      typename Param,
+      typename LaunchProperties = ext::oneapi::experimental::empty_properties_t,
+      typename = std::enable_if_t<
+          std::is_same_v<Param, info::kernel::max_num_work_groups_sync>, void>>
+  typename Param::return_type
+  ext_oneapi_get_info(const device &Dev, sycl::range<2> Range,
+                      LaunchProperties Props = {}, size_t Bytes = 0) const {
+    return queryMaxNumWorkGroups<2>(Dev, Range, Bytes,
+                                    ProcessLaunchProperties(Props));
+  }
+
+  /// Query the maximum number of work-groups that are allowed when the kernel
+  /// uses root-group synchronization, assuming the kernel is launched with a
+  /// work-group size of r, kernel launch properties props, and bytes bytes of
+  /// dynamic work-group local memory. If the kernel can be submitted to the
+  /// device without an error, the minimum value returned by this query is 1,
+  /// otherwise it is 0. Param must be equal to max_num_work_groups_sync.
+  /// \param Dev is a valid SYCL device
+  /// \param Range is the execution range of the kernel
+  /// \param Props are the launch properties of the kernel
+  /// \param Bytes is the amount of dynamic work group local memory allocated by
+  /// the kernel
+  template <
+      typename Param,
+      typename LaunchProperties = ext::oneapi::experimental::empty_properties_t,
+      typename = std::enable_if_t<
+          std::is_same_v<Param, info::kernel::max_num_work_groups_sync>, void>>
+  typename Param::return_type
+  ext_oneapi_get_info(const device &Dev, sycl::range<3> Range,
+                      LaunchProperties Props = {}, size_t Bytes = 0) const {
+    return queryMaxNumWorkGroups<3>(Dev, Range, Bytes,
+                                    ProcessLaunchProperties(Props));
+  }
+
 private:
   /// Constructs a SYCL kernel object from a valid kernel_impl instance.
   kernel(std::shared_ptr<detail::kernel_impl> Impl);
@@ -268,6 +340,47 @@ class __SYCL_EXPORT kernel : public detail::OwnerLessBase<kernel> {
   typename detail::ABINeutralT_t<
       typename detail::is_kernel_info_desc<Param>::return_type>
   get_info_impl() const;
+
+  // The type below bundles together some execution information for the kernel
+  // such as the amount of dynamic work group memory requested represented by
+  // DynamicWorkGroupMem and the cache config setting represented by
+  // CacheConfig.
+  struct KernelExecInfoTy {
+    size_t DynamicWorkGroupMem;
+    ur_kernel_cache_config_t CacheConfig;
+  };
+
+  template <int Dimensions>
+  size_t queryMaxNumWorkGroups(const device &Dev, sycl::range<Dimensions> Range,
+                               size_t Bytes, KernelExecInfoTy ExecInfo) const;
+
+  template <
+      typename LaunchProperties = ext::oneapi::experimental::empty_properties_t>
+  KernelExecInfoTy ProcessLaunchProperties(LaunchProperties Props) const {
+    ur_kernel_cache_config_t CacheConfig =
+        ur_kernel_cache_config_t::UR_KERNEL_CACHE_CONFIG_DEFAULT;
+    size_t DynamicLocalMem = 0;
+    if constexpr (LaunchProperties::template has_property<
+                      ext::intel::experimental::cache_config_key>()) {
+      auto prop = Props.template get_property<
+          ext::intel::experimental::cache_config_key>();
+      if (prop.value == ext::intel::experimental::cache_config_enum::large_slm)
+        CacheConfig =
+            ur_kernel_cache_config_t::UR_KERNEL_CACHE_CONFIG_LARGE_SLM;
+      else if (prop.value ==
+               ext::intel::experimental::cache_config_enum::large_data) {
+        CacheConfig =
+            ur_kernel_cache_config_t::UR_KERNEL_CACHE_CONFIG_LARGE_DATA;
+      }
+    }
+    if constexpr (LaunchProperties::template has_property<
+                      ext::oneapi::experimental::work_group_scratch_size>()) {
+      auto prop = Props.template get_property<
+          ext::oneapi::experimental::work_group_scratch_size>();
+      DynamicLocalMem = prop.size;
+    }
+    return {DynamicLocalMem, CacheConfig};
+  }
 };
 } // namespace _V1
 } // namespace sycl

sycl/source/detail/kernel_impl.hpp

@@ -243,6 +243,10 @@ class kernel_impl {
                : ProgramManager::getInstance().getDeviceKernelInfo(
                      std::string_view(getName()));
   }
+  template <int Dimensions>
+  size_t queryMaxNumWorkGroups(queue Queue,
+                               const range<Dimensions> &WorkGroupSize,
+                               size_t DynamicLocalMemorySize) const;
 
 private:
   Managed<ur_kernel_handle_t> MKernel;
@@ -273,10 +277,6 @@ class kernel_impl {
   bool exceedsOccupancyResourceLimits(const device &Device,
                                       const range<Dimensions> &WorkGroupSize,
                                       size_t DynamicLocalMemorySize) const;
-  template <int Dimensions>
-  size_t queryMaxNumWorkGroups(queue Queue,
-                               const range<Dimensions> &WorkGroupSize,
-                               size_t DynamicLocalMemorySize) const;
 
   void enableUSMIndirectAccess() const;
   std::optional<unsigned> getFreeFuncKernelArgSize() const;