GauS: Differentiable Scheduling Optimization via Gaussian Reparameterization

Abstract: Efficient operator scheduling is a fundamental challenge in software compilation and hardware synthesis. While recent differentiable approaches have sought to replace traditional ones like exact solvers or heuristics with gradient-based search, they typically rely on categorical distributions that fail to capture the ordinal nature of time and suffer from a parameter space that scales poorly. In this paper, we propose a novel differentiable framework, GauS, that models operator scheduling as a stochastic relaxation using Gaussian distributions, which fully utilize modern parallel computing devices like GPUs. By representing schedules as continuous Gaussian variables, we successfully capture the ordinal nature of time and reduce the optimization space by orders of magnitude. Our method is highly flexible to represent various objectives and constraints, which provides the first differentiable formulation for the complex pipelined scheduling problem. We evaluate our method on a range of benchmarks, demonstrating that Gaus achieves Pareto-optimal results.

• The paper introduces GauS, a Gaussian relaxation that represents each operator with only a mean and standard deviation, reducing scheduling parameters from O(D|V|) to O(|V|).
• The method combines differentiable objective expectations with augmented Lagrangian constraints, GPU vectorization, informed initialization, rounding, and lightweight legalization to optimize regular and modulo schedules.
• GauS improves geometric-mean quality over GS-Schedule by 71.8% where both are feasible, runs one to two orders of magnitude faster than commercial solvers, and avoids the memory failures seen on large graphs.

GauS is a differentiable framework for operator scheduling that replaces categorical relaxations with a Gaussian reparameterization of each operator's start time. The work targets the scalability bottleneck of prior differentiable schedulers—most notably GS-Schedule, which parameterizes each operator with a $D$-dimensional categorical distribution over time steps—and demonstrates that a two-parameter-per-node continuous relaxation achieves comparable or better solution quality while being one to two orders of magnitude faster and dramatically more memory-efficient.

Motivation and problem setting

Operator scheduling assigns operations of a DAG $G=(V,E)$ to discrete time steps under dependency and latency constraints, optimizing objectives such as peak resource usage, communication overhead, or memory footprint. Exact ILP/SMT formulations guarantee optimality but scale exponentially; heuristics such as list scheduling and force-directed scheduling (FDS) are fast but greedy and local. Differentiable combinatorial optimization offers a middle path, but the existing approach, GS-Schedule, has two structural weaknesses. First, it is ordinally blind: a categorical distribution treats step $d$ and step $d+1$ as unrelated nominal labels, so shifting probability mass across zero-valued buckets yields vanishing gradients. Second, its parameter space scales as $O(D\cdot|V|)$, which becomes prohibitive for modern workloads with tens of thousands of nodes and large depths, and it fails to exploit GPU parallelism.

Gaussian reparameterization

The core idea is to model each operator's execution step as an independent random variable $X_i \sim \mathcal{N}(\mu_i, \sigma_i^2)$, where $\mu_i$ encodes the expected schedule step and $\sigma_i$ controls the degree of relaxation. This reduces the optimization space from $\mathbb{R}^{D\cdot|V|}$ to just $2|V|$ parameters—a reduction of orders of magnitude for deep graphs. Because the Gaussian PDF is continuous and unimodal, gradient updates on $G=(V,E)$0 move probability mass smoothly through adjacent steps, providing informative local-refinement signals that categorical formulations lack.

Discrete placement probabilities are obtained by integrating the PDF over unit intervals centered at integer steps, expressed via the standard normal CDF $G=(V,E)$1. All objectives and constraints are then written as differentiable expectations over these probabilities:

• Dependency violations: summed products of placement probabilities across all invalid $G=(V,E)$2 pairs per edge.
• Memory footprint: expected register pressure at step $G=(V,E)$3 combines the probability that a producer has started with the probability that no successor has; the global peak uses a LogSumExp smooth-max with temperature $G=(V,E)$4.
• Resource usage and communication overhead: analogous expectation derivations, with ReLU-based violation terms when resource caps are imposed.
• Modulo scheduling: reservation-table probabilities are wrapped modulo the initiation interval (II), and recurrence constraints on back-edges are handled with expected-violation sums analogous to dependencies.

This yields what the authors state is the first differentiable formulation of pipelined (modulo) scheduling.

Optimization procedure

Constraint handling uses an augmented Lagrangian method (ALM) rather than static penalty weights: multipliers $G=(V,E)$5 are initialized near zero and updated as $G=(V,E)$6 after each iteration, progressively enforcing hard constraints. Discrete schedules are extracted by rounding $G=(V,E)$7; any residual legality violations are repaired by lightweight greedy legalization passes (topological-order clamping for regular scheduling, fixed-point relaxation for modulo scheduling), and legalized schedules re-initialize the continuous parameters. A notable advantage of the representation is informed initialization: means can be seeded at the midpoint of ASAP/ALAP bounds, with standard deviations proportional to each node's scheduling slack ($G=(V,E)$8). The entire pipeline is vectorized in PyTorch for GPU execution.

Evaluation

Experiments cover three formulations: (A) latency-constrained resource-plus-communication optimization, (B) latency-constrained memory-footprint minimization, and (C) resource- and recurrence-constrained modulo scheduling minimizing pipelined memory footprint. Benchmarks include EPFL circuits (up to 57,375 nodes, depth 4,373) and larger synthetic random workloads (up to ~9,400 nodes), with a uniform 15-minute limit against CPLEX/Gurobi on SDC+LP formulations, list scheduling, FDS, and GS-Schedule.

Key results:

• Formulation A: on instances where both methods produce feasible solutions, GauS improves geometric-mean solution quality over GS-Schedule by 71.8%; GS-Schedule frequently exceeds available CUDA memory (OOM) on large EPFL graphs such as square, multiplier, and div, while GauS completes all instances.
• Formulations B and C: GauS matches or exceeds heuristic quality, with list scheduling and FDS typically producing 20%–60% higher memory footprints; exact solvers remain competitive only on small graphs and often fail to find feasible solutions within the time limit at scale.
• Anytime behavior: GauS reaches equivalent quality one to two orders of magnitude faster than commercial solvers and sits on the Pareto frontier of quality versus runtime for most benchmarks.
• GPU efficiency: profiling shows GauS maintains near-100% average GPU utilization and roughly linear memory growth in $G=(V,E)$9, whereas GS-Schedule drops below 40% utilization as complexity increases.

The margin over heuristics narrows on large Formulation C instances, which the authors attribute to the proportional addition of back-edges constraining the feasible space and reducing degrees of freedom—an honest caveat about where the method's advantage diminishes.

Limitations and open questions

The framework models operators as independent Gaussians, ignoring correlations between nodes—particularly those sharing edges—which can cause convergence to suboptimal solutions. Capturing such correlations via a Gaussian process would incur $d$0 complexity, which the authors deem prohibitive; efficient correlation modeling remains open. Additionally, the legalization heuristics are greedy repairs that may perturb optimized schedules, and the evaluation relies on synthetically augmented recurrence constraints for modulo benchmarks due to the absence of public large-scale pipelined datasets. Whether ALM-based constraint enforcement generalizes robustly to tighter II targets or heterogeneous latency operators is not established here.

Conclusion

GauS demonstrates that a minimal Gaussian stochastic relaxation—two parameters per operator—suffices to make differentiable scheduling scalable to tens of thousands of nodes, delivering Pareto-optimal quality-speed tradeoffs, first-of-kind differentiable modulo scheduling, and substantially better GPU utilization than categorical predecessors. Its main unresolved question is how to incorporate inter-node dependence structure without sacrificing linear parameter scaling.