use sim_kernel::Symbol;
use sim_lib_numbers_tensor::{TensorExecError, TensorRequest};
use crate::{
WgpuKernelDType, WgpuTensorExecutor,
dispatch::{
WgpuDispatchBuffer, buffer_size, check_storage_buffer_limit, compiled_pipeline, f32_bytes,
pipeline_symbol, read_f32s, readback_buffer, tensor_f32_values, u32_count,
},
kernel_reductions::fixed_tree_sum_values,
kernel_support::{invalid, round, shape_error},
};
pub(crate) fn execute_reduction_dispatch(
executor: &WgpuTensorExecutor,
cx: &mut sim_kernel::Cx,
request: &TensorRequest,
op: crate::WgpuKernelOp,
dtype: WgpuKernelDType,
) -> std::result::Result<WgpuDispatchBuffer, TensorExecError> {
let Some(context) = &executor.context else {
return Err(invalid("wgpu device context is unavailable"));
};
let [tensor] = request.inputs.as_ref() else {
return Err(invalid(
"wgpu reduction dispatch expects exactly one tensor input",
));
};
if !request.output.shape().is_empty() {
return Err(shape_error("wgpu reduction output must be scalar"));
}
let len = tensor.len();
if len == 0 {
if matches!(op, crate::WgpuKernelOp::Min | crate::WgpuKernelOp::Max) {
return Err(invalid("wgpu min/max reductions require at least one cell"));
}
let symbol = pipeline_symbol(executor, context, op, dtype, 0);
return scalar_resident_buffer(executor, context, 0.0, symbol);
}
let input_values = tensor_f32_values(cx, tensor, dtype)?;
let partial_count = len.div_ceil(512);
check_storage_buffer_limit(executor, input_values.len(), "wgpu reduction input")?;
check_storage_buffer_limit(executor, partial_count, "wgpu reduction partial output")?;
let output_size = buffer_size(partial_count)?;
let input = upload_buffer(
executor,
context,
"sim-compute-wgpu-reduction-input",
wgpu::BufferUsages::STORAGE,
&f32_bytes(&input_values),
);
let partials = context.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("sim-compute-wgpu-reduction-partials"),
size: output_size,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
});
let readback = readback_buffer(&context.device, output_size, "reduction");
let params = upload_buffer(
executor,
context,
"sim-compute-wgpu-reduction-params",
wgpu::BufferUsages::UNIFORM,
&reduction_params_bytes(op, len)?,
);
let pipeline = compiled_pipeline(executor, context, op, dtype, 0);
let layout = pipeline.pipeline.get_bind_group_layout(0);
let bind_group = context
.device
.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("sim-compute-wgpu-reduction-bind-group"),
layout: &layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: input.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: partials.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 2,
resource: params.as_entire_binding(),
},
],
});
let mut encoder = context
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor {
label: Some("sim-compute-wgpu-reduction-encoder"),
});
{
let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some("sim-compute-wgpu-reduction-pass"),
timestamp_writes: None,
});
pass.set_pipeline(&pipeline.pipeline);
pass.set_bind_group(0, &bind_group, &[]);
pass.dispatch_workgroups(
u32_count(partial_count, "wgpu reduction partial count")?,
1,
1,
);
}
encoder.copy_buffer_to_buffer(&partials, 0, &readback, 0, output_size);
context.queue.submit([encoder.finish()]);
executor.physical_counters().record_scalar_sync();
let partials = read_f32s(context, &readback, partial_count)?;
let value = match op {
crate::WgpuKernelOp::Sum => fixed_tree_sum_values(partials, dtype),
crate::WgpuKernelOp::Norm => fixed_tree_sum_values(partials, dtype).sqrt(),
crate::WgpuKernelOp::Min => fixed_tree_min_max(partials, false)?,
crate::WgpuKernelOp::Max => fixed_tree_min_max(partials, true)?,
_ => {
return Err(invalid(
"wgpu reduction dispatch received a non-reduction op",
));
}
};
scalar_resident_buffer(
executor,
context,
round(dtype, value),
pipeline.record.symbol,
)
}
fn upload_buffer(
executor: &WgpuTensorExecutor,
context: &crate::site::WgpuExecutionContext,
label: &'static str,
usage: wgpu::BufferUsages,
bytes: &[u8],
) -> wgpu::Buffer {
let buffer = context.device.create_buffer(&wgpu::BufferDescriptor {
label: Some(label),
size: bytes.len().max(4) as u64,
usage: usage | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
context.queue.write_buffer(&buffer, 0, bytes);
executor
.physical_counters()
.record_upload(bytes.len() as u64);
buffer
}
fn scalar_resident_buffer(
executor: &WgpuTensorExecutor,
context: &crate::site::WgpuExecutionContext,
value: f32,
pipeline: Symbol,
) -> std::result::Result<WgpuDispatchBuffer, TensorExecError> {
let buffer = context.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("sim-compute-wgpu-reduction-scalar"),
size: 4,
usage: wgpu::BufferUsages::STORAGE
| wgpu::BufferUsages::COPY_DST
| wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
});
context.queue.write_buffer(&buffer, 0, &value.to_ne_bytes());
executor.physical_counters().record_upload(4);
Ok(WgpuDispatchBuffer {
buffer: std::sync::Arc::new(buffer),
pipeline,
len: 1,
})
}
fn fixed_tree_min_max(
mut values: Vec<f32>,
max: bool,
) -> std::result::Result<f32, TensorExecError> {
while values.len() > 1 {
let mut next = Vec::with_capacity(values.len().div_ceil(2));
for pair in values.chunks(2) {
next.push(if let [left, right] = pair {
if max {
left.max(*right)
} else {
left.min(*right)
}
} else {
pair[0]
});
}
values = next;
}
values
.into_iter()
.next()
.ok_or_else(|| invalid("wgpu min/max reductions require at least one cell"))
}
fn reduction_params_bytes(
op: crate::WgpuKernelOp,
len: usize,
) -> std::result::Result<[u8; 8], TensorExecError> {
let op = match op {
crate::WgpuKernelOp::Sum => 0_u32,
crate::WgpuKernelOp::Min => 1,
crate::WgpuKernelOp::Max => 2,
crate::WgpuKernelOp::Norm => 3,
_ => return Err(invalid("wgpu reduction params received a non-reduction op")),
};
let len = u32_count(len, "wgpu reduction length")?;
let mut bytes = [0_u8; 8];
bytes[..4].copy_from_slice(&op.to_ne_bytes());
bytes[4..].copy_from_slice(&len.to_ne_bytes());
Ok(bytes)
}