pyobjc-framework-metal

import Metal import objc import struct def run_metal_kernel(): # 1. Get the default Metal device device = Metal.MTLCreateSystemDefaultDevice() if device is None: print("Error: No Metal device found.") return print(f"Using Metal device: {device.name()}") # 2. Create a simple Metal shader (MSL) source # This kernel adds two numbers kernel_source = """ #include <metal_stdlib> kernel void add_numbers( device const float *inA [[buffer(0)]], device const float *inB [[buffer(1)]], device float *out [[buffer(2)]], uint id [[thread_position_in_grid]]) { out[id] = inA[id] + inB[id]; } """ # 3. Create a library from the source # Using newLibraryWithSource_options_error_ instead of newLibraryWithSource_options_error # as PyObjC usually appends '_' to methods with Objective-C error pointers. error_ptr = objc.nil library = device.newLibraryWithSource_options_error_(kernel_source, objc.nil, error_ptr) if library is None: # Check if error_ptr now points to an actual error object if error_ptr and error_ptr[0] is not objc.nil: # error_ptr is a C array of MTL_Error* in PyObjC error_obj = error_ptr[0] print(f"Failed to create Metal library: {error_obj.localizedDescription()}") else: print("Failed to create Metal library (unknown error).") return # 4. Get the kernel function function = library.newFunctionWithName_("add_numbers") if function is None: print("Error: Failed to find kernel function 'add_numbers'.") return # 5. Create a compute pipeline state pipeline_state = device.newComputePipelineStateWithFunction_error_(function, objc.nil) if pipeline_state is None: print("Error: Failed to create compute pipeline state.") return # 6. Prepare data data_size = 10 * struct.calcsize('f') # 10 floats input_a = [float(i) for i in range(10)] input_b = [float(i * 2) for i in range(10)] output_data = [0.0] * 10 # Create Metal buffers buffer_a = device.newBufferWithBytes_length_options_(bytes(struct.pack('f'*10, *input_a)), data_size, Metal.MTLResourceStorageModeManaged) buffer_b = device.newBufferWithBytes_length_options_(bytes(struct.pack('f'*10, *input_b)), data_size, Metal.MTLResourceStorageModeManaged) buffer_out = device.newBufferWithLength_options_(data_size, Metal.MTLResourceStorageModeManaged) # 7. Create a command queue command_queue = device.newCommandQueue_() if command_queue is None: print("Error: Failed to create command queue.") return # 8. Create a command buffer command_buffer = command_queue.commandBuffer_() # 9. Create a compute command encoder compute_encoder = command_buffer.computeCommandEncoder_() compute_encoder.setComputePipelineState_(pipeline_state) compute_encoder.setBuffer_offset_atIndex_(buffer_a, 0, 0) compute_encoder.setBuffer_offset_atIndex_(buffer_b, 0, 1) compute_encoder.setBuffer_offset_atIndex_(buffer_out, 0, 2) # 10. Dispatch threads grid_size = Metal.MTLSizeMake(10, 1, 1) thread_group_size = Metal.MTLSizeMake(min(10, pipeline_state.maxTotalThreadsPerThreadgroup()), 1, 1) compute_encoder.dispatchThreads_threadsPerThreadgroup_(grid_size, thread_group_size) compute_encoder.endEncoding_() # 11. Commit and wait for completion command_buffer.commit_() command_buffer.waitUntilCompleted_() # 12. Read results back to CPU (if using managed storage mode) buffer_out.didModifyRange_(Metal.NSMakeRange(0, data_size)) result_bytes = buffer_out.contents().tobytes() result = struct.unpack('f'*10, result_bytes) print("Input A:", input_a) print("Input B:", input_b) print("Output (A+B):", list(result)) expected_output = [input_a[i] + input_b[i] for i in range(10)] if all(abs(r - e) < 1e-5 for r, e in zip(result, expected_output)): print("✅ Output matches expected values.") else: print("❌ Output does not match expected values.") if __name__ == '__main__': run_metal_kernel()