Add Interleaved Complex Number Support to CubeCL
Motivation
CubeCL currently supports only real-valued scalar types (f16, bf16, f32, f64, integers). Scientific computing workloads — quantum mechanics, electromagnetism, signal processing, etc. — require native complex number tensor operations on GPU.
We are developing tenferro-rs, a tensor computation library with StableHLO-style ops, automatic differentiation, and complex linalg (SVD, QR, Cholesky, Eigh), and are evaluating CubeCL as a GPU backend for elementwise/structural/reduction operations (with cuBLAS/cuSOLVER handling GEMM and linalg). tenferro-rs natively supports Complex32/Complex64 throughout its type system. CubeCL's advantages over XLA for this use case:
- Pure Rust — no binary distribution problem (XLA requires shipping ~100MB+ libxla.so)
- Lightweight JIT — shape changes trigger fast PTX regeneration, not full LLVM recompilation (critical for iterative algorithms like TCI where tensor shapes change every step)
- Multi-vendor — CUDA, WebGPU from a single codebase
If CubeCL gains interleaved complex support for CUDA and WGPU backends, we can start integration testing immediately — CUDA for GPU-accelerated scientific workloads, WGPU for portability testing. Both backends have CI or local test coverage.
Related Work
Why Interleaved (Not Split)
Burn's current approach uses split representation (separate real/imag tensors) to avoid CubeCL IR changes. However, split representation is fundamentally inadequate for scientific computing:
- Complex linalg (SVD, QR, Cholesky, Eigh) requires interleaved data — cuSOLVER's
cgesvd, cgeqrf, cpotrf all expect cuComplex* (interleaved)
- Complex GEMM: split requires 4 real GEMMs + 2 adds; cuBLAS
cgemm does it in 1 call with Tensor Core support
- Memory transfers: every linalg call would need split→interleaved→split round-trip copies
- Cache locality: interleaved keeps re/im adjacent, better for SIMD and GPU memory coalescing
The proposal here is to add interleaved complex types (Complex<f32>, Complex<f64>) to CubeCL's IR and backends, covering elementwise + structural + reduction operations. Complex linalg and GEMM are delegated to cuBLAS/cuSOLVER/cuTENSOR.
Why CUDA + WGPU (not HIP)
- CUDA: Available locally (A100, CUDA 12.0), tested in upstream CI
- WGPU: Tested in upstream CI on Linux, testable locally without GPU
- HIP: No CI coverage in upstream, no local hardware for testing — deferred indefinitely
Scope
CubeCL handles:
- Elementwise: add, sub, mul, div, neg, conj, abs, exp, log, sin, cos, sqrt, pow, ...
- Structural: transpose, reshape, broadcast (element-type agnostic)
- Reduction: sum, prod
- Type conversion: real↔complex promotion/extraction
CubeCL does NOT handle (delegated to vendor libraries):
- GEMM → cuBLAS
cgemm/zgemm
- Linalg → cuSOLVER
- Tensor contraction → cuTENSOR
Feasibility Analysis
Key Enabler: thrust::complex<T>
CUDA's thrust::complex<float/double> provides operator overloading for +, -, *, /, ==, !=, plus thrust::abs, thrust::exp, thrust::log, thrust::sqrt, thrust::conj, etc.
Since CubeCL's code generator emits templates like {lhs} + {rhs}, substituting the type name is sufficient for most arithmetic ops:
// Current generated code (float)
float out = lhs + rhs;
// With complex (no codegen logic change needed)
thrust::complex<float> out = lhs + rhs; // operator+ defined by thrust
This dramatically reduces the CUDA backend effort.
Per-Layer Effort Estimate
1. IR Type System (cubecl-ir) — ~100 lines
- Add
Complex(ComplexKind) variant to ElemType with C32, C64
size() / size_bits(): return 2x base float size- Add validation: reject bitwise ops and ordering comparisons for complex types
2. CUDA Backend (cubecl-cpp/cuda) — ~220 lines
| Item |
Lines |
Notes |
| Include |
1 |
#include <thrust/complex.h> |
Elem enum |
~50 |
Add CF32, CF64; ident() → "thrust::complex<float>" |
| Binary ops |
~0 |
+, -, *, /, ==, != work via thrust overloading |
| Unary ops |
~30 |
thrust::abs, thrust::exp, thrust::log, thrust::sqrt, thrust::conj |
| Custom ops |
~60 |
pow, max/min, isnan/isinf (check re/im separately) |
| Type conversion |
~50 |
float→complex: thrust::complex<float>(x, 0.0f); complex→float: .real() |
| Reject paths |
~30 |
Bitwise, ordering comparisons → compile error |
3. WGPU Backend (cubecl-wgpu) — ~330 lines
- WGSL has no native complex type
- Requires
struct Complex32 { re: f32, im: f32 } and manual operation expansion
- Transcendentals need explicit formulas:
exp(z) = exp(re) * (cos(im), sin(im))
- Has CI coverage and is testable locally
4. CPU Backend (cubecl-cpu) — ~200 lines
- MLIR/LLVM struct type generation
num_complex formulas for lowering
5. Frontend / Proc Macros (cubecl-core, cubecl-macros) — ~180 lines
CubeType / CubePrimitive impl for Complex32/64Scalar trait: Complex is 2-field — expand generation needs adjustment- New frontend ops:
conj(), real(), imag()
6. Vectorization — ~30 lines
- Treat Complex as single element; disable auto-vectorization for complex types
- Scientific workloads are GEMM-dominated; vectorization is not critical
Operation Compatibility Matrix
| Category |
Works via thrust overloading |
Needs custom impl |
Not applicable |
| Binary arithmetic |
+, -, *, / |
pow, max, min, atan2, hypot |
%, saturating ops, mulhi |
| Unary math |
abs, exp, log, sqrt, sin, cos |
isnan, isinf |
ceil, floor, round, trunc |
| Comparison |
==, != |
— |
<, <=, >, >= (no total order on C) |
| Bitwise |
— |
— |
All (not defined for complex) |
| Structural |
All (element-type agnostic) |
— |
— |
| Reduction |
sum, prod |
— |
max, min |
Summary
| Scope |
Lines Changed |
Difficulty |
| CUDA only |
~500 |
Low–Medium |
| CUDA + WGPU |
~830 |
Medium |
| CUDA + WGPU + CPU |
~1030 |
Medium |
Proposed Phased Approach
Phase 1: CUDA backend
- IR type system:
Complex(C32), Complex(C64)
- CUDA codegen via
thrust::complex<T>
- Frontend API:
conj(), real(), imag()
- Validation: reject invalid ops for complex types
- Tests for elementwise, structural, reduction
Phase 2: WGPU backend
- WGSL struct-based emulation
- Manual transcendental expansion
Phase 3: CPU backend
Add Interleaved Complex Number Support to CubeCL
Motivation
CubeCL currently supports only real-valued scalar types (f16, bf16, f32, f64, integers). Scientific computing workloads — quantum mechanics, electromagnetism, signal processing, etc. — require native complex number tensor operations on GPU.
We are developing tenferro-rs, a tensor computation library with StableHLO-style ops, automatic differentiation, and complex linalg (SVD, QR, Cholesky, Eigh), and are evaluating CubeCL as a GPU backend for elementwise/structural/reduction operations (with cuBLAS/cuSOLVER handling GEMM and linalg). tenferro-rs natively supports Complex32/Complex64 throughout its type system. CubeCL's advantages over XLA for this use case:
If CubeCL gains interleaved complex support for CUDA and WGPU backends, we can start integration testing immediately — CUDA for GPU-accelerated scientific workloads, WGPU for portability testing. Both backends have CI or local test coverage.
Related Work
Why Interleaved (Not Split)
Burn's current approach uses split representation (separate real/imag tensors) to avoid CubeCL IR changes. However, split representation is fundamentally inadequate for scientific computing:
cgesvd,cgeqrf,cpotrfall expectcuComplex*(interleaved)cgemmdoes it in 1 call with Tensor Core supportThe proposal here is to add interleaved complex types (
Complex<f32>,Complex<f64>) to CubeCL's IR and backends, covering elementwise + structural + reduction operations. Complex linalg and GEMM are delegated to cuBLAS/cuSOLVER/cuTENSOR.Why CUDA + WGPU (not HIP)
Scope
CubeCL handles:
CubeCL does NOT handle (delegated to vendor libraries):
cgemm/zgemmFeasibility Analysis
Key Enabler:
thrust::complex<T>CUDA's
thrust::complex<float/double>provides operator overloading for+,-,*,/,==,!=, plusthrust::abs,thrust::exp,thrust::log,thrust::sqrt,thrust::conj, etc.Since CubeCL's code generator emits templates like
{lhs} + {rhs}, substituting the type name is sufficient for most arithmetic ops:This dramatically reduces the CUDA backend effort.
Per-Layer Effort Estimate
1. IR Type System (
cubecl-ir) — ~100 linesComplex(ComplexKind)variant toElemTypewithC32,C64size()/size_bits(): return 2x base float size2. CUDA Backend (
cubecl-cpp/cuda) — ~220 lines#include <thrust/complex.h>ElemenumCF32,CF64;ident()→"thrust::complex<float>"+,-,*,/,==,!=work via thrust overloadingthrust::abs,thrust::exp,thrust::log,thrust::sqrt,thrust::conjpow,max/min,isnan/isinf(check re/im separately)float→complex:thrust::complex<float>(x, 0.0f);complex→float:.real()3. WGPU Backend (
cubecl-wgpu) — ~330 linesstruct Complex32 { re: f32, im: f32 }and manual operation expansionexp(z) = exp(re) * (cos(im), sin(im))4. CPU Backend (
cubecl-cpu) — ~200 linesnum_complexformulas for lowering5. Frontend / Proc Macros (
cubecl-core,cubecl-macros) — ~180 linesCubeType/CubePrimitiveimpl for Complex32/64Scalartrait: Complex is 2-field — expand generation needs adjustmentconj(),real(),imag()6. Vectorization — ~30 lines
Operation Compatibility Matrix
+,-,*,/pow,max,min,atan2,hypot%, saturating ops,mulhiabs,exp,log,sqrt,sin,cosisnan,isinfceil,floor,round,trunc==,!=<,<=,>,>=(no total order on C)sum,prodmax,minSummary
Proposed Phased Approach
Phase 1: CUDA backend
Complex(C32),Complex(C64)thrust::complex<T>conj(),real(),imag()Phase 2: WGPU backend
Phase 3: CPU backend