Accelerating Exact Combinatorial Optimization via RL-based Initialization -- A Case Study in Scheduling

TL;DR

This paper introduces a two-phase RL-to-ILP scheduling framework that significantly accelerates exact combinatorial optimization for dataflow graph scheduling, achieving comparable performance with much faster runtimes on EdgeTPU platforms.

Contribution

The paper presents a novel RL-to-ILP framework that guarantees optimality and determinism while drastically reducing runtime in scheduling problems.

Findings

Achieves up to 128× speedup over traditional exact methods.

Maintains optimal scheduling performance comparable to exact methods.

Improves on-chip inference runtime and acceleration on EdgeTPU.

Abstract

Scheduling on dataflow graphs (also known as computation graphs) is an NP-hard problem. The traditional exact methods are limited by runtime complexity, while reinforcement learning (RL) and heuristic-based approaches struggle with determinism and solution quality. This research aims to develop an innovative approach that employs machine learning (ML) for addressing combinatorial optimization problems, using scheduling as a case study. The goal is to provide guarantees in optimality and determinism while maintaining the runtime cost of heuristic methods. Specifically, we introduce a novel two-phase RL-to-ILP scheduling framework, which includes three steps: 1) RL solver acts as coarse-grain scheduler, 2) solution relaxation and 3) exact solving via ILP. Our framework demonstrates the same scheduling performance compared with using exact scheduling methods while achieving up to 128 $\times$ speed improvements. This was conducted on actual EdgeTPU platforms, utilizing ImageNet DNN computation graphs as input. Additionally, the framework offers improved on-chip inference runtime and acceleration compared to the commercially available EdgeTPU compiler.