cs.AROct 10, 2026

MoSAIC: Learning DAG Scheduling Heuristics from Optimal Schedules through Motif-Aware Recovery and LLM-Guided Reasoning

Authors: Shiva Ahir, Hashmath Shaik, Alex Doboli, Nadeem Yaseen, Simona Doboli

Organizations: Department of Electrical & Computer Engineering, Stony Brook University, New York, United States of America · Barcelona Supercomputing Center, Barcelona, Catalonia, Spain · Department of Computer Science, Hofstra University, New York, United States of America

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Efficient task scheduling in heterogeneous computing systems remains a fundamental challenge in embedded and high-level synthesis flows. This paper presents MoSAIC (Motif-aware Scheduling with AI-driven Cognition), a learning-based framework that bridges the gap between exact and heuristic scheduling approaches for DAG-based applications. The system uses the reasoning capabilities of large language models (LLMs) to learn the relationships between DAG structures and scheduling priority functions. The framework first leverages the OR-Tools CP-SAT solver [4] to generate optimal schedules, which are then used to supervise the learning of compact priority-based scheduling functions through Bayesian Optimization. These learned heuristics capture relationships between scheduling decisions and structural graph features such as rank, depth, and fanout, enabling the construction of a reusable knowledge base across diverse DAG instances. Experimental results across synthetic graph families and real-world benchmarks, including PolyBench and MachSuite, show that the learned heuristics consistently improve over classical approaches such as HEFT [2], while maintaining low runtime overhead. The results indicate that combining optimal supervision with learning-based generalization provides a practical approach for scalable and adaptive task scheduling.

Figures & tables

Fig. 1: Representative recurring structural motifs in MoSAIC: fork-join, fan-out, and fan-in.
Fig. 2: CP-SAT stage in the MoSAIC pipeline. A weighted DAG is transformed into a constraint optimization model, solved for an optimal schedule, and used to generate supervision for downstream learning stages.
Fig. 3: Example of an optimal schedule generated by CP-SAT for a representative DAG instance, showing task assignments, timing, and processor-level execution.
Fig. 4: Overview of the MoSAIC framework. A corpus of synthetic and real benchmark DAGs is converted into structural task-graph representations containing node costs, dependencies, communication information, and motif descriptors. Each DAG is solved offline using CP-SAT to obtain optimal schedules, which provide supervision for Bayesian Optimization to recover interpretable priority weights θ∗\theta^{\ast} . The resulting knowledge base links graph structure, motif context, optimal scheduling traces, and recovered heuristics. A learning-based model, including an LLM refined through reinforcement learning, uses this supervision to predict scheduling behavior for unseen DAGs. At inference time, the predicted heuristic or next-task decision is used within a lightweight list scheduler to generate near-optimal schedules with much lower runtime than repeated exact optimization.
Fig. 5: Scheduling timeline comparison on the same DAG instance using HEFT, BO-recovered scheduling, and CP-SAT optimal scheduling. The learned policy reduces idle gaps and achieves makespan closer to the optimal solution.
Fig. 6: Example of the structured input representation provided to the LLM for scheduling decisions, combining graph-level features, ready-task descriptors, and heuristic context.
Fig. 7: Performance scaling across increasing graph sizes. Learned scheduling policies (BO and LLM+RL) consistently achieve lower optimality gaps than HEFT across all graph sizes.
Fig. 8: Motif-conditioned interpretation of learned priority behavior across graph families. Different local structures induce different scheduling preferences. BA-like hub-dominated regions emphasize critical-path urgency while discouraging premature fanout expansion. ER-like regions balance urgency and topological progress. WS-like clustered and merge-heavy regions assign stronger importance to depth and indegree to reduce synchronization delay. This illustrates that effective scheduling behavior is graph-structure dependent rather than governed by a single fixed priority rule.

Explore similar work

CardsList

  1. A Learning Method with Gap-Aware Generation for Heterogeneous DAG Scheduling

    Mar 24, 2026Ruisong Zhou, Haijun Zou, Li Zhou +2Neural Combinatorial OptimizationRL for Combinatorial Optimization

    School of Mathematical Science, Peking University, Beijing 100871, China · State Key Laboratory of Mathematical Sciences, Institute of Computational Mathematics and Scientific/Engineering Computing, Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing 100190, China · Theory Lab, Central Research Institute, 2012 Labs, Huawei Technologies Co., Ltd, China +1

  2. Decentralized Task Scheduling in Distributed Systems: A Deep Reinforcement Learning Approach

    Mar 25, 2026Daniel Benniah JohnDecentralized MARL

  3. PPO-STGNN: A Proximal Policy Optimization Approach with Spatio-Temporal Graph Neural Networks for DAG Task Scheduling in Cloud-Edge-End Computing

    Sep 3, 2026Yangshuo Qi, Chenwei Wang, Zihan Shen +1Edge-Cloud Computing

    Beijing University of Posts and Telecommunications, Beijing, China · Google, Mountain View, USA · NanJing University, NanJing, China