DAG scheduling based on diffusion model with partial-order feasibility guarantee

A central question for learning-based non-autoregressive directed acyclic graph (DAG) scheduling is how to represent and learn the solutions' global ordering structure, while guaranteeing that the solutions satisfy the partial-order constraints in the DAG. Existing approaches commonly predict the distribution of itemwise priorities, leaving the solution structure and partial-order feasibility to a post-processing feasibility repair procedure. In this paper, we propose a DAG scheduling framework based on a diffusion model in the doubly stochastic matrix (DSM) space with partial-order feasibility guarantee. The DSMs directly provide a continuous representation of scheduling solutions. We promote and ensure the feasibility of the DAG scheduling solutions generated by our framework from both the perspective of training and inference. At inference, we propose a feasible solution decoding procedure for DAG scheduling that can directly produce feasible solutions without additional post-processing repair, which applies our Birkhoff-von Neumann (BvN)-based feasible solution extraction algorithm to the diffusion model's DSM output. The efficiency of the decoding is supported by a successor interpretation of permutation matrices. At training, we further design a learning method guided by feasible components, to promote the DAG scheduling solutions' feasibility and quality. Comparison experiments and ablation studies evaluate the solution optimization performance of our framework across different DAG scheduling tasks.