#[cfg(feature = "gpu")]
use std::error::Error;
use crate::budget::GpuMemoryBudget;
#[cfg(feature = "gpu")]
use crate::budget::DEFAULT_SAFE_VRAM_BYTES;
use crate::config::ExecutorConfig;
use crate::error::ExecutorInitError;
pub const GPU_AVAILABLE_BYTES_ENV: &str = "REV_GPU_AVAILABLE_BYTES";
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct MemoryBudget {
available_bytes: usize,
target_batch_bytes: usize,
}
impl MemoryBudget {
pub fn from_config(config: &ExecutorConfig) -> Result<Self, ExecutorInitError> {
if !(0.0..=1.0).contains(&config.memory_fill_ratio) || config.memory_fill_ratio == 0.0 {
return Err(ExecutorInitError::InvalidConfig(
"memory_fill_ratio must be greater than 0.0 and no more than 1.0".to_string(),
));
}
if config.min_batch_bytes > config.max_batch_bytes {
return Err(ExecutorInitError::InvalidConfig(
"min_batch_bytes cannot exceed max_batch_bytes".to_string(),
));
}
let available_bytes = config.vram_override.unwrap_or(config.max_batch_bytes);
let fill_target = ((available_bytes as f64) * f64::from(config.memory_fill_ratio)) as usize;
let target_batch_bytes = fill_target
.max(config.min_batch_bytes)
.min(config.max_batch_bytes)
.max(1);
Ok(Self {
available_bytes,
target_batch_bytes,
})
}
pub fn available_bytes(self) -> usize {
self.available_bytes
}
pub fn target_batch_bytes(self) -> usize {
self.target_batch_bytes
}
pub fn gpu_memory_budget(
self,
fill_ratio: f32,
) -> Result<GpuMemoryBudget, crate::error::GpuError> {
GpuMemoryBudget::new(self.available_bytes, fill_ratio)
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct GpuMemoryAvailability {
pub budget_bytes: usize,
pub current_usage_bytes: usize,
pub available_bytes: usize,
pub available_for_reservation_bytes: usize,
pub current_reservation_bytes: usize,
pub source: &'static str,
pub is_estimate: bool,
pub probe_error: Option<String>,
}
#[cfg(feature = "gpu")]
pub fn resolve_gpu_memory_availability(
adapter_vendor: u32,
adapter_device: u32,
adapter_name: &str,
adapter_limits: &wgpu::Limits,
) -> GpuMemoryAvailability {
match query_gpu_memory_availability(adapter_vendor, adapter_device, adapter_name) {
Ok(availability) => availability,
Err(error) => availability_from_env()
.unwrap_or_else(|| estimate_gpu_memory_availability_from_limits(adapter_limits, error)),
}
}
#[cfg(feature = "gpu")]
pub fn estimate_gpu_memory_availability_from_limits(
adapter_limits: &wgpu::Limits,
probe_error: impl ToString,
) -> GpuMemoryAvailability {
let usable_storage_bindings = adapter_limits
.max_storage_buffers_per_shader_stage
.saturating_sub(1)
.max(1) as usize;
let storage_binding_window = (adapter_limits.max_storage_buffer_binding_size as usize)
.saturating_mul(usable_storage_bindings);
let buffer_window =
(adapter_limits.max_buffer_size as usize).saturating_mul(usable_storage_bindings);
let estimated_budget = storage_binding_window
.min(buffer_window)
.max(adapter_limits.max_storage_buffer_binding_size as usize)
.max(DEFAULT_SAFE_VRAM_BYTES);
GpuMemoryAvailability {
budget_bytes: estimated_budget,
current_usage_bytes: 0,
available_bytes: estimated_budget,
available_for_reservation_bytes: estimated_budget,
current_reservation_bytes: 0,
source: "wgpu.limits.estimate",
is_estimate: true,
probe_error: Some(probe_error.to_string()),
}
}
#[cfg(all(feature = "gpu", target_os = "windows"))]
pub fn query_gpu_memory_availability(
adapter_vendor: u32,
adapter_device: u32,
adapter_name: &str,
) -> Result<GpuMemoryAvailability, Box<dyn Error>> {
use windows::core::Interface;
use windows::Win32::Graphics::Dxgi::{
CreateDXGIFactory1, IDXGIAdapter3, IDXGIFactory1, DXGI_MEMORY_SEGMENT_GROUP_LOCAL,
DXGI_QUERY_VIDEO_MEMORY_INFO,
};
let factory: IDXGIFactory1 = unsafe { CreateDXGIFactory1()? };
let mut index = 0;
loop {
let adapter = match unsafe { factory.EnumAdapters1(index) } {
Ok(adapter) => adapter,
Err(error)
if error.code().0 == windows::Win32::Graphics::Dxgi::DXGI_ERROR_NOT_FOUND.0 =>
{
break;
}
Err(error) => return Err(error.into()),
};
let desc = unsafe { adapter.GetDesc1()? };
let dxgi_name = utf16z_to_string(&desc.Description);
let same_adapter = (desc.VendorId == adapter_vendor && desc.DeviceId == adapter_device)
|| names_match(&dxgi_name, adapter_name);
if same_adapter {
let adapter3: IDXGIAdapter3 = adapter.cast()?;
let mut info = DXGI_QUERY_VIDEO_MEMORY_INFO::default();
unsafe {
adapter3.QueryVideoMemoryInfo(0, DXGI_MEMORY_SEGMENT_GROUP_LOCAL, &mut info)?;
}
let budget = usize::try_from(info.Budget).unwrap_or(usize::MAX);
let current_usage = usize::try_from(info.CurrentUsage).unwrap_or(usize::MAX);
return Ok(GpuMemoryAvailability {
budget_bytes: budget,
current_usage_bytes: current_usage,
available_bytes: budget.saturating_sub(current_usage),
available_for_reservation_bytes: usize::try_from(info.AvailableForReservation)
.unwrap_or(usize::MAX),
current_reservation_bytes: usize::try_from(info.CurrentReservation)
.unwrap_or(usize::MAX),
source: "dxgi.local",
is_estimate: false,
probe_error: None,
});
}
index += 1;
}
Err(format!(
"no DXGI adapter matched vendor={adapter_vendor:#x} device={adapter_device:#x} name=\"{adapter_name}\""
)
.into())
}
#[cfg(all(feature = "gpu", not(target_os = "windows")))]
pub fn query_gpu_memory_availability(
_adapter_vendor: u32,
_adapter_device: u32,
_adapter_name: &str,
) -> Result<GpuMemoryAvailability, Box<dyn Error>> {
Err("runtime GPU memory availability probing is not implemented for this OS".into())
}
#[cfg(feature = "gpu")]
fn availability_from_env() -> Option<GpuMemoryAvailability> {
let bytes = std::env::var(GPU_AVAILABLE_BYTES_ENV)
.ok()
.and_then(|value| value.parse::<usize>().ok())
.filter(|bytes| *bytes > 0)?;
Some(GpuMemoryAvailability {
budget_bytes: bytes,
current_usage_bytes: 0,
available_bytes: bytes,
available_for_reservation_bytes: bytes,
current_reservation_bytes: 0,
source: "env.REV_GPU_AVAILABLE_BYTES",
is_estimate: true,
probe_error: None,
})
}
#[cfg(all(feature = "gpu", target_os = "windows"))]
fn utf16z_to_string(value: &[u16]) -> String {
let len = value
.iter()
.position(|code| *code == 0)
.unwrap_or(value.len());
String::from_utf16_lossy(&value[..len])
}
#[cfg(all(feature = "gpu", target_os = "windows"))]
fn names_match(dxgi_name: &str, adapter_name: &str) -> bool {
fn normalize(value: &str) -> String {
value
.chars()
.filter(|ch| ch.is_ascii_alphanumeric())
.flat_map(char::to_lowercase)
.collect()
}
normalize(dxgi_name) == normalize(adapter_name)
}