[SYCL] Parallel-for range correction to improve group size selection by GPU driver

clang/lib/Sema/SemaSYCL.cpp

@@ -510,6 +510,22 @@ class MarkDeviceFunction : public RecursiveASTVisitor<MarkDeviceFunction> {
       FunctionDecl *FD = WorkList.back().first;
       FunctionDecl *ParentFD = WorkList.back().second;
 
+      // To implement rounding-up of a parallel-for range (Jira 20239)
+      // a kernel call is modified like this:
+      // auto Wrapper = [=](TransformedArgType Arg) {
+      //  if (Arg[0] >= NumWorkItems[0])
+      //    return;
+      //  Arg.set_allowed_range(NumWorkItems);
+      //  KernelFunc(Arg);
+      // };
+      //
+      // This transformation leads to a condition where a kernel body
+      // function becomes callable from a new kernel body function.
+      // Hence this test.
+      if ((ParentFD == KernelBody) && isSYCLKernelBodyFunction(FD)) {
+        KernelBody = FD;
+      }
+
       if ((ParentFD == SYCLKernel) && isSYCLKernelBodyFunction(FD)) {
         assert(!KernelBody && "inconsistent call graph - only one kernel body "
                               "function can be called");

sycl/include/CL/sycl/handler.hpp

@@ -728,23 +728,58 @@ class __SYCL_EXPORT handler {
   void parallel_for_lambda_impl(range<Dims> NumWorkItems,
                                 KernelType KernelFunc) {
     throwIfActionIsCreated();
-    using NameT =
-        typename detail::get_kernel_name_t<KernelName, KernelType>::name;
     using LambdaArgType = sycl::detail::lambda_arg_type<KernelType, item<Dims>>;
+
+    // If 1D kernel argument is an integral type, convert it to sycl::item<1>
     using TransformedArgType =
         typename std::conditional<std::is_integral<LambdaArgType>::value &&
                                       Dims == 1,
                                   item<Dims>, LambdaArgType>::type;
+    using NameT =
+        typename detail::get_kernel_name_t<KernelName, KernelType>::name;
+    constexpr size_t GoodLocalSizeX = 32;
+    std::string KName = typeid(NameT *).name();
+    bool DisableRounding =
+        KName.find("SYCL_OPT_PFWGS_DISABLE") != std::string::npos;
+    if (!DisableRounding && NumWorkItems[0] % GoodLocalSizeX != 0) {
+      // Not a multiple
+      size_t NewValX =
+          ((NumWorkItems[0] + GoodLocalSizeX - 1) / GoodLocalSizeX) *
+          GoodLocalSizeX;
+      if (getenv("SYCL_OPT_PFWGS_TRACE") != nullptr)
+        std::cerr << "***** Adjusted size from " << NumWorkItems[0] << " to "
+                  << NewValX << " *****\n";
+      auto Wrapper = [=](TransformedArgType Arg) {
+        if (Arg[0] >= NumWorkItems[0])
+          return;
+        Arg.set_allowed_range(NumWorkItems);
+        KernelFunc(Arg);
+      };
+
+      using NameWT = NameT *;
+      range<Dims> AdjustedRange = NumWorkItems;
+      AdjustedRange.set_range(NewValX);
 #ifdef __SYCL_DEVICE_ONLY__
-    (void)NumWorkItems;
-    kernel_parallel_for<NameT, TransformedArgType>(KernelFunc);
+      kernel_parallel_for<NameWT, TransformedArgType>(Wrapper);
 #else
-    detail::checkValueRange<Dims>(NumWorkItems);
-    MNDRDesc.set(std::move(NumWorkItems));
-    StoreLambda<NameT, KernelType, Dims, TransformedArgType>(
-        std::move(KernelFunc));
-    MCGType = detail::CG::KERNEL;
+      detail::checkValueRange<Dims>(AdjustedRange);
+      MNDRDesc.set(std::move(AdjustedRange));
+      StoreLambda<NameWT, decltype(Wrapper), Dims, TransformedArgType>(
+          std::move(Wrapper));
+      MCGType = detail::CG::KERNEL;
 #endif
+    } else {
+#ifdef __SYCL_DEVICE_ONLY__
+      (void)NumWorkItems;
+      kernel_parallel_for<NameT, TransformedArgType>(KernelFunc);
+#else
+      detail::checkValueRange<Dims>(NumWorkItems);
+      MNDRDesc.set(std::move(NumWorkItems));
+      StoreLambda<NameT, KernelType, Dims, TransformedArgType>(
+          std::move(KernelFunc));
+      MCGType = detail::CG::KERNEL;
+#endif
+    }
   }
 
   /// Defines and invokes a SYCL kernel function for the specified range.

sycl/include/CL/sycl/id.hpp

@@ -94,6 +94,8 @@ template <int dimensions = 1> class id : public detail::array<dimensions> {
     return result;
   }
 
+  void set_allowed_range(range<dimensions> rnwi) { (void)rnwi[0]; }
+
 #ifndef __SYCL_DISABLE_ID_TO_INT_CONV__
   /* Template operator is not allowed because it disables further type
    * conversion. For example, the next code will not work in case of template

sycl/include/CL/sycl/item.hpp

@@ -104,6 +104,8 @@ template <int dimensions = 1, bool with_offset = true> class item {
 
   bool operator!=(const item &rhs) const { return rhs.MImpl != MImpl; }
 
+  void set_allowed_range(const range<dimensions> rnwi) { MImpl.MExtent = rnwi; }
+
 protected:
   template <bool has_offset = with_offset>
   item(detail::enable_if_t<has_offset, const range<dimensions>> &extent,

sycl/include/CL/sycl/range.hpp

@@ -62,6 +62,9 @@ template <int dimensions = 1> class range : public detail::array<dimensions> {
     return size;
   }
 
+  // Adjust the first dim of the range
+  void set_range(const size_t dim0) { this->common_array[0] = dim0; }
+
   range(const range<dimensions> &rhs) = default;
   range(range<dimensions> &&rhs) = default;
   range<dimensions> &operator=(const range<dimensions> &rhs) = default;