module Realization.Nvidia.SM86.StallCompaction

import Accelerator.SM86.Control
import Accelerator.SM86.Instruction
import Accelerator.SM86.Operands
import Accelerator.SM86.Types
import Std.List
import Std.Natural

-- A PROPOSER, not part of any argument: it rewrites a program's stall
-- counts to the smallest each instruction's successor allows, by the same
-- accounting SM86.Scoreboard's fixed-latency pass checks (issue one
-- instruction, wait its stall; a fixed-latency result is ready
-- sm86FixedLatency cycles after issue, a barrier sm86BarrierSetCycles after
-- the instruction that sets it).  Whatever it proposes is accepted only if
-- the checker accepts it (SM86.LinearStepCheck.linearStepProgramAcceptedWith
-- for the linear step); a mistake here is a rejected proposal, never a wrong
-- program.
--
-- Generated programs stall 15 cycles after every instruction (the maximum,
-- chosen when nothing modelled latency); most instructions need 1, none
-- less than its sm86MinimumStall.

-- The annotated program: each instruction's schedule keys computed once.
-- The old walk eliminated every body five times per step -- writes, latency
-- and set keys for the issue at this step, reads, writes, guard, predicate
-- and wait keys for the need of the previous one -- and rebuilt each
-- instruction with a sixth split.  Annotation does one 37-branch split per
-- instruction up front (each field the same branch body the old helpers
-- used), and carries the keys forward with the instruction itself, so the
-- walk never re-examines a body except to rebuild it with its new stall.
family CompactionAnnotated : Type 0
constructor CompactionAnnotatedEnd
constructor CompactionAnnotatedNext
field unrestricted compactionAnnotatedInstruction : (family SM86Instruction)
field unrestricted compactionAnnotatedNeeds : (family StdList Nat)
field unrestricted compactionAnnotatedIssue : (family StdList Nat)
field unrestricted compactionAnnotatedSetKeys : (family StdList Nat)
field unrestricted compactionAnnotatedLatency : Nat
field unrestricted compactionAnnotatedMinimumStall : Nat
recursive unrestricted compactionAnnotatedTail
end-family

-- keys ready at cycles: (key, ready)
family CompactionPending : Type 0
constructor CompactionPendingEnd
constructor CompactionPendingNext
field unrestricted compactionPendingKey : Nat
field unrestricted compactionPendingReady : Nat
recursive unrestricted compactionPendingTail
end-family

-- Ampere can read a variable-latency instruction's source registers after
-- issue. The RTX 3090 attention-key LDG overwrote its address before that
-- late read and faulted (c28712e8, driver 580.126.20). On that card, SB5
-- as the read barrier and an SB5 wait on every instruction completed a full
-- one-step update and checkpoint. This serial form is a conservative repair
-- for a schedule the scoreboard refuses; the caller preserves already
-- accepted schedules. The stall compactor below restores the shortest
-- schedule admitted by the fixed-latency model. A more precise register
-- liveness scheduler needs its own card qualification before replacing it.
def sm86ReadBarrierWait =
  (lambda unrestricted barrier : (family SM86Barrier) .
    (eliminate SM86Barrier (lambda unrestricted current : (family SM86Barrier) . Byte) barrier
      (branch SM86Barrier0 . (byte 1)) (branch SM86Barrier1 . (byte 2))
      (branch SM86Barrier2 . (byte 4)) (branch SM86Barrier3 . (byte 8))
      (branch SM86Barrier4 . (byte 16)) (branch SM86Barrier5 . (byte 32))
      (branch SM86Barrier6 . (byte 0)) (branch SM86BarrierNone . (byte 0))))

def sm86SerializeLateReadInstruction =
  (lambda unrestricted instruction : (family SM86Instruction) .
    (eliminate SM86Instruction
      (lambda unrestricted current : (family SM86Instruction) . (family SM86Instruction)) instruction
      (branch SM86InstructionValue guard body .
        (constructor SM86Instruction SM86InstructionValue guard
          (sm86BodyWithControl body
            (eliminate SM86Control
              (lambda unrestricted current : (family SM86Control) . (family SM86Control))
              (sm86BodyControlOf body)
              (branch SM86ControlValue stall yield write read wait reuse .
                (constructor SM86Control SM86ControlValue stall yield write
                  (eliminate SM86Barrier
                    (lambda unrestricted current : (family SM86Barrier) . (family SM86Barrier)) read
                    (branch SM86Barrier0 . (constructor SM86Barrier SM86Barrier0))
                    (branch SM86Barrier1 . (constructor SM86Barrier SM86Barrier1))
                    (branch SM86Barrier2 . (constructor SM86Barrier SM86Barrier2))
                    (branch SM86Barrier3 . (constructor SM86Barrier SM86Barrier3))
                    (branch SM86Barrier4 . (constructor SM86Barrier SM86Barrier4))
                    (branch SM86Barrier5 . (constructor SM86Barrier SM86Barrier5))
                    (branch SM86Barrier6 . (constructor SM86Barrier SM86Barrier6))
                    (branch SM86BarrierNone .
                      (nat-eliminate
                        (lambda unrestricted n : Nat . (family SM86Barrier))
                        (constructor SM86Barrier SM86Barrier5)
                        (lambda unrestricted p : Nat .
                          (lambda unrestricted ignored : (family SM86Barrier) .
                            (constructor SM86Barrier SM86BarrierNone)))
                        (sm86FixedLatency body))))
                  (byte-or wait (byte-or (byte 32) (sm86ReadBarrierWait read))) reuse))))))))

