Implemented dynamic_work_group_memory with lambdas

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

@@ -4846,6 +4846,10 @@ class SyclKernelIntHeaderCreator : public SyclKernelFieldHandler {
     } else if (SemaSYCL::isSyclType(FieldTy, SYCLTypeAttr::work_group_memory)) {
       addParam(FieldTy, SYCLIntegrationHeader::kind_work_group_memory,
                offsetOf(FD, FieldTy));
+    } else if (SemaSYCL::isSyclType(FieldTy,
+			            SYCLTypeAttr::dynamic_work_group_memory)) {
+      addParam(FieldTy, SYCLIntegrationHeader::kind_dynamic_work_group_memory,
+	       offsetOf(FD, FieldTy));
     } else if (SemaSYCL::isSyclType(FieldTy, SYCLTypeAttr::sampler) ||
                SemaSYCL::isSyclType(FieldTy, SYCLTypeAttr::annotated_ptr) ||
                SemaSYCL::isSyclType(FieldTy, SYCLTypeAttr::annotated_arg)) {
@@ -5993,6 +5997,7 @@ static const char *paramKind2Str(KernelParamKind K) {
     CASE(specialization_constants_buffer);
     CASE(pointer);
     CASE(work_group_memory);
+    CASE(dynamic_work_group_memory);
   }
   return "<ERROR>";
 

clang/test/SemaSYCL/Inputs/sycl/detail/kernel_desc.hpp

@@ -19,6 +19,7 @@ namespace detail {
     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-jit/common/include/Kernel.h

@@ -60,6 +60,7 @@ enum class ParameterKind : uint32_t {
   SpecConstBuffer = 4,
   Stream = 5,
   WorkGroupMemory = 6,
+  DynamicWorkGroupMemory = 7,
   Invalid = 0xF,
 };
 

sycl/include/sycl/detail/kernel_desc.hpp

@@ -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/ext/oneapi/experimental/graph.hpp

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#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 +18,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
@@ -501,6 +503,9 @@ class command_graph<graph_state::executable>
 namespace detail {
 class __SYCL_EXPORT dynamic_parameter_base {
 public:
+  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,
@@ -525,14 +530,86 @@ 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);
 };
+
+class dynamic_work_group_memory_base
+#ifndef __SYCL_DEVICE_ONLY__
+    : public dynamic_parameter_base
+#endif
+{
+public:
+  dynamic_work_group_memory_base() = default;
+  dynamic_work_group_memory_base(
+      experimental::command_graph<graph_state::modifiable> Graph, size_t Size)
+      :
+#ifndef __SYCL_DEVICE_ONLY__
+        dynamic_parameter_base(Graph),
+#endif
+        BufferSize(Size) {
+  }
+
+private:
+#ifdef __SYCL_DEVICE_ONLY__
+  [[maybe_unused]] unsigned char Padding[sizeof(dynamic_parameter_base)];
+#endif
+  size_t BufferSize{};
+  friend class sycl::handler;
+};
 } // namespace detail
 
+template <typename DataT,
+          typename = std::enable_if_t<detail::is_unbounded_array_v<DataT>>>
+class __SYCL_SPECIAL_CLASS
+__SYCL_TYPE(dynamic_work_group_memory) dynamic_work_group_memory
+    : public detail::dynamic_work_group_memory_base {
+private:
+  work_group_memory<DataT> WorkGroupMem;
+
+#ifdef __SYCL_DEVICE_ONLY__
+  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
+
+public:
+  /// 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)
+      : detail::dynamic_work_group_memory_base(
+            Graph, Num * sizeof(std::remove_extent_t<DataT>)) {}
+
+  /// Updates on the host this dynamic_work_group_memory and all registered
+  /// nodes with a new buffer size.
+  /// @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
+  }
+  work_group_memory<DataT> get() const { return WorkGroupMem; }
+
+  // Frontend requires special types to have a default constructor in order to
+  // have a uniform way of initializing an object of special type to then call
+  // the __init method on it. This is purely an implementation detail and not
+  // part of the spec.
+  // TODO: Revisit this once https://github.com/intel/llvm/issues/16061 is
+  // closed.
+  dynamic_work_group_memory() = default;
+};
+
 template <typename ValueT>
 class dynamic_parameter : public detail::dynamic_parameter_base {
   static constexpr bool IsAccessor =

sycl/include/sycl/ext/oneapi/experimental/work_group_memory.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>
@@ -115,6 +116,9 @@ 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, typename> friend class dynamic_work_group_memory;
+
   decoratedPtr ptr = nullptr;
 };
 } // namespace ext::oneapi::experimental

