10. Kernel Scheduling
The branch-free executor krun retires one resident-weight MAC per step with no
notion of latency: it reads its operands and updates the accumulator in the same
cycle. Real hardware reads weights and streamed operands from memory with a
fixed pipeline latency, so a product issued at cycle t is not accumulated
until cycle t + L. This chapter gives that pipeline its own cycle-level
scheduler and proves it computes exactly what the branch-free executor does, for
any latency L.
The scheduler is an in-order shift-register pipeline. Each cycle it issues one
product into the tail of a length-L register of in-flight products, and
retires the product falling off the head into the committed accumulator. The
correctness statement is a conservation law: at every cycle, the committed
accumulator plus the sum of the in-flight products equals the running dot
product. Nothing is lost and nothing is double-counted by the scheduling, so
reading the pipeline out at the end reproduces krun regardless of L.
namespace Janus
-- The product issued at kernel step `t`: the resident weight and the streamed
-- operand at the two pointers, advanced by `t`.
def kprod (s : KernelState) (t : Nat) : Int :=
s.weights (s.wptr + t) * s.stream (s.sptr + t)
def sumUpto (f : Nat -> Int) : Nat -> Int
| 0 => 0
| k + 1 => sumUpto f k + f k
def prodSum (s : KernelState) (n : Nat) : Int := sumUpto (kprod s) n
The branch-free accumulator, written as the sum of the issued products. This is
the same value the kernel chapter proves equals the source dot product; here it
is the specification the pipeline must match. The bridge below rewrites the
front-recursive fsumFrom of that chapter into this tail-recursive product sum.
theorem fsumFrom_succ (w a : Nat -> Int) (wp sp n : Nat) :
fsumFrom w a wp sp (n + 1)
= fsumFrom w a wp sp n + w (wp + n) * a (sp + n) := w:Nat → Inta:Nat → Intwp:Natsp:Natn:Nat⊢ fsumFrom w a wp sp (n + 1) = fsumFrom w a wp sp n + w (wp + n) * a (sp + n)
induction n generalizing wp sp with
w:Nat → Inta:Nat → Intwp:Natsp:Nat⊢ fsumFrom w a wp sp (0 + 1) = fsumFrom w a wp sp 0 + w (wp + 0) * a (sp + 0) All goals completed! 🐙
w:Nat → Inta:Nat → Intn:Natih:∀ (wp sp : Nat), fsumFrom w a wp sp (n + 1) = fsumFrom w a wp sp n + w (wp + n) * a (sp + n)wp:Natsp:Nat⊢ fsumFrom w a wp sp (n + 1 + 1) = fsumFrom w a wp sp (n + 1) + w (wp + (n + 1)) * a (sp + (n + 1))
w:Nat → Inta:Nat → Intn:Natih:∀ (wp sp : Nat), fsumFrom w a wp sp (n + 1) = fsumFrom w a wp sp n + w (wp + n) * a (sp + n)wp:Natsp:Nat⊢ w wp * a sp + fsumFrom w a (wp + 1) (sp + 1) (n + 1) =
w wp * a sp + fsumFrom w a (wp + 1) (sp + 1) n + w (wp + (n + 1)) * a (sp + (n + 1))
w:Nat → Inta:Nat → Intn:Natih:∀ (wp sp : Nat), fsumFrom w a wp sp (n + 1) = fsumFrom w a wp sp n + w (wp + n) * a (sp + n)wp:Natsp:Nat⊢ w wp * a sp + (fsumFrom w a (wp + 1) (sp + 1) n + w (wp + 1 + n) * a (sp + 1 + n)) =
w wp * a sp + fsumFrom w a (wp + 1) (sp + 1) n + w (wp + (n + 1)) * a (sp + (n + 1))
w:Nat → Inta:Nat → Intn:Natih:∀ (wp sp : Nat), fsumFrom w a wp sp (n + 1) = fsumFrom w a wp sp n + w (wp + n) * a (sp + n)wp:Natsp:Nathw:wp + 1 + n = wp + (n + 1)⊢ w wp * a sp + (fsumFrom w a (wp + 1) (sp + 1) n + w (wp + 1 + n) * a (sp + 1 + n)) =
w wp * a sp + fsumFrom w a (wp + 1) (sp + 1) n + w (wp + (n + 1)) * a (sp + (n + 1))
w:Nat → Inta:Nat → Intn:Natih:∀ (wp sp : Nat), fsumFrom w a wp sp (n + 1) = fsumFrom w a wp sp n + w (wp + n) * a (sp + n)wp:Natsp:Nathw:wp + 1 + n = wp + (n + 1)hs:sp + 1 + n = sp + (n + 1)⊢ w wp * a sp + (fsumFrom w a (wp + 1) (sp + 1) n + w (wp + 1 + n) * a (sp + 1 + n)) =
w wp * a sp + fsumFrom w a (wp + 1) (sp + 1) n + w (wp + (n + 1)) * a (sp + (n + 1))
w:Nat → Inta:Nat → Intn:Natih:∀ (wp sp : Nat), fsumFrom w a wp sp (n + 1) = fsumFrom w a wp sp n + w (wp + n) * a (sp + n)wp:Natsp:Nathw:wp + 1 + n = wp + (n + 1)hs:sp + 1 + n = sp + (n + 1)⊢ w wp * a sp + (fsumFrom w a (wp + 1) (sp + 1) n + w (wp + (n + 1)) * a (sp + (n + 1))) =
w wp * a sp + fsumFrom w a (wp + 1) (sp + 1) n + w (wp + (n + 1)) * a (sp + (n + 1))
All goals completed! 🐙
theorem prodSum_eq_fsumFrom (s : KernelState) (n : Nat) :
prodSum s n = fsumFrom s.weights s.stream s.wptr s.sptr n := s:KernelStaten:Nat⊢ prodSum s n = fsumFrom s.weights s.stream s.wptr s.sptr n
induction n with
s:KernelState⊢ prodSum s 0 = fsumFrom s.weights s.stream s.wptr s.sptr 0 All goals completed! 🐙
s:KernelStaten:Natih:prodSum s n = fsumFrom s.weights s.stream s.wptr s.sptr n⊢ prodSum s (n + 1) = fsumFrom s.weights s.stream s.wptr s.sptr (n + 1)
s:KernelStaten:Natih:prodSum s n = fsumFrom s.weights s.stream s.wptr s.sptr nh:prodSum s (n + 1) = prodSum s n + kprod s n⊢ prodSum s (n + 1) = fsumFrom s.weights s.stream s.wptr s.sptr (n + 1)
s:KernelStaten:Natih:prodSum s n = fsumFrom s.weights s.stream s.wptr s.sptr nh:prodSum s (n + 1) = prodSum s n + kprod s n⊢ fsumFrom s.weights s.stream s.wptr s.sptr n + kprod s n =
fsumFrom s.weights s.stream s.wptr s.sptr n + s.weights (s.wptr + n) * s.stream (s.sptr + n)
All goals completed! 🐙
theorem krun_acc_prodSum (n : Nat) (s : KernelState) :
(krun n s).acc = s.acc + prodSum s n := n:Nats:KernelState⊢ (krun n s).acc = s.acc + prodSum s n
All goals completed! 🐙
The pipeline state is a committed accumulator paired with the in-flight products
still in the read/multiply registers. fifoSum reads the total value currently
in flight.
def fifoSum : List Int -> Int
| [] => 0
| x :: xs => x + fifoSum xs
theorem fifoSum_append (l : List Int) (v : Int) :
fifoSum (l ++ [v]) = fifoSum l + v := l:List Intv:Int⊢ fifoSum (l ++ [v]) = fifoSum l + v
induction l with
v:Int⊢ fifoSum ([] ++ [v]) = fifoSum [] + v All goals completed! 🐙
v:Intx:Intxs:List Intih:fifoSum (xs ++ [v]) = fifoSum xs + v⊢ fifoSum (x :: xs ++ [v]) = fifoSum (x :: xs) + v
v:Intx:Intxs:List Intih:fifoSum (xs ++ [v]) = fifoSum xs + v⊢ x + fifoSum (xs ++ [v]) = x + fifoSum xs + v
v:Intx:Intxs:List Intih:fifoSum (xs ++ [v]) = fifoSum xs + v⊢ x + (fifoSum xs + v) = x + fifoSum xs + v
All goals completed! 🐙
abbrev Pipe := Int × List Int
-- One cycle: retire the head product into the accumulator, shift the register,
-- and issue `newProd` at the tail. An empty register (latency zero) retires the
-- issued product immediately.
def pipeCycle (newProd : Int) (p : Pipe) : Pipe :=
match p with
| (acc, []) => (acc + newProd, [])
| (acc, x :: xs) => (acc + x, xs ++ [newProd])
def pipeRun (feed : Nat -> Int) : Nat -> Pipe -> Pipe
| 0, p => p
| k + 1, p => pipeCycle (feed k) (pipeRun feed k p)
Each cycle conserves value: what leaves the register enters the accumulator, and what enters the register is exactly the issued product.
theorem pipeCycle_conserves (newProd : Int) (p : Pipe) :
(pipeCycle newProd p).1 + fifoSum (pipeCycle newProd p).2
= p.1 + fifoSum p.2 + newProd := newProd:Intp:Pipe⊢ (pipeCycle newProd p).fst + fifoSum (pipeCycle newProd p).snd = p.fst + fifoSum p.snd + newProd
newProd:Intacc:Intfifo:List Int⊢ (pipeCycle newProd (acc, fifo)).fst + fifoSum (pipeCycle newProd (acc, fifo)).snd =
(acc, fifo).fst + fifoSum (acc, fifo).snd + newProd
cases fifo with
newProd:Intacc:Int⊢ (pipeCycle newProd (acc, [])).fst + fifoSum (pipeCycle newProd (acc, [])).snd =
(acc, []).fst + fifoSum (acc, []).snd + newProd All goals completed! 🐙
newProd:Intacc:Intx:Intxs:List Int⊢ (pipeCycle newProd (acc, x :: xs)).fst + fifoSum (pipeCycle newProd (acc, x :: xs)).snd =
(acc, x :: xs).fst + fifoSum (acc, x :: xs).snd + newProd
newProd:Intacc:Intx:Intxs:List Int⊢ acc + x + (fifoSum xs + newProd) = acc + (x + fifoSum xs) + newProd
All goals completed! 🐙
theorem pipe_conservation (feed : Nat -> Int) (k : Nat) (p : Pipe) :
(pipeRun feed k p).1 + fifoSum (pipeRun feed k p).2
= p.1 + fifoSum p.2 + sumUpto feed k := feed:Nat → Intk:Natp:Pipe⊢ (pipeRun feed k p).fst + fifoSum (pipeRun feed k p).snd = p.fst + fifoSum p.snd + sumUpto feed k
induction k with
feed:Nat → Intp:Pipe⊢ (pipeRun feed 0 p).fst + fifoSum (pipeRun feed 0 p).snd = p.fst + fifoSum p.snd + sumUpto feed 0 All goals completed! 🐙
feed:Nat → Intp:Pipek:Natih:(pipeRun feed k p).fst + fifoSum (pipeRun feed k p).snd = p.fst + fifoSum p.snd + sumUpto feed k⊢ (pipeRun feed (k + 1) p).fst + fifoSum (pipeRun feed (k + 1) p).snd = p.fst + fifoSum p.snd + sumUpto feed (k + 1)
feed:Nat → Intp:Pipek:Natih:(pipeRun feed k p).fst + fifoSum (pipeRun feed k p).snd = p.fst + fifoSum p.snd + sumUpto feed k⊢ (pipeCycle (feed k) (pipeRun feed k p)).fst + fifoSum (pipeCycle (feed k) (pipeRun feed k p)).snd =
p.fst + fifoSum p.snd + (sumUpto feed k + feed k)
feed:Nat → Intp:Pipek:Natih:(pipeRun feed k p).fst + fifoSum (pipeRun feed k p).snd = p.fst + fifoSum p.snd + sumUpto feed k⊢ p.fst + fifoSum p.snd + sumUpto feed k + feed k = p.fst + fifoSum p.snd + (sumUpto feed k + feed k)
All goals completed! 🐙
The scheduler issues the n kernel products in order and then reads out the
pipeline. By the conservation law this reproduces the branch-free accumulator
for every latency L, so pipelining the MAC changes when a product is committed
but never what the kernel computes.
def schedResult (s : KernelState) (L n : Nat) : Int :=
let final := pipeRun (kprod s) n (s.acc, List.replicate L 0)
final.1 + fifoSum final.2
theorem fifoSum_replicate_zero (L : Nat) :
fifoSum (List.replicate L 0) = 0 := L:Nat⊢ fifoSum (List.replicate L 0) = 0
induction L with
⊢ fifoSum (List.replicate 0 0) = 0 All goals completed! 🐙
L:Natih:fifoSum (List.replicate L 0) = 0⊢ fifoSum (List.replicate (L + 1) 0) = 0 All goals completed! 🐙
theorem schedResult_eq_krun (s : KernelState) (L n : Nat) :
schedResult s L n = (krun n s).acc := s:KernelStateL:Natn:Nat⊢ schedResult s L n = (krun n s).acc
s:KernelStateL:Natn:Nathc:(pipeRun (kprod s) n (s.acc, List.replicate L 0)).fst + fifoSum (pipeRun (kprod s) n (s.acc, List.replicate L 0)).snd =
(s.acc, List.replicate L 0).fst + fifoSum (s.acc, List.replicate L 0).snd + sumUpto (kprod s) n⊢ schedResult s L n = (krun n s).acc
s:KernelStateL:Natn:Nathc:(pipeRun (kprod s) n (s.acc, List.replicate L 0)).fst + fifoSum (pipeRun (kprod s) n (s.acc, List.replicate L 0)).snd =
(s.acc, List.replicate L 0).fst + 0 + sumUpto (kprod s) n⊢ schedResult s L n = (krun n s).acc
s:KernelStateL:Natn:Nathc:(pipeRun (kprod s) n (s.acc, List.replicate L 0)).fst + fifoSum (pipeRun (kprod s) n (s.acc, List.replicate L 0)).snd =
(s.acc, List.replicate L 0).fst + 0 + sumUpto (kprod s) n⊢ (have final := pipeRun (kprod s) n (s.acc, List.replicate L 0);
final.fst + fifoSum final.snd) =
(krun n s).acc
s:KernelStateL:Natn:Nathc:(pipeRun (kprod s) n (s.acc, List.replicate L 0)).fst + fifoSum (pipeRun (kprod s) n (s.acc, List.replicate L 0)).snd =
(s.acc, List.replicate L 0).fst + 0 + sumUpto (kprod s) n⊢ (s.acc, List.replicate L 0).fst + 0 + sumUpto (kprod s) n = s.acc + prodSum s n
s:KernelStateL:Natn:Nathc:(pipeRun (kprod s) n (s.acc, List.replicate L 0)).fst + fifoSum (pipeRun (kprod s) n (s.acc, List.replicate L 0)).snd =
(s.acc, List.replicate L 0).fst + 0 + sumUpto (kprod s) n⊢ s.acc + 0 + sumUpto (kprod s) n = s.acc + prodSum s n
s:KernelStateL:Natn:Nathc:(pipeRun (kprod s) n (s.acc, List.replicate L 0)).fst + fifoSum (pipeRun (kprod s) n (s.acc, List.replicate L 0)).snd =
(s.acc, List.replicate L 0).fst + 0 + sumUpto (kprod s) n⊢ s.acc + 0 + sumUpto (kprod s) n = s.acc + sumUpto (kprod s) n
All goals completed! 🐙
-- The number of cycles to fully retire `n` MACs at latency `L`: `n` issue
-- cycles followed by an `L`-cycle drain of the register.
def macCycles (L n : Nat) : Nat := n + L
theorem macCycles_zero_latency (n : Nat) : macCycles 0 n = n := rfl
end Janus
The conservation law is the load-bearing statement. It holds cycle by cycle and
for any feed sequence, so it is exactly the invariant a synchronous-read cell
needs: whatever the memory-read latency of a bound SRAM macro, the committed
accumulator plus the in-flight products is always the running dot product, and
draining the pipeline commits the branch-free result. A latency-zero schedule
keeps nothing in flight, recovering krun cycle for cycle; a latency-L
schedule takes L extra cycles to fill and drain but retires the identical sum.
Making a full synchronous-read cell cycle-faithful is future work, but its
kernel path already has its correctness invariant here.