Accelerating Complex-Valued Tensors With Torch.compile
summary
Provides a working subclass storing real and imaginary parts that lets torch.compile succeed on complex dtypes today on PyTorch nightly, plus an open PR to Inductor to thread the subclass through compilation. Matmul performance forces keeping real/imag separate during kernels and reinterleaving at boundaries, which breaks view/mutation correctness unless views are disallowed or a second interleaved subclass is used. FFTs require backend support, and a planned JIT kernels on demand approach aims to avoid binary size blowup.
AI-written summary · about
chapters (9)
- 0:00Introduction
- 1:23Complex tensors in PyTorch
- 3:09Current solution and subclass
- 4:57Performance challenges
- 7:29View semantics and correctness
- 8:56FFT operations and backend
- 10:00Binary size and complex32
- 11:06Future JIT plans
- 12:08Q&A
related talks
07mComplex Numbers in PyTorch
12mLightning Talk: Accelerating Inference on CPU with Torch.Compile
23mBlazing Fast GenAI Inference With Torch.compile
36mDL Compiler Panel Discussion
15mLightning Talk: AOTInductor: Ahead-of-Time Compilation for PT2 Exported Models
09mTogether Goes Brrr: Threading Research & Production with Torch Compile
15mPyTorch 2.0: Unlocking the Power of Deep Learning with the Torch Compile API
24mAdvancements in High-Performance AI/ML through PyTorch's Python Compiler