Implemented dynamic_work_groups, with a hack and few issues

Author: konradkusiak97

clang/include/clang/Basic/Attr.td

@@ -1606,12 +1606,12 @@ def SYCLType: InheritableAttr {
   let Subjects = SubjectList<[CXXRecord, Enum], ErrorDiag>;
   let LangOpts = [SYCLIsDevice, SilentlyIgnoreSYCLIsHost];
   let Args = [EnumArgument<"Type", "SYCLType", /*is_string=*/true,
-                           ["accessor", "local_accessor", "work_group_memory",
+                           ["accessor", "local_accessor", "work_group_memory", "dynamic_work_group_memory",
                             "specialization_id", "kernel_handler", "buffer_location",
                             "no_alias", "accessor_property_list", "group",
                             "private_memory", "aspect", "annotated_ptr", "annotated_arg",
                             "stream", "sampler", "host_pipe", "multi_ptr"],
-                           ["accessor", "local_accessor", "work_group_memory",
+                           ["accessor", "local_accessor", "work_group_memory", "dynamic_work_group_memory",
                             "specialization_id", "kernel_handler", "buffer_location",
                             "no_alias", "accessor_property_list", "group",
                             "private_memory", "aspect", "annotated_ptr", "annotated_arg",

clang/include/clang/Sema/SemaSYCL.h

@@ -63,7 +63,8 @@ class SYCLIntegrationHeader {
     kind_specialization_constants_buffer,
     kind_stream,
     kind_work_group_memory,
-    kind_last = kind_work_group_memory
+    kind_dynamic_work_group_memory,
+    kind_last = kind_dynamic_work_group_memory
   };
 
 public:
@@ -666,7 +667,7 @@ class SemaSYCL : public SemaBase {
   // Used to check whether the function represented by FD is a SYCL
   // free function kernel or not.
   bool isFreeFunction(const FunctionDecl *FD);
-  
+
   StmtResult BuildSYCLKernelCallStmt(FunctionDecl *FD, CompoundStmt *Body);
 };
 

clang/lib/Sema/SemaSYCL.cpp

@@ -60,6 +60,7 @@ enum class ParameterKind : uint32_t {
   SpecConstBuffer = 4,
   Stream = 5,
   WorkGroupMemory = 6,
+  DynamicWorkGroupMemory = 7,
   Invalid = 0xF,
 };
 
@@ -239,8 +240,8 @@ class NDRange {
   NDRange(int Dimensions, const Indices &GlobalSize,
           const Indices &LocalSize = {1, 1, 1},
           const Indices &Offset = {0, 0, 0})
-      : Dimensions{Dimensions},
-        GlobalSize{GlobalSize}, LocalSize{LocalSize}, Offset{Offset} {
+      : Dimensions{Dimensions}, GlobalSize{GlobalSize}, LocalSize{LocalSize},
+        Offset{Offset} {
 #ifndef NDEBUG
     const auto CheckDim = [Dimensions](const Indices &Range) {
       return std::all_of(Range.begin() + Dimensions, Range.end(),

sycl-jit/common/include/Kernel.h

@@ -59,6 +59,7 @@ enum class kernel_param_kind_t {
   kind_specialization_constants_buffer = 4,
   kind_stream = 5,
   kind_work_group_memory = 6,
+  kind_dynamic_work_group_memory = 7,
   kind_invalid = 0xf, // not a valid kernel kind
 };
 

sycl/include/sycl/detail/kernel_desc.hpp

@@ -8,6 +8,8 @@
 
 #pragma once
 
+#include "sycl/ext/oneapi/experimental/graph.hpp"
+#include <cstddef>
 #include <sycl/accessor.hpp>               // for detail::AccessorBaseHost
 #include <sycl/context.hpp>                // for context
 #include <sycl/detail/export.hpp>          // for __SYCL_EXPORT
@@ -17,8 +19,9 @@
 #ifdef __INTEL_PREVIEW_BREAKING_CHANGES
 #include <sycl/detail/string_view.hpp>
 #endif
-#include <sycl/device.hpp>                     // for device
+#include <sycl/device.hpp> // for device
 #include <sycl/ext/oneapi/experimental/detail/properties/graph_properties.hpp> // for graph properties classes
+#include <sycl/ext/oneapi/experimental/work_group_memory.hpp> // for dynamic_work_group_memory
 #include <sycl/nd_range.hpp>                   // for range, nd_range
 #include <sycl/properties/property_traits.hpp> // for is_property, is_property_of
 #include <sycl/property_list.hpp>              // for property_list
@@ -485,6 +488,10 @@ class command_graph<graph_state::executable>
 namespace detail {
 class __SYCL_EXPORT dynamic_parameter_base {
 public:
+  dynamic_parameter_base() = default;
+  dynamic_parameter_base(
+      sycl::ext::oneapi::experimental::command_graph<graph_state::modifiable>
+          Graph);
   dynamic_parameter_base(
       sycl::ext::oneapi::experimental::command_graph<graph_state::modifiable>
           Graph,
@@ -509,14 +516,72 @@ class __SYCL_EXPORT dynamic_parameter_base {
   void updateValue(const raw_kernel_arg *NewRawValue, size_t Size);
 
   void updateAccessor(const sycl::detail::AccessorBaseHost *Acc);
+
+  void updateWorkGroupMem(size_t BufferSize);
+
   std::shared_ptr<dynamic_parameter_impl> impl;
 
   template <class Obj>
   friend const decltype(Obj::impl) &
   sycl::detail::getSyclObjImpl(const Obj &SyclObject);
 };
+
 } // namespace detail
 
+template <typename T> struct is_unbounded_array : std::false_type {};
+
+template <typename T> struct is_unbounded_array<T[]> : std::true_type {};
+
+template <typename T>
+inline constexpr bool is_unbounded_array_v = is_unbounded_array<T>::value;
+
+template <typename DataT,
+          typename = std::enable_if_t<is_unbounded_array_v<DataT>>>
+
+class __SYCL_SPECIAL_CLASS
+__SYCL_TYPE(dynamic_work_group_memory) dynamic_work_group_memory
+#ifdef __SYCL_DEVICE_ONLY__
+    : detail::dynamic_parameter_base
+#else
+    : public detail::dynamic_parameter_base
+#endif
+{
+public:
+  dynamic_work_group_memory() = default;
+  /// Constructs a new dynamic_work_group_memory object.
+  /// @param Graph The graph associated with this object.
+  /// @param Num Number of elements in the unbounded array DataT.
+  dynamic_work_group_memory(
+      experimental::command_graph<graph_state::modifiable> Graph, size_t Num) {
+    auto &WorkGroupMemImpl =
+        static_cast<detail::work_group_memory_impl &>(WorkGroupMem);
+    WorkGroupMemImpl.buffer_size = Num * sizeof(std::remove_extent_t<DataT>);
+  }
+
+  /// Updates this dynamic_work_group_memory and all registered nodes with a new
+  /// number of elements.
+  /// @param Num The new number of elements in the unbounded array.
+  void update(size_t Num) {
+#ifndef __SYCL_DEVICE_ONLY__
+    detail::dynamic_parameter_base::updateWorkGroupMem(
+        Num * sizeof(std::remove_extent_t<DataT>));
+#endif
+  }
+
+  const work_group_memory<DataT> &get() const { return WorkGroupMem; }
+
+private:
+  work_group_memory<DataT> WorkGroupMem;
+#ifdef __SYCL_DEVICE_ONLY__
+  // [[maybe_unused]] char padding[sizeof(detail::dynamic_parameter_base)];
+  using value_type = std::remove_all_extents_t<DataT>;
+  using decoratedPtr = typename sycl::detail::DecoratedType<
+      value_type, access::address_space::local_space>::type *;
+
+  void __init(decoratedPtr Ptr) { this->WorkGroupMem.__init(Ptr); }
+#endif
+};
+
 template <typename ValueT>
 class dynamic_parameter : public detail::dynamic_parameter_base {
   static constexpr bool IsAccessor =

sycl/include/sycl/ext/oneapi/experimental/graph.hpp

@@ -7,6 +7,7 @@
 
 #pragma once
 
+#include "sycl/ext/oneapi/experimental/graph.hpp"
 #include <sycl/access/access.hpp>
 #include <sycl/detail/defines.hpp>
 #include <sycl/ext/oneapi/properties/properties.hpp>
@@ -19,6 +20,14 @@ namespace sycl {
 inline namespace _V1 {
 class handler;
 
+namespace ext {
+namespace oneapi {
+namespace experimental {
+template <typename DataT, typename Enable> class dynamic_work_group_memory;
+}
+} // namespace oneapi
+} // namespace ext
+
 namespace detail {
 template <typename T> struct is_unbounded_array : std::false_type {};
 
@@ -38,6 +47,9 @@ class work_group_memory_impl {
 private:
   size_t buffer_size;
   friend class sycl::handler;
+
+  template <typename DataT, typename Enable>
+  friend class sycl::ext::oneapi::experimental::dynamic_work_group_memory;
 };
 
 } // namespace detail
@@ -115,6 +127,10 @@ class __SYCL_SPECIAL_CLASS __SYCL_TYPE(work_group_memory) work_group_memory
   friend class sycl::handler; // needed in order for handler class to be aware
                               // of the private inheritance with
                               // work_group_memory_impl as base class
+                              //
+  template <typename T, typename Enable>
+  friend class sycl::ext::oneapi::experimental::dynamic_work_group_memory;
+
   decoratedPtr ptr = nullptr;
 };
 } // namespace ext::oneapi::experimental

sycl/include/sycl/ext/oneapi/experimental/work_group_memory.hpp

@@ -152,8 +152,7 @@ class pipe;
 }
 
 namespace ext ::oneapi ::experimental {
-template <typename, typename>
-class work_group_memory;
+template <typename, typename> class work_group_memory;
 struct image_descriptor;
 } // namespace ext::oneapi::experimental
 
@@ -514,7 +513,8 @@ class __SYCL_EXPORT handler {
 
   /// Saves the location of user's code passed in \p CodeLoc for future usage in
   /// finalize() method.
-  /// TODO: remove the first version of this func (the one without the IsTopCodeLoc arg)
+  /// TODO: remove the first version of this func (the one without the
+  /// IsTopCodeLoc arg)
   ///   at the next ABI breaking window since removing it breaks ABI on windows.
   void saveCodeLoc(detail::code_location CodeLoc);
   void saveCodeLoc(detail::code_location CodeLoc, bool IsTopCodeLoc);
@@ -724,8 +724,9 @@ class __SYCL_EXPORT handler {
         detail::KernelLambdaHasKernelHandlerArgT<KernelType,
                                                  LambdaArgType>::value;
 
-    MHostKernel = std::make_unique<
-        detail::HostKernel<KernelType, LambdaArgType, Dims>>(KernelFunc);
+    MHostKernel =
+        std::make_unique<detail::HostKernel<KernelType, LambdaArgType, Dims>>(
+            KernelFunc);
 
     constexpr bool KernelHasName =
         detail::getKernelName<KernelName>() != nullptr &&
@@ -3769,7 +3770,8 @@ class __SYCL_EXPORT handler {
             "A local accessor must not be used in a SYCL kernel function "
             "that is invoked via single_task or via the simple form of "
             "parallel_for that takes a range parameter.");
-      if (Kind == detail::kernel_param_kind_t::kind_work_group_memory)
+      if (Kind == detail::kernel_param_kind_t::kind_work_group_memory ||
+          Kind == detail::kernel_param_kind_t::kind_dynamic_work_group_memory)
         throw sycl::exception(
             make_error_code(errc::kernel_argument),
             "A work group memory object must not be used in a SYCL kernel "

sycl/include/sycl/handler.hpp

@@ -1937,6 +1937,11 @@ void executable_command_graph::update(const std::vector<node> &Nodes) {
   impl->update(NodeImpls);
 }
 
+dynamic_parameter_base::dynamic_parameter_base(
+    command_graph<graph_state::modifiable> Graph)
+    : impl(std::make_shared<dynamic_parameter_impl>(
+          sycl::detail::getSyclObjImpl(Graph))) {}
+
 dynamic_parameter_base::dynamic_parameter_base(
     command_graph<graph_state::modifiable> Graph, size_t ParamSize,
     const void *Data)
@@ -1957,6 +1962,10 @@ void dynamic_parameter_base::updateAccessor(
   impl->updateAccessor(Acc);
 }
 
+void dynamic_parameter_base::updateWorkGroupMem(size_t BufferSize) {
+  impl->updateWorkGroupMem(BufferSize);
+}
+
 void dynamic_parameter_impl::updateValue(const raw_kernel_arg *NewRawValue,
                                          size_t Size) {
   // Number of bytes is taken from member of raw_kernel_arg object rather
@@ -2012,6 +2021,40 @@ void dynamic_parameter_impl::updateAccessor(
               sizeof(sycl::detail::AccessorBaseHost));
 }
 
+void dynamic_parameter_impl::updateWorkGroupMem(size_t BufferSize) {
+  for (auto &[NodeWeak, ArgIndex] : MNodes) {
+    auto NodeShared = NodeWeak.lock();
+    if (NodeShared) {
+      dynamic_parameter_impl::updateCGWorkGroupMem(NodeShared->MCommandGroup,
+                                                   ArgIndex, BufferSize);
+    }
+  }
+
+  for (auto &DynCGInfo : MDynCGs) {
+    auto DynCG = DynCGInfo.DynCG.lock();
+    if (DynCG) {
+      auto &CG = DynCG->MKernels[DynCGInfo.CGIndex];
+      dynamic_parameter_impl::updateCGWorkGroupMem(CG, DynCGInfo.ArgIndex,
+                                                   BufferSize);
+    }
+  }
+}
+
+void dynamic_parameter_impl::updateCGWorkGroupMem(
+    std::shared_ptr<sycl::detail::CG> CG, int ArgIndex, size_t BufferSize) {
+
+  auto &Args = static_cast<sycl::detail::CGExecKernel *>(CG.get())->MArgs;
+  for (auto &Arg : Args) {
+    if (Arg.MIndex != ArgIndex) {
+      continue;
+    }
+    assert(Arg.MType ==
+           sycl::detail::kernel_param_kind_t::kind_dynamic_work_group_memory);
+    Arg.MSize = BufferSize;
+    break;
+  }
+}
+
 void dynamic_parameter_impl::updateCGArgValue(
     std::shared_ptr<sycl::detail::CG> CG, int ArgIndex, const void *NewValue,
     size_t Size) {

sycl/source/detail/graph_impl.cpp

@@ -1430,6 +1430,9 @@ class exec_graph_impl {
 
 class dynamic_parameter_impl {
 public:
+  dynamic_parameter_impl(std::shared_ptr<graph_impl> GraphImpl)
+      : MGraph(GraphImpl) {}
+
   dynamic_parameter_impl(std::shared_ptr<graph_impl> GraphImpl,
                          size_t ParamSize, const void *Data)
       : MGraph(GraphImpl), MValueStorage(ParamSize),
@@ -1497,6 +1500,22 @@ class dynamic_parameter_impl {
   /// @param Acc The new accessor value
   void updateAccessor(const sycl::detail::AccessorBaseHost *Acc);
 
+  /// Update the internal value of this dynamic parameter as well as the value
+  /// of this parameter in all registered nodes and dynamic CGs. Should only be
+  /// called for dynamic_work_group_memory arguments parameter.
+  /// @param BufferSize The total size in bytes of the new work_group_memory
+  /// array
+  void updateWorkGroupMem(size_t BufferSize);
+
+  /// Static helper function for updating command-group
+  /// dynamic_work_group_memory arguments.
+  /// @param CG The command-group to update the argument information for.
+  /// @param ArgIndex The argument index to update.
+  /// @param BufferSize The total size in bytes of the new work_group_memory
+  /// array
+  static void updateCGWorkGroupMem(std::shared_ptr<sycl::detail::CG> CG,
+                                   int ArgIndex, size_t BufferSize);
+
   /// Static helper function for updating command-group value arguments.
   /// @param CG The command-group to update the argument information for.
   /// @param ArgIndex The argument index to update.

sycl/source/detail/graph_impl.hpp

@@ -170,6 +170,8 @@ translateArgType(kernel_param_kind_t Kind) {
     return PK::Stream;
   case kind::kind_work_group_memory:
     return PK::WorkGroupMemory;
+  case kind::kind_dynamic_work_group_memory:
+    return PK::DynamicWorkGroupMemory;
   case kind::kind_invalid:
     return PK::Invalid;
   }

sycl/source/detail/jit_compiler.cpp

@@ -2352,6 +2352,8 @@ void SetArgBasedOnType(
     const std::function<void *(Requirement *Req)> &getMemAllocationFunc,
     const sycl::context &Context, detail::ArgDesc &Arg, size_t NextTrueIndex) {
   switch (Arg.MType) {
+  case kernel_param_kind_t::kind_dynamic_work_group_memory:
+    break;
   case kernel_param_kind_t::kind_work_group_memory:
     break;
   case kernel_param_kind_t::kind_stream:

sycl/source/detail/scheduler/commands.cpp

@@ -999,6 +999,16 @@ void handler::processArg(void *Ptr, const detail::kernel_param_kind_t &Kind,
     }
     break;
   }
+  case kernel_param_kind_t::kind_dynamic_work_group_memory: {
+
+    auto *DynBase = static_cast<
+        ext::oneapi::experimental::detail::dynamic_parameter_base *>(Ptr);
+
+    registerDynamicParameter(*DynBase, Index + IndexShift);
+
+    Ptr = static_cast<void *>(++DynBase);
+    [[fallthrough]];
+  }
   case kernel_param_kind_t::kind_work_group_memory: {
     addArg(kernel_param_kind_t::kind_std_layout, nullptr,
            static_cast<detail::work_group_memory_impl *>(Ptr)->buffer_size,