module Learning.DataParallel

import Component.Pointwise
import Component.PointwiseIdentity
import Model.Parameter
import Data.Bytes
import Std.Fin
import Std.List
import Std.Natural
import Std.Word

-- A data-parallel learner: one parameter vector and its optimizer state
-- (for instance momentum), both replicated on every owner; immutable samples
-- partitioned among owners in possibly unequal shares; a per-sample
-- contribution program; and an update program every owner applies
-- identically to the mean contribution.
--
-- One update, as a mathematical definition independent of placement:
--   numerator = sum over all samples of contribution(parameters, sample)
--   mass      = number of samples
--   mean      = numerator / mass
--   (parameters, state) := update(parameters, state, mean, rate)
-- The sum's evaluation order is not part of this definition. A realization
-- fixes it explicitly (an ordered combine tree, Runtime.OrderedCombine) and
-- its numerical contract names that order; a different order is a different
-- floating-point program, never a silent reassociation.
--
-- Slots, by position, with k parameters, m optimizer-state words and samples
-- of width s. The contribution program reads parameters [0, k) and one
-- sample [k, k + s) and writes only its contribution [k + s, 2k + s). The
-- update program reads parameters [0, k), the state [k, k + m), the mean
-- [k + m, 2k + m) and the rate at 2k + m, and writes only parameters and
-- state [0, k + m). The state is persistent learner state: checkpointed with
-- the parameters and required to agree across owners. The same typed
-- assignment semantics (Component.Pointwise) serve gradient and non-gradient
-- rules: a contribution need not be a gradient.
family DataParallelSystem : Type 0
parameter erased dataParallelContributionSlots : Nat
parameter erased dataParallelUpdateSlots : Nat
constructor DataParallelSystemOf
field unrestricted dataParallelParameterCount : Nat
field unrestricted dataParallelSampleWidth : Nat
field unrestricted dataParallelContribution : (family StdList (family PointwiseAssignment StdU64 dataParallelContributionSlots))
field unrestricted dataParallelUpdate : (family StdList (family PointwiseAssignment StdU64 dataParallelUpdateSlots))
field unrestricted dataParallelInitial : (family StdList StdU64)
-- the optimizer state's initial words (m = its length; none for plain steps)
field unrestricted dataParallelInitialState : (family StdList StdU64)
field unrestricted dataParallelRate : StdU64
-- per owner, per sample, the sample's s words
field unrestricted dataParallelSamples : (family StdList (family StdList (family StdList StdU64)))

end-family

def dataParallelParameters =
  (lambda erased contribution : Nat .
    (lambda erased update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (eliminate
          DataParallelSystem
          (lambda unrestricted current : (family DataParallelSystem contribution update) . Nat)
          system
          (branch
            DataParallelSystemOf
            parameters
            width
            program
            rule
            initial
            state
            rate
            samples
            .
            parameters)))))

def dataParallelWidth =
  (lambda erased contribution : Nat .
    (lambda erased update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (eliminate
          DataParallelSystem
          (lambda unrestricted current : (family DataParallelSystem contribution update) . Nat)
          system
          (branch
            DataParallelSystemOf
            parameters
            width
            program
            rule
            initial
            state
            rate
            samples
            .
            width)))))

def dataParallelContributionProgram =
  (lambda erased contribution : Nat .
    (lambda erased update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (eliminate
          DataParallelSystem
          (lambda unrestricted current : (family DataParallelSystem contribution update) .
            (family StdList (family PointwiseAssignment StdU64 contribution)))
          system
          (branch
            DataParallelSystemOf
            parameters
            width
            program
            rule
            initial
            state
            rate
            samples
            .
            program)))))

def dataParallelUpdateProgram =
  (lambda erased contribution : Nat .
    (lambda erased update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (eliminate
          DataParallelSystem
          (lambda unrestricted current : (family DataParallelSystem contribution update) .
            (family StdList (family PointwiseAssignment StdU64 update)))
          system
          (branch
            DataParallelSystemOf
            parameters
            width
            program
            rule
            initial
            state
            rate
            samples
            .
            rule)))))

