module Realization.Nvidia.SM86.PointwiseMemorySM86

import Accelerator.SM86.Control
import Accelerator.SM86.Immediate
import Accelerator.SM86.Instruction
import Accelerator.SM86.InstructionEncoding
import Accelerator.SM86.Operands
import Accelerator.SM86.Program
import Accelerator.SM86.Types
import Component.Pointwise
import Hardware.Nvidia.SM86.Command.WholeProgramPlan
import Model.Word64
import Realization.Nvidia.SM86.ScaledGradientUpdateSM86
import Std.Fin
import Std.List
import Std.Natural
import Std.Word

family MemoryPointwiseFragment : Type 0
constructor MemoryPointwiseFragmentOf
field unrestricted memoryPointwiseFragmentProgram : (family SM86Program)
field unrestricted memoryPointwiseFragmentResult : Nat
field unrestricted memoryPointwiseFragmentEnd : Nat

end-family

-- Scalar memory realizations used by persistent state schedules. Each region
-- launches one thread. Addresses are CB0 pointer arguments supplied by the
-- enclosing plan; this module neither chooses device addresses nor builds a
-- launch table. A submission fence makes stores visible to the next phase.
-- Arithmetic results stay in the caller's device allocation between launches.
def memory64WordBytes : Nat =
  4

def memory64PairWords : Nat =
  2

def memory64R =
  (lambda unrestricted index : Nat . (sm86Register (nat-to-byte index)))

def memory64P =
  (lambda unrestricted index : Nat . (sm86RegisterPair (nat-to-byte index)))

def memory64U =
  sm86Unsigned32FromNaturalTruncated

def memory64PairAt =
  (lambda unrestricted base : Nat .
    (lambda unrestricted index : Nat . (nat-add base (nat-multiply memory64PairWords index))))

def memory64Next =
  (lambda unrestricted body : (family SM86InstructionBody) .
    (lambda unrestricted tail : (family SM86Program) .
      (constructor SM86Program SM86ProgramNext (sm86Instruction body) tail)))

def memory64Wait =
  (constructor
    SM86Control
    SM86ControlValue
    (byte 15)
    (constructor SM86YieldMode SM86Continue)
    (constructor SM86Barrier SM86BarrierNone)
    (constructor SM86Barrier SM86BarrierNone)
    (byte 3)
    (byte 0))

def memory64StoreControl =
  (constructor
    SM86Control
    SM86ControlValue
    (byte 15)
    (constructor SM86YieldMode SM86Continue)
    (constructor SM86Barrier SM86BarrierNone)
    (constructor SM86Barrier SM86Barrier1)
    (byte 3)
    (byte 0))

def memory64Exit =
  (memory64Next
    (constructor SM86InstructionBody SM86Exit memory64Wait)
    (constructor SM86Program SM86ProgramEnd))

-- Both words are loaded with the same source-address pair. Every successor
-- waits for late reads and writes, including before a pointer is reused.
def memory64Load =
  (lambda unrestricted destination : Nat .
    (lambda unrestricted pointer : Nat .
      (lambda unrestricted tail : (family SM86Program) .
        (memory64Next
          (constructor
            SM86InstructionBody
            SM86LoadGlobal
            (memory64R destination)
            (memory64R pointer)
            sm86Unsigned32Zero
            scaledUpdate64Control)
          (memory64Next
            (constructor
              SM86InstructionBody
              SM86LoadGlobal
              (memory64R (succ destination))
              (memory64R pointer)
              (memory64U memory64WordBytes)
              scaledUpdate64Control)
            tail)))))

def memory64Store =
  (lambda unrestricted pointer : Nat .
    (lambda unrestricted source : Nat .
      (lambda unrestricted tail : (family SM86Program) .
        (memory64Next
          (constructor
            SM86InstructionBody
            SM86StoreGlobal64
            (memory64R pointer)
            (memory64R source)
            sm86Unsigned32Zero
            memory64StoreControl)
          tail))))

