[SYCL] Fix a bug when using no device split and reqd_work_group_size (#16236)

There was a bug (#13523) where a
kernel couldn't be launched when `-fsycl-device-code-split=off` was used
and multiple kernels with different required work group sizes were
present. This issue was fixed by ensuring that the required work group
size metadata is not attached to the device image when multiple required
work group sizes are detected in a single module.

However, there was a similar but related case that was not fixed by that
PR, which is now demonstrated in the new test
no-split-reqd-wg-size-2.cpp. This issue occurs when there is a single
kernel with a required work group size and another kernel without one.
In this case, the module doesn't contain multiple required work group
sizes, so the required work group size metadata is still attached. As a
result of the metadata being attached, the runtime cannot launch the
kernel without a required work group size.

This PR removes the logic of ensuring metadata is not attached when
there are multiple required work group sizes, and instead adds logic
that ensures the metadata is not attached when the split mode is
`SPLIT_NONE`. This covers the old cases from the previous PR and the new
case in this PR.

Author: jzc

llvm/include/llvm/Support/PropertySetIO.h

@@ -212,6 +212,9 @@ class PropertySetRegistry {
   static constexpr char SYCL_VIRTUAL_FUNCTIONS[] = "SYCL/virtual functions";
   static constexpr char SYCL_IMPLICIT_LOCAL_ARG[] = "SYCL/implicit local arg";
 
+  static constexpr char PROPERTY_REQD_WORK_GROUP_SIZE[] =
+      "reqd_work_group_size_uint64_t";
+
   /// Function for bulk addition of an entire property set in the given
   /// \p Category .
   template <typename MapTy> void add(StringRef Category, const MapTy &Props) {
@@ -231,6 +234,17 @@ class PropertySetRegistry {
     PropSet.insert({PropName, PropertyValue(PropVal)});
   }
 
+  void remove(StringRef Category, StringRef PropName) {
+    auto PropertySetIt = PropSetMap.find(Category);
+    if (PropertySetIt == PropSetMap.end())
+      return;
+    auto &PropertySet = PropertySetIt->second;
+    auto PropIt = PropertySet.find(PropName);
+    if (PropIt == PropertySet.end())
+      return;
+    PropertySet.erase(PropIt);
+  }
+
   /// Parses from the given \p Buf a property set registry.
   static Expected<std::unique_ptr<PropertySetRegistry>>
   read(const MemoryBuffer *Buf);

llvm/lib/SYCLLowerIR/SYCLDeviceRequirements.cpp

@@ -40,7 +40,6 @@ SYCLDeviceRequirements
 llvm::computeDeviceRequirements(const Module &M,
                                 const SetVector<Function *> &EntryPoints) {
   SYCLDeviceRequirements Reqs;
-  bool MultipleReqdWGSize = false;
   // Process all functions in the module
   for (const Function &F : M) {
     if (auto *MDN = F.getMetadata("sycl_used_aspects")) {
@@ -81,8 +80,6 @@ llvm::computeDeviceRequirements(const Module &M,
             ExtractUnsignedIntegerFromMDNodeOperand(MDN, I));
       if (!Reqs.ReqdWorkGroupSize.has_value())
         Reqs.ReqdWorkGroupSize = NewReqdWorkGroupSize;
-      if (Reqs.ReqdWorkGroupSize != NewReqdWorkGroupSize)
-        MultipleReqdWGSize = true;
     }
 
     if (auto *MDN = F.getMetadata("sycl_joint_matrix")) {
@@ -119,13 +116,6 @@ llvm::computeDeviceRequirements(const Module &M,
     }
   }
 
-  // Usually, we would only expect one ReqdWGSize, as the module passed to
-  // this function would be split according to that. However, when splitting
-  // is disabled, this cannot be guaranteed. In this case, we reset the value,
-  // which makes so that no value is reqd_work_group_size data is attached in
-  // in the device image.
-  if (MultipleReqdWGSize)
-    Reqs.ReqdWorkGroupSize.reset();
   return Reqs;
 }
 
@@ -152,7 +142,8 @@ std::map<StringRef, util::PropertyValue> SYCLDeviceRequirements::asMap() const {
   // reqd_work_group_size_uint64_t attribute. At the next ABI-breaking
   // window, this can be changed back to reqd_work_group_size.
   if (ReqdWorkGroupSize.has_value())
-    Requirements["reqd_work_group_size_uint64_t"] = *ReqdWorkGroupSize;
+    Requirements[util::PropertySetRegistry::PROPERTY_REQD_WORK_GROUP_SIZE] =
+        *ReqdWorkGroupSize;
 
   if (JointMatrix.has_value())
     Requirements["joint_matrix"] = *JointMatrix;

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

@@ -310,6 +310,15 @@ std::string saveModuleProperties(module_split::ModuleDesc &MD,
   auto PropSet =
       computeModuleProperties(MD.getModule(), MD.entries(), GlobProps);
 
+  // When the split mode is none, the required work group size will be added
+  // to the whole module, which will make the runtime unable to
+  // launch the other kernels in the module that have different
+  // required work group sizes or no required work group sizes. So we need to
+  // remove the required work group size metadata in this case.
+  if (SplitMode == module_split::SPLIT_NONE)
+    PropSet.remove(PropSetRegTy::SYCL_DEVICE_REQUIREMENTS,
+                   PropSetRegTy::PROPERTY_REQD_WORK_GROUP_SIZE);
+
   std::string NewSuff = Suff.str();
   if (!Target.empty()) {
     PropSet.add(PropSetRegTy::SYCL_DEVICE_REQUIREMENTS, "compile_target",

sycl/test-e2e/Regression/no-split-reqd-wg-size-2.cpp

@@ -0,0 +1,19 @@
+// This test checks that with -fsycl-device-code-split=off, kernels
+// with different reqd_work_group_size dimensions can be launched.
+
+// RUN: %{build} -fsycl-device-code-split=off -o %t.out
+// RUN: %{run} %t.out
+
+// UNSUPPORTED: hip
+
+#include <sycl/detail/core.hpp>
+
+using namespace sycl;
+
+int main(int argc, char **argv) {
+  queue q;
+  q.single_task([] {});
+  q.parallel_for(range<2>(24, 1),
+                 [=](auto) [[sycl::reqd_work_group_size(24, 1)]] {});
+  return 0;
+}

sycl/test-e2e/Regression/no-split-reqd-wg-size.cpp

@@ -1,7 +1,7 @@
 // This test checks that with -fsycl-device-code-split=off, kernels
 // with different reqd_work_group_size dimensions can be launched.
 
-// RUN: %{build} -fsycl -fsycl-device-code-split=off -o %t.out
+// RUN: %{build} -fsycl-device-code-split=off -o %t.out
 // RUN: %{run} %t.out
 
 // UNSUPPORTED: hip