def sm86SerializeLateReads =
  (lambda unrestricted program : (family SM86Program) .
    (eliminate SM86Program
      (lambda unrestricted current : (family SM86Program) . (family SM86Program)) program
      (branch SM86ProgramEnd . (constructor SM86Program SM86ProgramEnd))
      (branch SM86ProgramNext instruction tail induction .
        (constructor SM86Program SM86ProgramNext
          (sm86SerializeLateReadInstruction instruction) induction))))

def scMax = (lambda unrestricted a : Nat . (lambda unrestricted b : Nat .
  (nat-eliminate (lambda unrestricted current : Nat . Nat)
    a
    (lambda unrestricted flagPredecessor : Nat . (lambda unrestricted flagRest : Nat . b))
    (nat-less-than a b))))

-- the latest ready cycle among the keys (0 when none is pending); the
-- comparisons are the primitive's own, eliminated where they are made
def scReadyOf =
  (lambda unrestricted pending : (family CompactionPending) .
    (lambda unrestricted keys : (family StdList Nat) .
      (eliminate StdList (lambda unrestricted current : (family StdList Nat) . Nat) keys
        (branch StdListEmpty . zero)
        (branch StdListCons key rest induction .
          (scMax induction
            (eliminate CompactionPending (lambda unrestricted current : (family CompactionPending) . Nat) pending
              (branch CompactionPendingEnd . zero)
              (branch CompactionPendingNext bound ready tail inner .
                (nat-eliminate (lambda unrestricted current : Nat . Nat)
                  inner
                  (lambda unrestricted flagPredecessor : Nat . (lambda unrestricted flagRest : Nat .
                    (nat-eliminate (lambda unrestricted current : Nat . Nat)
                      ready
                      (lambda unrestricted flagPredecessor : Nat . (lambda unrestricted flagRest : Nat . inner))
                      (nat-less-than ready inner))))
                  (nat-eliminate (lambda unrestricted current : Nat . Nat)
                    (nat-eliminate (lambda unrestricted current : Nat . Nat)
                      (succ zero)
                      (lambda unrestricted flagPredecessor : Nat . (lambda unrestricted flagRest : Nat . zero))
                      (nat-less-than key bound))
                    (lambda unrestricted flagPredecessor : Nat . (lambda unrestricted flagRest : Nat . zero))
                    (nat-less-than bound key))))))))))

def scIssue =
  (lambda unrestricted keys : (family StdList Nat) .
    (lambda unrestricted ready : Nat .
      (lambda unrestricted pending : (family CompactionPending) .
        (eliminate StdList (lambda unrestricted current : (family StdList Nat) . (family CompactionPending)) keys
          (branch StdListEmpty . pending)
          (branch StdListCons key rest induction .
            (constructor CompactionPending CompactionPendingNext key ready induction))))))

-- the keys an instruction must find ready when it issues, precomputed once
-- per instruction instead of once per step as the next instruction
def scNeedsMake =
  (lambda unrestricted reads : (family StdList Nat) .
    (lambda unrestricted writes : (family StdList Nat) .
      (lambda unrestricted guardKeys : (family StdList Nat) .
        (lambda unrestricted predicateWrites : (family StdList Nat) .
          (lambda unrestricted waitKeys : (family StdList Nat) .
            (stdListAppend Nat reads
              (stdListAppend Nat writes
                (stdListAppend Nat guardKeys
                  (stdListAppend Nat predicateWrites waitKeys)))))))))

def scIssueMake =
  (lambda unrestricted writes : (family StdList Nat) .
    (lambda unrestricted predicateWrites : (family StdList Nat) .
      (stdListAppend Nat writes predicateWrites)))

-- one 37-branch split per instruction, up front: the summary answers reads,
-- writes, predicate writes, wait and set keys and latency together, and the
-- node carries the instruction itself for the rebuild with its new stall
def sm86AnnotateProgram =
  (lambda unrestricted program : (family SM86Program) .
    (eliminate SM86Program
      (lambda unrestricted current : (family SM86Program) . (family CompactionAnnotated))
      program
      (branch SM86ProgramEnd . (constructor CompactionAnnotated CompactionAnnotatedEnd))
      (branch SM86ProgramNext instruction tail induction .
        (eliminate SM86Instruction
          (lambda unrestricted current : (family SM86Instruction) . (family CompactionAnnotated))
          instruction
          (branch SM86InstructionValue guard body .
            (eliminate SM86OpSummary
              (lambda unrestricted current : (family SM86OpSummary) . (family CompactionAnnotated))
              (sm86OpSummaryOfBody body)
              (branch SM86OpSummaryValue reads writes predicateWrites waitKeys setKeys stall latency variable minimum control .
                (constructor CompactionAnnotated CompactionAnnotatedNext
                  instruction
                  (scNeedsMake reads writes (sm86GuardKeys instruction) predicateWrites waitKeys)
                  (scIssueMake writes predicateWrites)
                  setKeys
                  latency
                  minimum
                  induction))))))))

