Repository: gtn-org/gtn (language: python)

Issue:
Parallel compose for both CPU and GPU



Discussion / hints:

Fix patch:
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 6621c03..5194d3f 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -33,20 +33,24 @@ endif()
 # Core lib sources -- maybe make me a submodule eventually
 include(${CMAKE_SOURCE_DIR}/gtn/CMakeLists.txt)
 
+if (GTN_BUILD_CUDA)
+  # TODO, better way to set these arch flags
+  add_compile_options(
+    $<$<COMPILE_LANGUAGE:CUDA>:-arch=sm_80>
+    $<$<COMPILE_LANGUAGE:CUDA>:-default-stream=per-thread>
+    $<$<COMPILE_LANGUAGE:CUDA>:--extended-lambda>)
+endif()
+
 # Main lib
 add_library(
   gtn
   ${GTN_SOURCES}
 )
 
-# Experimental
+include(${CMAKE_SOURCE_DIR}/gtn/cuda/CMakeLists.txt)
 if (GTN_BUILD_CUDA)
-  include(${CMAKE_SOURCE_DIR}/gtn/cuda/CMakeLists.txt)
   target_sources(gtn PRIVATE ${GTN_CUDA_SOURCES})
-  # TODO, better way to set these arch flags
-  target_compile_options(gtn PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-arch=sm_80>)
 else()
-  include(${CMAKE_SOURCE_DIR}/gtn/cuda/CMakeLists.txt)
   target_sources(gtn PRIVATE ${GTN_NO_CUDA_SOURCES})
 endif()
 
diff --git a/benchmarks/CMakeLists.txt b/benchmarks/CMakeLists.txt
index f065a57..cfde48a 100644
--- a/benchmarks/CMakeLists.txt
+++ b/benchmarks/CMakeLists.txt
@@ -21,5 +21,8 @@ if (GTN_BUILD_BENCHMARKS)
   build_benchmark(${PROJECT_SOURCE_DIR}/benchmarks/graph.cpp)
   build_benchmark(${PROJECT_SOURCE_DIR}/benchmarks/functions.cpp)
   build_benchmark(${PROJECT_SOURCE_DIR}/benchmarks/parallel.cpp)
+  if (GTN_BUILD_CUDA)
+    build_benchmark(${PROJECT_SOURCE_DIR}/benchmarks/parallel_cuda.cu)
+  endif ()
 endif ()
 
diff --git a/benchmarks/functions.cpp b/benchmarks/functions.cpp
index 44a7f6c..84d2ea3 100644
--- a/benchmarks/functions.cpp
+++ b/benchmarks/functions.cpp
@@ -70,33 +70,35 @@ void timeSimpleOps(Device device = Device::CPU) {
   TIME(concatBackward);
 }
 
