oxgpupy

Python bindings for oxgpu, a lightweight GPU compute library built on wgpu.

Installation

pip install oxgpupy

Features

• 🐍 Pythonic API: Easy-to-use Python interface for GPU computing
• 🚀 High Performance: Zero-copy data transfer where possible
• ⚡ Async Support: Full async/await support with asyncio
• 🔧 Type Safety: Leverages Python type hints
• 📝 WGSL Shaders: Write compute shaders in WebGPU Shading Language

Quick Start

import asyncio
import oxgpupy

async def main():
    # Create GPU context
    ctx = await oxgpupy.Context.new()

    # Create buffers
    data = [1.0, 2.0, 3.0, 4.0, 5.0]
    buffer = await oxgpupy.Buffer.from_list(ctx, data)

    # Read data back
    result = await buffer.to_list(ctx)
    print(f"Result: {result}")

asyncio.run(main())

Compute Shader Example

import asyncio
import oxgpupy

async def saxpy_example():
    # Initialize
    ctx = await oxgpupy.Context.new()

    # Create buffers
    x_data = [1.0, 2.0, 3.0, 4.0, 5.0]
    y_data = [2.0, 4.0, 6.0, 8.0, 10.0]

    x_buf = await oxgpupy.Buffer.from_list(ctx, x_data)
    y_buf = await oxgpupy.Buffer.from_list(ctx, y_data)

    # WGSL shader (SAXPY: y = x + y)
    shader = """
    @group(0) @binding(0) var<storage, read> x: array<f32>;
    @group(0) @binding(1) var<storage, read_write> y: array<f32>;

    @compute @workgroup_size(64)
    fn main(@builtin(global_invocation_id) id: vec3<u32>) {
        y[id.x] = x[id.x] + y[id.x];
    }
    """

    # Build kernel
    builder = oxgpupy.ComputeKernelBuilder.new(shader, "main")
    builder.bind(oxgpupy.KernelBinding.new(0, oxgpupy.BindingType.storage(read_only=True)))
    builder.bind(oxgpupy.KernelBinding.new(1, oxgpupy.BindingType.storage(read_only=False)))

    kernel = await builder.build(ctx)

    # Execute kernel
    kernel.run(ctx, (1, 1, 1), [x_buf, y_buf])

    # Get results
    result = await y_buf.to_list(ctx)
    print(f"Result: {result}")
    # Output: Result: [3.0, 6.0, 9.0, 12.0, 15.0]

asyncio.run(saxpy_example())

API Reference

Core Classes

Context

GPU context managing device and queue.

ctx = await Context.new()

Buffer

Typed GPU buffer for f32 data.

# Create from Python list
buf = await Buffer.from_list(ctx, [1.0, 2.0, 3.0])

# Create with custom usage
buf = await Buffer.new(ctx, data, BufferUsage.STORAGE)

# Create zero-initialized
buf = await Buffer.zeros(ctx, 100)

# Read data
data = await buf.to_list(ctx)

# Write data
buf.write(ctx, [4.0, 5.0, 6.0])

BufferUsage

Buffer usage flags.

• BufferUsage.STORAGE - Storage buffer
• BufferUsage.UNIFORM - Uniform buffer
• BufferUsage.COPY_SRC - Copy source
• BufferUsage.COPY_DST - Copy destination
• BufferUsage.MAP_READ - CPU readable
• BufferUsage.MAP_WRITE - CPU writable

BindingType

Shader binding types.

# Storage buffer (read-only or read-write)
ty = BindingType.storage(read_only=True)

# Uniform buffer
ty = BindingType.uniform()

KernelBinding

Binding configuration for shaders.

binding = KernelBinding.new(0, BindingType.storage(read_only=True))

ComputeKernelBuilder

Builder for creating compute kernels.

builder = ComputeKernelBuilder.new(shader_source, "main")
builder.bind(KernelBinding.new(0, BindingType.storage(read_only=True)))
kernel = await builder.build(ctx)

ComputeKernel

Compiled compute shader.

# Execute kernel
# groups: (x, y, z) workgroup dimensions
# buffers: list of Buffer objects
kernel.run(ctx, (1, 1, 1), [buffer1, buffer2])

Requirements

• Python 3.7+
• A GPU with WebGPU support

Development

Building from Source

# Clone repository
git clone https://github.com/vulkanic-labs/oxgpupy
cd oxgpupy

# Create virtual environment
python -m venv .venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate

# Install maturin
pip install maturin

# Build and install
maturin develop

# Run tests
python test.py

License

This project is licensed under the MIT License - see the LICENSE file for details.

Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

Links

• GitHub Repository
• oxgpu (Rust crate)