Elamur Research⌕
☰
home / papers / 2019 / Learning to Optimize Tensor Programs

Learning to Optimize Tensor Programs

Paper · 2019 · Tianqi Chen, Lianmin Zheng, Eddie Yan et al.

Summary

Deep learning systems depend on fast tensor operators such as matrix multiplication and convolution. Today these come mostly from hand-tuned vendor libraries like cuDNN, which cover only a narrow set of hardware and make every new device expensive to support. The authors treat operator implementation as a search problem. An operator is written as an index expression, and the goal is to find the schedule (loop order, tiling, threading, etc.) whose compiled code runs fastest on real hardware.

Their framework, AutoTVM, trains a statistical cost model that predicts run time from the generated low-level program. This model steers exploration of a space containing billions of program variants. Because a full system needs many similar operators, the model is transferred across workloads, which speeds up tuning by 2–10×.

Without relying on any external operator library, AutoTVM performs as well as or better than cuDNN, TFLite and the ARM Compute Library across three platforms: a GPU for servers, a CPU for embedded devices, and a GPU for mobile devices. End-to-end inference gets 1.2–3.8× faster. This matters because it cuts the engineering cost of moving DL models to new hardware.

Key points

  • The task is framed as minimizing measured run time over a schedule space, with schedules built from Halide-style transformation primitives. It resembles hyper-parameter optimization but differs in three ways:
    • each experiment is cheap, taking seconds;
    • programs have rich internal structure the model can use;
    • there are many similar operators to tune.
  • Hardware and baselines:
    • NVIDIA TITAN X vs. cuDNN v7 and TensorComprehensions;
    • ARM Cortex-A53 vs. TFLite;
    • ARM Mali-T860 vs. ARM Compute Library v18.03.
  • End-to-end workloads were ResNet-18, MobileNet, LSTM Language Model, DQN and DCGAN, compared against MXNet v1.1 and Tensorflow v1.7 (including XLA) on GPU. Speedups ranged from 1.2× to 3.8×, partly thanks to operator fusion that fixed libraries cannot offer.
  • Three program representations were compared:
    • configuration-space features worked only within a single domain;
    • flattened AST features transferred across convolution workloads;
    • a context relation representation also transferred across operator types.
  • A preliminary experiment showed transfer from a Mali GPU to a Cortex-A53 CPU. The code is published at tvm.ai.

Interesting ideas

  • I1. Learned cost model as a cheap stand-in for real measurement. A learned predictor ranks candidates, and only the promising ones are actually run on hardware. This pattern fits any setting where trials are cheap but the search space is enormous, such as tuning data pipelines or tokenizer settings.
  • I2. Match the feature representation to how far you want to transfer. Simple configuration features suffice within one task, while more invariant structural features are needed to cross task types. This is a useful principle when deciding how abstract features should be for cross-domain transfer in general.
  • I3. Reuse models across many related jobs. Instead of tuning each operator shape from scratch, a model trained on earlier workloads warm-starts new ones. The same idea could speed up repeated tuning across similar datasets or languages.
  • I4. Exploit problem structure rather than treating it as a black box. Feeding the model the actual program structure, instead of opaque parameters, made the search far more effective. Domain-aware surrogates may likewise beat generic Bayesian optimization in other areas.
  • I5. Evaluate by wall-clock time, not by trial count. Plotting speedup against real tuning seconds, against random search and other baselines, fairly accounts for the overhead of training the model itself. This is a good evaluation habit for any optimization method that adds its own compute cost.

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, topics and concepts

📚 Shelves 3The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksComputer Vision and Pattern Recognition · 3,486 worksHardware and Architecture · 167 works

Cite

Cite as Tianqi Chen, Lianmin Zheng, Eddie Yan, Ziheng Jiang, Thierry Moreau, Luis Ceze, Carlos Guestrin and Arvind Krishnamurthy. 2019. Learning to Optimize Tensor Programs. arXiv:1805.08166.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.1805.08166 · arXiv 1805.08166

@misc{chen2019learning,
  title = {{Learning to Optimize Tensor Programs}},
  author = {Tianqi Chen, Lianmin Zheng, Eddie Yan, Ziheng Jiang, Thierry Moreau, Luis Ceze, Carlos Guestrin and Arvind Krishnamurthy},
  year = {2019},
  doi = {10.48550/arxiv.1805.08166},
  eprint = {1805.08166},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.1805.08166}
}

References

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

  1. Martin 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 1 Accordin… — TensorFlow: A system for large-scale machine learning
  2. Amit Agarwal, Eldar Akchurin, Chris Basoglu, Guoguo Chen, Scott Cyphers, Jasha Droppo, Adam Eversole, Brian Guenter, Mark Hillebrand, Ryan Hoens, Xuedong Huang, Zhiheng Huang, Vladimir Ivanov, Alexey Kamenev, Philipp Kranen, Oleksii Kuchaiev, Wolfgang Manousek, Avner May, Bhaskar Mitra, Olivier Nano… find
  3. Miltiadis Allamanis, Marc Brockschmidt, and Mahmoud Khademi. Learning to represent programs with graphs. In International Conference on Learning Representations, 2018.
  4. Frédéric Bastien, Pascal Lamblin, Razvan Pascanu, James Bergstra, Ian J. Goodfellow, Arnaud Bergeron, Nicolas Bouchard, and Yoshua Bengio. Theano: new features and speed improvements. Deep Learning and Unsupervised Feature Learning NIPS 2012 Workshop, 2012.
  5. Uday Bondhugula, Albert Hartono, J. Ramanujam, and P. Sadayappan. A practical automatic polyhedral parallelizer and locality optimizer. In Proceedings of the 29th ACM SIGPLAN Conference on Programming Language Design and Implementation, PLDI ’08, pages 101–113. ACM, 2008. find
  6. Chris Burges, Tal Shaked, Erin Renshaw, Ari Lazier, Matt Deeds, Nicole Hamilton, and Greg Hullender. Learning to rank using gradient descent. In Proceedings of the 22Nd International Conference on Machine Learning, ICML ’05, pages 89–96, New York, NY, USA, 2005. ACM.
  7. Tianqi Chen and Carlos Guestrin. XGBoost: A scalable tree boosting system. In Proceedings of the 22Nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, KDD ’16, pages 785–794, New York, NY, USA, 2016. ACM. find
  8. Tianqi Chen, Mu Li, Yutian Li, Min Lin, Naiyan Wang, Minjie Wang, Tianjun Xiao, Bing Xu, Chiyuan Zhang, , and Zheng Zhang. MXNet: A flexible and efficient machine learning library for heterogeneous distributed systems. In Neural Information Processing Systems, Workshop on Machine Learning Systems (L…
  9. Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Meghan Cowan, Haichen Shen, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. Tvm: An automated end-to-end optimizing compiler for deep learning. In 13th USENIX Symposium on Operating Systems Design and I…
  10. Xinyun Chen, Chang Liu, and Dawn Song. Tree-to-tree neural networks for program translation. CoRR, abs/1802.03691, 2018. find
  11. J.H. Friedman. Greedy function approximation: a gradient boosting machine. Annals of Statistics, 29(5):1189–1232, 2001. find
  12. M. Frigo and S. G. Johnson. Fftw: an adaptive software architecture for the fft. In Acoustics, Speech and Signal Processing, 1998. Proceedings of the 1998 IEEE International Conference on, volume 3, pages 1381–1384 vol.3, May 1998. find
  13. Daniel Golovin, Benjamin Solnik, Subhodeep Moitra, Greg Kochanski, John Karro, and D. Sculley. Google vizier: A service for black-box optimization. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, KDD ’17, pages 1487–1495. ACM, 2017. find
  14. Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. arXiv preprint arXiv:1603.05027, 2016. find
  15. Troels Henriksen, Niels G. W. Serup, Martin Elsman, Fritz Henglein, and Cosmin E. Oancea. Futhark: Purely functional gpu-programming with nested parallelism and in-place array updates. In Proceedings of the 38th ACM SIGPLAN Conference on Programming Language Design and Implementation, PLDI 2017, pag… find
  16. Andrew G. Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam. Mobilenets: Efficient convolutional neural networks for mobile vision applications. CoRR, abs/1704.04861, 2017.
  17. Frank Hutter, Holger H. Hoos, and Kevin Leyton-Brown. Sequential model-based optimization for general algorithm configuration. In Proceedings of the 5th International Conference on Learning and Intelligent Optimization, LION’05, pages 507–523, Berlin, Heidelberg, 2011. Springer-Verlag. find
  18. Frank Hutter, Lin Xu, Holger Hoos, and Kevin Leyton-Brown. Algorithm runtime prediction: Methods and evaluation (extended abstract). In Proceedings of the Twenty-Fourth International Joint Conference on Artificial Intelligence, IJCAI 2015, Buenos Aires, Argentina, July 25-31, 2015, pages 4197–4201,… find
  19. S. Kirkpatrick, C. D. Gelatt, and M. P. Vecchi. 220(4598):671–680, 1983. find
  20. Fredrik Kjolstad, Shoaib Kamil, Stephen Chou, David Lugato, and Saman Amarasinghe. The tensor algebra compiler. Proc. ACM Program. Lang., 1(OOPSLA):77:1–77:29, October 2017. find
  21. Tim Kraska, Alex Beutel, Ed H. Chi, Jeffrey Dean, and Neoklis Polyzotis. The case for learned index structures. CoRR, abs/1712.01208, 2017. — [arXiv:1712.01208](https://arxiv.org/abs/1712.01208) find
  22. Andreas Krause and Daniel Golovin. Submodular function maximization. In Tractability: Practical Approaches to Hard Problems. Cambridge University Press, February 2014. find
  23. Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton. Imagenet classification with deep convolutional neural networks. In Advances in Neural Information Processing Systems 25, pages 1097–1105. 2012.
  24. Andrew Lavin and Scott Gray. Fast algorithms for convolutional neural networks. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pages 4013–4021, 2016.
  25. Lisha Li, Kevin G. Jamieson, Giulia DeSalvo, Afshin Rostamizadeh, and Ameet Talwalkar. Efficient hyperparameter optimization and infinitely many armed bandits. CoRR, abs/1603.06560, 2016. — [arXiv:1603.06560](https://arxiv.org/abs/1603.06560) find
  26. Azalia Mirhoseini, Hieu Pham, Quoc V. Le, Benoit Steiner, Rasmus Larsen, Yuefeng Zhou, Naveen Kumar, Mohammad Norouzi, Samy Bengio, and Jeff Dean. Device placement optimization with reinforcement learning. In Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NS… find
  27. Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 518(7540):529, 2015. find
  28. Ravi Teja Mullapudi, Andrew Adams, Dillon Sharlet, Jonathan Ragan-Kelley, and Kayvon Fatahalian. Automatically scheduling halide image processing pipelines. ACM Trans. Graph., 35(4):83:1–83:11, July 2016. find
  29. George L Nemhauser, Laurence A Wolsey, and Marshall L Fisher. An analysis of approximations for maximizing submodular set functions—i. Mathematical Programming, 14(1):265–294, 1978. find
  30. Shoumik Palkar, James J. Thomas, Deepak Narayanan, Anil Shanbhag, Rahul Palamuttam, Holger Pirk, Malte Schwarzkopf, Saman P. Amarasinghe, Samuel Madden, and Matei Zaharia. Weld: Rethinking the interface between data-intensive applications. CoRR, abs/1709.06416, 2017. — [arXiv:1709.06416](https://arxiv.org/abs/1709.06416) find
  31. Alec Radford, Luke Metz, and Soumith Chintala. Unsupervised representation learning with deep convolutional generative adversarial networks. arXiv preprint arXiv:1511.06434, 2015.
  32. Jonathan Ragan-Kelley, Connelly Barnes, Andrew Adams, Sylvain Paris, Frédo Durand, and Saman Amarasinghe. 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 Language D… find
  33. B. Shahriari, K. Swersky, Z. Wang, R. P. Adams, and N. de Freitas. Taking the human out of the loop: A review of bayesian optimization. Proceedings of the IEEE, 104(1):148–175, Jan 2016. find
  34. Jasper Snoek, Hugo Larochelle, and Ryan P. Adams. Practical bayesian optimization of machine learning algorithms. In Proceedings of the 25th International Conference on Neural Information Processing Systems - Volume 2, NIPS’12, pages 2951–2959, USA, 2012. find
  35. Jasper Snoek, Oren Rippel, Kevin Swersky, Ryan Kiros, Nadathur Satish, Narayanan Sundaram, Md. Mostofa Ali Patwary, Prabhat Prabhat, and Ryan P. Adams. Scalable bayesian optimization using deep neural networks. In Proceedings of the 32Nd International Conference on International Conference on Machin… find
  36. Michel Steuwer, Toomas Remmelg, and Christophe Dubach. Lift: A functional data-parallel ir for high-performance gpu code generation. In Proceedings of the 2017 International Symposium on Code Generation and Optimization, CGO ’17, pages 74–85, Piscataway, NJ, USA, 2017. IEEE Press. find
  37. Arvind K. Sujeeth, HyoukJoong Lee, Kevin J. Brown, Hassan Chafi, Michael Wu, Anand R. Atreya, Kunle Olukotun, Tiark Rompf, and Martin Odersky. Optiml: An implicitly parallel domain-specific language for machine learning. In Proceedings of the 28th International Conference on International Conference… find
  38. Ilya Sutskever, Oriol Vinyals, and Quoc V. Le. Sequence to sequence learning with neural networks. In Proceedings of the 27th International Conference on Neural Information Processing Systems - Volume 2, NIPS’14, pages 3104–3112, Cambridge, MA, USA, 2014. MIT Press.
  39. Kai Sheng Tai, Richard Socher, and Christopher D Manning. Improved semantic representations from tree-structured long short-term memory networks. arXiv preprint arXiv:1503.00075, 2015. find
  40. Nicolas Vasilache, Oleksandr Zinenko, Theodoros Theodoridis, Priya Goyal, Zachary DeVito, William S. Moses, Sven Verdoolaege, Andrew Adams, and Albert Cohen. Tensor comprehensions: Framework-agnostic high-performance machine learning abstractions. CoRR, abs/1802.04730, 2018.
  41. Sven Verdoolaege, Juan Carlos Juega, Albert Cohen, José Ignacio Gómez, Christian Tenllado, and Francky Catthoor. Polyhedral parallel code generation for cuda. ACM Trans. Archit. Code Optim., 9(4):54:1–54:23, January 2013. find
  42. R. Clint Whaley and Jack J. Dongarra. Automatically tuned linear algebra software. In Proceedings of the 1998 ACM/IEEE Conference on Supercomputing, SC ’98, pages 1–27, Washington, DC, USA, 1998. IEEE Computer Society. find
  43. Wojciech Zaremba, Ilya Sutskever, and Oriol Vinyals. Recurrent neural network regularization. arXiv preprint arXiv:1409.2329, 2014.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/chen-2019-learning-optimize-tensor-programs/