Component.PointwiseIdentity

1module Component.PointwiseIdentity
2
3import Component.Pointwise
4import Model.Parameter
5import Data.Bytes
6import Model.Word64
7import Std.Fin
8import Std.List
9import Std.Natural
10import Std.Word
11
12-- Versioned prefix encoding of typed expression trees. Tags encode structure;
13-- every number and literal occupies one canonical little-endian word. Input
14-- count, assignment count and destination order are part of the identity.
15-- This is semantic input to checkpoint hashing, not an executable packager.
16def pointwiseIdentityWord =
17  (lambda unrestricted value : Nat . (dataBytesWord64LE (modelWord64FromNaturalTruncated value)))
18
19def pointwiseUnaryIdentity =
20  (lambda unrestricted operation : (family PointwiseUnary) .
21    (eliminate
22      PointwiseUnary
23      (lambda unrestricted current : (family PointwiseUnary) . Bytes)
24      operation
25      (branch PointwiseNegate . b"n")
26      (branch PointwiseSquare . b"s")
27      (branch PointwiseRoundBF16 . b"b")
28      (branch PointwiseReciprocal . b"r")
29      (branch PointwiseReciprocalSquareRoot . b"i")
30      (branch PointwiseSquareRoot . b"q")
31      (branch PointwiseTanh . b"t")
32      (branch PointwiseExponentialBase2 . b"e")))
33
34def pointwiseBinaryIdentity =
35  (lambda unrestricted operation : (family PointwiseBinary) .
36    (eliminate
37      PointwiseBinary
38      (lambda unrestricted current : (family PointwiseBinary) . Bytes)
39      operation
40      (branch PointwiseAdd . b"+")
41      (branch PointwiseMultiply . b"*")
42      (branch PointwiseSubtract . b"-")
43      (branch PointwisePositiveGate . b">")))
44
45def pointwiseExpressionIdentity =
46  (lambda unrestricted slots : Nat .
47    (lambda unrestricted expression : (family PointwiseExpression StdU64 slots) .
48      (pointwiseFold
49        StdU64
50        Bytes
51        slots
52        (lambda unrestricted operation : (family PointwiseUnary) .
53          (lambda unrestricted child : Bytes .
54            (bytes-append b"u" (bytes-append (pointwiseUnaryIdentity operation) child))))
55        (lambda unrestricted operation : (family PointwiseBinary) .
56          (lambda unrestricted left : Bytes .
57            (lambda unrestricted right : Bytes .
58              (bytes-append
59                b"b"
60                (bytes-append (pointwiseBinaryIdentity operation) (bytes-append left right))))))
61        (lambda unrestricted bits : StdU64 . (bytes-append b"c" (dataBytesWord64LE bits)))
62        expression
63        (lambda unrestricted index : (family StdFin slots) .
64          (bytes-append b"i" (pointwiseIdentityWord (stdFinToNatural slots index)))))))
65
66def pointwiseProgramIdentity =
67  (lambda unrestricted slots : Nat .
68    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
69      (bytes-append
70        b"alpha-pointwise-v1;"
71        (bytes-append
72          (pointwiseIdentityWord slots)
73          (bytes-append
74            (pointwiseIdentityWord
75              (stdListLength (family PointwiseAssignment StdU64 slots) program))
76            (stdListFold
77              (family PointwiseAssignment StdU64 slots)
78              Bytes
79              (lambda unrestricted assignment : (family PointwiseAssignment StdU64 slots) .
80                (lambda unrestricted rest : Bytes .
81                  (eliminate
82                    PointwiseAssignment
83                    (lambda unrestricted current : (family PointwiseAssignment StdU64 slots) .
84                      Bytes)
85                    assignment
86                    (branch
87                      PointwiseAssign
88                      destination
89                      expression
90                      .
91                      (bytes-append
92                        (pointwiseIdentityWord (stdFinToNatural slots destination))
93                        (bytes-append (pointwiseExpressionIdentity slots expression) rest))))))
94              b""
95              program))))))