table_device_view.cuh

转到此文件的文档。

/*
 * Copyright (c) 2019-2024, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once

#include <cudf/column/column_device_view.cuh>
#include <cudf/detail/utilities/cuda_memcpy.hpp>
#include <cudf/detail/utilities/vector_factories.hpp>
#include <cudf/table/table_view.hpp>
#include <cudf/types.hpp>
#include <cudf/utilities/default_stream.hpp>

#include <rmm/cuda_stream_view.hpp>
#include <rmm/device_buffer.hpp>

#include <cassert>
#include <memory>
#include <numeric>

namespace CUDF_EXPORT cudf {
namespace detail {

template <typename ColumnDeviceView, typename HostTableView>
class table_device_view_base {
 public
 table_device_view_base() = delete;
 ~table_device_view_base() = default;
 table_device_view_base(table_device_view_base const&) = default;
 table_device_view_base(table_device_view_base&&) = default;
 table_device_view_base& operator=(table_device_view_base const&) = default;
 table_device_view_base& operator=(table_device_view_base&&) = default;

 __device__ ColumnDeviceView* begin() const noexcept { return _columns; }

 __device__ ColumnDeviceView* end() const noexcept { return _columns + _num_columns; }

 __device__ ColumnDeviceView const& column(size_type column_index) const noexcept
 {
 assert(column_index >= 0);
 assert(column_index < _num_columns);
 return _columns[column_index];
 }

 __device__ ColumnDeviceView& column(size_type column_index) noexcept
 {
 assert(column_index >= 0);
 assert(column_index < _num_columns);
 return _columns[column_index];
 }

 [[nodiscard]] __host__ __device__ size_type num_columns() const noexcept { return _num_columns; }

 [[nodiscard]] __host__ __device__ size_type num_rows() const noexcept { return _num_rows; }

 void destroy();

 private
 ColumnDeviceView* _columns{};
 size_type _num_rows{};
 size_type _num_columns{};

 protected
 table_device_view_base(HostTableView source_view, rmm::cuda_stream_view stream);

 rmm::device_buffer* _descendant_storage{};
};
} // namespace detail

class table_device_view : public detail::table_device_view_base<column_device_view, table_view> {
 public
 static auto create(table_view source_view,
 rmm::cuda_stream_view stream = cudf::get_default_stream())
 {
 auto deleter = [](table_device_view* t) { t->destroy(); };
 return std::unique_ptr<table_device_view, decltype(deleter)>{
 new table_device_view(source_view, stream), deleter};
 }

 private
 table_device_view(table_view source_view, rmm::cuda_stream_view stream)
 : detail::table_device_view_base<column_device_view, table_view>(source_view, stream)
 {
 }
};

class mutable_table_device_view
 : public detail::table_device_view_base<mutable_column_device_view, mutable_table_view> {
 public
 static auto create(mutable_table_view source_view,
 rmm::cuda_stream_view stream = cudf::get_default_stream())
 {
 auto deleter = [](mutable_table_device_view* t) { t->destroy(); };
 return std::unique_ptr<mutable_table_device_view, decltype(deleter)>{
 new mutable_table_device_view(source_view, stream), deleter};
 }

 private
 mutable_table_device_view(mutable_table_view source_view, rmm::cuda_stream_view stream)
 : detail::table_device_view_base<mutable_column_device_view, mutable_table_view>(source_view,
 stream)
 {
 }
};

template <typename ColumnDeviceView, typename HostTableView>
auto contiguous_copy_column_device_views(HostTableView source_view, rmm::cuda_stream_view stream)
{
 // First calculate the size of memory needed to hold the
 // table's ColumnDeviceViews. This is done by calling extent()
 // for each of the table's ColumnViews columns.
 std::size_t views_size_bytes = std::accumulate(
 source_view.begin(), source_view.end(), std::size_t{0}, [](std::size_t init, auto col) {
 return init + ColumnDeviceView::extent(col);
 });
 // pad the allocation for aligning the first pointer
 auto padded_views_size_bytes = views_size_bytes + std::size_t{alignof(ColumnDeviceView) - 1};
 // A buffer of CPU memory is allocated to hold the ColumnDeviceView
 // objects. Once filled, the CPU memory is then copied to device memory
 // and the pointer is set in the d_columns member.
 auto h_buffer = cudf::detail::make_host_vector<int8_t>(padded_views_size_bytes, stream);
 // Each ColumnDeviceView instance may have child objects which may
 // require setting some internal device pointers before being copied
 // from CPU to device.
 // Allocate the device memory to be used in the result.
 // We need this pointer in order to pass it down when creating the
 // ColumnDeviceViews so the column can set the pointer(s) for any
 // of its child objects.
 // align both h_ptr, d_ptr
 auto descendant_storage = std::make_unique<rmm::device_buffer>(padded_views_size_bytes, stream);
 void* h_ptr = detail::align_ptr_for_type<ColumnDeviceView>(h_buffer.data());
 void* d_ptr = detail::align_ptr_for_type<ColumnDeviceView>(descendant_storage->data());
 auto d_columns = detail::child_columns_to_device_array<ColumnDeviceView>(
 source_view.begin(), source_view.end(), h_ptr, d_ptr);

 auto const h_span = host_span<int8_t const>{h_buffer}.subspan(
 static_cast<int8_t const*>(h_ptr) - h_buffer.data(), views_size_bytes);
 auto const d_span = device_span<int8_t>{static_cast<int8_t*>(d_ptr), views_size_bytes};
 cudf::detail::cuda_memcpy(d_span, h_span, stream);
 return std::make_tuple(std::move(descendant_storage), d_columns);
}

} // namespace CUDF_EXPORT cudf