PyTorch Compilation Tutorial

This tutorial is designed for ECE6980 Spring 2023 @ Cornell. It covers the basic compilation techniques in the PyTorch ecosystem.

1. Basic PyTorch
   • Hierarchical module construction
2. TorchScript
   • Tracing mode
   • Scripting mode
3. torch.fx
   • Symbolic trace
   • Intermediate representation (IR)
   • Graph traversal and manipulation
   • limitation
4. slapo
   • Progressive optimization (manual + compiler optimization)
   • Subgraph matching
   • Operator fusion
   • Module replacement
   • Quantization (optional)
5. PyTorch 2.0 (optional)

Set up environment

Please use a clean environment and install the following package.

pip3 install torch numpy slapo tabulate transformers

Reading materials

Paper

• Adam Paszke et al., PyTorch: An Imperative Style, High-Performance Deep Learning Library, NeurIPS, 2019.
• James Reed et al., Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python, MLSys, 2022.
• Hongzheng Chen et al., Decoupled Model Schedule for Deep Learning Training, arXiv:2302.08005, 2023.

Blogs and Documents

• Yuxin Wu, TorchScript: Tracing vs. Scripting
• James Reed, Introduction to TorchScript
• PyTorch, JIT Technical Overview