Realization.Nvidia.SM86.StallCompaction

1module Realization.Nvidia.SM86.StallCompaction
2
3import Accelerator.SM86.Control
4import Accelerator.SM86.Instruction
5import Accelerator.SM86.Operands
6import Accelerator.SM86.Types
7import Std.List
8import Std.Natural
9
10-- A PROPOSER, not part of any argument: it rewrites a program's stall
11-- counts to the smallest each instruction's successor allows, by the same
12-- accounting SM86.Scoreboard's fixed-latency pass checks (issue one
13-- instruction, wait its stall; a fixed-latency result is ready
14-- sm86FixedLatency cycles after issue, a barrier sm86BarrierSetCycles after
15-- the instruction that sets it).  Whatever it proposes is accepted only if
16-- the checker accepts it (SM86.LinearStepCheck.linearStepProgramAcceptedWith
17-- for the linear step); a mistake here is a rejected proposal, never a wrong
18-- program.
19--
20-- Generated programs stall 15 cycles after every instruction (the maximum,
21-- chosen when nothing modelled latency); most instructions need 1, none
22-- less than its sm86MinimumStall.
23
24-- The annotated program: each instruction's schedule keys computed once.
25-- The old walk eliminated every body five times per step -- writes, latency
26-- and set keys for the issue at this step, reads, writes, guard, predicate
27-- and wait keys for the need of the previous one -- and rebuilt each
28-- instruction with a sixth split.  Annotation does one 37-branch split per
29-- instruction up front (each field the same branch body the old helpers
30-- used), and carries the keys forward with the instruction itself, so the
31-- walk never re-examines a body except to rebuild it with its new stall.
32family CompactionAnnotated : Type 0
33constructor CompactionAnnotatedEnd
34constructor CompactionAnnotatedNext
35field unrestricted compactionAnnotatedInstruction : (family SM86Instruction)
36field unrestricted compactionAnnotatedNeeds : (family StdList Nat)
37field unrestricted compactionAnnotatedIssue : (family StdList Nat)
38field unrestricted compactionAnnotatedSetKeys : (family StdList Nat)
39field unrestricted compactionAnnotatedLatency : Nat
40field unrestricted compactionAnnotatedMinimumStall : Nat
41recursive unrestricted compactionAnnotatedTail
42end-family
43
44-- keys ready at cycles: (key, ready)
45family CompactionPending : Type 0
46constructor CompactionPendingEnd
47constructor CompactionPendingNext
48field unrestricted compactionPendingKey : Nat
49field unrestricted compactionPendingReady : Nat
50recursive unrestricted compactionPendingTail
51end-family
52
53-- Ampere can read a variable-latency instruction's source registers after
54-- issue. The RTX 3090 attention-key LDG overwrote its address before that
55-- late read and faulted (c28712e8, driver 580.126.20). On that card, SB5
56-- as the read barrier and an SB5 wait on every instruction completed a full
57-- one-step update and checkpoint. This serial form is a conservative repair
58-- for a schedule the scoreboard refuses; the caller preserves already
59-- accepted schedules. The stall compactor below restores the shortest
60-- schedule admitted by the fixed-latency model. A more precise register
61-- liveness scheduler needs its own card qualification before replacing it.
62def sm86ReadBarrierWait =
63  (lambda unrestricted barrier : (family SM86Barrier) .
64    (eliminate SM86Barrier (lambda unrestricted current : (family SM86Barrier) . Byte) barrier
65      (branch SM86Barrier0 . (byte 1)) (branch SM86Barrier1 . (byte 2))
66      (branch SM86Barrier2 . (byte 4)) (branch SM86Barrier3 . (byte 8))
67      (branch SM86Barrier4 . (byte 16)) (branch SM86Barrier5 . (byte 32))
68      (branch SM86Barrier6 . (byte 0)) (branch SM86BarrierNone . (byte 0))))
69
70def sm86SerializeLateReadInstruction =
71  (lambda unrestricted instruction : (family SM86Instruction) .
72    (eliminate SM86Instruction
73      (lambda unrestricted current : (family SM86Instruction) . (family SM86Instruction)) instruction
74      (branch SM86InstructionValue guard body .
75        (constructor SM86Instruction SM86InstructionValue guard
76          (sm86BodyWithControl body
77            (eliminate SM86Control
78              (lambda unrestricted current : (family SM86Control) . (family SM86Control))
79              (sm86BodyControlOf body)
80              (branch SM86ControlValue stall yield write read wait reuse .
81                (constructor SM86Control SM86ControlValue stall yield write
82                  (eliminate SM86Barrier
83                    (lambda unrestricted current : (family SM86Barrier) . (family SM86Barrier)) read
84                    (branch SM86Barrier0 . (constructor SM86Barrier SM86Barrier0))
85                    (branch SM86Barrier1 . (constructor SM86Barrier SM86Barrier1))
86                    (branch SM86Barrier2 . (constructor SM86Barrier SM86Barrier2))
87                    (branch SM86Barrier3 . (constructor SM86Barrier SM86Barrier3))
88                    (branch SM86Barrier4 . (constructor SM86Barrier SM86Barrier4))
89                    (branch SM86Barrier5 . (constructor SM86Barrier SM86Barrier5))
90                    (branch SM86Barrier6 . (constructor SM86Barrier SM86Barrier6))
91                    (branch SM86BarrierNone .
92                      (nat-eliminate
93                        (lambda unrestricted n : Nat . (family SM86Barrier))
94                        (constructor SM86Barrier SM86Barrier5)
95                        (lambda unrestricted p : Nat .
96                          (lambda unrestricted ignored : (family SM86Barrier) .
97                            (constructor SM86Barrier SM86BarrierNone)))
98                        (sm86FixedLatency body))))
99                  (byte-or wait (byte-or (byte 32) (sm86ReadBarrierWait read))) reuse))))))))
100
101def sm86SerializeLateReads =
102  (lambda unrestricted program : (family SM86Program) .
103    (eliminate SM86Program
104      (lambda unrestricted current : (family SM86Program) . (family SM86Program)) program
105      (branch SM86ProgramEnd . (constructor SM86Program SM86ProgramEnd))
106      (branch SM86ProgramNext instruction tail induction .
107        (constructor SM86Program SM86ProgramNext
108          (sm86SerializeLateReadInstruction instruction) induction))))
109
110def scMax = (lambda unrestricted a : Nat . (lambda unrestricted b : Nat .
111  (nat-eliminate (lambda unrestricted current : Nat . Nat)
112    a
113    (lambda unrestricted flagPredecessor : Nat . (lambda unrestricted flagRest : Nat . b))
114    (nat-less-than a b))))
115
116-- the latest ready cycle among the keys (0 when none is pending); the
117-- comparisons are the primitive's own, eliminated where they are made
118def scReadyOf =
119  (lambda unrestricted pending : (family CompactionPending) .
120    (lambda unrestricted keys : (family StdList Nat) .
121      (eliminate StdList (lambda unrestricted current : (family StdList Nat) . Nat) keys
122        (branch StdListEmpty . zero)
123        (branch StdListCons key rest induction .
124          (scMax induction
125            (eliminate CompactionPending (lambda unrestricted current : (family CompactionPending) . Nat) pending
126              (branch CompactionPendingEnd . zero)
127              (branch CompactionPendingNext bound ready tail inner .
128                (nat-eliminate (lambda unrestricted current : Nat . Nat)
129                  inner
130                  (lambda unrestricted flagPredecessor : Nat . (lambda unrestricted flagRest : Nat .
131                    (nat-eliminate (lambda unrestricted current : Nat . Nat)
132                      ready
133                      (lambda unrestricted flagPredecessor : Nat . (lambda unrestricted flagRest : Nat . inner))
134                      (nat-less-than ready inner))))
135                  (nat-eliminate (lambda unrestricted current : Nat . Nat)
136                    (nat-eliminate (lambda unrestricted current : Nat . Nat)
137                      (succ zero)
138                      (lambda unrestricted flagPredecessor : Nat . (lambda unrestricted flagRest : Nat . zero))
139                      (nat-less-than key bound))
140                    (lambda unrestricted flagPredecessor : Nat . (lambda unrestricted flagRest : Nat . zero))
141                    (nat-less-than bound key))))))))))
142
143def scIssue =
144  (lambda unrestricted keys : (family StdList Nat) .
145    (lambda unrestricted ready : Nat .
146      (lambda unrestricted pending : (family CompactionPending) .
147        (eliminate StdList (lambda unrestricted current : (family StdList Nat) . (family CompactionPending)) keys
148          (branch StdListEmpty . pending)
149          (branch StdListCons key rest induction .
150            (constructor CompactionPending CompactionPendingNext key ready induction))))))
151
152-- the keys an instruction must find ready when it issues, precomputed once
153-- per instruction instead of once per step as the next instruction
154def scNeedsMake =
155  (lambda unrestricted reads : (family StdList Nat) .
156    (lambda unrestricted writes : (family StdList Nat) .
157      (lambda unrestricted guardKeys : (family StdList Nat) .
158        (lambda unrestricted predicateWrites : (family StdList Nat) .
159          (lambda unrestricted waitKeys : (family StdList Nat) .
160            (stdListAppend Nat reads
161              (stdListAppend Nat writes
162                (stdListAppend Nat guardKeys
163                  (stdListAppend Nat predicateWrites waitKeys)))))))))
164
165def scIssueMake =
166  (lambda unrestricted writes : (family StdList Nat) .
167    (lambda unrestricted predicateWrites : (family StdList Nat) .
168      (stdListAppend Nat writes predicateWrites)))
169
170-- one 37-branch split per instruction, up front: the summary answers reads,
171-- writes, predicate writes, wait and set keys and latency together, and the
172-- node carries the instruction itself for the rebuild with its new stall
173def sm86AnnotateProgram =
174  (lambda unrestricted program : (family SM86Program) .
175    (eliminate SM86Program
176      (lambda unrestricted current : (family SM86Program) . (family CompactionAnnotated))
177      program
178      (branch SM86ProgramEnd . (constructor CompactionAnnotated CompactionAnnotatedEnd))
179      (branch SM86ProgramNext instruction tail induction .
180        (eliminate SM86Instruction
181          (lambda unrestricted current : (family SM86Instruction) . (family CompactionAnnotated))
182          instruction
183          (branch SM86InstructionValue guard body .
184            (eliminate SM86OpSummary
185              (lambda unrestricted current : (family SM86OpSummary) . (family CompactionAnnotated))
186              (sm86OpSummaryOfBody body)
187              (branch SM86OpSummaryValue reads writes predicateWrites waitKeys setKeys stall latency variable minimum control .
188                (constructor CompactionAnnotated CompactionAnnotatedNext
189                  instruction
190                  (scNeedsMake reads writes (sm86GuardKeys instruction) predicateWrites waitKeys)
191                  (scIssueMake writes predicateWrites)
192                  setKeys
193                  latency
194                  minimum
195                  induction))))))))
196
197-- the entries still in flight at cycle `now` (an entry ready by then can
198-- raise no later stall: the program is linear in its length, not quadratic)
199def scPrune =
200  (lambda unrestricted pending : (family CompactionPending) .
201    (lambda unrestricted now : Nat .
202      (eliminate CompactionPending (lambda unrestricted current : (family CompactionPending) . (family CompactionPending)) pending
203        (branch CompactionPendingEnd . (constructor CompactionPending CompactionPendingEnd))
204        (branch CompactionPendingNext key ready tail induction .
205          (nat-eliminate (lambda unrestricted current : Nat . (family CompactionPending))
206            induction
207            (lambda unrestricted p : Nat . (lambda unrestricted ignored : (family CompactionPending) .
208              (constructor CompactionPending CompactionPendingNext key ready induction)))
209            (naturalLess now ready))))))
210
211-- the instruction with another stall
212def scWithStall =
213  (lambda unrestricted instruction : (family SM86Instruction) .
214    (lambda unrestricted stall : Nat .
215      (eliminate SM86Instruction (lambda unrestricted current : (family SM86Instruction) . (family SM86Instruction)) instruction
216        (branch SM86InstructionValue guard body .
217          (constructor SM86Instruction SM86InstructionValue guard
218            (sm86BodyWithControl body (sm86ControlWithStall (sm86BodyControlOf body) stall)))))))
219
220-- the latest ready cycle of anything in flight
221def scLatest =
222  (lambda unrestricted pending : (family CompactionPending) .
223    (eliminate CompactionPending (lambda unrestricted current : (family CompactionPending) . Nat) pending
224      (branch CompactionPendingEnd . zero)
225      (branch CompactionPendingNext key ready tail induction . (scMax ready induction))))
226
227-- `drained` nonzero: the last instruction also waits until every
228-- fixed-latency result in flight is ready, so whatever the program is placed
229-- before (a piece of a larger program, a loop's back edge) reads nothing
230-- early
231def sm86CompactStallsWith =
232  (lambda unrestricted drained : Nat .
233  (lambda unrestricted program : (family SM86Program) .
234    (app (app
235      (eliminate CompactionAnnotated
236        (lambda unrestricted current : (family CompactionAnnotated) .
237          (pi unrestricted now : Nat . (pi unrestricted pending : (family CompactionPending) . (family SM86Program))))
238        (sm86AnnotateProgram program)
239        (branch CompactionAnnotatedEnd .
240          (lambda unrestricted now : Nat . (lambda unrestricted pending : (family CompactionPending) .
241            (constructor SM86Program SM86ProgramEnd))))
242        (branch CompactionAnnotatedNext instruction needs issue setKeys latency minimum tail induction .
243          (lambda unrestricted now : Nat . (lambda unrestricted pending : (family CompactionPending) .
244            (let unrestricted issued =
245                  (scIssue issue (naturalAdd now latency)
246                    (scIssue setKeys (naturalAdd now sm86BarrierSetCycles) pending))
247              in (let unrestricted need =
248                    (eliminate CompactionAnnotated
249                      (lambda unrestricted current : (family CompactionAnnotated) . Nat)
250                      tail
251                      (branch CompactionAnnotatedEnd . (naturalSelect drained (scLatest issued) zero))
252                      (branch CompactionAnnotatedNext followingInstruction followingNeeds followingIssue followingSetKeys followingLatency followingMinimum followingTail followingInduction .
253                        (scReadyOf issued followingNeeds)))
254                in (let unrestricted stall = (scMax minimum
255                      (naturalSelect (naturalLess (naturalAdd now 1) need) (naturalSaturatingSubtract need now) 1))
256                  in (let unrestricted clamped = (naturalSelect (naturalLess 15 stall) 15 stall)
257                    in (constructor SM86Program SM86ProgramNext (scWithStall instruction clamped)
258                         (induction (naturalAdd now clamped) (scPrune issued (naturalAdd now clamped))))))))))))
259      0) (constructor CompactionPending CompactionPendingEnd))))
260
261def sm86CompactStalls = (sm86CompactStallsWith 0)
262def sm86CompactStallsDrained = (sm86CompactStallsWith 1)
263
264-- Recompute fixed-latency stalls after adding read-barrier waits; otherwise
265-- the two-instruction checker can reject a correct serialization at its new
266-- barrier-set dependency.
267def sm86GuardLateReads =
268  (lambda unrestricted program : (family SM86Program) .
269    (sm86CompactStalls (sm86SerializeLateReads program)))