INTRODUCTION TO METAL AND GPU COMPUTING ON APPLE PLATFORMS
Apple's Metal framework represents a low-level, high-performance interface for graphics and compute operations on Apple devices. Unlike higher-level frameworks, Metal provides direct access to the GPU hardware, enabling developers to achieve maximum performance for computationally intensive tasks. Metal is available on all modern Apple platforms including macOS, iOS, iPadOS, and tvOS, with particularly impressive performance on Apple Silicon devices that feature unified memory architecture.
The Metal Performance Shaders framework, commonly abbreviated as MPS, builds upon the Metal foundation by providing a collection of highly optimized compute and graphics shaders. These shaders implement common operations such as matrix multiplication, convolution, image processing, and neural network primitives. Apple's engineers have hand-tuned these implementations to extract maximum performance from Apple GPU architectures, making MPS an excellent choice for scientific computing, machine learning, and data processing tasks.
When we combine Python's Foreign Function Interface capabilities with Metal, we create a powerful development environment. Python provides rapid prototyping, easy data manipulation, and a rich ecosystem of scientific libraries, while Metal delivers the raw computational power of the GPU. This combination allows developers to write high-level Python code that orchestrates low-level GPU operations, achieving both developer productivity and runtime performance. You‘ll learn about how to compute matrix and vector implementations in thu Metal GPU and call them from your Python applications. Why is this important? For example, you could use these operations for implementing your own efficient artificial neural networks, graphics operations, scientific computations in Python while lifting the heavy load to the GPU.
WHY FOCUS ON METAL FOR APPLE PLATFORMS
There are several compelling reasons to use Metal for GPU computing on Apple platforms. First, Metal is the only officially supported GPU computing API on modern macOS systems. Apple deprecated OpenCL support and never supported CUDA on their platforms, making Metal the native choice for GPU acceleration. Second, Apple Silicon devices with their unified memory architecture provide unique advantages for Metal applications. The CPU and GPU share the same physical memory, eliminating the need for expensive data transfers between separate memory spaces. Third, Metal Performance Shaders provide highly optimized implementations that are specifically tuned for Apple GPU architectures, often outperforming generic implementations.
Fourth, Metal integrates seamlessly with other Apple frameworks and technologies. You can combine Metal compute operations with Core Image for image processing, Core ML for machine learning inference, or Metal rendering for visualization. Fifth, the Metal Shading Language is based on C++ and provides modern language features including templates, operator overloading, and lambda expressions, making it familiar to C++ developers while offering GPU-specific extensions.
UNDERSTANDING THE METAL ARCHITECTURE
Metal organizes GPU work around several key concepts that differ from traditional CPU programming models. Understanding these concepts is essential for writing efficient Metal code and properly interfacing with Metal from Python.
The device represents the GPU itself. On systems with multiple GPUs, you can enumerate available devices and select the appropriate one. Most commonly, you use the system default device, which Metal automatically selects based on system configuration and power management considerations.
Command queues manage the submission of work to the GPU. A command queue is associated with a specific device and maintains an ordered list of command buffers to execute. You typically create one or more command queues during initialization and reuse them throughout your application's lifetime. Command queues are thread-safe, allowing multiple threads to submit work concurrently.
Command buffers represent a sequence of GPU commands to execute. You create a command buffer from a command queue, record commands into it using encoders, then commit the buffer for execution. Command buffers execute asynchronously, allowing the CPU to continue working while the GPU processes commands. You can register completion handlers to receive notification when a command buffer finishes execution.
Command encoders write specific types of commands into command buffers. Compute command encoders record compute operations, which are the focus of our matrix and vector operations. Blit command encoders handle memory operations like copying and filling buffers. Render command encoders record graphics rendering commands, which we will not use in this tutorial.
Buffers represent regions of GPU memory. Metal buffers can store any type of data including vertex data, textures, or general-purpose arrays. Buffers have storage modes that control memory placement and access patterns. Shared storage mode creates buffers accessible from both CPU and GPU, which is convenient but may have performance implications on some systems. Private storage mode creates GPU-only buffers for maximum performance. Managed storage mode provides explicit synchronization between CPU and GPU copies.
Compute pipeline states represent compiled and optimized compute kernels ready for execution. You create a pipeline state from a Metal function, which is a compute kernel defined in Metal Shading Language. The pipeline state includes information about threadgroup sizes, memory usage, and other execution parameters that Metal uses for optimization.
METAL SHADING LANGUAGE FUNDAMENTALS
The Metal Shading Language, often abbreviated as MSL, is the programming language for writing GPU kernels in Metal. MSL is based on C++ fourteen and includes GPU-specific extensions for parallel programming. Understanding MSL syntax and semantics is necessary for implementing our matrix and vector operations.
MSL functions use qualifiers to indicate where they execute. The kernel qualifier marks a function as a compute kernel that can be dispatched from the host. Kernel functions must return void and can accept parameters from buffers, textures, or threadgroup memory. The vertex and fragment qualifiers mark graphics shader functions, which we will not use in this tutorial.
MSL provides address space qualifiers that specify where data resides. The device qualifier indicates GPU memory accessible by all threads. The threadgroup qualifier indicates fast memory shared among threads in the same threadgroup. The constant qualifier indicates read-only GPU memory optimized for uniform access patterns. The thread qualifier indicates private memory accessible only by a single thread, which is the default for local variables.
Attribute syntax in MSL uses double square brackets to bind parameters to resources and provide thread identification. The buffer attribute binds a parameter to a buffer at a specific index. The thread_position_in_grid attribute provides the global thread index across all threadgroups. The thread_position_in_threadgroup attribute provides the thread index within its threadgroup. The threadgroup_position_in_grid attribute provides the threadgroup index.
MSL includes built-in functions for common operations. Mathematical functions like sin, cos, sqrt, and pow operate on scalar and vector types. The threadgroup_barrier function synchronizes threads within a threadgroup, ensuring all threads reach the barrier before any proceed. Atomic operations like atomic_fetch_add provide thread-safe access to shared memory.
SETTING UP THE DEVELOPMENT ENVIRONMENT
Before implementing our Metal-accelerated operations, we need to set up the development environment properly. Metal development on macOS requires Xcode or at least the Xcode Command Line Tools, which include the Metal compiler and related utilities.
To install Xcode Command Line Tools, open Terminal and run the following command:
xcode-select --install
This command opens a dialog that guides you through the installation process. After installation completes, verify that the Metal compiler is available by running:
xcrun -sdk macosx metal --version
This command should display version information for the Metal compiler. If you see an error, ensure that Xcode Command Line Tools installed correctly and that your system meets the minimum requirements for Metal development.
On the Python side, you need Python 3.7 or later with NumPy installed. NumPy provides the array types we will use for data manipulation and validation. Install NumPy using pip if not already installed:
pip install numpy
You should also verify that your Mac supports Metal. All Macs from 2012 or later support Metal, but the best performance comes from Apple Silicon Macs with Mx chips. You can check Metal support by running a simple Python script:
import subprocess
result = subprocess.run(['system_profiler', 'SPDisplaysDataType'],
capture_output=True, text=True)
print(result.stdout)
Look for "Metal" in the output to confirm Metal support. Apple Silicon Macs will show "Metal: Supported, feature set macOS GPUFamily2 v1" or similar.
DESIGNING THE API FOR MATRIX AND VECTOR OPERATIONS
A well-designed API makes the library intuitive to use and easy to maintain. Our Metal-based matrix and vector operations library will provide a clean Python interface that hides the complexity of Metal while exposing powerful GPU-accelerated functionality.
We will implement the following vector operations. Vector addition computes element-wise sum of two vectors. Vector subtraction computes element-wise difference. Vector multiplication computes element-wise product, also known as Hadamard product. Vector scaling multiplies all elements by a scalar value. Dot product computes the sum of element-wise products, reducing two vectors to a single scalar.
For matrix operations, we will implement the following. Matrix-vector multiplication multiplies a matrix by a vector, producing a vector result. Matrix-matrix multiplication multiplies two matrices, producing a matrix result. Matrix transpose swaps rows and columns. Matrix addition computes element-wise sum of two matrices. Matrix scaling multiplies all elements by a scalar value.
Each operation will follow a consistent pattern. The Python interface accepts NumPy arrays as inputs and returns NumPy arrays as outputs. The implementation validates input types and dimensions, converts NumPy arrays to Metal buffers, executes the Metal compute kernel, and copies results back to NumPy arrays. This design provides a familiar NumPy-like interface while leveraging GPU acceleration transparently.
Error handling will be comprehensive. The implementation checks for invalid input types, incompatible dimensions, non-contiguous arrays, and Metal API failures. Each error condition raises an appropriate Python exception with a descriptive message, making debugging straightforward.
IMPLEMENTING VECTOR ADDITION WITH METAL
Let us begin with the simplest operation: vector addition. This operation takes two input vectors of equal length and produces an output vector where each element is the sum of corresponding input elements. Vector addition is perfectly parallel, meaning each output element can be computed independently.
Here is the Metal Shading Language kernel for vector addition:
#include <metal_stdlib>
using namespace metal;
kernel void vector_add(device const float* a [[buffer(0)]],
device const float* b [[buffer(1)]],
device float* c [[buffer(2)]],
constant uint& n [[buffer(3)]],
uint idx [[thread_position_in_grid]]) {
if (idx < n) {
c[idx] = a[idx] + b[idx];
}
}
This kernel demonstrates the basic structure of Metal compute kernels. The kernel qualifier indicates this function executes on the GPU and can be dispatched from the host. The function returns void, which is required for kernel functions.
The parameters use the device address space qualifier to indicate they point to GPU memory. The const qualifier on the input pointers indicates they are read-only. The buffer attribute with an index in double square brackets binds each parameter to a specific buffer slot. We pass the array length n as a constant buffer parameter rather than using a compile-time constant, allowing the same kernel to work with arrays of any size.
The thread_position_in_grid attribute provides the global thread index. Metal automatically assigns each thread a unique index when we dispatch the kernel. The index ranges from zero to the total number of threads minus one.
Inside the kernel, we check whether the thread index is within bounds. This check is necessary because we typically dispatch more threads than we have elements to ensure full GPU utilization. Threads with indices beyond the array length simply return without doing work.
The actual computation is straightforward: we read elements from the input arrays at the current index, add them, and store the result. This single line of code executes simultaneously across thousands of GPU threads, each processing a different element.
Now we need Objective-C++ wrapper code to interface with Metal from C, which Python can then call through ctypes. Here is the implementation:
#import <Metal/Metal.h>
#import <Foundation/Foundation.h>
#include <string.h>
#include <stdio.h>
static id<MTLDevice> g_device = nil;
static id<MTLCommandQueue> g_commandQueue = nil;
static id<MTLLibrary> g_library = nil;
static int initialize_metal() {
if (g_device != nil) {
return 0;
}
@autoreleasepool {
g_device = MTLCreateSystemDefaultDevice();
if (g_device == nil) {
fprintf(stderr, "Metal: Failed to create device\n");
return -1;
}
g_commandQueue = [g_device newCommandQueue];
if (g_commandQueue == nil) {
fprintf(stderr, "Metal: Failed to create command queue\n");
return -1;
}
NSError* error = nil;
NSString* libraryPath = @"./metal_ops.metallib";
if (![[NSFileManager defaultManager] fileExistsAtPath:libraryPath]) {
NSString* bundlePath = [[NSBundle mainBundle] pathForResource:@"metal_ops"
ofType:@"metallib"];
if (bundlePath != nil) {
libraryPath = bundlePath;
}
}
NSURL* libraryURL = [NSURL fileURLWithPath:libraryPath];
g_library = [g_device newLibraryWithURL:libraryURL error:&error];
if (g_library == nil) {
fprintf(stderr, "Metal: Failed to load library: %s\n",
[[error localizedDescription] UTF8String]);
return -1;
}
}
return 0;
}
static int dispatch_compute_1d(id<MTLComputeCommandEncoder> encoder,
id<MTLComputePipelineState> pipeline,
uint size) {
if (size == 0) {
return 0;
}
MTLSize gridSize = MTLSizeMake(size, 1, 1);
NSUInteger w = pipeline.threadExecutionWidth;
NSUInteger maxThreads = pipeline.maxTotalThreadsPerThreadgroup;
if (w == 0) w = 1;
if (maxThreads == 0) maxThreads = 256;
NSUInteger threadGroupSize = maxThreads;
threadGroupSize = (threadGroupSize / w) * w;
if (threadGroupSize > size) {
threadGroupSize = ((size + w - 1) / w) * w;
}
if (threadGroupSize == 0) {
threadGroupSize = w;
}
if (threadGroupSize > maxThreads) {
threadGroupSize = maxThreads;
}
MTLSize threadgroupSize = MTLSizeMake(threadGroupSize, 1, 1);
[encoder dispatchThreads:gridSize threadsPerThreadgroup:threadgroupSize];
return 0;
}
static int metal_vector_binary_op(const char* kernel_name,
const float* h_a, const float* h_b,
float* h_c, int n) {
if (n <= 0) {
return 0;
}
if (n > (int)UINT32_MAX) {
fprintf(stderr, "Metal: Array size exceeds maximum\n");
return -1;
}
@autoreleasepool {
if (initialize_metal() != 0) {
return -1;
}
NSError* error = nil;
NSString* functionName = [NSString stringWithUTF8String:kernel_name];
id<MTLFunction> function = [g_library newFunctionWithName:functionName];
if (function == nil) {
fprintf(stderr, "Metal: Function '%s' not found\n", kernel_name);
return -1;
}
id<MTLComputePipelineState> pipeline =
[g_device newComputePipelineStateWithFunction:function error:&error];
if (pipeline == nil) {
fprintf(stderr, "Metal: Failed to create pipeline: %s\n",
[[error localizedDescription] UTF8String]);
return -1;
}
size_t size = (size_t)n * sizeof(float);
uint32_t n_uint = (uint32_t)n;
id<MTLBuffer> bufferA = [g_device newBufferWithBytes:h_a
length:size
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferB = [g_device newBufferWithBytes:h_b
length:size
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferC = [g_device newBufferWithLength:size
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferN = [g_device newBufferWithBytes:&n_uint
length:sizeof(uint32_t)
options:MTLResourceStorageModeShared];
if (bufferA == nil || bufferB == nil || bufferC == nil || bufferN == nil) {
fprintf(stderr, "Metal: Failed to allocate buffers\n");
return -1;
}
id<MTLCommandBuffer> commandBuffer = [g_commandQueue commandBuffer];
if (commandBuffer == nil) {
fprintf(stderr, "Metal: Failed to create command buffer\n");
return -1;
}
id<MTLComputeCommandEncoder> encoder = [commandBuffer computeCommandEncoder];
if (encoder == nil) {
fprintf(stderr, "Metal: Failed to create encoder\n");
return -1;
}
[encoder setComputePipelineState:pipeline];
[encoder setBuffer:bufferA offset:0 atIndex:0];
[encoder setBuffer:bufferB offset:0 atIndex:1];
[encoder setBuffer:bufferC offset:0 atIndex:2];
[encoder setBuffer:bufferN offset:0 atIndex:3];
if (dispatch_compute_1d(encoder, pipeline, n) != 0) {
[encoder endEncoding];
return -1;
}
[encoder endEncoding];
[commandBuffer commit];
[commandBuffer waitUntilCompleted];
if (commandBuffer.status == MTLCommandBufferStatusError) {
fprintf(stderr, "Metal: Command buffer error: %s\n",
[[commandBuffer.error localizedDescription] UTF8String]);
return -1;
}
memcpy(h_c, [bufferC contents], size);
return 0;
}
}
extern "C" {
int metal_vector_add(const float* h_a, const float* h_b,
float* h_c, int n) {
return metal_vector_binary_op("vector_add", h_a, h_b, h_c, n);
}
}
This implementation demonstrates the complete Metal workflow. The initialize_metal function creates and caches the Metal device, command queue, and library. We use static variables to store these objects so they persist across function calls, avoiding repeated initialization overhead.
We obtain the default Metal device using MTLCreateSystemDefaultDevice. This function returns the most appropriate GPU for the current system configuration. We create a command queue from the device, which manages command buffer submission and execution. We load the compiled Metal library from the metallib file, checking both the current directory and the application bundle.
The dispatch_compute_1d helper function calculates appropriate threadgroup sizes based on the pipeline's capabilities and the problem size. This ensures we use the GPU efficiently while avoiding division by zero or other edge cases.
The metal_vector_binary_op function implements the common pattern for binary vector operations. It validates the input size, creates Metal buffers, sets up the compute encoder, dispatches the work, and copies results back to the host.
IMPLEMENTING VECTOR SUBTRACTION AND MULTIPLICATION
Vector subtraction and element-wise multiplication follow the same pattern as addition. Here are the Metal kernels:
kernel void vector_sub(device const float* a [[buffer(0)]],
device const float* b [[buffer(1)]],
device float* c [[buffer(2)]],
constant uint& n [[buffer(3)]],
uint idx [[thread_position_in_grid]]) {
if (idx < n) {
c[idx] = a[idx] - b[idx];
}
}
kernel void vector_mul(device const float* a [[buffer(0)]],
device const float* b [[buffer(1)]],
device float* c [[buffer(2)]],
constant uint& n [[buffer(3)]],
uint idx [[thread_position_in_grid]]) {
if (idx < n) {
c[idx] = a[idx] * b[idx];
}
}
The wrapper functions use the generic helper:
extern "C" {
int metal_vector_sub(const float* h_a, const float* h_b,
float* h_c, int n) {
return metal_vector_binary_op("vector_sub", h_a, h_b, h_c, n);
}
int metal_vector_mul(const float* h_a, const float* h_b,
float* h_c, int n) {
return metal_vector_binary_op("vector_mul", h_a, h_b, h_c, n);
}
}
IMPLEMENTING VECTOR SCALING
Vector scaling multiplies every element of a vector by a scalar value. The Metal kernel is:
kernel void vector_scale(device const float* a [[buffer(0)]],
constant float& scale [[buffer(1)]],
device float* b [[buffer(2)]],
constant uint& n [[buffer(3)]],
uint idx [[thread_position_in_grid]]) {
if (idx < n) {
b[idx] = a[idx] * scale;
}
}
The wrapper function:
extern "C" {
int metal_vector_scale(const float* h_a, float scale,
float* h_b, int n) {
if (n <= 0) {
return 0;
}
if (n > (int)UINT32_MAX) {
fprintf(stderr, "Metal: Array size exceeds maximum\n");
return -1;
}
@autoreleasepool {
if (initialize_metal() != 0) {
return -1;
}
NSError* error = nil;
id<MTLFunction> function = [g_library newFunctionWithName:@"vector_scale"];
if (function == nil) {
fprintf(stderr, "Metal: Function 'vector_scale' not found\n");
return -1;
}
id<MTLComputePipelineState> pipeline =
[g_device newComputePipelineStateWithFunction:function error:&error];
if (pipeline == nil) {
fprintf(stderr, "Metal: Failed to create pipeline\n");
return -1;
}
size_t size = (size_t)n * sizeof(float);
uint32_t n_uint = (uint32_t)n;
id<MTLBuffer> bufferA = [g_device newBufferWithBytes:h_a
length:size
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferScale = [g_device newBufferWithBytes:&scale
length:sizeof(float)
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferB = [g_device newBufferWithLength:size
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferN = [g_device newBufferWithBytes:&n_uint
length:sizeof(uint32_t)
options:MTLResourceStorageModeShared];
if (bufferA == nil || bufferScale == nil || bufferB == nil || bufferN == nil) {
fprintf(stderr, "Metal: Failed to allocate buffers\n");
return -1;
}
id<MTLCommandBuffer> commandBuffer = [g_commandQueue commandBuffer];
if (commandBuffer == nil) {
return -1;
}
id<MTLComputeCommandEncoder> encoder = [commandBuffer computeCommandEncoder];
if (encoder == nil) {
return -1;
}
[encoder setComputePipelineState:pipeline];
[encoder setBuffer:bufferA offset:0 atIndex:0];
[encoder setBuffer:bufferScale offset:0 atIndex:1];
[encoder setBuffer:bufferB offset:0 atIndex:2];
[encoder setBuffer:bufferN offset:0 atIndex:3];
if (dispatch_compute_1d(encoder, pipeline, n) != 0) {
[encoder endEncoding];
return -1;
}
[encoder endEncoding];
[commandBuffer commit];
[commandBuffer waitUntilCompleted];
if (commandBuffer.status == MTLCommandBufferStatusError) {
fprintf(stderr, "Metal: Command buffer error\n");
return -1;
}
memcpy(h_b, [bufferB contents], size);
return 0;
}
}
}
IMPLEMENTING VECTOR DOT PRODUCT
The dot product operation requires reduction: we must sum all element-wise products into a single scalar result. We implement this using a two-stage approach with partial reductions in threadgroup memory.
Here is the Metal kernel:
kernel void dot_product_partial(device const float* a [[buffer(0)]],
device const float* b [[buffer(1)]],
device float* partial [[buffer(2)]],
constant uint& n [[buffer(3)]],
uint gid [[thread_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint bid [[threadgroup_position_in_grid]]) {
threadgroup float shared_data[256];
float value = 0.0f;
if (gid < n) {
value = a[gid] * b[gid];
}
shared_data[tid] = value;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = 128; stride > 0; stride >>= 1) {
if (tid < stride) {
shared_data[tid] += shared_data[tid + stride];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
if (tid == 0) {
partial[bid] = shared_data[0];
}
}
The wrapper function:
extern "C" {
int metal_dot_product(const float* h_a, const float* h_b,
float* result, int n) {
if (n <= 0) {
*result = 0.0f;
return 0;
}
if (n > (int)UINT32_MAX) {
fprintf(stderr, "Metal: Array size exceeds maximum\n");
return -1;
}
@autoreleasepool {
if (initialize_metal() != 0) {
return -1;
}
NSError* error = nil;
id<MTLFunction> function = [g_library newFunctionWithName:
@"dot_product_partial"];
if (function == nil) {
fprintf(stderr, "Metal: Function 'dot_product_partial' not found\n");
return -1;
}
id<MTLComputePipelineState> pipeline =
[g_device newComputePipelineStateWithFunction:function error:&error];
if (pipeline == nil) {
fprintf(stderr, "Metal: Failed to create pipeline\n");
return -1;
}
size_t size = (size_t)n * sizeof(float);
uint32_t n_uint = (uint32_t)n;
NSUInteger threadGroupSize = 256;
NSUInteger numGroups = ((NSUInteger)n + threadGroupSize - 1) / threadGroupSize;
id<MTLBuffer> bufferA = [g_device newBufferWithBytes:h_a
length:size
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferB = [g_device newBufferWithBytes:h_b
length:size
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferPartial = [g_device newBufferWithLength:
numGroups * sizeof(float)
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferN = [g_device newBufferWithBytes:&n_uint
length:sizeof(uint32_t)
options:MTLResourceStorageModeShared];
if (bufferA == nil || bufferB == nil || bufferPartial == nil || bufferN == nil) {
fprintf(stderr, "Metal: Failed to allocate buffers\n");
return -1;
}
id<MTLCommandBuffer> commandBuffer = [g_commandQueue commandBuffer];
if (commandBuffer == nil) {
return -1;
}
id<MTLComputeCommandEncoder> encoder = [commandBuffer computeCommandEncoder];
if (encoder == nil) {
return -1;
}
[encoder setComputePipelineState:pipeline];
[encoder setBuffer:bufferA offset:0 atIndex:0];
[encoder setBuffer:bufferB offset:0 atIndex:1];
[encoder setBuffer:bufferPartial offset:0 atIndex:2];
[encoder setBuffer:bufferN offset:0 atIndex:3];
MTLSize gridSize = MTLSizeMake(n, 1, 1);
MTLSize threadgroupSize = MTLSizeMake(threadGroupSize, 1, 1);
[encoder dispatchThreads:gridSize threadsPerThreadgroup:threadgroupSize];
[encoder endEncoding];
[commandBuffer commit];
[commandBuffer waitUntilCompleted];
if (commandBuffer.status == MTLCommandBufferStatusError) {
fprintf(stderr, "Metal: Command buffer error\n");
return -1;
}
float* partial = (float*)[bufferPartial contents];
*result = 0.0f;
for (NSUInteger i = 0; i < numGroups; i++) {
*result += partial[i];
}
return 0;
}
}
}
IMPLEMENTING MATRIX-VECTOR MULTIPLICATION
Matrix-vector multiplication computes the product of a matrix and a vector. Here is the Metal kernel:
kernel void matvec(device const float* A [[buffer(0)]],
device const float* x [[buffer(1)]],
device float* y [[buffer(2)]],
constant uint& m [[buffer(3)]],
constant uint& n [[buffer(4)]],
uint row [[thread_position_in_grid]]) {
if (row < m) {
float sum = 0.0f;
for (uint col = 0; col < n; col++) {
sum += A[row * n + col] * x[col];
}
y[row] = sum;
}
}
The wrapper function:
extern "C" {
int metal_matvec(const float* h_A, const float* h_x,
float* h_y, int m, int n) {
if (m <= 0 || n <= 0) {
return 0;
}
if (m > (int)UINT32_MAX || n > (int)UINT32_MAX) {
fprintf(stderr, "Metal: Matrix dimensions exceed maximum\n");
return -1;
}
if ((long long)m * (long long)n > (long long)UINT32_MAX) {
fprintf(stderr, "Metal: Matrix size exceeds maximum\n");
return -1;
}
@autoreleasepool {
if (initialize_metal() != 0) {
return -1;
}
NSError* error = nil;
id<MTLFunction> function = [g_library newFunctionWithName:@"matvec"];
if (function == nil) {
fprintf(stderr, "Metal: Function 'matvec' not found\n");
return -1;
}
id<MTLComputePipelineState> pipeline =
[g_device newComputePipelineStateWithFunction:function error:&error];
if (pipeline == nil) {
fprintf(stderr, "Metal: Failed to create pipeline\n");
return -1;
}
uint32_t m_uint = (uint32_t)m;
uint32_t n_uint = (uint32_t)n;
id<MTLBuffer> bufferA = [g_device newBufferWithBytes:h_A
length:(size_t)m * (size_t)n * sizeof(float)
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferX = [g_device newBufferWithBytes:h_x
length:(size_t)n * sizeof(float)
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferY = [g_device newBufferWithLength:
(size_t)m * sizeof(float)
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferM = [g_device newBufferWithBytes:&m_uint
length:sizeof(uint32_t)
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferN = [g_device newBufferWithBytes:&n_uint
length:sizeof(uint32_t)
options:MTLResourceStorageModeShared];
if (bufferA == nil || bufferX == nil || bufferY == nil ||
bufferM == nil || bufferN == nil) {
fprintf(stderr, "Metal: Failed to allocate buffers\n");
return -1;
}
id<MTLCommandBuffer> commandBuffer = [g_commandQueue commandBuffer];
if (commandBuffer == nil) {
return -1;
}
id<MTLComputeCommandEncoder> encoder = [commandBuffer computeCommandEncoder];
if (encoder == nil) {
return -1;
}
[encoder setComputePipelineState:pipeline];
[encoder setBuffer:bufferA offset:0 atIndex:0];
[encoder setBuffer:bufferX offset:0 atIndex:1];
[encoder setBuffer:bufferY offset:0 atIndex:2];
[encoder setBuffer:bufferM offset:0 atIndex:3];
[encoder setBuffer:bufferN offset:0 atIndex:4];
if (dispatch_compute_1d(encoder, pipeline, m) != 0) {
[encoder endEncoding];
return -1;
}
[encoder endEncoding];
[commandBuffer commit];
[commandBuffer waitUntilCompleted];
if (commandBuffer.status == MTLCommandBufferStatusError) {
fprintf(stderr, "Metal: Command buffer error\n");
return -1;
}
memcpy(h_y, [bufferY contents], (size_t)m * sizeof(float));
return 0;
}
}
}
IMPLEMENTING MATRIX-MATRIX MULTIPLICATION
Matrix-matrix multiplication uses tiled computation for efficiency. Here is the Metal kernel:
constant uint TILE_SIZE = 16;
kernel void matmul_tiled(device const float* A [[buffer(0)]],
device const float* B [[buffer(1)]],
device float* C [[buffer(2)]],
constant uint& m [[buffer(3)]],
constant uint& k [[buffer(4)]],
constant uint& n [[buffer(5)]],
uint2 gid [[thread_position_in_grid]],
uint2 tid [[thread_position_in_threadgroup]]) {
threadgroup float As[TILE_SIZE][TILE_SIZE];
threadgroup float Bs[TILE_SIZE][TILE_SIZE];
uint row = gid.y;
uint col = gid.x;
float sum = 0.0f;
uint numTiles = (k + TILE_SIZE - 1) / TILE_SIZE;
for (uint t = 0; t < numTiles; t++) {
uint aCol = t * TILE_SIZE + tid.x;
if (row < m && aCol < k) {
As[tid.y][tid.x] = A[row * k + aCol];
} else {
As[tid.y][tid.x] = 0.0f;
}
uint bRow = t * TILE_SIZE + tid.y;
if (bRow < k && col < n) {
Bs[tid.y][tid.x] = B[bRow * n + col];
} else {
Bs[tid.y][tid.x] = 0.0f;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint i = 0; i < TILE_SIZE; i++) {
sum += As[tid.y][i] * Bs[i][tid.x];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
if (row < m && col < n) {
C[row * n + col] = sum;
}
}
The wrapper function:
extern "C" {
int metal_matmul(const float* h_A, const float* h_B,
float* h_C, int m, int k, int n) {
if (m <= 0 || k <= 0 || n <= 0) {
return 0;
}
if (m > (int)UINT32_MAX || k > (int)UINT32_MAX || n > (int)UINT32_MAX) {
fprintf(stderr, "Metal: Matrix dimensions exceed maximum\n");
return -1;
}
@autoreleasepool {
if (initialize_metal() != 0) {
return -1;
}
NSError* error = nil;
id<MTLFunction> function = [g_library newFunctionWithName:@"matmul_tiled"];
if (function == nil) {
fprintf(stderr, "Metal: Function 'matmul_tiled' not found\n");
return -1;
}
id<MTLComputePipelineState> pipeline =
[g_device newComputePipelineStateWithFunction:function error:&error];
if (pipeline == nil) {
fprintf(stderr, "Metal: Failed to create pipeline\n");
return -1;
}
uint32_t m_uint = (uint32_t)m;
uint32_t k_uint = (uint32_t)k;
uint32_t n_uint = (uint32_t)n;
id<MTLBuffer> bufferA = [g_device newBufferWithBytes:h_A
length:(size_t)m * (size_t)k * sizeof(float)
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferB = [g_device newBufferWithBytes:h_B
length:(size_t)k * (size_t)n * sizeof(float)
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferC = [g_device newBufferWithLength:
(size_t)m * (size_t)n * sizeof(float)
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferM = [g_device newBufferWithBytes:&m_uint
length:sizeof(uint32_t)
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferK = [g_device newBufferWithBytes:&k_uint
length:sizeof(uint32_t)
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferN = [g_device newBufferWithBytes:&n_uint
length:sizeof(uint32_t)
options:MTLResourceStorageModeShared];
if (bufferA == nil || bufferB == nil || bufferC == nil ||
bufferM == nil || bufferK == nil || bufferN == nil) {
fprintf(stderr, "Metal: Failed to allocate buffers\n");
return -1;
}
id<MTLCommandBuffer> commandBuffer = [g_commandQueue commandBuffer];
if (commandBuffer == nil) {
return -1;
}
id<MTLComputeCommandEncoder> encoder = [commandBuffer computeCommandEncoder];
if (encoder == nil) {
return -1;
}
[encoder setComputePipelineState:pipeline];
[encoder setBuffer:bufferA offset:0 atIndex:0];
[encoder setBuffer:bufferB offset:0 atIndex:1];
[encoder setBuffer:bufferC offset:0 atIndex:2];
[encoder setBuffer:bufferM offset:0 atIndex:3];
[encoder setBuffer:bufferK offset:0 atIndex:4];
[encoder setBuffer:bufferN offset:0 atIndex:5];
uint TILE_SIZE_CONST = 16;
MTLSize gridSize = MTLSizeMake(n, m, 1);
MTLSize threadgroupSize = MTLSizeMake(TILE_SIZE_CONST, TILE_SIZE_CONST, 1);
[encoder dispatchThreads:gridSize threadsPerThreadgroup:threadgroupSize];
[encoder endEncoding];
[commandBuffer commit];
[commandBuffer waitUntilCompleted];
if (commandBuffer.status == MTLCommandBufferStatusError) {
fprintf(stderr, "Metal: Command buffer error\n");
return -1;
}
memcpy(h_C, [bufferC contents], (size_t)m * (size_t)n * sizeof(float));
return 0;
}
}
}
IMPLEMENTING MATRIX TRANSPOSE
Matrix transpose swaps rows and columns. Here is the Metal kernel:
kernel void transpose(device const float* input [[buffer(0)]],
device float* output [[buffer(1)]],
constant uint& width [[buffer(2)]],
constant uint& height [[buffer(3)]],
uint2 gid [[thread_position_in_grid]]) {
if (gid.x < width && gid.y < height) {
uint inputIndex = gid.y * width + gid.x;
uint outputIndex = gid.x * height + gid.y;
output[outputIndex] = input[inputIndex];
}
}
The wrapper function:
extern "C" {
int metal_transpose(const float* h_input, float* h_output,
int width, int height) {
if (width <= 0 || height <= 0) {
return 0;
}
if (width > (int)UINT32_MAX || height > (int)UINT32_MAX) {
fprintf(stderr, "Metal: Matrix dimensions exceed maximum\n");
return -1;
}
@autoreleasepool {
if (initialize_metal() != 0) {
return -1;
}
NSError* error = nil;
id<MTLFunction> function = [g_library newFunctionWithName:@"transpose"];
if (function == nil) {
fprintf(stderr, "Metal: Function 'transpose' not found\n");
return -1;
}
id<MTLComputePipelineState> pipeline =
[g_device newComputePipelineStateWithFunction:function error:&error];
if (pipeline == nil) {
fprintf(stderr, "Metal: Failed to create pipeline\n");
return -1;
}
size_t size = (size_t)width * (size_t)height * sizeof(float);
uint32_t width_uint = (uint32_t)width;
uint32_t height_uint = (uint32_t)height;
id<MTLBuffer> bufferInput = [g_device newBufferWithBytes:h_input
length:size
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferOutput = [g_device newBufferWithLength:size
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferWidth = [g_device newBufferWithBytes:&width_uint
length:sizeof(uint32_t)
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferHeight = [g_device newBufferWithBytes:&height_uint
length:sizeof(uint32_t)
options:MTLResourceStorageModeShared];
if (bufferInput == nil || bufferOutput == nil ||
bufferWidth == nil || bufferHeight == nil) {
fprintf(stderr, "Metal: Failed to allocate buffers\n");
return -1;
}
id<MTLCommandBuffer> commandBuffer = [g_commandQueue commandBuffer];
if (commandBuffer == nil) {
return -1;
}
id<MTLComputeCommandEncoder> encoder = [commandBuffer computeCommandEncoder];
if (encoder == nil) {
return -1;
}
[encoder setComputePipelineState:pipeline];
[encoder setBuffer:bufferInput offset:0 atIndex:0];
[encoder setBuffer:bufferOutput offset:0 atIndex:1];
[encoder setBuffer:bufferWidth offset:0 atIndex:2];
[encoder setBuffer:bufferHeight offset:0 atIndex:3];
MTLSize gridSize = MTLSizeMake(width, height, 1);
NSUInteger w = pipeline.threadExecutionWidth;
if (w == 0) w = 1;
NSUInteger h = pipeline.maxTotalThreadsPerThreadgroup / w;
if (h == 0) h = 1;
if (h > (NSUInteger)height) h = height;
MTLSize threadgroupSize = MTLSizeMake(w, h, 1);
[encoder dispatchThreads:gridSize threadsPerThreadgroup:threadgroupSize];
[encoder endEncoding];
[commandBuffer commit];
[commandBuffer waitUntilCompleted];
if (commandBuffer.status == MTLCommandBufferStatusError) {
fprintf(stderr, "Metal: Command buffer error\n");
return -1;
}
memcpy(h_output, [bufferOutput contents], size);
return 0;
}
}
}
IMPLEMENTING MATRIX ADDITION AND SCALING
Matrix addition and scaling kernels:
kernel void matrix_add(device const float* A [[buffer(0)]],
device const float* B [[buffer(1)]],
device float* C [[buffer(2)]],
constant uint& size [[buffer(3)]],
uint idx [[thread_position_in_grid]]) {
if (idx < size) {
C[idx] = A[idx] + B[idx];
}
}
kernel void matrix_scale(device const float* A [[buffer(0)]],
constant float& scale [[buffer(1)]],
device float* B [[buffer(2)]],
constant uint& size [[buffer(3)]],
uint idx [[thread_position_in_grid]]) {
if (idx < size) {
B[idx] = A[idx] * scale;
}
}
The wrapper functions:
extern "C" {
int metal_matrix_add(const float* h_A, const float* h_B,
float* h_C, int rows, int cols) {
if (rows <= 0 || cols <= 0) {
return 0;
}
long long size = (long long)rows * (long long)cols;
if (size > (long long)INT32_MAX) {
fprintf(stderr, "Metal: Matrix size exceeds maximum\n");
return -1;
}
return metal_vector_binary_op("matrix_add", h_A, h_B, h_C, (int)size);
}
int metal_matrix_scale(const float* h_A, float scale,
float* h_B, int rows, int cols) {
if (rows <= 0 || cols <= 0) {
return 0;
}
long long total_size = (long long)rows * (long long)cols;
if (total_size > (long long)UINT32_MAX) {
fprintf(stderr, "Metal: Matrix size exceeds maximum\n");
return -1;
}
int size = (int)total_size;
@autoreleasepool {
if (initialize_metal() != 0) {
return -1;
}
NSError* error = nil;
id<MTLFunction> function = [g_library newFunctionWithName:@"matrix_scale"];
if (function == nil) {
fprintf(stderr, "Metal: Function 'matrix_scale' not found\n");
return -1;
}
id<MTLComputePipelineState> pipeline =
[g_device newComputePipelineStateWithFunction:function error:&error];
if (pipeline == nil) {
fprintf(stderr, "Metal: Failed to create pipeline\n");
return -1;
}
size_t byte_size = (size_t)size * sizeof(float);
uint32_t size_uint = (uint32_t)size;
id<MTLBuffer> bufferA = [g_device newBufferWithBytes:h_A
length:byte_size
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferScale = [g_device newBufferWithBytes:&scale
length:sizeof(float)
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferB = [g_device newBufferWithLength:byte_size
options:MTLResourceStorageModeShared];
id<MTLBuffer> bufferSize = [g_device newBufferWithBytes:&size_uint
length:sizeof(uint32_t)
options:MTLResourceStorageModeShared];
if (bufferA == nil || bufferScale == nil || bufferB == nil || bufferSize == nil) {
fprintf(stderr, "Metal: Failed to allocate buffers\n");
return -1;
}
id<MTLCommandBuffer> commandBuffer = [g_commandQueue commandBuffer];
if (commandBuffer == nil) {
return -1;
}
id<MTLComputeCommandEncoder> encoder = [commandBuffer computeCommandEncoder];
if (encoder == nil) {
return -1;
}
[encoder setComputePipelineState:pipeline];
[encoder setBuffer:bufferA offset:0 atIndex:0];
[encoder setBuffer:bufferScale offset:0 atIndex:1];
[encoder setBuffer:bufferB offset:0 atIndex:2];
[encoder setBuffer:bufferSize offset:0 atIndex:3];
if (dispatch_compute_1d(encoder, pipeline, size) != 0) {
[encoder endEncoding];
return -1;
}
[encoder endEncoding];
[commandBuffer commit];
[commandBuffer waitUntilCompleted];
if (commandBuffer.status == MTLCommandBufferStatusError) {
fprintf(stderr, "Metal: Command buffer error\n");
return -1;
}
memcpy(h_B, [bufferB contents], byte_size);
return 0;
}
}
}
COMPLETE METAL SHADERS FILE
Here is the complete metal_ops.metal file with all kernels:
#include <metal_stdlib>
using namespace metal;
kernel void vector_add(device const float* a [[buffer(0)]],
device const float* b [[buffer(1)]],
device float* c [[buffer(2)]],
constant uint& n [[buffer(3)]],
uint idx [[thread_position_in_grid]]) {
if (idx < n) {
c[idx] = a[idx] + b[idx];
}
}
kernel void vector_sub(device const float* a [[buffer(0)]],
device const float* b [[buffer(1)]],
device float* c [[buffer(2)]],
constant uint& n [[buffer(3)]],
uint idx [[thread_position_in_grid]]) {
if (idx < n) {
c[idx] = a[idx] - b[idx];
}
}
kernel void vector_mul(device const float* a [[buffer(0)]],
device const float* b [[buffer(1)]],
device float* c [[buffer(2)]],
constant uint& n [[buffer(3)]],
uint idx [[thread_position_in_grid]]) {
if (idx < n) {
c[idx] = a[idx] * b[idx];
}
}
kernel void vector_scale(device const float* a [[buffer(0)]],
constant float& scale [[buffer(1)]],
device float* b [[buffer(2)]],
constant uint& n [[buffer(3)]],
uint idx [[thread_position_in_grid]]) {
if (idx < n) {
b[idx] = a[idx] * scale;
}
}
kernel void dot_product_partial(device const float* a [[buffer(0)]],
device const float* b [[buffer(1)]],
device float* partial [[buffer(2)]],
constant uint& n [[buffer(3)]],
uint gid [[thread_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint bid [[threadgroup_position_in_grid]]) {
threadgroup float shared_data[256];
float value = 0.0f;
if (gid < n) {
value = a[gid] * b[gid];
}
shared_data[tid] = value;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = 128; stride > 0; stride >>= 1) {
if (tid < stride) {
shared_data[tid] += shared_data[tid + stride];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
if (tid == 0) {
partial[bid] = shared_data[0];
}
}
kernel void matvec(device const float* A [[buffer(0)]],
device const float* x [[buffer(1)]],
device float* y [[buffer(2)]],
constant uint& m [[buffer(3)]],
constant uint& n [[buffer(4)]],
uint row [[thread_position_in_grid]]) {
if (row < m) {
float sum = 0.0f;
for (uint col = 0; col < n; col++) {
sum += A[row * n + col] * x[col];
}
y[row] = sum;
}
}
constant uint TILE_SIZE = 16;
kernel void matmul_tiled(device const float* A [[buffer(0)]],
device const float* B [[buffer(1)]],
device float* C [[buffer(2)]],
constant uint& m [[buffer(3)]],
constant uint& k [[buffer(4)]],
constant uint& n [[buffer(5)]],
uint2 gid [[thread_position_in_grid]],
uint2 tid [[thread_position_in_threadgroup]]) {
threadgroup float As[TILE_SIZE][TILE_SIZE];
threadgroup float Bs[TILE_SIZE][TILE_SIZE];
uint row = gid.y;
uint col = gid.x;
float sum = 0.0f;
uint numTiles = (k + TILE_SIZE - 1) / TILE_SIZE;
for (uint t = 0; t < numTiles; t++) {
uint aCol = t * TILE_SIZE + tid.x;
if (row < m && aCol < k) {
As[tid.y][tid.x] = A[row * k + aCol];
} else {
As[tid.y][tid.x] = 0.0f;
}
uint bRow = t * TILE_SIZE + tid.y;
if (bRow < k && col < n) {
Bs[tid.y][tid.x] = B[bRow * n + col];
} else {
Bs[tid.y][tid.x] = 0.0f;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint i = 0; i < TILE_SIZE; i++) {
sum += As[tid.y][i] * Bs[i][tid.x];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
if (row < m && col < n) {
C[row * n + col] = sum;
}
}
kernel void transpose(device const float* input [[buffer(0)]],
device float* output [[buffer(1)]],
constant uint& width [[buffer(2)]],
constant uint& height [[buffer(3)]],
uint2 gid [[thread_position_in_grid]]) {
if (gid.x < width && gid.y < height) {
uint inputIndex = gid.y * width + gid.x;
uint outputIndex = gid.x * height + gid.y;
output[outputIndex] = input[inputIndex];
}
}
kernel void matrix_add(device const float* A [[buffer(0)]],
device const float* B [[buffer(1)]],
device float* C [[buffer(2)]],
constant uint& size [[buffer(3)]],
uint idx [[thread_position_in_grid]]) {
if (idx < size) {
C[idx] = A[idx] + B[idx];
}
}
kernel void matrix_scale(device const float* A [[buffer(0)]],
constant float& scale [[buffer(1)]],
device float* B [[buffer(2)]],
constant uint& size [[buffer(3)]],
uint idx [[thread_position_in_grid]]) {
if (idx < size) {
B[idx] = A[idx] * scale;
}
}
COMPILING THE METAL SHADERS
To compile the Metal shaders into a library:
xcrun -sdk macosx metal -c metal_ops.metal -o metal_ops.air
xcrun -sdk macosx metallib metal_ops.air -o metal_ops.metallib
For optimized builds:
xcrun -sdk macosx metal -O3 -c metal_ops.metal -o metal_ops.air
xcrun -sdk macosx metallib metal_ops.air -o metal_ops.metallib
COMPILING THE OBJECTIVE-C++ WRAPPER
To compile the wrapper into a dynamic library:
clang++ -std=c++17 -O3 -dynamiclib -o libmetal_ops.dylib metal_ops.mm \
-framework Metal -framework Foundation
For debug builds:
clang++ -std=c++17 -g -dynamiclib -o libmetal_ops.dylib metal_ops.mm \
-framework Metal -framework Foundation
PYTHON INTERFACE
Here is the complete Python interface in metal_ops.py:
"""
Metal-accelerated matrix and vector operations for Python.
This module provides GPU-accelerated linear algebra operations using
Apple's Metal framework. All operations are optimized for Apple Silicon
and Intel Macs with Metal support.
"""
import ctypes
import platform
import numpy as np
class MetalOpsError(Exception):
"""Exception raised for Metal operation errors"""
pass
def is_metal_available() -> bool:
"""
Check if Metal is available on the current system.
Returns:
True if Metal library can be loaded, False otherwise
"""
if platform.system() != 'Darwin':
return False
try:
ctypes.CDLL('./libmetal_ops.dylib')
return True
except (OSError, FileNotFoundError):
return False
class MetalVectorOps:
"""
Metal-accelerated vector operations.
All operations expect 1D NumPy arrays with dtype float32.
Arrays must be C-contiguous.
"""
def __init__(self):
"""Initialize Metal vector operations."""
if platform.system() != 'Darwin':
raise MetalOpsError("Metal is only available on macOS")
try:
self.lib = ctypes.CDLL('./libmetal_ops.dylib')
except (OSError, FileNotFoundError) as e:
raise MetalOpsError(f"Failed to load Metal library: {e}")
self._setup_function_signatures()
def _setup_function_signatures(self):
"""Setup ctypes function signatures"""
self.lib.metal_vector_add.argtypes = [
ctypes.POINTER(ctypes.c_float),
ctypes.POINTER(ctypes.c_float),
ctypes.POINTER(ctypes.c_float),
ctypes.c_int
]
self.lib.metal_vector_add.restype = ctypes.c_int
self.lib.metal_vector_sub.argtypes = [
ctypes.POINTER(ctypes.c_float),
ctypes.POINTER(ctypes.c_float),
ctypes.POINTER(ctypes.c_float),
ctypes.c_int
]
self.lib.metal_vector_sub.restype = ctypes.c_int
self.lib.metal_vector_mul.argtypes = [
ctypes.POINTER(ctypes.c_float),
ctypes.POINTER(ctypes.c_float),
ctypes.POINTER(ctypes.c_float),
ctypes.c_int
]
self.lib.metal_vector_mul.restype = ctypes.c_int
self.lib.metal_vector_scale.argtypes = [
ctypes.POINTER(ctypes.c_float),
ctypes.c_float,
ctypes.POINTER(ctypes.c_float),
ctypes.c_int
]
self.lib.metal_vector_scale.restype = ctypes.c_int
self.lib.metal_dot_product.argtypes = [
ctypes.POINTER(ctypes.c_float),
ctypes.POINTER(ctypes.c_float),
ctypes.POINTER(ctypes.c_float),
ctypes.c_int
]
self.lib.metal_dot_product.restype = ctypes.c_int
def _validate_vector_inputs(self, a: np.ndarray, b: np.ndarray):
"""Validate vector inputs for binary operations"""
if not isinstance(a, np.ndarray) or not isinstance(b, np.ndarray):
raise TypeError("Inputs must be numpy arrays")
if a.shape != b.shape:
raise ValueError(f"Input arrays must have the same shape: {a.shape} vs {b.shape}")
if a.dtype != np.float32 or b.dtype != np.float32:
raise TypeError(f"Input arrays must be float32, got {a.dtype} and {b.dtype}")
if a.ndim != 1:
raise ValueError(f"Inputs must be 1D arrays, got {a.ndim}D")
if not a.flags['C_CONTIGUOUS']:
raise ValueError("Array 'a' must be C-contiguous. Use np.ascontiguousarray() to fix.")
if not b.flags['C_CONTIGUOUS']:
raise ValueError("Array 'b' must be C-contiguous. Use np.ascontiguousarray() to fix.")
if a.size > 2**31 - 1:
raise ValueError(f"Array size {a.size} exceeds maximum supported size")
def _validate_single_vector(self, a: np.ndarray):
"""Validate single vector input"""
if not isinstance(a, np.ndarray):
raise TypeError("Input must be a numpy array")
if a.dtype != np.float32:
raise TypeError(f"Input array must be float32, got {a.dtype}")
if a.ndim != 1:
raise ValueError(f"Input must be 1D array, got {a.ndim}D")
if not a.flags['C_CONTIGUOUS']:
raise ValueError("Array must be C-contiguous. Use np.ascontiguousarray() to fix.")
if a.size > 2**31 - 1:
raise ValueError(f"Array size {a.size} exceeds maximum supported size")
def add(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
"""Element-wise vector addition: c = a + b"""
self._validate_vector_inputs(a, b)
n = a.size
c = np.empty_like(a)
a_ptr = a.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
b_ptr = b.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
c_ptr = c.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
result = self.lib.metal_vector_add(a_ptr, b_ptr, c_ptr, n)
if result != 0:
raise MetalOpsError("Metal vector addition failed")
return c
def subtract(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
"""Element-wise vector subtraction: c = a - b"""
self._validate_vector_inputs(a, b)
n = a.size
c = np.empty_like(a)
a_ptr = a.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
b_ptr = b.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
c_ptr = c.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
result = self.lib.metal_vector_sub(a_ptr, b_ptr, c_ptr, n)
if result != 0:
raise MetalOpsError("Metal vector subtraction failed")
return c
def multiply(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
"""Element-wise vector multiplication: c = a * b"""
self._validate_vector_inputs(a, b)
n = a.size
c = np.empty_like(a)
a_ptr = a.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
b_ptr = b.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
c_ptr = c.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
result = self.lib.metal_vector_mul(a_ptr, b_ptr, c_ptr, n)
if result != 0:
raise MetalOpsError("Metal vector multiplication failed")
return c
def scale(self, a: np.ndarray, scalar: float) -> np.ndarray:
"""Scalar multiplication: b = a * scalar"""
self._validate_single_vector(a)
n = a.size
b = np.empty_like(a)
a_ptr = a.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
b_ptr = b.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
result = self.lib.metal_vector_scale(a_ptr, float(scalar), b_ptr, n)
if result != 0:
raise MetalOpsError("Metal vector scaling failed")
return b
def dot(self, a: np.ndarray, b: np.ndarray) -> float:
"""Vector dot product: result = sum(a * b)"""
self._validate_vector_inputs(a, b)
n = a.size
result = ctypes.c_float()
a_ptr = a.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
b_ptr = b.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
ret = self.lib.metal_dot_product(a_ptr, b_ptr,
ctypes.byref(result), n)
if ret != 0:
raise MetalOpsError("Metal dot product failed")
return result.value
class MetalMatrixOps:
"""
Metal-accelerated matrix operations.
All operations expect 2D NumPy arrays with dtype float32.
Arrays must be C-contiguous.
"""
def __init__(self):
"""Initialize Metal matrix operations."""
if platform.system() != 'Darwin':
raise MetalOpsError("Metal is only available on macOS")
try:
self.lib = ctypes.CDLL('./libmetal_ops.dylib')
except (OSError, FileNotFoundError) as e:
raise MetalOpsError(f"Failed to load Metal library: {e}")
self._setup_function_signatures()
def _setup_function_signatures(self):
"""Setup ctypes function signatures"""
self.lib.metal_matvec.argtypes = [
ctypes.POINTER(ctypes.c_float),
ctypes.POINTER(ctypes.c_float),
ctypes.POINTER(ctypes.c_float),
ctypes.c_int,
ctypes.c_int
]
self.lib.metal_matvec.restype = ctypes.c_int
self.lib.metal_matmul.argtypes = [
ctypes.POINTER(ctypes.c_float),
ctypes.POINTER(ctypes.c_float),
ctypes.POINTER(ctypes.c_float),
ctypes.c_int,
ctypes.c_int,
ctypes.c_int
]
self.lib.metal_matmul.restype = ctypes.c_int
self.lib.metal_transpose.argtypes = [
ctypes.POINTER(ctypes.c_float),
ctypes.POINTER(ctypes.c_float),
ctypes.c_int,
ctypes.c_int
]
self.lib.metal_transpose.restype = ctypes.c_int
self.lib.metal_matrix_add.argtypes = [
ctypes.POINTER(ctypes.c_float),
ctypes.POINTER(ctypes.c_float),
ctypes.POINTER(ctypes.c_float),
ctypes.c_int,
ctypes.c_int
]
self.lib.metal_matrix_add.restype = ctypes.c_int
self.lib.metal_matrix_scale.argtypes = [
ctypes.POINTER(ctypes.c_float),
ctypes.c_float,
ctypes.POINTER(ctypes.c_float),
ctypes.c_int,
ctypes.c_int
]
self.lib.metal_matrix_scale.restype = ctypes.c_int
def _validate_matrix(self, A: np.ndarray):
"""Validate matrix input"""
if not isinstance(A, np.ndarray):
raise TypeError("Input must be a numpy array")
if A.dtype != np.float32:
raise TypeError(f"Input array must be float32, got {A.dtype}")
if A.ndim != 2:
raise ValueError(f"Input must be 2D array, got {A.ndim}D")
if not A.flags['C_CONTIGUOUS']:
raise ValueError("Array must be C-contiguous. Use np.ascontiguousarray() to fix.")
if A.shape[0] > 2**31 - 1 or A.shape[1] > 2**31 - 1:
raise ValueError(f"Matrix dimensions {A.shape} exceed maximum supported size")
total_size = A.shape[0] * A.shape[1]
if total_size > 2**31 - 1:
raise ValueError(f"Matrix total size {total_size} exceeds maximum supported size")
def matvec(self, A: np.ndarray, x: np.ndarray) -> np.ndarray:
"""Matrix-vector multiplication: y = A * x"""
self._validate_matrix(A)
if not isinstance(x, np.ndarray):
raise TypeError("Input must be a numpy array")
if x.ndim != 1:
raise ValueError(f"x must be 1D, got {x.ndim}D")
if A.shape[1] != x.shape[0]:
raise ValueError(f"Matrix columns {A.shape[1]} must match vector length {x.shape[0]}")
if x.dtype != np.float32:
raise TypeError(f"Vector must be float32, got {x.dtype}")
if not x.flags['C_CONTIGUOUS']:
raise ValueError("Vector must be C-contiguous")
m, n = A.shape
y = np.empty(m, dtype=np.float32)
A_ptr = A.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
x_ptr = x.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
y_ptr = y.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
result = self.lib.metal_matvec(A_ptr, x_ptr, y_ptr, m, n)
if result != 0:
raise MetalOpsError("Metal matrix-vector multiplication failed")
return y
def matmul(self, A: np.ndarray, B: np.ndarray) -> np.ndarray:
"""Matrix-matrix multiplication: C = A * B"""
self._validate_matrix(A)
self._validate_matrix(B)
if A.shape[1] != B.shape[0]:
raise ValueError(f"Matrix dimensions incompatible: {A.shape} and {B.shape}")
m, k = A.shape
k2, n = B.shape
C = np.empty((m, n), dtype=np.float32)
A_ptr = A.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
B_ptr = B.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
C_ptr = C.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
result = self.lib.metal_matmul(A_ptr, B_ptr, C_ptr, m, k, n)
if result != 0:
raise MetalOpsError("Metal matrix multiplication failed")
return C
def transpose(self, A: np.ndarray) -> np.ndarray:
"""Matrix transpose: B = A^T"""
self._validate_matrix(A)
height, width = A.shape
B = np.empty((width, height), dtype=np.float32)
A_ptr = A.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
B_ptr = B.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
result = self.lib.metal_transpose(A_ptr, B_ptr, width, height)
if result != 0:
raise MetalOpsError("Metal matrix transpose failed")
return B
def add(self, A: np.ndarray, B: np.ndarray) -> np.ndarray:
"""Element-wise matrix addition: C = A + B"""
self._validate_matrix(A)
self._validate_matrix(B)
if A.shape != B.shape:
raise ValueError(f"Matrix shapes must match: {A.shape} vs {B.shape}")
rows, cols = A.shape
C = np.empty_like(A)
A_ptr = A.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
B_ptr = B.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
C_ptr = C.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
result = self.lib.metal_matrix_add(A_ptr, B_ptr, C_ptr, rows, cols)
if result != 0:
raise MetalOpsError("Metal matrix addition failed")
return C
def scale(self, A: np.ndarray, scalar: float) -> np.ndarray:
"""Scalar multiplication: B = A * scalar"""
self._validate_matrix(A)
rows, cols = A.shape
B = np.empty_like(A)
A_ptr = A.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
B_ptr = B.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
result = self.lib.metal_matrix_scale(A_ptr, float(scalar), B_ptr, rows, cols)
if result != 0:
raise MetalOpsError("Metal matrix scaling failed")
return B
TESTING THE IMPLEMENTATION
Here is a comprehensive test suite:
import numpy as np
from metal_ops import is_metal_available, MetalVectorOps, MetalMatrixOps, MetalOpsError
def test_vector_add():
print("Testing vector addition...")
if not is_metal_available():
print(" SKIPPED: Metal not available")
return
ops = MetalVectorOps()
a = np.random.randn(1000).astype(np.float32)
b = np.random.randn(1000).astype(np.float32)
metal_result = ops.add(a, b)
numpy_result = a + b
np.testing.assert_allclose(metal_result, numpy_result, rtol=1e-5)
print(" PASSED")
def test_vector_subtract():
print("Testing vector subtraction...")
if not is_metal_available():
print(" SKIPPED: Metal not available")
return
ops = MetalVectorOps()
a = np.random.randn(1000).astype(np.float32)
b = np.random.randn(1000).astype(np.float32)
metal_result = ops.subtract(a, b)
numpy_result = a - b
np.testing.assert_allclose(metal_result, numpy_result, rtol=1e-5)
print(" PASSED")
def test_vector_multiply():
print("Testing vector multiplication...")
if not is_metal_available():
print(" SKIPPED: Metal not available")
return
ops = MetalVectorOps()
a = np.random.randn(1000).astype(np.float32)
b = np.random.randn(1000).astype(np.float32)
metal_result = ops.multiply(a, b)
numpy_result = a * b
np.testing.assert_allclose(metal_result, numpy_result, rtol=1e-5)
print(" PASSED")
def test_vector_scale():
print("Testing vector scaling...")
if not is_metal_available():
print(" SKIPPED: Metal not available")
return
ops = MetalVectorOps()
a = np.random.randn(1000).astype(np.float32)
scalar = 2.5
metal_result = ops.scale(a, scalar)
numpy_result = a * scalar
np.testing.assert_allclose(metal_result, numpy_result, rtol=1e-5)
print(" PASSED")
def test_dot_product():
print("Testing dot product...")
if not is_metal_available():
print(" SKIPPED: Metal not available")
return
ops = MetalVectorOps()
a = np.random.randn(1000).astype(np.float32)
b = np.random.randn(1000).astype(np.float32)
metal_result = ops.dot(a, b)
numpy_result = np.dot(a, b)
np.testing.assert_allclose(metal_result, numpy_result, rtol=1e-4)
print(" PASSED")
def test_matvec():
print("Testing matrix-vector multiplication...")
if not is_metal_available():
print(" SKIPPED: Metal not available")
return
ops = MetalMatrixOps()
A = np.random.randn(100, 50).astype(np.float32)
x = np.random.randn(50).astype(np.float32)
metal_result = ops.matvec(A, x)
numpy_result = np.matmul(A, x)
np.testing.assert_allclose(metal_result, numpy_result, rtol=1e-4)
print(" PASSED")
def test_matmul():
print("Testing matrix-matrix multiplication...")
if not is_metal_available():
print(" SKIPPED: Metal not available")
return
ops = MetalMatrixOps()
A = np.random.randn(100, 50).astype(np.float32)
B = np.random.randn(50, 80).astype(np.float32)
metal_result = ops.matmul(A, B)
numpy_result = np.matmul(A, B)
np.testing.assert_allclose(metal_result, numpy_result, rtol=1e-3)
print(" PASSED")
def test_transpose():
print("Testing matrix transpose...")
if not is_metal_available():
print(" SKIPPED: Metal not available")
return
ops = MetalMatrixOps()
A = np.random.randn(100, 80).astype(np.float32)
metal_result = ops.transpose(A)
numpy_result = A.T
np.testing.assert_allclose(metal_result, numpy_result, rtol=1e-5)
print(" PASSED")
def test_error_handling():
print("Testing error handling...")
if not is_metal_available():
print(" SKIPPED: Metal not available")
return
ops = MetalVectorOps()
try:
a = np.array([1, 2, 3], dtype=np.float64)
b = np.array([4, 5, 6], dtype=np.float64)
ops.add(a, b)
print(" FAILED: Should have raised TypeError")
return
except TypeError as e:
assert "float32" in str(e)
try:
a = np.array([1, 2, 3], dtype=np.float32)
b = np.array([4, 5], dtype=np.float32)
ops.add(a, b)
print(" FAILED: Should have raised ValueError")
return
except ValueError as e:
assert "shape" in str(e)
try:
a = np.array([[1, 2], [3, 4]], dtype=np.float32)
b = np.array([[5, 6], [7, 8]], dtype=np.float32)
ops.add(a, b)
print(" FAILED: Should have raised ValueError")
return
except ValueError as e:
assert "1D" in str(e)
print(" PASSED")
def test_edge_cases():
print("Testing edge cases...")
if not is_metal_available():
print(" SKIPPED: Metal not available")
return
ops = MetalVectorOps()
a = np.array([1.0], dtype=np.float32)
b = np.array([2.0], dtype=np.float32)
c = ops.add(a, b)
assert c[0] == 3.0
a = np.array([], dtype=np.float32)
b = np.array([], dtype=np.float32)
c = ops.add(a, b)
assert len(c) == 0
print(" PASSED")
def run_all_tests():
print("=" * 60)
print("Metal Operations Test Suite")
print("=" * 60)
test_vector_add()
test_vector_subtract()
test_vector_multiply()
test_vector_scale()
test_dot_product()
test_matvec()
test_matmul()
test_transpose()
test_error_handling()
test_edge_cases()
print("=" * 60)
print("All tests completed")
print("=" * 60)
if __name__ == '__main__':
run_all_tests()
BENCHMARKING PERFORMANCE
Here is a benchmark suite:
import numpy as np
import time
from metal_ops import is_metal_available, MetalVectorOps, MetalMatrixOps
def benchmark_vector_ops(sizes):
print("\nVector Operations Benchmark")
print("-" * 60)
if not is_metal_available():
print("Metal not available")
return
ops = MetalVectorOps()
for size in sizes:
a = np.random.randn(size).astype(np.float32)
b = np.random.randn(size).astype(np.float32)
_ = ops.add(a, b)
start = time.perf_counter()
for _ in range(10):
_ = ops.add(a, b)
metal_time = (time.perf_counter() - start) / 10
start = time.perf_counter()
for _ in range(10):
_ = a + b
numpy_time = (time.perf_counter() - start) / 10
speedup = numpy_time / metal_time if metal_time > 0 else 0
print(f"Size {size:8d}: Metal {metal_time*1000:6.2f} ms "
f"NumPy {numpy_time*1000:6.2f} ms "
f"Speedup {speedup:5.2f}x")
def benchmark_matmul(sizes):
print("\nMatrix Multiplication Benchmark")
print("-" * 60)
if not is_metal_available():
print("Metal not available")
return
ops = MetalMatrixOps()
for size in sizes:
A = np.random.randn(size, size).astype(np.float32)
B = np.random.randn(size, size).astype(np.float32)
_ = ops.matmul(A, B)
start = time.perf_counter()
for _ in range(5):
_ = ops.matmul(A, B)
metal_time = (time.perf_counter() - start) / 5
start = time.perf_counter()
for _ in range(5):
_ = np.matmul(A, B)
numpy_time = (time.perf_counter() - start) / 5
speedup = numpy_time / metal_time if metal_time > 0 else 0
print(f"Size {size:4d}x{size:4d}: Metal {metal_time*1000:7.2f} ms "
f"NumPy {numpy_time*1000:7.2f} ms "
f"Speedup {speedup:5.2f}x")
def run_all_benchmarks():
print("=" * 60)
print("Metal Operations Benchmark Suite")
print("=" * 60)
vector_sizes = [1000, 10000, 100000, 1000000]
matrix_sizes = [64, 128, 256, 512, 1024]
benchmark_vector_ops(vector_sizes)
benchmark_matmul(matrix_sizes)
print("\n" + "=" * 60)
print("Benchmarks completed")
print("=" * 60)
if __name__ == '__main__':
run_all_benchmarks()
CONCLUSION
This comprehensive tutorial has provided a complete guide to using Python's Foreign Function Interface to call GPU-accelerated code running on Apple's Metal framework. We explored the Metal architecture, implemented a full library of matrix and vector operations, and created a production-ready Python interface that provides clean, NumPy-like functionality while leveraging GPU acceleration.
The implementation demonstrates best practices including comprehensive error handling, proper resource management, thorough input validation, detailed documentation, and extensive testing. This foundation can serve as a starting point for more advanced GPU computing projects on Apple platforms.
No comments:
Post a Comment