Refactor how we find updatable partitions

Author: Ewan Crawford

sycl/source/detail/graph_impl.cpp

@@ -1401,6 +1401,16 @@ void exec_graph_impl::update(
   std::vector<sycl::detail::AccessorImplHost *> UpdateRequirements;
   bool NeedScheduledUpdate = needsScheduledUpdate(Nodes, UpdateRequirements);
   if (NeedScheduledUpdate) {
+    // Clean up any execution events which have finished so we don't pass them
+    // to the scheduler.
+    for (auto It = MExecutionEvents.begin(); It != MExecutionEvents.end();) {
+      if ((*It)->isCompleted()) {
+        It = MExecutionEvents.erase(It);
+        continue;
+      }
+      ++It;
+    }
+
     auto AllocaQueue = std::make_shared<sycl::detail::queue_impl>(
         sycl::detail::getSyclObjImpl(MGraphImpl->getDevice()),
         sycl::detail::getSyclObjImpl(MGraphImpl->getContext()),
@@ -1416,11 +1426,12 @@ void exec_graph_impl::update(
   } else {
     // For each partition in the executable graph, call UR update on the
     // command-buffer with the nodes to update.
-    auto PartitionedNodes = getPartitionForNodes(Nodes);
+    auto PartitionedNodes = getURUpdatableNodes(Nodes);
     for (auto It = PartitionedNodes.begin(); It != PartitionedNodes.end();
          It++) {
-      auto CommandBuffer = It->first->MCommandBuffers[MDevice];
-      updateKernelsImpl(CommandBuffer, It->second);
+      auto &Partition = MPartitions[It->first];
+      auto CommandBuffer = Partition->MCommandBuffers[MDevice];
+      updateURImpl(CommandBuffer, It->second);
     }
   }
 
@@ -1475,16 +1486,6 @@ bool exec_graph_impl::needsScheduledUpdate(
     }
   }
 
-  // Clean up any execution events which have finished so we don't pass them to
-  // the scheduler.
-  for (auto It = MExecutionEvents.begin(); It != MExecutionEvents.end();) {
-    if ((*It)->isCompleted()) {
-      It = MExecutionEvents.erase(It);
-      continue;
-    }
-    ++It;
-  }
-
   // If we have previous execution events do the update through the scheduler to
   // ensure it is ordered correctly.
   NeedScheduledUpdate |= MExecutionEvents.size() > 0;
@@ -1499,7 +1500,7 @@ void exec_graph_impl::populateURKernelUpdateStructs(
     std::vector<ur_exp_command_buffer_update_pointer_arg_desc_t> &PtrDescs,
     std::vector<ur_exp_command_buffer_update_value_arg_desc_t> &ValueDescs,
     sycl::detail::NDRDescT &NDRDesc,
-    ur_exp_command_buffer_update_kernel_launch_desc_t &UpdateDesc) {
+    ur_exp_command_buffer_update_kernel_launch_desc_t &UpdateDesc) const {
   auto ContextImpl = sycl::detail::getSyclObjImpl(MContext);
   const sycl::detail::AdapterPtr &Adapter = ContextImpl->getAdapter();
   auto DeviceImpl = sycl::detail::getSyclObjImpl(MGraphImpl->getDevice());
@@ -1656,97 +1657,84 @@ void exec_graph_impl::populateURKernelUpdateStructs(
   auto ExecNode = MIDCache.find(Node->MID);
   assert(ExecNode != MIDCache.end() && "Node ID was not found in ID cache");
 
-  ur_exp_command_buffer_command_handle_t Command =
-      MCommandMap[ExecNode->second];
-  UpdateDesc.hCommand = Command;
+  auto Command = MCommandMap.find(ExecNode->second);
+  assert(Command != MCommandMap.end());
+  UpdateDesc.hCommand = Command->second;
 
   // Update ExecNode with new values from Node, in case we ever need to
   // rebuild the command buffers
   ExecNode->second->updateFromOtherNode(Node);
 }
 
-std::map<std::shared_ptr<partition>, std::vector<std::shared_ptr<node_impl>>>
-exec_graph_impl::getPartitionForNodes(
-    const std::vector<std::shared_ptr<node_impl>> &Nodes) {
-  // Iterate over each partition in the executable graph, and find the nodes
-  // in "Nodes" that also exist in the partition.
-  std::map<std::shared_ptr<partition>, std::vector<std::shared_ptr<node_impl>>>
-      PartitionedNodes;
-  for (const auto &Partition : MPartitions) {
-    std::vector<std::shared_ptr<node_impl>> NodesForPartition;
-    const auto PartitionBegin = Partition->MSchedule.begin();
-    const auto PartitionEnd = Partition->MSchedule.end();
-    for (auto &Node : Nodes) {
-      auto ExecNode = MIDCache.find(Node->MID);
-      assert(ExecNode != MIDCache.end() && "Node ID was not found in ID cache");
-
-      if (std::find_if(PartitionBegin, PartitionEnd,
-                       [ExecNode](const auto &PartitionNode) {
-                         return PartitionNode->MID == ExecNode->second->MID;
-                       }) != PartitionEnd) {
-        NodesForPartition.push_back(Node);
-      }
-    }
-    if (!NodesForPartition.empty()) {
-      PartitionedNodes.insert({Partition, NodesForPartition});
+std::map<int, std::vector<std::shared_ptr<node_impl>>>
+exec_graph_impl::getURUpdatableNodes(
+    const std::vector<std::shared_ptr<node_impl>> &Nodes) const {
+  // Iterate over the list of nodes, and for every node that can
+  // be updated through UR, add it to the list of nodes for
+  // that can be updated for the UR command-buffer partition.
+  std::map<int, std::vector<std::shared_ptr<node_impl>>> PartitionedNodes;
+
+  // Initialize vector for each partition
+  for (size_t i = 0; i < MPartitions.size(); i++) {
+    PartitionedNodes[i] = {};
+  }
+
+  for (auto &Node : Nodes) {
+    // Kernel node update is the only command type supported in UR for update.
+    if (Node->MCGType != sycl::detail::CGType::Kernel) {
+      continue;
     }
+
+    auto ExecNode = MIDCache.find(Node->MID);
+    assert(ExecNode != MIDCache.end() && "Node ID was not found in ID cache");
+    auto PartitionIndex = MPartitionNodes.find(ExecNode->second);
+    assert(PartitionIndex != MPartitionNodes.end());
+    PartitionedNodes[PartitionIndex->second].push_back(Node);
   }
 
   return PartitionedNodes;
 }
 
 void exec_graph_impl::updateHostTasksImpl(
-    const std::vector<std::shared_ptr<node_impl>> &Nodes) {
+    const std::vector<std::shared_ptr<node_impl>> &Nodes) const {
   for (auto &Node : Nodes) {
     if (Node->MNodeType != node_type::host_task) {
       continue;
     }
     // Query the ID cache to find the equivalent exec node for the node passed
     // to this function.
-    // TODO: Handle subgraphs or any other cases where multiple nodes may be
-    // associated with a single key, once those node types are supported for
-    // update.
     auto ExecNode = MIDCache.find(Node->MID);
     assert(ExecNode != MIDCache.end() && "Node ID was not found in ID cache");
 
     ExecNode->second->updateFromOtherNode(Node);
   }
 }
 
-void exec_graph_impl::updateKernelsImpl(
+void exec_graph_impl::updateURImpl(
     ur_exp_command_buffer_handle_t CommandBuffer,
-    const std::vector<std::shared_ptr<node_impl>> &Nodes) {
-  // Kernel node update is the only command type supported in UR for update.
-  // Updating any other types of nodes, e.g. empty & barrier nodes is a no-op.
-  size_t NumKernelNodes = 0;
-  for (auto &N : Nodes) {
-    if (N->MCGType == sycl::detail::CGType::Kernel) {
-      NumKernelNodes++;
-    }
-  }
-
-  // Don't need to call through to UR if no kernel nodes to update
-  if (NumKernelNodes == 0) {
+    const std::vector<std::shared_ptr<node_impl>> &Nodes) const {
+  const size_t NumUpdatableNodes = Nodes.size();
+  if (NumUpdatableNodes == 0) {
     return;
   }
 
   std::vector<std::vector<ur_exp_command_buffer_update_memobj_arg_desc_t>>
-      MemobjDescsList(NumKernelNodes);
+      MemobjDescsList(NumUpdatableNodes);
   std::vector<std::vector<ur_exp_command_buffer_update_pointer_arg_desc_t>>
-      PtrDescsList(NumKernelNodes);
+      PtrDescsList(NumUpdatableNodes);
   std::vector<std::vector<ur_exp_command_buffer_update_value_arg_desc_t>>
-      ValueDescsList(NumKernelNodes);
-  std::vector<sycl::detail::NDRDescT> NDRDescList(NumKernelNodes);
+      ValueDescsList(NumUpdatableNodes);
+  std::vector<sycl::detail::NDRDescT> NDRDescList(NumUpdatableNodes);
   std::vector<ur_exp_command_buffer_update_kernel_launch_desc_t> UpdateDescList(
-      NumKernelNodes);
+      NumUpdatableNodes);
   std::vector<std::pair<ur_program_handle_t, ur_kernel_handle_t>>
-      KernelBundleObjList(NumKernelNodes);
+      KernelBundleObjList(NumUpdatableNodes);
 
   size_t StructListIndex = 0;
   for (auto &Node : Nodes) {
-    if (Node->MCGType != sycl::detail::CGType::Kernel) {
-      continue;
-    }
+    // This should be the case when getURUpdatableNodes() is used to
+    // create the list of nodes.
+    assert(Node->MCGType == sycl::detail::CGType::Kernel);
 
     auto &MemobjDescs = MemobjDescsList[StructListIndex];
     auto &KernelBundleObjs = KernelBundleObjList[StructListIndex];

sycl/source/detail/graph_impl.hpp

@@ -1303,18 +1303,30 @@ class exec_graph_impl {
   void update(std::shared_ptr<node_impl> Node);
   void update(const std::vector<std::shared_ptr<node_impl>> &Nodes);
 
-  /// Calls UR entry-point to update kernel nodes in command-buffer.
+  /// Calls UR entry-point to update nodes in command-buffer.
   /// @param CommandBuffer The UR command-buffer to update commands in.
-  /// @param Nodes List of nodes to update. May contain nodes of non-kernel
-  /// type, but only kernel nodes from the list will be used for update
-  void updateKernelsImpl(ur_exp_command_buffer_handle_t CommandBuffer,
-                         const std::vector<std::shared_ptr<node_impl>> &Nodes);
+  /// @param Nodes List of nodes to update. Only nodes which can be updated
+  /// through UR should be included in this list, currently this is only
+  /// nodes of kernel type.
+  void updateURImpl(ur_exp_command_buffer_handle_t CommandBuffer,
+                    const std::vector<std::shared_ptr<node_impl>> &Nodes) const;
 
-  /// Splits a list of nodes into separate lists depending on partition.
+  /// Update host-task nodes
+  /// @param Nodes List of nodes to update, any node that is not a host-task
+  /// will be ignored.
+  void updateHostTasksImpl(
+      const std::vector<std::shared_ptr<node_impl>> &Nodes) const;
+
+  /// Splits a list of nodes into separate lists of nodes for each
+  /// command-buffer partition.
+  ///
+  /// Only nodes that can be updated through the UR interface are included
+  /// in the list. Currently this is only kernel node types.
+  ///
   /// @param Nodes List of nodes to split
-  /// @return Map of partitions to nodes
-  std::map<std::shared_ptr<partition>, std::vector<std::shared_ptr<node_impl>>>
-  getPartitionForNodes(const std::vector<std::shared_ptr<node_impl>> &Nodes);
+  /// @return Map of partition indexes to nodes
+  std::map<int, std::vector<std::shared_ptr<node_impl>>> getURUpdatableNodes(
+      const std::vector<std::shared_ptr<node_impl>> &Nodes) const;
 
   unsigned long long getID() const { return MID; }
 
@@ -1408,13 +1420,7 @@ class exec_graph_impl {
       std::vector<ur_exp_command_buffer_update_pointer_arg_desc_t> &PtrDescs,
       std::vector<ur_exp_command_buffer_update_value_arg_desc_t> &ValueDescs,
       sycl::detail::NDRDescT &NDRDesc,
-      ur_exp_command_buffer_update_kernel_launch_desc_t &UpdateDesc);
-
-  /// Updates host-task nodes in the graph
-  /// @param Nodes List of nodes to update, any node that is not a host-task
-  /// will be ignored.
-  void
-  updateHostTasksImpl(const std::vector<std::shared_ptr<node_impl>> &Nodes);
+      ur_exp_command_buffer_update_kernel_launch_desc_t &UpdateDesc) const;
 
   /// Execution schedule of nodes in the graph.
   std::list<std::shared_ptr<node_impl>> MSchedule;

sycl/source/detail/scheduler/commands.cpp

@@ -3744,12 +3744,15 @@ ur_result_t UpdateCommandBufferCommand::enqueueImp() {
   }
 
   // Split list of nodes into nodes per UR command-buffer partition, then
-  // call UR update on each command-buffer partition
-  auto PartitionedNodes = MGraph->getPartitionForNodes(MNodes);
+  // call UR update on each command-buffer partition with those updatable
+  // nodes.
+  auto PartitionedNodes = MGraph->getURUpdatableNodes(MNodes);
   auto Device = MQueue->get_device();
+  auto &Partitions = MGraph->getPartitions();
   for (auto It = PartitionedNodes.begin(); It != PartitionedNodes.end(); It++) {
-    auto CommandBuffer = It->first->MCommandBuffers[Device];
-    MGraph->updateKernelsImpl(CommandBuffer, It->second);
+    const int PartitionIndex = It->first;
+    auto CommandBuffer = Partitions[PartitionIndex]->MCommandBuffers[Device];
+    MGraph->updateURImpl(CommandBuffer, It->second);
   }
 
   return UR_RESULT_SUCCESS;