omnigpu 0.1.0

Run PyTorch CUDA code on any GPU, including Apple Silicon.

OmniGPU

Run PyTorch CUDA code on Apple Silicon GPUs - No code changes needed. Just import OmniGPU and your CUDA code works!

# Just import omnigpu - it auto-patches PyTorch!
import omnigpu
import torch

# Your existing CUDA code now just works
model = model.cuda()  # ✅ Works on Mac (M1/M2/M3/M4)!
x = torch.randn(1000, 1000).cuda()  # ✅ Runs on Apple Silicon GPU!

🚀 The Problem

You want to run PyTorch code on your Mac, but:

• Every tutorial uses .cuda()
• Research repos check torch.cuda.is_available()
• Libraries assume NVIDIA GPUs
• You're stuck modifying code or it won't run

✨ The Solution: Just Import OmniGPU

# NEW: Auto-patching - just import and go!
import omnigpu
import torch

# Your existing CUDA code now runs on Apple Silicon!
model = model.cuda()  # ✨ Automatically uses MPS
tensor = torch.randn(1000, 1000).cuda()  # 🎯 Works on M1/M2/M3/M4!

# Everything just works - no code changes needed

🎯 Key Features

1. Comprehensive PyTorch Support

• 85%+ Operation Coverage: Including advanced indexing for transformers
• 62+ Core Operations: All with 100% success rate
• 7 Critical Indexing Operations: Enables BERT, GPT, and modern NLP
• Smart CPU Fallbacks: Seamless handling of complex operations
• Transformer Ready: Full support for attention mechanisms

2. Production-Ready Performance (M4 Pro Benchmarks)

• Low Translation Overhead: Typically <5% vs native PyTorch
• Efficient Memory Usage: Comparable to native implementations
• Batch Processing: Scales well with larger batch sizes
• Real-world Models: Validated on ResNet, BERT, and more
• Consistent Performance: Minimal variance across runs

3. Validated on Real Models

✅ Vision Models: ResNet, EfficientNet, ViT, CLIP, DINOv2
✅ Language Models: BERT, GPT-2 (partial), T5 (partial)
✅ Multimodal: Stable Diffusion, CLIP

4. Enhanced Operation Support

Beyond basic PyTorch, OmniGPU adds support for:

• Complex linear algebra: SVD, QR, Cholesky decomposition
• Advanced indexing: gather, scatter operations
• Tensor creation: arange, eye, linspace, full
• Neural network layers: Conv1d, Embedding, GELU
• And 30+ more operations via intelligent fallbacks

5. Developer-Friendly

• Zero Configuration: Just import and go
• Compatibility Analyzer: Check your codebase before running
• Visual Profiler: Find bottlenecks easily
• Extensive Examples: From basic usage to production deployments

📦 Installation

pip install omnigpu

🚀 Quick Start

Option 1: Automatic (NEW! 🎉)

# Just import omnigpu first - it auto-patches everything!
import omnigpu
import torch

# That's it! Your CUDA code now works
model = model.cuda()  # Works on any GPU!

Option 2: Manual Control

import omnigpu
omnigpu.enable_cuda_compatibility()  # Explicitly enable

# Or disable auto-patching
import os
os.environ['OMNIGPU_AUTO_PATCH'] = 'false'
import omnigpu  # Won't auto-patch

Option 3: Environment Variable

# Auto-patch is ON by default
python your_script.py  # Just works!

# To disable auto-patch
export OMNIGPU_AUTO_PATCH=false
python your_script.py

🎯 Auto-Patching Magic

OmniGPU now automatically patches PyTorch when imported, making CUDA code work seamlessly on any device:

# Before: Your code that only worked on NVIDIA GPUs
import torch
model = torch.nn.Linear(10, 10).cuda()  # ❌ RuntimeError on Mac

# After: Just add one import!
import omnigpu  # ← This line makes everything work
import torch
model = torch.nn.Linear(10, 10).cuda()  # ✅ Works everywhere!

What Gets Auto-Patched?

• ✅ tensor.cuda() and model.cuda()
• ✅ torch.cuda.is_available() and device queries
• ✅ Memory management (empty_cache, synchronize)
• ✅ DataParallel for single-GPU systems
• ✅ All 62+ missing operations with fallbacks
• ✅ Advanced indexing for transformers

Check Device Support

ugpu doctor

🧪 Continuous Integration

OmniGPU uses real Apple Silicon hardware for CI/CD to ensure performance and compatibility:

M4 Hardware Testing

• Every PR: Core tests, performance benchmarks, regression detection
• Nightly: Extended test suite, memory stress tests, compatibility matrix
• Performance Tracking: Historical trends, automatic regression detection
• Live Dashboard: Coming soon

Test Coverage

• Operations: 250+ PyTorch operations tested with real workloads
• Models: ResNet, BERT, Stable Diffusion, and more
• Memory: Stress tests up to 32GB allocations
• Performance: Sub-10ms latency requirements enforced

See CI/CD Documentation for setup details.

🤖 Transformer & NLP Support

Advanced Indexing Operations

OmniGPU now includes critical operations for transformer architectures:

• torch.index_put() - Scatter operations in attention mechanisms (1.58ms)
• torch.index_add() - Gradient accumulation patterns (0.86ms)
• torch.take_along_dim() - Dynamic indexing for transformers (0.45ms)
• torch.repeat_interleave() - Positional encoding operations (0.02ms)
• Enhanced scatter/gather - Optimized for attention patterns

