[SYCL] Add max work-group size kernel properties (#14518)

This patch adds two kernel properties to allow users to specify the
maximum work-group size that a kernel will be invoked with.

The `max_work_group_size` property corresponds to the
`intel::max_work_group_size` function attribute, but can be specified
with 1, 2, or 3 dimensions (unlike the attribute which accepts only 3).

The `max_linear_work_group_size` property is similar but is always a
single value which denotes the combined linear (total) work-group size.
This can be used when the user cannot guarantee a maximum bound in each
of the dimensions they wish to run the kernel, but can guarantee a
total. This acts similarly to CUDA's `maxThreadsPerBlock` launch bounds
property.

This patch also wires up the 'max_work_group_size' property to the
equivalent SPIR-V execution mode, which should hopefully improve certain
use cases.

Author: frasercrmck

llvm/lib/SYCLLowerIR/CompileTimePropertiesPass.cpp

@@ -362,19 +362,24 @@ attributeToExecModeMetadata(const Attribute &Attr, Function &F) {
       AddFPControlMetadataForWidth(SPIRV_DENORM_PRESERVE, 64);
   }
 
-  if (AttrKindStr == "sycl-work-group-size" ||
-      AttrKindStr == "sycl-work-group-size-hint") {
+  static constexpr std::tuple<const char *, const char *> SimpleWGAttrs[] = {
+      {"sycl-work-group-size", "reqd_work_group_size"},
+      {"sycl-work-group-size-hint", "work_group_size_hint"},
+      {"sycl-max-work-group-size", "max_work_group_size"},
+  };
+
+  for (auto &[AttrKind, MDStr] : SimpleWGAttrs) {
+    if (AttrKindStr != AttrKind)
+      continue;
     // Split values in the comma-separated list integers.
-    SmallVector<StringRef, 3> ValStrs;
-    Attr.getValueAsString().split(ValStrs, ',');
+    SmallVector<StringRef, 3> AttrValStrs;
+    Attr.getValueAsString().split(AttrValStrs, ',');
 
-    size_t NumDims = ValStrs.size();
-    assert(NumDims <= 3 &&
-           "sycl-work-group-size and sycl-work-group-size-hint currently only "
-           "support up to three values");
+    size_t NumDims = AttrValStrs.size();
+    assert(NumDims <= 3 && "Incorrect number of values for kernel property");
 
     // SYCL work-group sizes must be reversed for SPIR-V.
-    std::reverse(ValStrs.begin(), ValStrs.end());
+    std::reverse(AttrValStrs.begin(), AttrValStrs.end());
 
     // Use integer pointer size as closest analogue to size_t.
     IntegerType *IntPtrTy = DLayout.getIntPtrType(Ctx);
@@ -383,7 +388,7 @@ attributeToExecModeMetadata(const Attribute &Attr, Function &F) {
 
     // Get the integers from the strings.
     SmallVector<Metadata *, 3> MDVals;
-    for (StringRef ValStr : ValStrs)
+    for (StringRef ValStr : AttrValStrs)
       MDVals.push_back(ConstantAsMetadata::get(
           Constant::getIntegerValue(SizeTTy, APInt(SizeTBitSize, ValStr, 10))));
     while (MDVals.size() < 3)
@@ -397,10 +402,7 @@ attributeToExecModeMetadata(const Attribute &Attr, Function &F) {
                                            Type::getInt32Ty(Ctx), NumDims))));
     }
 
-    const char *MDName = (AttrKindStr == "sycl-work-group-size")
-                             ? "reqd_work_group_size"
-                             : "work_group_size_hint";
-    return std::pair<std::string, MDNode *>(MDName, MDNode::get(Ctx, MDVals));
+    return std::pair<std::string, MDNode *>(MDStr, MDNode::get(Ctx, MDVals));
   }
 
   if (AttrKindStr == "sycl-sub-group-size") {
@@ -413,6 +415,21 @@ attributeToExecModeMetadata(const Attribute &Attr, Function &F) {
                                             MDNode::get(Ctx, MD));
   }
 
