[SYCL] [ABI-Break] Partial implementation of sycl_ext_oneapi_cuda_cluster_group (#14113)

This PR is a partial implementation of
[`sycl_ext_oneapi_cuda_cluster_group`](#13594),
introducing the `cluster_size` property to launch a kernel with CUDA's
thread block clusters

Only a small part of the extension specification described in
#13594 is used in this implementation.
To be specific everything after the section "Launching a kernel with a
`cluster_group`" is not included in this PR. A very important point to
note is that this PR still fully represents a functional use case of
using Nvidia's cuda driver cluster launch feature for its primary
purpose which is to accelerate cross-work-group collective operations
(particularly for GEMM), leveraging cross-work group asynchronous
multi-casting of distributed shared memory across work-groups.
This is a high priority feature that is targeted for the next release.

The other parts of the extension specification described in
#13594, primarily related to the
"cluster_group" abstraction is a (user-facing) convenience abstraction
that is not required to be exposed in libraries that optimize such
library collective operations (GEMM). Please therefore focus reviews of
this PR on the relevant aspects of the extension that are required for
the implementation in this PR and the library based application of it as
described in this message.

---------

Signed-off-by: JackAKirk <[email protected]>

Author: AD2605

llvm/include/llvm/SYCLLowerIR/DeviceConfigFile.td

@@ -82,6 +82,7 @@ def AspectExt_oneapi_limited_graph : Aspect<"ext_oneapi_limited_graph">;
 def AspectExt_oneapi_private_alloca : Aspect<"ext_oneapi_private_alloca">;
 def AspectExt_oneapi_queue_profiling_tag : Aspect<"ext_oneapi_queue_profiling_tag">;
 def AspectExt_oneapi_virtual_mem : Aspect<"ext_oneapi_virtual_mem">;
+def AspectExt_oneapi_cuda_cluster_group : Aspect<"ext_oneapi_cuda_cluster_group">;
 // Deprecated aspects
 def AspectInt64_base_atomics : Aspect<"int64_base_atomics">;
 def AspectInt64_extended_atomics : Aspect<"int64_extended_atomics">;
@@ -139,7 +140,7 @@ def : TargetInfo<"__TestAspectList",
     AspectExt_oneapi_ballot_group, AspectExt_oneapi_fixed_size_group, AspectExt_oneapi_opportunistic_group,
     AspectExt_oneapi_tangle_group, AspectExt_intel_matrix, AspectExt_oneapi_is_composite, AspectExt_oneapi_is_component,
     AspectExt_oneapi_graph, AspectExt_intel_fpga_task_sequence, AspectExt_oneapi_limited_graph,
-    AspectExt_oneapi_private_alloca, AspectExt_oneapi_queue_profiling_tag, AspectExt_oneapi_virtual_mem],
+    AspectExt_oneapi_private_alloca, AspectExt_oneapi_queue_profiling_tag, AspectExt_oneapi_virtual_mem, AspectExt_oneapi_cuda_cluster_group],
     []>;
 // This definition serves the only purpose of testing whether the deprecated aspect list defined in here and in SYCL RT
 // match.
@@ -205,9 +206,9 @@ def : CudaTargetInfo<"nvidia_gpu_sm_87", !listconcat(CudaMinAspects, CudaBindles
 def : CudaTargetInfo<"nvidia_gpu_sm_89", !listconcat(CudaMinAspects, CudaBindlessImagesAspects,
     [AspectFp16, AspectAtomic64, AspectExt_oneapi_cuda_async_barrier])>;
 def : CudaTargetInfo<"nvidia_gpu_sm_90", !listconcat(CudaMinAspects, CudaBindlessImagesAspects,
-    [AspectFp16, AspectAtomic64, AspectExt_oneapi_cuda_async_barrier])>;
+    [AspectFp16, AspectAtomic64, AspectExt_oneapi_cuda_async_barrier, AspectExt_oneapi_cuda_cluster_group])>;
 def : CudaTargetInfo<"nvidia_gpu_sm_90a", !listconcat(CudaMinAspects, CudaBindlessImagesAspects,
-    [AspectFp16, AspectAtomic64, AspectExt_oneapi_cuda_async_barrier])>;
+    [AspectFp16, AspectAtomic64, AspectExt_oneapi_cuda_async_barrier, AspectExt_oneapi_cuda_cluster_group])>;
 
 //
 // HIP / AMDGPU device aspects

