Large-Scale TensorCircuit Contractions on GPUs: Disabling XLA GPU Autotuning
When running large-scale tensor-network contractions with TensorCircuit-NG and the JAX GPU backend, the following runtime configuration is worth testing: XLA_PYTHON_CLIENT_PREALLOCATE = false XLA_FLAGS = --xla_gpu_autotune_level = 0 python your_script.py Its main benefit is not speed, but lower persistent GPU memory usage from XLA GPU autotuning, which makes memory behavior during compilation and on the first visible GPU more predictable. In the TensorCircuit contraction workloads we tested, disabling autotuning also slightly improved steady-state runtime, but the memory savings were the more important result. Key takeaway XLA GPU autotuning evaluates alternative algorithms or workspace configurations for certain GPU kernels and custom calls, then selects an implementation. This can be valuable for convolutions, large GEMMs, and deep-learning workloads with fixed shapes. For large TensorCircuit contractions, however, the contraction path is already determined by OMECO or cotengra, leaving relatively little optimization freedom for autotuning while still potentially incurring substantial persistent memory overhead during compilation and tuning. For TensorCircuit contractions, run this A/B test by default: # Baseline XLA_PYTHON_CLIENT_PREALLOCATE = false python your_script.py # Test configuration XLA_PYTHON_CLIENT_PREALLOCATE = false XLA_FLAGS = --xla_gpu_autotune_level = 0 python your_script.py Both environment variables must be set before the Python process starts and before JAX is imported. This recommendation primarily concerns GPUs; CPU backends do not exhibit the same GPU-kernel autotuning behavior. Representative results All results below use a fixed contraction path so that path-search randomness does not affect the comparison. Workload Autotuning Post-compile memory Peak memory Steady-state runtime 100 qubits × 24 layers, amplitude Default 8.7 GiB 8.7 GiB 0.37 s 100 qubits × 24 layers, amplitude autotune=0 0.5 GiB 4.6 GiB 0.40 s 28 qubits × 12 layers, expecta