Refactor cuda code into .cu files (#97)

Author: RAMitchell

cpp/include/legate_dataframe/binaryop.hpp

@@ -17,13 +17,43 @@
 #pragma once
 
 #include <legate.h>
-
-#include <cudf/binaryop.hpp>
-
 #include <legate_dataframe/core/column.hpp>
+#include <legate_dataframe/core/library.hpp>
 
 namespace legate::dataframe {
 
+namespace task {
+const std::set<std::string> cudf_supported_binary_ops = {"add",
+                                                         "divide",
+                                                         "multiply",
+                                                         "power",
+                                                         "subtract",
+                                                         "bit_wise_and",
+                                                         "bit_wise_or",
+                                                         "bit_wise_xor",
+                                                         "shift_left",
+                                                         "shift_right",
+                                                         "logb",
+                                                         "atan2",
+                                                         "equal",
+                                                         "greater",
+                                                         "greater_equal",
+                                                         "less",
+                                                         "less_equal",
+                                                         "not_equal",
+                                                         // logical operators:
+                                                         "and",
+                                                         "or",
+                                                         "and_kleene",
+                                                         "or_kleene"};
+
+class BinaryOpColColTask : public Task<BinaryOpColColTask, OpCode::BinaryOpColCol> {
+ public:
+  static void cpu_variant(legate::TaskContext context);
+  static void gpu_variant(legate::TaskContext context);
+};
+
+}  // namespace task
 /**
  * @brief Performs a binary operation between two columns.
  *
@@ -43,15 +73,11 @@ namespace legate::dataframe {
  * @param output_type The desired data type of the output column
  * @return            Output column of `output_type` type containing the result of
  *                    the binary operation
- * @throw cudf::logic_error if @p lhs and @p rhs are different sizes
- * @throw cudf::logic_error if @p output_type dtype isn't boolean for comparison and logical
- * operations.
- * @throw cudf::logic_error if @p output_type dtype isn't fixed-width
- * @throw cudf::data_type_error if the operation is not supported for the types of @p lhs and @p rhs
+ * @throw std::invalid_argument if operator not supported
  */
 LogicalColumn binary_operation(const LogicalColumn& lhs,
                                const LogicalColumn& rhs,
                                std::string op,
-                               cudf::data_type output_type);
+                               std::shared_ptr<arrow::DataType> output_type);
 
 }  // namespace legate::dataframe

cpp/include/legate_dataframe/core/column.hpp

@@ -82,6 +82,31 @@ class LogicalColumn {
     }
   }
 
+  /**
+   * @brief Create a column with a legate array as the data
+   *
+   * @param array The logical array (zero copy)
+   * @param data_type The arrow data type of the column. If nullptr, the arrow data type is
+   * derived from the data type of `array`.
+   * @param scalar Whether to consider this column scalar.  WARNING: currently
+   * it is the callers responsibility to ensure the array is length 1 as this
+   * check could be blocking.
+   */
+  LogicalColumn(legate::LogicalArray array,
+                std::shared_ptr<arrow::DataType> data_type,
+                bool scalar = false)
+    : array_{std::move(array)}, scalar_{scalar}
+  {
+    if (array_->dim() != 1) { throw std::invalid_argument("array must be 1-D"); }
+    // Note: Checking the volume could be blocking, so assume that this is fine.
+    assert(!scalar || array_->unbound() || array_->volume() == 1);
+
+    if (!data_type) {
+      cudf_type_ = cudf::data_type{to_cudf_type_id(array_->type().code())};
+    } else {
+      cudf_type_ = to_cudf_type(*data_type);
+    }
+  }
   /*
    * Convenience constructor for tests
    */
@@ -259,6 +284,32 @@ class LogicalColumn {
     }
   }
 
+  /**
+   * @brief Create a new unbounded column from dtype and nullable
+   *
+   * @param dtype The data type of the new column
+   * @param nullable The nullable of the new column
+   * @param scalar Whether the result is a scalar column.
+   * @return The new unbounded column
+   */
+  static LogicalColumn empty_like(std::shared_ptr<arrow::DataType> dtype,
+                                  bool nullable,
+                                  bool scalar                = false,
+                                  std::optional<size_t> size = std::nullopt)
+  {
+    if (!size.has_value()) {
+      return LogicalColumn(
+        legate::Runtime::get_runtime()->create_array(to_legate_type(*dtype), 1, nullable),
+        dtype,
+        scalar);
+    } else {
+      return LogicalColumn(legate::Runtime::get_runtime()->create_array(
+                             Shape{size.value()}, to_legate_type(*dtype), nullable),
+                           dtype,
+                           scalar);
+    }
+  }
+
  public:
   LogicalColumn(const LogicalColumn& other)            = default;
   LogicalColumn& operator=(const LogicalColumn& other) = default;

