Collaborative Research: CNS Core: Small: A Compilation System for Mapping Deep Learning Models to Tensorized Instructions (DELITE)

随着机器学习（ML），尤其是深度神经网络（DNN）工作负载的迅速发展，许多现有的架构已经被丰富了针对这些工作负载的指令和/或处理能力。这些指令的示例包括来自Intel的AMX指令、来自NVIDIA的Tensor核心、来自AMD的DOT指令以及许多其他指令。这种张量化指令的出现导致了许多关于如何将其用于生产级现代DNN的常见和相关挑战。目前利用这些指令集进行DNN工作负载的最新技术非常有限，现有系统要么完全缺乏对这些指令集的关注，要么无法解决复杂DNN的全局优化问题，要么在其他方面受到限制。我们工作的前提是，一个认识到最新DNN趋势并可以在不同的张量化指令集上进行优化的编译系统，将为现代ML计算提供巨大的效率提升。由此产生的议程可能会产生重大的技术、经济和社会影响。从技术方面来看，这项工作影响了高性能计算（HPC）、嵌入式系统和支持AI/ML工作负载的系统等领域。随着DNN正在成为大多数人使用的应用程序的一个组成部分，这项工作将产生巨大的经济和社会影响。在教育方面，系统和机器学习交叉点的研究将被纳入多个课程，并有助于增加各级计算教育和研究的多样性，特别是通过让代表性不足的群体参与进来。该项目解决了与现代DNN和最近和新兴的张量化指令相关的以下挑战：1）用于密集模型的局部指令选择--为了提高每个运算符的执行效率，关键的第一个问题是选择张量化指令（和相关联的数据布局），这将针对任意形状的运算符来解决。2)DNN的全局优化-在局部运算符优化之后，每个运算符可能更喜欢自己的张量化指令和数据布局，因此在整个DNN的执行期间会产生显著的数据布局转换成本。该项目制定并解决了一个全局优化问题，该问题在本地操作员执行和数据转换成本之间选择了正确的权衡。3)动态DNN的优化-该项目还考虑了现代DNN模型中各种形式的动态性，包括动态输入形状，动态控制流和动态数据结构。它提出了新的优化，如有效的内存管理，同时重新审视其他人，如本地和全局指令选择，在这些形式的动态。4)将稀疏模型映射到新兴指令--该项目还计划在涉及稀疏性时提高使用各种类型的张量化指令的效率，建立在早期工作的基础上，用于优化GPU和SIMD指令集上的SpMM（和其他稀疏计算）等内核。5)（半）自动支持新指令-为了最大限度地减少优化和编程工作，该提案还引入了一个模块，用于使用新的张量化指令或功能自动优化DNN计算。除了解决上述问题之外，该项目的一个关键组成部分将是将它们的实现与多个后端的代码生成合并到一个可重用的系统中。该系统将以计算图形表示作为输入，并输出张量和LLVM IR，从而围绕行业中广泛使用的三种表示进行构建。该奖项反映了NSF的法定使命，并通过使用基金会的知识价值和更广泛的影响审查标准进行评估，被认为值得支持。

As Machine Learning (ML), and especially Deep Neural Network (DNN) workloads have rapidly become prominent, many existing architectures have been enriched with instructions and/or processing capabilities targeting these workloads. Examples of these instructions include AMX instructions from Intel, Tensor cores from NVIDIA, DOT instructions from AMD, and many others. The emergence of such tensorized instructions is leading to many common and related challenges regarding how they can be used for production-level modern DNNs. The current state-of-the-art for exploiting these instruction sets for DNN workloads is very limited, with existing systems either completely lacking attention on these, not addressing global optimizations for complex DNNs, or being limited in other ways. The premise of our work is that a compilation system that is cognizant of the latest DNN trends and can optimize across different tensorized instruction sets, will provide large efficiency gains for modern ML computations. The resulting agenda will likely result in significant technical, economic, and societal impacts. From the technical side, the work impacts areas like High-Performance Computing (HPC), Compilers, and systems supporting AI/ML workloads. As DNNs are becoming an integral part of applications that most humans use, this work is poised to have a large economic and societal impact. On the education side, the research at the intersection of systems and ML will be incorporated into multiple courses and help to increase diversity at all levels in computing education and research, particularly by involving members from underrepresented groups.This project addresses the following challenges associated with modern DNNs and recent and emerging tensorized instructions: 1) Local Instruction Selection for Dense Models -- To improve the execution efficiency of each operator, a critical first issue is selecting tensorized instructions (and associated data layouts), which will be addressed for arbitrary shapes of operators. 2) Global Optimizations for DNNs -- After local operator optimizations, each operator may prefer its own tensorized instruction and data layout, thus incurring significant data layout transformation costs during the execution of an entire DNN. This project formulates and solves a global optimization problem that chooses the right trade-off between the local operator execution and data transformation costs. 3) Optimizations for Dynamic DNNs -- This project also considers various forms of dynamism in modern DNN models including dynamic input shapes, dynamic control flows, and dynamic data structures. It proposes new optimizations such as those for effective memory management, while revisiting others like local and global instruction selection, in the presence of these forms of dynamism. 4) Mapping Sparse Models to Emerging Instructions -- This project also plans to improve the efficiency of using various types of tensorized instructions when sparsity is involved, building on top of earlier work for optimizing kernels like SpMM (and other sparse computations) on GPUs and SIMD instruction sets. 5) (Semi-) Automatic Support for New Instructions -- To minimize the optimization and programming effort, this proposal also introduces a module to automatically optimize DNN computations with new tensorized instructions or features. Besides addressing the above problems, one critical component of this project will be incorporating their implementations, together with code generation for multiple back-ends, in a reusable system. This system will take as the input the Computational Graph representation, and output Tensor and LLVM IRs, thus building around three representations widely used in the industry.This award reflects NSF's statutory mission and has been deemed worthy of support through evaluation using the Foundation's intellectual merit and broader impacts review criteria.