1module Component.Pointwise
2
3import Data.BFloat16
4import BinaryFloatModel
5import Float32Model
6import FloatLiteralSpec
7import Model.Config
8import Model.Parameter
9import Model.Word32
10import Std.Fin
11import Std.List
12import Std.Natural
13import Std.Word
14
15-- A scalar expression applied independently at each tensor element. The word
16-- type binds literal/input representation; a binary64 backend cannot consume
17-- a binary32 expression by reinterpreting its literals. Input indices are
18-- bounded in the type; hardware, pointers and thread counts do not occur here.
19-- The chosen arithmetic provider owns rounding. An explicit rounding operation
20-- remains a boundary that a backend may not fuse across.
21family PointwiseUnary : Type 0
22constructor PointwiseNegate
23constructor PointwiseSquare
24constructor PointwiseRoundBF16
25constructor PointwiseReciprocal
26constructor PointwiseReciprocalSquareRoot
27constructor PointwiseSquareRoot
28constructor PointwiseTanh
29constructor PointwiseExponentialBase2
30
31end-family
32
33family PointwiseBinary : Type 0
34constructor PointwiseAdd
35constructor PointwiseMultiply
36constructor PointwiseSubtract
37-- threshold(x, other): zero when x <= 0, otherwise other. NaNs take the
38-- other arm. This is the upstream ReLU/threshold-backward rule, including
39-- a zero gradient at the origin and masking a NaN gradient at negative x.
40constructor PointwisePositiveGate
41
42end-family
43
44family PointwiseExpression : Type 0
45parameter erased pointwiseWord : Type 0
46parameter erased pointwiseInputCount : Nat
47constructor PointwiseInput
48field unrestricted pointwiseInputIndex : (family StdFin pointwiseInputCount)
49constructor PointwiseConstant
50field unrestricted pointwiseConstantWord : pointwiseWord
51constructor PointwiseUnaryExpression
52field unrestricted pointwiseUnaryOperation : (family PointwiseUnary)
53recursive unrestricted pointwiseUnaryOperand
54constructor PointwiseBinaryExpression
55field unrestricted pointwiseBinaryOperation : (family PointwiseBinary)
56recursive unrestricted pointwiseLeftOperand
57recursive unrestricted pointwiseRightOperand
58
59end-family
60
61-- Assignments have sequential meaning: each expression reads the state left
62-- by its predecessor. A learner may explicitly share an intermediate by
63-- assigning a slot; subsequent expressions read that value without rebuilding
64-- its arithmetic tree. This is neither an optimizer catalogue nor a gradient
65-- requirement. Local update rules use exactly the same program structure.
66family PointwiseAssignment : Type 0
67parameter erased pointwiseAssignmentWord : Type 0
68parameter erased pointwiseAssignmentInputs : Nat
69constructor PointwiseAssign
70field unrestricted pointwiseDestination : (family StdFin pointwiseAssignmentInputs)
71field unrestricted pointwiseValue : (family PointwiseExpression pointwiseAssignmentWord pointwiseAssignmentInputs)
72
73end-family
74
75def pointwiseUnaryWith =
76 (lambda unrestricted operations : (family Float32Operations) .
77 (lambda unrestricted operation : (family PointwiseUnary) .
78 (lambda unrestricted value : Nat .
79 (eliminate
80 PointwiseUnary
81 (lambda unrestricted current : (family PointwiseUnary) . Nat)
82 operation
83 (branch PointwiseNegate . (f32Negate operations value))
84 (branch PointwiseSquare . (f32Multiply operations value value))
85 (branch PointwiseRoundBF16 . (bfloat16RoundWord value))
86 (branch PointwiseReciprocal . (f32Approximate operations 4 value))
87 (branch PointwiseReciprocalSquareRoot . (f32Approximate operations 5 value))
88 (branch PointwiseSquareRoot . (f32Approximate operations 6 value))
89 (branch PointwiseTanh . (f32Approximate operations 7 value))
90 (branch PointwiseExponentialBase2 . (f32Approximate operations 2 value))))))
91
92def pointwiseBinaryWith =
93 (lambda unrestricted operations : (family Float32Operations) .
94 (lambda unrestricted operation : (family PointwiseBinary) .
95 (lambda unrestricted left : Nat .
96 (lambda unrestricted right : Nat .
97 (eliminate
98 PointwiseBinary
99 (lambda unrestricted current : (family PointwiseBinary) . Nat)
100 operation
101 (branch PointwiseAdd . (f32Add operations left right))
102 (branch PointwiseMultiply . (f32Multiply operations left right))
103 (branch PointwiseSubtract . (f32Add operations left (f32Negate operations right)))
104 (branch
105 PointwisePositiveGate
106 .
107 (naturalSelect (naturalOr (modelF32Less left 0) (modelF32EqualValue left 0)) 0 right)))))))
108
109-- One structural fold serves semantic evaluators, use analysis and lowerings.
110-- The value domain may differ from the literal word representation: reference
111-- models use exact natural bit patterns, while a lowering uses typed fragments.
112def pointwiseFold =
113 (lambda erased word : Type 0 .
114 (lambda erased value : Type 0 .
115 (lambda erased inputs : Nat .
116 (lambda unrestricted unary : (pi unrestricted operation : (family PointwiseUnary) . (pi unrestricted x : value . value)) .
117 (lambda unrestricted binary : (pi unrestricted operation : (family PointwiseBinary) . (pi unrestricted x : value . (pi unrestricted y : value . value))) .
118 (lambda unrestricted literal : (pi unrestricted bits : word . value) .
119 (lambda unrestricted expression : (family PointwiseExpression word inputs) .
120 (lambda unrestricted values : (pi unrestricted index : (family StdFin inputs) . value) .
121 (eliminate
122 PointwiseExpression
123 (lambda unrestricted current : (family PointwiseExpression word inputs) . value)
124 expression
125 (branch PointwiseInput index . (values index))
126 (branch PointwiseConstant bits . (literal bits))
127 (branch
128 PointwiseUnaryExpression
129 operation
130 operand
131 result
132 .
133 (unary operation result))
134 (branch
135 PointwiseBinaryExpression
136 operation
137 left
138 right
139 leftResult
140 rightResult
141 .
142 (binary operation leftResult rightResult)))))))))))
143
144def pointwiseEvaluateWith =
145 (lambda unrestricted operations : (family Float32Operations) .
146 (lambda erased inputs : Nat .
147 (lambda unrestricted expression : (family PointwiseExpression StdU32 inputs) .
148 (lambda unrestricted values : (pi unrestricted i : (family StdFin inputs) . Nat) .
149 (pointwiseFold
150 StdU32
151 Nat
152 inputs
153 (pointwiseUnaryWith operations)
154 (pointwiseBinaryWith operations)
155 stdU32ToNatural
156 expression
157 values)))))
158
159def pointwiseExpressionReads =
160 (lambda erased word : Type 0 .
161 (lambda erased inputs : Nat .
162 (lambda unrestricted expression : (family PointwiseExpression word inputs) .
163 (lambda unrestricted index : Nat .
164 (pointwiseFold
165 word
166 Nat
167 inputs
168 (lambda unrestricted operation : (family PointwiseUnary) .
169 (lambda unrestricted x : Nat . x))
170 (lambda unrestricted operation : (family PointwiseBinary) . naturalOr)
171 (lambda unrestricted bits : word . 0)
172 expression
173 (lambda unrestricted source : (family StdFin inputs) .
174 (naturalEqual index (stdFinToNatural inputs source))))))))
175
176def pointwiseProgramReadsInitial =
177 (lambda erased word : Type 0 .
178 (lambda erased inputs : Nat .
179 (lambda unrestricted program : (family StdList (family PointwiseAssignment word inputs)) .
180 (lambda unrestricted index : Nat .
181 (eliminate
182 StdList
183 (lambda unrestricted current : (family StdList (family PointwiseAssignment word inputs)) .
184 Nat)
185 program
186 (branch StdListEmpty . 0)
187 (branch
188 StdListCons
189 head
190 tail
191 rest
192 .
193 (eliminate
194 PointwiseAssignment
195 (lambda unrestricted current : (family PointwiseAssignment word inputs) . Nat)
196 head
197 (branch
198 PointwiseAssign
199 destination
200 expression
201 .
202 (naturalOr
203 (pointwiseExpressionReads word inputs expression index)
204 (naturalAnd
205 (naturalIsZero (naturalEqual index (stdFinToNatural inputs destination)))
206 rest))))))))))
207
208def pointwiseProgramWrites =
209 (lambda erased word : Type 0 .
210 (lambda erased inputs : Nat .
211 (lambda unrestricted program : (family StdList (family PointwiseAssignment word inputs)) .
212 (lambda unrestricted index : Nat .
213 (eliminate
214 StdList
215 (lambda unrestricted current : (family StdList (family PointwiseAssignment word inputs)) .
216 Nat)
217 program
218 (branch StdListEmpty . 0)
219 (branch
220 StdListCons
221 head
222 tail
223 rest
224 .
225 (eliminate
226 PointwiseAssignment
227 (lambda unrestricted current : (family PointwiseAssignment word inputs) . Nat)
228 head
229 (branch
230 PointwiseAssign
231 destination
232 expression
233 .
234 (naturalOr (naturalEqual index (stdFinToNatural inputs destination)) rest)))))))))
235
236def pointwiseProgramEvaluate =
237 (lambda erased word : Type 0 .
238 (lambda erased value : Type 0 .
239 (lambda erased inputs : Nat .
240 (lambda unrestricted unary : (pi unrestricted operation : (family PointwiseUnary) . (pi unrestricted x : value . value)) .
241 (lambda unrestricted binary : (pi unrestricted operation : (family PointwiseBinary) . (pi unrestricted x : value . (pi unrestricted y : value . value))) .
242 (lambda unrestricted literal : (pi unrestricted bits : word . value) .
243 (lambda unrestricted program : (family StdList (family PointwiseAssignment word inputs)) .
244 (eliminate
245 StdList
246 (lambda unrestricted current : (family StdList (family PointwiseAssignment word inputs)) .
247 (pi unrestricted state : (pi unrestricted index : (family StdFin inputs) . value) .
248 (pi unrestricted index : (family StdFin inputs) . value)))
249 program
250 (branch
251 StdListEmpty
252 .
253 (lambda unrestricted state : (pi unrestricted index : (family StdFin inputs) . value) .
254 state))
255 (branch
256 StdListCons
257 head
258 tail
259 next
260 .
261 (lambda unrestricted state : (pi unrestricted index : (family StdFin inputs) . value) .
262 (eliminate
263 PointwiseAssignment
264 (lambda unrestricted current : (family PointwiseAssignment word inputs) .
265 (pi unrestricted index : (family StdFin inputs) . value))
266 head
267 (branch
268 PointwiseAssign
269 destination
270 expression
271 .
272 (let unrestricted result =
273 (pointwiseFold word value inputs unary binary literal expression state)
274 in
275 (next
276 (lambda unrestricted index : (family StdFin inputs) .
277 (nat-eliminate
278 (lambda unrestricted same : Nat . value)
279 (state index)
280 (lambda unrestricted previous : Nat .
281 (lambda unrestricted ignored : value . result))
282 (naturalEqual
283 (stdFinToNatural inputs destination)
284 (stdFinToNatural inputs index))))))))))))))))))
285
286-- Logical slot identities are independent of their eventual physical layout.
287-- Learner compositions use this relation to keep writable state separate from
288-- observations; the target planner separately owns physical non-overlap.
289def pointwiseSlotsSeparate =
290 (lambda erased slots : Nat .
291 (lambda unrestricted left : (family StdFin slots) .
292 (lambda unrestricted right : (family StdFin slots) .
293 (naturalIsZero (naturalEqual (stdFinToNatural slots left) (stdFinToNatural slots right))))))
294
295def pointwiseUnary =
296 (lambda erased word : Type 0 .
297 (lambda erased inputs : Nat .
298 (lambda unrestricted operation : (family PointwiseUnary) .
299 (lambda unrestricted value : (family PointwiseExpression word inputs) .
300 (constructor PointwiseExpression PointwiseUnaryExpression word inputs operation value)))))
301
302def pointwiseBinary =
303 (lambda erased word : Type 0 .
304 (lambda erased inputs : Nat .
305 (lambda unrestricted operation : (family PointwiseBinary) .
306 (lambda unrestricted left : (family PointwiseExpression word inputs) .
307 (lambda unrestricted right : (family PointwiseExpression word inputs) .
308 (constructor
309 PointwiseExpression
310 PointwiseBinaryExpression
311 word
312 inputs
313 operation
314 left
315 right))))))
316
317-- Each literal or arithmetic node defines one temporary value. Inputs refer
318-- to existing state. This expression traversal is representation-independent;
319-- a target maps values to register pairs or arena words and checks its own
320-- limit. Assignments reuse the largest expression workspace, not their sum.
321def pointwiseTemporaryValues =
322 (lambda erased word : Type 0 .
323 (lambda erased inputs : Nat .
324 (lambda unrestricted expression : (family PointwiseExpression word inputs) .
325 (pointwiseFold
326 word
327 Nat
328 inputs
329 (lambda unrestricted operation : (family PointwiseUnary) .
330 (lambda unrestricted child : Nat . (succ child)))
331 (lambda unrestricted operation : (family PointwiseBinary) .
332 (lambda unrestricted left : Nat .
333 (lambda unrestricted right : Nat . (succ (naturalAdd left right)))))
334 (lambda unrestricted bits : word . 1)
335 expression
336 (lambda unrestricted index : (family StdFin inputs) . 0)))))
337
338def pointwiseProgramTemporaryValues =
339 (lambda erased word : Type 0 .
340 (lambda erased inputs : Nat .
341 (lambda unrestricted program : (family StdList (family PointwiseAssignment word inputs)) .
342 (stdListFold
343 (family PointwiseAssignment word inputs)
344 Nat
345 (lambda unrestricted assignment : (family PointwiseAssignment word inputs) .
346 (lambda unrestricted rest : Nat .
347 (eliminate
348 PointwiseAssignment
349 (lambda unrestricted current : (family PointwiseAssignment word inputs) . Nat)
350 assignment
351 (branch
352 PointwiseAssign
353 destination
354 expression
355 .
356 (let unrestricted count =
357 (pointwiseTemporaryValues word inputs expression)
358 in
359 (naturalSelect (naturalLess count rest) rest count))))))
360 0
361 program))))
362
363-- BF16 squared ReLU and its backward are ordinary compositions. Upstream
364-- rounds both ReLU's result and the square; backward rounds 2*relu(x), then
365-- its product with dy, then applies threshold-backward. Multiplying by a
366-- zero mask would propagate NaNs and is not the same derivative program.
367def pointwiseReluSquared =
368 (lambda erased inputs : Nat .
369 (lambda unrestricted x : (family PointwiseExpression StdU32 inputs) .
370 (let unrestricted relu =
371 (pointwiseUnary
372 StdU32
373 inputs
374 (constructor PointwiseUnary PointwiseRoundBF16)
375 (pointwiseBinary StdU32 inputs (constructor PointwiseBinary PointwisePositiveGate) x x))
376 in
377 (pointwiseUnary
378 StdU32
379 inputs
380 (constructor PointwiseUnary PointwiseRoundBF16)
381 (pointwiseUnary StdU32 inputs (constructor PointwiseUnary PointwiseSquare) relu)))))
382
383def pointwiseReluSquaredBackward =
384 (lambda erased inputs : Nat .
385 (lambda unrestricted x : (family PointwiseExpression StdU32 inputs) .
386 (lambda unrestricted gradient : (family PointwiseExpression StdU32 inputs) .
387 (let unrestricted round =
388 (pointwiseUnary StdU32 inputs (constructor PointwiseUnary PointwiseRoundBF16))
389 in
390 (let unrestricted multiply =
391 (pointwiseBinary StdU32 inputs (constructor PointwiseBinary PointwiseMultiply))
392 in
393 (let unrestricted relu =
394 (round
395 (pointwiseBinary
396 StdU32
397 inputs
398 (constructor PointwiseBinary PointwisePositiveGate)
399 x
400 x))
401 in
402 (pointwiseBinary
403 StdU32
404 inputs
405 (constructor PointwiseBinary PointwisePositiveGate)
406 x
407 (round
408 (multiply
409 gradient
410 (round
411 (multiply
412 (constructor PointwiseExpression PointwiseConstant StdU32 inputs 0x40000000)
413 relu)))))))))))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.