Implement Dynamic Local Accessors

Author: fabiomestre

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

@@ -100,6 +100,8 @@ class node_impl;
 class graph_impl;
 class exec_graph_impl;
 class dynamic_parameter_impl;
+//template <typename DataT, int Dimensions>
+//class dynamic_local_accessor_impl;
 class dynamic_command_group_impl;
 } // namespace detail
 
@@ -484,6 +486,11 @@ 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,
@@ -498,6 +505,13 @@ class __SYCL_EXPORT dynamic_parameter_base {
   void updateValue(const raw_kernel_arg *NewRawValue, size_t Size);
 
   void updateAccessor(const sycl::detail::AccessorBaseHost *Acc);
+
+  sycl::detail::LocalAccessorImplPtr getLocalAccessor(handler* Handler);
+
+  void registerLocalAccessor(sycl::detail::LocalAccessorBaseHost* LocalAccBaseHost, handler* Handler);
+
+  void updateLocalAccessor(range<3> NewAllocationSize);
+
   std::shared_ptr<dynamic_parameter_impl> impl;
 
   template <class Obj>
@@ -535,6 +549,42 @@ class dynamic_parameter : public detail::dynamic_parameter_base {
   }
 };
 
+template <typename DataT, int Dimensions = 1>
+class dynamic_local_accessor : public detail::dynamic_parameter_base {
+public:
+  template <int Dims = Dimensions, typename = std::enable_if_t<(Dims > 0)>>
+  dynamic_local_accessor(command_graph<graph_state::modifiable> Graph,
+                         range<Dimensions> AllocationSize,
+                         const property_list &PropList = {})
+      : detail::dynamic_parameter_base(Graph), AllocationSize(AllocationSize) {
+    (void)PropList;
+  }
+
+  void update(range<Dimensions> NewAllocationSize) {
+    detail::dynamic_parameter_base::updateLocalAccessor(
+        ::sycl::detail::convertToArrayOfN<3, 1>(NewAllocationSize));
+  };
+
+  local_accessor<DataT, Dimensions> get(handler &CGH) {
+#ifndef __SYCL_DEVICE_ONLY__
+    ::sycl::detail::LocalAccessorImplPtr BaseLocalAcc = getLocalAccessor(&CGH);
+    if (BaseLocalAcc) {
+      return sycl::detail::createSyclObjFromImpl<local_accessor<DataT, Dimensions>>(BaseLocalAcc);
+    } else {
+      local_accessor<DataT, Dimensions> LocalAccessor(AllocationSize, CGH);
+      registerLocalAccessor(
+          static_cast<sycl::detail::LocalAccessorBaseHost *>(&LocalAccessor), &CGH);
+      return LocalAccessor;
+    }
+#else
+    return local_accessor<DataT, Dimensions>();
+#endif
+  };
+
+private:
+  range<Dimensions> AllocationSize;
+};
+
 /// Additional CTAD deduction guides.
 template <typename ValueT>
 dynamic_parameter(experimental::command_graph<graph_state::modifiable> Graph,

sycl/include/sycl/handler.hpp

@@ -647,6 +647,22 @@ class __SYCL_EXPORT handler {
     registerDynamicParameter(DynamicParam, ArgIndex);
   }
 
+  // setArgHelper for graph dynamic_local_accessors.
+  template <typename DataT, int Dims>
+  void
+  setArgHelper(int ArgIndex,
+               ext::oneapi::experimental::dynamic_local_accessor<DataT, Dims>
+                   &DynamicLocalAccessor) {
+#ifndef __SYCL_DEVICE_ONLY__
+    auto LocalAccessor = DynamicLocalAccessor.get(*this);
+    setArgHelper(ArgIndex, LocalAccessor);
+    registerDynamicParameter(DynamicLocalAccessor, ArgIndex);
+#else
+    (void)ArgIndex;
+    (void)DynamicLocalAccessor;
+#endif
+  }
+
   // setArgHelper for the raw_kernel_arg extension type.
   void setArgHelper(int ArgIndex,
                     sycl::ext::oneapi::experimental::raw_kernel_arg &&Arg) {
@@ -1838,6 +1854,13 @@ class __SYCL_EXPORT handler {
     setArgHelper(argIndex, dynamicParam);
   }
 
+  template <typename DataT, int Dims>
+  void set_arg(int argIndex,
+               ext::oneapi::experimental::dynamic_local_accessor<DataT, Dims>
+                   &DynamicLocalAccessor) {
+    setArgHelper(argIndex, DynamicLocalAccessor);
+  }
+
   // 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));

sycl/source/detail/graph_impl.cpp

@@ -1829,6 +1829,11 @@ dynamic_parameter_base::dynamic_parameter_base(
     : impl(std::make_shared<dynamic_parameter_impl>(
           sycl::detail::getSyclObjImpl(Graph), ParamSize, Data)) {}
 
+dynamic_parameter_base::dynamic_parameter_base(
+    command_graph<graph_state::modifiable> Graph)
+    : impl(std::make_shared<dynamic_parameter_impl>(
+          sycl::detail::getSyclObjImpl(Graph))) {}
+
 void dynamic_parameter_base::updateValue(const void *NewValue, size_t Size) {
   impl->updateValue(NewValue, Size);
 }
@@ -1843,6 +1848,20 @@ void dynamic_parameter_base::updateAccessor(
   impl->updateAccessor(Acc);
 }
 
+sycl::detail::LocalAccessorImplPtr
+dynamic_parameter_base::getLocalAccessor(handler *Handler) {
+  return impl->getLocalAccessor(Handler);
+}
+
+void dynamic_parameter_base::registerLocalAccessor(
+    sycl::detail::LocalAccessorBaseHost *LocalAccBaseHost, handler *Handler) {
+  impl->registerLocalAccessor(LocalAccBaseHost, Handler);
+}
+
+void dynamic_parameter_base::updateLocalAccessor(range<3> NewAllocationSize) {
+  impl->updateLocalAccessor(NewAllocationSize);
+}
+
 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
@@ -1898,6 +1917,53 @@ void dynamic_parameter_impl::updateAccessor(
               sizeof(sycl::detail::AccessorBaseHost));
 }
 
+sycl::detail::LocalAccessorImplPtr
+dynamic_parameter_impl::getLocalAccessor(handler *Handler) {
+  auto HandlerImpl = sycl::detail::getSyclObjImpl(*Handler);
+  auto FindLocalAcc = MHandlerToLocalAccMap.find(HandlerImpl);
+
+  if (FindLocalAcc != MHandlerToLocalAccMap.end()) {
+    auto LocalAccImpl = FindLocalAcc->second;
+    return LocalAccImpl;
+  }
+  return nullptr;
+}
+
+void dynamic_parameter_impl::registerLocalAccessor(
+    sycl::detail::LocalAccessorBaseHost *LocalAccBaseHost, handler *Handler) {
+
+  auto HandlerImpl = sycl::detail::getSyclObjImpl(*Handler);
+  auto LocalAccImpl = sycl::detail::getSyclObjImpl(*LocalAccBaseHost);
+
+  MHandlerToLocalAccMap.insert({HandlerImpl, LocalAccImpl});
+}
+
+void dynamic_parameter_impl::updateLocalAccessor(range<3> NewAllocationSize) {
+
+  for (auto &[NodeWeak, ArgIndex] : MNodes) {
+    auto NodeShared = NodeWeak.lock();
+    if (NodeShared) {
+      //  We can use the first local accessor in the map since the dimensions
+      //  and element type should be identical.
+      auto LocalAccessor = MHandlerToLocalAccMap.begin()->second;
+      dynamic_parameter_impl::updateCGLocalAccessor(
+          NodeShared->MCommandGroup, ArgIndex, NewAllocationSize,
+          LocalAccessor->MDims, LocalAccessor->MElemSize);
+    }
+  }
+
+  for (auto &DynCGInfo : MDynCGs) {
+    auto DynCG = DynCGInfo.DynCG.lock();
+    if (DynCG) {
+      auto &CG = DynCG->MKernels[DynCGInfo.CGIndex];
+      auto LocalAccessor = MHandlerToLocalAccMap.begin()->second;
+      dynamic_parameter_impl::updateCGLocalAccessor(
+          CG, DynCGInfo.ArgIndex, NewAllocationSize,
+          LocalAccessor->MDims, LocalAccessor->MElemSize);
+    }
+  }
+}
+
 void dynamic_parameter_impl::updateCGArgValue(
     std::shared_ptr<sycl::detail::CG> CG, int ArgIndex, const void *NewValue,
     size_t Size) {
@@ -1963,6 +2029,28 @@ void dynamic_parameter_impl::updateCGAccessor(
   }
 }
 
+void dynamic_parameter_impl::updateCGLocalAccessor(
+    std::shared_ptr<sycl::detail::CG> CG, int ArgIndex,
+    range<3> NewAllocationSize,
+    int Dims, int ElemSize) {
+  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);
+
+    int SizeInBytes = ElemSize;
+    for (int I = 0; I < Dims; ++I)
+      SizeInBytes *= NewAllocationSize[I];
+    SizeInBytes = std::max(SizeInBytes, 1);
+
+    Arg.MSize = SizeInBytes;
+    break;
+  }
+}
+
 dynamic_command_group_impl::dynamic_command_group_impl(
     const command_graph<graph_state::modifiable> &Graph)
     : MGraph{sycl::detail::getSyclObjImpl(Graph)}, MActiveCGF(0) {}
