a compute example

README.md

@@ -6,6 +6,7 @@ Unsafe Rusty [Metal](https://developer.apple.com/documentation/metal?language=ob
 Backfill List
 -
 **HIGH PRIORITY**:
+- [ ] Gol + sum examples
 - [ ] Tessellation drawing
 - [ ] Indirect buffers
 - [ ] Blit + resource state encoder

examples/sum/main.rs

@@ -0,0 +1,55 @@
+use iron_oxide::*;
+use std::os::raw::c_void;
+
+const IN_DATA: [u32; 9] = [1, 2, 3, 4, 5, 6, 7, 8, 9];
+
+fn main() {
+    unsafe {
+        let devices = MTLCopyAllDevices();
+        let device = devices.into_iter().find_map(|d| Some(d)).unwrap();
+
+        let queue = device.new_command_queue();
+
+        let library = device
+            .new_library_with_source(include_str!("sum.metal"), &MTLCompileOptions::new())
+            .unwrap();
+        let sum_kernel = library.new_function_with_name("sum").unwrap();
+        let sum_pipeline = device
+            .new_compute_pipeline_state_with_function(&sum_kernel)
+            .unwrap();
+
+        let out_buffer = device.new_buffer_with_length(
+            1,
+            MTLResourceOptions::new()
+                .set_storage_mode(MTLStorageMode::Shared)
+                .set_cpu_cache_mode(MTLCPUCacheMode::Default),
+        );
+
+        let command_buffer = queue.new_command_buffer(true);
+
+        let encoder = command_buffer.new_compute_encoder();
+        encoder.set_compute_pipeline_state(&sum_pipeline);
+        encoder.set_bytes(
+            IN_DATA.as_ptr() as *const c_void,
+            IN_DATA.len() as NSUInteger * 4,
+            0,
+        );
+        encoder.set_buffer(&out_buffer, 0, 1);
+        encoder.dispatch_threads(MTLSize {
+            width: IN_DATA.len() as NSUInteger,
+            height: 1,
+            depth: 1,
+        }, MTLSize {
+            width: IN_DATA.len() as NSUInteger,
+            height: 1,
+            depth: 1,
+        });
+        encoder.end_encoding();
+
+        command_buffer.commit();
+        command_buffer.wait_until_completed();
+
+        let sum: u32 = *(out_buffer.get_contents() as *const u32);
+        println!("{}", sum);
+    };
+}

examples/sum/sum.metal

@@ -0,0 +1,10 @@
+#include <metal_stdlib>
+
+using namespace metal;
+
+kernel void sum(device uint *data [[ buffer(0) ]],
+                volatile device atomic_uint *sum [[ buffer(1) ]],
+                uint gid [[ thread_position_in_grid ]])
+{
+    atomic_fetch_add_explicit(sum, data[gid], memory_order_relaxed);
+}

src/commandbuffer.rs

@@ -1,7 +1,7 @@
 use crate::import_objc_macros::*;
 use crate::{
-    handle, DeviceCreated, MTLDrawable, MTLParallelRenderCommandEncoder, MTLRenderCommandEncoder,
-    MTLRenderPassDescriptor, NSError, Object, ObjectPointer,
+    handle, DeviceCreated, MTLComputeCommandEncoder, MTLDrawable, MTLParallelRenderCommandEncoder,
+    MTLRenderCommandEncoder, MTLRenderPassDescriptor, NSError, Object, ObjectPointer,
 };
 
 #[repr(u64)]
@@ -84,6 +84,12 @@ impl MTLCommandBuffer {
             msg_send![k, retain]
         })
     }
+    pub unsafe fn new_compute_encoder(&self) -> MTLComputeCommandEncoder {
+        MTLComputeCommandEncoder::from_ptr({
+            let k = ObjectPointer(msg_send![self.get_ptr(), computeCommandEncoder]);
+            msg_send![k, retain]
+        })
+    }
 }
 
 impl Object for MTLCommandBuffer {

src/device.rs

@@ -214,7 +214,7 @@ impl MTLDevice {
     }
     pub unsafe fn new_compute_pipeline_state_with_function(
         &self,
-        function: MTLFunction,
+        function: &MTLFunction,
     ) -> Result<MTLComputePipelineState, NSError> {
         let mut err = ObjectPointer(std::ptr::null_mut());
         let b = ObjectPointer(msg_send![

src/encoder/compute.rs

@@ -1,5 +1,8 @@
 use crate::import_objc_macros::*;
-use crate::{handle, MTLBuffer, MTLCommandEncoder, MTLSamplerState, MTLSize, MTLTexture, NSRange, NSUInteger, NSUIntegerRange, Object, ObjectPointer, MTLComputePipelineState};
+use crate::{
+    handle, MTLBuffer, MTLCommandEncoder, MTLComputePipelineState, MTLSamplerState, MTLSize,
+    MTLTexture, NSRange, NSUInteger, NSUIntegerRange, Object, ObjectPointer,
+};
 use std::os::raw::c_void;
 
 pub struct MTLComputeCommandEncoder(ObjectPointer);