1module PartitionedRegression.Build.CoordinateDevice
2
3import Std.Numeric.Profile
4import Representation.ParticipantGroup
5import Learning.Numerics
6import PartitionedRegression.Coordinates
7import PartitionedRegression.CoordinateHost
8import Runtime.OrderedCombine
9import Data.SHA256Digest
10import Data.SHA256Schedule
11import Data.SHA256ScheduleExpand
12import Data.SHA256Compress
13import Data.SHA256Core
14import Compiler.Planning.OrderedCombine
15import PartitionedRegression.Build.ScalarEvaluation
16import PartitionedRegression.Build.CoordinateMemory
17import Platform.Linux.Nvidia.ProgramReadback
18import Checkpoint.NativeRegions
19import Compiler.Planning.RegionFrame
20import Hardware.Nvidia.SM86.Command.WholeProgramPlan
21import Platform.Linux.Nvidia.DeviceBinding
22import Platform.Linux.Nvidia.ResourcePlan
23import Platform.Linux.Nvidia.LinearStep
24import Platform.Linux.Nvidia.ProgramPlan
25import Platform.Linux.Nvidia.PlanHost
26import Platform.Linux.Nvidia.PlanHostRequest
27import Platform.RunPod.Nvidia.SM86
28import Runtime.ArenaCertificate
29import Runtime.NativeFloatPlan
30import Runtime.NativeMemoryPlan
31import Runtime.NativePeerEndpoint
32import Runtime.NativePhysicalEmbeddedArtifact
33import Runtime.NativePhysicalNativeELF
34import Runtime.NativePhysicalProgram
35import Runtime.NativeLaunchRecipeLoad
36import Runtime.RegionFrame
37import Runtime.SocketAddress
38import Model.Word32
39import Model.Word64
40import Std.Foundation
41import Std.Natural
42
43-- This policy is part of the frame/checkpoint identity. Each local product
44-- rounds its operands to binary32 on the explicit GPU realization and widens
45-- the device result. Residual summation and persistent updates use binary64.
46def cdNumerics : Bytes =
47 b"f32-products-f64-accumulation-v1"
48
49def cdIdentity =
50 (lambda unrestricted kind : (family NativePhysicalEmbeddedComponentKind) .
51 (let unrestricted name =
52 (eliminate
53 NativePhysicalEmbeddedComponentKind
54 (lambda unrestricted current : (family NativePhysicalEmbeddedComponentKind) . Bytes)
55 kind
56 (branch NativePhysicalEmbeddedHostELF . b"coordinate-scalar-sm86-host-v1")
57 (branch NativePhysicalEmbeddedProgramTable . b"coordinate-scalar-sm86-program-v1")
58 (branch NativePhysicalEmbeddedQMDTable . b"coordinate-scalar-sm86-qmd-v1")
59 (branch NativePhysicalEmbeddedPushbuffer . b"coordinate-scalar-sm86-push-v1")
60 (branch NativePhysicalEmbeddedGPFIFO . b"coordinate-scalar-sm86-ring-v1"))
61 in
62 (bytes-append
63 name
64 (phZeros
65 (naturalSaturatingSubtract nativePhysicalEmbeddedIdentityLength (bytes-length name))))))
66
67def cdBindingAdmitted =
68 (lambda unrestricted device : (family NvidiaDeviceBinding) .
69 (naturalAnd
70 (nvidiaDeviceBindingAdmitted device)
71 (naturalEqual (bytes-length (nvidiaDeviceBindingUUID device)) 16)))
72
73def cdCounter =
74 (lambda unrestricted end : Nat . (frameAlign end))
75
76def cdScratch =
77 (lambda unrestricted end : Nat . (naturalAdd (cdCounter end) coordinateWordBytes))
78
79def cdExtent =
80 (lambda unrestricted end : Nat . (naturalAdd (cdScratch end) coordinateWordBytes))
81
82-- State addresses are relative to the native interpreter arena, separate from
83-- the target's independently certified GPU/host virtual mappings. The native
84-- session prefix and the needed counter/scratch words cannot alias learning state.
85def cdSessionResidents =
86 (lambda unrestricted end : Nat .
87 (lambda unrestricted extent : Nat .
88 (lambda unrestricted residents : (family ArenaResidents) .
89 (constructor
90 ArenaResidents
91 ArenaResidentsNext
92 (constructor
93 ArenaResident
94 ArenaResidentValue
95 b"native-device-session"
96 0
97 phStateExtent
98 coordinateWordBytes)
99 (constructor
100 ArenaResidents
101 ArenaResidentsNext
102 (constructor
103 ArenaResident
104 ArenaResidentValue
105 b"device-local-state"
106 (cdCounter end)
107 (naturalSaturatingSubtract extent (cdCounter end))
108 coordinateWordBytes)
109 residents)))))
110
111-- Training needs only an issue counter; prediction also uses its explicit
112-- conversion scratch. A pairing retains only the host state it actually uses.
113def cdTrainingExtent =
114 (lambda unrestricted end : Nat . (naturalAdd (cdCounter end) coordinateWordBytes))
115
116def cdTrainingResidents =
117 (lambda unrestricted end : Nat . (cdSessionResidents end (cdTrainingExtent end)))
118
119def cdResidents =
120 (lambda unrestricted end : Nat . (cdSessionResidents end (cdExtent end)))
121
122def cdPredictionResidents =
123 (lambda unrestricted layout : (family CoordinatePredictionLayout) .
124 (constructor
125 ArenaResidents
126 ArenaResidentsNext
127 (constructor
128 ArenaResident
129 ArenaResidentValue
130 b"prediction-values"
131 (cpLeft layout)
132 (naturalSaturatingSubtract (cpExtent layout) (cpLeft layout))
133 coordinateWordBytes)
134 (nativeRegionResidents (cpCheckpoint layout))))
135
136-- A scalar prediction field is a product, suitable for both the forward
137-- partial and the backward residual-times-feature. Choosing the gradient
138-- field here would multiply by the input twice and change the learner.
139def cdProductResult : (family NvidiaLinearScalarResult) =
140 (constructor NvidiaLinearScalarResult NvidiaLinearScalarPrediction)
141
142def cdProduct =
143 (lambda unrestricted gpu : (family NvidiaPlanHostLayout) .
144 (lambda unrestricted entry : Nat .
145 (lambda unrestricted data : Nat .
146 (lambda unrestricted inputOffset : Nat .
147 (lambda unrestricted targetOffset : Nat .
148 (lambda unrestricted outputOffset : Nat .
149 (lambda unrestricted end : Nat .
150 (lambda unrestricted weight : Nat .
151 (lambda unrestricted input : (family NvidiaLinearScalarInput) .
152 (lambda unrestricted destination : Nat .
153 (nlsEvaluateScalar
154 b"coordinate-device-arithmetic"
155 gpu
156 entry
157 data
158 inputOffset
159 targetOffset
160 outputOffset
161 (cdCounter end)
162 (cdScratch end)
163 (constructor
164 NvidiaLinearScalarInput
165 NvidiaLinearScalarBinary64
166 (nativeMemoryLoad weight))
167 input
168 (constructor NvidiaLinearScalarInput NvidiaLinearScalarNatural 0)
169 cdProductResult
170 destination)))))))))))
171
172def cdTrainingAdmitted =
173 (lambda unrestricted worker : Nat .
174 (lambda unrestricted resume : Nat .
175 (lambda unrestricted epoch : Nat .
176 (lambda unrestricted steps : Nat .
177 (naturalAnd
178 (naturalLessOrEqual worker 1)
179 (naturalAnd
180 (naturalLessOrEqual resume 1)
181 (naturalAnd
182 (naturalNonzero epoch)
183 (naturalAnd
184 (naturalLessOrEqual (naturalAdd epoch resume) regionFrameWordLimit)
185 (naturalAnd
186 (naturalNonzero steps)
187 (naturalLessOrEqual (naturalMultiply cdmIssueCount steps) regionFrameWordLimit))))))))))
188
189-- Only this initial transfer writes parameter/accumulator from host state.
190-- A restored checkpoint follows the same path before the first device issue.
191def cdSeedMemory =
192 (lambda unrestricted layout : (family CoordinateHostLayout) .
193 (lambda unrestricted data : Nat .
194 (lambda unrestricted tail : (family NativePhysicalCommands) .
195 (phCopyStateToMapped
196 (chWeight layout)
197 (naturalAdd data cdmParameter)
198 coordinateWordBytes
199 (phCopyStateToMapped
200 (chAccumulator layout)
201 (naturalAdd data cdmAccumulator)
202 coordinateWordBytes
203 (phCopyStateToMapped
204 (chFeature layout)
205 (naturalAdd data cdmFeature)
206 coordinateWordBytes
207 (phCopyStateToMapped
208 (chDecay layout)
209 (naturalAdd data cdmDecay)
210 coordinateWordBytes
211 (phCopyStateToMapped
212 (chRate layout)
213 (naturalAdd data cdmRate)
214 coordinateWordBytes
215 tail))))))))
216
217-- Host values are validation/checkpoint shadows, never input to a later
218-- parameter update. Capture follows the device fence and precedes publication.
219def cdCaptureMemory =
220 (lambda unrestricted layout : (family CoordinateHostLayout) .
221 (lambda unrestricted data : Nat .
222 (lambda unrestricted tail : (family NativePhysicalCommands) .
223 (phCopyMappedToState
224 (naturalAdd data cdmParameter)
225 (chWeight layout)
226 coordinateWordBytes
227 (phCopyMappedToState
228 (naturalAdd data cdmAccumulator)
229 (chAccumulator layout)
230 coordinateWordBytes
231 tail)))))
232
233def cdMemoryPhase =
234 (lambda unrestricted layout : (family CoordinateHostLayout) .
235 (lambda unrestricted gpu : (family NvidiaPlanHostLayout) .
236 (lambda unrestricted component : (pi unrestricted kind : (family NvidiaWholeProgramComponent) . Bytes) .
237 (lambda unrestricted data : Nat .
238 (lambda unrestricted phase : (family CoordinateDevicePhase) .
239 (lambda unrestricted tail : (family NativePhysicalCommands) .
240 (let unrestricted issue =
241 (nvidiaProgramIssue gpu component (cdCounter (chExtent layout)) (cdmOrdinal phase))
242 in
243 (eliminate
244 CoordinateDevicePhase
245 (lambda unrestricted current : (family CoordinateDevicePhase) .
246 (family NativePhysicalCommands))
247 phase
248 (branch
249 CoordinateDeviceForward
250 .
251 (issue
252 (phCopyMappedToState
253 (naturalAdd data cdmPartial)
254 (chPartial layout)
255 coordinateWordBytes
256 tail)))
257 (branch
258 CoordinateDeviceBackward
259 .
260 (phCopyStateToMapped
261 (chError layout)
262 (naturalAdd data cdmResidual)
263 coordinateWordBytes
264 (issue
265 (phCopyMappedToState
266 (naturalAdd data cdmGradient)
267 (chGradient layout)
268 coordinateWordBytes
269 tail))))
270 (branch CoordinateDeviceUpdate . (issue (cdCaptureMemory layout data tail)))))))))))
271
272-- argv is the derived manifest/shard layout followed by an exclusive device
273-- readback path. The session uploads state once, then the three submissions
274-- share video memory. Final device capture/readback and release precede the
275-- global checkpoint publication protocol; no success bypasses either lifecycle.
276def coordinateDeviceArtifact =
277 (lambda unrestricted resources : (family NvidiaResourceContext) .
278 (lambda unrestricted profile : (family RunPodSM86TargetProfile) .
279 (lambda unrestricted device : (family NvidiaDeviceBinding) .
280 (lambda unrestricted worker : Nat .
281 (lambda unrestricted address : (family IPv4Address) .
282 (lambda unrestricted port : Nat .
283 (lambda unrestricted resume : Nat .
284 (lambda unrestricted epoch : Nat .
285 (lambda unrestricted steps : Nat .
286 (eliminate
287 NvidiaProgramPlan
288 (lambda unrestricted plan : (family NvidiaProgramPlan) . Bytes)
289 (cdmPlan profile)
290 (branch
291 NvidiaProgramPlanValue
292 gpu
293 component
294 recipes
295 data
296 dataGPU
297 dataBytes
298 admitted
299 .
300 (nat-eliminate
301 (lambda unrestricted valid : Nat . Bytes)
302 b""
303 (lambda unrestricted previous : Nat .
304 (lambda unrestricted ignored : Bytes .
305 (chWithLayout
306 cdNumerics
307 epoch
308 phStateExtent
309 (lambda unrestricted layout : (family CoordinateHostLayout) .
310 (eliminate
311 StdResult
312 (lambda unrestricted result : (family StdResult (family SocketAddressFailure) Bytes) .
313 Bytes)
314 (linuxIPv4SocketAddress address port)
315 (branch StdFailure error . b"")
316 (branch
317 StdSuccess
318 socketAddress
319 .
320 (nat-eliminate
321 (lambda unrestricted valid : Nat . Bytes)
322 b""
323 (lambda unrestricted previous : Nat .
324 (lambda unrestricted ignored : Bytes .
325 (nlsWithProgramRecipes
326 component
327 (lambda unrestricted program : Bytes .
328 (lambda unrestricted qmd : Bytes .
329 (lambda unrestricted push : Bytes .
330 (lambda unrestricted ring : Bytes .
331 (nlsWithPrivateSlots
332 nativePeerEndpointSlots
333 (nlsSessionPrelude
334 cdIdentity
335 b"coordinate-peer"
336 device
337 gpu
338 recipes
339 program
340 qmd
341 push
342 ring
343 (cdSeedMemory layout data chEnd))
344 (lambda unrestricted prepare : (family NativePhysicalCommands) .
345 (nlsWithPrivateSlots
346 nativePeerEndpointSlots
347 (cdMemoryPhase
348 layout
349 gpu
350 component
351 data
352 (constructor CoordinateDevicePhase CoordinateDeviceForward)
353 chEnd)
354 (lambda unrestricted forward : (family NativePhysicalCommands) .
355 (nlsWithPrivateSlots
356 nativePeerEndpointSlots
357 (cdMemoryPhase
358 layout
359 gpu
360 component
361 data
362 (constructor CoordinateDevicePhase CoordinateDeviceBackward)
363 chEnd)
364 (lambda unrestricted backward : (family NativePhysicalCommands) .
365 (nlsWithPrivateSlots
366 nativePeerEndpointSlots
367 (cdMemoryPhase
368 layout
369 gpu
370 component
371 data
372 (constructor CoordinateDevicePhase CoordinateDeviceUpdate)
373 chEnd)
374 (lambda unrestricted update : (family NativePhysicalCommands) .
375 (nlsWithPrivateSlots
376 nativePeerEndpointSlots
377 (cdCaptureMemory
378 layout
379 data
380 (nvidiaProgramReadback
381 (chArgument (succ chDirectoryIndex))
382 (prSemaphoreHost gpu)
383 data
384 dataBytes
385 cdmIssueCount
386 (cdCounter (chExtent layout))
387 (phTelemetry
388 3
389 b"coordinate-peer:device-updates-retired"
390 (phRelease chEnd))))
391 (lambda unrestricted finish : (family NativePhysicalCommands) .
392 (nlsEmbedHost
393 cdIdentity
394 (nlsGenerateResourceELF
395 resources
396 gpu
397 (chProgramFor
398 (naturalAdd nativePeerEndpointSlots phSlots)
399 (cdTrainingExtent (chExtent layout))
400 (chCommandsWithSession
401 layout
402 worker
403 socketAddress
404 resume
405 steps
406 (prSplice forward)
407 (prSplice backward)
408 (prSplice update)
409 (lambda unrestricted tail : (family NativePhysicalCommands) .
410 (chStore
411 (cdCounter (chExtent layout))
412 (phImm 0)
413 (prSplice prepare tail)))
414 (prSplice finish))))
415 program
416 qmd
417 push
418 ring))))))))))))))))))
419 (arenaCertificate
420 (cdTrainingExtent (chExtent layout))
421 (cdTrainingResidents (chExtent layout) (chResidents layout))))))))))
422 (naturalAnd
423 admitted
424 (naturalAnd
425 (cdBindingAdmitted device)
426 (cdTrainingAdmitted worker resume epoch steps)))))))))))))))
427
428-- Prediction uses the same numerical policy on a fresh bound device session.
429-- It restores both saved parameters, executes both products on that GPU, then
430-- combines the widened results. The extra argv follows the five input paths.
431def coordinateDevicePredictionArtifact =
432 (lambda unrestricted resources : (family NvidiaResourceContext) .
433 (lambda unrestricted profile : (family RunPodSM86TargetProfile) .
434 (lambda unrestricted device : (family NvidiaDeviceBinding) .
435 (lambda unrestricted epoch : Nat .
436 (eliminate
437 NvidiaLinearStepPlan
438 (lambda unrestricted plan : (family NvidiaLinearStepPlan) . Bytes)
439 (regressionScalarPlan profile)
440 (branch
441 NvidiaLinearStepPlanValue
442 gpu
443 component
444 recipes
445 data
446 inputOffset
447 targetOffset
448 outputOffset
449 dataBytes
450 admitted
451 .
452 (nat-eliminate
453 (lambda unrestricted valid : Nat . Bytes)
454 b""
455 (lambda unrestricted previous : Nat .
456 (lambda unrestricted ignored : Bytes .
457 (cpWithLayout
458 cdNumerics
459 epoch
460 phStateExtent
461 (lambda unrestricted layout : (family CoordinatePredictionLayout) .
462 (nat-eliminate
463 (lambda unrestricted valid : Nat . Bytes)
464 b""
465 (lambda unrestricted previous : Nat .
466 (lambda unrestricted ignored : Bytes .
467 (nlsWithProgramRecipes
468 component
469 (lambda unrestricted program : Bytes .
470 (lambda unrestricted qmd : Bytes .
471 (lambda unrestricted push : Bytes .
472 (lambda unrestricted ring : Bytes .
473 (let unrestricted entry =
474 (nvidiaPlanHostTableWord
475 (component
476 (constructor
477 NvidiaWholeProgramComponent
478 NvidiaWholeProgramGPFIFO))
479 0)
480 in
481 (nlsWithPrivateSlots
482 cpSlots
483 (nlsSessionPrelude
484 cdIdentity
485 b"coordinate-prediction"
486 device
487 gpu
488 recipes
489 program
490 qmd
491 push
492 ring
493 chEnd)
494 (lambda unrestricted prepare : (family NativePhysicalCommands) .
495 (nlsWithPrivateSlots
496 cpSlots
497 (cdProduct
498 gpu
499 entry
500 data
501 inputOffset
502 targetOffset
503 outputOffset
504 (cpExtent layout)
505 (cpLeft layout)
506 (constructor
507 NvidiaLinearScalarInput
508 NvidiaLinearScalarNatural
509 1)
510 (cpPrediction layout)
511 (cdProduct
512 gpu
513 entry
514 data
515 inputOffset
516 targetOffset
517 outputOffset
518 (cpExtent layout)
519 (cpRight layout)
520 (constructor
521 NvidiaLinearScalarInput
522 NvidiaLinearScalarNatural
523 2)
524 (cpFeature layout)
525 chEnd))
526 (lambda unrestricted products : (family NativePhysicalCommands) .
527 (nlsWithPrivateSlots
528 cpSlots
529 (nlsScalarReadback
530 (chArgument chOutputShardBase)
531 gpu
532 data
533 outputOffset
534 (cdCounter (cpExtent layout))
535 (phTelemetry
536 3
537 b"coordinate-prediction:device-products-retired"
538 (phRelease chEnd)))
539 (lambda unrestricted finish : (family NativePhysicalCommands) .
540 (cpWithRestore
541 layout
542 (chStore
543 (cdCounter (cpExtent layout))
544 (phImm 0)
545 (prSplice
546 prepare
547 (prSplice
548 products
549 (chBinary
550 chAdd
551 (cpPrediction layout)
552 (cpPrediction layout)
553 (cpFeature layout)
554 (prSplice finish (cpOutput layout chHalt))))))
555 (lambda unrestricted commands : (family NativePhysicalCommands) .
556 (nlsEmbedHost
557 cdIdentity
558 (nlsGenerateResourceELF
559 resources
560 gpu
561 (chProgramFor
562 (naturalAdd cpSlots phSlots)
563 (cdExtent (cpExtent layout))
564 commands))
565 program
566 qmd
567 push
568 ring)))))))))))))))))
569 (arenaCertificate
570 (cdExtent (cpExtent layout))
571 (cdResidents (cpExtent layout) (cdPredictionResidents layout))))))))
572 (naturalAnd admitted (cdBindingAdmitted device)))))))))The compiler supplied declaration spans and resolved links from this source snapshot. This page does not assert that this file belongs to a checked closure.