Skip to main content

08_cnn_convolution/
08_cnn_convolution.rs

1use apple_metal::MetalDevice;
2use apple_mps::{
3    cnn_accumulator_precision_option, cnn_convolution_flags, feature_channel_format,
4    CnnConvolution, CnnConvolutionDescriptor, Image, ImageDescriptor,
5};
6
7fn main() {
8    let device = MetalDevice::system_default().expect("no Metal device available");
9    let queue = device.new_command_queue().expect("command queue");
10
11    let descriptor = CnnConvolutionDescriptor::new(1, 1, 1, 1).expect("descriptor");
12    let convolution = CnnConvolution::new(
13        &device,
14        &descriptor,
15        &[2.0],
16        Some(&[0.5]),
17        cnn_convolution_flags::NONE,
18    )
19    .expect("convolution");
20    convolution.set_accumulator_precision_option(cnn_accumulator_precision_option::FLOAT);
21
22    let image_descriptor = ImageDescriptor::new(2, 2, 1, feature_channel_format::FLOAT32);
23    let source = Image::new(&device, image_descriptor).expect("source image");
24    let destination = Image::new(&device, image_descriptor).expect("destination image");
25    source.write_f32(&[1.0, 2.0, 3.0, 4.0]).expect("write source");
26
27    let command_buffer = queue.new_command_buffer().expect("command buffer");
28    convolution.encode_image(&command_buffer, &source, &destination);
29    command_buffer.commit();
30    command_buffer.wait_until_completed();
31
32    let output = destination.read_f32().expect("output");
33    println!("{output:?}");
34}