Translating CUDA Tile Operations from Python to Rust with AI Agents
Achieving 99.5% performance parity while maintaining Rust's safety guarantees.
1Rust ownership for GPU kernels
cuTile Rust extends Rust's ownership model to tile-based GPU kernels, splitting mutable outputs into disjoint pieces while preserving host-side ownership contracts across kernel launches. This maintains Rust's memory safety guarantees while enabling high-performance GPU programming. The system allows opting out for lower-level control when needed, providing flexibility without compromising safety.
2AI translation pipeline
An AI agent skill in the TileGym repository translates cuTile Python and Triton-TileIR kernels into cuTile Rust. The pipeline uses a bounded multi-agent workflow with machine-checkable verdicts at each stage. IR diffing against reference Tile IR verifies structural equivalence before functional testing, ensuring correctness throughout the translation process.
3Performance results
The team ported all 24 public TileGym operators to cuTile Rust, achieving 99.5% of cuTile Python performance on average. Benchmarking on NVIDIA DGX B200 shows all operators clear a 0.95 geomean speedup threshold versus cuTile Python, with overall parity at 0.995. The operators contain roughly 40 GPU kernels total, ranging from element-wise operations to complex models.
4Compilation differences
cuTile Rust requires explicit specialization in kernel signatures where cuTile Python specializes implicitly at call time. This makes type and shape constraints visible to rustc at compile time. While the GPU binary is still JIT-compiled, the implicitness is gone - nothing is specialized unless the kernel signature declares it.
5Shared IR foundation
All three frontends (cuTile Python, Triton-TileIR, cuTile Rust) use the same CUDA Tile IR (cuda_tile dialect), making translation practical and verifiable. This shared foundation enables checking the agent's output structurally through IR diffing, not just functionally through tests. The same compiler and performance model underlie all implementations.
6C-ABI integration
Converted kernels integrate with TileGym through a C-ABI layer that passes tensor descriptors without copying or allocating. This enables lazy compilation and seamless development cycles. The system uses ManuallyDrop to wrap PyTorch device pointers, allowing Rust to hand kernels tensors without taking ownership.
7Translation workflow
The translation process begins with an analyzer that solves the 'JIT hides the spec' problem by dumping Tile IR and writing a machine-readable spec. A kernel writer then produces Rust code, proving its work through functional tests and IR self-checks. Finally, a host/FFI builder makes kernels callable from TileGym.
8Example kernel comparison
The post shows a side-by-side comparison of a softmax kernel in both cuTile Python and cuTile Rust. While the Rust version requires more explicit declarations, both compile to the same Tile IR ops. The IR diff confirms identical operation inventory between implementations, validating the translation's accuracy.
9Specialization challenges
Most translation work comes from spelling out what Python source leaves implicit. The main cases involve constants baked into the DSL, vanished untaken branches, and launch parameters living in host code. The analyzer identifies these hidden specifications to guide the translation process.
10Agent architecture
The conversion skill uses specialized subagents, each loading only reference documents its stage needs. The top-level agent acts as pure orchestrator. This design ensures failures stay attributable to specific components and prevents any single agent from having too much authority.
11Development workflow
The system supports friction-free iteration in Rust. The backend tracks source freshness, automatically rebuilding shared libraries when kernel sources change. Developers can edit kernel.rs and run tests without explicit builds, maintaining Python's rapid development cycle while using Rust.
12Future implications
This work demonstrates how AI agents can bridge language gaps in high-performance computing while preserving each language's strengths. The approach maintains Python's productivity for algorithm development while bringing Rust's safety to GPU programming, potentially influencing future HPC toolchains.