[SYCL][Fusion] Take auxiliary resources from fused command groups

Reductions use auxiliary resources to handle reduction temporal
buffers. Assign those auxiliary resources from each fused command
group to the placeholder fusion event. This event will not be marked
as completed after the fused reductions finish execution either if
fusion is completed, cancelled or aborted.

Test is updated to check every algorithm that is selected
automatically by `sycl::reduction`, i.e., every supported
algorithm. We also cover both cases (fusion taking and not taking
place).

Signed-off-by: Victor Perez <[email protected]>

Author: victor-eds

sycl/source/detail/jit_compiler.cpp

@@ -667,8 +667,8 @@ jit_compiler::fuseKernels(QueueImplPtr Queue,
   unsigned KernelIndex = 0;
   ParamList FusedParams;
   PromotionMap PromotedAccs;
-  // TODO: Collect information about streams and auxiliary resources (which
-  // contain reductions) and figure out how to fuse them.
+  // TODO: Collect information about streams and figure out how
+  // to fuse them.
   for (auto &RawCmd : InputKernels) {
     auto *KernelCmd = static_cast<ExecCGCommand *>(RawCmd);
     auto &CG = KernelCmd->getCG();

sycl/source/detail/scheduler/graph_builder.cpp

@@ -1615,6 +1615,9 @@ Scheduler::GraphBuilder::completeFusion(QueueImplPtr Queue,
   auto FusedKernelCmd =
       std::make_unique<ExecCGCommand>(std::move(FusedCG), Queue);
 
+  // Inherit auxiliary resources from fused command groups
+  Scheduler::getInstance().takeAuxiliaryResources(FusedKernelCmd->getEvent(),
+                                                  PlaceholderCmd->getEvent());
   assert(PlaceholderCmd->MDeps.empty());
   // Next, backwards iterate over all the commands in the fusion list and remove
   // them from the graph to restore the state before starting fusion, so we can

sycl/source/detail/scheduler/scheduler.cpp

@@ -154,11 +154,11 @@ EventImplPtr Scheduler::addCG(
     for (const auto &StreamImplPtr : Streams) {
       StreamImplPtr->flush(NewEvent);
     }
-
-    if (!AuxiliaryResources.empty())
-      registerAuxiliaryResources(NewEvent, std::move(AuxiliaryResources));
   }
 
+  if (!AuxiliaryResources.empty())
+    registerAuxiliaryResources(NewEvent, std::move(AuxiliaryResources));
+
   return NewEvent;
 }
 
@@ -558,10 +558,35 @@ void Scheduler::cleanupDeferredMemObjects(BlockingT Blocking) {
   }
 }
 
+static void registerAuxiliaryResourcesNoLock(
+    std::unordered_map<EventImplPtr, std::vector<std::shared_ptr<const void>>>
+        &AuxiliaryResources,
+    const EventImplPtr &Event,
+    std::vector<std::shared_ptr<const void>> &&Resources) {
+  std::vector<std::shared_ptr<const void>> &StoredResources =
+      AuxiliaryResources[Event];
+  StoredResources.insert(StoredResources.end(),
+                         std::make_move_iterator(Resources.begin()),
+                         std::make_move_iterator(Resources.end()));
+}
+
+void Scheduler::takeAuxiliaryResources(const EventImplPtr &Dst,
+                                       const EventImplPtr &Src) {
+  std::unique_lock<std::mutex> Lock{MAuxiliaryResourcesMutex};
+  auto Iter = MAuxiliaryResources.find(Src);
+  if (Iter == MAuxiliaryResources.end()) {
+    return;
+  }
+  registerAuxiliaryResourcesNoLock(MAuxiliaryResources, Dst,
+                                   std::move(Iter->second));
+  MAuxiliaryResources.erase(Iter);
+}
+
 void Scheduler::registerAuxiliaryResources(
     EventImplPtr &Event, std::vector<std::shared_ptr<const void>> Resources) {
   std::unique_lock<std::mutex> Lock{MAuxiliaryResourcesMutex};
-  MAuxiliaryResources.insert({Event, std::move(Resources)});
+  registerAuxiliaryResourcesNoLock(MAuxiliaryResources, Event,
+                                   std::move(Resources));
 }
 
 void Scheduler::cleanupAuxiliaryResources(BlockingT Blocking) {

sycl/source/detail/scheduler/scheduler.hpp

@@ -543,6 +543,8 @@ class Scheduler {
     bool ShouldEnqueue;
   };
 
+  /// Assign \p Src's auxiliary resources to \p Dst.
+  void takeAuxiliaryResources(const EventImplPtr &Dst, const EventImplPtr &Src);
   void registerAuxiliaryResources(
       EventImplPtr &Event, std::vector<std::shared_ptr<const void>> Resources);
   void cleanupAuxiliaryResources(BlockingT Blocking);

sycl/test-e2e/KernelFusion/Reduction/group_reduce_and_last_wg_detection.cpp

@@ -0,0 +1,9 @@
+// RUN: %{build} -fsycl-embed-ir -o %t.out
+// RUN: %{run} %t.out
+// UNSUPPORTED: hip || cuda
+
+#include "./reduction.hpp"
+
+int main() {
+  test<detail::reduction::strategy::group_reduce_and_last_wg_detection>();
+}

sycl/test-e2e/KernelFusion/Reduction/local_atomic_and_atomic_cross_wg.cpp

@@ -0,0 +1,8 @@
+// RUN: %{build} -fsycl-embed-ir -o %t.out
+// RUN: %{run} %t.out
+
+#include "./reduction.hpp"
+
+int main() {
+  test<detail::reduction::strategy::local_atomic_and_atomic_cross_wg>();
+}

sycl/test-e2e/KernelFusion/Reduction/range_basic.cpp

@@ -0,0 +1,6 @@
+// RUN: %{build} -fsycl-embed-ir -o %t.out
+// RUN: %{run} %t.out
+
+#include "./reduction.hpp"
+
+int main() { test<detail::reduction::strategy::range_basic>(); }

sycl/test-e2e/KernelFusion/Reduction/reduction.cpp

@@ -0,0 +1,86 @@
+// Test fusion works with reductions. Some algorithms will lead to fusion being
+// cancelled in some devices. These should work properly anyway.
+
+#include <sycl/sycl.hpp>
+
+#include "../helpers.hpp"
+#include "sycl/detail/reduction_forward.hpp"
+
+using namespace sycl;
+
+constexpr inline size_t globalSize = 512;
+
+template <detail::reduction::strategy Strategy> struct is_fusion_supported {
+  constexpr static inline bool value =
+      detail::reduction::strategy::group_reduce_and_last_wg_detection <=
+          Strategy &&
+      Strategy < detail::reduction::strategy::group_reduce_and_atomic_cross_wg;
+};
+
+template <detail::reduction::strategy Strategy>
+constexpr inline bool is_fusion_supported_v =
+    is_fusion_supported<Strategy>::value;
+
+template <detail::reduction::strategy Strategy, bool Fuse>
+void test(nd_range<1> ndr) {
+  static_assert(is_fusion_supported_v<Strategy>,
+                "Testing unsupported algorithm");
+  std::array<int, globalSize> data;
+  int sumRes = 0;
+  int maxRes = 0;
+
+  {
+    queue q{ext::codeplay::experimental::property::queue::enable_fusion{}};
+
+    buffer<int> dataBuf{data};
+    buffer<int> sumBuf{&sumRes, 1};
+    buffer<int> maxBuf{&maxRes, 1};
+
+    ext::codeplay::experimental::fusion_wrapper fw{q};
+
+    fw.start_fusion();
+    iota(q, dataBuf, 0);
+
+    q.submit([&](handler &cgh) {
+      accessor in(dataBuf, cgh, read_only);
+      auto sumRed = reduction(sumBuf, cgh, plus<>{},
+                              property::reduction::initialize_to_identity{});
+      detail::reduction_parallel_for<detail::auto_name, Strategy>(
+          cgh, ndr, ext::oneapi::experimental::empty_properties_t{}, sumRed,
+          [=](nd_item<1> Item, auto &Red) {
+            Red.combine(in[Item.get_global_id()]);
+          });
+    });
+
+    q.submit([&](handler &cgh) {
+      accessor in(dataBuf, cgh, read_only);
+      auto maxRed = reduction(maxBuf, cgh, maximum<>{},
+                              property::reduction::initialize_to_identity{});
+      detail::reduction_parallel_for<detail::auto_name, Strategy>(
+          cgh, ndr, ext::oneapi::experimental::empty_properties_t{}, maxRed,
+          [=](nd_item<1> Item, auto &Red) {
+            Red.combine(in[Item.get_global_id()]);
+          });
+    });
+
+    if constexpr (Fuse) {
+      fw.complete_fusion(ext::codeplay::experimental::property::no_barriers{});
+    } else {
+      fw.cancel_fusion();
+    }
+  }
+
+  constexpr int expectedMax = globalSize - 1;
+  constexpr int expectedSum = globalSize * expectedMax / 2;
+
+  assert(sumRes == expectedSum);
+  assert(maxRes == expectedMax);
+}
+
+template <detail::reduction::strategy Strategy> void test() {
+  for (size_t localSize = 1; localSize <= globalSize; localSize *= 2) {
+    nd_range<1> ndr{globalSize, localSize};
+    test<Strategy, true>(ndr);
+    test<Strategy, false>(ndr);
+  }
+}