Change raw pointers in ast_plan to span.

rapidsai · May 3, 2021 · f0a5c66 · f0a5c66
1 parent 8931591
commit f0a5c66
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 16 deletions.
diff --git a/cpp/include/cudf/ast/detail/transform.cuh b/cpp/include/cudf/ast/detail/transform.cuh
@@ -26,6 +26,7 @@
 #include <cudf/table/table_device_view.cuh>
 #include <cudf/table/table_view.hpp>
 #include <cudf/types.hpp>
+#include <cudf/utilities/span.hpp>
 #include <cudf/utilities/traits.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
@@ -371,14 +372,19 @@ struct ast_plan {
 
     // Create device pointers to components of plan
     auto device_data_buffer_ptr = static_cast<const char*>(_device_data_buffer.data());
-    _device_data_references     = reinterpret_cast<const detail::device_data_reference*>(
-      device_data_buffer_ptr + buffer_offsets[0]);
-    _device_literals = reinterpret_cast<const cudf::detail::fixed_width_scalar_device_view_base*>(
-      device_data_buffer_ptr + buffer_offsets[1]);
-    _device_operators =
-      reinterpret_cast<const ast_operator*>(device_data_buffer_ptr + buffer_offsets[2]);
-    _device_operator_source_indices =
-      reinterpret_cast<const cudf::size_type*>(device_data_buffer_ptr + buffer_offsets[3]);
+    _device_data_references     = device_span<const detail::device_data_reference>(
+      reinterpret_cast<const detail::device_data_reference*>(device_data_buffer_ptr +
+                                                             buffer_offsets[0]),
+      _sizes[0]);
+    _device_literals = device_span<const cudf::detail::fixed_width_scalar_device_view_base>(
+      reinterpret_cast<const cudf::detail::fixed_width_scalar_device_view_base*>(
+        device_data_buffer_ptr + buffer_offsets[1]),
+      _sizes[1]);
+    _device_operators = device_span<const ast_operator>(
+      reinterpret_cast<const ast_operator*>(device_data_buffer_ptr + buffer_offsets[2]), _sizes[2]);
+    _device_operator_source_indices = device_span<const cudf::size_type>(
+      reinterpret_cast<const cudf::size_type*>(device_data_buffer_ptr + buffer_offsets[3]),
+      _sizes[3]);
   }
 
   /**
@@ -399,10 +405,10 @@ struct ast_plan {
   std::vector<const void*> _data_pointers;
 
   rmm::device_buffer _device_data_buffer;
-  const detail::device_data_reference* _device_data_references;
-  const cudf::detail::fixed_width_scalar_device_view_base* _device_literals;
-  const ast_operator* _device_operators;
-  const cudf::size_type* _device_operator_source_indices;
+  device_span<const detail::device_data_reference> _device_data_references;
+  device_span<const cudf::detail::fixed_width_scalar_device_view_base> _device_literals;
+  device_span<const ast_operator> _device_operators;
+  device_span<const cudf::size_type> _device_operator_source_indices;
 };
 
 /**

diff --git a/cpp/src/ast/transform.cu b/cpp/src/ast/transform.cu
@@ -123,11 +123,11 @@ std::unique_ptr<column> compute_column(table_view const table,
   cudf::ast::detail::compute_column_kernel<MAX_BLOCK_SIZE>
     <<<config.num_blocks, config.num_threads_per_block, shmem_size_per_block, stream.value()>>>(
       *table_device,
-      plan._device_literals,
+      plan._device_literals.data(),
       *mutable_output_device,
-      plan._device_data_references,
-      plan._device_operators,
-      plan._device_operator_source_indices,
+      plan._device_data_references.data(),
+      plan._device_operators.data(),
+      plan._device_operator_source_indices.data(),
       num_operators,
       num_intermediates);
   CHECK_CUDA(stream.value());