module PartitionedRegression.Build.CoordinateMemory

import PartitionedRegression.Coordinates
import Component.Pointwise
import Hardware.Nvidia.SM86.Command.WholeProgramPlan
import Realization.Nvidia.SM86.MemoryArithmeticSM86
import Std.Fin
import Realization.Nvidia.SM86.PointwiseMemorySM86
import Platform.Linux.Nvidia.ProgramPlan
import Platform.Linux.Nvidia.PlanHost
import Platform.Linux.Nvidia.PlanHostRequest
import Platform.RunPod.Nvidia.SM86
import Runtime.ArenaCertificate
import Model.Word64
import Model.Parameter
import Std.List
import Std.Natural

family CoordinateDevicePhase : Type 0
constructor CoordinateDeviceForward
constructor CoordinateDeviceBackward
constructor CoordinateDeviceUpdate

end-family

-- Relative regions of the device's persistent binary64 allocation. Their
-- addresses are supplied by the shared target plan after placement. The
-- parameter and accumulator are seeded once; only the residual changes on
-- the host between rounds. All arithmetic and intermediate stores are GPU work.
def cdmParameter : Nat =
  0

def cdmAccumulator : Nat =
  (naturalAdd cdmParameter coordinateWordBytes)

def cdmFeature : Nat =
  (naturalAdd cdmAccumulator coordinateWordBytes)

def cdmResidual : Nat =
  (naturalAdd cdmFeature coordinateWordBytes)

def cdmPartial : Nat =
  (naturalAdd cdmResidual coordinateWordBytes)

def cdmGradient : Nat =
  (naturalAdd cdmPartial coordinateWordBytes)

def cdmDecay : Nat =
  (naturalAdd cdmGradient coordinateWordBytes)

def cdmRate : Nat =
  (naturalAdd cdmDecay coordinateWordBytes)

def cdmBytes : Nat =
  (naturalAdd cdmRate coordinateWordBytes)

def cdmResident =
  (lambda unrestricted name : Bytes .
    (lambda unrestricted offset : Nat .
      (lambda unrestricted tail : (family ArenaResidents) .
        (constructor
          ArenaResidents
          ArenaResidentsNext
          (constructor
            ArenaResident
            ArenaResidentValue
            name
            offset
            coordinateWordBytes
            coordinateWordBytes)
          tail))))

def cdmResidents =
  (cdmResident
    b"parameter"
    cdmParameter
    (cdmResident
      b"accumulator"
      cdmAccumulator
      (cdmResident
        b"feature"
        cdmFeature
        (cdmResident
          b"residual"
          cdmResidual
          (cdmResident
            b"partial"
            cdmPartial
            (cdmResident
              b"gradient"
              cdmGradient
              (cdmResident
                b"decay"
                cdmDecay
                (cdmResident b"rate" cdmRate (constructor ArenaResidents ArenaResidentsEnd)))))))))

def cdmCertified : (equal Nat (arenaCertificate cdmBytes cdmResidents) 1) =
  (refl Nat 1)

def cdmPhases =
  (constructor
    StdList
    StdListCons
    (family CoordinateDevicePhase)
    (constructor CoordinateDevicePhase CoordinateDeviceForward)
    (constructor
      StdList
      StdListCons
      (family CoordinateDevicePhase)
      (constructor CoordinateDevicePhase CoordinateDeviceBackward)
      (constructor
        StdList
        StdListCons
        (family CoordinateDevicePhase)
        (constructor CoordinateDevicePhase CoordinateDeviceUpdate)
        (constructor StdList StdListEmpty (family CoordinateDevicePhase)))))

def cdmForwardOrdinal : Nat =
  0

def cdmBackwardOrdinal : Nat =
  (succ cdmForwardOrdinal)

def cdmUpdateOrdinal : Nat =
  (succ cdmBackwardOrdinal)

def cdmOrdinal =
  (lambda unrestricted phase : (family CoordinateDevicePhase) .
    (eliminate
      CoordinateDevicePhase
      (lambda unrestricted value : (family CoordinateDevicePhase) . Nat)
      phase
      (branch CoordinateDeviceForward . cdmForwardOrdinal)
      (branch CoordinateDeviceBackward . cdmBackwardOrdinal)
      (branch CoordinateDeviceUpdate . cdmUpdateOrdinal)))

