cudf/core/memory/
options.rs

1//! Physical-base options and validation for the device memory manager.
2
3use crate::core::error::{self, Result};
4
5/// Physical base used by the process-local device memory manager.
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum DeviceMemoryManagerBase {
8    /// Use CUDA's stream-ordered async memory pool.
9    CudaAsync {
10        /// Initial pool size in bytes; `None` or zero selects RMM's default.
11        initial_pool_size: Option<usize>,
12        /// Pool release threshold in bytes; `None` or `usize::MAX` selects RMM's default.
13        release_threshold: Option<usize>,
14    },
15    /// Use RMM's pool resource around the raw CUDA resource.
16    Pool {
17        /// Initial RMM pool size in bytes.
18        initial_pool_size: usize,
19        /// Maximum RMM pool size in bytes; `None` or zero selects the native default.
20        maximum_pool_size: Option<usize>,
21    },
22}
23
24impl Default for DeviceMemoryManagerBase {
25    fn default() -> Self {
26        Self::CudaAsync { initial_pool_size: None, release_threshold: None }
27    }
28}
29
30/// Options applied when a device memory manager is first created.
31#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
32pub struct DeviceMemoryManagerOptions {
33    /// Physical allocator selected when the process-local device manager is created.
34    pub base: DeviceMemoryManagerBase,
35    /// Immutable hard limit enforced by the process manager. `None` leaves the
36    /// selected physical base unpartitioned.
37    pub allocation_limit_bytes: Option<usize>,
38}
39
40impl DeviceMemoryManagerOptions {
41    fn normalized(self) -> Self {
42        let base = match self.base {
43            DeviceMemoryManagerBase::CudaAsync { initial_pool_size, release_threshold } => {
44                DeviceMemoryManagerBase::CudaAsync {
45                    initial_pool_size: initial_pool_size.filter(|size| *size != 0),
46                    release_threshold: release_threshold.filter(|threshold| *threshold != usize::MAX),
47                }
48            },
49            DeviceMemoryManagerBase::Pool { initial_pool_size, maximum_pool_size } => DeviceMemoryManagerBase::Pool {
50                initial_pool_size,
51                maximum_pool_size: maximum_pool_size.filter(|size| *size != 0),
52            },
53        };
54        Self { base, allocation_limit_bytes: self.allocation_limit_bytes }
55    }
56
57    /// Validate this configuration without consulting CUDA or process-global
58    /// allocator state.
59    pub fn validate(&self) -> Result<()> {
60        if self.allocation_limit_bytes == Some(0) {
61            return Err(invalid_memory_manager_options(
62                "allocation_limit_bytes",
63                "allocation limit must be greater than zero when configured",
64            ));
65        }
66        let normalized = (*self).normalized();
67        let Some(limit) = normalized.allocation_limit_bytes else {
68            if let DeviceMemoryManagerBase::Pool { initial_pool_size, maximum_pool_size: Some(maximum) } =
69                normalized.base
70                && maximum < initial_pool_size
71            {
72                return Err(invalid_memory_manager_options(
73                    "maximum_pool_size",
74                    "maximum pool size must not be below the initial pool size",
75                ));
76            }
77            return Ok(());
78        };
79        match normalized.base {
80            DeviceMemoryManagerBase::CudaAsync { initial_pool_size: Some(initial), .. } if initial > limit => Err(
81                invalid_memory_manager_options("initial_pool_size", "initial pool size exceeds the allocation limit"),
82            ),
83            DeviceMemoryManagerBase::Pool { initial_pool_size, .. } if initial_pool_size > limit => Err(
84                invalid_memory_manager_options("initial_pool_size", "initial pool size exceeds the allocation limit"),
85            ),
86            DeviceMemoryManagerBase::Pool { initial_pool_size, maximum_pool_size: Some(maximum) }
87                if maximum < initial_pool_size =>
88            {
89                Err(invalid_memory_manager_options(
90                    "maximum_pool_size",
91                    "maximum pool size must not be below the initial pool size",
92                ))
93            },
94            DeviceMemoryManagerBase::Pool { maximum_pool_size: Some(maximum), .. } if maximum > limit => Err(
95                invalid_memory_manager_options("maximum_pool_size", "maximum pool size exceeds the allocation limit"),
96            ),
97            _ => Ok(()),
98        }
99    }
100
101    pub(super) fn to_sys(self) -> Result<cudf_sys::cudf_memory_manager_options_t> {
102        self.validate()?;
103        let mut out = cudf_sys::cudf_memory_manager_options_t {
104            allocation_limit_bytes: self.allocation_limit_bytes.unwrap_or(0) as u64,
105            ..Default::default()
106        };
107        match self.base {
108            DeviceMemoryManagerBase::CudaAsync { initial_pool_size, release_threshold } => {
109                out.base_kind = cudf_sys::CUDF_MEMORY_MANAGER_BASE_CUDA_ASYNC;
110                out.initial_pool_size_bytes = initial_pool_size.unwrap_or(0) as u64;
111                out.release_threshold_bytes = release_threshold.unwrap_or(usize::MAX) as u64;
112            },
113            DeviceMemoryManagerBase::Pool { initial_pool_size, maximum_pool_size } => {
114                out.base_kind = cudf_sys::CUDF_MEMORY_MANAGER_BASE_POOL;
115                out.initial_pool_size_bytes = initial_pool_size as u64;
116                out.maximum_pool_size_bytes = maximum_pool_size.unwrap_or(0) as u64;
117            },
118        }
119        Ok(out)
120    }
121
122    pub(super) fn from_sys(value: cudf_sys::cudf_memory_manager_options_t) -> Result<Self> {
123        let initial_pool_size = memory_manager_option_bytes(value.initial_pool_size_bytes, "initial_pool_size_bytes")?;
124        let base = match value.base_kind {
125            cudf_sys::CUDF_MEMORY_MANAGER_BASE_CUDA_ASYNC => {
126                let release_threshold = if value.release_threshold_bytes == u64::MAX {
127                    None
128                } else {
129                    Some(memory_manager_option_bytes(value.release_threshold_bytes, "release_threshold_bytes")?)
130                };
131                DeviceMemoryManagerBase::CudaAsync {
132                    initial_pool_size: (initial_pool_size != 0).then_some(initial_pool_size),
133                    release_threshold,
134                }
135            },
136            cudf_sys::CUDF_MEMORY_MANAGER_BASE_POOL => {
137                let maximum_pool_size = if value.maximum_pool_size_bytes == 0 {
138                    None
139                } else {
140                    Some(memory_manager_option_bytes(value.maximum_pool_size_bytes, "maximum_pool_size_bytes")?)
141                };
142                DeviceMemoryManagerBase::Pool { initial_pool_size, maximum_pool_size }
143            },
144            other => {
145                return Err(error::Error::new_with(
146                    error::ErrorCode::ForeignInvariantViolation,
147                    "memory::device_memory_manager::options_for_device",
148                    "FFI returned an unknown memory-manager base kind",
149                )
150                .with_detail("base_kind", other.to_string()));
151            },
152        };
153        let allocation_limit_bytes = if value.allocation_limit_bytes == 0 {
154            None
155        } else {
156            Some(memory_manager_option_bytes(value.allocation_limit_bytes, "allocation_limit_bytes")?)
157        };
158        let options = Self { base, allocation_limit_bytes };
159        options.validate()?;
160        Ok(options)
161    }
162}
163
164fn invalid_memory_manager_options(field: &'static str, message: &'static str) -> error::Error {
165    error::Error::new_with(error::ErrorCode::ValueOutOfRange, "memory::device_memory_manager::options", message)
166        .with_detail("field", field)
167        .with_detail("reason", "invalid_device_memory_manager_options")
168}
169
170fn memory_manager_option_bytes(value: u64, field: &'static str) -> Result<usize> {
171    usize::try_from(value).map_err(|_| {
172        error::Error::new_with(
173            error::ErrorCode::ValueOutOfRange,
174            "memory::device_memory_manager::options_for_device",
175            "installed device memory manager option does not fit usize",
176        )
177        .with_detail("field", field)
178        .with_detail("value", value.to_string())
179    })
180}
181
182/// Options for an opt-in query memory domain.
183#[derive(Debug, Clone, Copy, PartialEq, Eq)]
184pub struct QueryMemoryDomainOptions {
185    pub grant_bytes: usize,
186    pub ceiling_bytes: usize,
187}
188
189impl QueryMemoryDomainOptions {
190    /// Create a domain whose initial grant and immutable logical ceiling are equal.
191    pub fn new(grant_bytes: usize) -> Self {
192        Self { grant_bytes, ceiling_bytes: grant_bytes }
193    }
194
195    /// Set the immutable logical ceiling for grant resizing.
196    pub fn with_ceiling_bytes(mut self, ceiling_bytes: usize) -> Self {
197        self.ceiling_bytes = ceiling_bytes;
198        self
199    }
200}
cudf/core/memory/options.rs

cudf/core/memory/
options.rs
