Add parallelTransformReduce and parallelForEachError

[rnk]

Summary

parallelTransformReduce is modelled on the C++17 pstl API of
std::transform_reduce, except our wrappers do not use execution policy
parameters.

parallelForEachError allows loops that contain potentially failing
operations to propagate errors out of the loop. This was one of the
major challenges I encountered while parallelizing PDB type merging in
LLD. Parallelizing a loop with parallelForEachError is not behavior
preserving: the loop will no longer stop on the first error, it will
continue working and report all errors it encounters in a list.

I plan to use this to propagate errors out of LLD's
coff::TpiSource::remapTpiWithGHashes, which currently stores errors an
error in the TpiSource object.

[rriddle]

Nice, this looks like a good addition. There's already at least one place in MLIR that I could see this used as well.

llvm/include/llvm/Support/Parallel.h

169

Can you check for the case where NumInputs == 0 and return Init? Any reason that we should require that there is always at least one task?

193

We could strip the first reduce given that we know that there will be at least one task, but I suppose the reduce function won't be costly in most cases.

[rnk]

• more tests, handle corner cases

[rnk]

More unit testing showed that I wasn't more than 1024 inputs well, so I fixed that. Take another look if you want to double check that. Thanks!

llvm/include/llvm/Support/Parallel.h

169

Let's just make it work. I think I hit this edge case somewhere in LLD for parallelForEach.

172

Hm, I guess this assert is wrong for empty inputs. Amusingly, LLVM hid the bug from me because integer divide by zero is UB. :)

193

Eh, let's do it.

llvm/include/llvm/Support/Parallel.h

 //===- llvm/Support/Parallel.h - Parallel algorithms ----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_SUPPORT_PARALLEL_H
 #define LLVM_SUPPORT_PARALLEL_H
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Config/llvm-config.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/Threading.h"
 
 #include <algorithm>
 #include <condition_variable>
 #include <functional>
 #include <mutex>
 
TG.spawn([=, &Fn] {
       for (IndexTy J = I, E = I + TaskSize; J != E; ++J)
         Fn(J);
     });
   }
   for (IndexTy J = I; J < End; ++J)
     Fn(J);
 }
 
+template <class IterTy, class ResultTy, class ReduceFuncTy,
+          class TransformFuncTy>
+ResultTy parallel_transform_reduce(IterTy Begin, IterTy End, ResultTy Init,
+                                   ReduceFuncTy Reduce,
+                                   TransformFuncTy Transform) {
+  // TaskGroup has a relatively high overhead, so we want to reduce
+  // the number of spawn() calls. We'll create up to 1024 tasks here.
+  // (Note that 1024 is an arbitrary number. This code probably needs
+  // improving to take the number of available cores into account.)
+  size_t NumInputs = std::distance(Begin, End);
+  size_t NumTasks = std::min(static_cast<size_t>(1024), NumInputs);
+  size_t TaskSize = NumInputs / NumTasks;
+  assert(TaskSize > 0);
+
+  std::vector<ResultTy> Results(NumTasks, Init);
+  {
+    TaskGroup TG;
+    for (size_t TaskId = 0; TaskId < NumTasks; ++TaskId) {
+      IterTy TBegin = Begin + TaskId * TaskSize;
+      IterTy TEnd = std::min(TBegin + TaskSize, End);
+      TG.spawn([=, &Transform, &Reduce, &Results] {
+        // Reduce the result of transformation eagerly within each task.
+        ResultTy R = Init;
+        for (IterTy It = TBegin; It != TEnd; ++It)
+          R = Reduce(R, Transform(*It));
+        Results[TaskId] = R;
+      });
+    }
+  }
+
+  // Do a final reduction. There are at most 1024 tasks, so this only adds
+  // constant single-threaded overhead for large inputs. Hopefully most
+  // reductions are cheaper than the transformation.
+  for (auto &R : Results)
+    Init = Reduce(Init, std::move(R));
+  return std::move(Init);
+}
+
 #endif
 
 } // namespace detail
 } // namespace parallel
 
 template <class RandomAccessIterator,
           class Comparator = std::less<
               typename std::iterator_traits<RandomAccessIterator>::value_type>>