def dataParallelInitialParameters =
  (lambda erased contribution : Nat .
    (lambda erased update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (eliminate
          DataParallelSystem
          (lambda unrestricted current : (family DataParallelSystem contribution update) .
            (family StdList StdU64))
          system
          (branch
            DataParallelSystemOf
            parameters
            width
            program
            rule
            initial
            state
            rate
            samples
            .
            initial)))))

def dataParallelInitialStateWords =
  (lambda erased contribution : Nat .
    (lambda erased update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (eliminate
          DataParallelSystem
          (lambda unrestricted current : (family DataParallelSystem contribution update) .
            (family StdList StdU64))
          system
          (branch
            DataParallelSystemOf
            parameters
            width
            program
            rule
            initial
            state
            rate
            samples
            .
            state)))))

-- m, the optimizer state's word count.
def dataParallelStateCount =
  (lambda erased contribution : Nat .
    (lambda erased update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (stdListLength StdU64 (dataParallelInitialStateWords contribution update system)))))

def dataParallelRateWord =
  (lambda erased contribution : Nat .
    (lambda erased update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (eliminate
          DataParallelSystem
          (lambda unrestricted current : (family DataParallelSystem contribution update) . StdU64)
          system
          (branch
            DataParallelSystemOf
            parameters
            width
            program
            rule
            initial
            state
            rate
            samples
            .
            rate)))))

def dataParallelOwnerSamples =
  (lambda erased contribution : Nat .
    (lambda erased update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (eliminate
          DataParallelSystem
          (lambda unrestricted current : (family DataParallelSystem contribution update) .
            (family StdList (family StdList (family StdList StdU64))))
          system
          (branch
            DataParallelSystemOf
            parameters
            width
            program
            rule
            initial
            state
            rate
            samples
            .
            samples)))))

def dataParallelOwners =
  (lambda erased contribution : Nat .
    (lambda erased update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (stdListLength
          (family StdList (family StdList StdU64))
          (dataParallelOwnerSamples contribution update system)))))

-- The samples one owner holds, in order (none for an owner past the list).
def dataParallelSamplesOf =
  (lambda erased contribution : Nat .
    (lambda erased update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (lambda unrestricted owner : Nat .
          (stdListHeadOr
            (family StdList (family StdList StdU64))
            (constructor StdList StdListEmpty (family StdList StdU64))
            (stdListDrop
              (family StdList (family StdList StdU64))
              owner
              (dataParallelOwnerSamples contribution update system)))))))

def dataParallelTotalMass =
  (lambda erased contribution : Nat .
    (lambda erased update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (stdListFold
          (family StdList (family StdList StdU64))
          Nat
          (lambda unrestricted owned : (family StdList (family StdList StdU64)) .
            (lambda unrestricted rest : Nat .
              (naturalAdd (stdListLength (family StdList StdU64) owned) rest)))
          0
          (dataParallelOwnerSamples contribution update system)))))

-- Every assignment of `program` writes a slot in [low, high).
def dataParallelWritesWithin =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
      (lambda unrestricted low : Nat .
        (lambda unrestricted high : Nat .
          (stdListFold
            (family PointwiseAssignment StdU64 slots)
            Nat
            (lambda unrestricted assignment : (family PointwiseAssignment StdU64 slots) .
              (lambda unrestricted rest : Nat .
                (eliminate
                  PointwiseAssignment
                  (lambda unrestricted current : (family PointwiseAssignment StdU64 slots) . Nat)
                  assignment
                  (branch
                    PointwiseAssign
                    destination
                    expression
                    .
                    (naturalAnd
                      rest
                      (naturalAnd
                        (naturalLessOrEqual low (stdFinToNatural slots destination))
                        (naturalLess (stdFinToNatural slots destination) high)))))))
            1
            program)))))

