use crate::{
buffer::MlxBuffer,
device::MlxDevice,
dtypes::DType,
encoder::CommandEncoder,
error::{MlxError, Result},
kernel_registry::KernelRegistry,
};
use metal::MTLSize;
use super::encode_helpers::{as_bytes, encode_with_args, KernelArg};
pub const DFLASH2_DYNAMIC_CONV_KERNEL: &str = "dflash2_dynamic_conv_f32";
pub static DFLASH2_DYNAMIC_CONV_SOURCE: &str =
include_str!("../shaders/dflash2_dynamic_conv.metal");
pub fn register(registry: &mut KernelRegistry) {
registry.register_source(DFLASH2_DYNAMIC_CONV_KERNEL, DFLASH2_DYNAMIC_CONV_SOURCE);
}
#[repr(u32)]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum DFlash2ConvolutionStage {
Prepare = 0,
Residual = 1,
}
#[repr(C)]
#[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct DFlash2DynamicConvParams {
pub lanes: u32,
pub rows_per_lane: u32,
pub hidden: u32,
pub groups: u32,
pub stage: u32,
}
fn product(dims: &[usize]) -> Result<usize> {
dims.iter().try_fold(1usize, |count, dim| {
count
.checked_mul(*dim)
.ok_or_else(|| MlxError::InvalidArgument("DFlash2 convolution shape overflow".into()))
})
}
fn physical_range(buffer: &MlxBuffer) -> Result<(usize, usize)> {
let start = (buffer.contents_ptr() as usize)
.checked_add(
usize::try_from(buffer.byte_offset())
.map_err(|_| MlxError::InvalidArgument("DFlash2 buffer offset overflow".into()))?,
)
.ok_or_else(|| MlxError::InvalidArgument("DFlash2 buffer start overflow".into()))?;
let end = start
.checked_add(buffer.data_byte_len())
.ok_or_else(|| MlxError::InvalidArgument("DFlash2 buffer end overflow".into()))?;
Ok((start, end))
}
fn require(buffer: &MlxBuffer, dtype: DType, shape: &[usize], label: &str) -> Result<()> {
if buffer.dtype() != dtype || buffer.shape() != shape {
return Err(MlxError::InvalidArgument(format!(
"DFlash2 convolution {label} must be {dtype} {shape:?}, got {} {:?}",
buffer.dtype(),
buffer.shape()
)));
}
let bytes = product(shape)?
.checked_mul(dtype.size_of())
.ok_or_else(|| MlxError::InvalidArgument("DFlash2 convolution byte overflow".into()))?;
if buffer.data_byte_len() != bytes {
return Err(MlxError::InvalidArgument(format!(
"DFlash2 convolution {label} has {} data bytes, expected {bytes}",
buffer.data_byte_len()
)));
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
pub fn dispatch_dflash2_dynamic_conv_f32(
encoder: &mut CommandEncoder,
registry: &mut KernelRegistry,
device: &MlxDevice,
input: &MlxBuffer,
dynamic_weights: &MlxBuffer,
base_kernel: &MlxBuffer,
residual: Option<&MlxBuffer>,
output: &MlxBuffer,
mut params: DFlash2DynamicConvParams,
stage: DFlash2ConvolutionStage,
) -> Result<()> {
params.stage = stage as u32;
if params.lanes == 0
|| params.rows_per_lane < 2
|| params.rows_per_lane > 16
|| params.hidden == 0
|| params.groups == 0
|| params.hidden % params.groups != 0
{
return Err(MlxError::InvalidArgument(format!(
"invalid DFlash2 convolution geometry: lanes={}, rows={}, hidden={}, groups={}",
params.lanes, params.rows_per_lane, params.hidden, params.groups
)));
}
let rows = (params.lanes as usize)
.checked_mul(params.rows_per_lane as usize)
.ok_or_else(|| MlxError::InvalidArgument("DFlash2 convolution row overflow".into()))?;
let activations = [rows, params.hidden as usize];
let dynamic = [rows, 4 * params.groups as usize];
let base = [2, 2, params.hidden as usize];
require(input, DType::F32, &activations, "input")?;
require(dynamic_weights, DType::F32, &dynamic, "dynamic weights")?;
require(base_kernel, DType::BF16, &base, "base kernel")?;
require(output, DType::F32, &activations, "output")?;
let residual = match stage {
DFlash2ConvolutionStage::Prepare => input,
DFlash2ConvolutionStage::Residual => residual.ok_or_else(|| {
MlxError::InvalidArgument(
"DFlash2 residual convolution requires a residual buffer".into(),
)
})?,
};
if stage == DFlash2ConvolutionStage::Residual {
require(residual, DType::F32, &activations, "residual")?;
}
if output.contents_ptr() == input.contents_ptr()
|| output.contents_ptr() == dynamic_weights.contents_ptr()
|| output.contents_ptr() == base_kernel.contents_ptr()
|| (stage == DFlash2ConvolutionStage::Residual
&& output.contents_ptr() == residual.contents_ptr())
{
return Err(MlxError::InvalidArgument(
"DFlash2 convolution output must not alias an input".into(),
));
}
if encoder.is_capturing() {
let reads = [input, dynamic_weights, base_kernel, residual]
.iter()
.map(|buffer| physical_range(buffer))
.collect::<Result<Vec<_>>>()?;
encoder.set_pending_buffer_ranges(reads, vec![physical_range(output)?]);
}
let pipeline = registry.get_pipeline(DFLASH2_DYNAMIC_CONV_KERNEL, device.metal_device())?;
let elements = product(&activations)?;
encode_with_args(
encoder,
pipeline,
&[
(0, KernelArg::Buffer(input)),
(1, KernelArg::Buffer(dynamic_weights)),
(2, KernelArg::Buffer(base_kernel)),
(3, KernelArg::Buffer(residual)),
(4, KernelArg::Buffer(output)),
(5, KernelArg::Bytes(as_bytes(¶ms))),
],
MTLSize::new(elements as u64, 1, 1),
MTLSize::new(256.min(elements) as u64, 1, 1),
);
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
fn reference(
input: &[f32],
dynamic: &[f32],
base: &[f32],
residual: &[f32],
params: DFlash2DynamicConvParams,
stage: DFlash2ConvolutionStage,
) -> Vec<f32> {
let hidden = params.hidden as usize;
let groups = params.groups as usize;
let rows_per_lane = params.rows_per_lane as usize;
let rows = params.lanes as usize * rows_per_lane;
let group_size = hidden / groups;
let stage_tap = stage as usize * 2;
let dynamic_width = 4 * groups;
let mut output = vec![0.0; rows * hidden];
for row in 0..rows {
for channel in 0..hidden {
let element = row * hidden + channel;
let group = channel / group_size;
let mut value = input[element]
* (base[stage_tap * hidden + channel]
+ dynamic[row * dynamic_width + stage_tap * groups + group]);
if row % rows_per_lane > 0 {
value += input[element - hidden]
* (base[(stage_tap + 1) * hidden + channel]
+ dynamic[row * dynamic_width + (stage_tap + 1) * groups + group]);
}
if stage == DFlash2ConvolutionStage::Residual {
value += residual[element];
}
output[element] = value;
}
}
output
}
#[test]
fn parameter_abi_is_five_u32_words() {
assert_eq!(std::mem::size_of::<DFlash2DynamicConvParams>(), 20);
}
#[test]
fn gpu_convolution_matches_lane_local_reference() {
let device = MlxDevice::new().unwrap();
let mut registry = KernelRegistry::new();
register(&mut registry);
let params = DFlash2DynamicConvParams {
lanes: 2,
rows_per_lane: 2,
hidden: 2,
groups: 1,
stage: 0,
};
let input_values = [1.0, 2.0, 3.0, 4.0, 10.0, 20.0, 30.0, 40.0];
let mut input = device.alloc_buffer(32, DType::F32, vec![4, 2]).unwrap();
input
.as_mut_slice::<f32>()
.unwrap()
.copy_from_slice(&input_values);
let dynamic = device
.alloc_buffer(4 * 4 * 4, DType::F32, vec![4, 4])
.unwrap();
let mut base = device.alloc_buffer(16, DType::BF16, vec![2, 2, 2]).unwrap();
let base_values = [1.0f32, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 2.0];
for (output, value) in base
.as_mut_slice::<u16>()
.unwrap()
.iter_mut()
.zip(base_values)
{
*output = (value.to_bits() >> 16) as u16;
}
let output = device.alloc_buffer(32, DType::F32, vec![4, 2]).unwrap();
let mut encoder = device.command_encoder().unwrap();
dispatch_dflash2_dynamic_conv_f32(
&mut encoder,
&mut registry,
&device,
&input,
&dynamic,
&base,
None,
&output,
params,
DFlash2ConvolutionStage::Prepare,
)
.unwrap();
encoder.commit_and_wait().unwrap();
let expected = reference(
&input_values,
&[0.0; 16],
&base_values,
&[0.0; 8],
params,
DFlash2ConvolutionStage::Prepare,
);
assert_eq!(output.as_slice::<f32>().unwrap(), expected);
}
#[test]
fn reference_does_not_cross_lane_boundaries() {
let params = DFlash2DynamicConvParams {
lanes: 2,
rows_per_lane: 2,
hidden: 2,
groups: 1,
stage: 0,
};
let input = [1.0, 2.0, 3.0, 4.0, 10.0, 20.0, 30.0, 40.0];
let dynamic = vec![0.0; 4 * 4];
let base = [1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 2.0];
let output = reference(
&input,
&dynamic,
&base,
&[0.0; 8],
params,
DFlash2ConvolutionStage::Prepare,
);
assert_eq!(output, [1.0, 2.0, 4.0, 6.0, 10.0, 20.0, 40.0, 60.0]);
}
#[test]
fn residual_stage_adds_declared_residual() {
let params = DFlash2DynamicConvParams {
lanes: 1,
rows_per_lane: 2,
hidden: 2,
groups: 1,
stage: 0,
};
let output = reference(
&[1.0, 2.0, 3.0, 4.0],
&[0.0; 8],
&[0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0],
&[5.0, 6.0, 7.0, 8.0],
params,
DFlash2ConvolutionStage::Residual,
);
assert_eq!(output, [6.0, 8.0, 10.0, 12.0]);
}
}