PyObjC Metal Framework

12.1 · active · verified Tue Apr 14

Wrappers for the “Metal” framework on macOS. PyObjC allows full-featured Cocoa applications to be written in pure Python, bridging Python and Objective-C. This specific package provides bindings for Apple's Metal framework, enabling GPU-accelerated computing and graphics. It is actively maintained with frequent updates tied to macOS SDK releases, currently at version 12.1.

Warnings

Install

Imports

Quickstart

This quickstart demonstrates how to execute a basic Metal compute kernel using `pyobjc-framework-metal`. It initializes a Metal device, compiles a simple shader to add two arrays of numbers, sets up input/output buffers, dispatches the compute command, and reads the results back.

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()

view raw JSON →