Hardware.Nvidia.SM86.Command.AffineLaunchSchedule

1module Hardware.Nvidia.SM86.Command.AffineLaunchSchedule
2
3import Hardware.Nvidia.SM86.Command.WholeProgramPlan
4import Model.Parameter
5import Model.Word64
6import Std.Natural
7import Std.Word
8
9-- A repeated body may change only its parameter words between iterations.
10-- Compare adjacent body instances by launch ordinal and word. Differences
11-- are modulo 2^64 because device pointers are encoded as 64-bit words. The
12-- body is retained once; WholeProgramPlan expands its affine parameters
13-- when it emits the launch-table recipe. Callers prove the body shape and
14-- pointer progression for the geometry they instantiate.
15def nvidiaAffineNoLaunches : (family NvidiaLaunchSchedule) =
16  (constructor NvidiaLaunchSchedule NvidiaLaunchScheduleEmpty)
17def nvidiaAffineScheduleIf =
18  (lambda unrestricted condition : Nat .
19    (lambda unrestricted whenTrue : (family NvidiaLaunchSchedule) .
20      (lambda unrestricted whenFalse : (family NvidiaLaunchSchedule) .
21        (nat-eliminate
22          (lambda unrestricted current : Nat . (family NvidiaLaunchSchedule))
23          whenFalse
24          (lambda unrestricted p : Nat .
25            (lambda unrestricted induction : (family NvidiaLaunchSchedule) . whenTrue))
26          condition))))
27
28def nvidiaAffineWordMaximum : Nat =
29  (modelWord64Natural stdU64AllOnes)
30
31def nvidiaAffineWordDifference =
32  (lambda unrestricted a : Nat .
33    (lambda unrestricted b : Nat .
34      (naturalSelect
35        (naturalLessOrEqual a b)
36        (naturalSaturatingSubtract b a)
37        (naturalSaturatingSubtract
38          nvidiaAffineWordMaximum
39          (naturalSaturatingSubtract (naturalSaturatingSubtract a b) 1)))))
40
41def nvidiaAffinePatchValue =
42  (lambda unrestricted patches : (family NvidiaParameterPatches) .
43    (lambda unrestricted word : Nat .
44      (eliminate
45        NvidiaParameterPatches
46        (lambda unrestricted current : (family NvidiaParameterPatches) . Nat)
47        patches
48        (branch NvidiaParameterPatchesEnd . 0)
49        (branch
50          NvidiaParameterPatchesNext
51          head
52          tail
53          induction
54          .
55          (eliminate
56            NvidiaParameterPatch
57            (lambda unrestricted current : (family NvidiaParameterPatch) . Nat)
58            head
59            (branch
60              NvidiaParameterPatchValue
61              at
62              value
63              .
64              (naturalSelect (naturalEqual at word) (modelWord64Natural value) induction)))))))
65
66def nvidiaAffinePatchPresent =
67  (lambda unrestricted patches : (family NvidiaParameterPatches) .
68    (lambda unrestricted word : Nat .
69      (eliminate
70        NvidiaParameterPatches
71        (lambda unrestricted current : (family NvidiaParameterPatches) . Nat)
72        patches
73        (branch NvidiaParameterPatchesEnd . 0)
74        (branch
75          NvidiaParameterPatchesNext
76          head
77          tail
78          induction
79          .
80          (eliminate
81            NvidiaParameterPatch
82            (lambda unrestricted current : (family NvidiaParameterPatch) . Nat)
83            head
84            (branch
85              NvidiaParameterPatchValue
86              at
87              value
88              .
89              (naturalOr (naturalEqual at word) induction)))))))
90
91def nvidiaAffinePatchesOf =
92  (lambda unrestricted template : (family NvidiaLaunchTemplate) .
93    (eliminate
94      NvidiaLaunchTemplate
95      (lambda unrestricted current : (family NvidiaLaunchTemplate) .
96        (family NvidiaParameterPatches))
97      template
98      (branch
99        NvidiaLaunchTemplateValue
100        identity
101        kernel
102        block
103        .
104        (eliminate
105          NvidiaParameterBlock
106          (lambda unrestricted current : (family NvidiaParameterBlock) .
107            (family NvidiaParameterPatches))
108          block
109          (branch NvidiaParameterBlockValue patches . patches)))))
110
111def nvidiaAffineAdjust =
112  (lambda unrestricted ordinal : Nat .
113    (lambda unrestricted word : Nat .
114      (lambda unrestricted delta : Nat .
115        (lambda unrestricted tail : (family NvidiaParameterAdjustments) .
116          (nat-eliminate
117            (lambda unrestricted current : Nat . (family NvidiaParameterAdjustments))
118            tail
119            (lambda unrestricted q : Nat .
120              (lambda unrestricted ignored : (family NvidiaParameterAdjustments) .
121                (constructor
122                  NvidiaParameterAdjustments
123                  NvidiaParameterAdjustmentsNext
124                  (constructor
125                    NvidiaParameterAdjustment
126                    NvidiaParameterAdjustmentValue
127                    (constructor
128                      NvidiaParameterAdjustmentScope
129                      NvidiaParameterAdjustmentBlock
130                      ordinal)
131                    word
132                    (modelWord64FromNaturalTruncated delta))
133                  tail)))
134            (naturalNonzero delta))))))
135
136-- Compare sparse words for one launch. An absent word is zero in the launch
137-- table, so additions and removals are deltas from or to zero. The ordinal
138-- scope prevents a pointer in one kernel from relocating the same ABI slot
139-- when that slot is a scalar in another kernel.
140def nvidiaAffineBlockDeltas =
141  (lambda unrestricted ordinal : Nat .
142    (lambda unrestricted first : (family NvidiaParameterPatches) .
143      (lambda unrestricted second : (family NvidiaParameterPatches) .
144        (lambda unrestricted tail : (family NvidiaParameterAdjustments) .
145          (eliminate
146            NvidiaParameterPatches
147            (lambda unrestricted current : (family NvidiaParameterPatches) .
148              (family NvidiaParameterAdjustments))
149            first
150            (branch
151              NvidiaParameterPatchesEnd
152              .
153              (eliminate
154                NvidiaParameterPatches
155                (lambda unrestricted current : (family NvidiaParameterPatches) .
156                  (family NvidiaParameterAdjustments))
157                second
158                (branch NvidiaParameterPatchesEnd . tail)
159                (branch
160                  NvidiaParameterPatchesNext
161                  head
162                  rest
163                  induction
164                  .
165                  (eliminate
166                    NvidiaParameterPatch
167                    (lambda unrestricted current : (family NvidiaParameterPatch) .
168                      (family NvidiaParameterAdjustments))
169                    head
170                    (branch
171                      NvidiaParameterPatchValue
172                      word
173                      value
174                      .
175                      (nat-eliminate
176                        (lambda unrestricted current : Nat . (family NvidiaParameterAdjustments))
177                        (nvidiaAffineAdjust ordinal word (modelWord64Natural value) induction)
178                        (lambda unrestricted q : Nat .
179                          (lambda unrestricted ignored : (family NvidiaParameterAdjustments) .
180                            induction))
181                        (nvidiaAffinePatchPresent first word)))))))
182            (branch
183              NvidiaParameterPatchesNext
184              head
185              rest
186              induction
187              .
188              (eliminate
189                NvidiaParameterPatch
190                (lambda unrestricted current : (family NvidiaParameterPatch) .
191                  (family NvidiaParameterAdjustments))
192                head
193                (branch
194                  NvidiaParameterPatchValue
195                  word
196                  value
197                  .
198                  (nvidiaAffineAdjust
199                    ordinal
200                    word
201                    (nvidiaAffineWordDifference (modelWord64Natural value) (nvidiaAffinePatchValue second word))
202                    induction)))))))))
203
204def nvidiaAffineOneOf =
205  (lambda unrestricted schedule : (family NvidiaLaunchSchedule) .
206    (eliminate
207      NvidiaLaunchSchedule
208      (lambda unrestricted current : (family NvidiaLaunchSchedule) .
209        (family NvidiaParameterPatches))
210      schedule
211      (branch
212        NvidiaLaunchScheduleEmpty
213        .
214        (constructor NvidiaParameterPatches NvidiaParameterPatchesEnd))
215      (branch NvidiaLaunchScheduleOne template . (nvidiaAffinePatchesOf template))
216      (branch
217        NvidiaLaunchScheduleAppend
218        left
219        right
220        il
221        ir
222        .
223        (constructor NvidiaParameterPatches NvidiaParameterPatchesEnd))
224      (branch
225        NvidiaLaunchScheduleRepeat
226        count
227        iteration
228        body
229        ib
230        .
231        (constructor NvidiaParameterPatches NvidiaParameterPatchesEnd))))
232
233def nvidiaAffineLeftOf =
234  (lambda unrestricted schedule : (family NvidiaLaunchSchedule) .
235    (eliminate
236      NvidiaLaunchSchedule
237      (lambda unrestricted current : (family NvidiaLaunchSchedule) . (family NvidiaLaunchSchedule))
238      schedule
239      (branch NvidiaLaunchScheduleEmpty . schedule)
240      (branch NvidiaLaunchScheduleOne template . schedule)
241      (branch NvidiaLaunchScheduleAppend left right il ir . left)
242      (branch NvidiaLaunchScheduleRepeat count iteration body ib . body)))
243
244def nvidiaAffineRightOf =
245  (lambda unrestricted schedule : (family NvidiaLaunchSchedule) .
246    (eliminate
247      NvidiaLaunchSchedule
248      (lambda unrestricted current : (family NvidiaLaunchSchedule) . (family NvidiaLaunchSchedule))
249      schedule
250      (branch NvidiaLaunchScheduleEmpty . schedule)
251      (branch NvidiaLaunchScheduleOne template . schedule)
252      (branch NvidiaLaunchScheduleAppend left right il ir . right)
253      (branch NvidiaLaunchScheduleRepeat count iteration body ib . body)))
254
255-- the deltas between two bodies of the same shape, blocks numbered in the
256-- expanded body from `offset`; a nested repeat's difference applies to
257-- every copy of its body
258def nvidiaAffineDeltas =
259  (lambda unrestricted first : (family NvidiaLaunchSchedule) .
260    (eliminate
261      NvidiaLaunchSchedule
262      (lambda unrestricted current : (family NvidiaLaunchSchedule) .
263        (pi unrestricted second : (family NvidiaLaunchSchedule) .
264          (pi unrestricted offset : Nat .
265            (pi unrestricted tail : (family NvidiaParameterAdjustments) .
266              (family NvidiaParameterAdjustments)))))
267      first
268      (branch
269        NvidiaLaunchScheduleEmpty
270        .
271        (lambda unrestricted second : (family NvidiaLaunchSchedule) .
272          (lambda unrestricted offset : Nat .
273            (lambda unrestricted tail : (family NvidiaParameterAdjustments) . tail))))
274      (branch
275        NvidiaLaunchScheduleOne
276        template
277        .
278        (lambda unrestricted second : (family NvidiaLaunchSchedule) .
279          (lambda unrestricted offset : Nat .
280            (lambda unrestricted tail : (family NvidiaParameterAdjustments) .
281              (nvidiaAffineBlockDeltas offset (nvidiaAffinePatchesOf template) (nvidiaAffineOneOf second) tail)))))
282      (branch
283        NvidiaLaunchScheduleAppend
284        left
285        right
286        il
287        ir
288        .
289        (lambda unrestricted second : (family NvidiaLaunchSchedule) .
290          (lambda unrestricted offset : Nat .
291            (lambda unrestricted tail : (family NvidiaParameterAdjustments) .
292              (il
293                (nvidiaAffineLeftOf second)
294                offset
295                (ir (nvidiaAffineRightOf second) (naturalAdd offset (nvidiaLaunchScheduleCount left)) tail))))))
296      (branch
297        NvidiaLaunchScheduleRepeat
298        count
299        iteration
300        body
301        ib
302        .
303        (lambda unrestricted second : (family NvidiaLaunchSchedule) .
304          (lambda unrestricted offset : Nat .
305            (lambda unrestricted tail : (family NvidiaParameterAdjustments) .
306              (nat-eliminate
307                (lambda unrestricted current : Nat . (family NvidiaParameterAdjustments))
308                tail
309                (lambda unrestricted p : Nat .
310                  (lambda unrestricted induction : (family NvidiaParameterAdjustments) .
311                    (ib
312                      (nvidiaAffineRightOf second)
313                      (naturalAdd offset (naturalMultiply p (nvidiaLaunchScheduleCount body)))
314                      induction)))
315                count)))))))
316
317-- `count` bodies of an affine loop, as one Repeat
318def nvidiaAffineLaunchRepeat =
319  (lambda unrestricted count : Nat .
320    (lambda unrestricted body : (pi unrestricted index : Nat . (family NvidiaLaunchSchedule)) .
321      (nvidiaAffineScheduleIf
322        (naturalLess count 2)
323        (nvidiaAffineScheduleIf (naturalIsZero count) nvidiaAffineNoLaunches (body 0))
324        (let unrestricted first =
325          (body 0)
326          in
327          (constructor
328            NvidiaLaunchSchedule
329            NvidiaLaunchScheduleRepeat
330            count
331            (constructor
332              NvidiaParameterIteration
333              NvidiaParameterIterationAffine
334              (nvidiaAffineDeltas
335                first
336                (body 1)
337                0
338                (constructor NvidiaParameterAdjustments NvidiaParameterAdjustmentsEnd)))
339            first)))))