-- Admitted when the slot spaces match the declared shapes (with the state's m
-- words), each program writes only its outputs, every sample has the declared width, at least two
-- owners exist and the total mass is positive and exactly representable as a
-- binary64 integer (below 2^53), since realizations carry it in binary64.
def dataParallelAdmitted =
  (lambda unrestricted contribution : Nat .
    (lambda unrestricted update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (let unrestricted k =
          (dataParallelParameters contribution update system)
          in
          (let unrestricted s =
            (dataParallelWidth contribution update system)
            in
            (let unrestricted m =
              (dataParallelStateCount contribution update system)
              in
              (naturalAnd
                (naturalNonzero k)
                (naturalAnd
                  (naturalEqual contribution (naturalAdd (naturalMultiply 2 k) s))
                  (naturalAnd
                    (naturalEqual update (succ (naturalAdd (naturalMultiply 2 k) m)))
                    (naturalAnd
                      (naturalEqual
                        (stdListLength
                          StdU64
                          (dataParallelInitialParameters contribution update system))
                        k)
                      (naturalAnd
                        (dataParallelWritesWithin
                          contribution
                          (dataParallelContributionProgram contribution update system)
                          (naturalAdd k s)
                          contribution)
                        (naturalAnd
                          (dataParallelWritesWithin
                            update
                            (dataParallelUpdateProgram contribution update system)
                            0
                            (naturalAdd k m))
                          (naturalAnd
                            (naturalLessOrEqual 2 (dataParallelOwners contribution update system))
                            (naturalAnd
                              (naturalNonzero (dataParallelTotalMass contribution update system))
                              (naturalAnd
                                (naturalBelowTwoPower
                                  53
                                  (dataParallelTotalMass contribution update system))
                                (stdListFold
                                  (family StdList (family StdList StdU64))
                                  Nat
                                  (lambda unrestricted owned : (family StdList (family StdList StdU64)) .
                                    (lambda unrestricted rest : Nat .
                                      (naturalAnd
                                        rest
                                        (stdListFold
                                        (family StdList StdU64)
                                        Nat
                                        (lambda unrestricted sample : (family StdList StdU64) .
                                        (lambda unrestricted more : Nat .
                                        (naturalAnd
                                        more
                                        (naturalEqual (stdListLength StdU64 sample) s))))
                                        1
                                        owned))))
                                  1
                                  (dataParallelOwnerSamples contribution update system)))))))))))))))))

def dataParallelWordsIdentity =
  (lambda unrestricted words : (family StdList StdU64) .
    (bytes-append
      (pointwiseIdentityWord (stdListLength StdU64 words))
      (stdListFold
        StdU64
        Bytes
        (lambda unrestricted word : StdU64 .
          (lambda unrestricted tail : Bytes . (bytes-append (dataBytesWord64LE word) tail)))
        b""
        words)))

-- The learner's rule: both programs, initial parameters and state, rate.
def dataParallelRuleIdentity =
  (lambda unrestricted contribution : Nat .
    (lambda unrestricted update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (bytes-append
          (pointwiseIdentityWord (dataParallelParameters contribution update system))
          (bytes-append
            (pointwiseIdentityWord (dataParallelWidth contribution update system))
            (bytes-append
              (pointwiseProgramIdentity
                contribution
                (dataParallelContributionProgram contribution update system))
              (bytes-append
                (pointwiseProgramIdentity
                  update
                  (dataParallelUpdateProgram contribution update system))
                (bytes-append
                  (dataParallelWordsIdentity
                    (dataParallelInitialParameters contribution update system))
                  (bytes-append
                    (dataParallelWordsIdentity
                      (dataParallelInitialStateWords contribution update system))
                    (dataBytesWord64LE (dataParallelRateWord contribution update system)))))))))))

def dataParallelSamplesIdentity =
  (lambda unrestricted samples : (family StdList (family StdList StdU64)) .
    (stdListFold
      (family StdList StdU64)
      Bytes
      (lambda unrestricted sample : (family StdList StdU64) .
        (lambda unrestricted more : Bytes . (bytes-append (dataParallelWordsIdentity sample) more)))
      b""
      samples))

-- The learner's identity: its rule and every owner's samples in order.
-- Evaluation order is not part of it; a realization adds its numerical
-- contract beside it.
def dataParallelIdentity =
  (lambda unrestricted contribution : Nat .
    (lambda unrestricted update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (bytes-append
          b"alpha-data-parallel-v1;"
          (bytes-append
            (dataParallelRuleIdentity contribution update system)
            (stdListFold
              (family StdList (family StdList StdU64))
              Bytes
              (lambda unrestricted owned : (family StdList (family StdList StdU64)) .
                (lambda unrestricted rest : Bytes .
                  (bytes-append
                    (pointwiseIdentityWord (stdListLength (family StdList StdU64) owned))
                    (bytes-append (dataParallelSamplesIdentity owned) rest))))
              b""
              (dataParallelOwnerSamples contribution update system)))))))

