FreeTensor: A Free-form DSL with Holistic Optimizations for Irregular Tensor Programs

on YouTube ↗16m

" summary

A Python-embedded DSL and compiler that compiles loop-based, irregular tensor code to native CPU/GPU kernels, using polyhedral analysis via isl to enable legal loop transforms, partial-evaluation of recursive dimension handling, and an AD strategy that recomputes only necessary intermediates. Evaluated on SubdivNet, LongForm, a differentiable renderer and GNNs against PyTorch/JAX/TVM/Julia/TGL, it reports up to 5.10× without AD, up to 12.7× with AD and an AD memory-reuse win of up to 6.83×, with far fewer GPU kernels and memory accesses.

AI-written summary · about

" chapters (11)

1. 0:00Intro
2. 0:31Frameworks for Tensor Programs
3. 1:24Tensor Operators vs. Irregularity
4. 3:35Core Idea: Reintroducing Loops!
5. 5:09Embedded Python DSL
6. 6:05Optimizing Compiler
7. 6:54Generic operations in Python DSL
8. 8:24Dependence-aware Transformations
9. 10:00Fine-grained Automatic Differentiation
10. 12:37Performance Evaluation: End to End
11. 13:20FreeTensor in a nutshell

" related talks

• 24mVerified Tensor-Program Optimization Via High-Level Scheduling RewritesAdam Chlipala, Jonathan Ragan-Kelley, Gilbert Louis Bernstein, Amanda Liu · POPL 2022
• 45mTensorly: A Flexible Python Framework for Machine LearningAnima Anandkumar · PyData Los Angeles 2018
• 28mFunctional Sparse Tensor CompilationAmir Shaikhha, Shideh Hashemian · FProPer 2024
• 20mCompiling Structured Tensor AlgebraMahdi Ghorbani, Amir Shaikhha, Mathieu Huot, Shideh Hashemian · OOPSLA 2023
• 15mEnabling Tensor Language Model to Assist in Generating High-Performance Tensor...Yu Zhang, Yi Zhai, Sijia Yang, Keyu Pan, Renwei Zhang, Shuo Liu, Chao Liu, Zichun Ye, Jianmin Ji, Jie Zhao, Yanyong Zhang · 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI '24)
• 15mCompiling and Compressing Structured TensorsEmilien Bauer · Sparse 2025
• 11mQiMeng-Xpiler: Transcompiling Tensor Programs for Deep Learning Systems with a...Qi Guo, Xuehai Zhou, Shouyang Dong, Yuanbo Wen, Jun Bi, Di Huang, Jiaming Guo, Jianxing Xu, Ruibai Xu, Xinkai Song, Yifan Hao, Ling Li, Tianshi Chen, Yunji Chen · OSDI '25
• 22mTileLoom: Automatic Dataflow Planning for Tile-Based Languages on Spatial Dataflow Accel…Wei Li, Cheng Tan, Zhenyu Bai, Heru Wang, Pranav Dangi, Zhiqiang Zhang, Huiying Lan, Weng-Fai Wong, Tulika Mitra · 20th USENIX Symposium on Operating Systems Design and Implementation (OSDI '26)