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

Abstract

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. Although AStitch serves as a stand-alone compiler engine that is portable to any version of TensorFlow, its basic ideas can be generally applied to other ML frameworks and optimization compilers. Experimental results show that AStitch can

Optimizing Direct Convolutions on ARM Multi-Cores

“…On ThunderX2, nDirect delivers slightly lower performance for the end-to-end inference compared to Ansor, with a speedup of 0.88× to 0.98×. The better performance of Ansor on the whole CNN is due to its ability to optimize across CNN layers through operator fusion [67,72]. This technique can write back operations for intermediate results and fetch operations in the CNNs pipeline, further reducing memory access latency and bandwidth pressure to improve CNNs end-toend performance.…”

Pame

“…We use blocks as the granularity to add exits, because the block is the basic unit/step to compress information to a low dimensional space mathematically [13]. In addition, many graph optimizations are done between layers in the same block [14,53]. Adding an exit in a block invalidates the graph optimization, and incurs the high cost of data movement between stages, as there are high data dependencies within a block.…”

RCM: Residue-aware Consolidation for Heterogeneous MLaaS Cluster

“…Mlmodelci [20] provided a one-stop platform for multimedia developers to provide efficient ML services, while NSML [21] created a collaborative environment for users to deploy their own commercial service. AStitch [22] improved the execution efficiency of ML tasks by means of compilation optimization, and avoided unnecessary repeated calculation. FedAMP [23] promoted pairwise collaboration between clients with similar data to significantly improve federated learning performance.…”