module Realization.Nvidia.SM86.MemoryArithmeticSM86

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

family MemoryArithmeticOperation : Type 0
constructor MemoryNarrowedProduct
constructor MemoryScaledUpdate

end-family

-- Register allocation: input pairs first, then pointer pairs. No high word
-- is repurposed while a conversion can still consume it.
def memoryProductInputCount : Nat =
  2

def memoryProductPointerCount : Nat =
  (succ memoryProductInputCount)

def memoryProductValueBase : Nat =
  0

def memoryProductPointerBase : Nat =
  (memory64PairAt memoryProductValueBase memoryProductInputCount)

def memoryProductLeft : Nat =
  memoryProductValueBase

def memoryProductRight : Nat =
  (memory64PairAt memoryProductValueBase 1)

def memoryProductOutputPointer : Nat =
  (memory64PairAt memoryProductPointerBase memoryProductInputCount)

def memoryProductIdentity : Bytes =
  b"narrow-binary64-product-binary32-widen-v1"

def memoryProductProgram =
  (memory64Pointers
    memoryProductPointerCount
    memoryProductPointerBase
    (memory64Loads
      memoryProductInputCount
      memoryProductValueBase
      memoryProductPointerBase
      (memory64Next
        (constructor
          SM86InstructionBody
          SM86FloatConvert
          (constructor
            SM86FloatConversion
            SM86NarrowBinary64
            (memory64R memoryProductLeft)
            (memory64P memoryProductLeft))
          scaledUpdate64Control)
        (memory64Next
          (constructor
            SM86InstructionBody
            SM86FloatConvert
            (constructor
              SM86FloatConversion
              SM86NarrowBinary64
              (memory64R memoryProductRight)
              (memory64P memoryProductRight))
            scaledUpdate64Control)
          (memory64Next
            (constructor
              SM86InstructionBody
              SM86FloatMultiply
              (memory64R memoryProductLeft)
              (memory64R memoryProductLeft)
              (memory64R memoryProductRight)
              memory64Wait)
            (memory64Next
              (constructor
                SM86InstructionBody
                SM86FloatConvert
                (constructor
                  SM86FloatConversion
                  SM86WidenBinary32
                  (memory64P memoryProductLeft)
                  (memory64R memoryProductLeft))
                scaledUpdate64Control)
              (memory64Store memoryProductOutputPointer memoryProductLeft memory64Exit)))))))

-- This ABI binds five state slots; the arithmetic is a typed composition from
-- the learner owner. Register allocation, initial reads and final stores are
-- derived by the shared pointwise compiler, rather than restated by this kernel.
def memoryUpdateInputCount : Nat =
  5

def memoryUpdateAssignments =
  (scaledGradientAssignments
    StdU64
    memoryUpdateInputCount
    (stdFinFirst 4)
    (stdFinNextOf 4 (stdFinFirst 3))
    (stdFinNextOf 4 (stdFinNextOf 3 (stdFinFirst 2)))
    (stdFinNextOf 4 (stdFinNextOf 3 (stdFinNextOf 2 (stdFinFirst 1))))
    (stdFinNextOf 4 (stdFinNextOf 3 (stdFinNextOf 2 (stdFinNextOf 1 (stdFinFirst 0)))))
    (refl Nat 1))

def memoryUpdateIdentity : Bytes =
  b"accumulate-scaled-gradient-subtract-scaled-update-binary64-v1"

def memoryUpdateProgram =
  (memoryPointwiseProgram memoryUpdateInputCount memoryUpdateAssignments (refl Nat 1))

def memoryArithmeticRegions =
  (constructor
    NvidiaDeviceRegions
    NvidiaDeviceRegionsNext
    (nvidiaSM86ScalarRegion
      memoryProductIdentity
      memoryProductProgram
      b"binary64 memory arithmetic requires an SM86 realization")
    (constructor
      NvidiaDeviceRegions
      NvidiaDeviceRegionsNext
      (nvidiaSM86ScalarRegion
        memoryUpdateIdentity
        memoryUpdateProgram
        b"binary64 memory arithmetic requires an SM86 realization")
      (constructor NvidiaDeviceRegions NvidiaDeviceRegionsEnd)))

-- The pointer-list length determines the constant block extent. A pairing
-- supplies the same semantic argument list that the kernel ABI declares.
def memoryArithmeticArgumentCount =
  (lambda unrestricted operation : (family MemoryArithmeticOperation) .
    (eliminate
      MemoryArithmeticOperation
      (lambda unrestricted value : (family MemoryArithmeticOperation) . Nat)
      operation
      (branch MemoryNarrowedProduct . memoryProductPointerCount)
      (branch MemoryScaledUpdate . memoryUpdateInputCount)))

def memoryArithmeticIdentity =
  (lambda unrestricted operation : (family MemoryArithmeticOperation) .
    (eliminate
      MemoryArithmeticOperation
      (lambda unrestricted value : (family MemoryArithmeticOperation) . Bytes)
      operation
      (branch MemoryNarrowedProduct . memoryProductIdentity)
      (branch MemoryScaledUpdate . memoryUpdateIdentity)))

def memoryArithmeticLaunch =
  (lambda unrestricted operation : (family MemoryArithmeticOperation) .
    (lambda unrestricted programBase : Nat .
      (lambda unrestricted pointers : (family StdList Nat) .
        (lambda erased arity : (equal Nat (stdListLength Nat pointers) (memoryArithmeticArgumentCount operation)) .
          (nvidiaSM86PointerKernelLaunch
            (memoryArithmeticIdentity operation)
            memoryArithmeticRegions
            programBase
            pointers)))))