Finding the best way to schedule operations in a computation graph is a classical NP-hard problem which is central to compiler optimization. However, evaluating the goodness of a schedule on the target hardware can be very time-consuming. Traditional approaches as well as previous machine learning ones typically optimize proxy metrics, which are fast to evaluate but can lead to bad schedules when tested on the target hardware. In this work, we propose a new approach to scheduling by sampling proportionally to the proxy metric using a novel GFlowNet method. We introduce a technique to control the trade-off between diversity and goodness of the proposed schedules at inference time and demonstrate empirically that the pure optimization baselines can lead to subpar performance with respect to our approach when tested on a target model. Furthermore, we show that conditioning the GFlowNet on the computation graph enables generalization to unseen scheduling problems for both synthetic and real-world compiler datasets.

翻译：在计算图中寻找最优操作调度方案是一个经典的NP难问题，也是编译器优化的核心任务。然而，在目标硬件上评估调度方案的质量往往耗时巨大。传统方法及先前的机器学习方法通常优化代理指标——这些指标评估速度快，但在目标硬件上测试时可能导致次优调度方案。本文提出一种基于新型GFlowNet方法的新调度策略，通过按代理指标比例进行采样。我们引入了一种推理阶段控制生成调度方案多样性与质量权衡的技术，并通过实验证明，在目标模型测试中，纯优化基线方法的性能显著逊色于我们的方法。此外，研究表明，通过将GFlowNet条件化于计算图，该方法能够泛化至未见过的调度问题，适用于合成数据集和真实编译器数据集。