[SYCL] Optional kernel features: implement split based on reqd-work-group-size

llvm/test/tools/sycl-post-link/sycl-esimd-large-grf.ll

@@ -9,16 +9,16 @@
 
 ; RUN: sycl-post-link -split=source -symbols -split-esimd -lower-esimd -S %s -o %t.table
 ; RUN: FileCheck %s -input-file=%t.table
-; RUN: FileCheck %s -input-file=%t_esimd_large_grf_1.ll --check-prefixes CHECK-ESIMD-LargeGRF-IR
-; RUN: FileCheck %s -input-file=%t_esimd_large_grf_1.prop --check-prefixes CHECK-ESIMD-LargeGRF-PROP
-; RUN: FileCheck %s -input-file=%t_0.sym --check-prefixes CHECK-SYCL-SYM
-; RUN: FileCheck %s -input-file=%t_esimd_0.sym --check-prefixes CHECK-ESIMD-SYM
-; RUN: FileCheck %s -input-file=%t_esimd_large_grf_1.sym --check-prefixes CHECK-ESIMD-LargeGRF-SYM
+; RUN: FileCheck %s -input-file=%t_esimd_large_grf_0.ll --check-prefixes CHECK-ESIMD-LargeGRF-IR
+; RUN: FileCheck %s -input-file=%t_esimd_large_grf_0.prop --check-prefixes CHECK-ESIMD-LargeGRF-PROP
+; RUN: FileCheck %s -input-file=%t_1.sym --check-prefixes CHECK-SYCL-SYM
+; RUN: FileCheck %s -input-file=%t_esimd_1.sym --check-prefixes CHECK-ESIMD-SYM
+; RUN: FileCheck %s -input-file=%t_esimd_large_grf_0.sym --check-prefixes CHECK-ESIMD-LargeGRF-SYM
 
 ; CHECK: [Code|Properties|Symbols]
 ; CHECK: {{.*}}_0.ll|{{.*}}_0.prop|{{.*}}_0.sym
-; CHECK: {{.*}}esimd_0.ll|{{.*}}esimd_0.prop|{{.*}}esimd_0.sym
-; CHECK: {{.*}}esimd_large_grf_1.ll|{{.*}}esimd_large_grf_1.prop|{{.*}}esimd_large_grf_1.sym
+; CHECK: {{.*}}esimd-large-grf.ll.tmp_1.ll|{{.*}}esimd-large-grf.ll.tmp_1.prop|{{.*}}esimd-large-grf.ll.tmp_1.sym
+; CHECK: {{.*}}esimd-large-grf.ll.tmp_esimd_1.ll|{{.*}}esimd-large-grf.ll.tmp_esimd_1.prop|{{.*}}esimd-large-grf.ll.tmp_esimd_1.sym
 
 ; CHECK-ESIMD-LargeGRF-PROP: isEsimdImage=1|1
 ; CHECK-ESIMD-LargeGRF-PROP: isLargeGRF=1|1

llvm/test/tools/sycl-post-link/sycl-large-grf.ll

@@ -9,14 +9,14 @@
 
 ; RUN: sycl-post-link -split=source -symbols -split-esimd -lower-esimd -S %s -o %t.table
 ; RUN: FileCheck %s -input-file=%t.table
-; RUN: FileCheck %s -input-file=%t_large_grf_1.ll --check-prefixes CHECK-LARGE-GRF-IR
-; RUN: FileCheck %s -input-file=%t_large_grf_1.prop --check-prefixes CHECK-LARGE-GRF-PROP
-; RUN: FileCheck %s -input-file=%t_0.sym --check-prefixes CHECK-SYCL-SYM
-; RUN: FileCheck %s -input-file=%t_large_grf_1.sym --check-prefixes CHECK-LARGE-GRF-SYM
+; RUN: FileCheck %s -input-file=%t_large_grf_0.ll --check-prefixes CHECK-LARGE-GRF-IR
+; RUN: FileCheck %s -input-file=%t_large_grf_0.prop --check-prefixes CHECK-LARGE-GRF-PROP
+; RUN: FileCheck %s -input-file=%t_1.sym --check-prefixes CHECK-SYCL-SYM
+; RUN: FileCheck %s -input-file=%t_large_grf_0.sym --check-prefixes CHECK-LARGE-GRF-SYM
 
 ; CHECK: [Code|Properties|Symbols]
 ; CHECK: {{.*}}_0.ll|{{.*}}_0.prop|{{.*}}_0.sym
