Metadata-Version: 2.4
Name: computeml
Version: 0.2.0
Summary: ML inference on any GPU through Vulkan compute shaders, from NumPy
Author: Jorge Cardona
License: MIT
Project-URL: Homepage, https://gitlab.com/cardona-oss/compute-ml
Project-URL: Issues, https://gitlab.com/cardona-oss/compute-ml/-/tree/main/docs/issues
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: numpy>=1.22
Provides-Extra: test
Requires-Dist: pytest; extra == "test"

# computeml

Run ML inference on any GPU from Python, through Vulkan compute shaders. You work with NumPy
arrays. The wheel ships the compiled library, so you only need NumPy and a Vulkan driver.

## Install

```bash
pip install computeml
```

```bash
uv add computeml
```

Requirements:

- Linux x86_64 (glibc 2.27 or newer) and Python 3.10 or newer.
- A Vulkan driver. NVIDIA, AMD and Intel GPU drivers include one. With no GPU, Mesa lavapipe
  runs on the CPU: `apt install mesa-vulkan-drivers`.

## Quickstart

```python
import numpy as np
import computeml as cml

ctx = cml.Context()
print(ctx.device_name)

a = ctx.asarray(np.random.rand(512, 256))
b = ctx.asarray(np.random.rand(256, 128))
c = cml.ops.relu(a @ b) * 2.0
print(c.shape, c.numpy()[0, :4])
```

`ctx.asarray` copies an array to the GPU. Ops queue work on the GPU. `.numpy()` waits for the
work and copies the result back.

## Coming from NumPy, PyTorch or CuPy

| You know | In computeml |
|---|---|
| `cp.asarray(x)`, `torch.as_tensor(x)` | `ctx.asarray(x)` |
| `np.zeros`, `ones`, `full`, `empty` | `ctx.zeros(shape)`, `ctx.ones`, `ctx.full`, `ctx.empty` |
| `a + b`, `a * 2`, `a @ b`, `-a` | the same |
| `.shape`, `.dtype`, `.ndim`, `.size` | the same |
| `t.cpu().numpy()`, `cp.asnumpy(t)` | `t.numpy()` |
| `torch.nn.functional.layer_norm` | `cml.ops.layer_norm` |
| `torch.nn.Linear(16, 32)` | `cml.nn.Linear(ctx, 16, 32)` |

Differences you will meet:

- No broadcasting. Shapes must match element for element, or the op raises `ValueError`.
- float32 only for inputs and outputs. `asarray` converts ints and float64 to float32.
- `tensor + number` is not supported yet. `tensor * number` and `tensor / number` work.
- `np.asarray(tensor)` raises `TypeError`. Call `.numpy()`, so every GPU-to-CPU copy is visible.
- Every layer takes the `Context` first. There is no default device.

## Layers

```python
import numpy as np
import computeml as cml

ctx = cml.Context()
model = cml.nn.Sequential(
    cml.nn.Linear(ctx, 16, 32, activation="gelu"),
    cml.nn.Linear(ctx, 32, 8),
    cml.nn.LayerNorm(ctx, 8),
)
x = ctx.asarray(np.random.rand(4, 16))
print(model(x).numpy().shape)
```

`Linear` weights use the PyTorch layout `(out, in)`. Load trained weights with
`cml.nn.Linear.from_numpy(ctx, weight, bias)`.

## Fused ops

One kernel does the work of several ops, so the data is read from GPU memory once.

```python
import numpy as np
import computeml as cml

ctx = cml.Context()
x = ctx.asarray(np.random.rand(8, 64))
w = ctx.asarray(np.random.rand(128, 64))
b = ctx.asarray(np.random.rand(128))
y = cml.ops.linear(x, w, b, activation="gelu")  # gelu(x @ w.T + b), one kernel
print(y.shape)
```

`cml.ops` also has softmax, layer and RMS norm, flash attention (multi-head, causal), RoPE,
embeddings, permute, concat and reductions. See
[`ops.py`](https://gitlab.com/cardona-oss/compute-ml/-/blob/main/python/computeml/ops.py).

## Capture and replay

`ctx.capture` records a function's GPU work on the first call. Later calls resubmit the
recorded work with new input values, with no Python per op.

```python
import numpy as np
import computeml as cml

ctx = cml.Context()
w = ctx.asarray(np.random.rand(64, 64))

@ctx.capture
def step(x):
    return cml.ops.relu(x @ w)

for _ in range(3):
    y = step(np.random.rand(32, 64))
    print(y.numpy().sum())
print(step.record_count)  # 1: recorded once, replayed twice
```

The output is the same tensor on every call, so read it before the next call. Any other op
on the context in between makes the next call record again.

## Quantized weights

```python
import numpy as np
import computeml as cml

ctx = cml.Context()
x = ctx.asarray(np.random.rand(16, 256))
w = ctx.asarray(np.random.rand(256, 64)).quantize("Q8_0")
print((x @ w).numpy().shape)  # w is dequantized inside the matmul kernel
```

## Errors and logs

- `cml.NoDeviceError`: no Vulkan driver or GPU. The message says what to install.
- `ImportError` on `import computeml`: the Vulkan loader `libvulkan.so.1` is missing.
- `ValueError` / `TypeError`: a wrong shape or dtype, raised before any GPU work.
- Set `CML_LOG=1` to print the chosen GPU and kernel class to stderr.
- Set `CML_GPU_INDEX=1` to use the second GPU.

## Build from source

Set `COMPUTEML_LIB` to use your own `libcomputeml.so`. To build it and run the tests from a
checkout of the repo:

```bash
cmake -B build -S . && cmake --build build --target computeml_shared
cd python && uv run --with pytest python -m pytest -q
```

A BERT-tiny example lives in
[`examples/python/bert_tiny.py`](https://gitlab.com/cardona-oss/compute-ml/-/blob/main/examples/python/bert_tiny.py).
