[SYCL] Implement true USM reductions (#1888)

This patch replaces the existing implementation of USM reductions with the
more reliable and more efficient implementation.

The previous implementation created placeholder buffer/accessor for USM pointer
passed to reduction. That caused some unnecessary overhead.
Even worse, that approach did not work on systems with detached memory, and could
cause the need in additional workaround - copy from USM to HOST mem before
the main reduction work and copy from HOST to USM memory after reduction work.
That would make the implementation even slower.

The new approach uses USM memory passed by user to reduction,
and it works 2x faster than the previous approach with workaround.

In order to avoid code duplication which could be caused by adding
true USM implementations, the significant rework in reduCGFunc and reduAuxCGFunc
was done. The re-work though did NOT change the actual implementation
and the algorithmic parts.

Signed-off-by: Vyacheslav N Klochkov <[email protected]>

Author: v-klochkov

sycl/include/CL/sycl/handler.hpp

@@ -179,33 +179,19 @@ class reduction_impl;
 
 using cl::sycl::detail::enable_if_t;
 
-template <typename KernelName, typename KernelType, int Dims, class Reduction>
-enable_if_t<Reduction::has_fast_reduce && Reduction::has_fast_atomics>
-reduCGFunc(handler &CGH, KernelType KernelFunc, const nd_range<Dims> &Range,
-           Reduction &Redu, typename Reduction::rw_accessor_type &Out);
-
-template <typename KernelName, typename KernelType, int Dims, class Reduction>
-enable_if_t<!Reduction::has_fast_reduce && Reduction::has_fast_atomics>
+template <typename KernelName, typename KernelType, int Dims, class Reduction,
+          typename OutputT>
+enable_if_t<Reduction::has_fast_atomics>
 reduCGFunc(handler &CGH, KernelType KernelFunc, const nd_range<Dims> &Range,
-           Reduction &Redu, typename Reduction::rw_accessor_type &Out);
-
-template <typename KernelName, typename KernelType, int Dims, class Reduction>
-enable_if_t<Reduction::has_fast_reduce && !Reduction::has_fast_atomics>
-reduCGFunc(handler &CGH, KernelType KernelFunc, const nd_range<Dims> &Range,
-           Reduction &Redu);
+           Reduction &Redu, OutputT Out);
 
 template <typename KernelName, typename KernelType, int Dims, class Reduction>
-enable_if_t<!Reduction::has_fast_reduce && !Reduction::has_fast_atomics>
+enable_if_t<!Reduction::has_fast_atomics>
 reduCGFunc(handler &CGH, KernelType KernelFunc, const nd_range<Dims> &Range,
            Reduction &Redu);
 
 template <typename KernelName, typename KernelType, int Dims, class Reduction>
-enable_if_t<Reduction::has_fast_reduce && !Reduction::has_fast_atomics>
-reduAuxCGFunc(handler &CGH, const nd_range<Dims> &Range, size_t NWorkItems,
-              Reduction &Redu);
-
-template <typename KernelName, typename KernelType, int Dims, class Reduction>
-enable_if_t<!Reduction::has_fast_reduce && !Reduction::has_fast_atomics>
+enable_if_t<!Reduction::has_fast_atomics>
 reduAuxCGFunc(handler &CGH, const nd_range<Dims> &Range, size_t NWorkItems,
               Reduction &Redu);
 } // namespace detail
@@ -958,24 +944,23 @@ class __SYCL_EXPORT handler {
   template <typename KernelName = detail::auto_name, typename KernelType,
             int Dims, typename Reduction>
   detail::enable_if_t<Reduction::accessor_mode == access::mode::read_write &&
-                      Reduction::has_fast_atomics>
+                      Reduction::has_fast_atomics && !Reduction::is_usm>
   parallel_for(nd_range<Dims> Range, Reduction Redu, KernelType KernelFunc) {
-    if (Reduction::is_usm)
-      Redu.associateWithHandler(*this);
-    shared_ptr_class<detail::queue_impl> QueueCopy = MQueue;
-    auto Acc = Redu.getUserAccessor();
-    intel::detail::reduCGFunc<KernelName>(*this, KernelFunc, Range, Redu, Acc);
-
-    // Submit non-blocking copy from reduction accessor to user's reduction
-    // variable.
-    if (Reduction::is_usm) {
-      this->finalize();
-      handler CopyHandler(QueueCopy, MIsHost);
-      CopyHandler.saveCodeLoc(MCodeLoc);
-      Redu.associateWithHandler(CopyHandler);
-      CopyHandler.copy(Acc, Redu.getUSMPointer());
-      MLastEvent = CopyHandler.finalize();
-    }
+    intel::detail::reduCGFunc<KernelName>(*this, KernelFunc, Range, Redu,
+                                          Redu.getUserAccessor());
+  }
+
+  /// Implements parallel_for() accepting nd_range and 1 reduction variable
+  /// having 'read_write' access mode.
+  /// This version uses fast sycl::atomic operations to update user's reduction
+  /// variable at the end of each work-group work.
+  template <typename KernelName = detail::auto_name, typename KernelType,
+            int Dims, typename Reduction>
+  detail::enable_if_t<Reduction::accessor_mode == access::mode::read_write &&
+                      Reduction::has_fast_atomics && Reduction::is_usm>
+  parallel_for(nd_range<Dims> Range, Reduction Redu, KernelType KernelFunc) {
+    intel::detail::reduCGFunc<KernelName>(*this, KernelFunc, Range, Redu,
+                                          Redu.getUSMPointer());
   }
 
   /// Implements parallel_for() accepting nd_range and 1 reduction variable
@@ -1045,8 +1030,6 @@ class __SYCL_EXPORT handler {
     //    necessary to reduce all partial sums into one final sum.
 
     // 1. Call the kernel that includes user's lambda function.
-    if (Reduction::is_usm && NWorkGroups == 1)
-      Redu.associateWithHandler(*this);
     intel::detail::reduCGFunc<KernelName>(*this, KernelFunc, Range, Redu);
     shared_ptr_class<detail::queue_impl> QueueCopy = MQueue;
     this->finalize();
@@ -1073,26 +1056,16 @@ class __SYCL_EXPORT handler {
       handler AuxHandler(QueueCopy, MIsHost);
       AuxHandler.saveCodeLoc(MCodeLoc);
 
-      // The last kernel DOES write to reduction's accessor.
+      // The last kernel DOES write to user's accessor passed to reduction.
       // Associate it with handler manually.
-      if (NWorkGroups == 1)
+      if (NWorkGroups == 1 && !Reduction::is_usm)
         Redu.associateWithHandler(AuxHandler);
       intel::detail::reduAuxCGFunc<KernelName, KernelType>(AuxHandler, Range,
                                                            NWorkItems, Redu);
       MLastEvent = AuxHandler.finalize();
 
       NWorkItems = NWorkGroups;
     } // end while (NWorkItems > 1)
-
-    // Submit non-blocking copy from reduction accessor to user's reduction
-    // variable.
-    if (Reduction::is_usm) {
-      handler CopyHandler(QueueCopy, MIsHost);
-      CopyHandler.saveCodeLoc(MCodeLoc);
-      Redu.associateWithHandler(CopyHandler);
-      CopyHandler.copy(Redu.getUserAccessor(), Redu.getUSMPointer());
-      MLastEvent = CopyHandler.finalize();
-    }
   }
 
   /// Hierarchical kernel invocation method of a kernel defined as a lambda