sycl/include/sycl/handler.hpp

@@ -150,6 +150,7 @@ class pipe;
 
 namespace ext ::oneapi ::experimental {
 template <typename, typename> class work_group_memory;
+template <typename, typename> class dynamic_work_group_memory;
 struct image_descriptor;
 __SYCL_EXPORT void async_free(sycl::handler &h, void *ptr);
 __SYCL_EXPORT void *async_malloc(sycl::handler &h, sycl::usm::alloc kind,
@@ -160,6 +161,8 @@ __SYCL_EXPORT void *async_malloc_from_pool(sycl::handler &h, size_t size,
 
 namespace ext::oneapi::experimental::detail {
 class graph_impl;
+class dynamic_work_group_memory_base;
+class dynamic_parameter_base;
 } // namespace ext::oneapi::experimental::detail
 namespace detail {
 
@@ -680,6 +683,12 @@ class __SYCL_EXPORT handler {
     registerDynamicParameter(DynamicParam, ArgIndex);
   }
 
+  // setArgHelper for graph dynamic_work_group_memory
+  void
+  setArgHelper(int ArgIndex,
+               ext::oneapi::experimental::detail::dynamic_work_group_memory_base
+                   &DynWorkGroupMemParam);
+
   // setArgHelper for the raw_kernel_arg extension type.
   void setArgHelper(int ArgIndex,
                     sycl::ext::oneapi::experimental::raw_kernel_arg &&Arg) {
@@ -1879,6 +1888,16 @@ class __SYCL_EXPORT handler {
     setArgHelper(argIndex, dynamicParam);
   }
 
+  // set_arg for graph dynamic_work_group_memory
+  template <typename DataT>
+  void set_arg(int argIndex,
+               ext::oneapi::experimental::dynamic_work_group_memory<DataT>
+                   &dynWorkGroupMem) {
+    ext::oneapi::experimental::detail::dynamic_work_group_memory_base
+        &dynWorkGroupMemImpl = dynWorkGroupMem;
+    setArgHelper(argIndex, dynWorkGroupMemImpl);
+  }
+
   // set_arg for the raw_kernel_arg extension type.
   void set_arg(int argIndex, ext::oneapi::experimental::raw_kernel_arg &&Arg) {
     setArgHelper(argIndex, std::move(Arg));
@@ -3771,7 +3790,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/source/detail/graph_impl.cpp

@@ -1994,6 +1994,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)
@@ -2014,6 +2019,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
@@ -2069,6 +2078,39 @@ 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->MCommandGroups[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_std_layout);
+    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.hpp

@@ -1479,6 +1479,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),
@@ -1546,6 +1549,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/handler_impl.hpp

@@ -199,6 +199,11 @@ class handler_impl {
   std::vector<std::shared_ptr<detail::work_group_memory_impl>>
       MWorkGroupMemoryObjects;
 
+  /// List of dynamic work group memory objects associated with this handler
+  std::vector<std::shared_ptr<
+      ext::oneapi::experimental::detail::dynamic_work_group_memory_base>>
+      MDynWorkGroupMemoryParams;
+
   /// Potential event mode for the result event of the command.
   ext::oneapi::experimental::event_mode_enum MEventMode =
       ext::oneapi::experimental::event_mode_enum::none;

sycl/source/detail/jit_compiler.cpp

@@ -180,6 +180,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/scheduler/commands.cpp

@@ -2319,6 +2319,8 @@ void SetArgBasedOnType(
     const ContextImplPtr &ContextImpl, 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: