Elamur Research
☰
home / papers / 2019 / Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations

Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations

Written for the Research library (claude-opus-5-5) from the paper text, in our own words: the license "author's copy — read only" does not allow redistributing the text itself. Original — eecs.harvard.edu~htk/publication/2019-mapl-tillet-kung-cox.pdf

Summary

Deep learning research often depends on fast GPU kernels. Vendor libraries such as cuBLAS and cuDNN cover only a fixed set of operations. Custom primitives therefore either run slowly or need expert-written code that does not port well. Existing DSLs (Tensor Comprehensions, Halide, TVM, PlaidML) frequently lag behind vendor code and cannot easily express non-affine indexing such as structured sparsity.

The authors propose Triton, an open-source language and compiler built around tiles: statically shaped, possibly parametric multi-dimensional sub-arrays. It has three parts:

  • Triton-C, a CUDA-like frontend with NumPy-style broadcasting and an SPMD programming model;
  • Triton-IR, an LLVM-based intermediate representation extended with tile-level data and control flow;
  • Triton-JIT, which applies machine-independent and GPU-specific optimization passes and uses an auto-tuner to choose meta-parameters like tile sizes.

In experiments, Triton matches cuBLAS on matrix multiplication and outperforms other DSLs. It reproduces cuDNN's IMPLICIT_GEMM convolution and enables an efficient fused shift-convolution. This shows that tile-aware compilation can deliver near-vendor performance while keeping custom kernels portable and fairly concise.

Key points

  • Triton consists of Triton-C, Triton-IR (LLVM-based) and Triton-JIT. Triton-JIT includes an auto-tuner over parameters marked with the tunable keyword.
  • On matrix multiplication tasks from DeepSpeech2 and transformer models, Triton is roughly on par with cuBLAS 10.0, exceeding 90% of peak device performance. It is up to 3× faster than Auto-TVM, Tensor Comprehensions and PlaidML.
  • Triton's re-implementation of cuDNN's IMPLICIT_GEMM beats cuDNN 7.0 on ResNet layers and ties on DeepSpeech2 layers.
  • A fused shift-conv kernel written in Triton-C almost completely hides the cost of shifting, compared with a naive shift kernel followed by a cuBLAS call.
  • The roofline comparison in Figure 1 was run on an NVIDIA GeForce GTX1070. Planned future work covers tensor cores, quantized kernels and integration with higher-level DSLs.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields: Computing & GPUs

Cited by

416 papers of our library cite it (of the 140673 whose bibliographies we read); 8 of them have a page here, 408 not yet.

Cite

The paper itself, not this page. BibTeX from the library's record of the work.

@inproceedings{tillet2019triton,
  title = {{Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations}},
  author = {Tillet, Philippe and Kung, H. T. and Cox, David},
  booktitle = {MAPL 2019},
  year = {2019},
  url = {https://www.eecs.harvard.edu/~htk/publication/2019-mapl-tillet-kung-cox.pdf}
}

References

48 references, 4 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, Manjunath Kudlur, Josh Levenberg, Rajat Monga, Sherry Moore, Derek G. Murray, Benoit Steiner, Paul Tucker, Vijay Vasudevan, Pete Warden, Martin Wicke, Yua… find
  2. R. C. Agarwal, S. M. Balle, F. G. Gustavson, M. Joshi, and P. Palkar. 1995. A three-dimensional approach to parallel matrix multiplication. IBM Journal of Research and Development 39, 5 (Sep. 1995), 575–582. https://doi.org/10.1147/rd.395.0575 — [doi:10.1147/rd.395.0575](https://doi.org/10.1147/rd.395.0575) ↗
  3. Dario Amodei, Rishita Anubhai, Eric Battenberg, Carl Case, Jared Casper, Bryan Catanzaro, Jingdong Chen, Mike Chrzanowski, Adam Coates, Greg Diamos, Erich Elsen, Jesse Engel, Linxi Fan, Christopher Fougner, Tony Han, Awni Y. Hannun, Billy Jun, Patrick LeGresley, Libby Lin, Sharan Narang, Andrew Y. N… find
  4. M. Auguin and F. Larbey. [n. d.]. Opsila: an advanced SIMD for numerical analysis and signal processing. ([n. d.]), 311–318. find
  5. Bin Bao and Chen Ding. 2013. Defensive Loop Tiling for Shared Cache. In Proceedings of the 2013 IEEE/ACM International Symposium on Code Generation and Optimization (CGO) (CGO ’13). IEEE Computer Society, Washington, DC, USA, 1–11. find
  6. Muthu Manikandan Baskaran, Albert Hartono, Sanket Tavarageri, Thomas Henretty, J. Ramanujam, and P. Sadayappan. 2010. Parameterized Tiling Revisited. In Proceedings of the 8th Annual IEEE/ACM International Symposium on Code Generation and Optimization (CGO ’10). ACM, New York, NY, USA, 200–209. http… — [doi:10.1145/1772954.1772983](https://doi.org/10.1145/1772954.1772983) find
  7. Matthias Braun, Sebastian Buchwald, Sebastian Hack, Roland Leissa, Christoph Mallon, and Andreas Zwinkau. 2013. Simple and Efficient Construction of Static Single Assignment Form. In Proceedings of the 22Nd International Conference on Compiler Construction (CC’13). 102– 122. find
  8. Lori Carter, Beth Simon, Brad Calder, Larry Carter, and Ferrante. [n. d.]. Predicated Static Single Assignment. In Proceedings of the PACT 1999 Conference on Parallel Architectures and Compilation Techniques. find
  9. Tianqi Chen, Mu Li, Yutian Li, Min Lin, Naiyan Wang, Minjie Wang, Tianjun Xiao, Bing Xu, Chiyuan Zhang, and Zheng Zhang. 2015. MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems. CoRR abs/1512.01274 (2015). find our card
  10. Tianqi Chen, Thierry Moreau, Ziheng Jiang, Haichen Shen, Eddie Q. Yan, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. 2018. TVM: End-to-End Optimization Stack for Deep Learning. CoRR abs/1802.04799 (2018). find our card
  11. Tianqi Chen, Thierry Moreau, Ziheng Jiang, Haichen Shen, Eddie Q. Yan, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. 2018. TVM: End-to-End Optimization Stack for Deep Learning. CoRR abs/1802.04799 (2018). find our card
  12. Jack W. Davidson and Sanjay Jinturkar. 1994. Memory Access Coalescing: A Technique for Eliminating Redundant Memory Accesses. In Proceedings of the ACM SIGPLAN 1994 Conference on Programming Language Design and Implementation (PLDI ’94). ACM, New York, NY, USA, 186–195. https://doi.org/10.1145/17824… — [doi:10.1145/178243.178259](https://doi.org/10.1145/178243.178259) ↗
  13. Greg Diamos, Shubho Sengupta, Bryan Catanzaro, Mike Chrzanowski, Adam Coates, Erich Elsen, Jesse Engel, Awni Hannun, and Sanjeev Satheesh. 2016. Persistent RNNs: Stashing Recurrent Weights On-Chip. In Proceedings of The 33rd International Conference on Machine Learning. find
  14. Venmugil Elango, Norm Rubin, Mahesh Ravishankar, Hariharan San-danagobalane, and Vinod Grover. 2018. Diesel: DSL for Linear Algebra and Neural Net Computations on GPUs. In Proceedings of the 2Nd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages (MAPL 2018). ACM, New Yo… — [doi:10.1145/3211346.3211354](https://doi.org/10.1145/3211346.3211354) find
  15. Jordan Gergov. 1999. Algorithms for Compile-time Memory Optimization. In Proceedings of the Tenth Annual ACM-SIAM Symposium on Discrete Algorithms (SODA ’99). find
  16. Google Inc. 2017. Tensorflow XLA. https://www.tensorflow.org/performance/xla/ ↗
  17. Scott Gray and Alex Radford ans Diederik P. Kingma. 2017. Gpu kernels for block-sparse weights. CoRR abs/1711.09224 (2017). find
  18. M. Griebl, C. Lengauer, and S. Wetzel. 1998. Code generation in the polytope model. In Proceedings. 1998 International Conference on Parallel Architectures and Compilation Techniques (Cat. No.98EX192). 106–111. https://doi.org/10.1109/PACT.1998.727179 — [doi:10.1109/pact.1998.727179](https://doi.org/10.1109/pact.1998.727179) ↗
  19. K. He, X. Zhang, S. Ren, and J. Sun. 2016. Deep Residual Learning for Image Recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR). 770–778. https://doi.org/10.1109/CVPR.2016.90 ↗
  20. Andrew G. Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam. 2017. MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications. CoRR abs/1704.04861 (2017). arXiv:1704.04861 http://arxiv.org/abs/1704.04861 ↗ our card
  21. Jianyu Huang and Robert A. van de Geijn. 2016. BLISlab: A Sandbox for Optimizing GEMM. FLAME Working Note \#80, TR-16-13. The University of Texas at Austin, Department of Computer Science. http: //arxiv.org/pdf/1609.00076v1.pdf — [arXiv:1609.00076](https://arxiv.org/abs/1609.00076) find
  22. Intel AI. 2018. PlaidML. https://www.intel.ai/reintroducing-plaidml ↗
  23. Andrew Kerr, Duane Merrill, Julien Demouth, and John Tran. [n. d.]. CUTLASS: Fast Linear Algebra in CUDA C++. ([n. d.]). https://devblogs. nvidia.com/cutlass-linear-algebra-cuda/ find
  24. Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton. 2012. ImageNet Classification with Deep Convolutional Neural Networks. In Proceedings of the 25th International Conference on Neural Information Processing Systems - Volume 1 (NIPS’12). find
  25. Andrew Lavin and Scott Gray. 2016. Fast Algorithms for Convolutional Neural Networks. (CVPR’16). find
  26. Darryl D. Lin, Sachin S. Talathi, and V. Sreekanth Annapureddy. 2016. Fixed Point Quantization of Deep Convolutional Networks. In Proceedings of the 33rd International Conference on International Conference on Machine Learning - Volume 48 (ICML’16). JMLR.org, 2849–2858. http://dl.acm.org/citation.cf… ↗
  27. Matt Martineau, Patrick Atkinson, and Simon McIntosh-Smith. 2017. Benchmarking the NVIDIA V100 GPU and Tensor Cores. Springer. find
  28. Bradley McDanel, Surat Teerapittayanon, and H.T. Kung. 2017. Embedded Binarized Neural Networks. In Proceedings of the 2017 International Conference on Embedded Wireless Systems and Networks (EWSN &\#8217;17). Junction Publishing, USA, 168–173. http://dl.acm.org/citation.cfm?id=3108009.3108031 ↗
  29. W. M. McKeeman. 1965. Peephole Optimization. Commun. ACM 8, 7 (July 1965), 443–444. https://doi.org/10.1145/364995.365000 — [doi:10.1145/364995.365000](https://doi.org/10.1145/364995.365000) ↗
  30. Sanyam Mehta, Rajat Garg, Nishad Trivedi, and Pen-Chung Yew. 2016. TurboTiling: Leveraging Prefetching to Boost Performance of Tiled Codes. In Proceedings of the 2016 International Conference on Supercomputing (ICS ’16). ACM, New York, NY, USA, Article 38, 12 pages. find
  31. Sharan Narang, Eric Undersander, and Gregory F. Diamos. 2017. Block-Sparse Recurrent Neural Networks. CoRR abs/1711.02782 (2017). — [arXiv:1711.02782](https://arxiv.org/abs/1711.02782) find
  32. Rajib Nath, Stanimire Tomov, and Jack Dongarra. 2011. Accelerating GPU Kernels for Dense Linear Algebra. In Proceedings of the 9th International Conference on High Performance Computing for Computational Science (VECPAR’10). Springer-Verlag, Berlin, Heidelberg, 83–92. http://dl.acm.org/citation.cfm?… ↗
  33. John Nickolls, Ian Buck, Michael Garland, and Kevin Skadron. 2008. Scalable Parallel Programming with CUDA. Queue 6, 2 (March 2008), 40–53. https://doi.org/10.1145/1365490.1365500 ↗
  34. Cedric Nugteren. 2017. CLBlast: A Tuned OpenCL BLAS Library. CoRR abs/1705.05249 (2017). arXiv:1705.05249 http://arxiv.org/abs/1705.05249 ↗
  35. Travis Oliphant. 2006–. NumPy: A guide to NumPy. USA: Trelgol Publishing. (2006–). http://www.numpy.org/ [Online; accessed \<today\>]. ↗
  36. PyTorch. 2016. . https://github.com/pytorch/pytorch ↗
  37. Jonathan Ragan-Kelley, Connelly Barnes, Andrew Adams, Sylvain Paris, Frédo Durand, and Saman Amarasinghe. 2013. Halide: A Language and Compiler for Optimizing Parallelism, Locality, and Recomputation in Image Processing Pipelines. In Proceedings of the 34th ACM SIGPLAN Conference on Programming Lang… find
  38. Nadav Rotem, Jordan Fix, Saleem Abdulrasool, Summer Deng, Roman Dzhabarov, James Hegeman, Roman Levenstein, Bert Maher, Nadathur Satish, Jakob Olesen, Jongsoo Park, Artem Rakhov, and Misha Smelyanskiy. 2018. Glow: Graph Lowering Compiler Techniques for Neural Networks. CoRR abs/1805.00907 (2018). find
  39. Arthur Stoutchinin and Francois de Ferriere. 2001. Efficient Static Single Assignment Form for Predication. In Proceedings of the 34th Annual ACM/IEEE International Symposium on Microarchitecture (MI-CRO 34). IEEE Computer Society, Washington, DC, USA, 172–181. http://dl.acm.org/citation.cfm?id=5639… ↗
  40. Philippe Tillet and David Cox. 2017. Input-aware Auto-tuning of Compute-bound HPC Kernels. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC ’17). ACM. find
  41. Didem Unat, Tan Nguyen, Weiqun Zhang, Muhammed Nufail Farooqi, Burak Bastem, George Michelogiannakis, Ann Almgren, editor="Kunkel Julian M. Shalf, John", Pavan Balaji, and Jack Dongarra. 2016". TiDA: High-Level Programming Abstractions for Data Locality Management. find
  42. Field G. Van Zee and Robert A. van de Geijn. 2015. BLIS: A Framework for Rapidly Instantiating BLAS Functionality. ACM Trans. Math. Softw. 41, 3 (June 2015). find
  43. Nicolas Vasilache, Oleksandr Zinenko, Theodoros Theodoridis, Priya Goyal, Zachary DeVito, William S. Moses, Sven Verdoolaege, Andrew Adams, and Albert Cohen. 2018. Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions. CoRR abs/1802.04730 (2018). find
  44. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention Is All You Need. CoRR abs/1706.03762 (2017). arXiv:1706.03762 http://arxiv.org/abs/1706.03762 ↗
  45. R. Clint Whaley, Antoine Petitet, and Jack J. Dongarra. 2000. Automated Empirical Optimization of Software and the ATLAS Project. PARALLEL COMPUTING 27 (2000), 2001. find
  46. Samuel Williams, Andrew Waterman, and David Patterson. 2009. Roofline: An Insightful Visual Performance Model for Multicore Architectures. Commun. ACM 52, 4 (April 2009), 65–76. https://doi.org/10.1145/1498765.1498785 ↗
  47. Bichen Wu, Alvin Wan, Xiangyu Yue, Peter H. Jin, Sicheng Zhao, Noah Golmant, Amir Gholaminejad, Joseph Gonzalez, and Kurt Keutzer. 2017. Shift: A Zero FLOP, Zero Parameter Alternative to Spatial Convolutions. CoRR abs/1711.08141 (2017). — [arXiv:1711.08141](https://arxiv.org/abs/1711.08141) find
  48. Jiacheng Zhao, Huimin Cui, Yalin Zhang, Jingling Xue, and Xiaobing Feng. 2018. Revisiting Loop Tiling for Datacenters: Live and Let Live. In Proceedings of the 2018 International Conference on Supercomputing (ICS ’18). ACM, New York, NY, USA, 328–340. find