Can I read AStitch: enabling a new multi-dimensional optimization space for memory-intensive ML training and inference on modern SIMT architectures on EtoBox?

AStitch: enabling a new multi-dimensional optimization space for memory-intensive ML training and inference on modern SIMT architectures by Zhen Zheng; Xuanda Yang; Pengzhan Zhao; Guoping Long; Kai Zhu; Feiwen Zhu; Wenyi Zhao; Xiaoyong Liu; Jun Yang; Jidong Zhai; Shuaiwen Leon Song; Wei Lin is a scholarly article available to read on EtoBox.

What is AStitch: enabling a new multi-dimensional optimization space for memory-intensive ML training and inference on modern SIMT architectures about?

This work reveals that memory-intensive computation is a rising performance-critical factor in recent machine learning models. Due to a unique set of new challenges, existing ML optimizing compilers cannot perform efficient fusion under complex two-level dependencies combined with just-in-time demand. They face the dilemma of either performing costly fusion due to heavy redundant computation, or skipping fusion which results in massive number of kernels. Furthermore, they often suffer from low parallelism due to the lack of support for real-world production workloads with irregular tensor shapes. To address these rising challenges, we propose AStitch, a machine learning optimizing compiler that opens a new multi-dimensional optimization space for memory-intensive ML computations. It systematically abstracts four operator-stitching schemes while considering multi-dimensional optimization objectives, tackles complex computation graph dependencies with novel hierarchical data reuse, and efficiently processes various tensor shapes via adaptive thread mapping. Finally, AStitch provides justin-time support incorporating our proposed optimizations for both ML training and inference. Altho

Author

Zhen Zheng; Xuanda Yang; Pengzhan Zhao; Guoping Long; Kai Zhu; Feiwen Zhu; Wenyi Zhao; Xiaoyong Liu; Jun Yang; Jidong Zhai; Shuaiwen Leon Song; Wei Lin

Publisher

ACM

Published

2022

Language

EN