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难问题，对编译器优化至关重要。然而，在目标硬件上评估调度方案优劣可能非常耗时。传统方法及以往的机器学习方法通常优化代理指标（proxy metrics），这些指标评估速度快，但在目标硬件上测试时可能导致不良调度。本文提出一种调度新方法，通过新型GFlowNet技术按代理指标比例进行采样。我们引入了一种推理阶段控制所生成调度方案多样性与质量之间权衡的技术，并通过实验证明，在目标模型上测试时，纯优化基线方法的性能可能逊于我们的方法。此外，研究表明，通过将GFlowNet对计算图进行条件化处理，可使其泛化至合成和真实编译器数据集中未见过的调度问题。