Merge pull request #62 from bdice/streams

Add stream support.

Author: karthikeyann

velox/experimental/cudf/connectors/parquet/ParquetDataSink.cpp

@@ -157,7 +157,9 @@ void ParquetDataSink::appendData(RowVectorPtr input) {
   checkRunning();
 
   // Convert the input RowVectorPtr to cudf::table
-  auto cudfInput = with_arrow::to_cudf_table(input, input->pool());
+  auto stream = cudfGlobalStreamPool().get_stream();
+  auto cudfInput = with_arrow::to_cudf_table(input, input->pool(), stream);
+  stream.synchronize();
   VELOX_CHECK_NOT_NULL(
       cudfInput, "Failed to convert input RowVectorPtr to cudf::table");
 

velox/experimental/cudf/connectors/parquet/ParquetDataSource.cpp

@@ -115,12 +115,13 @@ std::optional<RowVectorPtr> ParquetDataSource::next(
     }
 
     if (readTables.size()) {
-      auto readTable = concatenateTables(std::move(readTables));
+      auto stream = cudf::get_default_stream();
+      auto readTable = concatenateTables(std::move(readTables), stream);
       if (cudfTable_) {
         // Concatenate the current view ahead of the read table.
         auto tableViews = std::vector<cudf::table_view>{
             currentCudfTableView_, readTable->view()};
-        cudfTable_ = cudf::concatenate(tableViews, cudf::get_default_stream());
+        cudfTable_ = cudf::concatenate(tableViews, stream);
       } else {
         cudfTable_ = std::move(readTable);
       }
@@ -135,8 +136,10 @@ std::optional<RowVectorPtr> ParquetDataSource::next(
   // If the current table view has <= size rows, this is the last chunk.
   if (currentCudfTableView_.num_rows() <= size) {
     // Convert the current table view to RowVectorPtr.
-    output =
-        with_arrow::to_velox_column(currentCudfTableView_, pool_, columnNames);
+    auto stream = cudf::get_default_stream();
+    output = with_arrow::to_velox_column(
+        currentCudfTableView_, pool_, columnNames, stream);
+    stream.synchronize();
     // Reset internal tables
     resetCudfTableAndView();
   } else {
@@ -149,7 +152,10 @@ std::optional<RowVectorPtr> ParquetDataSource::next(
         tableSplits[0].num_rows(),
         "cudf::split yielded incorrect partitions");
     // Convert the first split view to RowVectorPtr.
-    output = with_arrow::to_velox_column(tableSplits[0], pool_, columnNames);
+    auto stream = cudf::get_default_stream();
+    output =
+        with_arrow::to_velox_column(tableSplits[0], pool_, columnNames, stream);
+    stream.synchronize();
     // Set the current view to the second split view.
     currentCudfTableView_ = tableSplits[1];
   }

velox/experimental/cudf/exec/CudfConversion.cpp

@@ -111,9 +111,14 @@ RowVectorPtr CudfFromVelox::getOutput() {
     return nullptr;
   }
 
+  // Get a stream from the global stream pool
+  auto stream = cudfGlobalStreamPool().get_stream();
+
   // Convert RowVector to cudf table
-  auto tbl = with_arrow::to_cudf_table(input, input->pool());
-  cudf::get_default_stream().synchronize();
+  auto tbl = with_arrow::to_cudf_table(input, input->pool(), stream);
+
+  stream.synchronize();
+
   VELOX_CHECK_NOT_NULL(tbl);
 
   if (cudfDebugEnabled()) {
@@ -126,7 +131,7 @@ RowVectorPtr CudfFromVelox::getOutput() {
   // Return a CudfVector that owns the cudf table
   auto const size = tbl->num_rows();
   return std::make_shared<CudfVector>(
-      input->pool(), outputType_, size, std::move(tbl));
+      input->pool(), outputType_, size, std::move(tbl), stream);
 }
 
 void CudfFromVelox::close() {
@@ -162,6 +167,7 @@ RowVectorPtr CudfToVelox::getOutput() {
     return nullptr;
   }
 
+  auto stream = inputs_.front()->stream();
   std::unique_ptr<cudf::table> tbl = inputs_.front()->release();
   inputs_.pop_front();
 
@@ -172,12 +178,12 @@ RowVectorPtr CudfToVelox::getOutput() {
     std::cout << "CudfToVelox table number of rows: " << tbl->num_rows()
               << std::endl;
   }
-
-  cudf::get_default_stream().synchronize();
   if (tbl->num_rows() == 0) {
     return nullptr;
   }
-  RowVectorPtr output = with_arrow::to_velox_column(tbl->view(), pool(), "");
+  RowVectorPtr output =
+      with_arrow::to_velox_column(tbl->view(), pool(), "", stream);
+  stream.synchronize();
   finished_ = noMoreInput_ && inputs_.empty();
   output->setType(outputType_);
   return output;

velox/experimental/cudf/exec/CudfFilterProject.cpp

@@ -258,6 +258,7 @@ RowVectorPtr CudfFilterProject::getOutput() {
 
   auto cudf_input = std::dynamic_pointer_cast<CudfVector>(input_);
   VELOX_CHECK_NOT_NULL(cudf_input);
+  auto stream = cudf_input->stream();
   auto input_table_columns = cudf_input->release()->release();
   // add ast unsupported precomputed columns to input_table
   // Works only directly on column in input table, not intermediate columns
@@ -267,13 +268,13 @@ RowVectorPtr CudfFilterProject::getOutput() {
       auto new_column = cudf::datetime::extract_datetime_component(
           input_table_columns[dependent_column_index]->view(),
           cudf::datetime::datetime_component::YEAR,
-          cudf::get_default_stream(),
+          stream,
           cudf::get_current_device_resource_ref());
       input_table_columns.emplace_back(std::move(new_column));
     } else if (ins_name == "length") {
       auto new_column = cudf::strings::count_characters(
           input_table_columns[dependent_column_index]->view(),
-          cudf::get_default_stream(),
+          stream,
           cudf::get_current_device_resource_ref());
       input_table_columns.emplace_back(std::move(new_column));
     } else {
@@ -289,7 +290,7 @@ RowVectorPtr CudfFilterProject::getOutput() {
     auto col = cudf::compute_column(
         cudf_table_view,
         tree.back(),
-        cudf::get_default_stream(),
+        stream,
         cudf::get_current_device_resource_ref());
     columns.emplace_back(std::move(col));
   }
@@ -308,6 +309,7 @@ RowVectorPtr CudfFilterProject::getOutput() {
   }
 
   auto output_table = std::make_unique<cudf::table>(std::move(output_columns));
+  stream.synchronize();
   auto const num_columns = output_table->num_columns();
   auto const size = output_table->num_rows();
   if (cudfDebugEnabled()) {
@@ -316,7 +318,7 @@ RowVectorPtr CudfFilterProject::getOutput() {
   }
 
   auto cudf_output = std::make_shared<CudfVector>(
-      input_->pool(), outputType_, size, std::move(output_table));
+      input_->pool(), outputType_, size, std::move(output_table), stream);
   input_.reset();
   if (num_columns == 0 or size == 0) {
     return nullptr;

velox/experimental/cudf/exec/CudfHashJoin.cpp

@@ -151,14 +151,22 @@ void CudfHashJoinBuild::noMoreInput() {
   };
 
   auto cudf_tables = std::vector<std::unique_ptr<cudf::table>>(inputs_.size());
+  auto input_streams = std::vector<rmm::cuda_stream_view>(inputs_.size());
   for (int i = 0; i < inputs_.size(); i++) {
     VELOX_CHECK_NOT_NULL(inputs_[i]);
+    input_streams[i] = inputs_[i]->stream();
     cudf_tables[i] = inputs_[i]->release();
   }
-  auto tbl = concatenateTables(std::move(cudf_tables));
+  auto stream = cudfGlobalStreamPool().get_stream();
+  cudf::detail::join_streams(input_streams, stream);
+  auto tbl = concatenateTables(std::move(cudf_tables), stream);
+
+  // Release input data after synchronizing
+  stream.synchronize();
+  input_streams.clear();
+  cudf_tables.clear();
 
   // Release input data
-  cudf::get_default_stream().synchronize();
   inputs_.clear();
 
   VELOX_CHECK_NOT_NULL(tbl);
@@ -246,6 +254,7 @@ RowVectorPtr CudfHashJoinProbe::getOutput() {
   }
   auto cudf_input = std::dynamic_pointer_cast<CudfVector>(input_);
   VELOX_CHECK_NOT_NULL(cudf_input);
+  auto stream = cudf_input->stream();
   auto tbl = cudf_input->release();
   if (cudfDebugEnabled()) {
     std::cout << "Probe table number of columns: " << tbl->num_columns()
@@ -307,8 +316,8 @@ RowVectorPtr CudfHashJoinProbe::getOutput() {
           hb.get(),
           hashObject_.has_value());
   }
-  auto const [left_join_indices, right_join_indices] =
-      hb->inner_join(tbl->view().select(probe_key_indices));
+  auto const [left_join_indices, right_join_indices] = hb->inner_join(
+      tbl->view().select(probe_key_indices), std::nullopt, stream);
   auto left_indices_span =
       cudf::device_span<cudf::size_type const>{*left_join_indices};
   auto right_indices_span =
@@ -361,8 +370,10 @@ RowVectorPtr CudfHashJoinProbe::getOutput() {
   auto left_indices_col = cudf::column_view{left_indices_span};
   auto right_indices_col = cudf::column_view{right_indices_span};
   auto constexpr oob_policy = cudf::out_of_bounds_policy::DONT_CHECK;
-  auto left_result = cudf::gather(left_input, left_indices_col, oob_policy);
-  auto right_result = cudf::gather(right_input, right_indices_col, oob_policy);
+  auto left_result =
+      cudf::gather(left_input, left_indices_col, oob_policy, stream);
+  auto right_result =
+      cudf::gather(right_input, right_indices_col, oob_policy, stream);
 
   if (cudfDebugEnabled()) {
     std::cout << "Left result number of columns: " << left_result->num_columns()
@@ -391,7 +402,7 @@ RowVectorPtr CudfHashJoinProbe::getOutput() {
     return nullptr;
   }
   return std::make_shared<CudfVector>(
-      pool(), outputType, size, std::move(cudf_output));
+      pool(), outputType, size, std::move(cudf_output), stream);
 }
 
 exec::BlockingReason CudfHashJoinProbe::isBlocked(ContinueFuture* future) {

velox/experimental/cudf/exec/CudfOrderBy.cpp

@@ -86,11 +86,20 @@ void CudfOrderBy::noMoreInput() {
     return;
   }
   auto cudf_tables = std::vector<std::unique_ptr<cudf::table>>(inputs_.size());
+  auto input_streams = std::vector<rmm::cuda_stream_view>(inputs_.size());
   for (int i = 0; i < inputs_.size(); i++) {
     VELOX_CHECK_NOT_NULL(inputs_[i]);
+    input_streams[i] = inputs_[i]->stream();
     cudf_tables[i] = inputs_[i]->release();
   }
-  auto tbl = concatenateTables(std::move(cudf_tables));
+  auto stream = cudfGlobalStreamPool().get_stream();
+  cudf::detail::join_streams(input_streams, stream);
+  auto tbl = concatenateTables(std::move(cudf_tables), stream);
+
+  // Release input data after synchronizing
+  stream.synchronize();
+  input_streams.clear();
+  cudf_tables.clear();
 
   // Release input data
   inputs_.clear();
@@ -105,10 +114,11 @@ void CudfOrderBy::noMoreInput() {
 
   auto keys = tbl->view().select(sort_keys_);
   auto values = tbl->view();
-  auto result = cudf::sort_by_key(values, keys, column_order_, null_order_);
+  auto result =
+      cudf::sort_by_key(values, keys, column_order_, null_order_, stream);
   auto const size = result->num_rows();
   outputTable_ = std::make_shared<CudfVector>(
-      pool(), outputType_, size, std::move(result));
+      pool(), outputType_, size, std::move(result), stream);
 }
 
 RowVectorPtr CudfOrderBy::getOutput() {

velox/experimental/cudf/exec/Utilities.cpp

@@ -18,7 +18,10 @@
 #include <memory>
 #include <string_view>
 
+#include "velox/experimental/cudf/exec/Utilities.h"
+
 #include <common/base/Exceptions.h>
+
 #include <rmm/mr/device/arena_memory_resource.hpp>
 #include <rmm/mr/device/cuda_async_memory_resource.hpp>
 #include <rmm/mr/device/cuda_memory_resource.hpp>
@@ -28,12 +31,11 @@
 #include <rmm/mr/device/pool_memory_resource.hpp>
 
 #include <cudf/concatenate.hpp>
+#include <cudf/detail/utilities/stream_pool.hpp>
 #include <cudf/utilities/default_stream.hpp>
 #include <cudf/utilities/error.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 
-#include "velox/experimental/cudf/exec/Utilities.h"
-
 namespace facebook::velox::cudf_velox {
 
 namespace {
@@ -84,13 +86,18 @@ std::shared_ptr<rmm::mr::device_memory_resource> create_memory_resource(
       "\nExpecting: cuda, pool, async, arena, managed, or managed_pool");
 }
 
+cudf::detail::cuda_stream_pool& cudfGlobalStreamPool() {
+  return cudf::detail::global_cuda_stream_pool();
+};
+
 bool cudfDebugEnabled() {
   const char* env_cudf_debug = std::getenv("VELOX_CUDF_DEBUG");
   return env_cudf_debug != nullptr && std::stoi(env_cudf_debug);
 }
 
 std::unique_ptr<cudf::table> concatenateTables(
-    std::vector<std::unique_ptr<cudf::table>> tables) {
+    std::vector<std::unique_ptr<cudf::table>> tables,
+    rmm::cuda_stream_view stream) {
   // Check for empty vector
   VELOX_CHECK_GT(tables.size(), 0);
 
@@ -105,9 +112,7 @@ std::unique_ptr<cudf::table> concatenateTables(
       std::back_inserter(tableViews),
       [&](auto const& tbl) { return tbl->view(); });
   return cudf::concatenate(
-      tableViews,
-      cudf::get_default_stream(),
-      cudf::get_current_device_resource_ref());
+      tableViews, stream, cudf::get_current_device_resource_ref());
 }
 
 } // namespace facebook::velox::cudf_velox

velox/experimental/cudf/exec/Utilities.h

@@ -19,18 +19,32 @@
 #include <memory>
 #include <string_view>
 
+#include <cudf/detail/utilities/stream_pool.hpp>
 #include <cudf/table/table.hpp>
 #include <rmm/mr/device/device_memory_resource.hpp>
 
 namespace facebook::velox::cudf_velox {
 
+/**
+ * @brief Creates a memory resource based on the given mode.
+ */
 [[nodiscard]] std::shared_ptr<rmm::mr::device_memory_resource>
 create_memory_resource(std::string_view mode);
 
+/**
+ * @brief Returns the global CUDA stream pool used by cudf.
+ */
+[[nodiscard]] cudf::detail::cuda_stream_pool& cudfGlobalStreamPool();
+
+/**
+ * @brief Returns true if the VELOX_CUDF_DEBUG environment variable is set to a
+ * nonzero value.
+ */
 bool cudfDebugEnabled();
 
 // Concatenate a vector of cuDF tables into a single table
 std::unique_ptr<cudf::table> concatenateTables(
-    std::vector<std::unique_ptr<cudf::table>> tables);
+    std::vector<std::unique_ptr<cudf::table>> tables,
+    rmm::cuda_stream_view stream);
 
 } // namespace facebook::velox::cudf_velox