def memory64Pointer =
  (lambda unrestricted destination : Nat .
    (lambda unrestricted argument : Nat .
      (lambda unrestricted tail : (family SM86Program) .
        (memory64Next
          (constructor
            SM86InstructionBody
            SM86MoveConstant
            (memory64R destination)
            (byte 0)
            (memory64U (nvidiaKernelPointerByteOffset argument))
            memory64Wait)
          (memory64Next
            (constructor
              SM86InstructionBody
              SM86MoveConstant
              (memory64R (succ destination))
              (byte 0)
              (memory64U (nat-add (nvidiaKernelPointerByteOffset argument) memory64WordBytes))
              memory64Wait)
            tail)))))

def memory64Pointers =
  (lambda unrestricted count : Nat .
    (lambda unrestricted base : Nat .
      (lambda unrestricted tail : (family SM86Program) .
        (nat-eliminate
          (lambda unrestricted n : Nat . (family SM86Program))
          tail
          (lambda unrestricted index : Nat .
            (lambda unrestricted rest : (family SM86Program) .
              (memory64Pointer (memory64PairAt base index) index rest)))
          count))))

def memory64Loads =
  (lambda unrestricted count : Nat .
    (lambda unrestricted base : Nat .
      (lambda unrestricted pointers : Nat .
        (lambda unrestricted tail : (family SM86Program) .
          (nat-eliminate
            (lambda unrestricted n : Nat . (family SM86Program))
            tail
            (lambda unrestricted index : Nat .
              (lambda unrestricted rest : (family SM86Program) .
                (memory64Load (memory64PairAt base index) (memory64PairAt pointers index) rest)))
            count)))))

-- Each assignment gets a temporary register range above the persistent value
-- pairs. That range is reused only after the result has been copied back and
-- all late reads/writes have retired. Pointer pairs follow the maximum range,
-- so expression depth cannot silently collide with an address register.
def memoryPointwiseEnd =
  (constructor SM86Program SM86ProgramEnd)

def memoryPointwiseCopyWord =
  (lambda unrestricted destination : Nat .
    (lambda unrestricted source : Nat .
      (lambda unrestricted tail : (family SM86Program) .
        (memory64Next
          (constructor
            SM86InstructionBody
            SM86IntegerMultiplyAddImmediate
            (memory64R destination)
            (memory64R source)
            sm86Unsigned32One
            sm86ZeroRegister
            memory64Wait)
          tail))))

def memoryPointwiseCopy =
  (lambda unrestricted destination : Nat .
    (lambda unrestricted source : Nat .
      (lambda unrestricted tail : (family SM86Program) .
        (memoryPointwiseCopyWord
          destination
          source
          (memoryPointwiseCopyWord (succ destination) (succ source) tail)))))

def memoryPointwiseImmediate =
  (lambda unrestricted destination : Nat .
    (lambda unrestricted value : Nat .
      (lambda unrestricted tail : (family SM86Program) .
        (memory64Next
          (constructor
            SM86InstructionBody
            SM86MoveImmediate
            (memory64R destination)
            (memory64U value)
            memory64Wait)
          tail))))

def memoryPointwiseUnarySupported =
  (lambda unrestricted operation : (family PointwiseUnary) .
    (eliminate
      PointwiseUnary
      (lambda unrestricted current : (family PointwiseUnary) . Nat)
      operation
      (branch PointwiseNegate . 1)
      (branch PointwiseSquare . 1)
      (branch PointwiseRoundBF16 . 0)
      (branch PointwiseReciprocal . 0)
      (branch PointwiseReciprocalSquareRoot . 0)
      (branch PointwiseSquareRoot . 0)
      (branch PointwiseTanh . 0)
      (branch PointwiseExponentialBase2 . 0)))

def memoryPointwiseBinarySupported =
  (lambda unrestricted operation : (family PointwiseBinary) .
    (eliminate
      PointwiseBinary
      (lambda unrestricted current : (family PointwiseBinary) . Nat)
      operation
      (branch PointwiseAdd . 1)
      (branch PointwiseMultiply . 1)
      (branch PointwiseSubtract . 1)
      (branch PointwisePositiveGate . 0)))

