Realization.Nvidia.SM86.MemoryArithmeticSM86

1module Realization.Nvidia.SM86.MemoryArithmeticSM86
2
3import Accelerator.SM86.Control
4import Accelerator.SM86.Immediate
5import Accelerator.SM86.Instruction
6import Accelerator.SM86.InstructionEncoding
7import Accelerator.SM86.Operands
8import Accelerator.SM86.Types
9import Component.Pointwise
10import Hardware.Nvidia.SM86.Command.WholeProgramPlan
11import Model.Config
12import Model.Word64
13import Learning.ScaledGradientUpdate
14import Realization.Nvidia.SM86.PointwiseMemorySM86
15import Realization.Nvidia.SM86.ScaledGradientUpdateSM86
16import Std.Fin
17import Std.Word
18import Std.List
19import Std.Natural
20
21family MemoryArithmeticOperation : Type 0
22constructor MemoryNarrowedProduct
23constructor MemoryScaledUpdate
24
25end-family
26
27-- Register allocation: input pairs first, then pointer pairs. No high word
28-- is repurposed while a conversion can still consume it.
29def memoryProductInputCount : Nat =
30  2
31
32def memoryProductPointerCount : Nat =
33  (succ memoryProductInputCount)
34
35def memoryProductValueBase : Nat =
36  0
37
38def memoryProductPointerBase : Nat =
39  (memory64PairAt memoryProductValueBase memoryProductInputCount)
40
41def memoryProductLeft : Nat =
42  memoryProductValueBase
43
44def memoryProductRight : Nat =
45  (memory64PairAt memoryProductValueBase 1)
46
47def memoryProductOutputPointer : Nat =
48  (memory64PairAt memoryProductPointerBase memoryProductInputCount)
49
50def memoryProductIdentity : Bytes =
51  b"narrow-binary64-product-binary32-widen-v1"
52
53def memoryProductProgram =
54  (memory64Pointers
55    memoryProductPointerCount
56    memoryProductPointerBase
57    (memory64Loads
58      memoryProductInputCount
59      memoryProductValueBase
60      memoryProductPointerBase
61      (memory64Next
62        (constructor
63          SM86InstructionBody
64          SM86FloatConvert
65          (constructor
66            SM86FloatConversion
67            SM86NarrowBinary64
68            (memory64R memoryProductLeft)
69            (memory64P memoryProductLeft))
70          scaledUpdate64Control)
71        (memory64Next
72          (constructor
73            SM86InstructionBody
74            SM86FloatConvert
75            (constructor
76              SM86FloatConversion
77              SM86NarrowBinary64
78              (memory64R memoryProductRight)
79              (memory64P memoryProductRight))
80            scaledUpdate64Control)
81          (memory64Next
82            (constructor
83              SM86InstructionBody
84              SM86FloatMultiply
85              (memory64R memoryProductLeft)
86              (memory64R memoryProductLeft)
87              (memory64R memoryProductRight)
88              memory64Wait)
89            (memory64Next
90              (constructor
91                SM86InstructionBody
92                SM86FloatConvert
93                (constructor
94                  SM86FloatConversion
95                  SM86WidenBinary32
96                  (memory64P memoryProductLeft)
97                  (memory64R memoryProductLeft))
98                scaledUpdate64Control)
99              (memory64Store memoryProductOutputPointer memoryProductLeft memory64Exit)))))))
100
101-- This ABI binds five state slots; the arithmetic is a typed composition from
102-- the learner owner. Register allocation, initial reads and final stores are
103-- derived by the shared pointwise compiler, rather than restated by this kernel.
104def memoryUpdateInputCount : Nat =
105  5
106
107def memoryUpdateAssignments =
108  (scaledGradientAssignments
109    StdU64
110    memoryUpdateInputCount
111    (stdFinFirst 4)
112    (stdFinNextOf 4 (stdFinFirst 3))
113    (stdFinNextOf 4 (stdFinNextOf 3 (stdFinFirst 2)))
114    (stdFinNextOf 4 (stdFinNextOf 3 (stdFinNextOf 2 (stdFinFirst 1))))
115    (stdFinNextOf 4 (stdFinNextOf 3 (stdFinNextOf 2 (stdFinNextOf 1 (stdFinFirst 0)))))
116    (refl Nat 1))
117
118def memoryUpdateIdentity : Bytes =
119  b"accumulate-scaled-gradient-subtract-scaled-update-binary64-v1"
120
121def memoryUpdateProgram =
122  (memoryPointwiseProgram memoryUpdateInputCount memoryUpdateAssignments (refl Nat 1))
123
124def memoryArithmeticRegions =
125  (constructor
126    NvidiaDeviceRegions
127    NvidiaDeviceRegionsNext
128    (nvidiaSM86ScalarRegion
129      memoryProductIdentity
130      memoryProductProgram
131      b"binary64 memory arithmetic requires an SM86 realization")
132    (constructor
133      NvidiaDeviceRegions
134      NvidiaDeviceRegionsNext
135      (nvidiaSM86ScalarRegion
136        memoryUpdateIdentity
137        memoryUpdateProgram
138        b"binary64 memory arithmetic requires an SM86 realization")
139      (constructor NvidiaDeviceRegions NvidiaDeviceRegionsEnd)))
140
141-- The pointer-list length determines the constant block extent. A pairing
142-- supplies the same semantic argument list that the kernel ABI declares.
143def memoryArithmeticArgumentCount =
144  (lambda unrestricted operation : (family MemoryArithmeticOperation) .
145    (eliminate
146      MemoryArithmeticOperation
147      (lambda unrestricted value : (family MemoryArithmeticOperation) . Nat)
148      operation
149      (branch MemoryNarrowedProduct . memoryProductPointerCount)
150      (branch MemoryScaledUpdate . memoryUpdateInputCount)))
151
152def memoryArithmeticIdentity =
153  (lambda unrestricted operation : (family MemoryArithmeticOperation) .
154    (eliminate
155      MemoryArithmeticOperation
156      (lambda unrestricted value : (family MemoryArithmeticOperation) . Bytes)
157      operation
158      (branch MemoryNarrowedProduct . memoryProductIdentity)
159      (branch MemoryScaledUpdate . memoryUpdateIdentity)))
160
161def memoryArithmeticLaunch =
162  (lambda unrestricted operation : (family MemoryArithmeticOperation) .
163    (lambda unrestricted programBase : Nat .
164      (lambda unrestricted pointers : (family StdList Nat) .
165        (lambda erased arity : (equal Nat (stdListLength Nat pointers) (memoryArithmeticArgumentCount operation)) .
166          (nvidiaSM86PointerKernelLaunch
167            (memoryArithmeticIdentity operation)
168            memoryArithmeticRegions
169            programBase
170            pointers)))))