+  if (AttrKindStr == "sycl-max-linear-work-group-size") {
+    auto MaxLinearSize = getAttributeAsInteger<uint64_t>(Attr);
+    // Use integer pointer size as closest analogue to size_t.
+    IntegerType *IntPtrTy = DLayout.getIntPtrType(Ctx);
+    IntegerType *SizeTTy = Type::getIntNTy(Ctx, IntPtrTy->getBitWidth());
+    unsigned SizeTBitSize = SizeTTy->getBitWidth();
+
+    // Get the integers from the strings.
+    Metadata *MD = ConstantAsMetadata::get(Constant::getIntegerValue(
+        SizeTTy, APInt(SizeTBitSize, MaxLinearSize, 10)));
+
+    return std::pair<std::string, MDNode *>("max_linear_work_group_size",
+                                            MDNode::get(Ctx, MD));
+  }
+
   // The sycl-single-task attribute currently only has an effect when targeting
   // SPIR FPGAs, in which case it will generate a "max_global_work_dim" MD node
   // with a 0 value, similar to applying [[intel::max_global_work_dim(0)]] to

llvm/lib/SYCLLowerIR/ComputeModuleRuntimeInfo.cpp

@@ -112,28 +112,35 @@ std::vector<StringRef> getKernelNamesUsingAssert(const Module &M) {
   return SPIRKernelNames;
 }
 
-// Gets reqd_work_group_size information for function Func.
-std::vector<uint32_t> getKernelReqdWorkGroupSizeMetadata(const Function &Func) {
-  MDNode *ReqdWorkGroupSizeMD = Func.getMetadata("reqd_work_group_size");
-  if (!ReqdWorkGroupSizeMD)
+// Gets 1- to 3-dimension work-group related information for function Func.
+// Returns an empty vector if not present.
+template <typename T>
+std::vector<T> getKernelWorkGroupMetadata(const Function &Func,
+                                          const char *MDName) {
+  MDNode *WorkGroupMD = Func.getMetadata(MDName);
+  if (!WorkGroupMD)
     return {};
-  size_t NumOperands = ReqdWorkGroupSizeMD->getNumOperands();
+  size_t NumOperands = WorkGroupMD->getNumOperands();
   assert(NumOperands >= 1 && NumOperands <= 3 &&
-         "reqd_work_group_size does not have between 1 and 3 operands.");
-  std::vector<uint32_t> OutVals;
+         "work-group metadata does not have between 1 and 3 operands.");
+  std::vector<T> OutVals;
   OutVals.reserve(NumOperands);
-  for (const MDOperand &MDOp : ReqdWorkGroupSizeMD->operands())
+  for (const MDOperand &MDOp : WorkGroupMD->operands())
     OutVals.push_back(mdconst::extract<ConstantInt>(MDOp)->getZExtValue());
   return OutVals;
 }
