feature: Metal timestamps

Cargo.toml

@@ -48,6 +48,9 @@ name = "window"
 [[example]]
 name = "headless-render"
 
+[[example]]
+name = "counters"
+
 [[example]]
 name = "library"
 

examples/counters/main.rs

@@ -0,0 +1,126 @@
+use metal::*;
+fn main() {
+    let device = Device::system_default().expect("No device found");
+
+    let counter_sample_buffer = create_counter_sample_buffer(&device);
+
+    //Apple silicon uses at stage boundary
+    let counter_sampling_point = MTLCounterSamplingPoint::AtStageBoundary;
+    assert!(device.supports_counter_sampling(counter_sampling_point));
+
+    let command_queue = device.new_command_queue();
+
+    let data = [1u32; 64 * 64];
+
+    let buffer = device.new_buffer_with_data(
+        unsafe { std::mem::transmute(data.as_ptr()) },
+        (data.len() * std::mem::size_of::<u32>()) as u64,
+        MTLResourceOptions::CPUCacheModeDefaultCache,
+    );
+
+    let sum = {
+        let data = [0u32];
+        device.new_buffer_with_data(
+            unsafe { std::mem::transmute(data.as_ptr()) },
+            (data.len() * std::mem::size_of::<u32>()) as u64,
+            MTLResourceOptions::CPUCacheModeDefaultCache,
+        )
+    };
+
+    let command_buffer = command_queue.new_command_buffer();
+
+    let compute_pass_descriptor = ComputePassDescriptor::new();
+    let sample_buffer_attachments = compute_pass_descriptor.sample_buffer_attachments();
+    let sample_buffer_attachment_descriptor = sample_buffer_attachments.object_at(0).unwrap();
+
+    sample_buffer_attachment_descriptor.set_sample_buffer(&counter_sample_buffer);
+    sample_buffer_attachment_descriptor.set_start_of_encoder_sample_index(0);
+    sample_buffer_attachment_descriptor.set_end_of_encoder_sample_index(1);
+
+    let encoder = command_buffer.compute_command_encoder_with_descriptor(&compute_pass_descriptor);
+    let library_path = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .join("examples/compute/shaders.metallib");
+
+    let library = device.new_library_with_file(library_path).unwrap();
+    let kernel = library.get_function("sum", None).unwrap();
+
+    let pipeline_state_descriptor = ComputePipelineDescriptor::new();
+    pipeline_state_descriptor.set_compute_function(Some(&kernel));
+
+    let pipeline_state = device
+        .new_compute_pipeline_state_with_function(
+            pipeline_state_descriptor.compute_function().unwrap(),
+        )
+        .unwrap();
+
+    encoder.set_compute_pipeline_state(&pipeline_state);
+    encoder.set_buffer(0, Some(&buffer), 0);
+    encoder.set_buffer(1, Some(&sum), 0);
+
+    let width = 16;
+
+    let thread_group_count = MTLSize {
+        width,
+        height: 1,
+        depth: 1,
+    };
+
+    let thread_group_size = MTLSize {
+        width: (data.len() as u64 + width) / width,
+        height: 1,
+        depth: 1,
+    };
+
+    encoder.dispatch_thread_groups(thread_group_count, thread_group_size);
+    encoder.end_encoding();
+
+    let blit_encoder = command_buffer.new_blit_command_encoder();
+    let destination_buffer = device.new_buffer(
+        (std::mem::size_of::<u64>() * 2) as u64,
+        MTLResourceOptions::StorageModeShared,
+    );
+    let range = crate::NSRange::new(0_u64, 2_u64);
+    blit_encoder.resolve_counters(&counter_sample_buffer, range, &destination_buffer, 0_u64);
+    blit_encoder.end_encoding();
+
+    command_buffer.commit();
+    command_buffer.wait_until_completed();
+
+    let timestamps =
+        unsafe { std::slice::from_raw_parts(destination_buffer.contents() as *const u64, 2) };
+    println!("Start timestamp: {}", timestamps[0]);
+    println!("End timestamp:   {}", timestamps[1]);
+    println!("Elapsed time:    {}", timestamps[1] - timestamps[0]);
+
+    let ptr = sum.contents() as *mut u32;
+    println!("Compute shader sum: {}", unsafe { *ptr });
+
+    unsafe {
+        assert_eq!(4096, *ptr);
+    }
+}
+
+fn create_counter_sample_buffer(device: &Device) -> CounterSampleBuffer {
+    let counter_sample_buffer_desc = metal::CounterSampleBufferDescriptor::new();
+    counter_sample_buffer_desc.set_storage_mode(metal::MTLStorageMode::Shared);
+    counter_sample_buffer_desc.set_sample_count(2_u64);
+    counter_sample_buffer_desc.set_counter_set(&fetch_timestamp_counter_set(device));
+
+    device
+        .new_counter_sample_buffer_with_descriptor(&counter_sample_buffer_desc)
+        .unwrap()
+}
+
+fn fetch_timestamp_counter_set(device: &Device) -> metal::CounterSet {
+    let counter_sets = device.counter_sets();
+    let mut timestamp_counter = None;
+    for cs in counter_sets.iter() {
+        if cs.name() == "timestamp" {
+            timestamp_counter = Some(cs);
+            break;
+        }
+    }
+    timestamp_counter
+        .expect("No timestamp counter found")
+        .clone()
+}

