PartitionedRegression.Build.CoordinateMemory

1module PartitionedRegression.Build.CoordinateMemory
2
3import PartitionedRegression.Coordinates
4import Component.Pointwise
5import Hardware.Nvidia.SM86.Command.WholeProgramPlan
6import Realization.Nvidia.SM86.MemoryArithmeticSM86
7import Std.Fin
8import Realization.Nvidia.SM86.PointwiseMemorySM86
9import Platform.Linux.Nvidia.ProgramPlan
10import Platform.Linux.Nvidia.PlanHost
11import Platform.Linux.Nvidia.PlanHostRequest
12import Platform.RunPod.Nvidia.SM86
13import Runtime.ArenaCertificate
14import Model.Word64
15import Model.Parameter
16import Std.List
17import Std.Natural
18
19family CoordinateDevicePhase : Type 0
20constructor CoordinateDeviceForward
21constructor CoordinateDeviceBackward
22constructor CoordinateDeviceUpdate
23
24end-family
25
26-- Relative regions of the device's persistent binary64 allocation. Their
27-- addresses are supplied by the shared target plan after placement. The
28-- parameter and accumulator are seeded once; only the residual changes on
29-- the host between rounds. All arithmetic and intermediate stores are GPU work.
30def cdmParameter : Nat =
31  0
32
33def cdmAccumulator : Nat =
34  (naturalAdd cdmParameter coordinateWordBytes)
35
36def cdmFeature : Nat =
37  (naturalAdd cdmAccumulator coordinateWordBytes)
38
39def cdmResidual : Nat =
40  (naturalAdd cdmFeature coordinateWordBytes)
41
42def cdmPartial : Nat =
43  (naturalAdd cdmResidual coordinateWordBytes)
44
45def cdmGradient : Nat =
46  (naturalAdd cdmPartial coordinateWordBytes)
47
48def cdmDecay : Nat =
49  (naturalAdd cdmGradient coordinateWordBytes)
50
51def cdmRate : Nat =
52  (naturalAdd cdmDecay coordinateWordBytes)
53
54def cdmBytes : Nat =
55  (naturalAdd cdmRate coordinateWordBytes)
56
57def cdmResident =
58  (lambda unrestricted name : Bytes .
59    (lambda unrestricted offset : Nat .
60      (lambda unrestricted tail : (family ArenaResidents) .
61        (constructor
62          ArenaResidents
63          ArenaResidentsNext
64          (constructor
65            ArenaResident
66            ArenaResidentValue
67            name
68            offset
69            coordinateWordBytes
70            coordinateWordBytes)
71          tail))))
72
73def cdmResidents =
74  (cdmResident
75    b"parameter"
76    cdmParameter
77    (cdmResident
78      b"accumulator"
79      cdmAccumulator
80      (cdmResident
81        b"feature"
82        cdmFeature
83        (cdmResident
84          b"residual"
85          cdmResidual
86          (cdmResident
87            b"partial"
88            cdmPartial
89            (cdmResident
90              b"gradient"
91              cdmGradient
92              (cdmResident
93                b"decay"
94                cdmDecay
95                (cdmResident b"rate" cdmRate (constructor ArenaResidents ArenaResidentsEnd)))))))))
96
97def cdmCertified : (equal Nat (arenaCertificate cdmBytes cdmResidents) 1) =
98  (refl Nat 1)
99
100def cdmPhases =
101  (constructor
102    StdList
103    StdListCons
104    (family CoordinateDevicePhase)
105    (constructor CoordinateDevicePhase CoordinateDeviceForward)
106    (constructor
107      StdList
108      StdListCons
109      (family CoordinateDevicePhase)
110      (constructor CoordinateDevicePhase CoordinateDeviceBackward)
111      (constructor
112        StdList
113        StdListCons
114        (family CoordinateDevicePhase)
115        (constructor CoordinateDevicePhase CoordinateDeviceUpdate)
116        (constructor StdList StdListEmpty (family CoordinateDevicePhase)))))
117
118def cdmForwardOrdinal : Nat =
119  0
120
121def cdmBackwardOrdinal : Nat =
122  (succ cdmForwardOrdinal)
123
124def cdmUpdateOrdinal : Nat =
125  (succ cdmBackwardOrdinal)
126
127def cdmOrdinal =
128  (lambda unrestricted phase : (family CoordinateDevicePhase) .
129    (eliminate
130      CoordinateDevicePhase
131      (lambda unrestricted value : (family CoordinateDevicePhase) . Nat)
132      phase
133      (branch CoordinateDeviceForward . cdmForwardOrdinal)
134      (branch CoordinateDeviceBackward . cdmBackwardOrdinal)
135      (branch CoordinateDeviceUpdate . cdmUpdateOrdinal)))
136
137def cdmPointer =
138  (lambda unrestricted base : Nat .
139    (lambda unrestricted offset : Nat .
140      (lambda unrestricted tail : (family StdList Nat) .
141        (constructor StdList StdListCons Nat (naturalAdd base offset) tail))))
142
143def cdmLaunch =
144  (lambda unrestricted program : Nat .
145    (lambda unrestricted data : Nat .
146      (lambda unrestricted phase : (family CoordinateDevicePhase) .
147        (let unrestricted pointer =
148          (cdmPointer data)
149          in
150          (let unrestricted nil =
151            (constructor StdList StdListEmpty Nat)
152            in
153            (eliminate
154              CoordinateDevicePhase
155              (lambda unrestricted current : (family CoordinateDevicePhase) .
156                (family NvidiaLaunchTemplate))
157              phase
158              (branch
159                CoordinateDeviceForward
160                .
161                (memoryArithmeticLaunch
162                  (constructor MemoryArithmeticOperation MemoryNarrowedProduct)
163                  program
164                  (pointer cdmParameter (pointer cdmFeature (pointer cdmPartial nil)))
165                  (refl Nat 3)))
166              (branch
167                CoordinateDeviceBackward
168                .
169                (memoryArithmeticLaunch
170                  (constructor MemoryArithmeticOperation MemoryNarrowedProduct)
171                  program
172                  (pointer cdmResidual (pointer cdmFeature (pointer cdmGradient nil)))
173                  (refl Nat 3)))
174              (branch
175                CoordinateDeviceUpdate
176                .
177                (memoryArithmeticLaunch
178                  (constructor MemoryArithmeticOperation MemoryScaledUpdate)
179                  program
180                  (pointer
181                    cdmParameter
182                    (pointer
183                      cdmAccumulator
184                      (pointer cdmGradient (pointer cdmDecay (pointer cdmRate nil)))))
185                  (refl Nat 5)))))))))
186
187def cdmLaunches =
188  (lambda unrestricted program : Nat .
189    (lambda unrestricted data : Nat .
190      (nvidiaLaunchSequence
191        (stdListMap
192          (family CoordinateDevicePhase)
193          (family NvidiaLaunchTemplate)
194          (cdmLaunch program data)
195          cdmPhases))))
196
197def cdmBatch =
198  (lambda unrestricted phase : (family CoordinateDevicePhase) .
199    (constructor
200      NvidiaSubmissionSchedule
201      NvidiaSubmissionScheduleOne
202      (constructor
203        NvidiaSubmissionBatch
204        NvidiaSubmissionBatchValue
205        b"coordinate-phase"
206        (naturalMultiply (cdmOrdinal phase) prSemaphoreStride)
207        (constructor NvidiaLaunchReferences NvidiaLaunchReferencesRange (cdmOrdinal phase) 1))))
208
209def cdmSubmissions =
210  (stdListFold
211    (family CoordinateDevicePhase)
212    (family NvidiaSubmissionSchedule)
213    (lambda unrestricted phase : (family CoordinateDevicePhase) .
214      (lambda unrestricted tail : (family NvidiaSubmissionSchedule) .
215        (constructor NvidiaSubmissionSchedule NvidiaSubmissionScheduleAppend (cdmBatch phase) tail)))
216    (constructor NvidiaSubmissionSchedule NvidiaSubmissionScheduleEmpty)
217    cdmPhases)
218
219def cdmIssueCount : Nat =
220  (nvidiaSubmissionScheduleCount cdmSubmissions)
221
222def cdmPlan =
223  (lambda unrestricted profile : (family RunPodSM86TargetProfile) .
224    (nvidiaProgramPlan
225      memoryArithmeticRegions
226      cdmLaunches
227      cdmSubmissions
228      cdmBytes
229      (constructor NvidiaPlanHostMemory NvidiaPlanHostVideoMemory)
230      runPodSM86QMDArchitecture
231      runPodSM86ComputeClass
232      runPodSM86SPAVersion
233      runPodSM86QMDSlotBytes
234      (runPodSM86ProductIdentityPadded profile)
235      nvidiaPlanHostAmpere580
236      (modelWord64Natural (runPodSM86TotalVideoMemory profile))))