-- the entries still in flight at cycle `now` (an entry ready by then can
-- raise no later stall: the program is linear in its length, not quadratic)
def scPrune =
  (lambda unrestricted pending : (family CompactionPending) .
    (lambda unrestricted now : Nat .
      (eliminate CompactionPending (lambda unrestricted current : (family CompactionPending) . (family CompactionPending)) pending
        (branch CompactionPendingEnd . (constructor CompactionPending CompactionPendingEnd))
        (branch CompactionPendingNext key ready tail induction .
          (nat-eliminate (lambda unrestricted current : Nat . (family CompactionPending))
            induction
            (lambda unrestricted p : Nat . (lambda unrestricted ignored : (family CompactionPending) .
              (constructor CompactionPending CompactionPendingNext key ready induction)))
            (naturalLess now ready))))))

-- the instruction with another stall
def scWithStall =
  (lambda unrestricted instruction : (family SM86Instruction) .
    (lambda unrestricted stall : Nat .
      (eliminate SM86Instruction (lambda unrestricted current : (family SM86Instruction) . (family SM86Instruction)) instruction
        (branch SM86InstructionValue guard body .
          (constructor SM86Instruction SM86InstructionValue guard
            (sm86BodyWithControl body (sm86ControlWithStall (sm86BodyControlOf body) stall)))))))

-- the latest ready cycle of anything in flight
def scLatest =
  (lambda unrestricted pending : (family CompactionPending) .
    (eliminate CompactionPending (lambda unrestricted current : (family CompactionPending) . Nat) pending
      (branch CompactionPendingEnd . zero)
      (branch CompactionPendingNext key ready tail induction . (scMax ready induction))))

-- `drained` nonzero: the last instruction also waits until every
-- fixed-latency result in flight is ready, so whatever the program is placed
-- before (a piece of a larger program, a loop's back edge) reads nothing
-- early
def sm86CompactStallsWith =
  (lambda unrestricted drained : Nat .
  (lambda unrestricted program : (family SM86Program) .
    (app (app
      (eliminate CompactionAnnotated
        (lambda unrestricted current : (family CompactionAnnotated) .
          (pi unrestricted now : Nat . (pi unrestricted pending : (family CompactionPending) . (family SM86Program))))
        (sm86AnnotateProgram program)
        (branch CompactionAnnotatedEnd .
          (lambda unrestricted now : Nat . (lambda unrestricted pending : (family CompactionPending) .
            (constructor SM86Program SM86ProgramEnd))))
        (branch CompactionAnnotatedNext instruction needs issue setKeys latency minimum tail induction .
          (lambda unrestricted now : Nat . (lambda unrestricted pending : (family CompactionPending) .
            (let unrestricted issued =
                  (scIssue issue (naturalAdd now latency)
                    (scIssue setKeys (naturalAdd now sm86BarrierSetCycles) pending))
              in (let unrestricted need =
                    (eliminate CompactionAnnotated
                      (lambda unrestricted current : (family CompactionAnnotated) . Nat)
                      tail
                      (branch CompactionAnnotatedEnd . (naturalSelect drained (scLatest issued) zero))
                      (branch CompactionAnnotatedNext followingInstruction followingNeeds followingIssue followingSetKeys followingLatency followingMinimum followingTail followingInduction .
                        (scReadyOf issued followingNeeds)))
                in (let unrestricted stall = (scMax minimum
                      (naturalSelect (naturalLess (naturalAdd now 1) need) (naturalSaturatingSubtract need now) 1))
                  in (let unrestricted clamped = (naturalSelect (naturalLess 15 stall) 15 stall)
                    in (constructor SM86Program SM86ProgramNext (scWithStall instruction clamped)
                         (induction (naturalAdd now clamped) (scPrune issued (naturalAdd now clamped))))))))))))
      0) (constructor CompactionPending CompactionPendingEnd))))

def sm86CompactStalls = (sm86CompactStallsWith 0)
def sm86CompactStallsDrained = (sm86CompactStallsWith 1)

-- Recompute fixed-latency stalls after adding read-barrier waits; otherwise
-- the two-instruction checker can reject a correct serialization at its new
-- barrier-set dependency.
def sm86GuardLateReads =
  (lambda unrestricted program : (family SM86Program) .
    (sm86CompactStalls (sm86SerializeLateReads program)))