if (parallel::strategy.ThreadsRequested != 1) {
     parallel::detail::parallel_for_each_n(Begin, End, Fn);
     return;
   }
 #endif
   for (size_t I = Begin; I != End; ++I)
     Fn(I);
 }
 
+template <class IterTy, class ResultTy, class ReduceFuncTy,
+          class TransformFuncTy>
+ResultTy parallelTransformReduce(IterTy Begin, IterTy End, ResultTy Init,
+                                 ReduceFuncTy Reduce,
+                                 TransformFuncTy Transform) {
+#if LLVM_ENABLE_THREADS
+  if (parallel::strategy.ThreadsRequested != 1) {
+    return parallel::detail::parallel_transform_reduce(Begin, End, Init, Reduce,
+                                                       Transform);
+  }
+#endif
+  for (IterTy I = Begin; I != End; ++I)
+    Init = Reduce(std::move(Init), Transform(*I));
+  return std::move(Init);
+}
+
 // Range wrappers.
 template <class RangeTy,
           class Comparator = std::less<decltype(*std::begin(RangeTy()))>>
 void parallelSort(RangeTy &&R, const Comparator &Comp = Comparator()) {
   parallelSort(std::begin(R), std::end(R), Comp);
 }
 
 template <class RangeTy, class FuncTy>
 void parallelForEach(RangeTy &&R, FuncTy Fn) {
   parallelForEach(std::begin(R), std::end(R), Fn);
 }
 
+template <class RangeTy, class ResultTy, class ReduceFuncTy,
+          class TransformFuncTy>
+ResultTy parallelTransformReduce(RangeTy &&R, ResultTy Init,
+                                 ReduceFuncTy Reduce,
+                                 TransformFuncTy Transform) {
+  return parallelTransformReduce(std::begin(R), std::end(R), Init, Reduce,
+                                 Transform);
+}
+
+// Parallel for-each, but with error handling.
+template <class RangeTy, class FuncTy>
+Error parallelForEachError(RangeTy &&R, FuncTy Fn) {
+  // The transform_reduce algorithm requires that the initial value be copyable.
+  // Error objects are uncopyable. We only need to copy initial success values,
+  // so work around this mismatch via the C API. The C API represents success
+  // values with a null pointer. The joinErrors discards null values and joins
+  // multiple errors into an ErrorList.
+  return unwrap(parallelTransformReduce(
+      std::begin(R), std::end(R), wrap(Error::success()),
+      [](LLVMErrorRef Lhs, LLVMErrorRef Rhs) {
+        return wrap(joinErrors(unwrap(Lhs), unwrap(Rhs)));
+      },
+      [&Fn](auto &&V) { return wrap(Fn(V)); }));
+}
+
 } // namespace llvm
 
 #endif // LLVM_SUPPORT_PARALLEL_H

llvm/unittests/Support/ParallelTest.cpp

TEST(Parallel, parallel_for) {
 
   uint32_t expected[2049];
   std::fill(expected, expected + 2049, 2);
   ASSERT_TRUE(std::equal(range, range + 2049, expected));
   // Check that we don't write past the end of the requested range.
   ASSERT_EQ(range[2049], 1u);
 }
 
+TEST(Parallel, TransformReduce) {
+  // Sum the lengths of these strings in parallel.
+  const char *strs[] = {"a", "ab", "abc", "abcd", "abcde", "abcdef"};
+  size_t lenSum =
+      parallelTransformReduce(strs, static_cast<size_t>(0), std::plus<size_t>(),
+                              [](const char *s) { return strlen(s); });
+  ASSERT_EQ(lenSum, static_cast<size_t>(21));
+}
+
+TEST(Parallel, ForEachError) {
+  int nums[] = {1, 2, 3, 4, 5, 6};
+  Error e = parallelForEachError(nums, [](int v) -> Error {
+    if ((v & 1) == 0)
+      return createStringError(std::errc::invalid_argument, "asdf");
+    return Error::success();
+  });
+  EXPECT_TRUE(e.isA<ErrorList>());
+  std::string errText = toString(std::move(e));
+  EXPECT_EQ(errText, std::string("asdf\nasdf\nasdf"));
+}
+
 #endif