def cdmPointer =
  (lambda unrestricted base : Nat .
    (lambda unrestricted offset : Nat .
      (lambda unrestricted tail : (family StdList Nat) .
        (constructor StdList StdListCons Nat (naturalAdd base offset) tail))))

def cdmLaunch =
  (lambda unrestricted program : Nat .
    (lambda unrestricted data : Nat .
      (lambda unrestricted phase : (family CoordinateDevicePhase) .
        (let unrestricted pointer =
          (cdmPointer data)
          in
          (let unrestricted nil =
            (constructor StdList StdListEmpty Nat)
            in
            (eliminate
              CoordinateDevicePhase
              (lambda unrestricted current : (family CoordinateDevicePhase) .
                (family NvidiaLaunchTemplate))
              phase
              (branch
                CoordinateDeviceForward
                .
                (memoryArithmeticLaunch
                  (constructor MemoryArithmeticOperation MemoryNarrowedProduct)
                  program
                  (pointer cdmParameter (pointer cdmFeature (pointer cdmPartial nil)))
                  (refl Nat 3)))
              (branch
                CoordinateDeviceBackward
                .
                (memoryArithmeticLaunch
                  (constructor MemoryArithmeticOperation MemoryNarrowedProduct)
                  program
                  (pointer cdmResidual (pointer cdmFeature (pointer cdmGradient nil)))
                  (refl Nat 3)))
              (branch
                CoordinateDeviceUpdate
                .
                (memoryArithmeticLaunch
                  (constructor MemoryArithmeticOperation MemoryScaledUpdate)
                  program
                  (pointer
                    cdmParameter
                    (pointer
                      cdmAccumulator
                      (pointer cdmGradient (pointer cdmDecay (pointer cdmRate nil)))))
                  (refl Nat 5)))))))))

def cdmLaunches =
  (lambda unrestricted program : Nat .
    (lambda unrestricted data : Nat .
      (nvidiaLaunchSequence
        (stdListMap
          (family CoordinateDevicePhase)
          (family NvidiaLaunchTemplate)
          (cdmLaunch program data)
          cdmPhases))))

def cdmBatch =
  (lambda unrestricted phase : (family CoordinateDevicePhase) .
    (constructor
      NvidiaSubmissionSchedule
      NvidiaSubmissionScheduleOne
      (constructor
        NvidiaSubmissionBatch
        NvidiaSubmissionBatchValue
        b"coordinate-phase"
        (naturalMultiply (cdmOrdinal phase) prSemaphoreStride)
        (constructor NvidiaLaunchReferences NvidiaLaunchReferencesRange (cdmOrdinal phase) 1))))

def cdmSubmissions =
  (stdListFold
    (family CoordinateDevicePhase)
    (family NvidiaSubmissionSchedule)
    (lambda unrestricted phase : (family CoordinateDevicePhase) .
      (lambda unrestricted tail : (family NvidiaSubmissionSchedule) .
        (constructor NvidiaSubmissionSchedule NvidiaSubmissionScheduleAppend (cdmBatch phase) tail)))
    (constructor NvidiaSubmissionSchedule NvidiaSubmissionScheduleEmpty)
    cdmPhases)

def cdmIssueCount : Nat =
  (nvidiaSubmissionScheduleCount cdmSubmissions)

def cdmPlan =
  (lambda unrestricted profile : (family RunPodSM86TargetProfile) .
    (nvidiaProgramPlan
      memoryArithmeticRegions
      cdmLaunches
      cdmSubmissions
      cdmBytes
      (constructor NvidiaPlanHostMemory NvidiaPlanHostVideoMemory)
      runPodSM86QMDArchitecture
      runPodSM86ComputeClass
      runPodSM86SPAVersion
      runPodSM86QMDSlotBytes
      (runPodSM86ProductIdentityPadded profile)
      nvidiaPlanHostAmpere580
      (modelWord64Natural (runPodSM86TotalVideoMemory profile))))