Learning Automatic Schedulers through Projective Reparameterization
Ajay JainSaman Amarasinghe
Challenges in ML for Systems
Learning automatic schedulers through projective reparameterization Ajay Jain
Challenges in ML for Systems
1. Discrete search spaces
Learning automatic schedulers through projective reparameterization Ajay Jain
Challenges in ML for Systems
1. Discrete search spaces
2. Large, combinatorial search spaces
Learning automatic schedulers through projective reparameterization Ajay Jain
Challenges in ML for Systems
1. Discrete search spaces
2. Large, combinatorial search spaces
3. Highly constrained feasible sets
Challenging for supervised learning!
Learning automatic schedulers through projective reparameterization Ajay Jain
Motivating application: Instruction Scheduling
Learning automatic schedulers through projective reparameterization Ajay Jain
N! POSSIBLESCHEDULES
FEASIBLESCHEDULES
EFFICIENT SCHEDULES
Motivating application: Instruction Scheduling
• List scheduling with heuristics• Stochastic search & superoptimization• Integer linear programming• Reinforcement learning
Challenging for supervised learning!
Learning automatic schedulers through projective reparameterization Ajay Jain
Baseline: Sinkhorn iteration from ranking literatureà 16% of schedules are invalid
This work
• Introduce EPOCS operator
• General approach for learning under dynamic constraints
• Formulate instruction scheduling as relaxed integer program
• Imitate GCC compiler instruction schedules
Learning automatic schedulers through projective reparameterization Ajay Jain
Permutation matrix representation
Learning automatic schedulers through projective reparameterization Ajay Jain
Fixed constraints
Input dependent partial order
constraints
0 1 2 3 41 0 2 3 4
ranking, scheduling, packet switching, matching…
Learning automatic schedulers through projective reparameterization Ajay Jain
Learning automatic schedulers through projective reparameterization Ajay Jain
Sinkhorniteration
Fixed constraintsDynamic constraints
Fixed constraintsDynamic constraintsEPOCS operator
Learning automatic schedulers through projective reparameterization Ajay Jain
1) Normalize*
2) Normalize*
3) ReLU
4) Project
One iteration of EPOCS for scheduling
Learning automatic schedulers through projective reparameterization Ajay Jain
Correct the relaxation with matching(Hungarian algorithm)
Evaluation
Train POCSNet to imitate GCC 4.9.4 schedules77,202 basic blocks from SPEC2006, SPEC2017 [Mendis et al 2019]
Evaluate data dependency violations, accuracy
Baseline: Sinkhorn iteration
Learning automatic schedulers through projective reparameterization Ajay Jain
Learning automatic schedulers through projective reparameterization Ajay Jain
Fixed constraints only16% of predicted schedules are infeasible
+ Dynamic constraints with EPOCS/POPOCS
Imposing dynamic constraints reduces data dependency violations
Imposing constraints improves accuracy (+4%)
POCSNet schedule latencies are on par with GCC latencies
Learning automatic schedulers through projective reparameterization Ajay Jain
0.356
0.397
0%
5%
10%
15%
20%
25%
30%
35%
40%
45%
Accuracy of schedules
Dynamic constraintsEPOCS/POPOCS
Fixed constraintsSinkhorn iteration
Fixed constraints
Sinkhorn iteration
Dynamic constraintsEPOCS/POPOCS
Accuracy 35.6% 39.7%
Kendall tau distance 0.238 0.222
Learning automatic schedulers through projective reparameterization Ajay Jain
mov rdi, rbpmov rsi, rbxadd rsp, 0x18pop rbxpop rbp
Shuffled input block
add rsp, 0x18mov rsi, rbxmov rdi, rbppop rbxpop rbp
POCSNet scheduled block
Takeaways
• EPOCS: General purpose op for dynamic constraints on NNs
• One application: Job scheduling problems
• A step toward correct-by-construction ML for Systems:
Enforce known constraints end-to-end for accuracy boost + guarantees
Learning automatic schedulers through projective reparameterization Ajay Jain
Contact: [email protected]
Learning automatic schedulers through projective reparameterization Ajay Jain
movsd xmm5, qword ptr [rsp+0x20]movsd xmm3, qword ptr [r13+0xa0]subsd xmm3, xmm5divsd xmm3, xmm8ucomisd xmm2, xmm3
movsd xmm3, qword ptr [r13+0xa0]movsd xmm5, qword ptr [rsp+0x20]subsd xmm3, xmm5divsd xmm3, xmm8ucomisd xmm2, xmm3
Shuffled input block
POCSNet scheduled block