Metadata-Version: 2.3
Name: spork-metal
Version: 0.3.0
Summary: A tracing Python DSL for writing Metal compute kernels for Apple GPUs.
Keywords: metal,apple,gpu,kernel,dsl,tracing,compute,mlx,machine-learning
Author: Sasha Krassovsky
Author-email: Sasha Krassovsky <krassovskysasha@gmail.com>
License: Copyright 2026 Sasha Krassovsky
         
         Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
         
         The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
         
         THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
         
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: MacOS
Classifier: Operating System :: MacOS :: MacOS X
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Software Development :: Code Generators
Classifier: Topic :: Software Development :: Compilers
Classifier: Typing :: Typed
Requires-Dist: numpy>=2.0
Requires-Dist: pyobjc>=11.1
Requires-Python: >=3.13
Project-URL: Homepage, https://github.com/save-buffer/spork
Project-URL: Issues, https://github.com/save-buffer/spork/issues
Project-URL: Repository, https://github.com/save-buffer/spork
Description-Content-Type: text/markdown

# Spork: A Tracing DSL for Apple GPUs

Spork is a DSL for writing kernels for Apple GPUs. You can think of it as a Python wrapper
on top of Metal, Apple's GPU language. This makes development more convenient, since it makes
the kernels live in the same source language as most ML code, and makes it easy to verify 
correctness against NumPy. 

Unlike Triton, Spork is a tracing DSL rather than a parsing-based DSL. This means that a Spork
program is fundamentally a Python program that creates a Metal program. Spork can use any Python
libraries to aid in metaprogramming. 

## Installation

Spork runs on Apple Silicon (or any Mac with Metal). Install it from PyPI:

```bash
pip install spork-metal
```

Or with [uv](https://docs.astral.sh/uv/):

```bash
uv add spork-metal
```

The package is published as `spork-metal` on PyPI; the import name is `spork`:

```python
import spork as sk
```

## Example: Matrix Addition
One of the most basic kernels is a matrix-addition kernel. In Numpy, you could write
```python
shape = (1024, 1024)
A = np.random.randn(*shape).astype(np.float32)
B = np.random.randn(*shape).astype(np.float32)
out = A + B
```

A Metal kernel to perform this operation would look like this:
```cpp
#include <metal_stdlib>
using namespace metal;

kernel void matrix_add(
    device float *out [[buffer(0)]],
    device const float *A [[buffer(1)]],
    device const float *B [[buffer(2)]],
    uint index [[thread_position_in_grid]])
{
    out[index] = A[index] + B[index];
}
```

Notice that it exists in a separate source file and requires special code to compile, link to,
and invoke. An equivalent Spork kernel looks like this:
```python
@sk.jit
def matrix_add(
    out   : sk.DevicePointer[sk.dt.float32],
    A     : sk.DevicePointer[sk.dt.float32],
    B     : sk.DevicePointer[sk.dt.float32],
    index : sk.Uint[sk.ThreadPositionInGrid],
):
    out[index] = A[index] + B[index]
```

Notice that we have a direct correspondance here between the Spork kernel and the Metal kernel. We
declare the types of our inputs, and we can take the attribute parameters that we take in Metal. 


Then to actually invoke and run kernel, we simply call it from Python with Numpy arrays, and 
like magic it runs your kernel!

```python
matrix_add[
    (int(np.prod(shape)) // 128, 1, 1),
    (128, 1, 1),
](
    C,
    A,
    B,
)
```
Notice that in the brackets, we provide two parameters. The first is the Grid size, and the second
is the Warpgroup size. In the parentheses, we supply the Numpy tensors we wish to use. Notice
that here we're taking a pointer to `C`, which is allocated via Numpy but written to from the 
Spork kernel.

