Implemented dynamic_work_group_memory with lambdas

Author: Konrad Kusiak

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

@@ -488,7 +488,6 @@ 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);
@@ -526,37 +525,55 @@ class __SYCL_EXPORT dynamic_parameter_base {
   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 {};
+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) {
+  }
 
-template <typename T>
-inline constexpr bool is_unbounded_array_v = is_unbounded_array<T>::value;
+private:
+#ifdef __SYCL_DEVICE_ONLY__
+  [[maybe_unused]] char padding[sizeof(dynamic_parameter_base)];
+#endif
+  size_t BufferSize{};
+  friend class sycl::handler;
+};
+} // namespace detail
 
 template <typename DataT,
-          typename = std::enable_if_t<is_unbounded_array_v<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;
+
+  using value_type = std::remove_all_extents_t<DataT>;
+  using decoratedPtr = typename sycl::detail::DecoratedType<
+      value_type, access::address_space::local_space>::type *;
+
 #ifdef __SYCL_DEVICE_ONLY__
-    : detail::dynamic_parameter_base
-#else
-    : public detail::dynamic_parameter_base
+  void __init(decoratedPtr Ptr) { this->WorkGroupMem.__init(Ptr); }
 #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>);
-  }
+      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 this dynamic_work_group_memory and all registered nodes with a new
   /// number of elements.
@@ -567,19 +584,15 @@ __SYCL_TYPE(dynamic_work_group_memory) dynamic_work_group_memory
         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
+  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>

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

@@ -20,14 +20,6 @@ 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 {};
 
@@ -47,13 +39,12 @@ 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
 namespace ext::oneapi::experimental {
+// Forward decleration to be able to befriend dynamic_work_group_memory.
+// template <typename, typename> class dynamic_work_group_memory;
 
 struct indeterminate_t {};
 inline constexpr indeterminate_t indeterminate;
@@ -127,9 +118,8 @@ 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;
+
+  template <typename, typename> friend class dynamic_work_group_memory;
 
   decoratedPtr ptr = nullptr;
 };

sycl/include/sycl/handler.hpp

@@ -153,11 +153,14 @@ class pipe;
 
 namespace ext ::oneapi ::experimental {
 template <typename, typename> class work_group_memory;
+template <typename, typename> class dynamic_work_group_memory;
 struct image_descriptor;
 } // namespace ext::oneapi::experimental
 
 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 {
 

sycl/source/detail/graph_impl.cpp

@@ -2033,7 +2033,7 @@ void dynamic_parameter_impl::updateWorkGroupMem(size_t BufferSize) {
   for (auto &DynCGInfo : MDynCGs) {
     auto DynCG = DynCGInfo.DynCG.lock();
     if (DynCG) {
-      auto &CG = DynCG->MKernels[DynCGInfo.CGIndex];
+      auto &CG = DynCG->MCommandGroups[DynCGInfo.CGIndex];
       dynamic_parameter_impl::updateCGWorkGroupMem(CG, DynCGInfo.ArgIndex,
                                                    BufferSize);
     }
@@ -2048,8 +2048,7 @@ void dynamic_parameter_impl::updateCGWorkGroupMem(
     if (Arg.MIndex != ArgIndex) {
       continue;
     }
-    assert(Arg.MType ==
-           sycl::detail::kernel_param_kind_t::kind_dynamic_work_group_memory);
+    assert(Arg.MType == sycl::detail::kernel_param_kind_t::kind_std_layout);
     Arg.MSize = BufferSize;
     break;
   }

sycl/source/handler.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "sycl/detail/helpers.hpp"
+#include "sycl/ext/oneapi/experimental/graph.hpp"
 #include "ur_api.h"
 #include <algorithm>
 
@@ -1004,10 +1005,15 @@ void handler::processArg(void *Ptr, const detail::kernel_param_kind_t &Kind,
     auto *DynBase = static_cast<
         ext::oneapi::experimental::detail::dynamic_parameter_base *>(Ptr);
 
+    auto *DynWorkGroupBase = static_cast<
+        ext::oneapi::experimental::detail::dynamic_work_group_memory_base *>(
+        Ptr);
+
     registerDynamicParameter(*DynBase, Index + IndexShift);
 
-    Ptr = static_cast<void *>(++DynBase);
-    [[fallthrough]];
+    addArg(kernel_param_kind_t::kind_std_layout, nullptr,
+           DynWorkGroupBase->BufferSize, Index + IndexShift);
+    break;
   }
   case kernel_param_kind_t::kind_work_group_memory: {
     addArg(kernel_param_kind_t::kind_std_layout, nullptr,