02_compile_matmul/
02_compile_matmul.rs

1use apple_metal::MetalDevice;
2use apple_mpsgraph::{data_type, FeedDescription, Graph, TensorData};
3
4fn main() {
5    let device = MetalDevice::system_default().expect("no Metal device available");
6    let queue = device
7        .new_command_queue()
8        .expect("failed to create command queue");
9    let graph = Graph::new().expect("failed to create MPSGraph");
10
11    let left = graph
12        .placeholder(Some(&[2, 3]), data_type::FLOAT32, Some("left"))
13        .expect("failed to create left placeholder");
14    let right = graph
15        .placeholder(Some(&[3, 2]), data_type::FLOAT32, Some("right"))
16        .expect("failed to create right placeholder");
17    let output = graph
18        .matrix_multiplication(&left, &right, Some("matmul"))
19        .expect("failed to create matrix multiplication op");
20
21    let executable = graph
22        .compile(
23            &device,
24            &[
25                FeedDescription::new(&left, &[2, 3], data_type::FLOAT32),
26                FeedDescription::new(&right, &[3, 2], data_type::FLOAT32),
27            ],
28            &[&output],
29        )
30        .expect("failed to compile executable");
31
32    let left_data = TensorData::from_f32_slice(&device, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3])
33        .expect("failed to create left tensor data");
34    let right_data =
35        TensorData::from_f32_slice(&device, &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0], &[3, 2])
36            .expect("failed to create right tensor data");
37
38    let results = executable
39        .run(&queue, &[&left_data, &right_data])
40        .expect("failed to run executable");
41    let values = results[0].read_f32().expect("failed to read tensor output");
42    let expected = [58.0_f32, 64.0, 139.0, 154.0];
43    for (actual, expected_value) in values.iter().zip(expected) {
44        assert!(
45            (actual - expected_value).abs() < 1.0e-4,
46            "unexpected matrix multiply result: {values:?}"
47        );
48    }
49
50    println!("compile+matmul smoke passed: {values:?}");
51}
02_compile_matmul/02_compile_matmul.rs

02_compile_matmul/
02_compile_matmul.rs
