Realization.Nvidia.SM86.LinearStepPlanSM86

1module Realization.Nvidia.SM86.LinearStepPlanSM86
2
3import Accelerator.SM86.Instruction
4import Accelerator.SM86.Operands
5import Accelerator.SM86.Types
6import Accelerator.SM86.InstructionEncoding
7import Hardware.Nvidia.SM86.Command.WholeProgramPlan
8import Model.Config
9import Model.Word64
10import Realization.Nvidia.SM86.LinearStepSM86
11import Runtime.ArenaCertificate
12import Std.Natural
13
14-- The device side of the checked linear step, declared the way Bob and
15-- Coppelius declare theirs: a device arena with a placement certificate, one
16-- typed code region, one launch whose parameter block names the four data
17-- regions and the learning rate, and one submission.  The NVIDIA backend
18-- (compiler-native-encode of the plan, at build time) derives the program
19-- table, the QMD, the pushbuffer and the GPFIFO from this; nothing here
20-- joins bytes or knows a wire layout.
21--
22-- The arena.  Every address below is arena base + resident offset, and the
23-- residents' placement is decided by Runtime.ArenaCertificate in
24-- Proof.CheckedLinearStep: aligned, inside the arena, pairwise disjoint.
25def linearStepPlanArenaBase : Nat = 0x0000_0008_0000_0000
26def linearStepPlanArenaExtent : Nat = 0x0000_0000_2000_0000
27
28def linearStepPlanPushbufferOffset : Nat = 0x0100_0000
29def linearStepPlanSemaphoreOffset : Nat = 0x0200_0000
30def linearStepPlanProgramOffset : Nat = 0x0300_0000
31def linearStepPlanQMDOffset : Nat = 0x0400_0000
32def linearStepPlanWeightsOffset : Nat = 0x1000_0000
33def linearStepPlanInputOffset : Nat = 0x1001_0000
34def linearStepPlanTargetOffset : Nat = 0x1002_0000
35def linearStepPlanOutputOffset : Nat = 0x1003_0000
36
37def linearStepPlanAddress =
38  (lambda unrestricted offset : Nat . (naturalAdd linearStepPlanArenaBase offset))
39
40def linearStepPlanWord64 =
41  (lambda unrestricted value : Nat . (modelWord64FromNaturalTruncated value))
42
43-- Code: 56 instructions of 16 bytes, padded to the region alignment, then
44-- one 512-byte parameter block for the one launch.
45def linearStepPlanRegionAlignment : Nat = 256
46def linearStepPlanInstructionBytes : Nat = 16
47def linearStepPlanParameterBytes : Nat = 512
48
49def linearStepPlanAlign =
50  (lambda unrestricted value : Nat .
51    (naturalMultiply
52      (naturalDivideUnchecked
53        (naturalAdd value (naturalSaturatingSubtract linearStepPlanRegionAlignment 1))
54        linearStepPlanRegionAlignment)
55      linearStepPlanRegionAlignment))
56
57def linearStepPlanCodeBytes : Nat =
58  (naturalMultiply linearStepSM86InstructionCount linearStepPlanInstructionBytes)
59
60def linearStepPlanCodeExtent : Nat =
61  (linearStepPlanAlign linearStepPlanCodeBytes)
62
63def linearStepPlanProgramBytes : Nat =
64  (naturalAdd linearStepPlanCodeExtent linearStepPlanParameterBytes)
65
66def linearStepPlanQMDTableBytes : Nat = 256
67def linearStepPlanSemaphoreStride : Nat = 32
68def linearStepPlanPushbufferBudget : Nat = 0x0100_0000
69
70-- Byte extents of the data residents for a shape: W m x k f32, x k f32,
71-- t m f32, the output record; the checked 2 x 2 has 16, 8, 8, 44.
72def linearStepPlanWeightsBytesFor = (lambda unrestricted m : Nat . (lambda unrestricted k : Nat . (naturalMultiply 4 (naturalMultiply m k))))
73def linearStepPlanInputBytesFor = (lambda unrestricted k : Nat . (naturalMultiply 4 k))
74def linearStepPlanTargetBytesFor = (lambda unrestricted m : Nat . (naturalMultiply 4 m))
75def linearStepPlanWeightsBytes : Nat = (linearStepPlanWeightsBytesFor linearStepSM86Outputs linearStepSM86Inputs)
76def linearStepPlanInputBytes : Nat = (linearStepPlanInputBytesFor linearStepSM86Inputs)
77def linearStepPlanTargetBytes : Nat = (linearStepPlanTargetBytesFor linearStepSM86Outputs)
78
79-- the residents and certificate for any shape: the same offsets, the
80-- shape's extents, the shape's program length
81def linearStepPlanProgramBytesFor =
82  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
83    (naturalAdd (linearStepPlanAlign (naturalMultiply (linearStepSM86InstructionCountFor m k) linearStepPlanInstructionBytes)) linearStepPlanParameterBytes)))
84
85def linearStepPlanResident =
86  (lambda unrestricted identity : Bytes .
87    (lambda unrestricted offset : Nat .
88      (lambda unrestricted extent : Nat .
89        (lambda unrestricted alignment : Nat .
90          (lambda unrestricted rest : (family ArenaResidents) .
91            (constructor ArenaResidents ArenaResidentsNext
92              (constructor ArenaResident ArenaResidentValue identity offset extent alignment)
93              rest))))))
94
95def linearStepPlanResidents : (family ArenaResidents) =
96  (linearStepPlanResident b"pushbuffer" linearStepPlanPushbufferOffset linearStepPlanPushbufferBudget 4096
97  (linearStepPlanResident b"semaphores" linearStepPlanSemaphoreOffset linearStepPlanSemaphoreStride 4096
98  (linearStepPlanResident b"program" linearStepPlanProgramOffset linearStepPlanProgramBytes linearStepPlanRegionAlignment
99  (linearStepPlanResident b"qmd" linearStepPlanQMDOffset linearStepPlanQMDTableBytes linearStepPlanRegionAlignment
100  (linearStepPlanResident b"weights" linearStepPlanWeightsOffset linearStepPlanWeightsBytes 4
101  (linearStepPlanResident b"input" linearStepPlanInputOffset linearStepPlanInputBytes 4
102  (linearStepPlanResident b"target" linearStepPlanTargetOffset linearStepPlanTargetBytes 4
103  (linearStepPlanResident b"output" linearStepPlanOutputOffset linearStepSM86OutputRecordBytes 4
104    (constructor ArenaResidents ArenaResidentsEnd)))))))))
105
106def linearStepPlanCertificate : Nat =
107  (arenaCertificate linearStepPlanArenaExtent linearStepPlanResidents)
108
109-- One typed region: the realization's machine bytes with its register and
110-- block requirements.
111def linearStepPlanRegions : (family NvidiaDeviceRegions) =
112  (constructor NvidiaDeviceRegions NvidiaDeviceRegionsNext
113    (constructor NvidiaDeviceRegion NvidiaDeviceRegionValue
114      linearStepSM86Identity
115      linearStepSM86MachineBytes
116      linearStepSM86Registers
117      linearStepSM86BlockWidth
118      zero)
119    (constructor NvidiaDeviceRegions NvidiaDeviceRegionsEnd))
120
121-- One launch: one block of one warp, the four addresses and eta at the
122-- realization's parameter offsets (64-bit words 44..47 and 50).
123def linearStepPlanPatch =
124  (lambda unrestricted byteOffset : Nat .
125    (lambda unrestricted value : Nat .
126      (lambda unrestricted rest : (family NvidiaParameterPatches) .
127        (constructor NvidiaParameterPatches NvidiaParameterPatchesNext
128          (constructor NvidiaParameterPatch NvidiaParameterPatchValue
129            (naturalDivideUnchecked byteOffset 8)
130            (linearStepPlanWord64 value))
131          rest))))
132
133-- Shape-independent parameter placement: data addresses come from the
134-- admitted arena, not from the historical checked-step fixture's offsets.
135def linearStepPlanParametersAt =
136  (lambda unrestricted weights : Nat .
137  (lambda unrestricted input : Nat .
138  (lambda unrestricted target : Nat .
139  (lambda unrestricted output : Nat .
140  (lambda unrestricted learningRate : Nat .
141    (constructor NvidiaParameterBlock NvidiaParameterBlockValue
142      (linearStepPlanPatch linearStepSM86WeightsOffset weights
143      (linearStepPlanPatch linearStepSM86InputOffset input
144      (linearStepPlanPatch linearStepSM86TargetOffset target
145      (linearStepPlanPatch linearStepSM86OutputOffset output
146      (linearStepPlanPatch linearStepSM86LearningRateOffset learningRate
147        (constructor NvidiaParameterPatches NvidiaParameterPatchesEnd))))))))))))
148
149def linearStepPlanParameters =
150  (linearStepPlanParametersAt
151    (linearStepPlanAddress linearStepPlanWeightsOffset)
152    (linearStepPlanAddress linearStepPlanInputOffset)
153    (linearStepPlanAddress linearStepPlanTargetOffset)
154    (linearStepPlanAddress linearStepPlanOutputOffset))
155
156def linearStepPlanKernel : (family NvidiaLaunchKernel) =
157  (constructor NvidiaLaunchKernel NvidiaLaunchKernelValue
158    (linearStepPlanWord64 (linearStepPlanAddress linearStepPlanProgramOffset))
159    1 1 1 1 1 linearStepPlanParameterBytes)
160
161def linearStepPlanLaunches =
162  (lambda unrestricted learningRate : Nat .
163    (constructor NvidiaLaunchSchedule NvidiaLaunchScheduleOne
164      (constructor NvidiaLaunchTemplate NvidiaLaunchTemplateValue
165        linearStepSM86Identity
166        linearStepPlanKernel
167        (linearStepPlanParameters learningRate))))
168
169def linearStepPlanLaunchCount : Nat = 1
170def linearStepPlanSubmissionCount : Nat = 1
171
172def linearStepPlanSubmissions : (family NvidiaSubmissionSchedule) =
173  (constructor NvidiaSubmissionSchedule NvidiaSubmissionScheduleOne
174    (constructor NvidiaSubmissionBatch NvidiaSubmissionBatchValue
175      b"linear-step-submission"
176      zero
177      (constructor NvidiaLaunchReferences NvidiaLaunchReferencesRange zero 1)))
178
179-- The same launch submitted `count` times, one semaphore slot per
180-- submission: the step is idempotent, so every submission rewrites the same
181-- record, and the host waits on the last slot.  The semaphore resident grows
182-- with the count; the certificate is decided for the count in use.
183def linearStepPlanSubmissionsRepeated =
184  (lambda unrestricted count : Nat .
185    (constructor NvidiaSubmissionSchedule NvidiaSubmissionScheduleRepeat
186      count zero linearStepPlanSemaphoreStride linearStepPlanSubmissions))
187
188def linearStepPlanResidentsFor =
189  (lambda unrestricted submissions : Nat .
190    (linearStepPlanResident b"pushbuffer" linearStepPlanPushbufferOffset linearStepPlanPushbufferBudget 4096
191    (linearStepPlanResident b"semaphores" linearStepPlanSemaphoreOffset (naturalMultiply linearStepPlanSemaphoreStride submissions) 4096
192    (linearStepPlanResident b"program" linearStepPlanProgramOffset linearStepPlanProgramBytes linearStepPlanRegionAlignment
193    (linearStepPlanResident b"qmd" linearStepPlanQMDOffset linearStepPlanQMDTableBytes linearStepPlanRegionAlignment
194    (linearStepPlanResident b"weights" linearStepPlanWeightsOffset linearStepPlanWeightsBytes 4
195    (linearStepPlanResident b"input" linearStepPlanInputOffset linearStepPlanInputBytes 4
196    (linearStepPlanResident b"target" linearStepPlanTargetOffset linearStepPlanTargetBytes 4
197    (linearStepPlanResident b"output" linearStepPlanOutputOffset linearStepSM86OutputRecordBytes 4
198      (constructor ArenaResidents ArenaResidentsEnd))))))))))
199
200def linearStepPlanCertificateFor =
201  (lambda unrestricted submissions : Nat .
202    (arenaCertificate linearStepPlanArenaExtent (linearStepPlanResidentsFor submissions)))
203
204def linearStepPlanResidentsForShape =
205  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
206    (lambda unrestricted submissions : Nat .
207      (linearStepPlanResident b"pushbuffer" linearStepPlanPushbufferOffset linearStepPlanPushbufferBudget 4096
208      (linearStepPlanResident b"semaphores" linearStepPlanSemaphoreOffset (naturalMultiply linearStepPlanSemaphoreStride submissions) 4096
209      (linearStepPlanResident b"program" linearStepPlanProgramOffset (linearStepPlanProgramBytesFor m k) linearStepPlanRegionAlignment
210      (linearStepPlanResident b"qmd" linearStepPlanQMDOffset linearStepPlanQMDTableBytes linearStepPlanRegionAlignment
211      (linearStepPlanResident b"weights" linearStepPlanWeightsOffset (linearStepPlanWeightsBytesFor m k) 4
212      (linearStepPlanResident b"input" linearStepPlanInputOffset (linearStepPlanInputBytesFor k) 4
213      (linearStepPlanResident b"target" linearStepPlanTargetOffset (linearStepPlanTargetBytesFor m) 4
214      (linearStepPlanResident b"output" linearStepPlanOutputOffset (linearStepSM86OutputRecordBytesFor m k) 4
215        (constructor ArenaResidents ArenaResidentsEnd))))))))))))
216
217def linearStepPlanCertificateForShape =
218  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
219    (lambda unrestricted submissions : Nat .
220      (arenaCertificate linearStepPlanArenaExtent (linearStepPlanResidentsForShape m k submissions)))))
221
222-- a program's machine bytes by the compiler's encoder (none when it refuses)
223def linearStepPlanMachineBytesOf =
224  (lambda unrestricted program : (family SM86Program) .
225    (eliminate SM86ProgramEncodingResult
226      (lambda unrestricted current : (family SM86ProgramEncodingResult) . Bytes)
227      (sm86EncodeProgram program)
228      (branch SM86ProgramEncodingSucceeded bytes telemetry . bytes)
229      (branch SM86ProgramEncodingFailed index failure telemetry . b"")))
230
231-- a program's instructions, and the program region it needs: its code
232-- aligned, then the parameter block -- for the generator's program, the
233-- shape's (Proof.LinearStepGeneratorForall: its length is the count)
234def linearStepPlanInstructionsOf =
235  (lambda unrestricted program : (family SM86Program) .
236    (eliminate SM86Program (lambda unrestricted current : (family SM86Program) . Nat) program
237      (branch SM86ProgramEnd . zero)
238      (branch SM86ProgramNext head tail induction . (succ induction))))
239
240def linearStepPlanProgramBytesOf =
241  (lambda unrestricted program : (family SM86Program) .
242    (naturalAdd (linearStepPlanAlign (naturalMultiply (linearStepPlanInstructionsOf program) linearStepPlanInstructionBytes)) linearStepPlanParameterBytes))
243
244-- the residents and certificate for a shape's data and a program's region
245def linearStepPlanResidentsForProgram =
246  (lambda unrestricted program : (family SM86Program) .
247  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
248    (lambda unrestricted submissions : Nat .
249      (linearStepPlanResident b"pushbuffer" linearStepPlanPushbufferOffset linearStepPlanPushbufferBudget 4096
250      (linearStepPlanResident b"semaphores" linearStepPlanSemaphoreOffset (naturalMultiply linearStepPlanSemaphoreStride submissions) 4096
251      (linearStepPlanResident b"program" linearStepPlanProgramOffset (linearStepPlanProgramBytesOf program) linearStepPlanRegionAlignment
252      (linearStepPlanResident b"qmd" linearStepPlanQMDOffset linearStepPlanQMDTableBytes linearStepPlanRegionAlignment
253      (linearStepPlanResident b"weights" linearStepPlanWeightsOffset (linearStepPlanWeightsBytesFor m k) 4
254      (linearStepPlanResident b"input" linearStepPlanInputOffset (linearStepPlanInputBytesFor k) 4
255      (linearStepPlanResident b"target" linearStepPlanTargetOffset (linearStepPlanTargetBytesFor m) 4
256      (linearStepPlanResident b"output" linearStepPlanOutputOffset (linearStepSM86OutputRecordBytesFor m k) 4
257        (constructor ArenaResidents ArenaResidentsEnd)))))))))))))
258
259def linearStepPlanCertificateForProgram =
260  (lambda unrestricted program : (family SM86Program) .
261  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
262    (lambda unrestricted submissions : Nat .
263      (arenaCertificate linearStepPlanArenaExtent (linearStepPlanResidentsForProgram program m k submissions))))))
264
265-- the plan for a shape and a program for it -- the generator's, one a
266-- proposer rewrote, or another realization's (the looped one): the program
267-- region is the program's; the same offsets and launch shape. Keep this typed
268-- until whole-program encoding: arena placement must be checkable without
269-- evaluating a build-only machine-image intrinsic.
270def linearStepPlanRegionsForProgram =
271  (lambda unrestricted program : (family SM86Program) .
272    (constructor NvidiaDeviceRegions NvidiaDeviceRegionsNext
273      (constructor NvidiaDeviceRegion NvidiaDeviceProgramRegion
274        linearStepSM86Identity
275        program
276        (sm86RegisterDemand program)
277        linearStepSM86BlockWidth
278        zero
279        (constructor NvidiaDeviceRealization NvidiaDeviceRealizationRefused
280          b"linear step requires an established SM86 realization"))
281      (constructor NvidiaDeviceRegions NvidiaDeviceRegionsEnd)))
282
283def linearStepPlanForProgram =
284  (lambda unrestricted program : (family SM86Program) .
285  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
286  (lambda unrestricted component : (family NvidiaWholeProgramComponent) .
287    (lambda unrestricted architecture : (family ModelWord32) .
288      (lambda unrestricted computeClass : (family ModelWord32) .
289        (lambda unrestricted spaVersion : (family ModelWord32) .
290          (lambda unrestricted learningRate : Nat .
291          (lambda unrestricted submissions : Nat .
292            (constructor NvidiaWholeProgramPlan NvidiaWholeProgramPlanValue
293              component
294              architecture
295              computeClass
296              spaVersion
297              (linearStepPlanWord64 (linearStepPlanAddress linearStepPlanProgramOffset))
298              (linearStepPlanWord64 (linearStepPlanAddress linearStepPlanQMDOffset))
299              (linearStepPlanWord64 (linearStepPlanAddress linearStepPlanPushbufferOffset))
300              (linearStepPlanWord64 (linearStepPlanAddress linearStepPlanSemaphoreOffset))
301              linearStepPlanRegionAlignment
302              (linearStepPlanProgramBytesOf program)
303              (linearStepPlanRegionsForProgram program)
304              (linearStepPlanLaunches learningRate)
305              linearStepPlanLaunchCount
306              linearStepPlanQMDTableBytes
307              (linearStepPlanSubmissionsRepeated submissions)
308              submissions
309              submissions))))))))))
310
311def linearStepPlanComponentForProgram =
312  (lambda unrestricted program : (family SM86Program) .
313  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
314  (lambda unrestricted component : (family NvidiaWholeProgramComponent) .
315    (lambda unrestricted architecture : (family ModelWord32) .
316      (lambda unrestricted computeClass : (family ModelWord32) .
317        (lambda unrestricted spaVersion : (family ModelWord32) .
318          (lambda unrestricted learningRate : Nat .
319          (lambda unrestricted submissions : Nat .
320            (compiler-native-encode
321              (family NvidiaWholeProgramPlan)
322              (linearStepPlanForProgram program m k component architecture computeClass spaVersion learningRate submissions)))))))))))
323
324-- the plan for any shape: the generator's program for it
325def linearStepPlanRegionsFor =
326  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
327    (linearStepPlanRegionsForProgram (linearStepSM86ProgramFor m k))))
328def linearStepPlanFor =
329  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
330    (linearStepPlanForProgram (linearStepSM86ProgramFor m k) m k)))
331def linearStepPlanComponentFor =
332  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
333    (linearStepPlanComponentForProgram (linearStepSM86ProgramFor m k) m k)))
334
335-- the checked 2 x 2: the generator's plan for that shape
336def linearStepPlan = (linearStepPlanFor linearStepSM86Outputs linearStepSM86Inputs)
337
338-- The backend boundary: the compiler realizes one component of the plan.
339def linearStepPlanComponent =
340  (lambda unrestricted component : (family NvidiaWholeProgramComponent) .
341    (lambda unrestricted architecture : (family ModelWord32) .
342      (lambda unrestricted computeClass : (family ModelWord32) .
343        (lambda unrestricted spaVersion : (family ModelWord32) .
344          (lambda unrestricted learningRate : Nat .
345          (lambda unrestricted submissions : Nat .
346            (compiler-native-encode
347              (family NvidiaWholeProgramPlan)
348              (linearStepPlan component architecture computeClass spaVersion learningRate submissions))))))))