use std::time::Instant;
#[cfg(feature = "gpu")]
use std::sync::mpsc;
#[cfg(feature = "gpu")]
use std::sync::Arc;
use crate::batch::BatchSpan;
use crate::budget::GpuMemoryBudget;
use crate::config::{ExecutionMode, ExecutorConfig};
use crate::error::{ExecutionError, ExecutorInitError, GpuError, GpuFailure};
#[cfg(feature = "gpu")]
use crate::job::BufferRole;
use crate::job::{EncodedBatch, GpuJob, JobError};
use crate::memory::MemoryBudget;
use crate::planner::{elapsed_ms, plan_batches};
use crate::report::{BatchPath, BatchReport, ExecutionPath, JobReport};
use crate::resources::ResourceCache;
use crate::telemetry::{TelemetryCollector, TelemetryReport};
#[derive(Debug)]
pub struct GpuExecutor {
config: ExecutorConfig,
#[cfg(feature = "gpu")]
device: Option<Arc<wgpu::Device>>,
#[cfg(feature = "gpu")]
queue: Option<Arc<wgpu::Queue>>,
budget: MemoryBudget,
resource_cache: ResourceCache,
telemetry: TelemetryCollector,
}
impl GpuExecutor {
pub fn new() -> Result<Self, ExecutorInitError> {
Self::with_config(ExecutorConfig::from_default_json_file_or_default()?)
}
pub fn with_config(config: ExecutorConfig) -> Result<Self, ExecutorInitError> {
let budget = MemoryBudget::from_config(&config)?;
#[cfg(feature = "gpu")]
let (device, queue) = maybe_init_gpu(&config)?;
#[cfg(not(feature = "gpu"))]
if matches!(config.execution_mode, ExecutionMode::GpuOnly) {
return Err(ExecutorInitError::NoGpuAvailable);
}
Ok(Self {
config,
#[cfg(feature = "gpu")]
device,
#[cfg(feature = "gpu")]
queue,
budget,
resource_cache: ResourceCache::default(),
telemetry: TelemetryCollector::default(),
})
}
pub fn execute<J: GpuJob>(&mut self, job: &J) -> Result<JobReport<J::Output>, ExecutionError> {
let total_started = Instant::now();
let gpu_budget = self
.budget
.gpu_memory_budget(self.config.memory_fill_ratio)
.map_err(ExecutionError::Gpu)?;
let plan = plan_batches(job.items(), gpu_budget).map_err(ExecutionError::Gpu)?;
let mut outputs = Vec::new();
let plan_pack_ms = plan.pack_ms();
let spans = plan.into_spans();
let mut batches = Vec::with_capacity(spans.len());
let mut failed = 0_usize;
for (span_index, span) in spans.into_iter().enumerate() {
let mut report = BatchReport::new(span.clone(), BatchPath::Skipped);
if span_index == 0 {
report.pack_ms = plan_pack_ms;
}
let wants_gpu = self.should_attempt_gpu(&span);
if wants_gpu {
let gpu_started = Instant::now();
match self.execute_gpu_batch(job, &span) {
Ok(batch_outputs) => {
report.path = BatchPath::Gpu;
report.execute_ms = elapsed_ms(gpu_started);
outputs.extend(batch_outputs);
}
Err(error) if matches!(self.config.execution_mode, ExecutionMode::GpuOnly) => {
return Err(ExecutionError::Gpu(error));
}
Err(error) => match self.execute_cpu_batch(job, &span) {
Ok((batch_outputs, post_ms)) => {
report.path = BatchPath::CpuFallback;
report.execute_ms = elapsed_ms(gpu_started);
report.post_ms = post_ms;
report.error = Some(error.to_string());
outputs.extend(batch_outputs);
}
Err(JobError::FallbackNotImplemented) => {
failed = failed.saturating_add(1);
report.path = BatchPath::Skipped;
report.error = Some(error.to_string());
}
Err(error) => return Err(ExecutionError::CpuFallback(error)),
},
}
} else {
match self.execute_cpu_batch(job, &span) {
Ok((batch_outputs, post_ms)) => {
report.path = BatchPath::CpuFallback;
report.post_ms = post_ms;
outputs.extend(batch_outputs);
}
Err(error) => return Err(ExecutionError::CpuFallback(error)),
}
}
batches.push(report);
}
if !batches.is_empty() && failed == batches.len() {
return Err(ExecutionError::AllBatchesFailed {
attempted: batches.len(),
failed,
});
}
let execution_path = execution_path_for(&batches, self.config.execution_mode);
let total_ms = elapsed_ms(total_started);
self.record_telemetry(job.label(), execution_path, &batches);
Ok(JobReport {
outputs,
execution_path,
batches,
total_ms,
})
}
#[cfg(feature = "async")]
pub async fn execute_async<J: GpuJob>(
&mut self,
job: &J,
) -> Result<JobReport<J::Output>, ExecutionError> {
self.execute(job)
}
pub fn reconfigure(&mut self, config: ExecutorConfig) {
if let Ok(budget) = MemoryBudget::from_config(&config) {
self.budget = budget;
self.config = config;
}
}
pub fn telemetry(&self) -> &TelemetryCollector {
&self.telemetry
}
pub fn clear_cache(&mut self) {
self.resource_cache.clear();
}
pub fn config(&self) -> &ExecutorConfig {
&self.config
}
pub fn budget(&self) -> MemoryBudget {
self.budget
}
pub fn gpu_memory_budget(&self) -> Result<GpuMemoryBudget, GpuError> {
self.budget.gpu_memory_budget(self.config.memory_fill_ratio)
}
fn should_attempt_gpu(&self, span: &BatchSpan) -> bool {
match self.config.execution_mode {
ExecutionMode::Auto => {
self.gpu_available() && span.estimated_bytes >= self.config.min_batch_bytes
}
ExecutionMode::PreferGpu | ExecutionMode::GpuOnly => self.gpu_available(),
ExecutionMode::PreferCpu => false,
}
}
fn gpu_available(&self) -> bool {
#[cfg(feature = "gpu")]
{
self.device.is_some() && self.queue.is_some()
}
#[cfg(not(feature = "gpu"))]
{
false
}
}
fn execute_gpu_batch<J: GpuJob>(
&mut self,
job: &J,
span: &BatchSpan,
) -> Result<Vec<J::Output>, GpuError> {
let encoded = job
.encode_batch(span)
.map_err(|error| GpuError::ShaderCompilation {
message: error.to_string(),
})?;
self.execute_encoded_gpu_batch(job, span, encoded)
}
#[cfg(feature = "gpu")]
fn execute_encoded_gpu_batch<J: GpuJob>(
&mut self,
job: &J,
span: &BatchSpan,
encoded: EncodedBatch,
) -> Result<Vec<J::Output>, GpuError> {
let device = self
.device
.as_ref()
.ok_or(GpuError::GpuFailure(GpuFailure::BackendUnavailable))?;
let queue = self
.queue
.as_ref()
.ok_or(GpuError::GpuFailure(GpuFailure::BackendUnavailable))?;
let shader = encoded
.shader
.as_ref()
.ok_or_else(|| GpuError::ShaderCompilation {
message: "encoded batch did not provide a compute shader".to_string(),
})?;
if encoded.buffers.is_empty() {
return Err(GpuError::InvalidBatchPlan(
"encoded GPU batch requires at least one buffer",
));
}
let output_index = encoded
.buffers
.iter()
.position(|buffer| buffer.readback)
.ok_or(GpuError::InvalidBatchPlan(
"encoded GPU batch requires one readback buffer",
))?;
let output_bytes_len = encoded.buffers[output_index].bytes.len().max(4);
let shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some(&shader.label),
source: wgpu::ShaderSource::Wgsl(shader.source.clone().into()),
});
let mut layout_entries = Vec::with_capacity(encoded.buffers.len());
for buffer in &encoded.buffers {
layout_entries.push(wgpu::BindGroupLayoutEntry {
binding: buffer.binding,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage {
read_only: buffer.read_only,
},
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
});
}
let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some(&format!(
"{}-bind-group-layout",
encoded.label.as_deref().unwrap_or(job.label())
)),
entries: &layout_entries,
});
let mut gpu_buffers = Vec::with_capacity(encoded.buffers.len());
for (index, buffer) in encoded.buffers.iter().enumerate() {
let mut usage = wgpu::BufferUsages::STORAGE;
if !buffer.bytes.is_empty() {
usage |= wgpu::BufferUsages::COPY_DST;
}
if buffer.readback || matches!(buffer.role, BufferRole::Output) {
usage |= wgpu::BufferUsages::COPY_SRC;
}
let size = u64::try_from(buffer.bytes.len().max(4)).unwrap_or(u64::MAX);
let gpu_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: buffer.label.as_deref().or(Some(match buffer.role {
BufferRole::Input => "cgpu-input",
BufferRole::Output => "cgpu-output",
BufferRole::Scratch => "cgpu-scratch",
BufferRole::Metadata => "cgpu-metadata",
})),
size,
usage,
mapped_at_creation: !buffer.bytes.is_empty(),
});
if !buffer.bytes.is_empty() {
let mut mapped = gpu_buffer.slice(..).get_mapped_range_mut();
mapped
.slice(0..buffer.bytes.len())
.copy_from_slice(&buffer.bytes);
drop(mapped);
gpu_buffer.unmap();
}
self.resource_cache.insert(
format!("{}:buffer:{index}", job.label()),
crate::resources::ResourceEntry {
label: buffer
.label
.clone()
.unwrap_or_else(|| format!("buffer-{index}")),
bytes: buffer.bytes.len(),
},
);
gpu_buffers.push(gpu_buffer);
}
let mut bind_entries = Vec::with_capacity(encoded.buffers.len());
for (buffer, gpu_buffer) in encoded.buffers.iter().zip(gpu_buffers.iter()) {
bind_entries.push(wgpu::BindGroupEntry {
binding: buffer.binding,
resource: gpu_buffer.as_entire_binding(),
});
}
let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some(&format!(
"{}-bind-group",
encoded.label.as_deref().unwrap_or(job.label())
)),
layout: &bind_group_layout,
entries: &bind_entries,
});
let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some(&format!(
"{}-pipeline-layout",
encoded.label.as_deref().unwrap_or(job.label())
)),
bind_group_layouts: &[Some(&bind_group_layout)],
immediate_size: 0,
});
let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some(&format!(
"{}-pipeline",
encoded.label.as_deref().unwrap_or(job.label())
)),
layout: Some(&pipeline_layout),
module: &shader_module,
entry_point: Some(&shader.entry_point),
compilation_options: wgpu::PipelineCompilationOptions::default(),
cache: None,
});
let readback = device.create_buffer(&wgpu::BufferDescriptor {
label: Some(&format!(
"{}-readback",
encoded.label.as_deref().unwrap_or(job.label())
)),
size: u64::try_from(output_bytes_len).unwrap_or(u64::MAX),
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
label: Some(&format!(
"{}-encoder",
encoded.label.as_deref().unwrap_or(job.label())
)),
});
{
let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some(&format!(
"{}-pass",
encoded.label.as_deref().unwrap_or(job.label())
)),
timestamp_writes: None,
});
pass.set_pipeline(&pipeline);
pass.set_bind_group(0, &bind_group, &[]);
pass.dispatch_workgroups(encoded.dispatch.x, encoded.dispatch.y, encoded.dispatch.z);
}
encoder.copy_buffer_to_buffer(
&gpu_buffers[output_index],
0,
&readback,
0,
u64::try_from(output_bytes_len).unwrap_or(u64::MAX),
);
queue.submit(Some(encoder.finish()));
let slice = readback.slice(..);
let (sender, receiver) = mpsc::channel();
slice.map_async(wgpu::MapMode::Read, move |result| {
let _ = sender.send(result);
});
device
.poll(wgpu::PollType::wait_indefinitely())
.map_err(|_| GpuError::GpuFailure(GpuFailure::Readback))?;
receiver
.recv()
.map_err(|_| GpuError::GpuFailure(GpuFailure::Readback))?
.map_err(|_| GpuError::GpuFailure(GpuFailure::Readback))?;
let mapped = slice.get_mapped_range();
let bytes = mapped.to_vec();
drop(mapped);
readback.unmap();
job.decode_batch(&bytes, span)
.map_err(|error| GpuError::GpuFailureDetail {
failure: GpuFailure::Readback,
message: error.to_string(),
})
}
#[cfg(not(feature = "gpu"))]
fn execute_encoded_gpu_batch<J: GpuJob>(
&mut self,
_job: &J,
_span: &BatchSpan,
_encoded: EncodedBatch,
) -> Result<Vec<J::Output>, GpuError> {
Err(GpuError::GpuFailure(GpuFailure::BackendUnavailable))
}
fn execute_cpu_batch<J: GpuJob>(
&self,
job: &J,
span: &BatchSpan,
) -> Result<(Vec<J::Output>, f64), JobError> {
let started = Instant::now();
let output = job.cpu_fallback(span)?;
Ok((output, elapsed_ms(started)))
}
fn record_telemetry(
&mut self,
label: &'static str,
execution_path: ExecutionPath,
batches: &[BatchReport],
) {
if !self.config.enable_telemetry {
return;
}
let mut report = TelemetryReport::new(label);
report.execution_mode = Some(self.config.execution_mode);
report.batch_count = batches.len();
report.bytes_in = batches.iter().map(|batch| batch.bytes_in).sum();
report.bytes_out = batches.iter().map(|batch| batch.bytes_out).sum();
report.plan_time = duration_from_ms(batches.iter().map(|batch| batch.pack_ms).sum());
report.gpu_time = duration_from_ms(
batches
.iter()
.filter(|batch| batch.path == BatchPath::Gpu)
.map(|batch| batch.execute_ms)
.sum(),
);
report.readback_time =
duration_from_ms(batches.iter().map(|batch| batch.readback_ms).sum());
report.cpu_fallback_time = duration_from_ms(
batches
.iter()
.filter(|batch| batch.path == BatchPath::CpuFallback)
.map(|batch| batch.post_ms)
.sum(),
);
if matches!(
execution_path,
ExecutionPath::CpuOnly | ExecutionPath::GpuWithFallback | ExecutionPath::Mixed
) {
report.decode_time = duration_from_ms(batches.iter().map(|batch| batch.post_ms).sum());
}
self.telemetry.record(report, &self.config.telemetry_sink);
}
}
#[cfg(feature = "gpu")]
fn maybe_init_gpu(
config: &ExecutorConfig,
) -> Result<(Option<Arc<wgpu::Device>>, Option<Arc<wgpu::Queue>>), ExecutorInitError> {
if matches!(config.execution_mode, ExecutionMode::PreferCpu) {
return Ok((None, None));
}
match pollster::block_on(init_gpu_device()) {
Ok((device, queue)) => Ok((Some(Arc::new(device)), Some(Arc::new(queue)))),
Err(error) if matches!(config.execution_mode, ExecutionMode::GpuOnly) => Err(error),
Err(_) => Ok((None, None)),
}
}
#[cfg(feature = "gpu")]
async fn init_gpu_device() -> Result<(wgpu::Device, wgpu::Queue), ExecutorInitError> {
let mut descriptor = wgpu::InstanceDescriptor::new_without_display_handle();
descriptor.backends = wgpu::Backends::PRIMARY;
let instance = wgpu::Instance::new(descriptor);
let adapter = instance
.request_adapter(&wgpu::RequestAdapterOptions {
power_preference: wgpu::PowerPreference::LowPower,
compatible_surface: None,
force_fallback_adapter: false,
})
.await
.map_err(|_| ExecutorInitError::NoGpuAvailable)?;
adapter
.request_device(&wgpu::DeviceDescriptor {
label: Some("cgpu-executor-device"),
required_features: wgpu::Features::empty(),
required_limits: wgpu::Limits::downlevel_defaults(),
..Default::default()
})
.await
.map_err(ExecutorInitError::DeviceRequestFailed)
}
fn execution_path_for(batches: &[BatchReport], mode: ExecutionMode) -> ExecutionPath {
if batches.is_empty() {
return ExecutionPath::CpuOnly;
}
let gpu_count = batches
.iter()
.filter(|batch| batch.path == BatchPath::Gpu)
.count();
let fallback_count = batches
.iter()
.filter(|batch| batch.path == BatchPath::CpuFallback)
.count();
let recovered_gpu_failure = batches
.iter()
.any(|batch| batch.path == BatchPath::CpuFallback && batch.error.is_some());
if gpu_count == batches.len() {
ExecutionPath::GpuOnly
} else if recovered_gpu_failure {
ExecutionPath::GpuWithFallback
} else if fallback_count == batches.len() || matches!(mode, ExecutionMode::PreferCpu) {
ExecutionPath::CpuOnly
} else {
ExecutionPath::Mixed
}
}
fn duration_from_ms(ms: f64) -> std::time::Duration {
std::time::Duration::from_secs_f64((ms / 1000.0).max(0.0))
}