[论文解读] On Optimizing Operator Fusion Plans for Large-Scale Machine Learning in SystemML

Matthias Böehm, Berthold Reinwald+3 more|arXiv (Cornell University)|Jan 2, 2018

Formal Methods in Verification参考文献 64被引用 5

一句话总结

本文提出了一种基于成本的优化框架，用于大规模机器学习系统中的算子融合，特别针对SystemML。该框架引入了系统化的算法，用于探索有效的融合计划、使用成本模型选择最优候选方案，并为密集、稀疏和压缩数据的本地与分布式操作生成高效代码——在最小开销下实现了相较于手写融合算子高达21倍的性能提升。

ABSTRACT

Many large-scale machine learning (ML) systems allow specifying custom ML algorithms by means of linear algebra programs, and then automatically generate efficient execution plans. In this context, optimization opportunities for fused operators---in terms of fused chains of basic operators---are ubiquitous. These opportunities include (1) fewer materialized intermediates, (2) fewer scans of input data, and (3) the exploitation of sparsity across chains of operators. Automatic operator fusion eliminates the need for hand-written fused operators and significantly improves performance for complex or previously unseen chains of operations. However, existing fusion heuristics struggle to find good fusion plans for complex DAGs or hybrid plans of local and distributed operations. In this paper, we introduce an optimization framework for systematically reason about fusion plans that considers materialization points in DAGs, sparsity exploitation, different fusion template types, as well as local and distributed operations. In detail, we contribute algorithms for (1) candidate exploration of valid fusion plans, (2) cost-based candidate selection, and (3) code generation of local and distributed operations over dense, sparse, and compressed data. Our experiments in SystemML show end-to-end performance improvements with optimized fusion plans of up to 21x compared to hand-written fused operators, with negligible optimization and code generation overhead.

研究动机与目标

• 解决在复杂DAG和混合本地/分布式工作负载中基于启发式方法的算子融合的局限性。
• 通过考虑物化点、稀疏性利用和多种融合模板类型，系统化地优化融合计划。
• 通过自动化融合避免为每种模式手动编码，从而降低开发开销。
• 在本地和分布式执行环境中，为密集、稀疏和压缩数据表示实现高效的代码生成。
• 通过有原则的、基于成本的融合规划，提升大规模机器学习工作负载的端到端性能。

提出的方法

• 提出三阶段框架：候选探索、基于成本的候选选择，以及融合算子的代码生成。
• 开发算法以在DAG中探索满足内存约束和数据布局特性的有效融合计划。
• 采用考虑物化、I/O、计算和操作链中稀疏性利用的成本模型。
• 支持多种融合模板类型，包括单遍、多聚合和稀疏性利用模式。
• 为本地和分布式执行生成优化代码，处理密集、稀疏和压缩数据表示。
• 将优化流水线集成到SystemML的编译器和运行 runtime 中，实现无缝部署。

实验结果

研究问题

• RQ1 如何在密集线性代数操作的DAG中系统化地探索融合计划，以避免指数级搜索空间？
• RQ2 何种成本模型能有效比较考虑物化、I/O、计算和稀疏性的融合计划？
• RQ3 如何在包含本地和分布式操作的混合工作负载中优化融合计划？
• RQ4 自动融合能否在性能和通用性方面超越手写融合算子？
• RQ5 如何在不物化密集中间结果的前提下，在操作链中有效利用稀疏性？

主要发现

• 所提出的优化框架在SystemML中实现了相较于手写融合算子高达21倍的端到端性能提升。
• 该框架在性能上匹配手写优化代码，同时消除了为每种融合模式手动实现的需要。
• 优化与代码生成的开销可忽略不计，使该方法在实际部署中具有可行性。
• 系统有效利用了操作链中的稀疏性，避免了昂贵的密集中间计算。
• 基于成本的选择机制在复杂DAG和混合执行计划中优于基于启发式的融合策略。
• 该框架可泛化应用于密集、稀疏和压缩数据表示，增强了其在多样化机器学习工作负载中的适用性。