-- Every sample in owner order: the learner's input sequence without its
-- placement among owners.
def dataParallelAllSamples =
  (lambda erased contribution : Nat .
    (lambda erased update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (stdListFold
          (family StdList (family StdList StdU64))
          (family StdList (family StdList StdU64))
          (lambda unrestricted owned : (family StdList (family StdList StdU64)) .
            (lambda unrestricted rest : (family StdList (family StdList StdU64)) .
              (stdListAppend (family StdList StdU64) owned rest)))
          (constructor StdList StdListEmpty (family StdList StdU64))
          (dataParallelOwnerSamples contribution update system)))))

-- The placement-independent identity: the rule and the input sequence, not
-- how the samples are split among owners. State saved under one split is a
-- state of the same learner under any other split of the same sequence.
def dataParallelLearnerIdentity =
  (lambda unrestricted contribution : Nat .
    (lambda unrestricted update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (let unrestricted samples =
          (dataParallelAllSamples contribution update system)
          in
          (bytes-append
            b"alpha-data-parallel-learner-v1;"
            (bytes-append
              (dataParallelRuleIdentity contribution update system)
              (bytes-append
                (pointwiseIdentityWord (stdListLength (family StdList StdU64) samples))
                (dataParallelSamplesIdentity samples))))))))

-- The same learner with its input sequence split into `owners` contiguous
-- shares whose sizes differ by at most one, larger first.
def dataParallelRepartitioned =
  (lambda erased contribution : Nat .
    (lambda erased update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (lambda unrestricted owners : Nat .
          (let unrestricted samples =
            (dataParallelAllSamples contribution update system)
            in
            (let unrestricted total =
              (stdListLength (family StdList StdU64) samples)
              in
              (let unrestricted base =
                (naturalDivideUnchecked total owners)
                in
                (let unrestricted extra =
                  (naturalModuloUnchecked total owners)
                  in
                  (eliminate
                    DataParallelSystem
                    (lambda unrestricted current : (family DataParallelSystem contribution update) .
                      (family DataParallelSystem contribution update))
                    system
                    (branch
                      DataParallelSystemOf
                      parameters
                      width
                      program
                      rule
                      initial
                      state
                      rate
                      held
                      .
                      (constructor
                        DataParallelSystem
                        DataParallelSystemOf
                        contribution
                        update
                        parameters
                        width
                        program
                        rule
                        initial
                        state
                        rate
                        (stdListMap
                          Nat
                          (family StdList (family StdList StdU64))
                          (lambda unrestricted share : Nat .
                            (stdListTake
                              (family StdList StdU64)
                              (naturalAdd base (naturalLess share extra))
                              (stdListDrop
                                (family StdList StdU64)
                                (naturalAdd
                                  (naturalMultiply share base)
                                  (naturalMinimum share extra))
                                samples)))
                          (stdListNaturalsBelow owners)))))))))))))

-- ---- coordinate separability ----
-- An update is separable when each coordinate's results depend only on that
-- coordinate and on shared values, so different owners can each update the
-- coordinates they own (optimizer-state partitioning). With q per-coordinate
-- state words (word j of coordinate i at state index j k + i) and the other
-- m - q k words shared, an update slot belongs to coordinate i (parameter i,
-- a state word of coordinate i, mean i) or is shared (a shared state word,
-- the rate). `dataParallelSlotClass` is succ i for coordinate i and 0 for
-- shared.
def dataParallelSlotClass =
  (lambda unrestricted k : Nat .
    (lambda unrestricted m : Nat .
      (lambda unrestricted q : Nat .
        (lambda unrestricted slot : Nat .
          (naturalSelect
            (naturalLess slot k)
            (succ slot)
            (naturalSelect
              (naturalLess slot (naturalAdd k (naturalMultiply q k)))
              (succ (naturalModuloUnchecked (naturalSaturatingSubtract slot k) k))
              (naturalSelect
                (naturalLess slot (naturalAdd k m))
                0
                (naturalSelect
                  (naturalLess slot (naturalAdd (naturalMultiply 2 k) m))
                  (succ (naturalSaturatingSubtract slot (naturalAdd k m)))
                  0))))))))

