Use our own implementation of parallel_memcopy (#7254)

src/ray/util/memory.cc

@@ -0,0 +1,50 @@
+#include "ray/util/memory.h"
+
+#include <cstring>
+#include <thread>
+#include <vector>
+
+namespace ray {
+
+uint8_t *pointer_logical_and(const uint8_t *address, uintptr_t bits) {
+  uintptr_t value = reinterpret_cast<uintptr_t>(address);
+  return reinterpret_cast<uint8_t *>(value & bits);
+}
+
+void parallel_memcopy(uint8_t *dst, const uint8_t *src, int64_t nbytes,
+                      uintptr_t block_size, int num_threads) {
+  std::vector<std::thread> threadpool(num_threads);
+  uint8_t *left = pointer_logical_and(src + block_size - 1, ~(block_size - 1));
+  uint8_t *right = pointer_logical_and(src + nbytes, ~(block_size - 1));
+  int64_t num_blocks = (right - left) / block_size;
+
+  // Update right address
+  right = right - (num_blocks % num_threads) * block_size;
+
+  // Now we divide these blocks between available threads. The remainder is
+  // handled on the main thread.
+  int64_t chunk_size = (right - left) / num_threads;
+  int64_t prefix = left - src;
+  int64_t suffix = src + nbytes - right;
+  // Now the data layout is | prefix | k * num_threads * block_size | suffix |.
+  // We have chunk_size = k * block_size, therefore the data layout is
+  // | prefix | num_threads * chunk_size | suffix |.
+  // Each thread gets a "chunk" of k blocks.
+
+  // Start all threads first and handle leftovers while threads run.
+  for (int i = 0; i < num_threads; i++) {
+    threadpool[i] = std::thread(std::memcpy, dst + prefix + i * chunk_size,
+                                left + i * chunk_size, chunk_size);
+  }
+
+  std::memcpy(dst, src, prefix);
+  std::memcpy(dst + prefix + num_threads * chunk_size, right, suffix);
+
+  for (auto &t : threadpool) {
+    if (t.joinable()) {
+      t.join();
+    }
+  }
+}
+
+}  // namespace ray

src/ray/util/memory.h

@@ -0,0 +1,15 @@
+#ifndef RAY_UTIL_MEMORY_H
+#define RAY_UTIL_MEMORY_H
+
+#include <stdint.h>
+
+namespace ray {
+
+// A helper function for doing memcpy with multiple threads. This is required
+// to saturate the memory bandwidth of modern cpus.
+void parallel_memcopy(uint8_t *dst, const uint8_t *src, int64_t nbytes,
+                      uintptr_t block_size, int num_threads);
+
+}  // namespace ray
+
+#endif  // RAY_UTIL_MEMORY_H