[SYCL][Graph] Initial implementation of graph-owned device allocations

- Allocations managed via new graph_mem_pool
- Device allocations use virtual memory
- Intercept async_alloc calls when adding nodes to graph
- New tests for functionality

Author: Bensuo

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

@@ -121,7 +121,9 @@ enum class node_type {
   memadvise = 7,
   ext_oneapi_barrier = 8,
   host_task = 9,
-  native_command = 10
+  native_command = 10,
+  async_malloc = 11,
+  async_free = 12
 };
 
 /// Class representing a node in the graph, returned by command_graph::add().
@@ -429,6 +431,10 @@ class __SYCL_EXPORT executable_command_graph
   /// @param Nodes The nodes to use for updating the graph.
   void update(const std::vector<node> &Nodes);
 
+  /// Return the total amount of memory required by this graph for graph-owned
+  /// memory allocations.
+  size_t get_required_mem_size() const;
+
   /// Common Reference Semantics
   friend bool operator==(const executable_command_graph &LHS,
                          const executable_command_graph &RHS) {

sycl/source/detail/async_alloc.cpp

@@ -44,23 +44,28 @@ void *async_malloc(sycl::handler &h, sycl::usm::alloc kind, size_t size) {
         sycl::make_error_code(sycl::errc::feature_not_supported),
         "Only device backed asynchronous allocations are supported!");
 
-  h.throwIfGraphAssociated<
-      ext::oneapi::experimental::detail::UnsupportedGraphFeatures::
-          sycl_ext_oneapi_async_alloc>();
-
   auto &Adapter = h.getContextImplPtr()->getAdapter();
-  auto &Q = h.MQueue->getHandleRef();
 
   // Get events to wait on.
   auto depEvents = getUrEvents(h.impl->CGData.MEvents);
   uint32_t numEvents = h.impl->CGData.MEvents.size();
 
   void *alloc = nullptr;
+
   ur_event_handle_t Event;
-  Adapter->call<sycl::errc::runtime,
-                sycl::detail::UrApiKind::urEnqueueUSMDeviceAllocExp>(
-      Q, (ur_usm_pool_handle_t)0, size, nullptr, numEvents, depEvents.data(),
-      &alloc, &Event);
+  // If a graph is present do the allocation from the graph memory pool instead.
+  if (auto Graph = h.getCommandGraph(); Graph) {
+    // size may be modified to reflect the aligned size based on device
+    // granularity.
+    alloc = Graph->getMemPool().malloc(size, kind);
+
+  } else {
+    auto &Q = h.MQueue->getHandleRef();
+    Adapter->call<sycl::errc::runtime,
+                  sycl::detail::UrApiKind::urEnqueueUSMDeviceAllocExp>(
+        Q, (ur_usm_pool_handle_t)0, size, nullptr, numEvents, depEvents.data(),
+        &alloc, &Event);
+  }
 
   // Async malloc must return a void* immediately.
   // Set up CommandGroup which is a no-op and pass the
@@ -90,25 +95,31 @@ __SYCL_EXPORT void *async_malloc(const sycl::queue &q, sycl::usm::alloc kind,
 __SYCL_EXPORT void *async_malloc_from_pool(sycl::handler &h, size_t size,
                                            const memory_pool &pool) {
 
-  h.throwIfGraphAssociated<
-      ext::oneapi::experimental::detail::UnsupportedGraphFeatures::
-          sycl_ext_oneapi_async_alloc>();
-
   auto &Adapter = h.getContextImplPtr()->getAdapter();
-  auto &Q = h.MQueue->getHandleRef();
   auto &memPoolImpl = sycl::detail::getSyclObjImpl(pool);
 
   // Get events to wait on.
   auto depEvents = getUrEvents(h.impl->CGData.MEvents);
   uint32_t numEvents = h.impl->CGData.MEvents.size();
 
   void *alloc = nullptr;
-  ur_event_handle_t Event;
-  Adapter->call<sycl::errc::runtime,
-                sycl::detail::UrApiKind::urEnqueueUSMDeviceAllocExp>(
-      Q, memPoolImpl.get()->get_handle(), size, nullptr, numEvents,
-      depEvents.data(), &alloc, &Event);
 
+  ur_event_handle_t Event;
+  // If a graph is present do the allocation from the graph memory pool instead.
+  if (auto Graph = h.getCommandGraph(); Graph) {
+    // size may be modified to reflect the aligned size based on device
+    // granularity.
+    // Memory pool is passed as the graph may use some properties of it.
+    alloc = Graph->getMemPool().malloc(size, pool.get_alloc_kind(),
+                                       sycl::detail::getSyclObjImpl(pool));
+
+  } else {
+    auto &Q = h.MQueue->getHandleRef();
+    Adapter->call<sycl::errc::runtime,
+                  sycl::detail::UrApiKind::urEnqueueUSMDeviceAllocExp>(
+        Q, memPoolImpl.get()->get_handle(), size, nullptr, numEvents,
+        depEvents.data(), &alloc, &Event);
+  }
   // Async malloc must return a void* immediately.
   // Set up CommandGroup which is a no-op and pass the event from the alloc.
   h.impl->MAsyncAllocEvent = Event;
@@ -135,9 +146,15 @@ async_malloc_from_pool(const sycl::queue &q, size_t size,
 }
 
 __SYCL_EXPORT void async_free(sycl::handler &h, void *ptr) {
-  h.throwIfGraphAssociated<
-      ext::oneapi::experimental::detail::UnsupportedGraphFeatures::
-          sycl_ext_oneapi_async_alloc>();
+  if (auto Graph = h.getCommandGraph(); Graph) {
+    // Check if the pointer to be freed has an associated allocation node, and
+    // error if not
+    if (!Graph->getMemPool().hasAllocation(ptr)) {
+      throw sycl::exception(sycl::make_error_code(sycl::errc::invalid),
+                            "Cannot add a free node to a graph for which "
+                            "there is no associated allocation node!");
+    }
+  }
 
   h.impl->MFreePtr = ptr;
   h.setType(detail::CGType::AsyncFree);

sycl/source/detail/graph_impl.cpp

@@ -340,7 +340,7 @@ graph_impl::graph_impl(const sycl::context &SyclContext,
                        const sycl::device &SyclDevice,
                        const sycl::property_list &PropList)
     : MContext(SyclContext), MDevice(SyclDevice), MRecordingQueues(),
-      MEventsMap(), MInorderQueueMap(),
+      MEventsMap(), MInorderQueueMap(), MGraphMemPool(SyclContext, SyclDevice),
       MID(NextAvailableID.fetch_add(1, std::memory_order_relaxed)) {
   checkGraphPropertiesAndThrow(PropList);
   if (PropList.has_property<property::graph::no_cycle_check>()) {
@@ -748,12 +748,12 @@ void graph_impl::beginRecording(
   }
 }
 
-// Check if nodes are empty and if so loop back through predecessors until we
-// find the real dependency.
+// Check if nodes do not require enqueueing and if so loop back through
+// predecessors until we find the real dependency.
 void exec_graph_impl::findRealDeps(
     std::vector<ur_exp_command_buffer_sync_point_t> &Deps,
     std::shared_ptr<node_impl> CurrentNode, int ReferencePartitionNum) {
-  if (CurrentNode->isEmpty()) {
+  if (!CurrentNode->requiresEnqueue()) {
     for (auto &N : CurrentNode->MPredecessors) {
       auto NodeImpl = N.lock();
       findRealDeps(Deps, NodeImpl, ReferencePartitionNum);
@@ -873,9 +873,9 @@ void exec_graph_impl::createCommandBuffers(
   Partition->MCommandBuffers[Device] = OutCommandBuffer;
 
   for (const auto &Node : Partition->MSchedule) {
-    // Empty nodes are not processed as other nodes, but only their
+    // Some nodes are not scheduled like other nodes, and only their
     // dependencies are propagated in findRealDeps
-    if (Node->isEmpty())
+    if (!Node->requiresEnqueue())
       continue;
 
     sycl::detail::CGType type = Node->MCGType;
@@ -941,6 +941,8 @@ exec_graph_impl::exec_graph_impl(sycl::context Context,
 
 exec_graph_impl::~exec_graph_impl() {
   try {
+    MGraphImpl->markExecGraphDestroyed();
+
     const sycl::detail::AdapterPtr &Adapter =
         sycl::detail::getSyclObjImpl(MContext)->getAdapter();
     MSchedule.clear();
@@ -950,6 +952,9 @@ exec_graph_impl::~exec_graph_impl() {
       Event->wait(Event);
     }
 
+    // Clean up any graph-owned allocations that were allocated
+    MGraphImpl->getMemPool().deallocateAndUnmapAll();
+
     for (const auto &Partition : MPartitions) {
       Partition->MSchedule.clear();
       for (const auto &Iter : Partition->MCommandBuffers) {
@@ -1867,6 +1872,14 @@ modifiable_command_graph::finalize(const sycl::property_list &PropList) const {
   // Graph is read and written in this scope so we lock
   // this graph with full priviledges.
   graph_impl::WriteLock Lock(impl->MMutex);
+  // If the graph uses graph-owned allocations and an executable graph already
+  // exists we must throw an error.
+  if (impl->getMemPool().hasAllocations() && impl->getExecGraphCount() > 0) {
+    throw sycl::exception(sycl::make_error_code(errc::invalid),
+                          "Graphs containing allocations can only have a "
+                          "single executable graph alive at any onc time.");
+  }
+
   return command_graph<graph_state::executable>{
       this->impl, this->impl->getContext(), PropList};
 }
@@ -1994,11 +2007,16 @@ executable_command_graph::executable_command_graph(
     const property_list &PropList)
     : impl(std::make_shared<detail::exec_graph_impl>(Ctx, Graph, PropList)) {
   finalizeImpl(); // Create backend representation for executable graph
+  // Mark that we have created an executable graph from the modifiable graph.
+  Graph->markExecGraphCreated();
 }
 
 void executable_command_graph::finalizeImpl() {
   impl->makePartitions();
 
+  // Handle any work required for graph-owned memory allocations
+  impl->finalizeMemoryAllocations();
+
   auto Device = impl->getGraphImpl()->getDevice();
   for (auto Partition : impl->getPartitions()) {
     if (!Partition->isHostTask()) {
@@ -2026,6 +2044,13 @@ void executable_command_graph::update(const std::vector<node> &Nodes) {
   impl->update(NodeImpls);
 }
 
+size_t executable_command_graph::get_required_mem_size() const {
+  // Since each graph has a unique mem pool, return the current memory usage for
+  // now. This call my change if we move to being able to share memory between
+  // unique graphs.
+  return impl->getGraphImpl()->getMemPool().getMemUseCurrent();
+}
+
 dynamic_parameter_base::dynamic_parameter_base(
     command_graph<graph_state::modifiable> Graph)
     : impl(std::make_shared<dynamic_parameter_impl>(

sycl/source/detail/graph_impl.hpp

@@ -17,6 +17,7 @@
 #include <detail/accessor_impl.hpp>
 #include <detail/cg.hpp>
 #include <detail/event_impl.hpp>
+#include <detail/graph_memory_pool.hpp>
 #include <detail/host_task.hpp>
 #include <detail/kernel_impl.hpp>
 #include <detail/sycl_mem_obj_t.hpp>
@@ -73,6 +74,11 @@ inline node_type getNodeTypeFromCG(sycl::detail::CGType CGType) {
     return node_type::subgraph;
   case sycl::detail::CGType::EnqueueNativeCommand:
     return node_type::native_command;
+  case sycl::detail::CGType::AsyncAlloc:
+    return node_type::async_malloc;
+  case sycl::detail::CGType::AsyncFree:
+    return node_type::async_free;
+
   default:
     assert(false && "Invalid Graph Node Type");
     return node_type::empty;
@@ -473,6 +479,15 @@ class node_impl : public std::enable_shared_from_this<node_impl> {
     }
   }
 
+  /// Returns true if this node should be enqueued to the backend, if not only
+  /// its dependencies are considered.
+  bool requiresEnqueue() const {
+    return MNodeType != node_type::empty &&
+           MNodeType != node_type::ext_oneapi_barrier &&
+           MNodeType != node_type::async_malloc &&
+           MNodeType != node_type::async_free;
+  }
+
 private:
   void rebuildArgStorage(std::vector<sycl::detail::ArgDesc> &Args,
                          const std::vector<std::vector<char>> &OldArgStorage,
@@ -919,6 +934,12 @@ class graph_impl : public std::enable_shared_from_this<graph_impl> {
   /// @return Context associated with graph.
   sycl::context getContext() const { return MContext; }
 
+  /// Query for the context impl tied to this graph.
+  /// @return shared_ptr ref for the context impl associated with graph.
+  const std::shared_ptr<sycl::detail::context_impl> &getContextImplPtr() const {
+    return sycl::detail::getSyclObjImpl(MContext);
+  }
+
   /// Query for the device_impl tied to this graph.
   /// @return device_impl shared ptr reference associated with graph.
   const DeviceImplPtr &getDeviceImplPtr() const {
@@ -1139,6 +1160,24 @@ class graph_impl : public std::enable_shared_from_this<graph_impl> {
 
   unsigned long long getID() const { return MID; }
 
+  /// Get the memory pool used for graph-owned allocations.
+  graph_mem_pool &getMemPool() { return MGraphMemPool; }
+
+  /// Mark that an executable graph was created from this modifiable graph, used
+  /// for tracking live graphs for graph-owned allocations.
+  void markExecGraphCreated() { MExecGraphCount++; }
+
+  /// Mark that an executable graph created from this modifiable graph was
+  /// destroyed, used for tracking live graphs for graph-owned allocations.
+  void markExecGraphDestroyed() {
+    assert(MExecGraphCount != 0);
+    MExecGraphCount--;
+  }
+
+  /// Get the number of unique executable graph instances currently alive for
+  /// this graph.
+  size_t getExecGraphCount() const { return MExecGraphCount; }
+
 private:
   /// Check the graph for cycles by performing a depth-first search of the
   /// graph. If a node is visited more than once in a given path through the
@@ -1206,10 +1245,17 @@ class graph_impl : public std::enable_shared_from_this<graph_impl> {
   std::map<std::weak_ptr<sycl::detail::queue_impl>, std::shared_ptr<node_impl>,
            std::owner_less<std::weak_ptr<sycl::detail::queue_impl>>>
       MBarrierDependencyMap;
+  /// Graph memory pool for handling graph-owned memory allocations for this
+  /// graph.
+  graph_mem_pool MGraphMemPool;
 
   unsigned long long MID;
   // Used for std::hash in order to create a unique hash for the instance.
   inline static std::atomic<unsigned long long> NextAvailableID = 0;
+
+  // The number of live executable graphs that have been created from this
+  // modifiable graph
+  size_t MExecGraphCount = 0;
 };
 
 /// Class representing the implementation of command_graph<executable>.
@@ -1334,6 +1380,14 @@ class exec_graph_impl {
 
   unsigned long long getID() const { return MID; }
 
+  /// Do any work required during finalization to finalize graph-owned memory
+  /// allocations.
+  void finalizeMemoryAllocations() {
+    // This call allocates physical memory and maps all virtual device
+    // allocations
+    MGraphImpl->getMemPool().allocateAndMapAll();
+  }
+
 private:
   /// Create a command-group for the node and add it to command-buffer by going
   /// through the scheduler.