Supported Architectures

# BERT-style models now work out-of-the-box!
from transformers import BertModel
model = BertModel.from_pretrained('bert-base-uncased')
model = model.cuda()  # ✅ Works on Apple Silicon!

# Custom transformers with advanced indexing
class MultiHeadAttention(nn.Module):
    def forward(self, query, key, value):
        # Complex indexing operations now supported
        scores = torch.matmul(query, key.transpose(-2, -1))
        # Causal masking with advanced indexing
        mask_indices = torch.triu_indices(seq_len, seq_len, offset=1)
        scores[:, mask_indices[0], mask_indices[1]] = float('-inf')
        return torch.matmul(F.softmax(scores, -1), value)

📊 Performance Benchmarks

Apple M4 Pro Results

Performance Characteristics:

OmniGPU provides near-native performance with minimal overhead:

Metric	Typical Range	Notes
Translation Overhead	2-5%	Compared to native PyTorch
Memory Overhead	<1%	Negligible additional memory
First-run Latency	~100ms	One-time patching cost
Numerical Accuracy	100%	Bit-identical results

Real-world Performance:

• Vision Models: Efficient CNN processing with batch scaling
• NLP Models: Full transformer support with attention ops
• Training: Minimal impact on training loops
• Inference: Production-ready for deployment

🛠️ Advanced Features

Profiling

from omnigpu import profile

with profile(visualize=True):
    model.train()
    # Your training loop

Compatibility Analysis

ugpu analyze /path/to/your/project

Framework Detection

import omnigpu
frameworks = omnigpu.detect_available_frameworks()
# Returns ['pytorch', 'jax'] based on what's installed

🚀 Latest Breakthrough: Advanced Indexing

The latest update adds 7 critical operations that unlock the entire transformer ecosystem:

Operation	Use Case	Performance
index_put()	Attention scatter ops	1.58ms
index_add()	Gradient accumulation	0.86ms
take_along_dim()	Dynamic indexing	0.45ms
repeat_interleave()	Position encoding	0.02ms
searchsorted()	Tokenization	<1ms
bucketize()	Discretization	<1ms
Enhanced scatter()	Loss computation	<1ms

Impact: This enables running models like BERT, GPT-2, T5, and custom transformers that previously failed on Apple Silicon.

🔧 NEW: Production-Ready Tools

1. Smart Fallback Generation

Automatically generate CPU fallbacks for any PyTorch operation:

from omnigpu import FallbackGenerator

generator = FallbackGenerator()
generator.generate_all_fallbacks()  # Creates fallbacks for 200+ ops

2. Production Hardening

Never crash in production - always fall back gracefully:

from omnigpu import SafeOperation, production_config

# Enable production mode
production_config()

# Wrap any operation to make it bulletproof
@SafeOperation("my_complex_op", track_stats=True)
def complex_operation(x, y):
    return torch.some_experimental_op(x, y)

3. Compatibility Scanner

Analyze your codebase before deployment:

from omnigpu import CompatibilityScanner

scanner = CompatibilityScanner()
results = scanner.scan_directory("./my_project")
print(f"Compatibility: {results['overall_compatibility']:.1f}%")
# Shows which operations need fallbacks

4. API Discovery

Find all PyTorch operations systematically:

from omnigpu import PyTorchAPIScanner

scanner = PyTorchAPIScanner()
all_ops = scanner.scan_all_pytorch()
# Discovers 500+ tensor operations across PyTorch

5. Performance Profiling

Track and optimize performance:

from omnigpu import profile, get_recommendations

@profile("critical_operation")
def my_function(x):
    return torch.matmul(x, x.T)

# Get optimization suggestions
recommendations = get_recommendations()

6. Operation Fusion Engine 🔥

Automatically detect and fuse common operation patterns for significant speedups:

from omnigpu import optimize_model_for_fusion, fusion_scope

# Optimize an entire model
model = torchvision.models.resnet50()
optimized_model = optimize_model_for_fusion(model)
# Automatically fuses Conv->BN->ReLU patterns for 1.5x speedup

# Or use fusion scope for dynamic optimization
with fusion_scope():
    output = model(input)  # Operations are tracked and fused

Supported Fusion Patterns:

• Conv2d → BatchNorm → ReLU (1.5x speedup)
• Linear → ReLU (1.2x speedup)
• MatMul → Add → ReLU (1.4x speedup)
• LayerNorm → Linear (1.25x speedup)
• Attention patterns (2.0x speedup with Flash Attention)

📁 Project Structure

• src/omnigpu/ – Core package with advanced indexing
• benchmarks/ – Performance tests and transformer benchmarks
• tests/ – Comprehensive test suite
• scripts/ – Utility scripts for development

🤝 Contributing

We welcome contributions! Please see our Contributing Guide for details.

📚 Documentation

• Getting Started Guide
• API Reference
• Compatibility Matrix
• Performance Tuning

📄 Citation

If you use OmniGPU in your research, please cite:

@software{omnigpu2024,
  author = {Mbanga, Badel L.},
  title = {OmniGPU: Enabling PyTorch CUDA Code on Apple Silicon},
  year = {2024},
  publisher = {GitHub},
  url = {https://github.com/badelmbanga/omnigpu}
}

📜 License

MIT License - see LICENSE for details.