Oven-Tensor

A PyTorch-style tensor library with GPU acceleration using CUDA kernels compiled by oven-compiler.

✨ Features

• 🚀 PyTorch-like Interface: Familiar tensor operations with .cpu(), .gpu(), @ operator
• ⚡ Automatic Kernel Compilation: Python kernels compiled to PTX using oven-compiler
• 💾 Smart Caching: Compiled kernels cached for fast subsequent loads
• 🔄 CPU/GPU Hybrid: Seamless switching between NumPy (CPU) and CUDA (GPU)
• 🔧 Custom Kernels: Easy to add and execute custom CUDA kernels
• 🎯 Dynamic Registration: Register kernels at runtime from code or files

📦 Installation

pip install oven-tensor

Requirements:

• Python 3.7+
• CUDA-capable GPU
• oven-compiler in PATH
• PyCUDA

🚀 Quick Start

import oven_tensor as ot

# Create tensors
x = ot.tensor([1.0, 2.0, 3.0, 4.0])
y = ot.tensor([2.0, 3.0, 4.0, 5.0])

# CPU operations (NumPy backend)
z_cpu = x + y
print(z_cpu)  # Tensor([3. 5. 7. 9.], device=cpu)

# GPU operations (CUDA kernels)
x_gpu = x.gpu()
y_gpu = y.gpu()
z_gpu = x_gpu + y_gpu
print(z_gpu.cpu())  # Tensor([3. 5. 7. 9.], device=cpu)

# Matrix multiplication
A = ot.tensor([[1.0, 2.0], [3.0, 4.0]])
B = ot.tensor([[5.0, 6.0], [7.0, 8.0]])
C = A @ B
print(C)  # Tensor([[19. 22.], [43. 50.]], device=cpu)

📚 API Reference

Tensor Creation

ot.tensor([1, 2, 3])      # From data
ot.zeros((2, 3))          # Zero tensor  
ot.ones((2, 3))           # Ones tensor
ot.randn((2, 3))          # Random normal
ot.linspace(0, 10, 5)     # Evenly spaced values

Basic Operations

# Unary operations
x.sigmoid(), x.exp(), x.sqrt(), x.abs()
x.sin(), x.cos(), x.log(), x.tanh()

# Binary operations  
x + y, x - y, x * y, x / y, x ** y, x % y

# Matrix operations
A @ B, A.matmul(B), ot.matmul(A, B)

Device Management

x.gpu()                   # Move to GPU
x.cpu()                   # Move to CPU
x.to(ot.device('gpu'))    # Explicit device transfer

🔧 Custom Kernels

Basic Usage

# Execute built-in custom kernels
x = ot.tensor([1, 2, 3, 4]).gpu()
result = ot.zeros((4,)).gpu()

ot.execute_kernel("vector_scale", x, result, scale=2.0)
ot.execute_kernel("vector_relu", x, result)

Writing Custom Kernels

Add kernels in oven_tensor/kernels/:

# my_kernels.py
import oven.language as ol

def my_kernel(x_ptr: ol.ptr, y_ptr: ol.ptr, factor: float):
    """Scale vector by factor"""
    idx = ol.get_global_id()
    x_val = ol.load(x_ptr, idx)
    y_val = x_val * factor
    ol.store(y_val, y_ptr, idx)

Dynamic Registration

Register kernels at runtime:

# From code string
kernel_code = '''
import oven.language as ol

def runtime_kernel(x_ptr: ol.ptr, y_ptr: ol.ptr, factor: float):
    idx = ol.get_global_id()
    x_val = ol.load(x_ptr, idx)
    y_val = x_val * factor
    ol.store(y_val, y_ptr, idx)
'''

functions = ot.register_kernel_from_code(kernel_code, "my_module", ["runtime_kernel"])
ot.execute_kernel("runtime_kernel", x, result, factor=3.0)

# From file
functions = ot.register_kernel_from_file("my_kernels.py")

# Cleanup
ot.unregister_kernel("runtime_kernel")

🎛️ Cache Management

# Command-line tool
oven-tensor-cache list    # List cached functions
oven-tensor-cache clear   # Clear cache
oven-tensor-cache info    # Show cache info

# Python API
ot.clear_kernel_cache()           # Clear cache
ot.reload_kernels()               # Reload kernels
ot.list_available_functions()     # List functions

🧪 Testing

# Run all tests
./scripts/run_tests.sh

# Specific categories
pytest tests/ -m "not gpu"       # Skip GPU tests
pytest tests/ -m "not slow"      # Skip slow tests
pytest tests/ --cov=oven_tensor  # With coverage

📄 License

MIT License