@@ -2084,6 +2172,7 @@ size_t dynamic_command_group::get_active_index() const {
 void dynamic_command_group::set_active_index(size_t Index) {
   return impl->setActiveIndex(Index);
 }
+
 } // namespace experimental
 } // namespace oneapi
 } // namespace ext

sycl/source/detail/graph_impl.hpp

@@ -1431,6 +1431,10 @@ class exec_graph_impl {
 
 class dynamic_parameter_impl {
 public:
+  /// Used for parameters that don't have data such as local_accessors.
+  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) {
@@ -1496,6 +1500,26 @@ class dynamic_parameter_impl {
   /// @param Acc The new accessor value
   void updateAccessor(const sycl::detail::AccessorBaseHost *Acc);
 
+  /// Updates the value of all local accessors in registered nodes and dynamic
+  /// CGs.
+  /// @param NewAllocationSize The new size for the update local accessors.
+  void updateLocalAccessor(range<3> NewAllocationSize);
+
+  /// Gets the implementation for the local accessor that is associated with
+  /// a specific handler.
+  /// @param The handler that the local accessor is associated with.
+  /// @return returns the impl object for the local accessor that is associated
+  /// with this handler. Or nullptr if no local accessor has been registered
+  /// for this handler.
+  sycl::detail::LocalAccessorImplPtr getLocalAccessor(handler *Handler);
+
+  /// Associates a local accessor with this dynamic local accessor for a
+  /// specific handler.
+  /// @param LocalAccBase the local accessor that needs to be registered.
+  /// @param Handler the handler that the LocalAccessor is associated with.
+  void registerLocalAccessor(sycl::detail::LocalAccessorBaseHost *LocalAccBase,
+                             handler *Handler);
+
   /// 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.
@@ -1512,13 +1536,29 @@ class dynamic_parameter_impl {
                                int ArgIndex,
                                const sycl::detail::AccessorBaseHost *Acc);
 
+  /// Static helper function for updating command-group local accessor
+  /// arguments.
+  /// @param CG The command-group to update the argument information for.
+  /// @param ArgIndex The argument index to update.
+  /// @param NewAllocationSize The new allocation size for the local accessor
+  /// argument.
+  /// @param Dims The dimensions of the local accessor argument.
+  /// @param ElemSize The size of each element in the local accessor.
+  static void updateCGLocalAccessor(std::shared_ptr<sycl::detail::CG> CG,
+                                    int ArgIndex, range<3> NewAllocationSize,
+                                    int Dims, int ElemSize);
+
   // Weak ptrs to node_impls which will be updated
   std::vector<std::pair<std::weak_ptr<node_impl>, int>> MNodes;
   // Dynamic command-groups which will be updated
   std::vector<DynamicCGInfo> MDynCGs;
 
   std::shared_ptr<graph_impl> MGraph;
   std::vector<std::byte> MValueStorage;
+
+  std::unordered_map<std::shared_ptr<sycl::detail::handler_impl>,
+                     sycl::detail::LocalAccessorImplPtr>
+      MHandlerToLocalAccMap;
 };
 
 class dynamic_command_group_impl

sycl/test-e2e/CMakeLists.txt

@@ -91,7 +91,5 @@ add_custom_target(check-sycl-e2e
   USES_TERMINAL
 )
 
-add_executable("local_memory_test" Graph/Explicit/compile_time_local_memory.cpp)
-
 add_subdirectory(External)
 add_subdirectory(ExtraTests)

sycl/test-e2e/Graph/Update/update_dynamic_local_accessor.cpp

@@ -0,0 +1,71 @@
+// RUN: %{build} -o %t.out
+// RUN: %{run} %t.out
+// Extra run to check for leaks in Level Zero using UR_L0_LEAKS_DEBUG
+// RUN: %if level_zero %{env SYCL_PI_LEVEL_ZERO_USE_IMMEDIATE_COMMANDLISTS=0 %{l0_leak_check} %{run} %t.out 2>&1 | FileCheck %s --implicit-check-not=LEAK %}
+// Extra run to check for immediate-command-list in Level Zero
+// RUN: %if level_zero %{env SYCL_PI_LEVEL_ZERO_USE_IMMEDIATE_COMMANDLISTS=1 %{l0_leak_check} %{run} %t.out 2>&1 | FileCheck %s --implicit-check-not=LEAK %}
+
+// Tests updating local accessor parameters.
+#include "../graph_common.hpp"
+
+int main() {
+  using T = int;
+
+  const size_t LocalMemSize = 128;
+
+  queue Queue{};
+
+  std::vector<T> HostDataBeforeUpdate(Size);
+  std::vector<T> HostDataAfterUpdate(Size);
+  std::iota(HostDataBeforeUpdate.begin(), HostDataBeforeUpdate.end(), 10);
+
+  T *PtrA = malloc_device<T>(Size, Queue);
+  Queue.copy(HostDataBeforeUpdate.data(), PtrA, Size);
+  Queue.wait_and_throw();
+
+  exp_ext::command_graph Graph{Queue.get_context(), Queue.get_device()};
+
+  exp_ext::dynamic_local_accessor<T, 1> DynLocalAccessor{Graph, LocalMemSize};
+
+  auto Node = Graph.add([&](handler &CGH) {
+    CGH.set_arg(0, DynLocalAccessor);
+    auto LocalMem = DynLocalAccessor.get(CGH);
+
+    CGH.parallel_for(nd_range({Size}, {LocalMemSize}), [=](nd_item<1> Item) {
+      LocalMem[Item.get_local_linear_id()] = Item.get_local_linear_id();
+      PtrA[Item.get_global_linear_id()] = LocalMem[Item.get_local_linear_id()];
+    });
+  });
+
+  auto GraphExec = Graph.finalize(exp_ext::property::graph::updatable{});
+
+  // Submit the graph before the update and save the results.
+  Queue.submit([&](handler &CGH) { CGH.ext_oneapi_graph(GraphExec); });
+  Queue.wait_and_throw();
+  Queue.copy(PtrA, HostDataBeforeUpdate.data(), Size);
+  Queue.wait_and_throw();
+
+  DynLocalAccessor.update(LocalMemSize * 2);
+  Node.update_nd_range(nd_range({Size}, {LocalMemSize * 2}));
+  GraphExec.update(Node);
+
+  // Submit the graph after the update and save the results.
+  Queue.submit([&](handler &CGH) { CGH.ext_oneapi_graph(GraphExec); });
+  Queue.wait_and_throw();
+  Queue.copy(PtrA, HostDataAfterUpdate.data(), Size);
+  Queue.wait_and_throw();
+
+  for (size_t i = 0; i < Size; i++) {
+    T Ref = i % LocalMemSize;
+    assert(check_value(i, Ref, HostDataBeforeUpdate[i], "PtrA Before Update"));
+  }
+
+  for (size_t i = 0; i < Size; i++) {
+    T Ref = i % (LocalMemSize * 2);
+    assert(check_value(i, Ref, HostDataAfterUpdate[i], "PtrA After Update"));
+  }
+
+  free(PtrA, Queue);
+
+  return 0;
+}