[论文解读] Stripe: Tensor Compilation via the Nested Polyhedral Model

Tim Zerrell, Jeremy Bruestle|arXiv (Cornell University)|Mar 14, 2019

Parallel Computing and Optimization Techniques参考文献 20被引用 25

一句话总结

本文提出 Stripe，一种基于嵌套多面体模型的机器学习编译器专用中间表示（IR），可实现跨多种硬件的自动高性能内核生成。通过使用分层多面体迭代空间建模张量计算，Stripe 将算法设计、优化阶段与硬件目标解耦，支持模块化、可重用的优化，以及面向现代机器学习加速器的可扩展代码生成。

ABSTRACT

Hardware architectures and machine learning (ML) libraries evolve rapidly. Traditional compilers often fail to generate high-performance code across the spectrum of new hardware offerings. To mitigate, engineers develop hand-tuned kernels for each ML library update and hardware upgrade. Unfortunately, this approach requires excessive engineering effort to scale or maintain with any degree of state-of-the-art performance. Here we present a Nested Polyhedral Model for representing highly parallelizable computations with limited dependencies between iterations. This model provides an underlying framework for an intermediate representation (IR) called Stripe, amenable to standard compiler techniques while naturally modeling key aspects of modern ML computing. Stripe represents parallelism, efficient memory layout, and multiple compute units at a level of abstraction amenable to automatic optimization. We describe how Stripe enables a compiler for ML in the style of LLVM that allows independent development of algorithms, optimizations, and hardware accelerators. We also discuss the design exploration advantages of Stripe over kernel libraries and schedule-based or schedule-space-based code generation.

研究动机与目标

• 解决为支持机器学习算法和硬件加速器快速演进而产生的手写内核组合爆炸问题。
• 将算法开发、优化阶段与硬件特定代码生成解耦，以降低维护开销并提高可扩展性。
• 通过使用结构化、可扩展的 IR 建模复杂并行性和内存访问模式，实现对机器学习工作负载的自动高性能代码生成。
• 为类似 LLVM 的编译器栈提供基础，使优化阶段可在不同操作和硬件平台之间复用和组合。
• 通过引入数学基础扎实、分层的张量计算表示，克服通用编译器和现有机器学习专用框架的局限性。

提出的方法

• 提出嵌套多面体模型作为形式化方法，用于表示具有有限迭代间依赖关系的高度可并行张量计算。
• 设计 Stripe 作为领域专用 IR，利用多面体约束和嵌套结构捕获迭代空间、内存布局和计算单元抽象。
• 构建 IR 以支持分层分块和调度，实现对多级内存和计算单元的优化。
• 将 Stripe 集成到 PlaidML 编译器栈中，接收来自 Tile（一种高层 IR）的下探后张量操作，随后应用与硬件无关的优化阶段。
• 在 Stripe IR 上使用标准编译技术（如分块、向量化和内存布局变换），且优化独立于目标硬件执行。
• 通过允许基于 Stripe IR 构建新优化阶段，利用嵌套多面体模型的形式化结构，实现可扩展性。

实验结果

研究问题

• RQ1 如何设计一种编译器 IR，以建模现代机器学习工作负载中复杂的分层并行性和内存访问模式，同时保持对自动优化的友好性？
• RQ2 嵌套多面体模型能否有效表示卷积等操作中常见的非矩形迭代空间和边界条件？
• RQ3 在编译器栈中，能否在不牺牲性能的前提下，实现算法、优化和硬件之间的充分关注分离？
• RQ4 与 Halide、TVM 或 MLIR 等现有方法相比，Stripe IR 在表达能力、模块化和性能可移植性方面表现如何？
• RQ5 Stripe IR 是否能实现对多种硬件平台的自动高性能代码生成，而无需为每个操作与架构组合手动调优内核？

主要发现

• 嵌套多面体模型通过在多面体框架内建模保护区域（halo regions）和边界约束，成功捕获了复杂且非矩形的迭代空间（如非逐点卷积中的情况）。
• Stripe 的嵌套块结构天然支持多级分块，可高效映射工作负载至分层内存系统和多个计算单元。
• 该 IR 允许多个机器学习操作和硬件目标之间独立开发和重用优化阶段，显著减少对手动优化内核的需求。
• Stripe 实现了与手写实现性能相当的自动内核生成，初步的 GPU 基准测试已验证其有效性。
• 将 Stripe 集成到 PlaidML 的编译流程中，实现了与硬件无关的编译，且优化阶段可复用、可组合。
• 作者预计未来将 Stripe 作为 MLIR 的方言集成，将进一步增强互操作性，并支持在不同编译器生态间共享更广泛的优化阶段。