Janus

9. Kernel Execution🔗

Janus has ordinary control flow, but the core throughput path is intentionally branch-free: a straight sequence of resident-weight MACs over a streamed operand block. This chapter gives that path its own compact executor and proves that it computes the source dot product for all inputs.

namespace Janus

structure KernelState where
  acc : Int
  weights : Nat -> Int
  stream : Nat -> Int
  wptr : Nat
  sptr : Nat

def fsumFrom (w a : Nat -> Int) (wp sp : Nat) : Nat -> Int
  | 0 => 0
  | n + 1 => w wp * a sp + fsumFrom w a (wp + 1) (sp + 1) n

def kstep (s : KernelState) : KernelState :=
  let x := s.weights s.wptr
  let y := s.stream s.sptr
  { s with acc := s.acc + x * y, wptr := s.wptr + 1, sptr := s.sptr + 1 }

def krun : Nat -> KernelState -> KernelState
  | 0, s => s
  | n + 1, s => krun n (kstep s)

def srun : Nat -> KernelState -> KernelState := krun

theorem run_eq_srun_branchFree (n : Nat) (s : KernelState) :
    krun n s = srun n s := rfl

The accumulator invariant says exactly what has happened after n kernel cycles: the initial accumulator plus the dot product of the n resident and streamed operands starting at the two pointers.

theorem krun_acc : forall (n : Nat) (s : KernelState),
    (krun n s).acc =
      s.acc + fsumFrom s.weights s.stream s.wptr s.sptr n
  | 0, s =>s:KernelState⊢ (krun 0 s).acc = s.acc + fsumFrom s.weights s.stream s.wptr s.sptr 0 bys:KernelState⊢ (krun 0 s).acc = s.acc + fsumFrom s.weights s.stream s.wptr s.sptr 0
      simp [krun, fsumFrom]All goals completed! 🐙
  | n + 1, s =>n:Nats:KernelState⊢ (krun (n + 1) s).acc = s.acc + fsumFrom s.weights s.stream s.wptr s.sptr (n + 1) byn:Nats:KernelState⊢ (krun (n + 1) s).acc = s.acc + fsumFrom s.weights s.stream s.wptr s.sptr (n + 1)
      rw [krun,n:Nats:KernelState⊢ (krun n (kstep s)).acc = s.acc + fsumFrom s.weights s.stream s.wptr s.sptr (n + 1) krun_acc nn:Nats:KernelState⊢ (kstep s).acc + fsumFrom (kstep s).weights (kstep s).stream (kstep s).wptr (kstep s).sptr n =
  s.acc + fsumFrom s.weights s.stream s.wptr s.sptr (n + 1)]n:Nats:KernelState⊢ (kstep s).acc + fsumFrom (kstep s).weights (kstep s).stream (kstep s).wptr (kstep s).sptr n =
  s.acc + fsumFrom s.weights s.stream s.wptr s.sptr (n + 1)
      change s.acc + s.weights s.wptr * s.stream s.sptr +
          fsumFrom s.weights s.stream (s.wptr + 1) (s.sptr + 1) n =
        s.acc + (s.weights s.wptr * s.stream s.sptr +
          fsumFrom s.weights s.stream (s.wptr + 1) (s.sptr + 1) n)n:Nats:KernelState⊢ s.acc + s.weights s.wptr * s.stream s.sptr + fsumFrom s.weights s.stream (s.wptr + 1) (s.sptr + 1) n =
  s.acc + (s.weights s.wptr * s.stream s.sptr + fsumFrom s.weights s.stream (s.wptr + 1) (s.sptr + 1) n)
      omegaAll goals completed! 🐙

def dotSpec (w a : Nat -> Int) (n : Nat) : Int :=
  fsumFrom w a 0 0 n