cpp/include/legate_dataframe/csv.hpp

@@ -18,13 +18,30 @@
 
 #include <string>
 
-#include <cudf/types.hpp>
-
-#include <legate_dataframe/core/column.hpp>
+#include <legate.h>
+#include <legate_dataframe/core/library.hpp>
 #include <legate_dataframe/core/table.hpp>
 
 namespace legate::dataframe {
 
+namespace task {
+
+class CSVRead : public Task<CSVRead, OpCode::CSVRead> {
+ public:
+  static void cpu_variant(legate::TaskContext context);
+  static void gpu_variant(legate::TaskContext context);
+};
+
+class CSVWrite : public Task<CSVWrite, OpCode::CSVWrite> {
+ public:
+  static constexpr auto GPU_VARIANT_OPTIONS = legate::VariantOptions{}
+                                                .with_has_allocations(true)
+                                                .with_elide_device_ctx_sync(true)
+                                                .with_has_side_effect(true);
+  static void cpu_variant(legate::TaskContext context);
+  static void gpu_variant(legate::TaskContext context);
+};
+}  // namespace task
 /**
  * @brief Write table to csv files.
  *
@@ -69,7 +86,7 @@ void csv_write(LogicalTable& tbl, const std::string& dirpath, char delimiter = '
  * @return The read LogicalTable.
  */
 LogicalTable csv_read(const std::vector<std::string>& files,
-                      const std::vector<cudf::data_type>& dtypes,
+                      const std::vector<std::shared_ptr<arrow::DataType>>& dtypes,
                       bool na_filter                                       = true,
                       char delimiter                                       = ',',
                       const std::optional<std::vector<std::string>>& names = std::nullopt,

cpp/src/binaryop.cpp

@@ -14,12 +14,6 @@
  * limitations under the License.
  */
 
-#include <legate.h>
-
-#include <cudf/binaryop.hpp>
-#include <cudf/types.hpp>
-#include <cudf/unary.hpp>
-
 #include <arrow/compute/api.h>
 #include <legate_dataframe/binaryop.hpp>
 #include <legate_dataframe/core/column.hpp>
@@ -30,140 +24,61 @@
 
 namespace legate::dataframe::task {
 
-cudf::binary_operator arrow_to_cudf_binary_op(std::string op, legate::Type output_type)
+/*static*/ void BinaryOpColColTask::cpu_variant(legate::TaskContext context)
 {
-  // Arrow binary operators taken from the below list,
-  // where an equivalent cudf binary operator exists.
-  // https://arrow.apache.org/docs/cpp/compute.html#element-wise-scalar-functions
-  // https://docs.rapids.ai/api/libcudf/stable/group__transformation__binaryops
-  std::unordered_map<std::string, cudf::binary_operator> arrow_to_cudf_ops = {
-    {"add", cudf::binary_operator::ADD},
-    {"divide", cudf::binary_operator::DIV},
-    {"multiply", cudf::binary_operator::MUL},
-    {"power", cudf::binary_operator::POW},
-    {"subtract", cudf::binary_operator::SUB},
-    {"bit_wise_and", cudf::binary_operator::BITWISE_AND},
-    {"bit_wise_or", cudf::binary_operator::BITWISE_OR},
-    {"bit_wise_xor", cudf::binary_operator::BITWISE_XOR},
-    {"shift_left", cudf::binary_operator::SHIFT_LEFT},
-    {"shift_right", cudf::binary_operator::SHIFT_RIGHT},
-    {"logb", cudf::binary_operator::LOG_BASE},
-    {"atan2", cudf::binary_operator::ATAN2},
-    {"equal", cudf::binary_operator::EQUAL},
-    {"greater", cudf::binary_operator::GREATER},
-    {"greater_equal", cudf::binary_operator::GREATER_EQUAL},
-    {"less", cudf::binary_operator::LESS},
-    {"less_equal", cudf::binary_operator::LESS_EQUAL},
-    {"not_equal", cudf::binary_operator::NOT_EQUAL},
-    // logical operators:
-    {"and", cudf::binary_operator::LOGICAL_AND},
-    {"or", cudf::binary_operator::LOGICAL_OR},
-    {"and_kleene", cudf::binary_operator::NULL_LOGICAL_AND},
-    {"or_kleene", cudf::binary_operator::NULL_LOGICAL_OR},
-  };
-
-  // Cudf has a special case for powers with integers
-  // https://github.com/rapidsai/cudf/issues/10178#issuecomment-3004143727
-  if (op == "power" && output_type.to_string().find("int") != std::string::npos) {
-    return cudf::binary_operator::INT_POW;
+  TaskContext ctx{context};
+  auto op        = argument::get_next_scalar<std::string>(ctx);
+  const auto lhs = argument::get_next_input<PhysicalColumn>(ctx);
+  const auto rhs = argument::get_next_input<PhysicalColumn>(ctx);
+  auto output    = argument::get_next_output<PhysicalColumn>(ctx);
+
+  std::vector<arrow::Datum> args(2);
+  if (lhs.num_rows() == 1) {
+    auto scalar = ARROW_RESULT(lhs.arrow_array_view()->GetScalar(0));
+    args[0]     = scalar;
+  } else {
+    args[0] = lhs.arrow_array_view();
+  }
+  if (rhs.num_rows() == 1) {
+    auto scalar = ARROW_RESULT(rhs.arrow_array_view()->GetScalar(0));
+    args[1]     = scalar;
+  } else {
+    args[1] = rhs.arrow_array_view();
   }
 
-  if (arrow_to_cudf_ops.find(op) != arrow_to_cudf_ops.end()) { return arrow_to_cudf_ops[op]; }
-  throw std::invalid_argument("Could not find cudf binary operator matching: " + op);
-  return cudf::binary_operator::INVALID_BINARY;
-}
-
-class BinaryOpColColTask : public Task<BinaryOpColColTask, OpCode::BinaryOpColCol> {
- public:
-  static void cpu_variant(legate::TaskContext context)
-  {
-    TaskContext ctx{context};
-    auto op        = argument::get_next_scalar<std::string>(ctx);
-    const auto lhs = argument::get_next_input<PhysicalColumn>(ctx);
-    const auto rhs = argument::get_next_input<PhysicalColumn>(ctx);
-    auto output    = argument::get_next_output<PhysicalColumn>(ctx);
-
-    std::vector<arrow::Datum> args(2);
-    if (lhs.num_rows() == 1) {
-      auto scalar = ARROW_RESULT(lhs.arrow_array_view()->GetScalar(0));
-      args[0]     = scalar;
-    } else {
-      args[0] = lhs.arrow_array_view();
-    }
-    if (rhs.num_rows() == 1) {
-      auto scalar = ARROW_RESULT(rhs.arrow_array_view()->GetScalar(0));
-      args[1]     = scalar;
-    } else {
-      args[1] = rhs.arrow_array_view();
-    }
-
-    if (output.cudf_type().id() == cudf::type_id::BOOL8 &&
-        (op == "and" || op == "or" || op == "and_kleene" || op == "or_kleene")) {
-      // arrow doesn't seem to cast for the user for logical ops.
-      args[0] = ARROW_RESULT(arrow::compute::Cast(args[0], arrow::boolean()));
-      args[1] = ARROW_RESULT(arrow::compute::Cast(args[1], arrow::boolean()));
-    }
-
-    // Result may be scalar or array
-    auto datum_result = ARROW_RESULT(arrow::compute::CallFunction(op, args));
+  if (output.cudf_type().id() == cudf::type_id::BOOL8 &&
+      (op == "and" || op == "or" || op == "and_kleene" || op == "or_kleene")) {
+    // arrow doesn't seem to cast for the user for logical ops.
+    args[0] = ARROW_RESULT(arrow::compute::Cast(args[0], arrow::boolean()));
+    args[1] = ARROW_RESULT(arrow::compute::Cast(args[1], arrow::boolean()));
+  }
 
-    // Coerce the output type if necessary
-    auto arrow_result_type = to_arrow_type(output.cudf_type().id());
-    if (datum_result.type() != arrow_result_type) {
-      auto coerced_result = ARROW_RESULT(arrow::compute::Cast(
-        datum_result, arrow_result_type, arrow::compute::CastOptions::Unsafe()));
-      datum_result        = std::move(coerced_result);
-    }
+  // Result may be scalar or array
+  auto datum_result = ARROW_RESULT(arrow::compute::CallFunction(op, args));
 
-    if (datum_result.is_scalar()) {
-      auto as_array = ARROW_RESULT(arrow::MakeArrayFromScalar(*datum_result.scalar(), 1));
-      if (get_prefer_eager_allocations()) {
-        output.copy_into(std::move(as_array));
-      } else {
-        output.move_into(std::move(as_array));
-      }
-    } else {
-      if (get_prefer_eager_allocations()) {
-        output.copy_into(std::move(datum_result.make_array()));
-      } else {
-        output.move_into(std::move(datum_result.make_array()));
-      }
-    }
+  // Coerce the output type if necessary
+  auto arrow_result_type = to_arrow_type(output.cudf_type().id());
+  if (datum_result.type() != arrow_result_type) {
+    auto coerced_result = ARROW_RESULT(
+      arrow::compute::Cast(datum_result, arrow_result_type, arrow::compute::CastOptions::Unsafe()));
+    datum_result = std::move(coerced_result);
   }
 
-  static void gpu_variant(legate::TaskContext context)
-  {
-    TaskContext ctx{context};
-    auto arrow_op  = argument::get_next_scalar<std::string>(ctx);
-    const auto lhs = argument::get_next_input<PhysicalColumn>(ctx);
-    const auto rhs = argument::get_next_input<PhysicalColumn>(ctx);
-    auto output    = argument::get_next_output<PhysicalColumn>(ctx);
-    auto op        = arrow_to_cudf_binary_op(arrow_op, output.type());
-
-    std::unique_ptr<cudf::column> ret;
-    /*
-     * If one (not both) are length 1, use scalars as cudf doesn't allow
-     * broadcast binary operations.
-     */
-    if (lhs.num_rows() == 1 && rhs.num_rows() != 1) {
-      auto lhs_scalar = lhs.cudf_scalar();
-      ret             = cudf::binary_operation(
-        *lhs_scalar, rhs.column_view(), op, output.cudf_type(), ctx.stream(), ctx.mr());
-    } else if (rhs.num_rows() == 1 && lhs.num_rows() != 1) {
-      auto rhs_scalar = rhs.cudf_scalar();
-      ret             = cudf::binary_operation(
-        lhs.column_view(), *rhs_scalar, op, output.cudf_type(), ctx.stream(), ctx.mr());
+  if (datum_result.is_scalar()) {
+    auto as_array = ARROW_RESULT(arrow::MakeArrayFromScalar(*datum_result.scalar(), 1));
+    if (get_prefer_eager_allocations()) {
+      output.copy_into(std::move(as_array));
     } else {
-      ret = cudf::binary_operation(
-        lhs.column_view(), rhs.column_view(), op, output.cudf_type(), ctx.stream(), ctx.mr());
+      output.move_into(std::move(as_array));
     }
+  } else {
     if (get_prefer_eager_allocations()) {
-      output.copy_into(std::move(ret));
+      output.copy_into(std::move(datum_result.make_array()));
     } else {
-      output.move_into(std::move(ret));
+      output.move_into(std::move(datum_result.make_array()));
     }
   }
-};
+}
 
 }  // namespace legate::dataframe::task
 
@@ -181,15 +96,16 @@ namespace legate::dataframe {
 LogicalColumn binary_operation(const LogicalColumn& lhs,
                                const LogicalColumn& rhs,
                                std::string op,
-                               cudf::data_type output_type)
+                               std::shared_ptr<arrow::DataType> output_type)
 {
   auto runtime = legate::Runtime::get_runtime();
 
   // Check if the op is valid before we enter the task
   // This allows us to to throw nicely
   if (runtime->get_machine().count(legate::mapping::TaskTarget::GPU) > 0) {
-    // Throws if op doesn't exist
-    task::arrow_to_cudf_binary_op(op, to_legate_type(output_type.id()));
+    if (task::cudf_supported_binary_ops.count(op) == 0) {
+      throw std::invalid_argument("Unsupported binary operator: " + op);
+    }
   } else {
     auto result = arrow::compute::GetFunctionRegistry()->GetFunction(op);
     if (!result.ok()) {
@@ -201,7 +117,7 @@ LogicalColumn binary_operation(const LogicalColumn& lhs,
   auto scalar_result = lhs.is_scalar() && rhs.is_scalar();
   std::optional<size_t> size{};
   if (get_prefer_eager_allocations()) { size = lhs.is_scalar() ? rhs.num_rows() : lhs.num_rows(); }
-  auto ret = LogicalColumn::empty_like(std::move(output_type), nullable, scalar_result, size);
+  auto ret = LogicalColumn::empty_like(output_type, nullable, scalar_result, size);
   legate::AutoTask task =
     runtime->create_task(get_library(), task::BinaryOpColColTask::TASK_CONFIG.task_id());
   argument::add_next_scalar(task, op);