-; CHECK: {{.*}}_large_grf_1.ll|{{.*}}_large_grf_1.prop|{{.*}}_large_grf_1.sym
+; CHECK: {{.*}}-large-grf.ll.tmp_1.ll|{{.*}}-large-grf.ll.tmp_1.prop|{{.*}}-large-grf.ll.tmp_1.sym
 
 ; CHECK-LARGE-GRF-PROP: isLargeGRF=1|1
 

llvm/tools/sycl-post-link/ModuleSplitter.cpp

@@ -13,6 +13,7 @@
 #include "Support.h"
 
 #include "llvm/ADT/SetVector.h"
+#include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
@@ -712,8 +713,8 @@ namespace {
 struct UsedOptionalFeatures {
   SmallVector<int, 4> Aspects;
   bool UsesLargeGRF = false;
-  // TODO: extend this further with reqd-sub-group-size, reqd-work-group-size
-  // and other properties
+  SmallVector<int, 3> ReqdWorkGroupSize;
+  // TODO: extend this further with reqd-sub-group-size and other properties
 
   UsedOptionalFeatures() = default;
 
@@ -734,17 +735,37 @@ struct UsedOptionalFeatures {
     if (F->hasFnAttribute(::sycl::kernel_props::ATTR_LARGE_GRF))
       UsesLargeGRF = true;
 
+    if (const MDNode *MDN = F->getMetadata("reqd_work_group_size")) {
+      size_t NumOperands = MDN->getNumOperands();
+      assert(NumOperands >= 1 && NumOperands <= 3 &&
+             "reqd_work_group_size does not have between 1 and 3 operands.");
+      ReqdWorkGroupSize.reserve(NumOperands);
+      for (const MDOperand &MDOp : MDN->operands())
+        ReqdWorkGroupSize.push_back(
+            mdconst::extract<ConstantInt>(MDOp)->getZExtValue());
+    }
+
     llvm::hash_code AspectsHash =
         llvm::hash_combine_range(Aspects.begin(), Aspects.end());
     llvm::hash_code LargeGRFHash = llvm::hash_value(UsesLargeGRF);
-    Hash = static_cast<unsigned>(llvm::hash_combine(AspectsHash, LargeGRFHash));
+    llvm::hash_code ReqdWorkGroupSizeHash = llvm::hash_combine_range(
+        ReqdWorkGroupSize.begin(), ReqdWorkGroupSize.end());
+    Hash = static_cast<unsigned>(
+        llvm::hash_combine(AspectsHash, LargeGRFHash, ReqdWorkGroupSizeHash));
   }
 
   std::string generateModuleName(StringRef BaseName) const {
+    std::string Ret = BaseName.str();
+    if (!ReqdWorkGroupSize.empty()) {
+      Ret += "-reqd-wg-size";
+      for (int V : ReqdWorkGroupSize)
+        Ret += "-" + std::to_string(V);
+    }
+
     if (Aspects.empty())
-      return BaseName.str() + "-no-aspects";
+      return Ret + "-no-aspects";
 
-    std::string Ret = BaseName.str() + "-aspects";
+    Ret += "-aspects";
     for (int A : Aspects) {
       Ret += "-" + std::to_string(A);
     }

llvm/tools/sycl-post-link/SYCLDeviceRequirements.cpp

@@ -30,7 +30,9 @@ void llvm::getSYCLDeviceRequirements(
   // Scan the module and if the metadata is present fill the corresponing
   // property with metadata's aspects
   constexpr std::pair<const char *, const char *> ReqdMDs[] = {
-      {"sycl_used_aspects", "aspects"}, {"sycl_fixed_targets", "fixed_target"}};
+      {"sycl_used_aspects", "aspects"},
+      {"sycl_fixed_targets", "fixed_target"},
+      {"reqd_work_group_size", "reqd_work_group_size"}};
 
   for (const auto &MD : ReqdMDs) {
     std::set<uint32_t> Aspects;

sycl/source/detail/program_manager/program_manager.cpp

@@ -1711,7 +1711,7 @@ ProgramManager::getSYCLDeviceImagesWithCompatibleState(
 
     for (const sycl::device &Dev : Devs) {
       if (!compatibleWithDevice(BinImage, Dev) ||
-          !doesDevSupportImgAspects(Dev, *BinImage))
+          !doesDevSupportDeviceRequirements(Dev, *BinImage))
         continue;
 
       std::shared_ptr<std::vector<sycl::kernel_id>> KernelIDs;
@@ -2178,26 +2178,57 @@ std::pair<RT::PiKernel, std::mutex *> ProgramManager::getOrCreateKernel(
                         &(BuildResult->MBuildResultMutex));
 }
 
-bool doesDevSupportImgAspects(const device &Dev,
-                              const RTDeviceBinaryImage &Img) {
-  const RTDeviceBinaryImage::PropertyRange &PropRange =
-      Img.getDeviceRequirements();
-  RTDeviceBinaryImage::PropertyRange::ConstIterator PropIt = std::find_if(
-      PropRange.begin(), PropRange.end(),
-      [](RTDeviceBinaryImage::PropertyRange::ConstIterator &&Prop) {
-        using namespace std::literals;
-        return (*Prop)->Name == "aspects"sv;
-      });
-  if (PropIt == PropRange.end())
+bool doesDevSupportDeviceRequirements(const device &Dev,
+                                      const RTDeviceBinaryImage &Img) {
+  auto getPropIt = [&Img](const std::string &PropName) {
+    const RTDeviceBinaryImage::PropertyRange &PropRange =
+        Img.getDeviceRequirements();
+    RTDeviceBinaryImage::PropertyRange::ConstIterator PropIt = std::find_if(
+        PropRange.begin(), PropRange.end(),
+        [&PropName](RTDeviceBinaryImage::PropertyRange::ConstIterator &&Prop) {
+          return (*Prop)->Name == PropName;
+        });
+    return (PropIt == PropRange.end())
+               ? std::nullopt
+               : std::optional<
+                     RTDeviceBinaryImage::PropertyRange::ConstIterator>{PropIt};
+  };
+
+  auto AspectsPropIt = getPropIt("aspects");
+  auto ReqdWGSizePropIt = getPropIt("reqd_work_group_size");
+
+  if (!AspectsPropIt && !ReqdWGSizePropIt) {
     return true;
-  ByteArray Aspects = DeviceBinaryProperty(*PropIt).asByteArray();
-  // Drop 8 bytes describing the size of the byte array.
-  Aspects.dropBytes(8);
-  while (!Aspects.empty()) {
-    aspect Aspect = Aspects.consume<aspect>();
-    if (!Dev.has(Aspect))
+  }
+
+  // Checking if device supports defined aspects
+  if (AspectsPropIt) {
+    ByteArray Aspects =
+        DeviceBinaryProperty(*(AspectsPropIt.value())).asByteArray();
+    // Drop 8 bytes describing the size of the byte array.
+    Aspects.dropBytes(8);
+    while (!Aspects.empty()) {
+      aspect Aspect = Aspects.consume<aspect>();
+      if (!Dev.has(Aspect))
+        return false;
+    }
+  }
+
+  // Checking if device supports defined required work group size
+  if (ReqdWGSizePropIt) {
+    ByteArray ReqdWGSize =
+        DeviceBinaryProperty(*(ReqdWGSizePropIt.value())).asByteArray();
+    // Drop 8 bytes describing the size of the byte array.
+    ReqdWGSize.dropBytes(8);
+    int ReqdWGSizeAllDimsTotal = 1;
+    while (!ReqdWGSize.empty()) {
+      ReqdWGSizeAllDimsTotal *= ReqdWGSize.consume<int>();
+    }
+    if (static_cast<long unsigned int>(ReqdWGSizeAllDimsTotal) >
+        Dev.get_info<info::device::max_work_group_size>())
       return false;
   }
+
   return true;
 }
 

sycl/source/detail/program_manager/program_manager.hpp

@@ -46,8 +46,8 @@ __SYCL_INLINE_VER_NAMESPACE(_V1) {
 class context;
 namespace detail {
 
-bool doesDevSupportImgAspects(const device &Dev,
-                              const RTDeviceBinaryImage &BinImages);
+bool doesDevSupportDeviceRequirements(const device &Dev,
+                                      const RTDeviceBinaryImage &BinImages);
 
 // This value must be the same as in libdevice/device_itt.h.
 // See sycl/doc/design/ITTAnnotations.md for more info.

sycl/source/kernel_bundle.cpp

@@ -294,7 +294,7 @@ bool is_compatible(const std::vector<kernel_id> &KernelIDs, const device &Dev) {
       detail::ProgramManager::getInstance().getRawDeviceImages(KernelIDs);
   return std::all_of(BinImages.begin(), BinImages.end(),
                      [&Dev](const detail::RTDeviceBinaryImage *Img) {
-                       return doesDevSupportImgAspects(Dev, *Img);
+                       return doesDevSupportDeviceRequirements(Dev, *Img);
                      });
 }
 

sycl/unittests/SYCL2020/IsCompatible.cpp

@@ -6,6 +6,8 @@
 #include <gtest/gtest.h>
 
 class TestKernelCPU;
+class TestKernelCPUInvalidReqdWGSize1D;
+class TestKernelCPUInvalidReqdWGSize2D;
 class TestKernelGPU;
 class TestKernelACC;
 
@@ -29,6 +31,46 @@ template <> struct KernelInfo<TestKernelCPU> {
 } // __SYCL_INLINE_VER_NAMESPACE(_V1)
 } // namespace sycl
 
+namespace sycl {
+__SYCL_INLINE_VER_NAMESPACE(_V1) {
+namespace detail {
+template <> struct KernelInfo<TestKernelCPUInvalidReqdWGSize1D> {
+  static constexpr unsigned getNumParams() { return 0; }
+  static const kernel_param_desc_t &getParamDesc(int) {
+    static kernel_param_desc_t Dummy;
+    return Dummy;
+  }
+  static constexpr const char *getName() { return "TestKernelCPUInvalidReqdWGSize1D"; }
+  static constexpr bool isESIMD() { return false; }
+  static constexpr bool callsThisItem() { return false; }
+  static constexpr bool callsAnyThisFreeFunction() { return false; }
+  static constexpr int64_t getKernelSize() { return 1; }
+};
+
+} // namespace detail
+} // __SYCL_INLINE_VER_NAMESPACE(_V1)
+} // namespace sycl
+
+namespace sycl {
+__SYCL_INLINE_VER_NAMESPACE(_V1) {
+namespace detail {
+template <> struct KernelInfo<TestKernelCPUInvalidReqdWGSize2D> {
+  static constexpr unsigned getNumParams() { return 0; }
+  static const kernel_param_desc_t &getParamDesc(int) {
+    static kernel_param_desc_t Dummy;
+    return Dummy;
+  }
+  static constexpr const char *getName() { return "TestKernelCPUInvalidReqdWGSize2D"; }
+  static constexpr bool isESIMD() { return false; }
+  static constexpr bool callsThisItem() { return false; }
+  static constexpr bool callsAnyThisFreeFunction() { return false; }
+  static constexpr int64_t getKernelSize() { return 1; }
+};
+
+} // namespace detail
+} // __SYCL_INLINE_VER_NAMESPACE(_V1)
+} // namespace sycl
+
 namespace sycl {
 __SYCL_INLINE_VER_NAMESPACE(_V1) {
 namespace detail {
@@ -71,11 +113,11 @@ template <> struct KernelInfo<TestKernelACC> {
 
 static sycl::unittest::PiImage
 generateDefaultImage(std::initializer_list<std::string> KernelNames,
-                     const std::vector<sycl::aspect> &Aspects) {
+                     const std::vector<sycl::aspect> &Aspects, const std::vector<int> &ReqdWGSize = {}) {
   using namespace sycl::unittest;
 
   PiPropertySet PropSet;
-  addAspects(PropSet, Aspects);
+  addDeviceRequirementsProps(PropSet, Aspects, ReqdWGSize);
 
   std::vector<unsigned char> Bin{0, 1, 2, 3, 4, 5}; // Random data
 
@@ -92,12 +134,16 @@ generateDefaultImage(std::initializer_list<std::string> KernelNames,
   return Img;
 }
 
-static sycl::unittest::PiImage Imgs[3] = {
-    generateDefaultImage({"TestKernelCPU"}, {sycl::aspect::cpu}),
+static sycl::unittest::PiImage Imgs[5] = {
+    generateDefaultImage({"TestKernelCPU"}, {sycl::aspect::cpu}, {32}),
+    generateDefaultImage({"TestKernelCPUInvalidReqdWGSize1D"},
+                         {sycl::aspect::cpu}, {257}), // 257 > 256
+    generateDefaultImage({"TestKernelCPUInvalidReqdWGSize2D"},
+                         {sycl::aspect::cpu}, {32, 9}), // 32*9=288 > 256
     generateDefaultImage({"TestKernelGPU"}, {sycl::aspect::gpu}),
     generateDefaultImage({"TestKernelACC"}, {sycl::aspect::accelerator})};
 
-static sycl::unittest::PiImageArray<3> ImgArray{Imgs};
+static sycl::unittest::PiImageArray<5> ImgArray{Imgs};
 
 static pi_result redefinedDeviceGetInfoCPU(pi_device device,
                                            pi_device_info param_name,
@@ -108,6 +154,10 @@ static pi_result redefinedDeviceGetInfoCPU(pi_device device,
     auto *Result = reinterpret_cast<_pi_device_type *>(param_value);
     *Result = PI_DEVICE_TYPE_CPU;
   }
+  if (param_name == PI_DEVICE_INFO_MAX_WORK_GROUP_SIZE) {
+    auto *Result = reinterpret_cast<int *>(param_value);
+    *Result = 256;
+  }
   return PI_SUCCESS;
 }
 
@@ -150,6 +200,26 @@ TEST(IsCompatible, CPU) {
   EXPECT_FALSE(sycl::is_compatible<TestKernelACC>(Dev));
 }
 
+TEST(IsCompatible, CPUInvalidReqdWGSize1D) {
+  sycl::unittest::PiMock Mock;
+  Mock.redefineAfter<sycl::detail::PiApiKind::piDeviceGetInfo>(
+      redefinedDeviceGetInfoCPU);
+  sycl::platform Plt = Mock.getPlatform();
+  const sycl::device Dev = Plt.get_devices()[0];
+
+  EXPECT_FALSE(sycl::is_compatible<TestKernelCPUInvalidReqdWGSize1D>(Dev));
+}
+
+TEST(IsCompatible, CPUInvalidReqdWGSize2D) {
+  sycl::unittest::PiMock Mock;
+  Mock.redefineAfter<sycl::detail::PiApiKind::piDeviceGetInfo>(
+      redefinedDeviceGetInfoCPU);
+  sycl::platform Plt = Mock.getPlatform();
+  const sycl::device Dev = Plt.get_devices()[0];
+
+  EXPECT_FALSE(sycl::is_compatible<TestKernelCPUInvalidReqdWGSize2D>(Dev));
+}
+
 TEST(IsCompatible, GPU) {
   sycl::unittest::PiMock Mock;
   Mock.redefineAfter<sycl::detail::PiApiKind::piDeviceGetInfo>(

sycl/unittests/SYCL2020/KernelBundle.cpp

@@ -74,7 +74,7 @@ generateDefaultImage(std::initializer_list<std::string> KernelNames,
 
   PiPropertySet PropSet;
   if (!Aspects.empty())
-    addAspects(PropSet, Aspects);
+    addDeviceRequirementsProps(PropSet, Aspects);
 
   std::vector<unsigned char> Bin{0, 1, 2, 3, 4, 5}; // Random data
 

sycl/unittests/helpers/PiImage.hpp

@@ -479,8 +479,7 @@ inline PiProperty makeDeviceGlobalInfo(const std::string &Name,
 }
 
 /// Utility function to add aspects to property set.
-inline void addAspects(PiPropertySet &Props,
-                       const std::vector<sycl::aspect> &Aspects) {
+inline PiProperty makeAspectsProp(const std::vector<sycl::aspect> &Aspects) {
   const size_t BYTES_FOR_SIZE = 8;
   std::vector<char> ValData(BYTES_FOR_SIZE +
                             Aspects.size() * sizeof(sycl::aspect));
@@ -490,8 +489,29 @@ inline void addAspects(PiPropertySet &Props,
   auto *AspectsPtr = reinterpret_cast<const unsigned char *>(&Aspects[0]);
   std::uninitialized_copy(AspectsPtr, AspectsPtr + Aspects.size(),
                           ValData.data() + BYTES_FOR_SIZE);
-  PiProperty Prop{"aspects", ValData, PI_PROPERTY_TYPE_BYTE_ARRAY};
-  PiArray<PiProperty> Value{std::move(Prop)};
+  return {"aspects", ValData, PI_PROPERTY_TYPE_BYTE_ARRAY};
+}
+
+inline PiProperty makeReqdWGSizeProp(const std::vector<int> &ReqdWGSize) {
+  const size_t BYTES_FOR_SIZE = 8;
+  std::vector<char> ValData(BYTES_FOR_SIZE + ReqdWGSize.size() * sizeof(int));
+  uint64_t ValDataSize = ValData.size();
+  std::uninitialized_copy(&ValDataSize, &ValDataSize + sizeof(uint64_t),
+                          ValData.data());
+  auto *ReqdWGSizePtr = reinterpret_cast<const unsigned char *>(&ReqdWGSize[0]);
+  std::uninitialized_copy(ReqdWGSizePtr,
+                          ReqdWGSizePtr + ReqdWGSize.size() * sizeof(int),
+                          ValData.data() + BYTES_FOR_SIZE);
+  return {"reqd_work_group_size", ValData, PI_PROPERTY_TYPE_BYTE_ARRAY};
+}
+
+inline void
+addDeviceRequirementsProps(PiPropertySet &Props,
+                           const std::vector<sycl::aspect> &Aspects,
+                           const std::vector<int> &ReqdWGSize = {}) {
+  PiArray<PiProperty> Value{makeAspectsProp(Aspects)};
+  if (!ReqdWGSize.empty())
+    Value.push_back(makeReqdWGSizeProp(ReqdWGSize));
   Props.insert(__SYCL_PI_PROPERTY_SET_SYCL_DEVICE_REQUIREMENTS,
                std::move(Value));
 }