-- Every input `expression` reads is shared or of class `allowed`.
def dataParallelReadsClass =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted k : Nat .
      (lambda unrestricted m : Nat .
        (lambda unrestricted q : Nat .
          (lambda unrestricted allowed : Nat .
            (lambda unrestricted expression : (family PointwiseExpression StdU64 slots) .
              (pointwiseFold
                StdU64
                Nat
                slots
                (lambda unrestricted operation : (family PointwiseUnary) .
                  (lambda unrestricted child : Nat . child))
                (lambda unrestricted operation : (family PointwiseBinary) .
                  (lambda unrestricted left : Nat .
                    (lambda unrestricted right : Nat . (naturalAnd left right))))
                (lambda unrestricted word : StdU64 . 1)
                expression
                (lambda unrestricted index : (family StdFin slots) .
                  (let unrestricted class =
                    (dataParallelSlotClass k m q (stdFinToNatural slots index))
                    in
                    (naturalOr (naturalIsZero class) (naturalEqual class allowed)))))))))))

-- The update is separable under q per-coordinate state words: a coordinate's
-- assignment reads only its coordinate and shared values, and a shared
-- assignment reads only shared values.
def dataParallelSeparable =
  (lambda unrestricted contribution : Nat .
    (lambda unrestricted update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (lambda unrestricted q : Nat .
          (let unrestricted k =
            (dataParallelParameters contribution update system)
            in
            (let unrestricted m =
              (dataParallelStateCount contribution update system)
              in
              (naturalAnd
                (naturalLessOrEqual (naturalMultiply q k) m)
                (stdListFold
                  (family PointwiseAssignment StdU64 update)
                  Nat
                  (lambda unrestricted assignment : (family PointwiseAssignment StdU64 update) .
                    (lambda unrestricted rest : Nat .
                      (eliminate
                        PointwiseAssignment
                        (lambda unrestricted current : (family PointwiseAssignment StdU64 update) .
                          Nat)
                        assignment
                        (branch
                          PointwiseAssign
                          destination
                          expression
                          .
                          (naturalAnd
                            rest
                            (dataParallelReadsClass
                              update
                              k
                              m
                              q
                              (dataParallelSlotClass k m q (stdFinToNatural update destination))
                              expression))))))
                  1
                  (dataParallelUpdateProgram contribution update system)))))))))

-- The assignments one owner runs: every shared assignment and those of the
-- coordinates `owns` names (owns i = 1), in program order.
def dataParallelOwnedProgram =
  (lambda unrestricted contribution : Nat .
    (lambda unrestricted update : Nat .
      (lambda unrestricted system : (family DataParallelSystem contribution update) .
        (lambda unrestricted q : Nat .
          (lambda unrestricted owns : (pi unrestricted coordinate : Nat . Nat) .
            (let unrestricted k =
              (dataParallelParameters contribution update system)
              in
              (let unrestricted m =
                (dataParallelStateCount contribution update system)
                in
                (stdListFold
                  (family PointwiseAssignment StdU64 update)
                  (family StdList (family PointwiseAssignment StdU64 update))
                  (lambda unrestricted assignment : (family PointwiseAssignment StdU64 update) .
                    (lambda unrestricted rest : (family StdList (family PointwiseAssignment StdU64 update)) .
                      (eliminate
                        PointwiseAssignment
                        (lambda unrestricted current : (family PointwiseAssignment StdU64 update) .
                          (family StdList (family PointwiseAssignment StdU64 update)))
                        assignment
                        (branch
                          PointwiseAssign
                          destination
                          expression
                          .
                          (let unrestricted class =
                            (dataParallelSlotClass k m q (stdFinToNatural update destination))
                            in
                            (nat-eliminate
                              (lambda unrestricted current : Nat .
                                (family StdList (family PointwiseAssignment StdU64 update)))
                              rest
                              (lambda unrestricted previous : Nat .
                                (lambda unrestricted ignored : (family StdList (family PointwiseAssignment StdU64 update)) .
                                  (constructor
                                    StdList
                                    StdListCons
                                    (family PointwiseAssignment StdU64 update)
                                    assignment
                                    rest)))
                              (naturalOr (naturalIsZero class) (owns (naturalPredecessor class)))))))))
                  (constructor StdList StdListEmpty (family PointwiseAssignment StdU64 update))
                  (dataParallelUpdateProgram contribution update system)))))))))