Realization.Nvidia.SM86.PointwiseMemorySM86

1module Realization.Nvidia.SM86.PointwiseMemorySM86
2
3import Accelerator.SM86.Control
4import Accelerator.SM86.Immediate
5import Accelerator.SM86.Instruction
6import Accelerator.SM86.InstructionEncoding
7import Accelerator.SM86.Operands
8import Accelerator.SM86.Program
9import Accelerator.SM86.Types
10import Component.Pointwise
11import Hardware.Nvidia.SM86.Command.WholeProgramPlan
12import Model.Word64
13import Realization.Nvidia.SM86.ScaledGradientUpdateSM86
14import Std.Fin
15import Std.List
16import Std.Natural
17import Std.Word
18
19family MemoryPointwiseFragment : Type 0
20constructor MemoryPointwiseFragmentOf
21field unrestricted memoryPointwiseFragmentProgram : (family SM86Program)
22field unrestricted memoryPointwiseFragmentResult : Nat
23field unrestricted memoryPointwiseFragmentEnd : Nat
24
25end-family
26
27-- Scalar memory realizations used by persistent state schedules. Each region
28-- launches one thread. Addresses are CB0 pointer arguments supplied by the
29-- enclosing plan; this module neither chooses device addresses nor builds a
30-- launch table. A submission fence makes stores visible to the next phase.
31-- Arithmetic results stay in the caller's device allocation between launches.
32def memory64WordBytes : Nat =
33  4
34
35def memory64PairWords : Nat =
36  2
37
38def memory64R =
39  (lambda unrestricted index : Nat . (sm86Register (nat-to-byte index)))
40
41def memory64P =
42  (lambda unrestricted index : Nat . (sm86RegisterPair (nat-to-byte index)))
43
44def memory64U =
45  sm86Unsigned32FromNaturalTruncated
46
47def memory64PairAt =
48  (lambda unrestricted base : Nat .
49    (lambda unrestricted index : Nat . (nat-add base (nat-multiply memory64PairWords index))))
50
51def memory64Next =
52  (lambda unrestricted body : (family SM86InstructionBody) .
53    (lambda unrestricted tail : (family SM86Program) .
54      (constructor SM86Program SM86ProgramNext (sm86Instruction body) tail)))
55
56def memory64Wait =
57  (constructor
58    SM86Control
59    SM86ControlValue
60    (byte 15)
61    (constructor SM86YieldMode SM86Continue)
62    (constructor SM86Barrier SM86BarrierNone)
63    (constructor SM86Barrier SM86BarrierNone)
64    (byte 3)
65    (byte 0))
66
67def memory64StoreControl =
68  (constructor
69    SM86Control
70    SM86ControlValue
71    (byte 15)
72    (constructor SM86YieldMode SM86Continue)
73    (constructor SM86Barrier SM86BarrierNone)
74    (constructor SM86Barrier SM86Barrier1)
75    (byte 3)
76    (byte 0))
77
78def memory64Exit =
79  (memory64Next
80    (constructor SM86InstructionBody SM86Exit memory64Wait)
81    (constructor SM86Program SM86ProgramEnd))
82
83-- Both words are loaded with the same source-address pair. Every successor
84-- waits for late reads and writes, including before a pointer is reused.
85def memory64Load =
86  (lambda unrestricted destination : Nat .
87    (lambda unrestricted pointer : Nat .
88      (lambda unrestricted tail : (family SM86Program) .
89        (memory64Next
90          (constructor
91            SM86InstructionBody
92            SM86LoadGlobal
93            (memory64R destination)
94            (memory64R pointer)
95            sm86Unsigned32Zero
96            scaledUpdate64Control)
97          (memory64Next
98            (constructor
99              SM86InstructionBody
100              SM86LoadGlobal
101              (memory64R (succ destination))
102              (memory64R pointer)
103              (memory64U memory64WordBytes)
104              scaledUpdate64Control)
105            tail)))))
106
107def memory64Store =
108  (lambda unrestricted pointer : Nat .
109    (lambda unrestricted source : Nat .
110      (lambda unrestricted tail : (family SM86Program) .
111        (memory64Next
112          (constructor
113            SM86InstructionBody
114            SM86StoreGlobal64
115            (memory64R pointer)
116            (memory64R source)
117            sm86Unsigned32Zero
118            memory64StoreControl)
119          tail))))
120
121def memory64Pointer =
122  (lambda unrestricted destination : Nat .
123    (lambda unrestricted argument : Nat .
124      (lambda unrestricted tail : (family SM86Program) .
125        (memory64Next
126          (constructor
127            SM86InstructionBody
128            SM86MoveConstant
129            (memory64R destination)
130            (byte 0)
131            (memory64U (nvidiaKernelPointerByteOffset argument))
132            memory64Wait)
133          (memory64Next
134            (constructor
135              SM86InstructionBody
136              SM86MoveConstant
137              (memory64R (succ destination))
138              (byte 0)
139              (memory64U (nat-add (nvidiaKernelPointerByteOffset argument) memory64WordBytes))
140              memory64Wait)
141            tail)))))
142
143def memory64Pointers =
144  (lambda unrestricted count : Nat .
145    (lambda unrestricted base : Nat .
146      (lambda unrestricted tail : (family SM86Program) .
147        (nat-eliminate
148          (lambda unrestricted n : Nat . (family SM86Program))
149          tail
150          (lambda unrestricted index : Nat .
151            (lambda unrestricted rest : (family SM86Program) .
152              (memory64Pointer (memory64PairAt base index) index rest)))
153          count))))
154
155def memory64Loads =
156  (lambda unrestricted count : Nat .
157    (lambda unrestricted base : Nat .
158      (lambda unrestricted pointers : Nat .
159        (lambda unrestricted tail : (family SM86Program) .
160          (nat-eliminate
161            (lambda unrestricted n : Nat . (family SM86Program))
162            tail
163            (lambda unrestricted index : Nat .
164              (lambda unrestricted rest : (family SM86Program) .
165                (memory64Load (memory64PairAt base index) (memory64PairAt pointers index) rest)))
166            count)))))
167
168-- Each assignment gets a temporary register range above the persistent value
169-- pairs. That range is reused only after the result has been copied back and
170-- all late reads/writes have retired. Pointer pairs follow the maximum range,
171-- so expression depth cannot silently collide with an address register.
172def memoryPointwiseEnd =
173  (constructor SM86Program SM86ProgramEnd)
174
175def memoryPointwiseCopyWord =
176  (lambda unrestricted destination : Nat .
177    (lambda unrestricted source : Nat .
178      (lambda unrestricted tail : (family SM86Program) .
179        (memory64Next
180          (constructor
181            SM86InstructionBody
182            SM86IntegerMultiplyAddImmediate
183            (memory64R destination)
184            (memory64R source)
185            sm86Unsigned32One
186            sm86ZeroRegister
187            memory64Wait)
188          tail))))
189
190def memoryPointwiseCopy =
191  (lambda unrestricted destination : Nat .
192    (lambda unrestricted source : Nat .
193      (lambda unrestricted tail : (family SM86Program) .
194        (memoryPointwiseCopyWord
195          destination
196          source
197          (memoryPointwiseCopyWord (succ destination) (succ source) tail)))))
198
199def memoryPointwiseImmediate =
200  (lambda unrestricted destination : Nat .
201    (lambda unrestricted value : Nat .
202      (lambda unrestricted tail : (family SM86Program) .
203        (memory64Next
204          (constructor
205            SM86InstructionBody
206            SM86MoveImmediate
207            (memory64R destination)
208            (memory64U value)
209            memory64Wait)
210          tail))))
211
212def memoryPointwiseUnarySupported =
213  (lambda unrestricted operation : (family PointwiseUnary) .
214    (eliminate
215      PointwiseUnary
216      (lambda unrestricted current : (family PointwiseUnary) . Nat)
217      operation
218      (branch PointwiseNegate . 1)
219      (branch PointwiseSquare . 1)
220      (branch PointwiseRoundBF16 . 0)
221      (branch PointwiseReciprocal . 0)
222      (branch PointwiseReciprocalSquareRoot . 0)
223      (branch PointwiseSquareRoot . 0)
224      (branch PointwiseTanh . 0)
225      (branch PointwiseExponentialBase2 . 0)))
226
227def memoryPointwiseBinarySupported =
228  (lambda unrestricted operation : (family PointwiseBinary) .
229    (eliminate
230      PointwiseBinary
231      (lambda unrestricted current : (family PointwiseBinary) . Nat)
232      operation
233      (branch PointwiseAdd . 1)
234      (branch PointwiseMultiply . 1)
235      (branch PointwiseSubtract . 1)
236      (branch PointwisePositiveGate . 0)))
237
238def memoryPointwiseExpressionSupported =
239  (lambda erased slots : Nat .
240    (lambda unrestricted expression : (family PointwiseExpression StdU64 slots) .
241      (pointwiseFold
242        StdU64
243        Nat
244        slots
245        (lambda unrestricted operation : (family PointwiseUnary) .
246          (lambda unrestricted child : Nat .
247            (naturalAnd (memoryPointwiseUnarySupported operation) child)))
248        (lambda unrestricted operation : (family PointwiseBinary) .
249          (lambda unrestricted left : Nat .
250            (lambda unrestricted right : Nat .
251              (naturalAnd (memoryPointwiseBinarySupported operation) (naturalAnd left right)))))
252        (lambda unrestricted bits : StdU64 . 1)
253        expression
254        (lambda unrestricted index : (family StdFin slots) . 1))))
255
256-- Unsupported branches below are inaccessible through memoryPointwiseProgram's
257-- admission proof. They contain no host or alternative arithmetic realization.
258-- Adding a branch requires its numeric contract and checked-model laws first.
259def memoryPointwiseUnaryProgram =
260  (lambda unrestricted operation : (family PointwiseUnary) .
261    (lambda unrestricted destination : Nat .
262      (lambda unrestricted source : Nat .
263        (eliminate
264          PointwiseUnary
265          (lambda unrestricted current : (family PointwiseUnary) . (family SM86Program))
266          operation
267          (branch
268            PointwiseNegate
269            .
270            (memoryPointwiseCopyWord
271              destination
272              source
273              (memoryPointwiseImmediate
274                (succ destination)
275                0x80000000
276                (memory64Next
277                  (constructor
278                    SM86InstructionBody
279                    SM86LogicThreeInputTruthTable
280                    (memory64R (succ destination))
281                    (memory64R (succ source))
282                    (memory64R (succ destination))
283                    (byte 0x3c)
284                    memory64Wait)
285                  memoryPointwiseEnd))))
286          (branch
287            PointwiseSquare
288            .
289            (scaledUpdate64Next
290              (constructor SM86Float64Operation SM86Float64Multiply)
291              (memory64P destination)
292              (memory64P source)
293              (memory64P source)
294              memoryPointwiseEnd))
295          (branch PointwiseRoundBF16 . memoryPointwiseEnd)
296          (branch PointwiseReciprocal . memoryPointwiseEnd)
297          (branch PointwiseReciprocalSquareRoot . memoryPointwiseEnd)
298          (branch PointwiseSquareRoot . memoryPointwiseEnd)
299          (branch PointwiseTanh . memoryPointwiseEnd)
300          (branch PointwiseExponentialBase2 . memoryPointwiseEnd)))))
301
302def memoryPointwiseBinaryProgram =
303  (lambda unrestricted operation : (family PointwiseBinary) .
304    (lambda unrestricted destination : Nat .
305      (lambda unrestricted left : Nat .
306        (lambda unrestricted right : Nat .
307          (let unrestricted emit =
308            (lambda unrestricted numeric : (family SM86Float64Operation) .
309              (scaledUpdate64Next
310                numeric
311                (memory64P destination)
312                (memory64P left)
313                (memory64P right)
314                memoryPointwiseEnd))
315            in
316            (eliminate
317              PointwiseBinary
318              (lambda unrestricted current : (family PointwiseBinary) . (family SM86Program))
319              operation
320              (branch PointwiseAdd . (emit (constructor SM86Float64Operation SM86Float64Add)))
321              (branch
322                PointwiseMultiply
323                .
324                (emit (constructor SM86Float64Operation SM86Float64Multiply)))
325              (branch
326                PointwiseSubtract
327                .
328                (emit (constructor SM86Float64Operation SM86Float64Subtract)))
329              (branch PointwisePositiveGate . memoryPointwiseEnd)))))))
330
331-- Observe each value at its production/use boundary, before a subsequent
332-- assignment can overwrite it. The IEEE realization supplies an empty observer;
333-- stricter numerical policies reuse this traversal without changing its order.
334def memoryPointwiseCompile =
335  (lambda unrestricted observe : (pi unrestricted result : Nat . (family SM86Program)) .
336    (lambda unrestricted slots : Nat .
337      (lambda unrestricted expression : (family PointwiseExpression StdU64 slots) .
338        (eliminate
339          PointwiseExpression
340          (lambda unrestricted current : (family PointwiseExpression StdU64 slots) .
341            (pi unrestricted next : Nat . (family MemoryPointwiseFragment)))
342          expression
343          (branch
344            PointwiseInput
345            index
346            .
347            (lambda unrestricted next : Nat .
348              (constructor
349                MemoryPointwiseFragment
350                MemoryPointwiseFragmentOf
351                (observe (memory64PairAt 0 (stdFinToNatural slots index)))
352                (memory64PairAt 0 (stdFinToNatural slots index))
353                next)))
354          (branch
355            PointwiseConstant
356            bits
357            .
358            (lambda unrestricted next : Nat .
359              (let unrestricted value =
360                (modelWord64Natural bits)
361                in
362                (constructor
363                  MemoryPointwiseFragment
364                  MemoryPointwiseFragmentOf
365                  (memoryPointwiseImmediate
366                    next
367                    (naturalModuloUnchecked value 0x100000000)
368                    (memoryPointwiseImmediate
369                      (succ next)
370                      (naturalDivideUnchecked value 0x100000000)
371                      (observe next)))
372                  next
373                  (naturalAdd next memory64PairWords)))))
374          (branch
375            PointwiseUnaryExpression
376            operation
377            operand
378            compile
379            .
380            (lambda unrestricted next : Nat .
381              (eliminate
382                MemoryPointwiseFragment
383                (lambda unrestricted current : (family MemoryPointwiseFragment) .
384                  (family MemoryPointwiseFragment))
385                (compile next)
386                (branch
387                  MemoryPointwiseFragmentOf
388                  program
389                  result
390                  end
391                  .
392                  (constructor
393                    MemoryPointwiseFragment
394                    MemoryPointwiseFragmentOf
395                    (sm86ProgramAppend
396                      program
397                      (sm86ProgramAppend
398                        (memoryPointwiseUnaryProgram operation end result)
399                        (observe end)))
400                    end
401                    (naturalAdd end memory64PairWords))))))
402          (branch
403            PointwiseBinaryExpression
404            operation
405            left
406            right
407            compileLeft
408            compileRight
409            .
410            (lambda unrestricted next : Nat .
411              (eliminate
412                MemoryPointwiseFragment
413                (lambda unrestricted current : (family MemoryPointwiseFragment) .
414                  (family MemoryPointwiseFragment))
415                (compileLeft next)
416                (branch
417                  MemoryPointwiseFragmentOf
418                  leftProgram
419                  leftResult
420                  leftEnd
421                  .
422                  (eliminate
423                    MemoryPointwiseFragment
424                    (lambda unrestricted current : (family MemoryPointwiseFragment) .
425                      (family MemoryPointwiseFragment))
426                    (compileRight leftEnd)
427                    (branch
428                      MemoryPointwiseFragmentOf
429                      rightProgram
430                      rightResult
431                      rightEnd
432                      .
433                      (constructor
434                        MemoryPointwiseFragment
435                        MemoryPointwiseFragmentOf
436                        (sm86ProgramAppend
437                          leftProgram
438                          (sm86ProgramAppend
439                            rightProgram
440                            (sm86ProgramAppend
441                              (memoryPointwiseBinaryProgram
442                                operation
443                                rightEnd
444                                leftResult
445                                rightResult)
446                              (observe rightEnd))))
447                        rightEnd
448                        (naturalAdd rightEnd memory64PairWords))))))))))))
449
450def memoryPointwiseValueEnd =
451  (lambda unrestricted slots : Nat .
452    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
453      (memory64PairAt 0 (naturalAdd slots (pointwiseProgramTemporaryValues StdU64 slots program)))))
454
455def memoryPointwiseSupported =
456  (lambda erased slots : Nat .
457    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
458      (eliminate
459        StdList
460        (lambda unrestricted current : (family StdList (family PointwiseAssignment StdU64 slots)) .
461          Nat)
462        program
463        (branch StdListEmpty . 1)
464        (branch
465          StdListCons
466          head
467          tail
468          rest
469          .
470          (eliminate
471            PointwiseAssignment
472            (lambda unrestricted current : (family PointwiseAssignment StdU64 slots) . Nat)
473            head
474            (branch
475              PointwiseAssign
476              destination
477              expression
478              .
479              (naturalAnd (memoryPointwiseExpressionSupported slots expression) rest)))))))
480
481def memoryPointwiseAdmitted =
482  (lambda unrestricted slots : Nat .
483    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
484      (naturalAnd
485        (memoryPointwiseSupported slots program)
486        (sm86RegisterSpanAdmitted (memory64PairAt (memoryPointwiseValueEnd slots program) slots)))))
487
488def memoryPointwiseAssignments =
489  (lambda unrestricted observe : (pi unrestricted result : Nat . (family SM86Program)) .
490    (lambda unrestricted slots : Nat .
491      (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
492        (lambda unrestricted tail : (family SM86Program) .
493          (eliminate
494            StdList
495            (lambda unrestricted current : (family StdList (family PointwiseAssignment StdU64 slots)) .
496              (family SM86Program))
497            program
498            (branch StdListEmpty . tail)
499            (branch
500              StdListCons
501              head
502              rest
503              next
504              .
505              (eliminate
506                PointwiseAssignment
507                (lambda unrestricted current : (family PointwiseAssignment StdU64 slots) .
508                  (family SM86Program))
509                head
510                (branch
511                  PointwiseAssign
512                  destination
513                  expression
514                  .
515                  (eliminate
516                    MemoryPointwiseFragment
517                    (lambda unrestricted current : (family MemoryPointwiseFragment) .
518                      (family SM86Program))
519                    (memoryPointwiseCompile observe slots expression (memory64PairAt 0 slots))
520                    (branch
521                      MemoryPointwiseFragmentOf
522                      code
523                      result
524                      end
525                      .
526                      (sm86ProgramAppend
527                        code
528                        (memoryPointwiseCopy
529                          (memory64PairAt 0 (stdFinToNatural slots destination))
530                          result
531                          next))))))))))))
532
533-- The use analysis preserves the distinction between an initialized input and
534-- a slot first defined by this program. An output-only slot is never loaded.
535-- Stores happen after all assignments, so every expression observes the same
536-- sequential register state and no later read sees a half-published update.
537def memoryPointwiseProgramWith =
538  (lambda unrestricted observe : (pi unrestricted result : Nat . (family SM86Program)) .
539    (lambda unrestricted store : (pi unrestricted pointer : Nat . (pi unrestricted source : Nat . (pi unrestricted tail : (family SM86Program) . (family SM86Program)))) .
540      (lambda unrestricted finish : (family SM86Program) .
541        (lambda unrestricted slots : Nat .
542          (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
543            (let unrestricted pointers =
544              (memoryPointwiseValueEnd slots program)
545              in
546              (let unrestricted stores =
547                (nat-eliminate
548                  (lambda unrestricted count : Nat . (family SM86Program))
549                  finish
550                  (lambda unrestricted index : Nat .
551                    (lambda unrestricted next : (family SM86Program) .
552                      (nat-eliminate
553                        (lambda unrestricted needed : Nat . (family SM86Program))
554                        next
555                        (lambda unrestricted previous : Nat .
556                          (lambda unrestricted ignored : (family SM86Program) .
557                            (store (memory64PairAt pointers index) (memory64PairAt 0 index) next)))
558                        (pointwiseProgramWrites StdU64 slots program index))))
559                  slots)
560                in
561                (let unrestricted body =
562                  (memoryPointwiseAssignments observe slots program stores)
563                  in
564                  (nat-eliminate
565                    (lambda unrestricted count : Nat . (family SM86Program))
566                    body
567                    (lambda unrestricted index : Nat .
568                      (lambda unrestricted next : (family SM86Program) .
569                        (let unrestricted read =
570                          (pointwiseProgramReadsInitial StdU64 slots program index)
571                          in
572                          (let unrestricted write =
573                            (pointwiseProgramWrites StdU64 slots program index)
574                            in
575                            (nat-eliminate
576                              (lambda unrestricted needed : Nat . (family SM86Program))
577                              next
578                              (lambda unrestricted previous : Nat .
579                                (lambda unrestricted ignored : (family SM86Program) .
580                                  (memory64Pointer
581                                    (memory64PairAt pointers index)
582                                    index
583                                    (nat-eliminate
584                                      (lambda unrestricted initial : Nat . (family SM86Program))
585                                      next
586                                      (lambda unrestricted before : Nat .
587                                        (lambda unrestricted unused : (family SM86Program) .
588                                        (memory64Load
589                                        (memory64PairAt 0 index)
590                                        (memory64PairAt pointers index)
591                                        next)))
592                                      read))))
593                              (naturalOr read write))))))
594                    slots)))))))))
595
596-- Existing IEEE consumers retain precisely the same instruction stream.
597def memoryPointwiseProgram =
598  (lambda unrestricted slots : Nat .
599    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
600      (lambda erased admitted : (equal Nat (memoryPointwiseAdmitted slots program) 1) .
601        (memoryPointwiseProgramWith
602          (lambda unrestricted result : Nat . memoryPointwiseEnd)
603          memory64Store
604          memory64Exit
605          slots
606          program))))