theorem kernelDot_correct (w a : Nat -> Int) (n : Nat) :
    (krun n (KernelState.mk 0 w a 0 0)).acc = dotSpec w a n := byw:Nat → Inta:Nat → Intn:Nat⊢ (krun n { acc := 0, weights := w, stream := a, wptr := 0, sptr := 0 }).acc = dotSpec w a n
  rw [krun_accw:Nat → Inta:Nat → Intn:Nat⊢ { acc := 0, weights := w, stream := a, wptr := 0, sptr := 0 }.acc +
    fsumFrom { acc := 0, weights := w, stream := a, wptr := 0, sptr := 0 }.weights
      { acc := 0, weights := w, stream := a, wptr := 0, sptr := 0 }.stream
      { acc := 0, weights := w, stream := a, wptr := 0, sptr := 0 }.wptr
      { acc := 0, weights := w, stream := a, wptr := 0, sptr := 0 }.sptr n =
  dotSpec w a n]w:Nat → Inta:Nat → Intn:Nat⊢ { acc := 0, weights := w, stream := a, wptr := 0, sptr := 0 }.acc +
    fsumFrom { acc := 0, weights := w, stream := a, wptr := 0, sptr := 0 }.weights
      { acc := 0, weights := w, stream := a, wptr := 0, sptr := 0 }.stream
      { acc := 0, weights := w, stream := a, wptr := 0, sptr := 0 }.wptr
      { acc := 0, weights := w, stream := a, wptr := 0, sptr := 0 }.sptr n =
  dotSpec w a n
  simp [dotSpec]All goals completed! 🐙

The integer proof above is the source-level claim. The fixed-width hardware claim needs an overflow obligation: every prefix sum that may enter the accumulator must fit the configured signed accumulator width. This is stronger than a final-result-only check, because wrapping at an intermediate cycle would already have changed the later computation.

def kernelPrefixFits (cfg : Config) (s : KernelState) (n : Nat) : Prop :=
  forall k, k <= n ->
    signedFits cfg.accBits
      (s.acc + fsumFrom s.weights s.stream s.wptr s.sptr k)

theorem kernel_final_fits_of_prefixFits (cfg : Config)
    (s : KernelState) (n : Nat)
    (hfit : kernelPrefixFits cfg s n) :
    signedFits cfg.accBits (krun n s).acc := bycfg:Configs:KernelStaten:Nathfit:kernelPrefixFits cfg s n⊢ signedFits cfg.accBits (krun n s).acc
  rw [krun_acccfg:Configs:KernelStaten:Nathfit:kernelPrefixFits cfg s n⊢ signedFits cfg.accBits (s.acc + fsumFrom s.weights s.stream s.wptr s.sptr n)]cfg:Configs:KernelStaten:Nathfit:kernelPrefixFits cfg s n⊢ signedFits cfg.accBits (s.acc + fsumFrom s.weights s.stream s.wptr s.sptr n)
  exact hfit n (Nat.le_refl n)All goals completed! 🐙

theorem kernel_accumulator_exact_when_prefixFits (cfg : Config)
    (hbits : 0 < cfg.accBits) (s : KernelState) (n : Nat)
    (hfit : kernelPrefixFits cfg s n) :
    (accOfInt cfg (krun n s).acc).toInt = (krun n s).acc := bycfg:Confighbits:0 < cfg.accBitss:KernelStaten:Nathfit:kernelPrefixFits cfg s n⊢ BitVec.toInt (accOfInt cfg (krun n s).acc) = (krun n s).acc
  exact accOfInt_exact cfg hbits
    (kernel_final_fits_of_prefixFits cfg s n hfit)All goals completed! 🐙

A small executable smoke test stays in the book. If the semantics change, this example changes with it.

def wDemo : Nat -> Int :=
  fun i => match i with | 0 => 1 | 1 => 2 | 2 => 3 | _ => 0
def aDemo : Nat -> Int :=
  fun i => match i with | 0 => 4 | 1 => 5 | 2 => 6 | _ => 0

theorem kernelDot_demo : dotSpec wDemo aDemo 3 = 32 := by⊢ dotSpec wDemo aDemo 3 = 32
  native_decideAll goals completed! 🐙

end Janus