07_optimizer_and_state/
07_optimizer_and_state.rs

1use apple_metal::{resource_options, MetalBuffer, MetalDevice};
2use apple_mps::{
3    data_type, nn_regularization_type, state_batch_increment_read_count, state_resource_type,
4    NNOptimizerStochasticGradientDescent, State, StateResourceList, Vector, VectorDescriptor,
5};
6
7fn as_bytes<T>(values: &[T]) -> &[u8] {
8    unsafe {
9        core::slice::from_raw_parts(values.as_ptr().cast::<u8>(), core::mem::size_of_val(values))
10    }
11}
12
13fn buffer_with_f32_values(device: &MetalDevice, values: &[f32]) -> MetalBuffer {
14    let buffer = device
15        .new_buffer(
16            core::mem::size_of_val(values),
17            resource_options::STORAGE_MODE_SHARED,
18        )
19        .expect("buffer");
20    let _ = buffer.write_bytes(as_bytes(values));
21    buffer
22}
23
24fn read_f32_values(buffer: &MetalBuffer, len: usize) -> Vec<f32> {
25    let ptr = buffer.contents().expect("buffer contents").cast::<f32>();
26    unsafe { core::slice::from_raw_parts(ptr, len).to_vec() }
27}
28
29fn vector_with_values(device: &MetalDevice, values: &[f32]) -> (MetalBuffer, Vector) {
30    let buffer = buffer_with_f32_values(device, values);
31    let descriptor = VectorDescriptor::contiguous(values.len(), data_type::FLOAT32).expect("vector desc");
32    let vector = Vector::new_with_buffer(&buffer, descriptor).expect("vector");
33    (buffer, vector)
34}
35
36fn main() {
37    let device = MetalDevice::system_default().expect("no Metal device available");
38    let queue = device.new_command_queue().expect("command queue");
39
40    let command_buffer = queue.new_command_buffer().expect("command buffer");
41    let temporary_a = State::temporary_with_buffer_size(&command_buffer, 32).expect("temporary state a");
42    let temporary_b = State::temporary_with_buffer_size(&command_buffer, 64).expect("temporary state b");
43    let unique_count = state_batch_increment_read_count(&[&temporary_a, &temporary_a, &temporary_b], 1);
44    assert_eq!(unique_count, 2);
45    assert_eq!(temporary_a.resource_type_at_index(0), state_resource_type::BUFFER);
46    command_buffer.commit();
47    command_buffer.wait_until_completed();
48
49    let resource_list = StateResourceList::new().expect("resource list");
50    resource_list.append_buffer(16);
51    let persistent_state = State::new_with_resource_list(&device, &resource_list).expect("persistent state");
52    assert_eq!(persistent_state.resource_count(), 1);
53    assert_eq!(persistent_state.buffer_size_at_index(0), 16);
54
55    let (gradient_buffer, gradient_vector) = vector_with_values(&device, &[0.1, -0.2]);
56    let (values_buffer, values_vector) = vector_with_values(&device, &[1.0, -1.0]);
57    let (result_buffer, result_vector) = vector_with_values(&device, &[0.0, 0.0]);
58    let optimizer = NNOptimizerStochasticGradientDescent::new(&device, 0.5).expect("sgd");
59    let base = optimizer.as_optimizer().expect("optimizer base");
60    assert_eq!(base.regularization_type(), nn_regularization_type::NONE);
61
62    let command_buffer = queue.new_command_buffer().expect("command buffer");
63    optimizer.encode_vector(
64        &command_buffer,
65        &gradient_vector,
66        &values_vector,
67        None,
68        &result_vector,
69    );
70    command_buffer.commit();
71    command_buffer.wait_until_completed();
72
73    let _ = gradient_buffer;
74    let _ = values_buffer;
75    let output = read_f32_values(&result_buffer, 2);
76    println!("{output:?}");
77}
07_optimizer_and_state/07_optimizer_and_state.rs

07_optimizer_and_state/
07_optimizer_and_state.rs