-// Gets work_group_num_dim information for function Func, conviniently 0 if
-// metadata is not present.
-uint32_t getKernelWorkGroupNumDim(const Function &Func) {
-  MDNode *MaxDimMD = Func.getMetadata("work_group_num_dim");
-  if (!MaxDimMD)
-    return 0;
-  assert(MaxDimMD->getNumOperands() == 1 && "Malformed node.");
-  return mdconst::extract<ConstantInt>(MaxDimMD->getOperand(0))->getZExtValue();
+
+// Gets a single-dimensional piece of information for function Func.
+// Returns std::nullopt if metadata is not present.
+template <typename T>
+std::optional<T> getKernelSingleEltMetadata(const Function &Func,
+                                            const char *MDName) {
+  if (MDNode *MaxDimMD = Func.getMetadata(MDName)) {
+    assert(MaxDimMD->getNumOperands() == 1 && "Malformed node.");
+    return mdconst::extract<ConstantInt>(MaxDimMD->getOperand(0))
+        ->getZExtValue();
+  }
+  return std::nullopt;
 }
 
 PropSetRegTy computeModuleProperties(const Module &M,
@@ -249,22 +256,40 @@ PropSetRegTy computeModuleProperties(const Module &M,
   SmallVector<std::string, 4> MetadataNames;
 
   if (GlobProps.EmitProgramMetadata) {
-    // Add reqd_work_group_size and work_group_num_dim information to
-    // program metadata.
+    // Add various pieces of function metadata to program metadata.
     for (const Function &Func : M.functions()) {
-      std::vector<uint32_t> KernelReqdWorkGroupSize =
-          getKernelReqdWorkGroupSizeMetadata(Func);
-      if (!KernelReqdWorkGroupSize.empty()) {
+      // Note - we're implicitly truncating 64-bit work-group data to 32 bits in
+      // all work-group related metadata. All current consumers of this program
+      // metadata format only support SYCL ID queries that fit within MAX_INT.
+      if (auto KernelReqdWorkGroupSize = getKernelWorkGroupMetadata<uint32_t>(
+              Func, "reqd_work_group_size");
+          !KernelReqdWorkGroupSize.empty()) {
         MetadataNames.push_back(Func.getName().str() + "@reqd_work_group_size");
         PropSet.add(PropSetRegTy::SYCL_PROGRAM_METADATA, MetadataNames.back(),
                     KernelReqdWorkGroupSize);
       }
 
-      uint32_t WorkGroupNumDim = getKernelWorkGroupNumDim(Func);
-      if (WorkGroupNumDim) {
+      if (auto WorkGroupNumDim = getKernelSingleEltMetadata<uint32_t>(
+              Func, "work_group_num_dim")) {
         MetadataNames.push_back(Func.getName().str() + "@work_group_num_dim");
         PropSet.add(PropSetRegTy::SYCL_PROGRAM_METADATA, MetadataNames.back(),
-                    WorkGroupNumDim);
+                    *WorkGroupNumDim);
+      }
+
+      if (auto KernelMaxWorkGroupSize =
+              getKernelWorkGroupMetadata<uint32_t>(Func, "max_work_group_size");
+          !KernelMaxWorkGroupSize.empty()) {
+        MetadataNames.push_back(Func.getName().str() + "@max_work_group_size");
+        PropSet.add(PropSetRegTy::SYCL_PROGRAM_METADATA, MetadataNames.back(),
+                    KernelMaxWorkGroupSize);
+      }
+
+      if (auto MaxLinearWGSize = getKernelSingleEltMetadata<uint64_t>(
+              Func, "max_linear_work_group_size")) {
+        MetadataNames.push_back(Func.getName().str() +
+                                "@max_linear_work_group_size");
+        PropSet.add(PropSetRegTy::SYCL_PROGRAM_METADATA, MetadataNames.back(),
+                    *MaxLinearWGSize);
       }
     }
 

llvm/lib/SYCLLowerIR/SYCLCreateNVVMAnnotations.cpp

@@ -97,7 +97,8 @@ SYCLCreateNVVMAnnotationsPass::run(Module &M, ModuleAnalysisManager &MAM) {
 
     constexpr static std::pair<const char *, const char *>
         SingleValAnnotations[] = {{"min_work_groups_per_cu", "minctasm"},
-                                  {"max_work_groups_per_mp", "maxclusterrank"}};
+                                  {"max_work_groups_per_mp", "maxclusterrank"},
+                                  {"max_linear_work_group_size", "maxntidx"}};
 
     for (auto &[MDName, AnnotationName] : SingleValAnnotations) {
       if (MDNode *Node = F.getMetadata(MDName)) {

sycl/cmake/modules/FetchUnifiedRuntime.cmake

@@ -117,13 +117,13 @@ if(SYCL_UR_USE_FETCH_CONTENT)
   endfunction()
 
   set(UNIFIED_RUNTIME_REPO "https://github.com/oneapi-src/unified-runtime.git")
-  # commit 45ad7c52a75e6d3e52f658e38e796563744914c7
-  # Merge: 7a2caca5 3bf2becb
+  # commit 7ecf64d60c31cd72bd88588498536d067bad59d6
+  # Merge: 17aa04d3 6eb5208b
   # Author: aarongreig <[email protected]>
-  # Date:   Tue Sep 24 08:04:54 2024 -0700
-  #     Merge pull request #2116 from RossBrunton/ross/morewarn
-  #     More warning squishing
-  set(UNIFIED_RUNTIME_TAG 45ad7c52a75e6d3e52f658e38e796563744914c7)
+  # Date:   Wed Sep 25 11:14:47 2024 +0100
+  #     Merge pull request #1996 from frasercrmck/ur-max-wg-size-props
+  #     Add two new properties to ur_kernel_group_info_t
+  set(UNIFIED_RUNTIME_TAG 7ecf64d60c31cd72bd88588498536d067bad59d6)
 
   set(UMF_BUILD_EXAMPLES OFF CACHE INTERNAL "EXAMPLES")
   # Due to the use of dependentloadflag and no installer for UMF and hwloc we need

sycl/doc/extensions/experimental/sycl_ext_oneapi_kernel_properties.asciidoc

@@ -58,7 +58,8 @@ Jessica Davies, Intel +
 Joe Garvey, Intel +
 Greg Lueck, Intel +
 John Pennycook, Intel +
-Roland Schulz, Intel
+Roland Schulz, Intel +
+Fraser Cormack, Codeplay
 
 == Overview
 
@@ -232,6 +233,68 @@ SYCL implementations may introduce additional kernel properties. If any
 combinations of kernel attributes are invalid, this must be clearly documented
 as part of the new kernel property definition.
 
+=== Kernel Properties for the CUDA backend
+
+The kernel properties specified in this section may only be used to decorate
+kernels that are submitted to the CUDA backend. Attempting to submit a kernel
+with these properties to another backend has undefined behavior.
+
+```c++
+namespace sycl {
+namespace ext {
+namespace oneapi {
+namespace experimental {
+
+struct max_work_group_size_key {
+  template <size_t... Dims>
+  using value_t = property_value<max_work_group_size_key, std::integral_constant<size_t, Dims>...>;
+}; // max_work_group_size_key
+
+struct max_linear_work_group_size_key {
+  template <size_t Size>
+  using value_t = property_value<max_linear_work_group_size_key, std::integral_constant<size_t, Size>>;
+}; // max_linear_work_group_size_key
+
+template <size_t... Dims>
+inline constexpr max_work_group_size_key::value_t<Dims...> max_work_group_size;
+
+template <size_t Size>
+inline constexpr max_linear_work_group_size_key::value_t<Size> max_linear_work_group_size;
+
+template <> struct is_property_key<max_work_group_size_key> : std::true_type {};
+template <> struct is_property_key<max_linear_work_group_size_key> : std::true_type {};
+
+} // namespace experimental
+} // namespace oneapi
+} // namespace ext
+} // namespace sycl
+```
+
+|===
+|Property|Description
+
+|`max_work_group_size`
+|The `max_work_group_size` property provides a promise to the compiler
+that the kernel will never be launched with a larger work-group than the
+specified size. The number of template arguments in the `Dims` parameter pack
+must match the dimensionality of the work-group used to invoke the kernel. The
+order of the template arguments matches the constructor of the `range` class.
+
+If the kernel is submitted with an `nd_range` that exceeds the size specified
+by the property, the implementation must throw a synchronous exception with the
+`errc::nd_range` error code.
+
+|`max_linear_work_group_size`
+|The `max_linear_work_group_size` property provides a promise to the compiler
+that the kernel will never be launched with a work-group for which the return
+value of `group::get_local_linear_range()` exceeds the specified amount.
+
+If the kernel is submitted with an `nd_range` that exceeds the size specified
+by the property, the implementation must throw a synchronous exception with the
+`errc::nd_range` error code.
+
+|===
+
 === Adding a Property List to a Kernel Launch
 
 To enable properties to be associated with kernels, this extension adds

sycl/include/sycl/ext/oneapi/kernel_properties/properties.hpp

@@ -71,6 +71,21 @@ struct single_task_kernel_key {
   using value_t = property_value<single_task_kernel_key>;
 };
 
+struct max_work_group_size_key
+    : detail::compile_time_property_key<detail::PropKind::MaxWorkGroupSize> {
+  template <size_t... Dims>
+  using value_t = property_value<max_work_group_size_key,
+                                 std::integral_constant<size_t, Dims>...>;
+};
+
+struct max_linear_work_group_size_key
+    : detail::compile_time_property_key<
+          detail::PropKind::MaxLinearWorkGroupSize> {
+  template <size_t Size>
+  using value_t = property_value<max_linear_work_group_size_key,
+                                 std::integral_constant<size_t, Size>>;
+};
+
 template <size_t Dim0, size_t... Dims>
 struct property_value<work_group_size_key, std::integral_constant<size_t, Dim0>,
                       std::integral_constant<size_t, Dims>...> {
@@ -138,6 +153,28 @@ template <> struct property_value<single_task_kernel_key> {
   using key_t = single_task_kernel_key;
 };
 
+template <size_t Dim0, size_t... Dims>
+struct property_value<max_work_group_size_key,
+                      std::integral_constant<size_t, Dim0>,
+                      std::integral_constant<size_t, Dims>...> {
+  static_assert(sizeof...(Dims) + 1 <= 3,
+                "max_work_group_size property currently "
+                "only supports up to three values.");
+  static_assert(
+      detail::AllNonZero<Dim0, Dims...>::value,
+      "max_work_group_size property must only contain non-zero values.");
+
+  using key_t = max_work_group_size_key;
+
+  constexpr size_t operator[](int Dim) const {
+    return std::array<size_t, sizeof...(Dims) + 1>{Dim0, Dims...}[Dim];
+  }
+};
+
+template <> struct property_value<max_linear_work_group_size_key> {
+  using key_t = max_linear_work_group_size_key;
+};
+
 template <size_t Dim0, size_t... Dims>
 inline constexpr work_group_size_key::value_t<Dim0, Dims...> work_group_size;
 
@@ -156,6 +193,14 @@ inline constexpr nd_range_kernel_key::value_t<Dims> nd_range_kernel;
 
 inline constexpr single_task_kernel_key::value_t single_task_kernel;
 
+template <size_t Dim0, size_t... Dims>
+inline constexpr max_work_group_size_key::value_t<Dim0, Dims...>
+    max_work_group_size;
+
+template <size_t Size>
+inline constexpr max_linear_work_group_size_key::value_t<Size>
+    max_linear_work_group_size;
+
 struct work_group_progress_key
     : detail::compile_time_property_key<detail::PropKind::WorkGroupProgress> {
   template <forward_progress_guarantee Guarantee,
@@ -283,6 +328,16 @@ template <> struct PropertyMetaInfo<single_task_kernel_key::value_t> {
   static constexpr const char *name = "sycl-single-task-kernel";
   static constexpr int value = 0;
 };
+template <size_t Dim0, size_t... Dims>
+struct PropertyMetaInfo<max_work_group_size_key::value_t<Dim0, Dims...>> {
+  static constexpr const char *name = "sycl-max-work-group-size";
+  static constexpr const char *value = SizeListToStr<Dim0, Dims...>::value;
+};
+template <size_t Size>
+struct PropertyMetaInfo<max_linear_work_group_size_key::value_t<Size>> {
+  static constexpr const char *name = "sycl-max-linear-work-group-size";
+  static constexpr size_t value = Size;
+};
 
 template <typename T, typename = void>
 struct HasKernelPropertiesGetMethod : std::false_type {};

sycl/include/sycl/ext/oneapi/properties/property.hpp

@@ -212,8 +212,10 @@ enum PropKind : uint32_t {
   Balanced = 71,
   InvocationCapacity = 72,
   ResponseCapacity = 73,
+  MaxWorkGroupSize = 74,
+  MaxLinearWorkGroupSize = 75,
   // PropKindSize must always be the last value.
-  PropKindSize = 74,
+  PropKindSize = 76,
 };
 
 struct property_key_base_tag {};