sycl/doc/design/DeviceConfigFile.md

@@ -180,6 +180,7 @@ def AspectExt_intel_free_memory : Aspect<"ext_intel_free_memory">;
 def AspectExt_intel_device_id : Aspect<"ext_intel_device_id">;
 def AspectExt_intel_memory_clock_rate : Aspect<"ext_intel_memory_clock_rate">;
 def AspectExt_intel_memory_bus_width : Aspect<"ext_intel_memory_bus_width">;
+def AspectExt_oneapi_cuda_cluster_group : Aspect<"ext_oneapi_cuda_cluster_group">;
 def AspectEmulated : Aspect<"emulated">;
     
 def TargetTable : DynamicTable { 

sycl/include/sycl/detail/cg.hpp

@@ -178,6 +178,7 @@ class CGExecKernel : public CG {
   std::vector<std::shared_ptr<const void>> MAuxiliaryResources;
   sycl::detail::pi::PiKernelCacheConfig MKernelCacheConfig;
   bool MKernelIsCooperative = false;
+  bool MKernelUsesClusterLaunch = false;
 
   CGExecKernel(NDRDescT NDRDesc, std::shared_ptr<HostKernelBase> HKernel,
                std::shared_ptr<detail::kernel_impl> SyclKernel,
@@ -188,15 +189,17 @@ class CGExecKernel : public CG {
                std::vector<std::shared_ptr<const void>> AuxiliaryResources,
                CGTYPE Type,
                sycl::detail::pi::PiKernelCacheConfig KernelCacheConfig,
-               bool KernelIsCooperative, detail::code_location loc = {})
+               bool KernelIsCooperative, bool MKernelUsesClusterLaunch,
+               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)),
-        MKernelIsCooperative(KernelIsCooperative) {
+        MKernelIsCooperative(KernelIsCooperative),
+        MKernelUsesClusterLaunch(MKernelUsesClusterLaunch) {
     assert(getType() == Kernel && "Wrong type of exec kernel CG.");
   }
 

sycl/include/sycl/detail/cg_types.hpp

@@ -50,8 +50,8 @@ class ArgDesc {
   int MIndex;
 };
 
-// The structure represents NDRange - global, local sizes, global offset and
-// number of dimensions.
+// The structure represents NDRange - global, local sizes, global offset,
+// number of dimensions, and the cluster dimensions if applicable.
 class NDRDescT {
   // The method initializes all sizes for dimensions greater than the passed one
   // to the default values, so they will not affect execution.
@@ -128,13 +128,25 @@ class NDRDescT {
     Dims = Dims_;
   }
 
+  template <int Dims_> void setClusterDimensions(sycl::range<Dims_> N) {
+    if (Dims_ != Dims) {
+      throw std::runtime_error(
+          "Dimensionality of cluster, global and local ranges must be same");
+    }
+
+    for (int I = 0; I < Dims_; ++I) {
+      ClusterDimensions[I] = N[I];
+    }
+  }
+
   sycl::range<3> GlobalSize;
   sycl::range<3> LocalSize;
   sycl::id<3> GlobalOffset;
   /// Number of workgroups, used to record the number of workgroups from the
   /// simplest form of parallel_for_work_group. If set, all other fields must be
   /// zero
   sycl::range<3> NumWorkGroups;
+  sycl::range<3> ClusterDimensions{1, 1, 1};
   size_t Dims;
 };
 

sycl/include/sycl/detail/pi.def

@@ -230,4 +230,7 @@ _PI_API(piextVirtualMemGetInfo)
 // Enqueue native command
 _PI_API(piextEnqueueNativeCommand)
 
+// Kernel Launch Properties
+_PI_API(piextEnqueueKernelLaunchCustom)
+
 #undef _PI_API

sycl/include/sycl/detail/pi.h

@@ -197,9 +197,11 @@
 // pi_virtual_access_flags bit flags.
 // 15.55 Added piextEnqueueNativeCommand as well as associated types and enums
 // 16.56 Replaced piextUSMEnqueueMemset with piextUSMEnqueueFill
+// 16.57 Added mappings to UR launch properties extension
+// (piextEnqueueKernelLaunchCustom)
 
 #define _PI_H_VERSION_MAJOR 16
-#define _PI_H_VERSION_MINOR 56
+#define _PI_H_VERSION_MINOR 57
 
 #define _PI_STRING_HELPER(a) #a
 #define _PI_CONCAT(a, b) _PI_STRING_HELPER(a.b)
@@ -514,8 +516,12 @@ typedef enum {
 
   // Virtual memory support
   PI_EXT_ONEAPI_DEVICE_INFO_SUPPORTS_VIRTUAL_MEM = 0x2011E,
+
   // Native enqueue
   PI_EXT_ONEAPI_DEVICE_INFO_ENQUEUE_NATIVE_COMMAND_SUPPORT = 0x2011F,
+
+  // Return whether cluster launch is supported by device
+  PI_EXT_ONEAPI_DEVICE_INFO_CLUSTER_LAUNCH = 0x2021,
 } _pi_device_info;
 
 typedef enum {
@@ -1317,8 +1323,28 @@ typedef enum {
           ///< P2P link, otherwise such operations are not supported.
 } _pi_peer_attr;
 
+typedef enum {
+  PI_LAUNCH_PROPERTY_IGNORE = 0x0,
+  PI_LAUNCH_PROPERTY_COOPERATIVE = 0x1,
+  PI_LAUNCH_PROPERTY_CLUSTER_DIMENSION = 0x2,
+} _pi_launch_property_id;
+
+typedef union {
+  int cooperative;
+  int32_t cluster_dims[3];
+} _pi_launch_property_value;
+
 using pi_mem_info = _pi_mem_info;
 using pi_peer_attr = _pi_peer_attr;
+using pi_launch_property_id = _pi_launch_property_id;
+using pi_launch_property_value = _pi_launch_property_value;
+
+typedef struct {
+  pi_launch_property_id id;
+  pi_launch_property_value value;
+} _pi_launch_property;
+
+using pi_launch_property = _pi_launch_property;
 
 //
 // Following section contains SYCL RT Plugin Interface (PI) functions.
@@ -1933,6 +1959,14 @@ __SYCL_EXPORT pi_result piextEnqueueCooperativeKernelLaunch(
     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 piextEnqueueKernelLaunchCustom(
+    pi_queue queue, pi_kernel kernel, pi_uint32 work_dim,
+    const size_t *global_work_size, const size_t *local_work_size,
+    pi_uint32 num_props_in_launch_prop_list,
+    const pi_launch_property *launch_prop_list,
+    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/pi.hpp

@@ -148,6 +148,7 @@ using PiImageOffset = ::pi_image_offset_struct;
 using PiImageRegion = ::pi_image_region_struct;
 using PiPhysicalMem = ::pi_physical_mem;
 using PiVirtualAccessFlags = ::pi_virtual_access_flags;
+using PiLaunchProperty = ::pi_launch_property;
 
 __SYCL_EXPORT void contextSetExtendedDeleter(const sycl::context &constext,
                                              pi_context_extended_deleter func,

sycl/include/sycl/device_aspect_macros.hpp

@@ -381,6 +381,11 @@
 #define __SYCL_ALL_DEVICES_HAVE_ext_oneapi_virtual_mem__ 0
 #endif
 
+#ifndef __SYCL_ALL_DEVICES_HAVE_ext_oneapi_cuda_cluster_group__
+// __SYCL_ASPECT(ext_oneapi_cuda_cluster_group, 75)
+#define __SYCL_ALL_DEVICES_HAVE_ext_oneapi_cuda_cluster_group__ 0
+#endif
+
 #ifndef __SYCL_ANY_DEVICE_HAS_host__
 // __SYCL_ASPECT(host, 0)
 #define __SYCL_ANY_DEVICE_HAS_host__ 0
@@ -750,3 +755,8 @@
 // __SYCL_ASPECT(ext_oneapi_virtual_mem, 74)
 #define __SYCL_ANY_DEVICE_HAS_ext_oneapi_virtual_mem__ 0
 #endif
+
+#ifndef __SYCL_ANY_DEVICE_HAS_ext_oneapi_cuda_cluster_group__
+// __SYCL_ASPECT(ext_oneapi_cuda_cluster_group, 75)
+#define __SYCL_ANY_DEVICE_HAS_ext_oneapi_cuda_cluster_group__ 0
+#endif

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

@@ -0,0 +1,94 @@
+//==--- cluster_group_prop.hpp --- SYCL extension for cuda cluster group ---==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include <sycl/ext/oneapi/properties/properties.hpp>
+#include <sycl/range.hpp>
+
+namespace sycl {
+inline namespace _V1 {
+namespace ext::oneapi::experimental {
+
+namespace cuda {
+template <int Dim>
+struct cluster_size
+    : ::sycl::ext::oneapi::experimental::detail::run_time_property_key<
+          ::sycl::ext::oneapi::experimental::detail::ClusterLaunch> {
+  cluster_size(const range<Dim> &size) : size(size) {}
+  sycl::range<Dim> get_cluster_size() { return size; }
+
+private:
+  range<Dim> size;
+};
+
+template <int Dim> using cluster_size_key = cluster_size<Dim>;
+
+} // namespace cuda
+
+template <>
+struct is_property_key<cuda::cluster_size_key<1>> : std::true_type {};
+template <>
+struct is_property_key<cuda::cluster_size_key<2>> : std::true_type {};
+template <>
+struct is_property_key<cuda::cluster_size_key<3>> : std::true_type {};
+
+template <typename T>
+struct is_property_key_of<cuda::cluster_size_key<1>, T> : std::true_type {};
+
+template <typename T>
+struct is_property_key_of<cuda::cluster_size_key<2>, T> : std::true_type {};
+
+template <typename T>
+struct is_property_key_of<cuda::cluster_size_key<3>, T> : std::true_type {};
+
+template <>
+struct is_property_value<cuda::cluster_size_key<1>>
+    : is_property_key<cuda::cluster_size_key<1>> {};
+template <>
+struct is_property_value<cuda::cluster_size_key<2>>
+    : is_property_key<cuda::cluster_size_key<2>> {};
+template <>
+struct is_property_value<cuda::cluster_size_key<3>>
+    : is_property_key<cuda::cluster_size_key<3>> {};
+
+template <typename O>
+struct is_property_value_of<cuda::cluster_size_key<1>, O>
+    : is_property_key_of<cuda::cluster_size_key<1>, O> {};
+
+template <typename O>
+struct is_property_value_of<cuda::cluster_size_key<2>, O>
+    : is_property_key_of<cuda::cluster_size_key<2>, O> {};
+
+template <typename O>
+struct is_property_value_of<cuda::cluster_size_key<3>, O>
+    : is_property_key_of<cuda::cluster_size_key<3>, O> {};
+
+namespace detail {
+template <typename PropertiesT> constexpr std::size_t getClusterDim() {
+  if constexpr (PropertiesT::template has_property<
+                    sycl::ext::oneapi::experimental::cuda::cluster_size_key<
+                        1>>()) {
+    return 1;
+  }
+  if constexpr (PropertiesT::template has_property<
+                    sycl::ext::oneapi::experimental::cuda::cluster_size_key<
+                        2>>()) {
+    return 2;
+  }
+  if constexpr (PropertiesT::template has_property<
+                    sycl::ext::oneapi::experimental::cuda::cluster_size_key<
+                        3>>()) {
+    return 3;
+  }
+  return 0;
+}
+} // namespace detail
+} // namespace ext::oneapi::experimental
+} // namespace _V1
+} // namespace sycl

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

@@ -149,9 +149,9 @@ void parallel_for(handler &CGH,
   ext::oneapi::experimental::detail::LaunchConfigAccess<range<Dimensions>,
                                                         Properties>
       ConfigAccess(Config);
-  CGH.parallel_for<KernelName>(ConfigAccess.getRange(),
-                               std::forward<ReductionsT>(Reductions)...,
-                               KernelObj);
+  CGH.parallel_for<KernelName>(
+      ConfigAccess.getRange(), ConfigAccess.getProperties(),
+      std::forward<ReductionsT>(Reductions)..., KernelObj);
 }
 
 template <typename KernelName = sycl::detail::auto_name, int Dimensions,
@@ -225,9 +225,9 @@ void nd_launch(handler &CGH,
   ext::oneapi::experimental::detail::LaunchConfigAccess<nd_range<Dimensions>,
                                                         Properties>
       ConfigAccess(Config);
-  CGH.parallel_for<KernelName>(ConfigAccess.getRange(),
-                               std::forward<ReductionsT>(Reductions)...,
-                               KernelObj);
+  CGH.parallel_for<KernelName>(
+      ConfigAccess.getRange(), ConfigAccess.getProperties(),
+      std::forward<ReductionsT>(Reductions)..., KernelObj);
 }
 
 template <typename KernelName = sycl::detail::auto_name, int Dimensions,