Source/Packages

Platform.Linux.Nvidia.LinearStep

packages/hardware/platforms/linux-nvidia/src/Platform/Linux/Nvidia/LinearStep.alpha

309 lines26 declarations14.6 KiBSHA-256 6ba9f89817ca

Complete file · line 69

LinearStep.alpha

Definition view
1module Platform.Linux.Nvidia.LinearStep
2
3import Accelerator.SM86.Instruction
4import Platform.Linux.Nvidia.ProgramPlan
5import Accelerator.SM86.InstructionEncoding
6import Accelerator.SM86.Types
7import Compiler.Planning.VA
8import Hardware.Nvidia.SM86.Command.WholeProgramPlan
9import Model.Config
10import Model.Parameter
11import Model.Word64
12import Platform.Linux.Nvidia.PlanHost
13import Platform.Linux.Nvidia.DeviceBinding
14import Platform.Linux.Nvidia.PlanHostRequest
15import Platform.Linux.Nvidia.ResourcePlan
16import Runtime.NativePhysicalMemory
17import Runtime.NativeResourcePlan
18import Execution.ResourceGraph
19import Std.Foundation
20import Std.List
21import Realization.Nvidia.SM86.LinearStepPlanSM86
22import Realization.Nvidia.SM86.LinearStepSM86
23import Runtime.ArenaCertificate
24import Runtime.NativePhysicalEmbeddedArtifact
25import Runtime.NativePhysicalProgram
26import Runtime.NativeFloatPlan
27import Runtime.NativeFilePlan
28import Data.Bytes
29import Runtime.NativeLaunchRecipeLoad
30import Runtime.NativePhysicalProgramComposer
31import Std.Natural
32import Std.Numeric.Profile
33
34-- A shape-derived, reusable scalar/vector linear-step placement. The same
35-- checked kernel can evaluate gradients from runtime data repeatedly. System
36-- semantics choose samples and combine gradients; this target owns memory,
37-- launch tables, submission retirement and readback. All buffers are explicit:
38-- host-visible system memory for this small realization, video memory for the
39-- separate direct-RM USERD allocation. The system-memory ring and data use
40-- the admitted direct mappings. This is a declared route, not a runtime retry.
41family NvidiaLinearStepPlan : Type 0
42constructor NvidiaLinearStepPlanValue
43field unrestricted nlsLayout : (family NvidiaPlanHostLayout)
44field unrestricted nlsComponents : (pi unrestricted component : (family NvidiaWholeProgramComponent) . Bytes)
45field unrestricted nlsRecipes : (family NvidiaPlanHostRecipeFills)
46field unrestricted nlsDataHost : Nat
47field unrestricted nlsInputOffset : Nat
48field unrestricted nlsTargetOffset : Nat
49field unrestricted nlsOutputOffset : Nat
50field unrestricted nlsDataBytes : Nat
51field unrestricted nlsAdmitted : Nat
52
53end-family
54
55-- Scalar evaluation chooses a semantic result field, never an unchecked byte
56-- offset. These cases apply to the admitted 1-output, 1-input specialization.
57family NvidiaLinearScalarResult : Type 0
58constructor NvidiaLinearScalarPrediction
59constructor NvidiaLinearScalarLoss
60constructor NvidiaLinearScalarGradient
61constructor NvidiaLinearScalarUpdatedWeight
62
63end-family
64
65family NvidiaLinearScalarInput : Type 0
66constructor NvidiaLinearScalarNatural
67field unrestricted nlsNaturalValue : Nat
68constructor NvidiaLinearScalarBinary64
69field unrestricted nlsBinary64Value : (family NativePhysicalOperand)
70
71end-family
72
73def nlsScalarResultOffset =
74  (lambda unrestricted result : (family NvidiaLinearScalarResult) .
75    (eliminate
76      NvidiaLinearScalarResult
77      (lambda unrestricted current : (family NvidiaLinearScalarResult) . Nat)
78      result
79      (branch NvidiaLinearScalarPrediction . 0)
80      (branch NvidiaLinearScalarLoss . (lsLossOffset 1))
81      (branch NvidiaLinearScalarGradient . (lsGradientOffset 1 1 0 0))
82      (branch NvidiaLinearScalarUpdatedWeight . (lsUpdatedOffset 1 1 0 0))))
83
84-- A runtime input is rounded once to binary32. A natural input first uses the
85-- explicit binary64 conversion, preserving the host arithmetic contract. The
86-- temporary is private and must not alias another input's still-live state.
87def nlsUploadScalar =
88  (lambda unrestricted error : Bytes .
89    (lambda unrestricted input : (family NvidiaLinearScalarInput) .
90      (lambda unrestricted scratch : Nat .
91        (lambda unrestricted destination : Nat .
92          (lambda unrestricted tail : (family NativePhysicalCommands) .
93            (eliminate
94              NvidiaLinearScalarInput
95              (lambda unrestricted current : (family NvidiaLinearScalarInput) .
96                (family NativePhysicalCommands))
97              input
98              (branch
99                NvidiaLinearScalarNatural
100                value
101                .
102                (nativeFloat64Natural
103                  error
104                  scratch
105                  value
106                  (nativeNumericConvert
107                    error
108                    (constructor StdNumericConversion StdNumericRoundBinary64)
109                    scratch
110                    (phLoad scratch)
111                    (phCopyStateToMapped scratch destination (linearStepPlanInputBytesFor 1) tail))))
112              (branch
113                NvidiaLinearScalarBinary64
114                value
115                .
116                (nativeNumericConvert
117                  error
118                  (constructor StdNumericConversion StdNumericRoundBinary64)
119                  scratch
120                  value
121                  (phCopyStateToMapped scratch destination (linearStepPlanInputBytesFor 1) tail)))))))))
122
123-- Upload all runtime operands, submit one checked scalar kernel, wait for its
124-- own fence, then widen the selected device result. No host expression computes
125-- that result. The caller supplies the matching scalar plan and disjoint state;
126-- its fresh-channel counter is independent of checkpointed learner versions.
127def nlsEvaluateScalar =
128  (lambda unrestricted error : Bytes .
129    (lambda unrestricted layout : (family NvidiaPlanHostLayout) .
130      (lambda unrestricted entry : Nat .
131        (lambda unrestricted data : Nat .
132          (lambda unrestricted inputOffset : Nat .
133            (lambda unrestricted targetOffset : Nat .
134              (lambda unrestricted outputOffset : Nat .
135                (lambda unrestricted counter : Nat .
136                  (lambda unrestricted scratch : Nat .
137                    (lambda unrestricted weight : (family NvidiaLinearScalarInput) .
138                      (lambda unrestricted input : (family NvidiaLinearScalarInput) .
139                        (lambda unrestricted target : (family NvidiaLinearScalarInput) .
140                          (lambda unrestricted result : (family NvidiaLinearScalarResult) .
141                            (lambda unrestricted destination : Nat .
142                              (lambda unrestricted tail : (family NativePhysicalCommands) .
143                                (nlsUploadScalar
144                                  error
145                                  weight
146                                  scratch
147                                  data
148                                  (nlsUploadScalar
149                                    error
150                                    input
151                                    scratch
152                                    (naturalAdd data inputOffset)
153                                    (nlsUploadScalar
154                                      error
155                                      target
156                                      scratch
157                                      (naturalAdd data targetOffset)
158                                      (prReissueCounted
159                                        layout
160                                        counter
161                                        0
162                                        entry
163                                        (phCopyMappedToState
164                                        (naturalAdd
165                                        data
166                                        (naturalAdd outputOffset (nlsScalarResultOffset result)))
167                                        scratch
168                                        (linearStepPlanInputBytesFor 1)
169                                        (nativeNumericConvert
170                                        error
171                                        (constructor StdNumericConversion StdNumericWidenBinary32)
172                                        destination
173                                        (phLoad scratch)
174                                        tail)))))))))))))))))))))
175
176-- Learning rate is a runtime scalar in the launch's typed parameter block.
177-- A zero rate asks this realization for a gradient without changing W. The
178-- output still contains y, loss, gradient and proposed W; callers select only
179-- the fields their learning transition needs.
180def nvidiaLinearStepPlan =
181  (lambda unrestricted m : Nat .
182    (lambda unrestricted k : Nat .
183      (lambda unrestricted learningRate : Nat .
184        (lambda unrestricted architecture : (family ModelWord32) .
185          (lambda unrestricted computeClass : (family ModelWord32) .
186            (lambda unrestricted spaVersion : (family ModelWord32) .
187              (lambda unrestricted qmdBytes : Nat .
188                (lambda unrestricted product : Bytes .
189                  (lambda unrestricted abi : (family NvidiaPlanHostABI) .
190                    (lambda unrestricted availableVideo : Nat .
191                      (let unrestricted inputOffset =
192                        (linearStepPlanWeightsBytesFor m k)
193                        in
194                        (let unrestricted targetOffset =
195                          (naturalAdd inputOffset (linearStepPlanInputBytesFor k))
196                          in
197                          (let unrestricted outputOffset =
198                            (naturalAdd targetOffset (linearStepPlanTargetBytesFor m))
199                            in
200                            (let unrestricted dataBytes =
201                              (naturalAdd outputOffset (linearStepSM86OutputRecordBytesFor m k))
202                              in
203                              (eliminate
204                                NvidiaProgramPlan
205                                (lambda unrestricted value : (family NvidiaProgramPlan) .
206                                  (family NvidiaLinearStepPlan))
207                                (nvidiaProgramPlan
208                                  (linearStepPlanRegionsForProgram (linearStepSM86ProgramFor m k))
209                                  (lambda unrestricted programGPU : Nat .
210                                    (lambda unrestricted dataGPU : Nat .
211                                      (constructor
212                                        NvidiaLaunchSchedule
213                                        NvidiaLaunchScheduleOne
214                                        (constructor
215                                        NvidiaLaunchTemplate
216                                        NvidiaLaunchTemplateValue
217                                        linearStepSM86Identity
218                                        (constructor
219                                        NvidiaLaunchKernel
220                                        NvidiaLaunchKernelValue
221                                        (phWord programGPU)
222                                        1
223                                        1
224                                        1
225                                        1
226                                        1
227                                        linearStepPlanParameterBytes)
228                                        (linearStepPlanParametersAt
229                                        dataGPU
230                                        (naturalAdd dataGPU inputOffset)
231                                        (naturalAdd dataGPU targetOffset)
232                                        (naturalAdd dataGPU outputOffset)
233                                        learningRate)))))
234                                  linearStepPlanSubmissions
235                                  dataBytes
236                                  (constructor NvidiaPlanHostMemory NvidiaPlanHostSystemMemory)
237                                  architecture
238                                  computeClass
239                                  spaVersion
240                                  qmdBytes
241                                  product
242                                  abi
243                                  availableVideo)
244                                (branch
245                                  NvidiaProgramPlanValue
246                                  layout
247                                  component
248                                  recipes
249                                  dataHost
250                                  dataGPU
251                                  extent
252                                  admitted
253                                  .
254                                  (constructor
255                                    NvidiaLinearStepPlan
256                                    NvidiaLinearStepPlanValue
257                                    layout
258                                    component
259                                    recipes
260                                    dataHost
261                                    inputOffset
262                                    targetOffset
263                                    outputOffset
264                                    extent
265                                    (naturalAnd (linearStepSM86ShapeAdmitted m k) admitted))))))))))))))))))
266
267-- Keep the complete last scalar result, timestamped semaphore and fresh-channel
268-- issue count as separate device evidence. The destination is exclusive. The
269-- record is never a learning input or an implied whole-run timing measurement.
270def nlsScalarReadback =
271  (lambda unrestricted path : (family NativePhysicalOperand) .
272    (lambda unrestricted layout : (family NvidiaPlanHostLayout) .
273      (lambda unrestricted data : Nat .
274        (lambda unrestricted outputOffset : Nat .
275          (lambda unrestricted counter : Nat .
276            (lambda unrestricted tail : (family NativePhysicalCommands) .
277              (nativeFileOpen
278                path
279                nativeFileCandidateFlags
280                prSlotResult
281                prSlotStatus
282                b"device-readback-open"
283                (nativeFileWrite
284                  (phSlot prSlotResult)
285                  (phImm (naturalAdd data outputOffset))
286                  (linearStepSM86OutputRecordBytesFor 1 1)
287                  (phImm 0)
288                  prSlotStatus
289                  b"device-readback-output"
290                  (nativeFileWrite
291                    (phSlot prSlotResult)
292                    (phImm (prSemaphoreHost layout))
293                    prSemaphoreStride
294                    (phImm (linearStepSM86OutputRecordBytesFor 1 1))
295                    prSlotStatus
296                    b"device-readback-fence"
297                    (nativeFileWrite
298                      (phSlot prSlotResult)
299                      (phState counter)
300                      (bytes-length (dataBytesWord64LE (modelWord64FromNaturalTruncated 0)))
301                      (phImm
302                        (naturalAdd (linearStepSM86OutputRecordBytesFor 1 1) prSemaphoreStride))
303                      prSlotStatus
304                      b"device-readback-issues"
305                      (nativeFileClose
306                        (phSlot prSlotResult)
307                        prSlotStatus
308                        b"device-readback-close"
309                        tail)))))))))))

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.