def memoryPointwiseExpressionSupported =
  (lambda erased slots : Nat .
    (lambda unrestricted expression : (family PointwiseExpression StdU64 slots) .
      (pointwiseFold
        StdU64
        Nat
        slots
        (lambda unrestricted operation : (family PointwiseUnary) .
          (lambda unrestricted child : Nat .
            (naturalAnd (memoryPointwiseUnarySupported operation) child)))
        (lambda unrestricted operation : (family PointwiseBinary) .
          (lambda unrestricted left : Nat .
            (lambda unrestricted right : Nat .
              (naturalAnd (memoryPointwiseBinarySupported operation) (naturalAnd left right)))))
        (lambda unrestricted bits : StdU64 . 1)
        expression
        (lambda unrestricted index : (family StdFin slots) . 1))))

-- Unsupported branches below are inaccessible through memoryPointwiseProgram's
-- admission proof. They contain no host or alternative arithmetic realization.
-- Adding a branch requires its numeric contract and checked-model laws first.
def memoryPointwiseUnaryProgram =
  (lambda unrestricted operation : (family PointwiseUnary) .
    (lambda unrestricted destination : Nat .
      (lambda unrestricted source : Nat .
        (eliminate
          PointwiseUnary
          (lambda unrestricted current : (family PointwiseUnary) . (family SM86Program))
          operation
          (branch
            PointwiseNegate
            .
            (memoryPointwiseCopyWord
              destination
              source
              (memoryPointwiseImmediate
                (succ destination)
                0x80000000
                (memory64Next
                  (constructor
                    SM86InstructionBody
                    SM86LogicThreeInputTruthTable
                    (memory64R (succ destination))
                    (memory64R (succ source))
                    (memory64R (succ destination))
                    (byte 0x3c)
                    memory64Wait)
                  memoryPointwiseEnd))))
          (branch
            PointwiseSquare
            .
            (scaledUpdate64Next
              (constructor SM86Float64Operation SM86Float64Multiply)
              (memory64P destination)
              (memory64P source)
              (memory64P source)
              memoryPointwiseEnd))
          (branch PointwiseRoundBF16 . memoryPointwiseEnd)
          (branch PointwiseReciprocal . memoryPointwiseEnd)
          (branch PointwiseReciprocalSquareRoot . memoryPointwiseEnd)
          (branch PointwiseSquareRoot . memoryPointwiseEnd)
          (branch PointwiseTanh . memoryPointwiseEnd)
          (branch PointwiseExponentialBase2 . memoryPointwiseEnd)))))

def memoryPointwiseBinaryProgram =
  (lambda unrestricted operation : (family PointwiseBinary) .
    (lambda unrestricted destination : Nat .
      (lambda unrestricted left : Nat .
        (lambda unrestricted right : Nat .
          (let unrestricted emit =
            (lambda unrestricted numeric : (family SM86Float64Operation) .
              (scaledUpdate64Next
                numeric
                (memory64P destination)
                (memory64P left)
                (memory64P right)
                memoryPointwiseEnd))
            in
            (eliminate
              PointwiseBinary
              (lambda unrestricted current : (family PointwiseBinary) . (family SM86Program))
              operation
              (branch PointwiseAdd . (emit (constructor SM86Float64Operation SM86Float64Add)))
              (branch
                PointwiseMultiply
                .
                (emit (constructor SM86Float64Operation SM86Float64Multiply)))
              (branch
                PointwiseSubtract
                .
                (emit (constructor SM86Float64Operation SM86Float64Subtract)))
              (branch PointwisePositiveGate . memoryPointwiseEnd)))))))

