torch_tvarant

Open-source PyTorch backend for the Tvarant RISC-V SIMT GPGPU (Xilinx Alveo U55C).

Develop on a CPU simulator, deploy through OpenCL → POCL → libtvarant when hardware is available. Includes a JIT graph compiler for LLM inference fusion.

📖 Documentation · 🐛 Issues · 🗺️ Roadmap · 🤝 Contributing

Quick start

import torch
import torch_tvarant

x = torch.randn(4, 8, device="tvarant")
w = torch.randn(8, 8, device="tvarant")
y = torch.nn.functional.relu(x @ w)
assert y.device.type == "tvarant"

pip install torch pytest ninja
pip install -e .
pytest tests/ -v

See Installation and Getting Started. API: docs/api.

Features

Feature	Description
Device backend	torch.device("tvarant") via PrivateUse1
CPU simulator	Host-accessible memory for fast iteration
OpenCL runtime	Same API on POCL / FPGA path
LLM ops	silu, softmax, bmm, layer_norm, embedding, …
Fused GEMM	Single kernel for matmul + bias + relu/silu
JIT compiler	FX fusion + OpenCL kernel codegen for pointwise chains
torch.compile	Registered backend="tvarant"

JIT compiler

import torch
import torch.nn as nn
import torch_tvarant

model = nn.Sequential(nn.Linear(768, 768), nn.ReLU()).to("tvarant")
compiled = torch_tvarant.compiler.compile(model)
y = compiled(torch.randn(4, 768, device="tvarant"))

Fuses linear + relu/silu into one kernel and collapses pointwise chains. Details: JIT Compiler docs.

FPGA / OpenCL

USE_OPENCL=1 pip install -e .
TVARANT_BACKEND=opencl python your_script.py

See FPGA / OpenCL docs.

Project structure

csrc/aten/          ATen op registrations
csrc/jit/           Pointwise JIT IR + OpenCL codegen
csrc/kernels/       Host + OpenCL kernels
csrc/runtime/       Sim and OpenCL device runtime
torch_tvarant/      Python package + graph compiler
tests/              Pytest suite
docs/               MkDocs documentation

Contributing

We welcome contributions! See CONTRIBUTING.md.

1. Fork the repo
2. Create a feature branch
3. Add tests
4. Open a pull request

Please read our Code of Conduct.

Roadmap

• Fused attention kernel (#1)
• KV-cache + decode GEMM (#2)
• RMSNorm, RoPE, SwiGLU (#3)
• Shape-specialized GEMM JIT (#4)

Full roadmap: docs/roadmap.md

License

Licensed under the Apache License 2.0.

Security

Report vulnerabilities via GitHub Security Advisories. See SECURITY.md.

Changelog

See CHANGELOG.md.