[SYCL][Graph] Implement Dynamic Local Accessors

sycl/cmake/modules/FetchUnifiedRuntime.cmake

@@ -116,7 +116,7 @@ if(SYCL_UR_USE_FETCH_CONTENT)
       CACHE PATH "Path to external '${name}' adapter source dir" FORCE)
   endfunction()
 
-  set(UNIFIED_RUNTIME_REPO "https://github.com/oneapi-src/unified-runtime.git")
+  set(UNIFIED_RUNTIME_REPO "https://github.com/Bensuo/unified-runtime.git")
   include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/UnifiedRuntimeTag.cmake)
 
   set(UMF_BUILD_EXAMPLES OFF CACHE INTERNAL "EXAMPLES")

sycl/cmake/modules/UnifiedRuntimeTag.cmake

@@ -1,7 +1,7 @@
-# commit 9e48f543b8dd39d45563169433bb529583625dfe
-# Merge: 6a3fece6 1a1108b3
+# commit 6d4eec8cdcfe8a5d359ed05092797c429c2ca878
+# Merge: 40d28e7bd84a 800b452d67c5
 # Author: Martin Grant <[email protected]>
-# Date:   Wed Jan 15 14:33:29 2025 +0000
-#     Merge pull request #2540 from martygrant/martin/program-info-unswitch
-#     Move urProgramGetInfo success test from a switch to individual tests.
-set(UNIFIED_RUNTIME_TAG 9e48f543b8dd39d45563169433bb529583625dfe)
+# Date:   Thu Dec 12 16:00:13 2024 +0000
+#     Merge pull request #2272 from martygrant/martin/virtual-memory-cts-spec-gap
+#     Improvements to align CTS and Spec for Virtual Memory
+set(UNIFIED_RUNTIME_TAG bf6b6f9df5cd7c1e3dda4af8e4b3546c7109f24f)

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

@@ -17,7 +17,7 @@
 #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/nd_range.hpp>                   // for range, nd_range
 #include <sycl/properties/property_traits.hpp> // for is_property, is_property_of
@@ -447,6 +447,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,
+      const property_list &PropList);
+
   dynamic_parameter_base(
       sycl::ext::oneapi::experimental::command_graph<graph_state::modifiable>
           Graph,
@@ -461,6 +466,15 @@ 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) const;
+
+  void
+  registerLocalAccessor(sycl::detail::LocalAccessorBaseHost *LocalAccBaseHost,
+                        handler *Handler);
+
+  void updateLocalAccessor(range<3> NewAllocationSize);
+
   std::shared_ptr<dynamic_parameter_impl> impl;
 
   template <class Obj>
@@ -498,6 +512,44 @@ 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, PropList),
+        AllocationSize(AllocationSize) {}
+
+  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
+    (void)CGH;
+    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) {
@@ -1834,14 +1850,21 @@ class __SYCL_EXPORT handler {
 
   // set_arg for graph dynamic_parameters
   template <typename T>
-  void set_arg(int argIndex,
-               ext::oneapi::experimental::dynamic_parameter<T> &dynamicParam) {
-    setArgHelper(argIndex, dynamicParam);
+  void set_arg(int ArgIndex,
+               ext::oneapi::experimental::dynamic_parameter<T> &DynamicParam) {
+    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));
+  void set_arg(int ArgIndex, ext::oneapi::experimental::raw_kernel_arg &&Arg) {
+    setArgHelper(ArgIndex, std::move(Arg));
   }
 
   /// Sets arguments for OpenCL interoperability kernels.

sycl/source/detail/graph_impl.cpp

@@ -1893,6 +1893,13 @@ 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, const property_list &PropList)
+    : impl(std::make_shared<dynamic_parameter_impl>(
+          sycl::detail::getSyclObjImpl(Graph))) {
+  checkGraphPropertiesAndThrow(PropList);
+}
+
 dynamic_parameter_base::dynamic_parameter_base(
     command_graph<graph_state::modifiable> Graph, size_t ParamSize,
     const void *Data)
@@ -1913,6 +1920,20 @@ void dynamic_parameter_base::updateAccessor(
   impl->updateAccessor(Acc);
 }
 
+sycl::detail::LocalAccessorImplPtr
+dynamic_parameter_base::getLocalAccessor(handler *Handler) const {
+  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
@@ -1968,6 +1989,53 @@ void dynamic_parameter_impl::updateAccessor(
               sizeof(sycl::detail::AccessorBaseHost));
 }
 
+sycl::detail::LocalAccessorImplPtr
+dynamic_parameter_impl::getLocalAccessor(handler *Handler) const {
+  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) {
@@ -2033,6 +2101,27 @@ 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) {}
@@ -2154,6 +2243,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

@@ -1412,6 +1412,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) {
@@ -1477,6 +1481,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) const;
+
+  /// 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.
@@ -1493,13 +1517,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