module Component.Pointwise import Data.BFloat16 import BinaryFloatModel import Float32Model import FloatLiteralSpec import Model.Config import Model.Parameter import Model.Word32 import Std.Fin import Std.List import Std.Natural import Std.Word -- A scalar expression applied independently at each tensor element. The word -- type binds literal/input representation; a binary64 backend cannot consume -- a binary32 expression by reinterpreting its literals. Input indices are -- bounded in the type; hardware, pointers and thread counts do not occur here. -- The chosen arithmetic provider owns rounding. An explicit rounding operation -- remains a boundary that a backend may not fuse across. family PointwiseUnary : Type 0 constructor PointwiseNegate constructor PointwiseSquare constructor PointwiseRoundBF16 constructor PointwiseReciprocal constructor PointwiseReciprocalSquareRoot constructor PointwiseSquareRoot constructor PointwiseTanh constructor PointwiseExponentialBase2 end-family family PointwiseBinary : Type 0 constructor PointwiseAdd constructor PointwiseMultiply constructor PointwiseSubtract -- threshold(x, other): zero when x <= 0, otherwise other. NaNs take the -- other arm. This is the upstream ReLU/threshold-backward rule, including -- a zero gradient at the origin and masking a NaN gradient at negative x. constructor PointwisePositiveGate end-family family PointwiseExpression : Type 0 parameter erased pointwiseWord : Type 0 parameter erased pointwiseInputCount : Nat constructor PointwiseInput field unrestricted pointwiseInputIndex : (family StdFin pointwiseInputCount) constructor PointwiseConstant field unrestricted pointwiseConstantWord : pointwiseWord constructor PointwiseUnaryExpression field unrestricted pointwiseUnaryOperation : (family PointwiseUnary) recursive unrestricted pointwiseUnaryOperand constructor PointwiseBinaryExpression field unrestricted pointwiseBinaryOperation : (family PointwiseBinary) recursive unrestricted pointwiseLeftOperand recursive unrestricted pointwiseRightOperand end-family -- Assignments have sequential meaning: each expression reads the state left -- by its predecessor. A learner may explicitly share an intermediate by -- assigning a slot; subsequent expressions read that value without rebuilding -- its arithmetic tree. This is neither an optimizer catalogue nor a gradient -- requirement. Local update rules use exactly the same program structure. family PointwiseAssignment : Type 0 parameter erased pointwiseAssignmentWord : Type 0 parameter erased pointwiseAssignmentInputs : Nat constructor PointwiseAssign field unrestricted pointwiseDestination : (family StdFin pointwiseAssignmentInputs) field unrestricted pointwiseValue : (family PointwiseExpression pointwiseAssignmentWord pointwiseAssignmentInputs) end-family def pointwiseUnaryWith = (lambda unrestricted operations : (family Float32Operations) . (lambda unrestricted operation : (family PointwiseUnary) . (lambda unrestricted value : Nat . (eliminate PointwiseUnary (lambda unrestricted current : (family PointwiseUnary) . Nat) operation (branch PointwiseNegate . (f32Negate operations value)) (branch PointwiseSquare . (f32Multiply operations value value)) (branch PointwiseRoundBF16 . (bfloat16RoundWord value)) (branch PointwiseReciprocal . (f32Approximate operations 4 value)) (branch PointwiseReciprocalSquareRoot . (f32Approximate operations 5 value)) (branch PointwiseSquareRoot . (f32Approximate operations 6 value)) (branch PointwiseTanh . (f32Approximate operations 7 value)) (branch PointwiseExponentialBase2 . (f32Approximate operations 2 value)))))) def pointwiseBinaryWith = (lambda unrestricted operations : (family Float32Operations) . (lambda unrestricted operation : (family PointwiseBinary) . (lambda unrestricted left : Nat . (lambda unrestricted right : Nat . (eliminate PointwiseBinary (lambda unrestricted current : (family PointwiseBinary) . Nat) operation (branch PointwiseAdd . (f32Add operations left right)) (branch PointwiseMultiply . (f32Multiply operations left right)) (branch PointwiseSubtract . (f32Add operations left (f32Negate operations right))) (branch PointwisePositiveGate . (naturalSelect (naturalOr (modelF32Less left 0) (modelF32EqualValue left 0)) 0 right))))))) -- One structural fold serves semantic evaluators, use analysis and lowerings. -- The value domain may differ from the literal word representation: reference -- models use exact natural bit patterns, while a lowering uses typed fragments. def pointwiseFold = (lambda erased word : Type 0 . (lambda erased value : Type 0 . (lambda erased inputs : Nat . (lambda unrestricted unary : (pi unrestricted operation : (family PointwiseUnary) . (pi unrestricted x : value . value)) . (lambda unrestricted binary : (pi unrestricted operation : (family PointwiseBinary) . (pi unrestricted x : value . (pi unrestricted y : value . value))) . (lambda unrestricted literal : (pi unrestricted bits : word . value) . (lambda unrestricted expression : (family PointwiseExpression word inputs) . (lambda unrestricted values : (pi unrestricted index : (family StdFin inputs) . value) . (eliminate PointwiseExpression (lambda unrestricted current : (family PointwiseExpression word inputs) . value) expression (branch PointwiseInput index . (values index)) (branch PointwiseConstant bits . (literal bits)) (branch PointwiseUnaryExpression operation operand result . (unary operation result)) (branch PointwiseBinaryExpression operation left right leftResult rightResult . (binary operation leftResult rightResult))))))))))) def pointwiseEvaluateWith = (lambda unrestricted operations : (family Float32Operations) . (lambda erased inputs : Nat . (lambda unrestricted expression : (family PointwiseExpression StdU32 inputs) . (lambda unrestricted values : (pi unrestricted i : (family StdFin inputs) . Nat) . (pointwiseFold StdU32 Nat inputs (pointwiseUnaryWith operations) (pointwiseBinaryWith operations) stdU32ToNatural expression values))))) def pointwiseExpressionReads = (lambda erased word : Type 0 . (lambda erased inputs : Nat . (lambda unrestricted expression : (family PointwiseExpression word inputs) . (lambda unrestricted index : Nat . (pointwiseFold word Nat inputs (lambda unrestricted operation : (family PointwiseUnary) . (lambda unrestricted x : Nat . x)) (lambda unrestricted operation : (family PointwiseBinary) . naturalOr) (lambda unrestricted bits : word . 0) expression (lambda unrestricted source : (family StdFin inputs) . (naturalEqual index (stdFinToNatural inputs source)))))))) def pointwiseProgramReadsInitial = (lambda erased word : Type 0 . (lambda erased inputs : Nat . (lambda unrestricted program : (family StdList (family PointwiseAssignment word inputs)) . (lambda unrestricted index : Nat . (eliminate StdList (lambda unrestricted current : (family StdList (family PointwiseAssignment word inputs)) . Nat) program (branch StdListEmpty . 0) (branch StdListCons head tail rest . (eliminate PointwiseAssignment (lambda unrestricted current : (family PointwiseAssignment word inputs) . Nat) head (branch PointwiseAssign destination expression . (naturalOr (pointwiseExpressionReads word inputs expression index) (naturalAnd (naturalIsZero (naturalEqual index (stdFinToNatural inputs destination))) rest)))))))))) def pointwiseProgramWrites = (lambda erased word : Type 0 . (lambda erased inputs : Nat . (lambda unrestricted program : (family StdList (family PointwiseAssignment word inputs)) . (lambda unrestricted index : Nat . (eliminate StdList (lambda unrestricted current : (family StdList (family PointwiseAssignment word inputs)) . Nat) program (branch StdListEmpty . 0) (branch StdListCons head tail rest . (eliminate PointwiseAssignment (lambda unrestricted current : (family PointwiseAssignment word inputs) . Nat) head (branch PointwiseAssign destination expression . (naturalOr (naturalEqual index (stdFinToNatural inputs destination)) rest))))))))) def pointwiseProgramEvaluate = (lambda erased word : Type 0 . (lambda erased value : Type 0 . (lambda erased inputs : Nat . (lambda unrestricted unary : (pi unrestricted operation : (family PointwiseUnary) . (pi unrestricted x : value . value)) . (lambda unrestricted binary : (pi unrestricted operation : (family PointwiseBinary) . (pi unrestricted x : value . (pi unrestricted y : value . value))) . (lambda unrestricted literal : (pi unrestricted bits : word . value) . (lambda unrestricted program : (family StdList (family PointwiseAssignment word inputs)) . (eliminate StdList (lambda unrestricted current : (family StdList (family PointwiseAssignment word inputs)) . (pi unrestricted state : (pi unrestricted index : (family StdFin inputs) . value) . (pi unrestricted index : (family StdFin inputs) . value))) program (branch StdListEmpty . (lambda unrestricted state : (pi unrestricted index : (family StdFin inputs) . value) . state)) (branch StdListCons head tail next . (lambda unrestricted state : (pi unrestricted index : (family StdFin inputs) . value) . (eliminate PointwiseAssignment (lambda unrestricted current : (family PointwiseAssignment word inputs) . (pi unrestricted index : (family StdFin inputs) . value)) head (branch PointwiseAssign destination expression . (let unrestricted result = (pointwiseFold word value inputs unary binary literal expression state) in (next (lambda unrestricted index : (family StdFin inputs) . (nat-eliminate (lambda unrestricted same : Nat . value) (state index) (lambda unrestricted previous : Nat . (lambda unrestricted ignored : value . result)) (naturalEqual (stdFinToNatural inputs destination) (stdFinToNatural inputs index)))))))))))))))))) -- Logical slot identities are independent of their eventual physical layout. -- Learner compositions use this relation to keep writable state separate from -- observations; the target planner separately owns physical non-overlap. def pointwiseSlotsSeparate = (lambda erased slots : Nat . (lambda unrestricted left : (family StdFin slots) . (lambda unrestricted right : (family StdFin slots) . (naturalIsZero (naturalEqual (stdFinToNatural slots left) (stdFinToNatural slots right)))))) def pointwiseUnary = (lambda erased word : Type 0 . (lambda erased inputs : Nat . (lambda unrestricted operation : (family PointwiseUnary) . (lambda unrestricted value : (family PointwiseExpression word inputs) . (constructor PointwiseExpression PointwiseUnaryExpression word inputs operation value))))) def pointwiseBinary = (lambda erased word : Type 0 . (lambda erased inputs : Nat . (lambda unrestricted operation : (family PointwiseBinary) . (lambda unrestricted left : (family PointwiseExpression word inputs) . (lambda unrestricted right : (family PointwiseExpression word inputs) . (constructor PointwiseExpression PointwiseBinaryExpression word inputs operation left right)))))) -- Each literal or arithmetic node defines one temporary value. Inputs refer -- to existing state. This expression traversal is representation-independent; -- a target maps values to register pairs or arena words and checks its own -- limit. Assignments reuse the largest expression workspace, not their sum. def pointwiseTemporaryValues = (lambda erased word : Type 0 . (lambda erased inputs : Nat . (lambda unrestricted expression : (family PointwiseExpression word inputs) . (pointwiseFold word Nat inputs (lambda unrestricted operation : (family PointwiseUnary) . (lambda unrestricted child : Nat . (succ child))) (lambda unrestricted operation : (family PointwiseBinary) . (lambda unrestricted left : Nat . (lambda unrestricted right : Nat . (succ (naturalAdd left right))))) (lambda unrestricted bits : word . 1) expression (lambda unrestricted index : (family StdFin inputs) . 0))))) def pointwiseProgramTemporaryValues = (lambda erased word : Type 0 . (lambda erased inputs : Nat . (lambda unrestricted program : (family StdList (family PointwiseAssignment word inputs)) . (stdListFold (family PointwiseAssignment word inputs) Nat (lambda unrestricted assignment : (family PointwiseAssignment word inputs) . (lambda unrestricted rest : Nat . (eliminate PointwiseAssignment (lambda unrestricted current : (family PointwiseAssignment word inputs) . Nat) assignment (branch PointwiseAssign destination expression . (let unrestricted count = (pointwiseTemporaryValues word inputs expression) in (naturalSelect (naturalLess count rest) rest count)))))) 0 program)))) -- BF16 squared ReLU and its backward are ordinary compositions. Upstream -- rounds both ReLU's result and the square; backward rounds 2*relu(x), then -- its product with dy, then applies threshold-backward. Multiplying by a -- zero mask would propagate NaNs and is not the same derivative program. def pointwiseReluSquared = (lambda erased inputs : Nat . (lambda unrestricted x : (family PointwiseExpression StdU32 inputs) . (let unrestricted relu = (pointwiseUnary StdU32 inputs (constructor PointwiseUnary PointwiseRoundBF16) (pointwiseBinary StdU32 inputs (constructor PointwiseBinary PointwisePositiveGate) x x)) in (pointwiseUnary StdU32 inputs (constructor PointwiseUnary PointwiseRoundBF16) (pointwiseUnary StdU32 inputs (constructor PointwiseUnary PointwiseSquare) relu))))) def pointwiseReluSquaredBackward = (lambda erased inputs : Nat . (lambda unrestricted x : (family PointwiseExpression StdU32 inputs) . (lambda unrestricted gradient : (family PointwiseExpression StdU32 inputs) . (let unrestricted round = (pointwiseUnary StdU32 inputs (constructor PointwiseUnary PointwiseRoundBF16)) in (let unrestricted multiply = (pointwiseBinary StdU32 inputs (constructor PointwiseBinary PointwiseMultiply)) in (let unrestricted relu = (round (pointwiseBinary StdU32 inputs (constructor PointwiseBinary PointwisePositiveGate) x x)) in (pointwiseBinary StdU32 inputs (constructor PointwiseBinary PointwisePositiveGate) x (round (multiply gradient (round (multiply (constructor PointwiseExpression PointwiseConstant StdU32 inputs 0x40000000) relu)))))))))))