GPU tropical path: keep tensors device-resident + device-side permute (replicate Julia's CuArray data flow)

[isPANN]

Summary

Our GPU contraction path is slow not because the device kernels are slow, but because the data flow leaves the device. Julia's OMEinsum.jl is fast on GPU precisely because it writes almost no GPU executor: its contraction logic is written once against AbstractArray, and CuArray (CUDA.jl) supplies the two things that actually make it fast — device residency and a device-side permutedims — for free. We don't get that free lunch in Rust, so we have to implement the same data flow by hand.

This issue tracks closing that gap.

Why Julia is "naturally" fast

Julia gets for free (via CuArray)	What we must build in Rust
Intermediates stay resident (CuArray never leaves VRAM)	Stop calling to_vec(); let CudaStorage flow through the whole contraction tree (the engine already supports this structurally)
permutedims(::CuArray) → auto-dispatched GPU kernel	New code: a device-side permute/gather — see "Resolved design" below
mul! → CuTropicalGEMM on resident CuArrays	Use the already-available device-resident batched API (tropical_matmul_gpu_strided_batched / GpuMatrix) instead of per-slice from_host/to_host
One driver context / module	Cache the context + kernel module once instead of rebuilding per node

Current state in the code (verified)

The tropical GPU path round-trips host↔device for operand prep and permutation:

• Cuda::plan_tropical_operands (src/backend/cuda/mod.rs:528) downloads both operands with to_vec() and runs the host CPU helpers gather_contiguous, reduce_trace, materialize_strided from backend::contract_plan (mod.rs:549-578).
• Cuda::run_tropical_gemm / ..._with_argmax (mod.rs:756, mod.rs:791) do a per-slice GpuMatrix::from_host(...) upload and to_host(...) download inside the batch loop (mod.rs:762-768, mod.rs:798-810).
• permute_tropical_output (mod.rs:707) permutes the result on the host.
• tropical_gemm_cuda::CudaContext::from_device(...) is rebuilt per call (mod.rs:607, mod.rs:663) — the kernel module is effectively reloaded for every node.

The doc comment at mod.rs:435-438 already acknowledges this: "Operand materialization … currently runs on the host via a download/upload roundtrip … a device-side permute is a later optimization."

Note: the standard cuTENSOR path (contract_cutensor, mod.rs:342) already feeds strides straight into the cuTENSOR descriptor and absorbs the permutation into cutensorContract, so it does not round-trip. The host-roundtrip permute is specific to the tropical path, because our CPU contraction helpers (materialize_strided / gather_contiguous / reduce_trace) only accept host &[T].

So "replicating Julia" here is not a line-by-line translation — it's reimplementing, by hand, the data flow that CuArray donated to Julia. The bulk of the new work is the device-side permute.

Proposed plan (incremental, each step independently verifiable)

1. Context + module caching — smallest change, ~9% for free. Cache the tropical_gemm_cuda::CudaContext (and any compiled kernel module) once instead of rebuilding per node. Do this first to confirm the pipeline still produces correct values.
2. Device-resident GEMM — switch run_tropical_gemm[_with_argmax] from per-slice from_host/to_host to the existing device-resident batched API (tropical_matmul_gpu_strided_batched / resident GpuMatrix), keeping the GEMM output in VRAM. (We are already ahead of Julia here — see findings — since CuTropicalGEMM.jl only specializes the 2-D mul! and OMEinsum loops it per batch slice, whereas we have a genuinely strided-batched tropical kernel.)
3. Device-side permute — the real work and the bulk of the win (~85%). Scope resolved below.

Steps 1 and 2 are "use an API that already exists"; step 3 is the new kernel work.

Resolved design for step 3 (follows Julia's approach)

The original open question — "cuTENSOR permute API vs. a hand-written CUDA kernel" — is resolved: a custom kernel is mandatory, because cuTENSOR 2.x provably cannot do two things this path needs:

1. No diagonal/trace. cuTENSOR requires "each mode may appear in each tensor at most once", so a repeated-label trace (e.g. iij->ij, our reduce_trace) cannot be expressed in any cuTENSOR permute or reduction descriptor. (cuTENSOR API docs)
2. No integer dtypes in permute/reduce. cuTENSOR permutation supports only FP16/BF16/FP32/FP64 + complex — not integers. Our u32 argmax buffer (tropical winner-routing) cannot go through cutensorPermute. (cuTENSOR types)

This mirrors Julia exactly: CUDA.jl's permutedims(::CuArray) is not cuTENSOR — it is a generic hand-written one-thread-per-element stride-arithmetic scatter kernel in GPUArrays.jl (src/host/linalg.jl), and OMEinsum resolves diagonals/traces as a separate unary pre-pass (the Diag/Tr rules in src/unaryrules.jl) before the tensor enters the permute → reshape → batched-GEMM → permute pipeline.

So we adopt the same shape — the whole tropical path uses our own kernels, not a cuTENSOR/kernel hybrid. Rationale for uniform-over-hybrid: the u32 argmax forces a custom kernel regardless; and the float result and u32 argmax must be permuted with the bit-for-bit identical index mapping (winner-routing correctness), which a single shared kernel guarantees and two separate code paths would risk. Permute is memory-bandwidth-bound and the tropical GEMM dominates wall-clock, so cuTENSOR's tuned-permute edge is moot here.

Two dtype-generic kernels (compiled once, via cudarc 0.19.x):

• (a) Generic N-d permute / strided gather — one thread per output element: decode the linear index with divrem against the shape, recombine with the permuted strides (the GPUArrays.jl algorithm). Templated on dtype so it serves f32/f64/c32/c64 operands/results and the u32 argmax. Covers operand canonicalization ([left, contracted, batch]) and the output permute. A naive index-arithmetic version is an acceptable, value-testable first cut (~40% of copy bandwidth); a tiled+shared-memory+padded transpose is a later ~2.6× optimization, not required initially.
• (b) Diagonal/trace gather + max/min reduce — same index machinery with shared-index axes, fusing the tropical max/min into the gather. Applied as a unary pre-pass (Julia's Diag/Tr). This is the correctness hotspot (repeated-label lowering) — add the value/gradient regression alongside it.

Compilation & caching (cudarc 0.19.x): compile the kernels once via nvrtc::compile_ptx or build.rs + nvcc embedding PTX (candle's pattern), and cache the resulting Arc<CudaModule> in an Arc<RwLock<HashMap<_, _>>> on the device wrapper — never compile or load_module per contraction. This is the same caching that fixes the per-node context rebuild in step 1. Launch via stream.launch_builder(&f).arg(..).launch(cfg).

What stays on cuTENSOR: the standard (+,×) path only — it already absorbs strided views and permutation into cutensorContract (mod.rs:342), no change needed. cuTENSOR permute remains a possible later optimization for the float operand permute if profiling shows the naive kernel is a bottleneck.

Not adopted: cuTT / LibreTT / TTLG — all are dense-only (no strided/diagonal input), so they would still require a separate gather and add a dependency for no gain over cuTENSOR on an NVIDIA-only build.

Structural lessons from OMEinsum we follow: keep diagonal/trace as a separate unary pre-pass so the GEMM core only ever sees clean [left, contracted, batch] tensors; and use GEMM transpose flags to avoid one permute where the layout allows.

Acceptance criteria

• Tropical GPU contractions keep intermediates resident in VRAM across the contraction tree (no per-node to_vec() round-trips for operand prep / permute).
• The driver/kernel context + compiled module is created once, not per node.
• Forward, backward, and argmax results remain bit-for-bit consistent with the CPU path (regressions on concrete values + gradients + device preservation, per the testing conventions). In particular, the float result and the u32 argmax must be permuted with identical index mappings.
• A benchmark demonstrating the wall-clock improvement on a representative tropical network.

Related

• Cuda backend silently returns standard (+, ×) results for tropical algebras — dispatch ignores the algebra parameter (follow-up to #6) #48 — Cuda backend silently returns standard (+, ×) results for tropical algebras (correctness, distinct from this performance/architecture issue).

References (step 3 design)

• cuTENSOR permute/reduce constraints (each mode once; supported dtypes): https://docs.nvidia.com/cuda/cutensor/latest/api/cutensor.html , https://docs.nvidia.com/cuda/cutensor/2.0.0/api/types.html
• CUDA.jl/GPUArrays.jl device permutedims kernel: https://github.com/JuliaGPU/GPUArrays.jl/blob/master/src/host/linalg.jl
• OMEinsum.jl unary (Diag/Tr) + binary GEMM lowering: https://github.com/under-Peter/OMEinsum.jl/blob/master/src/unaryrules.jl , https://github.com/under-Peter/OMEinsum.jl/blob/master/src/binaryrules.jl
• CuTropicalGEMM.jl mul! interception (2-D only, not batched): https://github.com/TensorBFS/CuTropicalGEMM.jl/blob/main/src/tropical_gemms.jl
• cudarc kernel compile/launch/caching (candle pattern): https://github.com/coreylowman/cudarc , https://github.com/huggingface/candle
• Transpose perf / tiled+padded ~2.6×: https://developer.nvidia.com/blog/efficient-matrix-transpose-cuda-cc/ ; cuTT (dense-only): https://arxiv.org/abs/1705.01598