1module Realization.Nvidia.SM86.Elementwise.GeneralSM86
2
3import Accelerator.SM86.Control
4import Accelerator.SM86.Immediate
5import Accelerator.SM86.Instruction
6import Accelerator.SM86.InstructionEncoding
7import Accelerator.SM86.Program
8import Accelerator.SM86.Types
9import Compiler.ApplicationBuilder
10import Data.SHA256Digest
11import Std.Natural
12
13family Elementwise256 : Type 0
14constructor Add256
15constructor Multiply256
16constructor Scale256
17
18end-family
19
20family Elementwise256ErrorCode : Type 0
21constructor Elementwise256InstructionCountMismatch
22constructor Elementwise256EncodingFailed
23constructor Elementwise256ByteCountMismatch
24constructor Elementwise256IdentityInvalid
25constructor Elementwise256HostFallbackObserved
26
27end-family
28
29family Elementwise256Telemetry : Type 0
30constructor Elementwise256TelemetryValue
31field unrestricted elementwise256TelemetryOperation : (family Elementwise256)
32field unrestricted elementwise256TelemetryExpectedInstructions : Nat
33field unrestricted elementwise256TelemetryActualInstructions : Nat
34field unrestricted elementwise256TelemetryExpectedBytes : Nat
35field unrestricted elementwise256TelemetryActualBytes : Nat
36field unrestricted elementwise256TelemetryRegisters : Nat
37field unrestricted elementwise256TelemetryGlobalLoads : Nat
38field unrestricted elementwise256TelemetryGlobalStores : Nat
39field unrestricted elementwise256TelemetryFloatAdds : Nat
40field unrestricted elementwise256TelemetryFloatMultiplies : Nat
41field unrestricted elementwise256TelemetryHostFallbackCalls : Nat
42field unrestricted elementwise256TelemetryApplicationBuilderContract : Nat
43
44end-family
45
46family Elementwise256BuildResult : Type 0
47constructor Elementwise256BuildReady
48field unrestricted elementwise256ReadyOperation : (family Elementwise256)
49field unrestricted elementwise256ReadyProgram : (family SM86Program)
50field unrestricted elementwise256ReadyImage : Bytes
51field unrestricted elementwise256ReadySHA256 : Bytes
52field unrestricted elementwise256ReadyTelemetry : (family Elementwise256Telemetry)
53constructor Elementwise256BuildRejected
54field unrestricted elementwise256RejectedCode : (family Elementwise256ErrorCode)
55field unrestricted elementwise256RejectedOrdinal : Nat
56field unrestricted elementwise256RejectedDetail : Bytes
57field unrestricted elementwise256RejectedTelemetry : (family Elementwise256Telemetry)
58
59end-family
60
61def elementwise256N1 =
62 (succ zero)
63
64def elementwise256N2 =
65 (byte-to-nat (byte 2))
66
67def elementwise256N12 =
68 (byte-to-nat (byte 12))
69
70def elementwise256N13 =
71 (byte-to-nat (byte 13))
72
73def elementwise256N16 =
74 (byte-to-nat (byte 16))
75
76def elementwise256N24 =
77 (byte-to-nat (byte 24))
78
79def elementwise256Word32 =
80 (lambda unrestricted b0 : Byte .
81 (lambda unrestricted b1 : Byte .
82 (lambda unrestricted b2 : Byte .
83 (lambda unrestricted b3 : Byte . (sm86Unsigned32 b0 b1 b2 b3)))))
84
85def elementwise256Zero32 =
86 (elementwise256Word32 (byte 0) (byte 0) (byte 0) (byte 0))
87
88def elementwise256Four32 =
89 (elementwise256Word32 (byte 4) (byte 0) (byte 0) (byte 0))
90
91def elementwise256ConstantBase =
92 (elementwise256Word32 (byte 40) (byte 0) (byte 0) (byte 0))
93
94def elementwise256Output0 =
95 (elementwise256Word32 (byte 96) (byte 1) (byte 0) (byte 0))
96
97def elementwise256Input0 =
98 (elementwise256Word32 (byte 104) (byte 1) (byte 0) (byte 0))
99
100def elementwise256Input1 =
101 (elementwise256Word32 (byte 112) (byte 1) (byte 0) (byte 0))
102
103def elementwise256Scalar0 =
104 (elementwise256Word32 (byte 144) (byte 1) (byte 0) (byte 0))
105
106def elementwise256Register =
107 (lambda unrestricted index : Byte . (sm86Register index))
108
109def elementwise256Instruction =
110 (lambda unrestricted body : (family SM86InstructionBody) .
111 (constructor
112 SM86Instruction
113 SM86InstructionValue
114 (constructor SM86InstructionGuard SM86InstructionAlways)
115 body))
116
117def elementwise256Next =
118 (lambda unrestricted body : (family SM86InstructionBody) .
119 (lambda unrestricted tail : (family SM86Program) .
120 (constructor SM86Program SM86ProgramNext (elementwise256Instruction body) tail)))
121
122def elementwise256End =
123 (constructor SM86Program SM86ProgramEnd)
124
125def elementwise256Set0 =
126 (sm86SetBarrierControl (constructor SM86Barrier SM86Barrier0))
127
128def elementwise256Set1 =
129 (sm86SetBarrierControl (constructor SM86Barrier SM86Barrier1))
130
131def elementwise256Wait0 =
132 (sm86WaitBarrierControl (constructor SM86WaitBarrier SM86WaitBarrier0))
133
134def elementwise256WaitBoth =
135 (constructor
136 SM86Control
137 SM86ControlValue
138 (byte 7)
139 (constructor SM86YieldMode SM86Continue)
140 (constructor SM86Barrier SM86BarrierNone)
141 (constructor SM86Barrier SM86BarrierNone)
142 (byte 3)
143 (byte 0))
144
145-- The prologue's schedule is a choice of controls. Existing kernels retain
146-- their qualified controls; generated expression kernels use the fixed-
147-- latency model's conservative schedule and can later propose a shorter one.
148def elementwise256PrefixWith =
149 (lambda unrestricted fixed : (family SM86Control) .
150 (lambda unrestricted set : (family SM86Control) .
151 (lambda unrestricted wait : (family SM86Control) .
152 (lambda unrestricted tail : (family SM86Program) .
153 (elementwise256Next
154 (constructor
155 SM86InstructionBody
156 SM86MoveConstant
157 (elementwise256Register (byte 1))
158 (byte 0)
159 elementwise256ConstantBase
160 fixed)
161 (elementwise256Next
162 (constructor
163 SM86InstructionBody
164 SM86SpecialToRegister
165 (elementwise256Register (byte 0))
166 (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdX)
167 set)
168 (elementwise256Next
169 (constructor
170 SM86InstructionBody
171 SM86MoveImmediate
172 (elementwise256Register (byte 3))
173 elementwise256Four32
174 wait)
175 (elementwise256Next
176 (constructor
177 SM86InstructionBody
178 SM86SpecialToRegister
179 (elementwise256Register (byte 2))
180 (constructor SM86SpecialRegister SM86ThreadIdX)
181 set)
182 (elementwise256Next
183 (constructor
184 SM86InstructionBody
185 SM86IntegerMultiplyAddConstant
186 (elementwise256Register (byte 0))
187 (elementwise256Register (byte 0))
188 (byte 0)
189 elementwise256Zero32
190 (elementwise256Register (byte 2))
191 wait)
192 tail)))))))))
193
194def elementwise256Prefix =
195 (elementwise256PrefixWith sm86SafeControl elementwise256Set0 elementwise256Wait0)
196
197def elementwise256BinaryTail =
198 (lambda unrestricted operation : (family Elementwise256) .
199 (elementwise256Next
200 (constructor
201 SM86InstructionBody
202 SM86IntegerMultiplyAddWideConstant
203 (elementwise256Register (byte 6))
204 (elementwise256Register (byte 0))
205 (elementwise256Register (byte 3))
206 (byte 0)
207 elementwise256Input0
208 sm86SafeControl)
209 (elementwise256Next
210 (constructor
211 SM86InstructionBody
212 SM86LoadGlobal
213 (elementwise256Register (byte 8))
214 (elementwise256Register (byte 6))
215 elementwise256Zero32
216 elementwise256Set0)
217 (elementwise256Next
218 (constructor
219 SM86InstructionBody
220 SM86IntegerMultiplyAddWideConstant
221 (elementwise256Register (byte 10))
222 (elementwise256Register (byte 0))
223 (elementwise256Register (byte 3))
224 (byte 0)
225 elementwise256Input1
226 sm86SafeControl)
227 (elementwise256Next
228 (constructor
229 SM86InstructionBody
230 SM86LoadGlobal
231 (elementwise256Register (byte 12))
232 (elementwise256Register (byte 10))
233 elementwise256Zero32
234 elementwise256Set1)
235 (elementwise256Next
236 (constructor
237 SM86InstructionBody
238 SM86IntegerMultiplyAddWideConstant
239 (elementwise256Register (byte 14))
240 (elementwise256Register (byte 0))
241 (elementwise256Register (byte 3))
242 (byte 0)
243 elementwise256Output0
244 sm86SafeControl)
245 (elementwise256Next
246 (eliminate
247 Elementwise256
248 (lambda unrestricted current : (family Elementwise256) .
249 (family SM86InstructionBody))
250 operation
251 (branch
252 Add256
253 .
254 (constructor
255 SM86InstructionBody
256 SM86FloatAdd
257 (elementwise256Register (byte 16))
258 (elementwise256Register (byte 8))
259 (elementwise256Register (byte 12))
260 elementwise256WaitBoth))
261 (branch
262 Multiply256
263 .
264 (constructor
265 SM86InstructionBody
266 SM86FloatMultiply
267 (elementwise256Register (byte 16))
268 (elementwise256Register (byte 8))
269 (elementwise256Register (byte 12))
270 elementwise256WaitBoth))
271 (branch
272 Scale256
273 .
274 (constructor
275 SM86InstructionBody
276 SM86FloatMultiply
277 (elementwise256Register (byte 16))
278 (elementwise256Register (byte 8))
279 (elementwise256Register (byte 12))
280 elementwise256WaitBoth)))
281 (elementwise256Next
282 (constructor
283 SM86InstructionBody
284 SM86StoreGlobal
285 (elementwise256Register (byte 14))
286 (elementwise256Register (byte 16))
287 elementwise256Zero32
288 sm86SafeControl)
289 (elementwise256Next
290 (constructor SM86InstructionBody SM86Exit sm86SafeControl)
291 elementwise256End)))))))))
292
293def elementwise256ScaleTail : (family SM86Program) =
294 (elementwise256Next
295 (constructor
296 SM86InstructionBody
297 SM86IntegerMultiplyAddWideConstant
298 (elementwise256Register (byte 6))
299 (elementwise256Register (byte 0))
300 (elementwise256Register (byte 3))
301 (byte 0)
302 elementwise256Input0
303 sm86SafeControl)
304 (elementwise256Next
305 (constructor
306 SM86InstructionBody
307 SM86LoadGlobal
308 (elementwise256Register (byte 8))
309 (elementwise256Register (byte 6))
310 elementwise256Zero32
311 elementwise256Set0)
312 (elementwise256Next
313 (constructor
314 SM86InstructionBody
315 SM86IntegerMultiplyAddWideConstant
316 (elementwise256Register (byte 10))
317 (elementwise256Register (byte 0))
318 (elementwise256Register (byte 3))
319 (byte 0)
320 elementwise256Output0
321 sm86SafeControl)
322 (elementwise256Next
323 (constructor
324 SM86InstructionBody
325 SM86MoveConstant
326 (elementwise256Register (byte 12))
327 (byte 0)
328 elementwise256Scalar0
329 sm86SafeControl)
330 (elementwise256Next
331 (constructor
332 SM86InstructionBody
333 SM86FloatMultiply
334 (elementwise256Register (byte 14))
335 (elementwise256Register (byte 8))
336 (elementwise256Register (byte 12))
337 elementwise256Wait0)
338 (elementwise256Next
339 (constructor
340 SM86InstructionBody
341 SM86StoreGlobal
342 (elementwise256Register (byte 10))
343 (elementwise256Register (byte 14))
344 elementwise256Zero32
345 sm86SafeControl)
346 (elementwise256Next
347 (constructor SM86InstructionBody SM86Exit sm86SafeControl)
348 elementwise256End)))))))
349
350def emitElementwise256SM86 =
351 (lambda unrestricted operation : (family Elementwise256) .
352 (eliminate
353 Elementwise256
354 (lambda unrestricted current : (family Elementwise256) . (family SM86Program))
355 operation
356 (branch Add256 . (elementwise256Prefix (elementwise256BinaryTail operation)))
357 (branch Multiply256 . (elementwise256Prefix (elementwise256BinaryTail operation)))
358 (branch Scale256 . (elementwise256Prefix elementwise256ScaleTail))))
359
360def elementwise256InstructionCount =
361 (lambda unrestricted operation : (family Elementwise256) .
362 (eliminate
363 Elementwise256
364 (lambda unrestricted current : (family Elementwise256) . Nat)
365 operation
366 (branch Add256 . elementwise256N13)
367 (branch Multiply256 . elementwise256N13)
368 (branch Scale256 . elementwise256N12)))
369
370def elementwise256RegisterCount : Nat =
371 elementwise256N24
372
373def elementwise256HostFallbackCalls : Nat =
374 zero
375
376def elementwise256ApplicationBuilderContract : Nat =
377 Compiler.ApplicationBuilder/applicationBuilderNativeOnly
378
379def elementwise256ErrorStableCode =
380 (lambda unrestricted code : (family Elementwise256ErrorCode) .
381 (eliminate
382 Elementwise256ErrorCode
383 (lambda unrestricted current : (family Elementwise256ErrorCode) . Bytes)
384 code
385 (branch Elementwise256InstructionCountMismatch . b"ELT-001")
386 (branch Elementwise256EncodingFailed . b"ELT-002")
387 (branch Elementwise256ByteCountMismatch . b"ELT-003")
388 (branch Elementwise256IdentityInvalid . b"ELT-004")
389 (branch Elementwise256HostFallbackObserved . b"ELT-005")))
390
391def elementwise256TelemetryFor =
392 (lambda unrestricted operation : (family Elementwise256) .
393 (lambda unrestricted actualInstructions : Nat .
394 (lambda unrestricted actualBytes : Nat .
395 (constructor
396 Elementwise256Telemetry
397 Elementwise256TelemetryValue
398 operation
399 (elementwise256InstructionCount operation)
400 actualInstructions
401 (naturalMultiply (elementwise256InstructionCount operation) elementwise256N16)
402 actualBytes
403 elementwise256RegisterCount
404 (eliminate
405 Elementwise256
406 (lambda unrestricted current : (family Elementwise256) . Nat)
407 operation
408 (branch Add256 . elementwise256N2)
409 (branch Multiply256 . elementwise256N2)
410 (branch Scale256 . elementwise256N1))
411 elementwise256N1
412 (eliminate
413 Elementwise256
414 (lambda unrestricted current : (family Elementwise256) . Nat)
415 operation
416 (branch Add256 . elementwise256N1)
417 (branch Multiply256 . zero)
418 (branch Scale256 . zero))
419 (eliminate
420 Elementwise256
421 (lambda unrestricted current : (family Elementwise256) . Nat)
422 operation
423 (branch Add256 . zero)
424 (branch Multiply256 . elementwise256N1)
425 (branch Scale256 . elementwise256N1))
426 elementwise256HostFallbackCalls
427 elementwise256ApplicationBuilderContract))))
428
429def elementwise256Reject =
430 (lambda unrestricted operation : (family Elementwise256) .
431 (lambda unrestricted code : (family Elementwise256ErrorCode) .
432 (lambda unrestricted ordinal : Nat .
433 (lambda unrestricted actualInstructions : Nat .
434 (lambda unrestricted actualBytes : Nat .
435 (constructor
436 Elementwise256BuildResult
437 Elementwise256BuildRejected
438 code
439 ordinal
440 (elementwise256ErrorStableCode code)
441 (elementwise256TelemetryFor operation actualInstructions actualBytes)))))))
442
443def elementwise256ImageSHA256 =
444 (lambda unrestricted operation : (family Elementwise256) .
445 (app
446 (lambda unrestricted program : (family SM86Program) .
447 (app
448 (lambda unrestricted actualInstructions : Nat .
449 (nat-eliminate
450 (lambda unrestricted countMatches : Nat . (family Elementwise256BuildResult))
451 (elementwise256Reject
452 operation
453 (constructor Elementwise256ErrorCode Elementwise256InstructionCountMismatch)
454 actualInstructions
455 actualInstructions
456 zero)
457 (lambda unrestricted countPredecessor : Nat .
458 (lambda unrestricted countInduction : (family Elementwise256BuildResult) .
459 (eliminate
460 SM86ProgramEncodingResult
461 (lambda unrestricted current : (family SM86ProgramEncodingResult) .
462 (family Elementwise256BuildResult))
463 (sm86EncodeProgram program)
464 (branch
465 SM86ProgramEncodingSucceeded
466 image
467 encodingTelemetry
468 .
469 (app
470 (lambda unrestricted actualBytes : Nat .
471 (nat-eliminate
472 (lambda unrestricted byteMatches : Nat .
473 (family Elementwise256BuildResult))
474 (elementwise256Reject
475 operation
476 (constructor Elementwise256ErrorCode Elementwise256ByteCountMismatch)
477 actualBytes
478 actualInstructions
479 actualBytes)
480 (lambda unrestricted bytePredecessor : Nat .
481 (lambda unrestricted byteInduction : (family Elementwise256BuildResult) .
482 (app
483 (lambda unrestricted identity : Bytes .
484 (nat-eliminate
485 (lambda unrestricted valid : Nat .
486 (family Elementwise256BuildResult))
487 (elementwise256Reject
488 operation
489 (constructor
490 Elementwise256ErrorCode
491 Elementwise256IdentityInvalid)
492 zero
493 actualInstructions
494 actualBytes)
495 (lambda unrestricted validPredecessor : Nat .
496 (lambda unrestricted validInduction : (family Elementwise256BuildResult) .
497 (constructor
498 Elementwise256BuildResult
499 Elementwise256BuildReady
500 operation
501 program
502 image
503 identity
504 (elementwise256TelemetryFor
505 operation
506 actualInstructions
507 actualBytes))))
508 (naturalEqual (bytes-length identity) (byte-to-nat (byte 64)))))
509 (sha256HexBytesOrEmpty (sha256Hex image)))))
510 (naturalEqual
511 actualBytes
512 (naturalMultiply
513 (elementwise256InstructionCount operation)
514 elementwise256N16))))
515 (bytes-length image)))
516 (branch
517 SM86ProgramEncodingFailed
518 ordinal
519 failure
520 encodingTelemetry
521 .
522 (elementwise256Reject
523 operation
524 (constructor Elementwise256ErrorCode Elementwise256EncodingFailed)
525 ordinal
526 actualInstructions
527 zero)))))
528 (naturalEqual actualInstructions (elementwise256InstructionCount operation))))
529 (sm86ProgramCount program)))
530 (emitElementwise256SM86 operation)))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.