-- Observe each value at its production/use boundary, before a subsequent
-- assignment can overwrite it. The IEEE realization supplies an empty observer;
-- stricter numerical policies reuse this traversal without changing its order.
def memoryPointwiseCompile =
  (lambda unrestricted observe : (pi unrestricted result : Nat . (family SM86Program)) .
    (lambda unrestricted slots : Nat .
      (lambda unrestricted expression : (family PointwiseExpression StdU64 slots) .
        (eliminate
          PointwiseExpression
          (lambda unrestricted current : (family PointwiseExpression StdU64 slots) .
            (pi unrestricted next : Nat . (family MemoryPointwiseFragment)))
          expression
          (branch
            PointwiseInput
            index
            .
            (lambda unrestricted next : Nat .
              (constructor
                MemoryPointwiseFragment
                MemoryPointwiseFragmentOf
                (observe (memory64PairAt 0 (stdFinToNatural slots index)))
                (memory64PairAt 0 (stdFinToNatural slots index))
                next)))
          (branch
            PointwiseConstant
            bits
            .
            (lambda unrestricted next : Nat .
              (let unrestricted value =
                (modelWord64Natural bits)
                in
                (constructor
                  MemoryPointwiseFragment
                  MemoryPointwiseFragmentOf
                  (memoryPointwiseImmediate
                    next
                    (naturalModuloUnchecked value 0x100000000)
                    (memoryPointwiseImmediate
                      (succ next)
                      (naturalDivideUnchecked value 0x100000000)
                      (observe next)))
                  next
                  (naturalAdd next memory64PairWords)))))
          (branch
            PointwiseUnaryExpression
            operation
            operand
            compile
            .
            (lambda unrestricted next : Nat .
              (eliminate
                MemoryPointwiseFragment
                (lambda unrestricted current : (family MemoryPointwiseFragment) .
                  (family MemoryPointwiseFragment))
                (compile next)
                (branch
                  MemoryPointwiseFragmentOf
                  program
                  result
                  end
                  .
                  (constructor
                    MemoryPointwiseFragment
                    MemoryPointwiseFragmentOf
                    (sm86ProgramAppend
                      program
                      (sm86ProgramAppend
                        (memoryPointwiseUnaryProgram operation end result)
                        (observe end)))
                    end
                    (naturalAdd end memory64PairWords))))))
          (branch
            PointwiseBinaryExpression
            operation
            left
            right
            compileLeft
            compileRight
            .
            (lambda unrestricted next : Nat .
              (eliminate
                MemoryPointwiseFragment
                (lambda unrestricted current : (family MemoryPointwiseFragment) .
                  (family MemoryPointwiseFragment))
                (compileLeft next)
                (branch
                  MemoryPointwiseFragmentOf
                  leftProgram
                  leftResult
                  leftEnd
                  .
                  (eliminate
                    MemoryPointwiseFragment
                    (lambda unrestricted current : (family MemoryPointwiseFragment) .
                      (family MemoryPointwiseFragment))
                    (compileRight leftEnd)
                    (branch
                      MemoryPointwiseFragmentOf
                      rightProgram
                      rightResult
                      rightEnd
                      .
                      (constructor
                        MemoryPointwiseFragment
                        MemoryPointwiseFragmentOf
                        (sm86ProgramAppend
                          leftProgram
                          (sm86ProgramAppend
                            rightProgram
                            (sm86ProgramAppend
                              (memoryPointwiseBinaryProgram
                                operation
                                rightEnd
                                leftResult
                                rightResult)
                              (observe rightEnd))))
                        rightEnd
                        (naturalAdd rightEnd memory64PairWords))))))))))))

def memoryPointwiseValueEnd =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
      (memory64PairAt 0 (naturalAdd slots (pointwiseProgramTemporaryValues StdU64 slots program)))))

def memoryPointwiseSupported =
  (lambda erased slots : Nat .
    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
      (eliminate
        StdList
        (lambda unrestricted current : (family StdList (family PointwiseAssignment StdU64 slots)) .
          Nat)
        program
        (branch StdListEmpty . 1)
        (branch
          StdListCons
          head
          tail
          rest
          .
          (eliminate
            PointwiseAssignment
            (lambda unrestricted current : (family PointwiseAssignment StdU64 slots) . Nat)
            head
            (branch
              PointwiseAssign
              destination
              expression
              .
              (naturalAnd (memoryPointwiseExpressionSupported slots expression) rest)))))))

def memoryPointwiseAdmitted =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
      (naturalAnd
        (memoryPointwiseSupported slots program)
        (sm86RegisterSpanAdmitted (memory64PairAt (memoryPointwiseValueEnd slots program) slots)))))