src/commandbuffer.rs

@@ -140,6 +140,13 @@ impl CommandBufferRef {
         unsafe { msg_send![self, computeCommandEncoderWithDispatchType: ty] }
     }
 
+    pub fn compute_command_encoder_with_descriptor(
+        &self,
+        descriptor: &ComputePassDescriptorRef,
+    ) -> &ComputeCommandEncoderRef {
+        unsafe { msg_send![self, computeCommandEncoderWithDescriptor: descriptor] }
+    }
+
     pub fn encode_signal_event(&self, event: &EventRef, new_value: u64) {
         unsafe {
             msg_send![self,

src/computepass.rs

@@ -0,0 +1,103 @@
+use super::*;
+
+/// See <https://developer.apple.com/documentation/metal/mtlcomputepassdescriptor>
+pub enum MTLComputePassDescriptor {}
+
+foreign_obj_type! {
+    type CType = MTLComputePassDescriptor;
+    pub struct ComputePassDescriptor;
+}
+
+impl ComputePassDescriptor {
+    /// Creates a default compute pass descriptor with no attachments.
+    pub fn new<'a>() -> &'a ComputePassDescriptorRef {
+        unsafe { msg_send![class!(MTLComputePassDescriptor), computePassDescriptor] }
+    }
+}
+
+impl ComputePassDescriptorRef {
+    pub fn sample_buffer_attachments(
+        &self,
+    ) -> &ComputePassSampleBufferAttachmentDescriptorArrayRef {
+        unsafe { msg_send![self, sampleBufferAttachments] }
+    }
+}
+
+/// See <https://developer.apple.com/documentation/metal/mtlcomputepasssamplebufferattachmentdescriptorarray>
+pub enum MTLComputePassSampleBufferAttachmentDescriptorArray {}
+
+foreign_obj_type! {
+    type CType = MTLComputePassSampleBufferAttachmentDescriptorArray;
+    pub struct ComputePassSampleBufferAttachmentDescriptorArray;
+}
+
+impl ComputePassSampleBufferAttachmentDescriptorArrayRef {
+    pub fn object_at(
+        &self,
+        index: NSUInteger,
+    ) -> Option<&ComputePassSampleBufferAttachmentDescriptorRef> {
+        unsafe { msg_send![self, objectAtIndexedSubscript: index] }
+    }
+
+    pub fn set_object_at(
+        &self,
+        index: NSUInteger,
+        attachment: Option<&ComputePassSampleBufferAttachmentDescriptorRef>,
+    ) {
+        unsafe {
+            msg_send![self, setObject:attachment
+                     atIndexedSubscript:index]
+        }
+    }
+}
+
+/// See <https://developer.apple.com/documentation/metal/mtlcomputepasssamplebufferattachmentdescriptor>
+pub enum MTLComputePassSampleBufferAttachmentDescriptor {}
+
+foreign_obj_type! {
+    type CType = MTLComputePassSampleBufferAttachmentDescriptor;
+    pub struct ComputePassSampleBufferAttachmentDescriptor;
+}
+
+impl ComputePassSampleBufferAttachmentDescriptor {
+    pub fn new() -> Self {
+        let class = class!(MTLComputePassSampleBufferAttachmentDescriptor);
+        unsafe { msg_send![class, new] }
+    }
+}
+
+impl ComputePassSampleBufferAttachmentDescriptorRef {
+    pub fn sample_buffer(&self) -> &CounterSampleBufferRef {
+        unsafe { msg_send![self, sampleBuffer] }
+    }
+
+    pub fn set_sample_buffer(&self, sample_buffer: &CounterSampleBufferRef) {
+        unsafe { msg_send![self, setSampleBuffer: sample_buffer] }
+    }
+
+    pub fn start_of_encoder_sample_index(&self) -> u64 {
+        unsafe { msg_send![self, startOfEncoderSampleIndex] }
+    }
+
+    pub fn set_start_of_encoder_sample_index(&self, start_of_encoder_sample_index: u64) {
+        unsafe {
+            msg_send![
+                self,
+                setStartOfEncoderSampleIndex: start_of_encoder_sample_index
+            ]
+        }
+    }
+
+    pub fn end_of_encoder_sample_index(&self) -> u64 {
+        unsafe { msg_send![self, endOfEncoderSampleIndex] }
+    }
+
+    pub fn set_end_of_encoder_sample_index(&self, end_of_encoder_sample_index: u64) {
+        unsafe {
+            msg_send![
+                self,
+                setEndOfEncoderSampleIndex: end_of_encoder_sample_index
+            ]
+        }
+    }
+}

src/counters.rs

@@ -0,0 +1,96 @@
+use crate::MTLStorageMode;
+
+/// See <https://developer.apple.com/documentation/metal/mtlcountersamplebufferdescriptor>
+pub enum MTLCounterSampleBufferDescriptor {}
+
+foreign_obj_type! {
+    type CType = MTLCounterSampleBufferDescriptor;
+    pub struct CounterSampleBufferDescriptor;
+}
+
+impl CounterSampleBufferDescriptor {
+    pub fn new() -> Self {
+        let class = class!(MTLCounterSampleBufferDescriptor);
+        unsafe { msg_send![class, new] }
+    }
+}
+
+impl CounterSampleBufferDescriptorRef {
+    pub fn counter_set(&self) -> &CounterSetRef {
+        unsafe { msg_send![self, counterSet] }
+    }
+
+    pub fn set_counter_set(&self, counter_set: &CounterSetRef) {
+        unsafe { msg_send![self, setCounterSet: counter_set] }
+    }
+
+    pub fn label(&self) -> &str {
+        unsafe { msg_send![self, label] }
+    }
+
+    pub fn set_label(&self, label: &str) {
+        unsafe { msg_send![self, setLabel: label] }
+    }
+
+    pub fn sample_count(&self) -> u64 {
+        unsafe { msg_send![self, sampleCount] }
+    }
+
+    pub fn set_sample_count(&self, sample_count: u64) {
+        unsafe { msg_send![self, setSampleCount: sample_count] }
+    }
+
+    pub fn storage_mode(&self) -> MTLStorageMode {
+        unsafe { msg_send![self, storageMode] }
+    }
+
+    pub fn set_storage_mode(&self, storage_mode: MTLStorageMode) {
+        unsafe { msg_send![self, setStorageMode: storage_mode] }
+    }
+}
+
+/// See <https://developer.apple.com/documentation/metal/mtlcountersamplebuffer>
+pub enum MTLCounterSampleBuffer {}
+
+foreign_obj_type! {
+    type CType = MTLCounterSampleBuffer;
+    pub struct CounterSampleBuffer;
+}
+
+/// See <https://developer.apple.com/documentation/metal/mtlcounter>
+pub enum MTLCounter {}
+
+foreign_obj_type! {
+    type CType = MTLCounter;
+    pub struct Counter;
+}
+
+impl CounterRef {}
+
+/// See <https://developer.apple.com/documentation/metal/mtlcounterset>
+pub enum MTLCounterSet {}
+
+foreign_obj_type! {
+    type CType = MTLCounterSet;
+    pub struct CounterSet;
+}
+
+impl CounterSetRef {
+    pub fn name(&self) -> &str {
+        unsafe {
+            let name = msg_send![self, name];
+            crate::nsstring_as_str(name)
+        }
+    }
+}
+
+/// See <https://developer.apple.com/documentation/metal/mtlcommoncounterset>
+pub enum MTLCommonCounterSet {}
+
+/// See <https://developer.apple.com/documentation/metal/mtlcommoncounter>
+pub enum MTLCommonCounter {}
+
+foreign_obj_type! {
+    type CType = MTLCommonCounter;
+    pub struct CommonCounter;
+}

src/device.rs

@@ -1987,6 +1987,20 @@ impl DeviceRef {
         }
     }
 
+    pub fn new_counter_sample_buffer_with_descriptor(
+        &self,
+        descriptor: &CounterSampleBufferDescriptorRef,
+    ) -> Result<CounterSampleBuffer, String> {
+        unsafe {
+            let counter_sample_buffer: *mut MTLCounterSampleBuffer = try_objc! { err =>
+                msg_send![self, newCounterSampleBufferWithDescriptor: descriptor error:&mut err]
+            };
+
+            assert!(!counter_sample_buffer.is_null());
+            Ok(CounterSampleBuffer::from_ptr(counter_sample_buffer))
+        }
+    }
+
     pub fn new_texture(&self, descriptor: &TextureDescriptorRef) -> Texture {
         unsafe { msg_send![self, newTextureWithDescriptor: descriptor] }
     }
@@ -2149,4 +2163,22 @@ impl DeviceRef {
     ) -> accelerator_structure::AccelerationStructure {
         unsafe { msg_send![self, newAccelerationStructureWithSize: size] }
     }
+
+    pub fn sample_timestamps(&self, cpu_timestamp: &mut u64, gpu_timestamp: &mut u64) {
+        unsafe { msg_send![self, sampleTimestamps: cpu_timestamp gpuTimestamp: gpu_timestamp] }
+    }
+
+    pub fn counter_sets(&self) -> Vec<CounterSet> {
+        unsafe {
+            let counter_sets: *mut Object = msg_send![self, counterSets];
+            let count: NSUInteger = msg_send![counter_sets, count];
+            let ret = (0..count)
+                .map(|i| {
+                    let a = msg_send![counter_sets, objectAtIndex: i];
+                    CounterSet::from_ptr(a)
+                })
+                .collect();
+            ret
+        }
+    }
 }