[SYCL][Level Zero] Implement sycl_ext_intel_cslice extension

sycl/include/sycl/detail/pi.h

@@ -630,6 +630,8 @@ static constexpr pi_device_partition_property
     PI_DEVICE_PARTITION_BY_COUNTS_LIST_END = 0x0;
 static constexpr pi_device_partition_property
     PI_DEVICE_PARTITION_BY_AFFINITY_DOMAIN = 0x1088;
+static constexpr pi_device_partition_property
+    PI_DEVICE_EXT_INTEL_PARTITION_BY_CSLICE = 0x1089;
 
 // For compatibility with OpenCL define this not as enum.
 using pi_device_affinity_domain = pi_bitfield;

sycl/include/sycl/device.hpp

@@ -178,6 +178,19 @@ class __SYCL_EXPORT device : public detail::OwnerLessBase<device> {
   std::vector<device>
   create_sub_devices(info::partition_affinity_domain AffinityDomain) const;
 
+  /// Partition device into sub devices
+  ///
+  /// Available only when prop is
+  /// info::partition_property::ext_intel_partition_by_cslice. If this SYCL
+  /// device does not support
+  /// info::partition_property::ext_intel_partition_by_cslice a
+  /// feature_not_supported exception must be thrown.
+  ///
+  /// \return a vector class of sub devices partitioned from this SYCL
+  /// device at a granularity of "cslice" (compute slice).
+  template <info::partition_property prop>
+  std::vector<device> create_sub_devices() const;
+
   /// Queries this SYCL device for information requested by the template
   /// parameter param
   ///

sycl/include/sycl/info/info_desc.hpp

@@ -55,7 +55,8 @@ enum class partition_property : pi_device_partition_property {
   no_partition = 0,
   partition_equally = PI_DEVICE_PARTITION_EQUALLY,
   partition_by_counts = PI_DEVICE_PARTITION_BY_COUNTS,
-  partition_by_affinity_domain = PI_DEVICE_PARTITION_BY_AFFINITY_DOMAIN
+  partition_by_affinity_domain = PI_DEVICE_PARTITION_BY_AFFINITY_DOMAIN,
+  ext_intel_partition_by_cslice = PI_DEVICE_EXT_INTEL_PARTITION_BY_CSLICE
 };
 
 enum class partition_affinity_domain : pi_device_affinity_domain {

sycl/plugins/level_zero/pi_level_zero.cpp

@@ -2605,29 +2605,36 @@ pi_result _pi_platform::populateDeviceCacheIfNeeded() {
             Ordinals.push_back(i);
           }
         }
-
-        // Create PI sub-sub-devices with the sub-device for all the ordinals.
-        // Each {ordinal, index} points to a specific CCS which constructs
-        // a sub-sub-device at this point.
-        // FIXME: Level Zero creates multiple PiDevices for a single physical
-        // device when sub-device is partitioned into sub-sub-devices.
-        // Sub-sub-device is technically a command queue and we should not build
-        // program for each command queue. PiDevice is probably not the right
-        // abstraction for a Level Zero command queue.
-        for (uint32_t J = 0; J < Ordinals.size(); ++J) {
-          for (uint32_t K = 0; K < QueueGroupProperties[Ordinals[J]].numQueues;
-               ++K) {
-            std::unique_ptr<_pi_device> PiSubSubDevice(
-                new _pi_device(ZeSubdevices[I], this, PiSubDevice.get()));
-            pi_result Result = PiSubSubDevice->initialize(Ordinals[J], K);
-            if (Result != PI_SUCCESS) {
-              return Result;
+        bool IsPVC =
+            (PiSubDevice->ZeDeviceProperties->deviceId & 0xff0) == 0xbd0;
+
+        // If isn't PVC, then submissions to different CCS can be executed on
+        // the same EUs still, so we cannot treat them as sub-sub-devices.
+        if (IsPVC) {
+          // Create PI sub-sub-devices with the sub-device for all the ordinals.
+          // Each {ordinal, index} points to a specific CCS which constructs
+          // a sub-sub-device at this point.
+          //
+          // FIXME: Level Zero creates multiple PiDevices for a single physical
+          // device when sub-device is partitioned into sub-sub-devices.
+          // Sub-sub-device is technically a command queue and we should not
+          // build program for each command queue. PiDevice is probably not the
+          // right abstraction for a Level Zero command queue.
+          for (uint32_t J = 0; J < Ordinals.size(); ++J) {
+            for (uint32_t K = 0;
+                 K < QueueGroupProperties[Ordinals[J]].numQueues; ++K) {
+              std::unique_ptr<_pi_device> PiSubSubDevice(
+                  new _pi_device(ZeSubdevices[I], this, PiSubDevice.get()));
+              pi_result Result = PiSubSubDevice->initialize(Ordinals[J], K);
+              if (Result != PI_SUCCESS) {
+                return Result;
+              }
+
+              // save pointers to sub-sub-devices for quick retrieval in the
+              // future.
+              PiSubDevice->SubDevices.push_back(PiSubSubDevice.get());
+              PiDevicesCache.push_back(std::move(PiSubSubDevice));
             }
-
-            // save pointers to sub-sub-devices for quick retrieval in the
-            // future.
-            PiSubDevice->SubDevices.push_back(PiSubSubDevice.get());
-            PiDevicesCache.push_back(std::move(PiSubSubDevice));
           }
         }
 
@@ -2862,31 +2869,44 @@ pi_result piDeviceGetInfo(pi_device Device, pi_device_info ParamName,
     if (ZeSubDeviceCount < 2) {
       return ReturnValue(pi_device_partition_property{0});
     }
+    bool PartitionedByCSlice = Device->SubDevices[0]
+                                   ->QueueGroup[_pi_queue::queue_type::Compute]
+                                   .ZeIndex >= 0;
+
     // It is debatable if SYCL sub-device and partitioning APIs sufficient to
     // expose Level Zero sub-devices?  We start with support of
     // "partition_by_affinity_domain" and "next_partitionable" but if that
-    // doesn't seem to be a good fit we could look at adding a more descriptive
-    // partitioning type.
+    // doesn't seem to be a good fit we could look at adding a more
+    // descriptive partitioning type.
     struct {
       pi_device_partition_property Arr[2];
-    } PartitionProperties = {{PI_DEVICE_PARTITION_BY_AFFINITY_DOMAIN, 0}};
+    } PartitionProperties = {{PartitionedByCSlice
+                                  ? PI_DEVICE_EXT_INTEL_PARTITION_BY_CSLICE
+                                  : PI_DEVICE_PARTITION_BY_AFFINITY_DOMAIN,
+                              0}};
     return ReturnValue(PartitionProperties);
   }
   case PI_DEVICE_INFO_PARTITION_AFFINITY_DOMAIN:
     return ReturnValue(pi_device_affinity_domain{
         PI_DEVICE_AFFINITY_DOMAIN_NUMA |
         PI_DEVICE_AFFINITY_DOMAIN_NEXT_PARTITIONABLE});
   case PI_DEVICE_INFO_PARTITION_TYPE: {
-    if (Device->isSubDevice()) {
+    // For root-device there is no partitioning to report.
+    if (!Device->isSubDevice())
+      return ReturnValue(pi_device_partition_property{0});
+
+    if (Device->QueueGroup[_pi_queue::queue_type::Compute].ZeIndex >= 0) {
       struct {
-        pi_device_partition_property Arr[3];
-      } PartitionProperties = {{PI_DEVICE_PARTITION_BY_AFFINITY_DOMAIN,
-                                PI_DEVICE_AFFINITY_DOMAIN_NEXT_PARTITIONABLE,
-                                0}};
+        pi_device_partition_property Arr[2];
+      } PartitionProperties = {{PI_DEVICE_EXT_INTEL_PARTITION_BY_CSLICE, 0}};
       return ReturnValue(PartitionProperties);
     }
-    // For root-device there is no partitioning to report.
-    return ReturnValue(pi_device_partition_property{0});
+
+    struct {
+      pi_device_partition_property Arr[3];
+    } PartitionProperties = {{PI_DEVICE_PARTITION_BY_AFFINITY_DOMAIN,
+                              PI_DEVICE_AFFINITY_DOMAIN_NEXT_PARTITIONABLE, 0}};
+    return ReturnValue(PartitionProperties);
   }
 
     // Everything under here is not supported yet
@@ -3258,15 +3278,19 @@ pi_result piDevicePartition(pi_device Device,
                             const pi_device_partition_property *Properties,
                             pi_uint32 NumDevices, pi_device *OutDevices,
                             pi_uint32 *OutNumDevices) {
+  PI_ASSERT(Device, PI_ERROR_INVALID_DEVICE);
   // Other partitioning ways are not supported by Level Zero
-  if (Properties[0] != PI_DEVICE_PARTITION_BY_AFFINITY_DOMAIN ||
-      (Properties[1] != PI_DEVICE_AFFINITY_DOMAIN_NEXT_PARTITIONABLE &&
-       Properties[1] != PI_DEVICE_AFFINITY_DOMAIN_NUMA)) {
+  if (Properties[0] == PI_DEVICE_PARTITION_BY_AFFINITY_DOMAIN) {
+    if ((Properties[1] != PI_DEVICE_AFFINITY_DOMAIN_NEXT_PARTITIONABLE &&
+         Properties[1] != PI_DEVICE_AFFINITY_DOMAIN_NUMA))
+      return PI_ERROR_INVALID_VALUE;
+  } else if (Properties[0] == PI_DEVICE_EXT_INTEL_PARTITION_BY_CSLICE) {
+    if (Properties[1] != 0)
+      return PI_ERROR_INVALID_VALUE;
+  } else {
     return PI_ERROR_INVALID_VALUE;
   }
 
-  PI_ASSERT(Device, PI_ERROR_INVALID_DEVICE);
-
   // Devices cache is normally created in piDevicesGet but still make
   // sure that cache is populated.
   //
@@ -3275,16 +3299,36 @@ pi_result piDevicePartition(pi_device Device,
     return Res;
   }
 
+  auto EffectiveSize = [&]() -> decltype(Device->SubDevices.size()) {
+    if (Device->SubDevices.size() == 0)
+      return 0;
+
+    // Sub-Sub-Devices are partitioned by CSlices, not by affinity domain.
+    if (Properties[0] == PI_DEVICE_PARTITION_BY_AFFINITY_DOMAIN) {
+      if (Device->isSubDevice())
+        return 0;
+    }
+    if (Properties[0] == PI_DEVICE_EXT_INTEL_PARTITION_BY_CSLICE) {
+      // Not a CSlice-based partitioning.
+      if (Device->SubDevices[0]
+              ->QueueGroup[_pi_queue::queue_type::Compute]
+              .ZeIndex < 0)
+        return 0;
+    }
+
+    return Device->SubDevices.size();
+  }();
+
   if (OutNumDevices) {
-    *OutNumDevices = Device->SubDevices.size();
+    *OutNumDevices = EffectiveSize;
   }
 
   if (OutDevices) {
     // TODO: Consider support for partitioning to <= total sub-devices.
     // Currently supported partitioning (by affinity domain/numa) would always
     // partition to all sub-devices.
     //
-    PI_ASSERT(NumDevices == Device->SubDevices.size(), PI_ERROR_INVALID_VALUE);
+    PI_ASSERT(NumDevices == EffectiveSize, PI_ERROR_INVALID_VALUE);
 
     for (uint32_t I = 0; I < NumDevices; I++) {
       OutDevices[I] = Device->SubDevices[I];

sycl/source/detail/device_impl.cpp

@@ -250,6 +250,28 @@ std::vector<device> device_impl::create_sub_devices(
   return create_sub_devices(Properties, SubDevicesCount);
 }
 
+std::vector<device> device_impl::create_sub_devices() const {
+  assert(!MIsHostDevice && "Partitioning is not supported on host.");
+
+  if (!is_partition_supported(
+          info::partition_property::ext_intel_partition_by_cslice)) {
+    throw sycl::feature_not_supported(
+        "Device does not support "
+        "sycl::info::partition_property::ext_intel_partition_by_cslice.",
+        PI_ERROR_INVALID_OPERATION);
+  }
+
+  const pi_device_partition_property Properties[2] = {
+      PI_DEVICE_EXT_INTEL_PARTITION_BY_CSLICE, 0};
+
+  pi_uint32 SubDevicesCount = 0;
+  const detail::plugin &Plugin = getPlugin();
+  Plugin.call<sycl::errc::invalid, PiApiKind::piDevicePartition>(
+      MDevice, Properties, 0, nullptr, &SubDevicesCount);
+
+  return create_sub_devices(Properties, SubDevicesCount);
+}
+
 pi_native_handle device_impl::getNative() const {
   auto Plugin = getPlugin();
   if (Plugin.getBackend() == backend::opencl)

sycl/source/detail/device_impl.hpp

@@ -172,6 +172,16 @@ class device_impl {
   std::vector<device>
   create_sub_devices(info::partition_affinity_domain AffinityDomain) const;
 
+  /// Partition device into sub devices
+  ///
+  /// If this SYCL device does not support
+  /// info::partition_property::ext_intel_partition_by_cslice a
+  /// feature_not_supported exception must be thrown.
+  ///
+  /// \return a vector class of sub devices partitioned from this SYCL
+  /// device at a granularity of "cslice" (compute slice).
+  std::vector<device> create_sub_devices() const;
+
   /// Check if desired partition property supported by device
   ///
   /// \param Prop is one of info::partition_property::(partition_equally,

sycl/source/detail/device_info.hpp

@@ -358,6 +358,7 @@ static bool is_sycl_partition_property(info::partition_property PP) {
   case info::partition_property::partition_equally:
   case info::partition_property::partition_by_counts:
   case info::partition_property::partition_by_affinity_domain:
+  case info::partition_property::ext_intel_partition_by_cslice:
     return true;
   }
   return false;

sycl/source/device.cpp

@@ -117,6 +117,14 @@ template __SYCL_EXPORT std::vector<device> device::create_sub_devices<
     info::partition_property::partition_by_affinity_domain>(
     info::partition_affinity_domain AffinityDomain) const;
 
+template <info::partition_property prop>
+std::vector<device> device::create_sub_devices() const {
+  return impl->create_sub_devices();
+}
+
+template __SYCL_EXPORT std::vector<device> device::create_sub_devices<
+    info::partition_property::ext_intel_partition_by_cslice>() const;
+
 bool device::has_extension(const std::string &extension_name) const {
   return impl->has_extension(extension_name);
 }

sycl/test/abi/sycl_symbols_linux.dump

@@ -4173,6 +4173,7 @@ _ZNK4sycl3_V16device14is_acceleratorEv
 _ZNK4sycl3_V16device18create_sub_devicesILNS0_4info18partition_propertyE4230EEESt6vectorIS1_SaIS1_EEm
 _ZNK4sycl3_V16device18create_sub_devicesILNS0_4info18partition_propertyE4231EEESt6vectorIS1_SaIS1_EERKS5_ImSaImEE
 _ZNK4sycl3_V16device18create_sub_devicesILNS0_4info18partition_propertyE4232EEESt6vectorIS1_SaIS1_EENS3_25partition_affinity_domainE
+_ZNK4sycl3_V16device18create_sub_devicesILNS0_4info18partition_propertyE4233EEESt6vectorIS1_SaIS1_EEv
 _ZNK4sycl3_V16device3getEv
 _ZNK4sycl3_V16device3hasENS0_6aspectE
 _ZNK4sycl3_V16device6is_cpuEv