def memoryPointwiseAssignments =
  (lambda unrestricted observe : (pi unrestricted result : Nat . (family SM86Program)) .
    (lambda unrestricted slots : Nat .
      (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
        (lambda unrestricted tail : (family SM86Program) .
          (eliminate
            StdList
            (lambda unrestricted current : (family StdList (family PointwiseAssignment StdU64 slots)) .
              (family SM86Program))
            program
            (branch StdListEmpty . tail)
            (branch
              StdListCons
              head
              rest
              next
              .
              (eliminate
                PointwiseAssignment
                (lambda unrestricted current : (family PointwiseAssignment StdU64 slots) .
                  (family SM86Program))
                head
                (branch
                  PointwiseAssign
                  destination
                  expression
                  .
                  (eliminate
                    MemoryPointwiseFragment
                    (lambda unrestricted current : (family MemoryPointwiseFragment) .
                      (family SM86Program))
                    (memoryPointwiseCompile observe slots expression (memory64PairAt 0 slots))
                    (branch
                      MemoryPointwiseFragmentOf
                      code
                      result
                      end
                      .
                      (sm86ProgramAppend
                        code
                        (memoryPointwiseCopy
                          (memory64PairAt 0 (stdFinToNatural slots destination))
                          result
                          next))))))))))))

-- The use analysis preserves the distinction between an initialized input and
-- a slot first defined by this program. An output-only slot is never loaded.
-- Stores happen after all assignments, so every expression observes the same
-- sequential register state and no later read sees a half-published update.
def memoryPointwiseProgramWith =
  (lambda unrestricted observe : (pi unrestricted result : Nat . (family SM86Program)) .
    (lambda unrestricted store : (pi unrestricted pointer : Nat . (pi unrestricted source : Nat . (pi unrestricted tail : (family SM86Program) . (family SM86Program)))) .
      (lambda unrestricted finish : (family SM86Program) .
        (lambda unrestricted slots : Nat .
          (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
            (let unrestricted pointers =
              (memoryPointwiseValueEnd slots program)
              in
              (let unrestricted stores =
                (nat-eliminate
                  (lambda unrestricted count : Nat . (family SM86Program))
                  finish
                  (lambda unrestricted index : Nat .
                    (lambda unrestricted next : (family SM86Program) .
                      (nat-eliminate
                        (lambda unrestricted needed : Nat . (family SM86Program))
                        next
                        (lambda unrestricted previous : Nat .
                          (lambda unrestricted ignored : (family SM86Program) .
                            (store (memory64PairAt pointers index) (memory64PairAt 0 index) next)))
                        (pointwiseProgramWrites StdU64 slots program index))))
                  slots)
                in
                (let unrestricted body =
                  (memoryPointwiseAssignments observe slots program stores)
                  in
                  (nat-eliminate
                    (lambda unrestricted count : Nat . (family SM86Program))
                    body
                    (lambda unrestricted index : Nat .
                      (lambda unrestricted next : (family SM86Program) .
                        (let unrestricted read =
                          (pointwiseProgramReadsInitial StdU64 slots program index)
                          in
                          (let unrestricted write =
                            (pointwiseProgramWrites StdU64 slots program index)
                            in
                            (nat-eliminate
                              (lambda unrestricted needed : Nat . (family SM86Program))
                              next
                              (lambda unrestricted previous : Nat .
                                (lambda unrestricted ignored : (family SM86Program) .
                                  (memory64Pointer
                                    (memory64PairAt pointers index)
                                    index
                                    (nat-eliminate
                                      (lambda unrestricted initial : Nat . (family SM86Program))
                                      next
                                      (lambda unrestricted before : Nat .
                                        (lambda unrestricted unused : (family SM86Program) .
                                        (memory64Load
                                        (memory64PairAt 0 index)
                                        (memory64PairAt pointers index)
                                        next)))
                                      read))))
                              (naturalOr read write))))))
                    slots)))))))))

-- Existing IEEE consumers retain precisely the same instruction stream.
def memoryPointwiseProgram =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
      (lambda erased admitted : (equal Nat (memoryPointwiseAdmitted slots program) 1) .
        (memoryPointwiseProgramWith
          (lambda unrestricted result : Nat . memoryPointwiseEnd)
          memory64Store
          memory64Exit
          slots
          program))))