module Realization.Nvidia.SM86.ScaledGradientUpdateSM86

import Accelerator.SM86.Control
import Accelerator.SM86.Instruction
import Accelerator.SM86.Types

-- Realizations of Learning.ScaledGradientUpdate's two independent operations.
-- Every arithmetic instruction rounds separately to binary64, nearest-even.
-- No contraction or conversion to binary32 is permitted. Values are supplied
-- as register pairs by the enclosing kernel's layout; no model facts live here.
--
-- This fragment reserves SB0 (write) and SB1 (read). Entry waits retire prior
-- work on both; each issue replaces both reservations. The continuation must
-- wait on mask 3 before consuming a result or overwriting a source. A caller
-- using other barriers must first retire them. The scoreboard checks the
-- complete enclosing program, including the continuation.
def scaledUpdate64Control =
  (constructor
    SM86Control
    SM86ControlValue
    (byte 15)
    (constructor SM86YieldMode SM86Continue)
    (constructor SM86Barrier SM86Barrier0)
    (constructor SM86Barrier SM86Barrier1)
    (byte 3)
    (byte 0))

def scaledUpdate64Next =
  (lambda unrestricted operation : (family SM86Float64Operation) .
    (lambda unrestricted destination : (family SM86RegisterPair) .
      (lambda unrestricted left : (family SM86RegisterPair) .
        (lambda unrestricted right : (family SM86RegisterPair) .
          (lambda unrestricted tail : (family SM86Program) .
            (constructor
              SM86Program
              SM86ProgramNext
              (sm86Instruction
                (constructor
                  SM86InstructionBody
                  SM86Float64Binary
                  operation
                  destination
                  left
                  right
                  scaledUpdate64Control))
              tail))))))

-- Admitted pairs are even, so distinct bases imply disjoint word ranges.
def scaledUpdate64Disjoint =
  (lambda unrestricted left : (family SM86RegisterPair) .
    (lambda unrestricted right : (family SM86RegisterPair) .
      (nat-add
        (nat-less-than (sm86RegisterPairIndex left) (sm86RegisterPairIndex right))
        (nat-less-than (sm86RegisterPairIndex right) (sm86RegisterPairIndex left)))))

def accumulateScaledGradient64Admitted =
  (lambda unrestricted destination : (family SM86RegisterPair) .
    (lambda unrestricted decay : (family SM86RegisterPair) .
      (lambda unrestricted previous : (family SM86RegisterPair) .
        (lambda unrestricted gradient : (family SM86RegisterPair) .
          (nat-multiply
            (sm86Float64RegistersAdmitted destination decay previous)
            (nat-multiply
              (sm86RegisterPairAdmitted gradient)
              (scaledUpdate64Disjoint destination gradient)))))))

def accumulateScaledGradient64SM86 =
  (lambda unrestricted destination : (family SM86RegisterPair) .
    (lambda unrestricted decay : (family SM86RegisterPair) .
      (lambda unrestricted previous : (family SM86RegisterPair) .
        (lambda unrestricted gradient : (family SM86RegisterPair) .
          (lambda erased admitted : (equal Nat (accumulateScaledGradient64Admitted destination decay previous gradient) 1) .
            (lambda unrestricted tail : (family SM86Program) .
              (scaledUpdate64Next
                (constructor SM86Float64Operation SM86Float64Multiply)
                destination
                decay
                previous
                (scaledUpdate64Next
                  (constructor SM86Float64Operation SM86Float64Add)
                  destination
                  destination
                  gradient
                  tail))))))))

def subtractScaledUpdate64Admitted =
  (lambda unrestricted destination : (family SM86RegisterPair) .
    (lambda unrestricted rate : (family SM86RegisterPair) .
      (lambda unrestricted parameter : (family SM86RegisterPair) .
        (lambda unrestricted update : (family SM86RegisterPair) .
          (lambda unrestricted scratch : (family SM86RegisterPair) .
            (nat-multiply
              (sm86Float64RegistersAdmitted destination rate parameter)
              (nat-multiply
                (sm86Float64RegistersAdmitted scratch rate update)
                (scaledUpdate64Disjoint scratch parameter))))))))

def subtractScaledUpdate64SM86 =
  (lambda unrestricted destination : (family SM86RegisterPair) .
    (lambda unrestricted rate : (family SM86RegisterPair) .
      (lambda unrestricted parameter : (family SM86RegisterPair) .
        (lambda unrestricted update : (family SM86RegisterPair) .
          (lambda unrestricted scratch : (family SM86RegisterPair) .
            (lambda erased admitted : (equal Nat (subtractScaledUpdate64Admitted destination rate parameter update scratch) 1) .
              (lambda unrestricted tail : (family SM86Program) .
                (scaledUpdate64Next
                  (constructor SM86Float64Operation SM86Float64Multiply)
                  scratch
                  rate
                  update
                  (scaledUpdate64Next
                    (constructor SM86Float64Operation SM86Float64Subtract)
                    destination
                    parameter
                    scratch
                    tail)))))))))