-void timeForward() {
-  auto graph = linearGraph(1000, 1000);
+void timeForward(Device device = Device::CPU) {
+  auto graph = linearGraph(100, 20000);
   std::vector<float> weights(graph.numArcs());
   std::generate(weights.begin(), weights.end(), std::rand);
   graph.setWeights(weights.data());
+  graph = graph.to(device);
+
   auto forwardScoreLinearForward = [&graph]() {
     auto out = forwardScore(graph);
   };
-  TIME(forwardScoreLinearForward);
+  TIME_DEVICE(forwardScoreLinearForward, device);
 
   auto forwardScoreLinearBackward = [&graph, out = forwardScore(graph)] {
     graph.zeroGrad();
     backward(out, true);
   };
-  TIME(forwardScoreLinearBackward);
+  TIME_DEVICE(forwardScoreLinearBackward, device);
 
-  graph = makeRandomDAG(20000, 200000);
+  graph = makeRandomDAG(500, 400000).to(device);
   auto forwardScoreRandDAGForward = [&graph]() {
     auto out = forwardScore(graph);
   };
-  TIME(forwardScoreRandDAGForward);
+  TIME_DEVICE(forwardScoreRandDAGForward, device);
 
   auto forwardScoreRandDAGBackward = [&graph, out = forwardScore(graph)] {
     graph.zeroGrad();
     backward(out, true);
   };
-  TIME(forwardScoreRandDAGBackward);
+  TIME_DEVICE(forwardScoreRandDAGBackward, device);
 }
 
 void timeCompose(Device device = Device::CPU) {
@@ -144,6 +146,7 @@ int main() {
   timeCompose();
   if (cuda::isAvailable()) {
     timeSimpleOps(Device::CUDA);
+    timeForward(Device::CUDA);
     timeCompose(Device::CUDA);
   }
 }
diff --git a/benchmarks/parallel.cpp b/benchmarks/parallel.cpp
index 0716733..b185ad9 100644
--- a/benchmarks/parallel.cpp
+++ b/benchmarks/parallel.cpp
@@ -12,7 +12,7 @@
 
 using namespace gtn;
 
-void timeParallelCompose(const int B) {
+void timeParallelCompose(const int B, Device device = Device::CPU) {
   const int N1 = 100;
   const int N2 = 50;
   const int A1 = 20;
@@ -29,15 +29,15 @@ void timeParallelCompose(const int B) {
         s.addArc(i, i, j);
       }
     }
-    graphs1.push_back(f);
-    graphs2.push_back(s);
+    graphs1.push_back(f.to(device));
+    graphs2.push_back(s.to(device));
   }
 
   auto composeParallel = [&graphs1, &graphs2]() {
     return parallelMap(compose, graphs1, graphs2);
   };
 
-  TIME(composeParallel);
+  TIME_DEVICE(composeParallel, device);
 
   auto out = parallelMap(compose, graphs1, graphs2);
   std::vector<bool> retainGraph({true});
@@ -46,7 +46,34 @@ void timeParallelCompose(const int B) {
         static_cast<void (*)(Graph, bool)>(&backward), out, retainGraph);
   };
 
-  TIME(backwardParallel);
+  TIME_DEVICE(backwardParallel, device);
+}
+
+void timeParallelForward(const int B, Device device = Device::CPU) {
+  std::vector<Graph> graphs;
+  for (int b = 0; b < B; b++) {
+    graphs.push_back(linearGraph(50, 1000, device));
+  }
+
+  auto forwardScoreParallel = [&graphs]() {
+    parallelMap(forwardScore, graphs);
+  };
+
+  TIME_DEVICE(forwardScoreParallel, device);
+}
+
+
+void timeParallelClone(const int B, Device device = Device::CPU) {
+  std::vector<Graph> graphs;
+  for (int b = 0; b < B; b++) {
+    graphs.push_back(linearGraph(10, 1000, device));
+  }
+
+  auto cloneParallel = [&graphs]() {
+    parallelMap(projectInput, graphs);
+  };
+
+  TIME_DEVICE(cloneParallel, device);
 }
 
 int main(int argc, char** argv) {
@@ -59,5 +86,14 @@ int main(int argc, char** argv) {
   if (argc > 1) {
     B = std::stoi(argv[1]);
   }
+  std::cout << "Batch size " << B << " with " << detail::getNumViableThreads(B)
+    << " threads." << std::endl;
+  timeParallelClone(B);
   timeParallelCompose(B);
+  timeParallelForward(B);
+  if (cuda::isAvailable()) {
+    timeParallelClone(B, Device::CUDA);
+    timeParallelCompose(B, Device::CUDA);
+    timeParallelForward(B, Device::CUDA);
+  }
 }
diff --git a/benchmarks/parallel_cuda.cu b/benchmarks/parallel_cuda.cu
new file mode 100644
index 0000000..8923017
--- /dev/null
+++ b/benchmarks/parallel_cuda.cu
@@ -0,0 +1,49 @@
+#include <iostream>
+#include <string>
+
+#include "gtn/gtn.h"
+#include "benchmarks/time_utils.h"
+
+using namespace gtn;
+
+__global__
+void longKernel(float *x, int n) {
+  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  for (int i = tid; i < n; i += blockDim.x * gridDim.x) {
+    x[i] = sqrt(pow(3.14159,i));
+  }
+}
+
+void launchKernel(int size) {
+  float *data;
+  cudaMallocAsync(&data, size * sizeof(float), 0);
+  longKernel<<<1, 64>>>(data, size);
+  cudaFreeAsync(data, 0);
+}
+
+void timeParallelCuda(const int B) {
+  std::vector<int> sizes(B, 1 << 20);
+
+  auto timeLongCudaKernel = [&sizes]() {
+    parallelMapDevice(Device::CUDA, launchKernel, sizes); 
+  };
+
+  TIME_DEVICE(timeLongCudaKernel, Device::CUDA);
+}
+
+int main(int argc, char** argv) {
+  /**
+   * Usage:
+   *   `./benchmark_parallel_cuda <batch_size (default 8)>`
+   */
+
+  int B = 8; // batch size
+  if (argc > 1) {
+    B = std::stoi(argv[1]);
+  }
+  std::cout << "Batch size " << B << " with " << detail::getNumViableThreads(B)
+    << " threads." << std::endl;
+  if (cuda::isAvailable()) {
+    timeParallelCuda(B);
+  }
+}
diff --git a/benchmarks/time_utils.h b/benchmarks/time_utils.h
index 911cddf..1f4d0cb 100644
--- a/benchmarks/time_utils.h
+++ b/benchmarks/time_utils.h
@@ -21,7 +21,7 @@ using namespace gtn;
 
 #define TIME_DEVICE(FUNC, DEVICE) \
   { \
-    auto deviceName = device.isCuda() ? "(cuda)" : "(cpu)"; \
+    auto deviceName = Device(DEVICE).isCuda() ? "(cuda)" : "(cpu)"; \
     std::cout << "Timing " << #FUNC << " " << deviceName << " ...  " << std::flush; \
     std::cout << std::setprecision(5) << timeit(FUNC, true) << " msec" << std::endl; \
   }
@@ -31,7 +31,7 @@ using namespace gtn;
   std::cout << std::setprecision(5) << timeit(FUNC) << " msec" << std::endl;
 
 #define milliseconds(x) \
-  std::chrono::duration_cast<std::chrono::milliseconds>(x).count()
+  (std::chrono::duration_cast<std::chrono::nanoseconds>(x).count() / 1e6)
 #define timeNow() std::chrono::high_resolution_clock::now()
 
 double timeit(std::function<void()> fn, bool isCuda = false) {
@@ -47,9 +47,9 @@ double timeit(std::function<void()> fn, bool isCuda = false) {
   auto start = timeNow();
   for (int i = 0; i < numIters; i++) {
     fn();
-  }
-  if (isCuda) {
-    cuda::synchronize();
+    if (isCuda) {
+      cuda::synchronize();
+    }
   }
   auto end = timeNow();
   return milliseconds(end - start) / static_cast<double>(numIters);
@@ -61,9 +61,8 @@ Graph makeRandomDAG(int num_nodes, int num_arcs) {
   graph.addNode(true);
   for (int n = 1; n < num_nodes; n++) {
     graph.addNode(false, n == num_nodes - 1);
-    graph.addArc(n - 1, n, 0); // assure graph is connected
   }
-  for (int i = 0; i < num_arcs - num_nodes + 1; i++) {
+  for (int i = 0; i < num_arcs; i++) {
     // To preserve DAG property, select src then select dst to be
     // greater than source.
     // select from [0, num_nodes-2]:
diff --git a/gtn/CMakeLists.txt b/gtn/CMakeLists.txt
index 2dd61f4..2290e5b 100644
--- a/gtn/CMakeLists.txt
+++ b/gtn/CMakeLists.txt
@@ -26,6 +26,7 @@ set(
   ${CMAKE_CURRENT_SOURCE_DIR}/gtn/hd_span.h
   ${CMAKE_CURRENT_SOURCE_DIR}/gtn/parallel/parallel_map.h
   ${CMAKE_CURRENT_SOURCE_DIR}/gtn/parallel/parallel_map.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/gtn/parallel/thread_pool.h
   ${CMAKE_CURRENT_SOURCE_DIR}/gtn/rand.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/gtn/rand.h
   ${CMAKE_CURRENT_SOURCE_DIR}/gtn/utils.cpp
diff --git a/gtn/cuda/CMakeLists.txt b/gtn/cuda/CMakeLists.txt
index 878b0a2..c9cdde2 100644
--- a/gtn/cuda/CMakeLists.txt
+++ b/gtn/cuda/CMakeLists.txt
@@ -15,4 +15,5 @@ set(
   ${CMAKE_CURRENT_SOURCE_DIR}/gtn/cuda/cuda.cu
   ${CMAKE_CURRENT_SOURCE_DIR}/gtn/cuda/functions.h
   ${CMAKE_CURRENT_SOURCE_DIR}/gtn/cuda/functions.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/gtn/cuda/shortest.cu
 )
diff --git a/gtn/cuda/compose.cu b/gtn/cuda/compose.cu
index 57615d3..cc270f5 100644
--- a/gtn/cuda/compose.cu
+++ b/gtn/cuda/compose.cu
@@ -941,4 +941,4 @@ Graph compose(const Graph& first, const Graph& second) {
 
 } // namespace detail
 } // namespace cuda 
-} // namespace gtn
\ No newline at end of file
+} // namespace gtn
diff --git a/gtn/cuda/cuda.cu b/gtn/cuda/cuda.cu
index 43acef7..df93767 100644
--- a/gtn/cuda/cuda.cu
+++ b/gtn/cuda/cuda.cu
@@ -37,23 +37,51 @@ void synchronize(int device) {
   synchronize();
 }
 
+void synchronizeStream() {
+  CUDA_CHECK(cudaStreamSynchronize(0));
+}
+
+Event::Event() {
+  CUDA_CHECK(cudaEventCreateWithFlags(&event_, cudaEventDisableTiming));
+}
+
+Event::~Event() {
+  CUDA_CHECK(cudaEventDestroy(event_));
+}
+
+void Event::record() {
+  CUDA_CHECK(cudaEventRecord(event_));
+}
+
+void Event::synchronize() {
+  CUDA_CHECK(cudaEventSynchronize(event_));
+}
+
+void Event::wait() {
+  CUDA_CHECK(cudaStreamWaitEvent(0, event_));
+}
+
 namespace detail {
 
 void negate(const float* in, float* out, size_t size) {
   thrust::transform(
-      thrust::device, in, in + size, out, thrust::negate<float>());
+    thrust::device, in, in + size, out, thrust::negate<float>());
 }
 
 void add(const float* a, const float* b, float* out, size_t size) {
-  thrust::device_ptr<const float> aPtr(a);
-  thrust::device_ptr<const float> bPtr(b);
-  thrust::device_ptr<float> outPtr(out);
-  thrust::transform(aPtr, aPtr + size, bPtr, outPtr, thrust::plus<float>());
+  transform(
+    a, a + size, b, out,
+    [] __device__ (const float lhs, const float rhs) {
+      return lhs + rhs;
+    });
 }
 
 void subtract(const float* a, const float* b, float* out, size_t size) {
-  thrust::transform(
-      thrust::device, a, a + size, b, out, thrust::minus<float>());
+  transform(
+    a, a + size, b, out,
+    [] __device__ (const float lhs, const float rhs) {
+      return lhs - rhs;
+    });
 }
 
 void copy(void* dst, const void* src, size_t size) {
@@ -63,12 +91,12 @@ void copy(void* dst, const void* src, size_t size) {
 void* allocate(size_t size, int device) {
   DeviceManager dm(device);
   void *res;
-  CUDA_CHECK(cudaMalloc(&res, size));
+  CUDA_CHECK(cudaMallocAsync(&res, size, 0));
   return res;
 }
 
 void free(void* ptr) {
-  CUDA_CHECK(cudaFree(ptr));
+  CUDA_CHECK(cudaFreeAsync(ptr, 0));
 }
 
 void cudaCheck(cudaError_t err, const char* file, int line) {
diff --git a/gtn/cuda/cuda.h b/gtn/cuda/cuda.h
index 402d2bb..d06000f 100644
--- a/gtn/cuda/cuda.h
+++ b/gtn/cuda/cuda.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <cstddef>
+#include <functional>
 
 #if defined(_CUDA_)
 #include <cuda.h>
@@ -37,11 +38,49 @@ int getDevice();
 void setDevice(int device);
 
 /**
- * Waits unitl all kernels in all streams are complete.
+ * Waits in the currently active device until all kernels in all streams are
+ * complete.
  */
 void synchronize();
+
+/**
+ * Waits in the given device until all kernels in all streams are complete.
+ */
 void synchronize(int device);
 
+/**
+ * Synchronize the default stream.
+ */
+void synchronizeStream();
+
+/**
+ * A class wrapper for a CUDA event.
+ */
+class Event {
+ public:
+  Event();
+  ~Event();
+
+  /**
+   * Record the event in the default stream.
+   */
+  void record();
+
+  /**
+   * Synchronize the host thread with the event.
+   */
+  void synchronize();
+
+  /**
+   * Wait for the event in the defualt stream.
+   */
+  void wait();
+#if defined(_CUDA_)
+ private:
+  cudaEvent_t event_;
+#endif
+};
+
 /** @} */
 
 namespace detail {
@@ -60,6 +99,12 @@ class DeviceManager {
   int device_;
 };
 
+void transform(
+    const float* aBegin,
+    const float* aEnd,
+    const float* bBegin,
+    float* outBegin,
+    const std::function<float>& op);
 void negate(const float* in, float* out, size_t size);
 void add(const float* a, const float* b, float* out, size_t size);
 void subtract(const float* a, const float* b, float* out, size_t size);
@@ -77,6 +122,38 @@ void fill(bool* dst, bool val, size_t size);
 
 #if defined(_CUDA_)
 void cudaCheck(cudaError_t err, const char* file, int line);
+#if defined(__NVCC__)
+namespace {
+template <typename T, typename F>
+__global__ void transformKernel(
+    const T* a,
+    const T* aEnd,
+    const T* b,
+    T* out,
+    const F op) {
+  const int gTid = blockIdx.x * blockDim.x + threadIdx.x;
+  int size = aEnd - a;
+  if (gTid < size) {
+    out[gTid] = op(a[gTid], b[gTid]);
+  }
+}
+template <typename T, typename F>
+void transform(
+    const T* aBegin,
+    const T* aEnd,
+    const T* bBegin,
+    T* outBegin,
+    const F op) {
+  int size = aEnd - aBegin;
+  if (size == 0) {
+    return;
+  }
+  const int NT = 128;
+  const int blocks = (size + NT - 1) / NT;
+  transformKernel<<<blocks, NT>>>(aBegin, aEnd, bBegin, outBegin, op);
+}
+}
+#endif
 #endif
 
 } // namespace detail
diff --git a/gtn/cuda/functions.cpp b/gtn/cuda/functions.cpp
index cc997de..fa69381 100644
--- a/gtn/cuda/functions.cpp
+++ b/gtn/cuda/functions.cpp
@@ -5,6 +5,7 @@ namespace cuda {
 
 namespace detail {
   Graph compose(const Graph& g1, const Graph& g2);
+  Graph shortestDistance(const Graph& g, bool tropical);
 } // namespace detail
 
 Graph concat(const std::vector<Graph>& graphs) {
@@ -32,11 +33,11 @@ Graph remove(const Graph& g, int ilabel, int olabel) {
 }
 
 Graph forwardScore(const Graph& g) {
-  throw std::logic_error("[cuda::forwardScore] GPU function not implemented.");
+  return cuda::detail::shortestDistance(g, false);
 }
 
 Graph viterbiScore(const Graph& g) {
-  throw std::logic_error("[cuda::viterbiScore] GPU function not implemented.");
+  return cuda::detail::shortestDistance(g, true);
 }
 
 Graph viterbiPath(const Graph& g) {
diff --git a/gtn/cuda/nocuda.cpp b/gtn/cuda/nocuda.cpp
index 348319c..722942a 100644
--- a/gtn/cuda/nocuda.cpp
+++ b/gtn/cuda/nocuda.cpp
@@ -35,14 +35,31 @@ void synchronize(int device) {
   throw std::logic_error("[cuda::synchronize] CUDA not available.");
 }
 
+void synchronizeStream() {
+  throw std::logic_error("[cuda::synchronizeStream] CUDA not available.");
+}
+
+Event::Event() {
+  throw std::logic_error("[cuda::Event] CUDA not available.");
+}
+
+Event::~Event() { }
+void Event::record() { }
+void Event::synchronize() { }
+void Event::wait() { }
+
 namespace detail {
 
 void add(const float* a, const float* b, float* out, size_t size) {
   throw std::logic_error("[cuda::detail::add] CUDA not available.");
 }
 
-float* ones(size_t size, int device) {
-  throw std::logic_error("[cuda::detail::ones] CUDA not available.");
+void subtract(const float* a, const float* b, float* out, size_t size) {
+  throw std::logic_error("[cuda::detail::subtract] CUDA not available.");
+}
+
+void negate(const float* in, float* out, size_t size) {
+  throw std::logic_error("[cuda::detail::negate] CUDA not available.");
 }
 
 void copy(void* dst, const void* src, size_t size) {
@@ -124,6 +141,14 @@ Graph viterbiPath(const Graph& g) {
   throw std::logic_error("[cuda::viterbiPath] CUDA not available.");
 }
 
+Graph scalarGraph(float val, bool calcGrad, Device device) {
+  throw std::logic_error("[cuda::scalarGraph] CUDA not available.");
+}
+
+Graph linearGraph(int M, int N, bool calcGrad, Device device) {
+  throw std::logic_error("[cuda::linearGraph] CUDA not available.");
+}
+
 
 } // namespace cuda
 } // namespace gtn
diff --git a/gtn/hd_span.h b/gtn/hd_span.h
index 6e02506..7258974 100644
--- a/gtn/hd_span.h
+++ b/gtn/hd_span.h
@@ -139,8 +139,8 @@ class HDSpan {
   };
 
   void resize(size_t size) {
-    // Smaller or same size is a no-op with capacity_ unchanged.
-    if (size <= size_) {
+    // If size is smaller than capacity just update the size.
+    if (size <= capacity_) {
       size_ = size;
       return;
     }
@@ -214,6 +214,9 @@ class HDSpan {
     capacity_ = 0;
   }
 
+  template <typename U>
+  friend void swap(HDSpan<U>& lhs, HDSpan<U>& rhs);
+
  private:
   T* data_{nullptr};
   size_t size_{0};
@@ -236,6 +239,18 @@ bool isSameDevice(const HDSpan<T>& a, const HDSpan<T>& b, const HDSpan<T>& c) {
 
 } // namespace
 
+template <typename T>
+void swap(HDSpan<T>& lhs, HDSpan<T>& rhs) {
+  if (!isSameDevice(lhs, rhs)) {
+    throw std::invalid_argument(
+        "[swap] Objects must be on the same device");
+  }
+  std::swap(lhs.data_, rhs.data_);
+  std::swap(lhs.size_, rhs.size_);
+  std::swap(lhs.capacity_, rhs.capacity_);
+}
+
+
 template <typename T>
 bool operator==(const HDSpan<T>& lhs, const HDSpan<T>& rhs) {
   if (!isSameDevice(lhs, rhs)) {
diff --git a/gtn/parallel/parallel_map.h b/gtn/parallel/parallel_map.h
index 02650ae..bc9e0e8 100644
--- a/gtn/parallel/parallel_map.h
+++ b/gtn/parallel/parallel_map.h
@@ -17,6 +17,7 @@
 #include <vector>
 
 #include "gtn/graph.h"
+#include "gtn/device.h"
 #include "gtn/parallel/thread_pool.h"
 
 namespace gtn {
@@ -136,11 +137,38 @@ struct OutPayload<void> {
   void value() const {}
 };
 
+template <typename T>
+bool anyCuda(const T&) {
+  return false;
+}
+
+template <>
+bool anyCuda<std::vector<Graph>>(const std::vector<Graph>& graphs) {
+  bool ret = false;
+  for (auto& g : graphs) {
+    ret |= g.device().isCuda();
+  }
+  return ret;
+}
+
+// Base case
+template <typename T>
+bool cudaCheck(T isCuda) {
+  return isCuda;
+}
+
+// Find out if any of a variable number of possible inputs are CUDA
+template <typename T, typename... Rest>
+bool cudaCheck(T isCuda, Rest... rest) {
+  return isCuda || cudaCheck(rest...);
+}
+
 } // namespace
 
 /**
- * Executes a function in parallel.
+ * Executes a function in parallel on the given device.
  *
+ * @param[in] device The specified device (e.g. `Device::CUDA`)
  * @param[in] function A function pointer to execute in parallel
  * @param[in] ...inputs variadic arguments of iterable/indexable containers,
  * such as `std::vector`s, i.e. `vector<T1>, vector<T2>,...`. Types must match
@@ -151,7 +179,10 @@ struct OutPayload<void> {
  * `function`. If the given function returns `void`, the return type is `void`.
  */
 template <typename FuncType, typename... Args>
-auto parallelMap(FuncType&& function, Args&&... inputs) {
+auto parallelMapDevice(
+    const Device& device, FuncType&& function, Args&&... inputs) {
+  bool isCuda = device.isCuda();
+
   // Maximum input size in number of elements
   const auto size = max(getSize(inputs)...);
 
@@ -165,28 +196,62 @@ auto parallelMap(FuncType&& function, Args&&... inputs) {
   std::mutex eMutex;
 
   for (size_t i = 0; i < size; ++i) {
-    futures[i] = threadPool.get().enqueue(
-        [size, i, &function, &eMutex, &eQueue](Args&&... inputs) -> OutType {
-          try {
-            return function(getIdxOrBroadcast(size, i, inputs)...);
-          } catch (...) {
-            std::unique_lock<std::mutex> m(eMutex);
-            eQueue.push(std::current_exception());
-            return OutType();
-          }
-        },
-        std::forward<Args>(inputs)...);
+    auto task = [size, i, &function, &eMutex, &eQueue](
+        Args&&... inputs) -> OutType {
+            try {
+              return function(getIdxOrBroadcast(size, i, inputs)...);
+            } catch (...) {
+              std::unique_lock<std::mutex> m(eMutex);
+              eQueue.push(std::current_exception());
+              return OutType();
+            }
+        };
+    if (isCuda) {
+      futures[i] = threadPool.get().enqueueIndex(
+          i, task, std::forward<Args>(inputs)...);
+    } else {
+      futures[i] = threadPool.get().enqueue(
+          task, std::forward<Args>(inputs)...);
+    }
   }
-
   // Waits until work is done
   auto out = OutPayload<OutType>(futures);
   while (!eQueue.empty()) {
     std::rethrow_exception(eQueue.front());
     eQueue.pop();
   }
+  if (isCuda) {
+    threadPool.get().syncStreams();
+  }
   return out.value();
 }
 
+/**
+ * Executes a function in parallel.
+ *
+ * @param[in] function A function pointer to execute in parallel
+ * @param[in] ...inputs variadic arguments of iterable/indexable containers,
+ * such as `std::vector`s, i.e. `vector<T1>, vector<T2>,...`. Types must match
+ * the input types of function exactly, i.e. `function` must take arguments
+ * ``T1, T2,...``.
+ *
+ * @return a vector of type `T` where `T` is the type of the output type of
+ * `function`. If the given function returns `void`, the return type is `void`.
+ *
+ * If any of the input types are graphs then the device is inferred from the
+ * device of the graph. The execution will prefer using `Device::CUDA` if the
+ * devices are mixed.
+ */
+template <typename FuncType, typename... Args>
+auto parallelMap(FuncType&& function, Args&&... inputs) {
+  // True if any of the inputs are Graphs on a CUDA device
+  bool isCuda = cudaCheck(anyCuda(inputs)...);
+  return parallelMapDevice(
+      isCuda ? Device(Device::CUDA, 0) : Device::CPU,
+      std::forward<FuncType&&>(function),
+      std::forward<Args&&>(inputs)...);
+}
+
 /**
  * @}
  */
diff --git a/gtn/parallel/thread_pool.h b/gtn/parallel/thread_pool.h
index 3e260a5..809e5ff 100644
--- a/gtn/parallel/thread_pool.h
+++ b/gtn/parallel/thread_pool.h
@@ -17,8 +17,10 @@
 #include <thread>
 #include <vector>
 
-namespace gtn {
+#include "gtn/cuda/cuda.h"
 
+namespace gtn {
+namespace detail {
 /**
  * A simple thread pool implementation from
  * https://github.com/progschj/ThreadPool for use in benchmarking
@@ -27,29 +29,75 @@ namespace gtn {
 class ThreadPool {
  public:
   ThreadPool(size_t threads) : stop(false) {
+    tasks.resize(threads);
+    if (cuda::isAvailable()) {
+      events.resize(threads);
+    }
     for (size_t i = 0; i < threads; ++i)
-      workers.emplace_back([this] {
+      workers.emplace_back([this, i] {
         for (;;) {
           std::function<void()> task;
-
           {
-            std::unique_lock<std::mutex> lock(this->queue_mutex);
-            this->condition.wait(
-                lock, [this] { return this->stop || !this->tasks.empty(); });
-            if (this->stop && this->tasks.empty())
+            std::unique_lock<std::mutex> lock(queue_mutex);
+            condition.wait(
+              lock, [this, i] {
+                return stop || !tasks[i].empty() || !sharedTasks.empty();
+              });
+            if (stop && tasks[i].empty() && sharedTasks.empty()) {
+              if (cuda::isAvailable()) {
+                cuda::synchronizeStream();
+              }
               return;
-            task = std::move(this->tasks.front());
-            this->tasks.pop();
+            }
+            if (!tasks[i].empty()) {
+              task = std::move(tasks[i].front());
+              tasks[i].pop();
+            } else {
+              task = std::move(sharedTasks.front());
+              sharedTasks.pop();
+            }
           }
-
           task();
         }
       });
   }
+  void syncStreams() {
+    // Every thread records an event, then the stream of each
+    // thread waits on all events
+    std::vector<std::future<void>> results;
+    for (int i = 0; i < events.size(); i++) {
+      auto func = [&event = events[i]]() { event.record(); };
+      auto task = std::make_shared<std::packaged_task<void()>>(func);
+      results.push_back(task->get_future());
+      std::unique_lock<std::mutex> lock(queue_mutex);
+      tasks[i].emplace([task]() { (*task)(); });
+    }
+    condition.notify_all();
+
+    for (auto& f : results) {
+      f.get();
+    }
+
+    for (int i = 0; i < events.size(); i++) {
+      auto task = [this]() {
+        for (auto& event : events) {
+          event.wait();
+        }
+      };
+      // Wait on events in the main stream as well
+      task();
+      std::unique_lock<std::mutex> lock(queue_mutex);
+      tasks[i].push(task);
+    }
+    condition.notify_all();
+  }
 
   template <class F, class... Args>
-  auto enqueue(F&& f, Args&&... args)
+  auto enqueueIndex(int taskNum, F&& f, Args&&... args)
       -> std::future<typename std::result_of<F(Args...)>::type> {
+    if (cuda::isAvailable()) {
+      cuda::synchronizeStream();
+    }
     using return_type = typename std::result_of<F(Args...)>::type;
 
     auto task = std::make_shared<std::packaged_task<return_type()>>(
@@ -63,31 +111,46 @@ class ThreadPool {
       if (stop)
         throw std::runtime_error("enqueue on stopped ThreadPool");
 
-      tasks.emplace([task]() { (*task)(); });
+      if (taskNum >= 0) {
+        auto tIdx = taskNum % tasks.size();
+        tasks[tIdx].emplace([task]() { (*task)(); });
+      } else {
+        sharedTasks.emplace([task]() { (*task)(); });
+      }
     }
-    condition.notify_one();
+    condition.notify_all();
     return res;
   }
+
+  template <class F, class... Args>
+  auto enqueue(F&& f, Args&&... args)
+      -> std::future<typename std::result_of<F(Args...)>::type> {
+    return enqueueIndex(-1, std::forward<F>(f), std::forward<Args>(args)...);
+  }
   ~ThreadPool() {
     {
       std::unique_lock<std::mutex> lock(queue_mutex);
       stop = true;
     }
     condition.notify_all();
-    for (std::thread& worker : workers)
+    for (std::thread& worker : workers) {
       worker.join();
+    }
   }
 
  private:
   // need to keep track of threads so we can join them
   std::vector<std::thread> workers;
-  // the task queue
-  std::queue<std::function<void()>> tasks;
+  // the task queues
+  std::queue<std::function<void()>> sharedTasks;
+  std::vector<std::queue<std::function<void()>>> tasks;
 
   // synchronization
   std::mutex queue_mutex;
   std::condition_variable condition;
+  std::vector<cuda::Event> events;
   bool stop;
 };
 
+} // namespace detail
 } // namespace gtn


Test patch:
diff --git a/bindings/python/test/test_cuda.py b/bindings/python/test/test_cuda.py
index 2f7a2ef..cd5b9a9 100644
--- a/bindings/python/test/test_cuda.py
+++ b/bindings/python/test/test_cuda.py
@@ -32,7 +32,7 @@ def test_graph_cuda(self):
       self.assertEqual(gdev.num_arcs(), g.num_arcs())
       self.assertTrue(gdev.is_cuda())
       self.assertEqual(gdev.item(), 0.5)
-      # self.assertRaises(ValueError, gdev.arc_sort)
+      self.assertRaises(ValueError, gdev.arc_sort)
 
       if gtn.cuda.device_count() > 1:
         gpu1 = gtn.Device(gtn.CUDA, 1) 
diff --git a/gtn/cuda/shortest.cu b/gtn/cuda/shortest.cu
new file mode 100644
index 0000000..47bdb52
--- /dev/null
+++ b/gtn/cuda/shortest.cu
@@ -0,0 +1,415 @@
+#include <thrust/device_vector.h>
+#include <thrust/fill.h>
+
+#include "gtn/creations.h"
+#include "gtn/cuda/cuda.h"
+#include "gtn/cuda/functions.h"
+#include "gtn/hd_span.h"
+
+using namespace gtn::detail;
+
+namespace gtn {
+namespace cuda {
+namespace detail {
+
+namespace{
+
+typedef Graph::SharedGraph GraphData;
+
+__constant__ constexpr float kInf = std::numeric_limits<float>::infinity();
+__constant__ constexpr float kNegInf = -std::numeric_limits<float>::infinity();
+constexpr int kNT = 128;
+
+typedef union {
+  float2 val; // val.x stores max
+  int2 idx; // val.y stores argmax
+  unsigned long long int data; // used for atomics
+} ArgMaxT;
+
+__device__ unsigned long long int atomicArgMax(
+    ArgMaxT* address, float val, int idx) {
+  auto laddr = (unsigned long long int*) address;
+  ArgMaxT old = *address;
+  ArgMaxT curr;
+  curr.val.x = val;
+  curr.idx.y = idx;
+  while (curr.val.x > (*address).val.x) {
+    old.data = atomicCAS(laddr, old.data, curr.data);
+  }
+  return old.data;
+}
+
+struct GradInfo {
+
+  HDSpan<float> scores;
+  HDSpan<ArgMaxT> maxes;
+  HDSpan<float> score;
+  HDSpan<ArgMaxT> maxAndIdx;
+
+  GradInfo(
+      const HDSpan<float>& scores,
+      const HDSpan<ArgMaxT>& maxes,
+      const HDSpan<float>& score,
+      const HDSpan<ArgMaxT>& maxAndIdx) :
+    scores(scores), maxes(maxes), score(score), maxAndIdx(maxAndIdx) { };
+
+  ~GradInfo() {
+    scores.clear();
+    maxes.clear();
+    score.clear();
+    maxAndIdx.clear();
+  };
+};
+
+__host__ __device__
+inline int divUp(int x, int y) {
+  return (x + y - 1) / y;
+}
+
+__global__
+void initExploreKernel(
+    const HDSpan<bool> flags,
+    const HDSpan<int> arcOffsets,
+    HDSpan<int> degrees,
+    HDSpan<int> toExplore,
+    HDSpan<int> count) {
+  const int gTid = blockIdx.x * blockDim.x + threadIdx.x;
+  if (gTid < degrees.size()) {
+    degrees[gTid] = arcOffsets[gTid + 1] - arcOffsets[gTid];
+    if (flags[gTid] && degrees[gTid] == 0) {
+      int oldCount = atomicAdd(count.data(), 1);
+      toExplore[oldCount] = gTid;
+    }
+  } 
+}
+
+void initExplore(
+    const HDSpan<bool>& flags,
+    const HDSpan<int>& arcOffsets,
+    HDSpan<int>& degrees,
+    HDSpan<int>& toExplore,
+    HDSpan<int>& numNodes) {
+  int numNodesH = 0;
+  if (flags.size() > 0) {
+    int blocks = divUp(flags.size(), kNT);
+    initExploreKernel<<<blocks, kNT>>>(
+        flags, arcOffsets, degrees, toExplore, numNodes);
+    CUDA_CHECK(
+      cudaMemcpy((void*)&numNodesH,
+      (void*)numNodes.begin(),
+      sizeof(int),
+      cudaMemcpyDeviceToHost));
+  }
+  toExplore.resize(numNodesH);
+} 
+
+__global__
+void reduceAccept(
+    const HDSpan<float> scores,
+    const HDSpan<int> acceptIds,
+    HDSpan<float> score,
+    HDSpan<ArgMaxT> maxAndIdx,
+    bool tropical) {
+  score[0] = 0.0;
+  maxAndIdx[0] = {kNegInf, -1};
+  __syncthreads();
+
+  for (int gTid = threadIdx.x; gTid < acceptIds.size(); gTid += blockDim.x) {
+    int nodeIdx = acceptIds[gTid];
+    atomicArgMax(maxAndIdx.data(), scores[nodeIdx], gTid);
+  }
+  __syncthreads();
+  float maxScore = maxAndIdx[0].val.x;
+  if (tropical || maxScore == kInf || maxScore == kNegInf) {
+    score[0] = maxScore;
+  } else {
+    for (int gTid = threadIdx.x; gTid < acceptIds.size(); gTid += blockDim.x) {
+      int nodeIdx = acceptIds[gTid];
+      atomicAdd(score.data(), expf(scores[nodeIdx] - maxScore));
+    }
+    __syncthreads();
+    if (threadIdx.x == 0) {
+      score[0] = logf(score[0]) + maxScore;
+    }
+  }
+}
+
+__global__
+void reduceArcScoresKernel(
+    const GraphData g,
+    const HDSpan<int> currExplore,
+    HDSpan<float> scores,
+    HDSpan<ArgMaxT> maxes,
+    const float* weights,
+    bool tropical) {
+  // max and score should be in shared memory
+  int nodeIdx = currExplore[blockIdx.x];
+  int numArcs = g.inArcOffset[nodeIdx + 1] - g.inArcOffset[nodeIdx]; 
+  numArcs += g.start[nodeIdx];
+  // Compute the max score
+  ArgMaxT maxAndIdx = {kNegInf, -1};
+  for (int gTid = threadIdx.x; gTid < numArcs; gTid += blockDim.x) {
+    float score;
+    if (gTid == numArcs - 1 && g.start[nodeIdx]) {
+      score = 0;
+    } else {
+      auto arcIdx = g.inArcs[g.inArcOffset[nodeIdx] + gTid];
+      int srcNodeIdx = g.srcNodes[arcIdx];
+      score = weights[arcIdx] + scores[srcNodeIdx];
+    }
+    if (score > maxAndIdx.val.x) {
+      maxAndIdx.val.x = score;
+      maxAndIdx.idx.y = gTid;
+    }
+  }
+  if (threadIdx.x < numArcs) {
+    atomicArgMax(maxes.data() + nodeIdx, maxAndIdx.val.x, maxAndIdx.idx.y);
+  }
+  __syncthreads();
+  float maxScore = maxes[nodeIdx].val.x;
+  if (tropical || maxScore == kInf || maxScore == kNegInf) {
+    scores[nodeIdx] = maxScore;
+  } else {
+    float expScore = 0.0;
+    for (int gTid = threadIdx.x; gTid < numArcs; gTid += blockDim.x) {
+      float score;
+      if (gTid == numArcs - 1 && g.start[nodeIdx]) {
+        score = 0.0;
+      } else {
+        auto arcIdx = g.inArcs[g.inArcOffset[nodeIdx] + gTid];
+        int srcNodeIdx = g.srcNodes[arcIdx];
+        score = weights[arcIdx] + scores[srcNodeIdx];
+      }
+      expScore += expf(score - maxScore);
+    }
+    if (threadIdx.x < numArcs) {
+      atomicAdd(scores.data() + nodeIdx, expScore);
+    }
+    __syncthreads();
+    if (threadIdx.x == 0) {
+      scores[nodeIdx] = logf(scores[nodeIdx]) + maxScore;
+    }
+  }
+}
+
+__global__
+void decrementDegreesKernel(
+    const GraphData g,
+    const HDSpan<int> currExplore,
+    HDSpan<int> nextExplore,
+    HDSpan<int> count,
+    HDSpan<int> degrees) {
+  int nodeIdx = currExplore[blockIdx.x];
+  int numArcs = g.outArcOffset[nodeIdx + 1] - g.outArcOffset[nodeIdx];
+  for (int gTid = threadIdx.x; gTid < numArcs; gTid += blockDim.x) {
+    auto arcIdx = g.outArcs[g.outArcOffset[nodeIdx] + gTid];
+    auto dstNodeIdx = g.dstNodes[arcIdx];
+    auto oldDegrees = atomicSub(degrees.data() + dstNodeIdx, 1);
+    if (oldDegrees == 1) {
+      int oldCount = atomicAdd(count.data(), 1);
+      nextExplore[oldCount] = dstNodeIdx;
+    }
+  }
+}
+
+__global__
+void setAcceptGradsKernel(
+    const HDSpan<int> acceptIds,
+    const HDSpan<float> score,
+    const HDSpan<ArgMaxT> maxAndIdx,
+    const HDSpan<float> nodeScores,
+    HDSpan<float> nodeGrads,
+    bool tropical) {
+  const int gTid = blockIdx.x * blockDim.x + threadIdx.x;
+  if (gTid >= acceptIds.size()) {
+    return;
+  }
+  int nodeIdx = acceptIds[gTid];
+  if (tropical) {
+    nodeGrads[nodeIdx] = gTid == maxAndIdx[0].idx.y ? 1.0 : 0.0;
+  } else {
+    float maxScore = maxAndIdx[0].val.x;
+    float denom = expf(score[0] - maxScore);
+    nodeGrads[nodeIdx] = expf(nodeScores[nodeIdx] - maxScore) / denom;
+  }
+}
+
+__global__
+void updateGradientsKernel(
+    GraphData g,
+    const float* weights,
+    const float* deltas,
+    const HDSpan<ArgMaxT> maxes,
+    const HDSpan<float> nodeScores,
+    HDSpan<float> nodeGrads,
+    HDSpan<float> arcGrads,
+    HDSpan<int> degrees,
+    const HDSpan<int> currExplore,
+    HDSpan<int> nextExplore,
+    HDSpan<int> count,
+    bool tropical) {
+  int nodeIdx = currExplore[blockIdx.x];
+  int numArcs = g.inArcOffset[nodeIdx + 1] - g.inArcOffset[nodeIdx]; 
+  for (int gTid = threadIdx.x; gTid < numArcs; gTid += blockDim.x) {
+    auto arcIdx = g.inArcs[g.inArcOffset[nodeIdx] + gTid];
+    auto srcNode = g.srcNodes[arcIdx];
+    float nodeGradUpdate;
+    if (tropical) {
+      nodeGradUpdate = (gTid == maxes[nodeIdx].idx.y) ? 1.0 : 0.0;
+    } else {
+      float maxScore = maxes[nodeIdx].val.x;
+      auto denom = expf(nodeScores[nodeIdx] - maxScore);
+      nodeGradUpdate = 
+        expf(nodeScores[srcNode] + weights[arcIdx] - maxScore) / denom;
+    }
+    nodeGradUpdate *= nodeGrads[nodeIdx];
+    atomicAdd(nodeGrads.data() + srcNode, nodeGradUpdate);
+    arcGrads[arcIdx] = nodeGradUpdate * deltas[0];
+    auto oldDegrees = atomicSub(degrees.data() + srcNode, 1);
+    if (oldDegrees == 1) {
+      int oldCount = atomicAdd(count.data(), 1);
+      nextExplore[oldCount] = srcNode;
+    }
+  }
+}
+
+void shortestDistanceGrad(
+    Graph& g,
+    const Graph& deltas,
+    const HDSpan<float>& nodeScores,
+    const HDSpan<ArgMaxT>& maxScores,
+    const HDSpan<float>& score,
+    const HDSpan<ArgMaxT>& maxAndIdx,
+    bool tropical) {
+  auto gData = g.getData();
+  HDSpan<int> degrees(g.numNodes(), Device::CUDA);
+  HDSpan<int> currExplore(g.numNodes(), Device::CUDA); 
+  HDSpan<int> nextExplore(g.numNodes(), Device::CUDA); 
+  HDSpan<int> numToExplore(1, 0, Device::CUDA);
+  HDSpan<float> nodeGrads(g.numNodes(), 0.0, Device::CUDA);
+  HDSpan<float> arcGrads(g.numArcs(), 0.0, Device::CUDA);
+
+  initExplore(
+      gData.accept,
+      gData.outArcOffset,
+      degrees,
+      currExplore,
+      numToExplore);
+
+  // Initialize the accept node gradients
+  if (g.numAccept() > 0) {
+    int blocks = divUp(g.numAccept(), kNT);
+    setAcceptGradsKernel<<<blocks, kNT>>>(
+        gData.acceptIds, score, maxAndIdx, nodeScores, nodeGrads, tropical);
+  }
+
+  while (currExplore.size() > 0) {
+    int arcsPerNode = divUp(g.numArcs(), g.numNodes());
+    int NT = std::min(1024, divUp(arcsPerNode, 32) * 32);
+    CUDA_CHECK(cudaMemsetAsync(numToExplore.begin(), 0, sizeof(int)));
+    updateGradientsKernel<<<currExplore.size(), NT>>>(
+        gData,
+        g.weights(),
+        deltas.weights(),
+        maxScores,
+        nodeScores,
+        nodeGrads,
+        arcGrads,
+        degrees,
+        currExplore,
+        nextExplore,
+        numToExplore,
+        tropical);
+    swap(currExplore, nextExplore);
+    int countHost;
+    CUDA_CHECK(cudaMemcpy(
+        (void*)&countHost,
+        (void*)numToExplore.begin(),
+        sizeof(int),
+        cudaMemcpyDeviceToHost));
+    currExplore.resize(countHost);
+  }
+
+  g.addGrad(arcGrads.data());
+  nodeGrads.clear();
+  arcGrads.clear();
+  currExplore.clear();
+  nextExplore.clear();
+  numToExplore.clear();
+}
+
+} // namespace
+
+Graph shortestDistance(const Graph& g, bool tropical) {
+  auto gData = g.getData();
+  HDSpan<float> scores(g.numNodes(), 0.0f, Device::CUDA);
+  HDSpan<ArgMaxT> maxes(g.numNodes(), Device::CUDA);
+  HDSpan<int> degrees(g.numNodes(), Device::CUDA);
+  HDSpan<int> currExplore(g.numNodes(), Device::CUDA); 
+  HDSpan<int> nextExplore(g.numNodes(), Device::CUDA); 
+  HDSpan<int> numToExplore(1, 0, Device::CUDA);
+
+  thrust::fill(
+      thrust::device, maxes.begin(), maxes.end(), ArgMaxT{kNegInf, -1});
+
+  initExplore(
+      gData.start, gData.inArcOffset, degrees, currExplore, numToExplore);
+
+  cudaStream_t str;
+  cudaStreamCreate(&str);
+  while (currExplore.size() > 0) {
+    int arcsPerNode = divUp(g.numArcs(), g.numNodes());
+    int NT = std::min(1024, divUp(arcsPerNode, 32) * 32);
+    reduceArcScoresKernel<<<currExplore.size(), NT, 0, str>>>(
+        gData, currExplore, scores, maxes, g.weights(), tropical);
+    CUDA_CHECK(cudaMemsetAsync(numToExplore.begin(), 0, sizeof(int)));
+    decrementDegreesKernel<<<currExplore.size(), NT>>>(
+        gData, currExplore, nextExplore, numToExplore, degrees);
+    swap(currExplore, nextExplore);
+    int countHost;
+    CUDA_CHECK(cudaMemcpyAsync(
+        (void*)&countHost,
+        (void*)numToExplore.begin(),
+        sizeof(int),
+        cudaMemcpyDeviceToHost));
+    cudaStreamSynchronize(0);
+    currExplore.resize(countHost);
+  }
+  cudaStreamSynchronize(str);
+  cudaStreamDestroy(str);
+  currExplore.clear();
+  nextExplore.clear();
+  numToExplore.clear();
+  degrees.clear();
+  
+  // Gather and log add the accept scores
+  HDSpan<float> score(1, Device::CUDA);;
+  HDSpan<ArgMaxT> maxAndIdx(1, Device::CUDA);
+  int NT = std::min(divUp(g.numAccept() + 1, 32) * 32, 1024);
+  reduceAccept<<<1, NT>>>(
+      scores, gData.acceptIds, score, maxAndIdx, tropical);
+
+  auto gradInfo = std::make_shared<GradInfo>(
+      scores, maxes, score, maxAndIdx);
+  auto gradFunc = [gradInfo, tropical](
+      std::vector<Graph>& inputs, Graph deltas) {
+    shortestDistanceGrad(
+        inputs[0],
+        deltas,
+        gradInfo->scores,
+        gradInfo->maxes,
+        gradInfo->score,
+        gradInfo->maxAndIdx,
+        tropical);
+  };
+
+  // Make the result graph
+  auto ngraph = scalarGraph(0, Device::CUDA, g.calcGrad());
+  ngraph.getWeights() = score;
+  ngraph.setInputs({g});
+  ngraph.setGradFunc(gradFunc);
+  return ngraph;
+}
+} // namespace detail
+} // namespace cuda 
+} // namespace gtn
diff --git a/test/cuda_functions_test.cpp b/test/cuda_functions_test.cpp
index 319df26..c970d1c 100644
--- a/test/cuda_functions_test.cpp
+++ b/test/cuda_functions_test.cpp
@@ -6,6 +6,10 @@
 
 using namespace gtn;
 
+namespace {
+  constexpr float inf = std::numeric_limits<float>::infinity();
+}
+
 // *NB* num_arcs is assumed to be greater than num_nodes.
 Graph makeRandomDAG(int num_nodes, int num_arcs) {
   Graph graph;
@@ -545,7 +549,7 @@ TEST_CASE("test edit distance", "[cuda functions]") {
       edits = edits.cuda();
     }
     auto outG = compose(xG, compose(edits, yG));
-    auto score = viterbiScore(outG.cpu());
+    auto score = viterbiScore(outG);
     return -score.item();
   };
 
@@ -601,7 +605,7 @@ TEST_CASE("test ngrams", "[cuda functions]") {
       ngramG = ngramG.cuda();
     }
     auto outG = compose(inputG, compose(ngramCounter, ngramG));
-    auto score = forwardScore(outG.cpu());
+    auto score = forwardScore(outG);
     return round(std::exp(score.item()));
   };
 
@@ -680,3 +684,403 @@ TEST_CASE("test cuda project and clone", "[cuda functions]") {
                                 "3 4 2 2 2\n"));
   CHECK(equal(projectOutput(graph).cpu(), outputExpected));
 }
+
+TEST_CASE("test cuda forward score", "[cuda functions]") {
+  {
+    // Check score of empty graph
+    Graph g;
+    CHECK(forwardScore(g.cuda()).item() == -inf);
+  }
+
+  {
+    // Handles negative infinity
+    Graph g;
+    g.addNode(true);
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, -inf);
+    g.addArc(0, 1, 1, 1, -inf);
+    CHECK(forwardScore(g.cuda()).item() == -inf);
+  }
+
+  {
+    // Handles infinity
+    Graph g;
+    g.addNode(true);
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, inf);
+    g.addArc(0, 1, 1, 1, 0);
+    CHECK(forwardScore(g.cuda()).item() == inf);
+  }
+
+  {
+    // Handles positive and negative infinity
+    Graph g;
+    g.addNode(true);
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, inf);
+    g.addArc(0, 1, 1, 1, -inf);
+    CHECK(forwardScore(g.cuda()).item() == inf);
+  }
+
+
+  {
+    // Single Node
+    Graph g;
+    g.addNode(true, true);
+    CHECK(forwardScore(g.cuda()).item() == 0.0);
+  }
+
+  {
+    // A simple test case
+    Graph g;
+    g.addNode(true);
+    g.addNode();
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, 1);
+    g.addArc(0, 1, 1, 1, 2);
+    g.addArc(0, 1, 2, 2, 3);
+    g.addArc(1, 2, 0, 0, 1);
+    g.addArc(1, 2, 1, 1, 2);
+    g.addArc(1, 2, 2, 2, 3);
+    CHECK(forwardScore(g.cuda()).item() == Approx(6.8152));
+  }
+
+  {
+    // Handle two start nodes
+    Graph g;
+    g.addNode(true);
+    g.addNode(true);
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, -5);
+    g.addArc(0, 2, 0, 0, 1);
+    g.addArc(1, 2, 0, 0, 2);
+    float expected = std::log(std::exp(1) + std::exp(-5 + 2) + std::exp(2));
+    CHECK(forwardScore(g.cuda()).item() == Approx(expected));
+  }
+
+  {
+    // Handle two accept nodes
+    Graph g;
+    g.addNode(true);
+    g.addNode(false, true);
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, 2);
+    g.addArc(0, 2, 0, 0, 2);
+    g.addArc(1, 2, 0, 0, 2);
+    float expected = std::log(2 * std::exp(2) + std::exp(4));
+    CHECK(forwardScore(g.cuda()).item() == Approx(expected));
+  }
+
+  {
+    // Handle case where some arcs don't lead to accepting states
+    Graph g;
+    g.addNode(true);
+    g.addNode(false, false);
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, 2);
+    g.addArc(0, 2, 0, 0, 2);
+    CHECK(forwardScore(g.cuda()).item() == 2.0);
+  }
+
+  {
+    // A more complex test case
+    std::stringstream in(
+        "0 1\n"
+        "3 4\n"
+        "0 1 0 0 2\n"
+        "0 2 1 1 1\n"
+        "1 2 0 0 2\n"
+        "2 3 0 0 1\n"
+        "2 3 1 1 1\n"
+        "1 4 0 0 2\n"
+        "2 4 1 1 3\n"
+        "3 4 0 0 2\n");
+    Graph g = loadTxt(in);
+    CHECK(forwardScore(g.cuda()).item() == Approx(8.36931));
+  }
+}
+
+TEST_CASE("test cuda forward score grad", "[cuda functions]") {
+  auto checkGrad = [](const Graph& g) {
+    auto gDev = g.cuda();
+    backward(forwardScore(g));
+    backward(forwardScore(gDev));
+    auto gradW = g.grad().weights();
+    auto gradDev = gDev.grad().cpu();
+    auto gradDevW = gradDev.weights();
+    float diff = 0, tot = 0;
+    for (int i = 0; i < g.numArcs(); ++i) {
+      diff += std::pow(gradW[i] - gradDevW[i], 2);
+      tot += std::pow(gradW[i], 2);
+    }
+    return (diff / tot) < 1e-3;
+  };
+
+  {
+    Graph g;
+    g.addNode(true);
+    g.addNode();
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, 1);
+    g.addArc(0, 1, 1, 1, 2);
+    g.addArc(0, 1, 2, 2, 3);
+    g.addArc(1, 2, 0, 0, 1);
+    g.addArc(1, 2, 1, 1, 2);
+    g.addArc(1, 2, 2, 2, 3);
+    CHECK(checkGrad(g));
+  }
+
+  {
+    // Handle two start nodes
+    Graph g;
+    g.addNode(true);
+    g.addNode(true);
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, -5);
+    g.addArc(0, 2, 0, 0, 1);
+    g.addArc(1, 2, 0, 0, 2);
+    g = g.cuda();
+    backward(forwardScore(g));
+    auto grad = g.grad().cpu();
+    double denom = 1 / (std::exp(-3) + std::exp(1) + std::exp(2));
+    CHECK(grad.weight(0) == Approx(denom * std::exp(-3)));
+    CHECK(grad.weight(1) == Approx(denom * std::exp(1)));
+    CHECK(grad.weight(2) == Approx(denom * (std::exp(-3) + std::exp(2))));
+  }
+
+  {
+    // Handle two accept nodes
+    Graph g;
+    g.addNode(true);
+    g.addNode(false, true);
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, 2);
+    g.addArc(0, 2, 0, 0, 2);
+    g.addArc(1, 2, 0, 0, 2);
+    g = g.cuda();
+    backward(forwardScore(g));
+    auto grad = g.grad().cpu();
+    double denom = 1 / (2 * std::exp(2) + std::exp(4));
+    CHECK(grad.weight(0) == Approx(denom * (std::exp(2) + std::exp(4))));
+    CHECK(grad.weight(1) == Approx(denom * std::exp(2)));
+    CHECK(grad.weight(2) == Approx(denom * std::exp(4)));
+  }
+
+  {
+    // Handle case where some arcs don't lead to accepting states
+    Graph g;
+    g.addNode(true);
+    g.addNode(false, false);
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, 2);
+    g.addArc(0, 2, 0, 0, 2);
+    g = g.cuda();
+    backward(forwardScore(g));
+    auto grad = g.grad().cpu();
+    CHECK(grad.weight(0) == Approx(0.0));
+    CHECK(grad.weight(1) == Approx(1.0));
+  }
+
+  {
+    // Handles negative infinity
+    Graph g;
+    g.addNode(true);
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, -inf);
+    g.addArc(0, 1, 1, 1, -inf);
+    g = g.cuda();
+    backward(forwardScore(g));
+    auto grad = g.grad().cpu();
+    CHECK(std::isnan(grad.weight(0)));
+    CHECK(std::isnan(grad.weight(1)));
+
+    Graph g2;
+    g2.addNode(true);
+    g2.addNode(false, true);
+    g2.addArc(0, 1, 0, 0, -inf);
+    g2.addArc(0, 1, 1, 1, 1.0);
+    g2 = g2.cuda();
+    backward(forwardScore(g2));
+    auto grad2 = g2.grad().cpu();
+    CHECK(grad2.weight(0) == Approx(0.0));
+    CHECK(grad2.weight(1) == Approx(1.0));
+  }
+
+  {
+    // Handles infinity
+    Graph g;
+    g.addNode(true);
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, inf);
+    g.addArc(0, 1, 1, 1, inf);
+    g = g.cuda();
+    backward(forwardScore(g));
+    auto grad = g.grad().cpu();
+    CHECK(std::isnan(grad.weight(0)));
+    CHECK(std::isnan(grad.weight(1)));
+
+    Graph g2;
+    g2.addNode(true);
+    g2.addNode(false, true);
+    g2.addArc(0, 1, 0, 0, inf);
+    g2.addArc(0, 1, 1, 1, 1.0);
+    g2 = g2.cuda();
+    backward(forwardScore(g2));
+    auto grad2 = g2.grad().cpu();
+    CHECK(std::isnan(grad2.weight(0)));
+    CHECK(std::isnan(grad2.weight(1)));
+  }
+
+  {
+    // A more complex test case
+    std::stringstream in(
+        "0 1\n"
+        "3 4\n"
+        "0 1 0 0 2\n"
+        "0 2 1 1 1\n"
+        "1 2 0 0 2\n"
+        "2 3 0 0 1\n"
+        "2 3 1 1 1\n"
+        "1 4 0 0 2\n"
+        "2 4 1 1 3\n"
+        "3 4 0 0 2\n");
+    Graph g = loadTxt(in);
+    CHECK(checkGrad(g));
+  }
+}
+
+TEST_CASE("test cuda viterbi score", "[cuda functions]") {
+  {
+    // Check score of empty graph
+    Graph g;
+    CHECK(viterbiScore(g.cuda()).item() == -inf);
+  }
+
+  {
+    // A simple test case
+    Graph g;
+    g.addNode(true);
+    g.addNode();
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, 1);
+    g.addArc(0, 1, 1, 1, 2);
+    g.addArc(0, 1, 2, 2, 3);
+    g.addArc(1, 2, 0, 0, 1);
+    g.addArc(1, 2, 1, 1, 2);
+    g.addArc(1, 2, 2, 2, 3);
+    CHECK(viterbiScore(g.cuda()).item() == 6.0f);
+  }
+
+  {
+    // Handle two start nodes
+    Graph g;
+    g.addNode(true);
+    g.addNode(true);
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, -5);
+    g.addArc(0, 2, 0, 0, 1);
+    g.addArc(1, 2, 0, 0, 2);
+    CHECK(viterbiScore(g.cuda()).item() == 2.0f);
+  }
+
+  {
+    // Handle two accept nodes
+    Graph g;
+    g.addNode(true);
+    g.addNode(false, true);
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, 2);
+    g.addArc(0, 2, 0, 0, 2);
+    g.addArc(1, 2, 0, 0, 2);
+    CHECK(viterbiScore(g.cuda()).item() == 4.0f);
+  }
+
+  {
+    // A more complex test case
+    std::stringstream in(
+        "0 1\n"
+        "3 4\n"
+        "0 1 0 0 2\n"
+        "0 2 1 1 1\n"
+        "1 2 0 0 2\n"
+        "2 3 0 0 1\n"
+        "2 3 1 1 1\n"
+        "1 4 0 0 2\n"
+        "2 4 1 1 3\n"
+        "3 4 0 0 2\n");
+    Graph g = loadTxt(in);
+    CHECK(viterbiScore(g.cuda()).item() == 7.0f);
+  }
+}
+
+TEST_CASE("test cuda viterbi score grad", "[cuda functions]") {
+  auto gradsToVec = [](Graph g) {
+    g = g.cuda();
+    backward(viterbiScore(g));
+    auto grad = g.grad().cpu();
+    std::vector<float> v(grad.numArcs());
+    std::copy(grad.weights(), grad.weights() + grad.numArcs(), v.begin());
+    return v;
+  };
+
+  {
+    Graph g;
+    g.addNode(true);
+    g.addNode();
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, 1);
+    g.addArc(0, 1, 1, 1, 2);
+    g.addArc(0, 1, 2, 2, 3);
+    g.addArc(1, 2, 0, 0, 1);
+    g.addArc(1, 2, 1, 1, 2);
+    g.addArc(1, 2, 2, 2, 3);
+    std::vector<float> expected = {0.0, 0.0, 1.0, 0.0, 0.0, 1.0};
+    CHECK(gradsToVec(g) == expected);
+  }
+
+  {
+    // Handle two start nodes
+    Graph g;
+    g.addNode(true);
+    g.addNode(true);
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, -5);
+    g.addArc(0, 2, 0, 0, 1);
+    g.addArc(1, 2, 0, 0, 2);
+    std::vector<float> expected = {0.0, 0.0, 1.0};
+    CHECK(gradsToVec(g) == expected);
+  }
+
+  {
+    // Handle two accept nodes
+    Graph g;
+    g.addNode(true);
+    g.addNode(false, true);
+    g.addNode(false, true);
+    g.addArc(0, 1, 0, 0, 2);
+    g.addArc(0, 2, 0, 0, 2);
+    g.addArc(1, 2, 0, 0, 2);
+    std::vector<float> expected = {1.0, 0.0, 1.0};
+    CHECK(gradsToVec(g) == expected);
+  }
+
+  {
+    // A more complex test case
+    std::stringstream in(
+        "0 1\n"
+        "3 4\n"
+        "0 1 0 0 2\n"
+        "0 2 1 1 1\n"
+        "1 2 0 0 2\n"
+        "2 3 0 0 1\n"
+        "2 3 1 1 1\n"
+        "1 4 0 0 2\n"
+        "2 4 1 1 3\n"
+        "3 4 0 0 2\n");
+    Graph g = loadTxt(in);
+    // two possible paths with same viterbi score
+    std::vector<float> expected1 = {1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0};
+    std::vector<float> expected2 = {1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 1.0};
+    CHECK(((gradsToVec(g) == expected1) || (gradsToVec(g) == expected2)));
+  }
+}
diff --git a/test/cuda_test.cpp b/test/cuda_test.cpp
index d1712a3..e7316b8 100644
--- a/test/cuda_test.cpp
+++ b/test/cuda_test.cpp
@@ -1,4 +1,5 @@
 #include "catch.hpp"
+#include <thread>
 
 #include <gtn/gtn.h>
 
@@ -15,6 +16,23 @@ TEST_CASE("test cuda utils", "[cuda]") {
     cuda::setDevice(numDevices - 1);
     device = cuda::getDevice();
     CHECK(device == numDevices - 1);
+
+    // Launch work in main thread stream
+    detail::HDSpan<float> a(1 << 10, 1.0, Device::CUDA);
+    detail::HDSpan<float> b(1 << 10, 0.0, Device::CUDA);
+    detail::HDSpan<float> c(1 << 10, 1000.0, Device::CUDA);
+    auto event = cuda::Event();
+    for (int i = 0; i < 1000; ++i) {
+      add(a, b, b);
+    }
+    event.record();
+
+    // Equality check in a different stream
+    auto eq = std::async([&b, &c, &event](){
+        event.wait();
+        return b == c;
+      });
+    CHECK(eq.get());
 }
 
 TEST_CASE("test graph cuda", "[cuda]") {
diff --git a/test/hd_span_test.cpp b/test/hd_span_test.cpp
index b2f1487..ebcebe4 100644
--- a/test/hd_span_test.cpp
+++ b/test/hd_span_test.cpp
@@ -140,6 +140,14 @@ TEST_CASE("test hd_span", "[hd_span]") {
     CHECK(h1 != h2);
   }
 
+  // Test swap
+  {
+    HDSpan<int> h1(2, 1);
+    HDSpan<int> h2(3, 2);
+    swap(h1, h2);
+    CHECK(h1 == HDSpan<int>(3, 2));
+    CHECK(h2 == HDSpan<int>(2, 1));
+  }
 }
 
 TEST_CASE("test hd_span cuda", "[hd_span]") {
diff --git a/test/no_cuda_test.cpp b/test/no_cuda_test.cpp
index 22e6916..4bc1275 100644
--- a/test/no_cuda_test.cpp
+++ b/test/no_cuda_test.cpp
@@ -12,6 +12,9 @@ TEST_CASE("test cuda utils", "[cuda]") {
     CHECK_THROWS(cuda::getDevice());
     CHECK_THROWS(cuda::setDevice(0));
     CHECK_THROWS(cuda::synchronize());
+    CHECK_THROWS(cuda::synchronizeStream());
+    CHECK_THROWS(cuda::Event());
+
     CHECK_THROWS(cuda::detail::allocate(1, 0));
     CHECK_THROWS(cuda::detail::free((void*)0));
     std::vector<float> a = {1, 0, 1};
diff --git a/test/parallel_test.cpp b/test/parallel_test.cpp
index 26abfdb..d29d2f6 100644
--- a/test/parallel_test.cpp
+++ b/test/parallel_test.cpp
@@ -18,6 +18,68 @@
 
 using namespace gtn;
 
+TEST_CASE("test thread pool", "[parallel]") {
+  gtn::detail::ThreadPool pool(2);
+
+  {
+    std::vector<std::future<int>> results(10);
+    for (int i = 0; i < results.size(); ++i) {
+      results[i] = pool.enqueue([](int idx){ return idx; }, i);
+    }
+    for (int i = 0; i < results.size(); ++i) {
+      CHECK(results[i].get() == i);
+    }
+  }
+
+  {
+    std::vector<std::future<int>> results(10);
+    for (int i = 0; i < results.size(); ++i) {
+      results[i] = pool.enqueueIndex(i, [](int idx){ return idx; }, i);
+    }
+    for (int i = 0; i < results.size(); ++i) {
+      CHECK(results[i].get() == i);
+    }
+  }
+
+  // Check main stream synchronizes with thread pool
+  if (cuda::isAvailable()) {
+    detail::HDSpan<float> a(1 << 20, 1.0, Device::CUDA);
+    detail::HDSpan<float> b(1 << 20, 0.0, Device::CUDA);
+    detail::HDSpan<float> c(1 << 20, 1000.0, Device::CUDA);
+    for (int i = 0; i < 1000; i++) {
+      detail::add(a, b, b);
+    }
+    std::future<bool> fut = pool.enqueue([&b, &c]() { return b == c; });
+    CHECK(fut.get());
+  }
+
+  // Check thread pool stream synchronization works
+  if (cuda::isAvailable()) {
+    detail::HDSpan<float> a(1 << 20, 1.0, Device::CUDA);
+    detail::HDSpan<float> b(1 << 20, 0.0, Device::CUDA);
+    detail::HDSpan<float> c(1 << 20, 100.0, Device::CUDA);
+    for (int i = 0; i < 100; ++i) {
+      pool.enqueueIndex(0, [&a, &b]() { detail::add(a, b, b); });
+    }
+    pool.syncStreams();
+    std::future<bool> fut = pool.enqueueIndex(1, [&b, &c]() { return b == c; });
+    CHECK(fut.get());
+  }
+
+  // Check that the main stream synchronizes with the worker streams
+  if (cuda::isAvailable()) {
+    detail::HDSpan<float> a(1 << 20, 0.0, Device::CUDA);
+    detail::HDSpan<float> b(1 << 20, 100.0, Device::CUDA);
+    for (int i = 0; i < 1000; ++i) {
+      pool.enqueueIndex(0, [&]() { a = b;});
+    }
+    pool.syncStreams();
+    detail::add(b, b, b);
+    detail::HDSpan<float> expected(1 << 20, 100, Device::CUDA);
+    CHECK((a == expected));
+  }
+}
+
 TEST_CASE("test parallel map one arg", "[parallel]") {
   const int B = 4;
 
@@ -25,7 +87,6 @@ TEST_CASE("test parallel map one arg", "[parallel]") {
   for (size_t i = 0; i < B; ++i) {
     inputs.push_back(scalarGraph(static_cast<float>(i)));
   }
-
   auto outputs = parallelMap(negate, inputs);
 
   std::vector<Graph> expectedOutputs;
@@ -237,3 +298,66 @@ TEST_CASE("test parallel map throws", "[parallel]") {
       std::logic_error,
       Catch::Message("[gtn::negate] input must have only one arc"));
 }
+
+TEST_CASE("test parallel map cuda", "[parallel]") {
+  if (!cuda::isAvailable()) {
+    return;
+  }
+  {
+    const int B = 4;
+    std::vector<Graph> inputs;
+    std::vector<Graph> expectedOutputs;
+    for (size_t i = 0; i < B; ++i) {
+      inputs.push_back(scalarGraph(static_cast<float>(i), Device::CUDA));
+      expectedOutputs.push_back(negate(inputs[i]));
+    }
+
+    auto outputs = parallelMap(negate, inputs);
+
+    for (size_t i = 0; i < B; ++i) {
+      CHECK(equal(outputs[i], expectedOutputs[i]));
+    }
+  }
+
+  {
+    std::vector<Graph> inputs;
+    inputs.push_back(linearGraph(100, 10000, Device::CUDA));
+    inputs.push_back(linearGraph(1, 1, Device::CUDA));
+    for (int i = 0; i < 100; ++i) {
+      parallelMap(projectInput, inputs);
+    }
+    std::vector<Graph> inputs1;
+    std::vector<Graph> inputs2;
+    std::vector<Graph> expectedOutputs;
+    inputs1.push_back(scalarGraph(2.0, Device::CUDA));
+    inputs1.push_back(scalarGraph(1.0, Device::CUDA));
+    inputs2.push_back(scalarGraph(0.0, Device::CUDA));
+    inputs2.push_back(scalarGraph(0.0, Device::CUDA));
+    expectedOutputs.push_back(scalarGraph(1.0, Device::CUDA));
+    expectedOutputs.push_back(scalarGraph(2.0, Device::CUDA));
+    auto outputs = parallelMap(add, inputs1, inputs2);
+    std::swap(outputs[0], outputs[1]);
+    auto results = parallelMap(equal, outputs, expectedOutputs);
+    CHECK(results == std::vector<bool>(2, true));
+  }
+
+  // Changing the size of the batch shouldn't cause synchronization issues
+  {
+    std::vector<Graph> inputs1;
+    std::vector<Graph> inputs2;
+    inputs1.push_back(scalarGraph(1.0, Device::CUDA));
+    inputs1.push_back(scalarGraph(1.0, Device::CUDA));
+    inputs2.push_back(scalarGraph(0.0, Device::CUDA));
+    inputs2.push_back(scalarGraph(0.0, Device::CUDA));
+    for (int i = 0; i < 100; i++) {
+      inputs2 = parallelMap(add, inputs1, inputs2);
+    }
+    for (int i = 0; i < 10; i++) {
+      inputs1.push_back(scalarGraph(0.0, Device::CUDA));
+      inputs2.push_back(scalarGraph(0.0, Device::CUDA));
+    }
+    auto results = parallelMap(add, inputs1, inputs2);
+    CHECK(equal(results[0], scalarGraph(101.0, Device::CUDA)));
+    CHECK(equal(results[1], scalarGraph(101.0, Device::CUDA)));
+  }
+}