Realization.Nvidia.SM86.AdamWHalfSM86

1module Realization.Nvidia.SM86.AdamWHalfSM86
2
3import Accelerator.SM86.Control
4import Accelerator.SM86.Immediate
5import Accelerator.SM86.Instruction
6import Accelerator.SM86.NumericSemantics
7import Accelerator.SM86.Program
8import Accelerator.SM86.Types
9import Realization.Nvidia.SM86.StreamingAttentionSM86
10import Std.Natural
11
12-- AdamW over the whole flat bank with the parameters' half copies written
13-- as it goes: the half cast that followed the
14-- update, fused into it.  Each element's update is
15-- Realization.Nvidia.SM86.AdamWSM86's (adamWF32MathProgram), instruction
16-- for instruction, and each half pair is the cast's (F2FP, element 2i + 1
17-- high, in the realization's half format), so the parameters, the moments
18-- and the halves are bit-identical to the update followed by the cast:
19--
20--   g = g one;  m = m beta1 + g (1 - beta1);  v = v beta2 + g^2 (1 - beta2)
21--   p = p decay + -(m stepSize (1 / (sqrt v + epsilon)))
22--   half[i] = half(p[i])
23--
24-- The half bank is parallel to the parameter bank: element i's half at byte
25-- 2 i of it.  A thread takes the quad 4T .. 4T + 3 (T = block 256 + thread):
26-- each bank's four elements in one 128-bit load and store, the four halves
27-- in one 64-bit store -- the bank is memory-bound, and wide accesses keep
28-- more of it in flight a thread -- then T + stride, ... `iterations` times;
29-- quads at or past `quads` are untouched.  The gradient is read, never
30-- written (every gradient the step reads was written whole by the backward
31-- before it).
32--
33-- Parameters (constant bank 0 from the SM86 parameter base): the pointers
34-- P, G, M, V (arguments 0-3, 16-byte aligned); the scalars one, beta1,
35-- 1 - beta1, beta2, 1 - beta2, decay, stepSize, epsilon (words from argument
36-- 4, where AdamWSM86's ABI puts them, so the host writes the step's scalars
37-- at the same offsets); the half bank's pointer (argument 9, 8-byte
38-- aligned).
39--
40-- Registers: R0 thread, R1 block, R2 T, R3 sixteen, R4 eight; R6:R7, R8:R9,
41-- R10:R11, R12:R13 the quad's P, G, M, V; R18..R25 the scalars; R29 the
42-- iterations left; R30:R31 the halves' address; R32..R35, R36..R39,
43-- R40..R43, R44..R47 the quad's p, g, m, v; element e's temporaries
44-- R48 + 3e .. R50 + 3e; R60, R61 the half pairs.  Loads under SB0..SB3 (P,
45-- G, M, V), the square roots and reciprocals under SB4; the stores' reads
46-- under SB5, which the next iteration's first load waits for.
47
48def ahThreads : Nat = 256
49def adamWHalfSM86ParameterArgument : Nat = 0
50def adamWHalfSM86GradientArgument : Nat = 1
51def adamWHalfSM86FirstMomentArgument : Nat = 2
52def adamWHalfSM86SecondMomentArgument : Nat = 3
53def adamWHalfSM86ScalarPairArgument =
54  (lambda unrestricted pair : Nat . (naturalAdd 4 pair))
55def adamWHalfSM86HalfArgument : Nat = 9
56def adamWHalfSM86ArgumentCount : Nat = 10
57def ahScalar = (lambda unrestricted j : Nat .
58  (naturalAdd (saArgument (adamWHalfSM86ScalarPairArgument 0))
59    (naturalMultiply 4 j)))
60def ahHalfPointer : Nat = (saArgument adamWHalfSM86HalfArgument)
61
62-- element e's p, g, m, v (bank b = 0..3) and temporaries (i = 0..2)
63def ahValue = (lambda unrestricted e : Nat . (lambda unrestricted b : Nat . (naturalAdd (naturalAdd 32 (naturalMultiply 4 b)) e)))
64def ahTemporary = (lambda unrestricted e : Nat . (lambda unrestricted i : Nat . (naturalAdd (naturalAdd 48 (naturalMultiply 3 e)) i)))
65def ahHalfPairs : Nat = 60
66def ahCounter : Nat = 29
67
68def ahIterations = (lambda unrestricted quads : Nat . (lambda unrestricted stride : Nat .
69  (naturalDivideUnchecked (naturalAdd quads (naturalSaturatingSubtract stride 1)) stride)))
70
71-- ---- the prologue: T, the constants, the scalars, the count ----
72-- `scaled` 1: the gradient factor (the first scalar, `one`) is instead the
73-- word at argument 10's pointer -- a device-computed scale, such as the
74-- clip factor -- loaded once (data under SB4, retired with the specials).
75-- g = g s is then exactly the separate scaling pass's product, and that
76-- pass over the bank is not needed.
77def adamWHalfSM86ScaleArgument : Nat = 10
78def adamWHalfSM86ScaledConstantBytes : Nat = (saArgument 11)
79def ahChoose = (lambda unrestricted condition : Nat .
80  (lambda unrestricted yes : (family SM86Program) . (lambda unrestricted no : (family SM86Program) .
81    (nat-eliminate (lambda unrestricted n : Nat . (family SM86Program)) no
82      (lambda unrestricted p : Nat . (lambda unrestricted ignored : (family SM86Program) . yes)) condition))))
83-- the scalars from the constant bank into R18 + j, from j = `first` on
84def ahScalarsFrom = (lambda unrestricted first : Nat . (lambda unrestricted tail : (family SM86Program) .
85  (saFor (naturalSaturatingSubtract 8 first)
86    (lambda unrestricted j : Nat .
87      (saMovConst (naturalAdd 18 (naturalAdd first j)) (ahScalar (naturalAdd first j))))
88    tail)))
89def ahPrologueWith = (lambda unrestricted scaled : Nat . (lambda unrestricted iterations : Nat . (lambda unrestricted tail : (family SM86Program) .
90  (saS2R 0 (constructor SM86SpecialRegister SM86ThreadIdX)
91  (saS2R 1 (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdX)
92  (saMovImm 3 16
93  (saMovImm 4 8
94  (saMovImm ahCounter iterations
95  (ahChoose scaled
96    (saMovConst 14 (saArgument adamWHalfSM86ScaleArgument)
97    (saMovConst 15 (naturalAdd (saArgument adamWHalfSM86ScaleArgument) 4)
98    (saOp (constructor SM86InstructionBody SM86LoadGlobal (saR 18) (saR 14) (saU 0) (saCtl saSB4 saWaitNone))
99      (ahScalarsFrom 1 (saMovImmAfter 5 0 (naturalAdd saWait5 saWait4)
100        (saImad 2 1 ahThreads 0 tail))))))
101    (ahScalarsFrom 0 (saMovImmAfter 5 0 saWait5
102      (saImad 2 1 ahThreads 0 tail))))))))))))
103def ahPrologue = (ahPrologueWith 0)
104-- the learner state (Realization's caller's: step size and epsilon words 0
105-- and 1, pending word 2) through argument 11's pointer: the step's two
106-- scalars into R24 and R25 (where ahScalarsFrom puts scalars 6 and 7) and
107-- the pending word into R27, under SB4
108def adamWHalfSM86StateArgument : Nat = 11
109def adamWHalfSM86StateConstantBytes : Nat = (saArgument 12)
110def ahStateLoads = (lambda unrestricted tail : (family SM86Program) .
111  (saMovConst 16 (saArgument adamWHalfSM86StateArgument)
112  (saMovConst 17 (naturalAdd (saArgument adamWHalfSM86StateArgument) 4)
113  (saOp (constructor SM86InstructionBody SM86LoadGlobal (saR 24) (saR 16) (saU 0) (saCtl saSB4 saWaitNone))
114  (saOp (constructor SM86InstructionBody SM86LoadGlobal (saR 25) (saR 16) (saU 4) (saCtl saSB4 saWaitNone))
115  (saOp (constructor SM86InstructionBody SM86LoadGlobal (saR 27) (saR 16) (saU 8) (saCtl saSB4 saWaitNone))
116    tail))))))
117-- scalars j from `first` below `last` from the constant bank into R18 + j
118def ahScalarsRange = (lambda unrestricted first : Nat . (lambda unrestricted last : Nat . (lambda unrestricted tail : (family SM86Program) .
119  (saFor (naturalSaturatingSubtract last first)
120    (lambda unrestricted j : Nat .
121      (saMovConst (naturalAdd 18 (naturalAdd first j)) (ahScalar (naturalAdd first j))))
122    tail))))
123-- nothing pending (word 2 zero, once the loads have landed): the thread exits
124def ahPendingOrExit = (lambda unrestricted tail : (family SM86Program) .
125  (saGreater saP0 27 0
126  (saUnless saP0 (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) tail)))
127-- T for side thread tid - first of the CTA's linear index
128def ahSideThread = (lambda unrestricted first : Nat . (lambda unrestricted threads : Nat .
129  (lambda unrestricted gridX : Nat . (lambda unrestricted tail : (family SM86Program) .
130    (saAddImm 0 0 (naturalSaturatingSubtract 4294967296 first)
131    (saImad 1 26 gridX 1
132    (saImad 2 1 threads 0 tail)))))))
133-- A side warp's prologue (GroupedLaunchSM86.sm86ProgramWithSideWarps): the
134-- CTA's threads past its first `first` run the update beside the program
135-- the first run. Its block is the CTA's linear index over a grid `gridX`
136-- wide (CTA y in R26), of `threads` side threads (thread tid - first); the
137-- elements a thread takes are the standalone image's for that T, so each
138-- element's update is the same instructions on the same words.
139def ahSidePrologueWith = (lambda unrestricted scaled : Nat . (lambda unrestricted first : Nat .
140  (lambda unrestricted threads : Nat . (lambda unrestricted gridX : Nat .
141  (lambda unrestricted iterations : Nat . (lambda unrestricted tail : (family SM86Program) .
142  (saS2R 0 (constructor SM86SpecialRegister SM86ThreadIdX)
143  (saS2R 1 (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdX)
144  (saS2R 26 (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdY)
145  (saMovImm 3 16
146  (saMovImm 4 8
147  (saMovImm ahCounter iterations
148  (ahChoose scaled
149    (saMovConst 14 (saArgument adamWHalfSM86ScaleArgument)
150    (saMovConst 15 (naturalAdd (saArgument adamWHalfSM86ScaleArgument) 4)
151    (saOp (constructor SM86InstructionBody SM86LoadGlobal (saR 18) (saR 14) (saU 0) (saCtl saSB4 saWaitNone))
152      (ahStateLoads (ahScalarsRange 1 6 (saMovImmAfter 5 0 (naturalAdd saWait5 saWait4)
153        (ahPendingOrExit (ahSideThread first threads gridX tail))))))))
154    (ahStateLoads (ahScalarsRange 0 6 (saMovImmAfter 5 0 (naturalAdd saWait5 saWait4)
155      (ahPendingOrExit (ahSideThread first threads gridX tail))))))))))))))))))
156def ahGuarded = (lambda unrestricted body : (family SM86InstructionBody) . (lambda unrestricted tail : (family SM86Program) .
157  (saUnless saP0 body tail)))
158-- a store whose register read SB5 counts
159def ahStoreControl : (family SM86Control) =
160  (constructor SM86Control SM86ControlValue (byte 15) (constructor SM86YieldMode SM86Continue)
161    saNone saSB5 (byte 0) (byte 0))
162
163-- bounds, the four pointers and the halves'
164def ahAddresses = (lambda unrestricted quads : Nat . (lambda unrestricted tail : (family SM86Program) .
165  (saGreater saP0 2 (naturalPredecessor quads)
166  (saWide 6 2 3 (saArgument adamWHalfSM86ParameterArgument)
167  (saWide 8 2 3 (saArgument adamWHalfSM86GradientArgument)
168  (saWide 10 2 3 (saArgument adamWHalfSM86FirstMomentArgument)
169  (saWide 12 2 3 (saArgument adamWHalfSM86SecondMomentArgument)
170  (saWide 30 2 4 ahHalfPointer tail))))))))
171
172-- bank b's quad (pointer 6 + 2b) into R32 + 4b, under SB b; the first waits
173-- for the last iteration's stores' reads
174def ahLoadBank = (lambda unrestricted b : Nat . (lambda unrestricted barrier : (family SM86Barrier) .
175  (lambda unrestricted wait : Nat . (lambda unrestricted tail : (family SM86Program) .
176    (ahGuarded (constructor SM86InstructionBody SM86LoadGlobalWide (saR (ahValue 0 b)) (saR (naturalAdd 6 (naturalMultiply 2 b))) (saU 0)
177      (saCtl barrier wait)) tail)))))
178def ahLoads = (lambda unrestricted tail : (family SM86Program) .
179  (ahLoadBank 0 saSB0 saWait5
180  (ahLoadBank 1 saSB1 saWaitNone
181  (ahLoadBank 2 saSB2 saWaitNone
182  (ahLoadBank 3 saSB3 saWaitNone tail)))))
183
184-- element e's update: adamWF32MathProgram's instructions on its registers
185def ahUpdate = (lambda unrestricted e : Nat . (lambda unrestricted tail : (family SM86Program) .
186  (let unrestricted p = (ahValue e 0) in (let unrestricted g = (ahValue e 1) in
187  (let unrestricted m = (ahValue e 2) in (let unrestricted v = (ahValue e 3) in
188  (let unrestricted t0 = (ahTemporary e 0) in (let unrestricted t1 = (ahTemporary e 1) in
189  (let unrestricted t2 = (ahTemporary e 2) in
190  (saFmul g g 18 (naturalSelect e saWaitNone (naturalAdd saWait0 (naturalAdd saWait1 (naturalAdd saWait2 saWait3))))
191  (saFmul t0 g 20 saWaitNone
192  (saFfma m m 19 t0
193  (saFmul t2 g g saWaitNone
194  (saFmul t2 t2 22 saWaitNone
195  (saFfma v v 21 t2
196  (saMufu t1 v (constructor SM86MultiFunction SM86SquareRoot) saWaitNone
197  (saFadd t1 t1 25 saWait4
198  (saMufu t1 t1 (constructor SM86MultiFunction SM86Reciprocal) saWaitNone
199  (saFmul t0 m 24 saWaitNone
200  (saFmul t0 t0 t1 saWait4
201  (saFneg t0 t0
202  (saFmul p p 23 saWaitNone
203  (saFadd p p t0 saWaitNone tail)))))))))))))))))))))))
204
205-- the quad's P, M, V (128 bits each) and its two half pairs (64 bits)
206def ahStores = (lambda unrestricted tail : (family SM86Program) .
207  (saPack ahHalfPairs (ahValue 1 0) (ahValue 0 0)
208  (saPack (succ ahHalfPairs) (ahValue 3 0) (ahValue 2 0)
209  (ahGuarded (constructor SM86InstructionBody SM86StoreGlobalWide (saR 6) (saR (ahValue 0 0)) (saU 0) ahStoreControl)
210  (ahGuarded (constructor SM86InstructionBody SM86StoreGlobalWide (saR 10) (saR (ahValue 0 2)) (saU 0) ahStoreControl)
211  (ahGuarded (constructor SM86InstructionBody SM86StoreGlobalWide (saR 12) (saR (ahValue 0 3)) (saU 0) ahStoreControl)
212  (ahGuarded (constructor SM86InstructionBody SM86StoreGlobal64 (saR 30) (saR ahHalfPairs) (saU 0) ahStoreControl) tail)))))))
213
214def ahBody = (lambda unrestricted quads : Nat .
215  (ahAddresses quads (ahLoads (ahUpdate 0 (ahUpdate 1 (ahUpdate 2 (ahUpdate 3 (ahStores sm86ProgramEmpty))))))))
216
217-- ---- several quads a thread in flight (side warps) ----
218-- `wide` quads a thread an iteration: T, T + S, .., T + (wide - 1) S (S the
219-- threads), every bank's loads issued before any update, so a thread keeps
220-- `wide` times the bytes in flight -- what a side warp, a few to an SM,
221-- needs to stream at the memory's rate. Each element's update is ahUpdate's
222-- instructions (ahWideUpdate), its quad's stores ahStores'. Registers: quad
223-- q's banks at 32 + 16 q + 4 b + e, the temporaries after the last quad's
224-- (shared: the updates run one after another), each quad's half pairs after
225-- those; quad q's guard P0 (q 0) or P(q + 1) (P1 counts the loop). At most
226-- four quads (P2..P4).
227def ahwValue = (lambda unrestricted q : Nat . (lambda unrestricted e : Nat . (lambda unrestricted b : Nat .
228  (naturalAdd 32 (naturalAdd (naturalMultiply 16 q) (naturalAdd (naturalMultiply 4 b) e))))))
229def ahwTemporary = (lambda unrestricted wide : Nat . (lambda unrestricted e : Nat . (lambda unrestricted i : Nat .
230  (naturalAdd 32 (naturalAdd (naturalMultiply 16 wide) (naturalAdd (naturalMultiply 3 e) i))))))
231def ahwHalfPairs = (lambda unrestricted wide : Nat . (lambda unrestricted q : Nat .
232  (naturalAdd 44 (naturalAdd (naturalMultiply 16 wide) (naturalMultiply 2 q)))))
233def ahwPredicate = (lambda unrestricted q : Nat .
234  (nat-eliminate (lambda unrestricted n : Nat . (family SM86Predicate)) saP0
235    (lambda unrestricted p : Nat . (lambda unrestricted ignored : (family SM86Predicate) .
236      (nat-eliminate (lambda unrestricted n : Nat . (family SM86Predicate)) (constructor SM86Predicate SM86Predicate2)
237        (lambda unrestricted p2 : Nat . (lambda unrestricted ignored2 : (family SM86Predicate) .
238          (nat-eliminate (lambda unrestricted n : Nat . (family SM86Predicate)) (constructor SM86Predicate SM86Predicate3)
239            (lambda unrestricted p3 : Nat . (lambda unrestricted ignored3 : (family SM86Predicate) .
240              (constructor SM86Predicate SM86Predicate4))) p2))) p))) q))
241def ahwBankBarrier = (lambda unrestricted b : Nat .
242  (nat-eliminate (lambda unrestricted n : Nat . (family SM86Barrier)) saSB0
243    (lambda unrestricted p : Nat . (lambda unrestricted ignored : (family SM86Barrier) .
244      (nat-eliminate (lambda unrestricted n : Nat . (family SM86Barrier)) saSB1
245        (lambda unrestricted p2 : Nat . (lambda unrestricted ignored2 : (family SM86Barrier) .
246          (nat-eliminate (lambda unrestricted n : Nat . (family SM86Barrier)) saSB2
247            (lambda unrestricted p3 : Nat . (lambda unrestricted ignored3 : (family SM86Barrier) . saSB3)) p2))) p))) b))
248-- bounds (quad q's guard: T + q S at or past `quads`), then the pointers
249def ahwAddresses = (lambda unrestricted wide : Nat . (lambda unrestricted stride : Nat . (lambda unrestricted quads : Nat .
250  (lambda unrestricted tail : (family SM86Program) .
251  (saGreater saP0 2 (naturalPredecessor quads)
252  (saFor (naturalPredecessor wide) (lambda unrestricted j : Nat . (lambda unrestricted rest : (family SM86Program) .
253    (saAddImm 27 2 (naturalMultiply (succ j) stride)
254    (saGreater (ahwPredicate (succ j)) 27 (naturalPredecessor quads) rest))))
255  (saWide 6 2 3 (saArgument adamWHalfSM86ParameterArgument)
256  (saWide 8 2 3 (saArgument adamWHalfSM86GradientArgument)
257  (saWide 10 2 3 (saArgument adamWHalfSM86FirstMomentArgument)
258  (saWide 12 2 3 (saArgument adamWHalfSM86SecondMomentArgument)
259  (saWide 30 2 4 ahHalfPointer tail)))))))))))
260-- bank b of every quad, under SB b; the first waits for the last
261-- iteration's stores' reads
262def ahwLoads = (lambda unrestricted wide : Nat . (lambda unrestricted stride : Nat . (lambda unrestricted tail : (family SM86Program) .
263  (saFor (naturalMultiply 4 wide) (lambda unrestricted i : Nat .
264    (let unrestricted b = (naturalDivideUnchecked i wide) in (let unrestricted q = (naturalModuloUnchecked i wide) in
265      (saUnless (ahwPredicate q)
266        (constructor SM86InstructionBody SM86LoadGlobalWide (saR (ahwValue q 0 b)) (saR (naturalAdd 6 (naturalMultiply 2 b)))
267          (saU (naturalMultiply q (naturalMultiply 16 stride)))
268          (saCtl (ahwBankBarrier b) (naturalSelect i saWaitNone saWait5)))))))
269    tail))))
270-- element e of quad q: ahUpdate's instructions on its registers, the first
271-- waiting for every bank's loads
272def ahwUpdate = (lambda unrestricted wide : Nat . (lambda unrestricted q : Nat . (lambda unrestricted e : Nat .
273  (lambda unrestricted tail : (family SM86Program) .
274  (let unrestricted p = (ahwValue q e 0) in (let unrestricted g = (ahwValue q e 1) in
275  (let unrestricted m = (ahwValue q e 2) in (let unrestricted v = (ahwValue q e 3) in
276  (let unrestricted t0 = (ahwTemporary wide e 0) in (let unrestricted t1 = (ahwTemporary wide e 1) in
277  (let unrestricted t2 = (ahwTemporary wide e 2) in
278  (saFmul g g 18 (naturalSelect (naturalAdd q e) saWaitNone (naturalAdd saWait0 (naturalAdd saWait1 (naturalAdd saWait2 saWait3))))
279  (saFmul t0 g 20 saWaitNone
280  (saFfma m m 19 t0
281  (saFmul t2 g g saWaitNone
282  (saFmul t2 t2 22 saWaitNone
283  (saFfma v v 21 t2
284  (saMufu t1 v (constructor SM86MultiFunction SM86SquareRoot) saWaitNone
285  (saFadd t1 t1 25 saWait4
286  (saMufu t1 t1 (constructor SM86MultiFunction SM86Reciprocal) saWaitNone
287  (saFmul t0 m 24 saWaitNone
288  (saFmul t0 t0 t1 saWait4
289  (saFneg t0 t0
290  (saFmul p p 23 saWaitNone
291  (saFadd p p t0 saWaitNone tail)))))))))))))))))))))))))
292-- quad q's P, M, V and half pairs, at its offsets
293def ahwStores = (lambda unrestricted wide : Nat . (lambda unrestricted stride : Nat . (lambda unrestricted tail : (family SM86Program) .
294  (saFor wide (lambda unrestricted q : Nat . (lambda unrestricted rest : (family SM86Program) .
295    (let unrestricted pred = (ahwPredicate q) in
296    (let unrestricted wideOffset = (naturalMultiply q (naturalMultiply 16 stride)) in
297    (let unrestricted halves = (ahwHalfPairs wide q) in
298    (saPack halves (ahwValue q 1 0) (ahwValue q 0 0)
299    (saPack (succ halves) (ahwValue q 3 0) (ahwValue q 2 0)
300    (saUnless pred (constructor SM86InstructionBody SM86StoreGlobalWide (saR 6) (saR (ahwValue q 0 0)) (saU wideOffset) ahStoreControl)
301    (saUnless pred (constructor SM86InstructionBody SM86StoreGlobalWide (saR 10) (saR (ahwValue q 0 2)) (saU wideOffset) ahStoreControl)
302    (saUnless pred (constructor SM86InstructionBody SM86StoreGlobalWide (saR 12) (saR (ahwValue q 0 3)) (saU wideOffset) ahStoreControl)
303    (saUnless pred (constructor SM86InstructionBody SM86StoreGlobal64 (saR 30) (saR halves)
304      (saU (naturalMultiply q (naturalMultiply 8 stride))) ahStoreControl) rest)))))))))))
305    tail))))
306def ahWideBody = (lambda unrestricted wide : Nat . (lambda unrestricted stride : Nat . (lambda unrestricted quads : Nat .
307  (ahwAddresses wide stride quads
308    (ahwLoads wide stride
309      (saFor (naturalMultiply 4 wide)
310        (lambda unrestricted k : Nat . (ahwUpdate wide (naturalDivideUnchecked k 4) (naturalModuloUnchecked k 4)))
311        (ahwStores wide stride sm86ProgramEmpty)))))))
312
313-- ---- the loop's control: the next quad, one iteration fewer ----
314def ahControl = (lambda unrestricted stride : Nat . (lambda unrestricted tail : (family SM86Program) .
315  (saAddImm 2 2 stride
316  (saAddImm ahCounter ahCounter 4294967295
317  (saGreater saP1 ahCounter 0 tail)))))
318def ahControlCount : Nat = 3
319-- back to the body's first instruction (the displacement counts the body,
320-- the control and the branch, from the branch's successor)
321def ahBranch = (lambda unrestricted body : Nat .
322  (saWhen saP1 (constructor SM86InstructionBody SM86Branch
323    (saU (naturalSaturatingSubtract 4294967296 (naturalMultiply 16 (naturalAdd body (naturalAdd ahControlCount 1)))))
324    (sm86Unsigned32 (byte 255) (byte 255) (byte 131) (byte 3))
325    sm86SafeControl)
326    sm86ProgramEmpty))
327
328def ahExit : (family SM86Program) =
329  (saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty)
330
331-- `finish` is applied to the prologue and the body on their own (e.g. the
332-- stalls compacted and drained); the control keeps the largest stalls.
333-- looped 1: the loop with its back edge; looped 0: what the schedule checks
334-- read, the body twice and no branch (it covers both ways into the body,
335-- and every iteration leaves the same things in flight: the stores' reads
336-- under SB5, nothing else)
337-- `prologue` iterations tail: ahPrologueWith scaled, or a side job's
338-- `body` the loop body (ahBody quads, or ahWideBody's), `stride` the
339-- quads the loop's T advances an iteration
340def ahAssembleBody = (lambda unrestricted prologue : (pi unrestricted iterations : Nat . (pi unrestricted tail : (family SM86Program) . (family SM86Program))) .
341  (lambda unrestricted program : (family SM86Program) .
342  (lambda unrestricted finish : (pi unrestricted piece : (family SM86Program) . (family SM86Program)) .
343  (lambda unrestricted looped : Nat .
344  (lambda unrestricted quads : Nat . (lambda unrestricted stride : Nat .
345    (let unrestricted body = (finish program) in
346    (let unrestricted once = (sm86ProgramAppend body (ahControl stride sm86ProgramEmpty)) in
347    (sm86ProgramAppend (finish (prologue (ahIterations quads stride) sm86ProgramEmpty))
348      (sm86ProgramAppend once
349        (nat-eliminate (lambda unrestricted n : Nat . (family SM86Program))
350          (sm86ProgramAppend once ahExit)
351          (lambda unrestricted predecessor : Nat . (lambda unrestricted induction : (family SM86Program) .
352            (sm86ProgramAppend (ahBranch (sm86ProgramCount body)) ahExit)))
353          looped)))))))))))
354def ahAssembleFrom = (lambda unrestricted prologue : (pi unrestricted iterations : Nat . (pi unrestricted tail : (family SM86Program) . (family SM86Program))) .
355  (lambda unrestricted finish : (pi unrestricted piece : (family SM86Program) . (family SM86Program)) .
356  (lambda unrestricted looped : Nat .
357  (lambda unrestricted quads : Nat . (lambda unrestricted stride : Nat .
358    (ahAssembleBody prologue (ahBody quads) finish looped quads stride))))))
359def ahAssembleWith = (lambda unrestricted scaled : Nat . (ahAssembleFrom (ahPrologueWith scaled)))
360-- thread T 0 stores the step's scalars and 1 (pending) to the learner state
361-- (argument 11's pointer): the update the next forward's side warps finish
362def ahLeadStore = (lambda unrestricted tail : (family SM86Program) .
363  (saGreater (constructor SM86Predicate SM86Predicate2) 2 0
364  (saMovConst 16 (saArgument adamWHalfSM86StateArgument)
365  (saMovConst 17 (naturalAdd (saArgument adamWHalfSM86StateArgument) 4)
366  (saMovImm 27 1
367  (saUnless (constructor SM86Predicate SM86Predicate2)
368    (constructor SM86InstructionBody SM86StoreGlobal64 (saR 16) (saR 24) (saU 0) ahStoreControl)
369  (saUnless (constructor SM86Predicate SM86Predicate2)
370    (constructor SM86InstructionBody SM86StoreGlobal (saR 16) (saR 27) (saU 8) ahStoreControl)
371    tail)))))))
372def ahLeadPrologueWith = (lambda unrestricted iterations : Nat . (lambda unrestricted tail : (family SM86Program) .
373  (ahPrologueWith 1 iterations (ahLeadStore tail))))
374def ahAssemble = (ahAssembleWith 0)
375
376-- the image for `quads` element quads (the bank's elements over four) over
377-- a grid whose threads number `stride`
378def adamWHalfSM86 = (lambda unrestricted finish : (pi unrestricted piece : (family SM86Program) . (family SM86Program)) .
379  (ahAssemble finish 1))
380def adamWHalfSM86Checked = (lambda unrestricted finish : (pi unrestricted piece : (family SM86Program) . (family SM86Program)) .
381  (ahAssemble finish 0))
382
383-- the image whose gradient factor comes from argument 10's pointer
384def adamWHalfSM86Scaled = (lambda unrestricted finish : (pi unrestricted piece : (family SM86Program) . (family SM86Program)) .
385  (ahAssembleWith 1 finish 1))
386def adamWHalfSM86ScaledChecked = (lambda unrestricted finish : (pi unrestricted piece : (family SM86Program) . (family SM86Program)) .
387  (ahAssembleWith 1 finish 0))
388
389-- the scaled image whose step size and epsilon are also the device's: words
390-- 0 and 1 at argument 11's pointer (ahStateLoads, under SB4, which every
391-- element's update waits for before it reads them), as a device-side
392-- schedule leaves them; word 2 is loaded and not read
393def ahDevicePrologue = (lambda unrestricted iterations : Nat . (lambda unrestricted tail : (family SM86Program) .
394  (ahPrologueWith 1 iterations (ahStateLoads tail))))
395def adamWHalfSM86ScaledDevice = (lambda unrestricted finish : (pi unrestricted piece : (family SM86Program) . (family SM86Program)) .
396  (ahAssembleFrom ahDevicePrologue finish 1))
397def adamWHalfSM86ScaledDeviceChecked = (lambda unrestricted finish : (pi unrestricted piece : (family SM86Program) . (family SM86Program)) .
398  (ahAssembleFrom ahDevicePrologue finish 0))
399
400-- the scaled image on side warps: `threads` threads a CTA past its first
401-- `first`, over a grid `gridX` wide, `stride` the side threads' count
402def adamWHalfSM86ScaledSide = (lambda unrestricted finish : (pi unrestricted piece : (family SM86Program) . (family SM86Program)) .
403  (lambda unrestricted wide : Nat .
404  (lambda unrestricted first : Nat . (lambda unrestricted threads : Nat . (lambda unrestricted gridX : Nat .
405  (lambda unrestricted quads : Nat . (lambda unrestricted stride : Nat .
406    (ahAssembleBody (ahSidePrologueWith 1 first threads gridX) (ahWideBody wide stride quads)
407      finish 1 quads (naturalMultiply wide stride)))))))))
408-- at most four quads a thread (the guard predicates), within the
409-- registers: the last half pair at 45 + 18 wide
410def adamWHalfSM86SideWideAdmitted = (lambda unrestricted wide : Nat .
411  (naturalAnd (naturalLess 0 wide) (naturalLessOrEqual wide 4)))
412
413-- the scaled image that also leaves its step pending for the side warps
414-- (ahLeadStore); its constants run through argument 11
415def adamWHalfSM86ScaledLead = (lambda unrestricted finish : (pi unrestricted piece : (family SM86Program) . (family SM86Program)) .
416  (ahAssembleFrom ahLeadPrologueWith finish 1))
417
418def adamWHalfSM86Threads : Nat = ahThreads
419def adamWHalfSM86Registers : Nat = 64
420-- the elements a thread takes an iteration
421def adamWHalfSM86Elements : Nat = 4
422def adamWHalfSM86Iterations = ahIterations
423-- through the half pointer, in whole 8-byte parameter words
424def adamWHalfSM86ConstantBytes : Nat = (saArgument adamWHalfSM86ArgumentCount)
425-- where the step-dependent scalars sit in the parameter block (the host
426-- writes each update's there)
427def adamWHalfSM86StepSizeOffset : Nat = (ahScalar 6)
428def adamWHalfSM86EpsilonOffset : Nat = (ahScalar 7)