Implementation:Rapidsai Cuml Raft Proto Buffer
| Knowledge Sources | |
|---|---|
| Domains | Machine_Learning, Memory_Management |
| Last Updated | 2026-02-08 12:00 GMT |
Overview
Provides a unified memory buffer abstraction in the raft_proto namespace that can hold data on either CPU or GPU, supporting owning and non-owning semantics with transparent cross-device copy operations.
Description
The buffer.hpp header defines the raft_proto::buffer<T> template class, a versatile container used throughout cuML's Forest Inference Library (FIL) for managing data that may reside on host (CPU) or device (GPU) memory.
Key Design Features:
- Dual memory support: The buffer can hold data on CPU or GPU, tracked by a
device_typeenum and adevice_id_variant. - Owning and non-owning modes: The internal
data_storeis astd::variantof four types: non-owning CPU/GPU buffers (wrapping existing pointers) and owning CPU/GPU buffers (managing their own allocations). - Cross-device copy construction: A buffer can be constructed from another buffer in a different memory location, automatically performing a host-to-device, device-to-host, or device-to-device copy.
- Move semantics with optional copy: Move construction avoids copies when the source and destination are on the same device; otherwise it performs a copy to the target device.
- Iterator construction: Buffers can be constructed from iterator pairs, copying elements to CPU and optionally to GPU.
Accessors:
size(): Number of elements.data(): Raw pointer to the underlying data (works on both host and device via HOST/DEVICE annotations).memory_type(): Returnsdevice_type::cpuordevice_type::gpu.device_index(): Returns the device ordinal.
Free Functions:
copy: Template functions for copying between buffers with optional bounds checking (bounds_checktemplate parameter), offset support, and CUDA stream specification. Multiple overloads support lvalue references, rvalue references, and various parameter combinations.
Usage
Use the buffer class when writing code in cuML's FIL that needs to work transparently with both CPU and GPU data. It is particularly useful for model loading and inference pipelines where data may need to be transferred between host and device. The non-owning mode enables zero-copy wrapping of existing allocations, while the owning mode handles memory lifecycle automatically.
Code Reference
Source Location
- Repository: Rapidsai_Cuml
- File:
cpp/include/cuml/fil/detail/raft_proto/buffer.hpp
Signature
namespace raft_proto {
template <typename T>
struct buffer {
using index_type = std::size_t;
using value_type = T;
// Default constructor
buffer();
// Owning buffer (uninitialized allocation)
buffer(index_type size, device_type mem_type = device_type::cpu,
int device = 0, cuda_stream stream = 0);
// Non-owning buffer (wraps existing pointer)
buffer(T* input_data, index_type size,
device_type mem_type = device_type::cpu, int device = 0);
// Copy-construct to different memory location
buffer(buffer<T> const& other, device_type mem_type,
int device = 0, cuda_stream stream = cuda_stream{});
// Copy-construct (same memory type)
buffer(buffer<T> const& other, cuda_stream stream = cuda_stream{});
// Move with optional cross-device copy
buffer(buffer<T>&& other, device_type mem_type, int device, cuda_stream stream);
buffer(buffer<T>&& other, device_type mem_type, int device);
buffer(buffer<T>&& other, device_type mem_type);
buffer(buffer<T>&& other) noexcept;
// Iterator construction
template <typename iter_t>
buffer(iter_t const& begin, iter_t const& end);
template <typename iter_t>
buffer(iter_t const& begin, iter_t const& end, device_type mem_type);
template <typename iter_t>
buffer(iter_t const& begin, iter_t const& end, device_type mem_type,
int device, cuda_stream stream = cuda_stream{});
auto size() const noexcept;
auto* data() const noexcept;
auto memory_type() const noexcept;
auto device() const noexcept;
auto device_index() const noexcept;
buffer<T>& operator=(buffer<T> const& other);
buffer<T>& operator=(buffer<T>&& other) noexcept;
friend void swap(buffer<T>& first, buffer<T>& second);
};
// Copy functions with bounds checking
template <bool bounds_check, typename T, typename U>
void copy(buffer<T>& dst, buffer<U> const& src,
typename buffer<T>::index_type dst_offset,
typename buffer<U>::index_type src_offset,
typename buffer<T>::index_type size,
cuda_stream stream);
template <bool bounds_check, typename T, typename U>
void copy(buffer<T>& dst, buffer<U> const& src, cuda_stream stream);
template <bool bounds_check, typename T, typename U>
void copy(buffer<T>& dst, buffer<U> const& src);
} // namespace raft_proto
Import
#include <cuml/fil/detail/raft_proto/buffer.hpp>
I/O Contract
Inputs
Owning constructor
| Name | Type | Required | Description |
|---|---|---|---|
| size | index_type | Yes | Number of elements to allocate |
| mem_type | device_type | No | Memory location: cpu or gpu (default: cpu) |
| device | int | No | Device ordinal (default: 0) |
| stream | cuda_stream | No | CUDA stream for GPU allocation (default: 0) |
Non-owning constructor
| Name | Type | Required | Description |
|---|---|---|---|
| input_data | T* | Yes | Pointer to existing data |
| size | index_type | Yes | Number of elements |
| mem_type | device_type | No | Memory type of the existing pointer (default: cpu) |
| device | int | No | Device ordinal (default: 0) |
copy function
| Name | Type | Required | Description |
|---|---|---|---|
| dst | buffer<T>& | Yes | Destination buffer |
| src | buffer const& | Yes | Source buffer |
| dst_offset | index_type | No | Offset in destination (default: 0) |
| src_offset | index_type | No | Offset in source (default: 0) |
| size | index_type | No | Number of elements to copy (default: src.size()) |
| stream | cuda_stream | No | CUDA stream (default: default stream) |
Outputs
| Name | Type | Description |
|---|---|---|
| data() | T* | Raw pointer to the buffer's data (host or device) |
| size() | index_type | Number of elements in the buffer |
| memory_type() | device_type | Whether data resides on cpu or gpu |
| device_index() | int | Ordinal of the device holding the data |
Usage Examples
#include <cuml/fil/detail/raft_proto/buffer.hpp>
using raft_proto::buffer;
using raft_proto::device_type;
// Create an owning CPU buffer
buffer<float> cpu_buf(1000, device_type::cpu);
// Create an owning GPU buffer
buffer<float> gpu_buf(1000, device_type::gpu, 0);
// Wrap existing device pointer (non-owning)
float* d_ptr = /* existing GPU allocation */;
buffer<float> gpu_view(d_ptr, 500, device_type::gpu, 0);
// Copy CPU buffer to GPU
buffer<float> gpu_copy(cpu_buf, device_type::gpu, 0);
// Copy between buffers with bounds checking
raft_proto::copy<true>(gpu_buf, cpu_buf);
// Copy with offset
raft_proto::copy<true>(gpu_buf, cpu_buf, 100, 0, 500, cuda_stream{});
// Create from iterators
std::vector<float> host_data = {1.0f, 2.0f, 3.0f};
buffer<float> from_iter(host_data.begin(), host_data.end());
// Move to GPU
buffer<float> on_gpu(std::move(from_iter), device_type::gpu);