[SYCL] Add support for native root group implementation

sycl/include/sycl/detail/cg.hpp

@@ -172,6 +172,7 @@ class CGExecKernel : public CG {
   std::vector<std::shared_ptr<detail::stream_impl>> MStreams;
   std::vector<std::shared_ptr<const void>> MAuxiliaryResources;
   sycl::detail::pi::PiKernelCacheConfig MKernelCacheConfig;
+  bool MKernelIsCooperative = false;
 
   CGExecKernel(NDRDescT NDRDesc, std::shared_ptr<HostKernelBase> HKernel,
                std::shared_ptr<detail::kernel_impl> SyclKernel,
@@ -182,14 +183,15 @@ class CGExecKernel : public CG {
                std::vector<std::shared_ptr<const void>> AuxiliaryResources,
                CGTYPE Type,
                sycl::detail::pi::PiKernelCacheConfig KernelCacheConfig,
-               detail::code_location loc = {})
+               bool KernelIsCooperative, detail::code_location loc = {})
       : CG(Type, std::move(CGData), std::move(loc)),
         MNDRDesc(std::move(NDRDesc)), MHostKernel(std::move(HKernel)),
         MSyclKernel(std::move(SyclKernel)),
         MKernelBundle(std::move(KernelBundle)), MArgs(std::move(Args)),
         MKernelName(std::move(KernelName)), MStreams(std::move(Streams)),
         MAuxiliaryResources(std::move(AuxiliaryResources)),
-        MKernelCacheConfig(std::move(KernelCacheConfig)) {
+        MKernelCacheConfig(std::move(KernelCacheConfig)),
+        MKernelIsCooperative(KernelIsCooperative) {
     assert(getType() == Kernel && "Wrong type of exec kernel CG.");
   }
 

sycl/include/sycl/detail/pi.def

@@ -87,6 +87,7 @@ _PI_API(piextKernelSetArgPointer)
 _PI_API(piKernelSetExecInfo)
 _PI_API(piextKernelCreateWithNativeHandle)
 _PI_API(piextKernelGetNativeHandle)
+_PI_API(piextKernelSuggestMaxCooperativeGroupCount)
 // Event
 _PI_API(piEventCreate)
 _PI_API(piEventGetInfo)
@@ -105,6 +106,7 @@ _PI_API(piSamplerRetain)
 _PI_API(piSamplerRelease)
 // Queue commands
 _PI_API(piEnqueueKernelLaunch)
+_PI_API(piextEnqueueCooperativeKernelLaunch)
 _PI_API(piEnqueueEventsWait)
 _PI_API(piEnqueueEventsWaitWithBarrier)
 _PI_API(piEnqueueMemBufferRead)

sycl/include/sycl/detail/pi.h

@@ -1632,6 +1632,14 @@ __SYCL_EXPORT pi_result piextKernelCreateWithNativeHandle(
 __SYCL_EXPORT pi_result
 piextKernelGetNativeHandle(pi_kernel kernel, pi_native_handle *nativeHandle);
 
+/// Gets the max work group count for a cooperative kernel.
+///
+/// \param kernel is the PI kernel being queried.
+/// \param group_count_ret is a pointer to where the query result will be
+/// stored.
+__SYCL_EXPORT pi_result piextKernelSuggestMaxCooperativeGroupCount(
+    pi_kernel kernel, pi_uint32 *group_count_ret);
+
 //
 // Events
 //
@@ -1714,6 +1722,12 @@ __SYCL_EXPORT pi_result piEnqueueKernelLaunch(
     const size_t *local_work_size, pi_uint32 num_events_in_wait_list,
     const pi_event *event_wait_list, pi_event *event);
 
+__SYCL_EXPORT pi_result piextEnqueueCooperativeKernelLaunch(
+    pi_queue queue, pi_kernel kernel, pi_uint32 work_dim,
+    const size_t *global_work_offset, const size_t *global_work_size,
+    const size_t *local_work_size, pi_uint32 num_events_in_wait_list,
+    const pi_event *event_wait_list, pi_event *event);
+
 __SYCL_EXPORT pi_result piEnqueueEventsWait(pi_queue command_queue,
                                             pi_uint32 num_events_in_wait_list,
                                             const pi_event *event_wait_list,

sycl/include/sycl/detail/spirv.hpp

@@ -23,6 +23,7 @@ struct sub_group;
 namespace experimental {
 template <typename ParentGroup> class ballot_group;
 template <size_t PartitionSize, typename ParentGroup> class fixed_size_group;
+template <int Dimensions> class root_group;
 template <typename ParentGroup> class tangle_group;
 class opportunistic_group;
 } // namespace experimental
@@ -51,6 +52,11 @@ namespace spirv {
 
 template <typename Group> struct group_scope {};
 
+template <int Dimensions>
+struct group_scope<sycl::ext::oneapi::experimental::root_group<Dimensions>> {
+  static constexpr __spv::Scope::Flag value = __spv::Scope::Flag::Device;
+};
+
 template <int Dimensions> struct group_scope<group<Dimensions>> {
   static constexpr __spv::Scope::Flag value = __spv::Scope::Flag::Workgroup;
 };

sycl/include/sycl/detail/type_traits.hpp

@@ -41,9 +41,9 @@ template <class T>
 inline constexpr bool is_fixed_topology_group_v =
     is_fixed_topology_group<T>::value;
 
-#ifdef SYCL_EXT_ONEAPI_ROOT_GROUP
-template <> struct is_fixed_topology_group<root_group> : std::true_type {};
-#endif
+template <int Dimensions> class root_group;
+template <int Dimensions>
+struct is_fixed_topology_group<root_group<Dimensions>> : std::true_type {};
 
 template <int Dimensions>
 struct is_fixed_topology_group<sycl::group<Dimensions>> : std::true_type {};

sycl/include/sycl/ext/oneapi/experimental/root_group.hpp

@@ -10,10 +10,10 @@
 
 #include <sycl/builtins.hpp>
 #include <sycl/ext/oneapi/properties/properties.hpp>
+#include <sycl/group.hpp>
 #include <sycl/memory_enums.hpp>
-#include <sycl/queue.hpp>
-
-#define SYCL_EXT_ONEAPI_ROOT_GROUP 1
+#include <sycl/nd_item.hpp>
+#include <sycl/sub_group.hpp>
 
 namespace sycl {
 inline namespace _V1 {
@@ -106,31 +106,18 @@ template <int Dimensions> root_group<Dimensions> get_root_group() {
 
 } // namespace ext::oneapi::experimental
 
-template <>
-typename ext::oneapi::experimental::info::kernel_queue_specific::
-    max_num_work_group_sync::return_type
-    kernel::ext_oneapi_get_info<
-        ext::oneapi::experimental::info::kernel_queue_specific::
-            max_num_work_group_sync>(const queue &q) const {
-  // TODO: query the backend to return a value >= 1.
-  return 1;
-}
-
 template <int dimensions>
 void group_barrier(ext::oneapi::experimental::root_group<dimensions> G,
                    memory_scope FenceScope = decltype(G)::fence_scope) {
-  (void)G;
-  (void)FenceScope;
 #ifdef __SYCL_DEVICE_ONLY__
-  // TODO: Change __spv::Scope::Workgroup to __spv::Scope::Device once backends
-  // support device scope. __spv::Scope::Workgroup is only valid when
-  // max_num_work_group_sync is 1, so that all work items in a root group will
-  // also be in the same work group.
-  __spirv_ControlBarrier(__spv::Scope::Workgroup, __spv::Scope::Workgroup,
-                         __spv::MemorySemanticsMask::SubgroupMemory |
-                             __spv::MemorySemanticsMask::WorkgroupMemory |
-                             __spv::MemorySemanticsMask::CrossWorkgroupMemory);
+  // Root group barrier first synchronizes using a work group barrier. This
+  // allows backends to ignore the second ControlBarrier (with Device scope) if
+  // their maximum number of work groups is 1.
+  group_barrier(get_child_group(G));
+  detail::spirv::ControlBarrier(G, FenceScope, memory_order::seq_cst);
 #else
+  (void)G;
+  (void)FenceScope;
   throw sycl::runtime_error("Barriers are not supported on host device",
                             PI_ERROR_INVALID_DEVICE);
 #endif

sycl/include/sycl/handler.hpp

@@ -36,6 +36,7 @@
 #include <sycl/ext/oneapi/device_global/device_global.hpp>
 #include <sycl/ext/oneapi/device_global/properties.hpp>
 #include <sycl/ext/oneapi/experimental/graph.hpp>
+#include <sycl/ext/oneapi/experimental/root_group.hpp>
 #include <sycl/ext/oneapi/kernel_properties/properties.hpp>
 #include <sycl/ext/oneapi/properties/properties.hpp>
 #include <sycl/group.hpp>
@@ -921,6 +922,10 @@ class __SYCL_EXPORT handler {
     } else {
       std::ignore = Props;
     }
+
+    constexpr bool UsesRootSync = PropertiesT::template has_property<
+        sycl::ext::oneapi::experimental::use_root_sync_key>();
+    setKernelIsCooperative(UsesRootSync);
   }
 
   /// Checks whether it is possible to copy the source shape to the destination
@@ -3604,6 +3609,8 @@ class __SYCL_EXPORT handler {
 
   // Set value of the gpu cache configuration for the kernel.
   void setKernelCacheConfig(sycl::detail::pi::PiKernelCacheConfig);
+  // Set value of the kernel is cooperative flag
+  void setKernelIsCooperative(bool);
 
   template <
       ext::oneapi::experimental::detail::UnsupportedGraphFeatures FeatureT>

sycl/include/sycl/sycl.hpp

@@ -79,6 +79,7 @@
 #include <sycl/ext/oneapi/experimental/cuda/barrier.hpp>
 #include <sycl/ext/oneapi/experimental/fixed_size_group.hpp>
 #include <sycl/ext/oneapi/experimental/opportunistic_group.hpp>
+#include <sycl/ext/oneapi/experimental/root_group.hpp>
 #include <sycl/ext/oneapi/experimental/tangle_group.hpp>
 #include <sycl/ext/oneapi/filter_selector.hpp>
 #include <sycl/ext/oneapi/functional.hpp>

sycl/plugins/cuda/pi_cuda.cpp

@@ -538,6 +538,16 @@ piEnqueueKernelLaunch(pi_queue Queue, pi_kernel Kernel, pi_uint32 WorkDim,
       NumEventsInWaitList, EventWaitList, OutEvent);
 }
 
+pi_result piextEnqueueCooperativeKernelLaunch(
+    pi_queue Queue, pi_kernel Kernel, pi_uint32 WorkDim,
+    const size_t *GlobalWorkOffset, const size_t *GlobalWorkSize,
+    const size_t *LocalWorkSize, pi_uint32 NumEventsInWaitList,
+    const pi_event *EventWaitList, pi_event *OutEvent) {
+  return pi2ur::piEnqueueKernelLaunch(
+      Queue, Kernel, WorkDim, GlobalWorkOffset, GlobalWorkSize, LocalWorkSize,
+      NumEventsInWaitList, EventWaitList, OutEvent);
+}
+
 pi_result piextKernelCreateWithNativeHandle(pi_native_handle NativeHandle,
                                             pi_context Context,
                                             pi_program Program,
@@ -553,6 +563,13 @@ pi_result piextKernelGetNativeHandle(pi_kernel Kernel,
   return pi2ur::piextKernelGetNativeHandle(Kernel, NativeHandle);
 }
 
+pi_result piextKernelSuggestMaxCooperativeGroupCount(pi_kernel Kernel,
+                                                     pi_uint32 *GroupCountRet) {
+  (void)Kernel;
+  *GroupCountRet = 1;
+  return PI_SUCCESS;
+}
+
 pi_result piEventCreate(pi_context Context, pi_event *RetEvent) {
   return pi2ur::piEventCreate(Context, RetEvent);
 }

sycl/plugins/hip/pi_hip.cpp

@@ -546,6 +546,16 @@ piEnqueueKernelLaunch(pi_queue Queue, pi_kernel Kernel, pi_uint32 WorkDim,
       NumEventsInWaitList, EventWaitList, OutEvent);
 }
 
+pi_result piextEnqueueCooperativeKernelLaunch(
+    pi_queue Queue, pi_kernel Kernel, pi_uint32 WorkDim,
+    const size_t *GlobalWorkOffset, const size_t *GlobalWorkSize,
+    const size_t *LocalWorkSize, pi_uint32 NumEventsInWaitList,
+    const pi_event *EventWaitList, pi_event *OutEvent) {
+  return pi2ur::piEnqueueKernelLaunch(
+      Queue, Kernel, WorkDim, GlobalWorkOffset, GlobalWorkSize, LocalWorkSize,
+      NumEventsInWaitList, EventWaitList, OutEvent);
+}
+
 pi_result piextKernelCreateWithNativeHandle(pi_native_handle NativeHandle,
                                             pi_context Context,
                                             pi_program Program,
@@ -561,6 +571,13 @@ pi_result piextKernelGetNativeHandle(pi_kernel Kernel,
   return pi2ur::piextKernelGetNativeHandle(Kernel, NativeHandle);
 }
 
+pi_result piextKernelSuggestMaxCooperativeGroupCount(pi_kernel Kernel,
+                                                     pi_uint32 *GroupCountRet) {
+  (void)Kernel;
+  *GroupCountRet = 1;
+  return PI_SUCCESS;
+}
+
 pi_result piEventCreate(pi_context Context, pi_event *RetEvent) {
   return pi2ur::piEventCreate(Context, RetEvent);
 }

sycl/plugins/level_zero/pi_level_zero.cpp

@@ -558,6 +558,16 @@ piEnqueueKernelLaunch(pi_queue Queue, pi_kernel Kernel, pi_uint32 WorkDim,
       NumEventsInWaitList, EventWaitList, OutEvent);
 }
 
+pi_result piextEnqueueCooperativeKernelLaunch(
+    pi_queue Queue, pi_kernel Kernel, pi_uint32 WorkDim,
+    const size_t *GlobalWorkOffset, const size_t *GlobalWorkSize,
+    const size_t *LocalWorkSize, pi_uint32 NumEventsInWaitList,
+    const pi_event *EventWaitList, pi_event *OutEvent) {
+  return pi2ur::piEnqueueKernelLaunch(
+      Queue, Kernel, WorkDim, GlobalWorkOffset, GlobalWorkSize, LocalWorkSize,
+      NumEventsInWaitList, EventWaitList, OutEvent);
+}
+
 pi_result piextKernelCreateWithNativeHandle(pi_native_handle NativeHandle,
                                             pi_context Context,
                                             pi_program Program,
@@ -573,6 +583,13 @@ pi_result piextKernelGetNativeHandle(pi_kernel Kernel,
   return pi2ur::piextKernelGetNativeHandle(Kernel, NativeHandle);
 }
 
+pi_result piextKernelSuggestMaxCooperativeGroupCount(pi_kernel Kernel,
+                                                     pi_uint32 *GroupCountRet) {
+  (void)Kernel;
+  *GroupCountRet = 1;
+  return PI_SUCCESS;
+}
+
 //
 // Events
 //

sycl/plugins/opencl/pi_opencl.cpp

@@ -2599,6 +2599,14 @@ pi_result piEventGetInfo(pi_event event, pi_event_info param_name,
   return static_cast<pi_result>(result);
 }
 
+pi_result
+piextKernelSuggestMaxCooperativeGroupCount(pi_kernel kernel,
+                                           pi_uint32 *group_count_ret) {
+  (void)kernel;
+  *group_count_ret = 1;
+  return PI_SUCCESS;
+}
+
 const char SupportedVersion[] = _PI_OPENCL_PLUGIN_VERSION_STRING;
 
 pi_result piPluginInit(pi_plugin *PluginInit) {
@@ -2683,6 +2691,8 @@ pi_result piPluginInit(pi_plugin *PluginInit) {
   _PI_CL(piextKernelSetArgPointer, piextKernelSetArgPointer)
   _PI_CL(piextKernelCreateWithNativeHandle, piextKernelCreateWithNativeHandle)
   _PI_CL(piextKernelGetNativeHandle, piextKernelGetNativeHandle)
+  _PI_CL(piextKernelSuggestMaxCooperativeGroupCount,
+         piextKernelSuggestMaxCooperativeGroupCount)
   // Event
   _PI_CL(piEventCreate, piEventCreate)
   _PI_CL(piEventGetInfo, piEventGetInfo)
@@ -2701,6 +2711,7 @@ pi_result piPluginInit(pi_plugin *PluginInit) {
   _PI_CL(piSamplerRelease, clReleaseSampler)
   // Queue commands
   _PI_CL(piEnqueueKernelLaunch, clEnqueueNDRangeKernel)
+  _PI_CL(piextEnqueueCooperativeKernelLaunch, clEnqueueNDRangeKernel)
   _PI_CL(piEnqueueEventsWait, clEnqueueMarkerWithWaitList)
   _PI_CL(piEnqueueEventsWaitWithBarrier, clEnqueueBarrierWithWaitList)
   _PI_CL(piEnqueueMemBufferRead, clEnqueueReadBuffer)

sycl/source/detail/graph_impl.cpp

@@ -651,7 +651,7 @@ exec_graph_impl::enqueue(const std::shared_ptr<sycl::detail::queue_impl> &Queue,
             // TODO: Pass accessor mem allocations
             nullptr,
             // TODO: Extract from handler
-            PI_EXT_KERNEL_EXEC_INFO_CACHE_DEFAULT);
+            PI_EXT_KERNEL_EXEC_INFO_CACHE_DEFAULT, CG->MKernelIsCooperative);
         if (Res != pi_result::PI_SUCCESS) {
           throw sycl::exception(
               sycl::make_error_code(sycl::errc::kernel),

sycl/source/detail/handler_impl.hpp

@@ -105,6 +105,8 @@ class handler_impl {
   sycl::detail::pi::PiKernelCacheConfig MKernelCacheConfig =
       PI_EXT_KERNEL_EXEC_INFO_CACHE_DEFAULT;
 
+  bool MKernelIsCooperative = false;
+
   // Extra information for bindless image copy
   sycl::detail::pi::PiMemImageDesc MImageDesc;
   sycl::detail::pi::PiMemImageFormat MImageFormat;

sycl/source/detail/jit_compiler.cpp

@@ -648,6 +648,10 @@ jit_compiler::fuseKernels(QueueImplPtr Queue,
     auto &CG = KernelCmd->getCG();
     assert(CG.getType() == CG::Kernel);
     auto *KernelCG = static_cast<CGExecKernel *>(&CG);
+    if (KernelCG->MKernelIsCooperative) {
+      printPerformanceWarning("Cannot fuse cooperative kernel");
+      return nullptr;
+    }
 
     auto KernelName = KernelCG->MKernelName;
     if (KernelName.empty()) {
@@ -878,7 +882,7 @@ jit_compiler::fuseKernels(QueueImplPtr Queue,
   FusedCG.reset(new detail::CGExecKernel(
       NDRDesc, nullptr, nullptr, std::move(KernelBundleImplPtr),
       std::move(CGData), std::move(FusedArgs), FusedKernelInfo.Name, {}, {},
-      CG::CGTYPE::Kernel, KernelCacheConfig));
+      CG::CGTYPE::Kernel, KernelCacheConfig, false));
   return FusedCG;
 }
 

sycl/source/detail/kernel_impl.hpp

@@ -16,6 +16,7 @@
 #include <sycl/detail/pi.h>
 #include <sycl/detail/pi.hpp>
 #include <sycl/device.hpp>
+#include <sycl/ext/oneapi/experimental/root_group.hpp>
 #include <sycl/info/info_desc.hpp>
 
 #include <cassert>
@@ -141,6 +142,9 @@ class kernel_impl {
   typename Param::return_type get_info(const device &Device,
                                        const range<3> &WGSize) const;
 
+  template <typename Param>
+  typename Param::return_type ext_oneapi_get_info(const queue &q) const;
+
   /// Get a reference to a raw kernel object.
   ///
   /// \return a reference to a valid PiKernel instance with raw kernel object.
@@ -255,6 +259,20 @@ kernel_impl::get_info(const device &Device,
       getPlugin());
 }
 
+template <>
+inline typename ext::oneapi::experimental::info::kernel_queue_specific::
+    max_num_work_group_sync::return_type
+    kernel_impl::ext_oneapi_get_info<
+        ext::oneapi::experimental::info::kernel_queue_specific::
+            max_num_work_group_sync>(const queue &Queue) const {
+  const auto &Plugin = getPlugin();
+  const auto &Handle = getHandleRef();
+  pi_uint32 GroupCount = 0;
+  Plugin->call<PiApiKind::piextKernelSuggestMaxCooperativeGroupCount>(
+      Handle, &GroupCount);
+  return GroupCount;
+}
+
 } // namespace detail
 } // namespace _V1
 } // namespace sycl