module Coppelius.Build.NativeHost

import Coppelius.ArenaPlan
import Coppelius.Build.DeviceImages
import Coppelius.Build.PhysicalComponents
import Coppelius.Build.Placement
import Coppelius.Build.Graph
import Coppelius.Learner
import Coppelius.TrainingRun
import Data.Bytes
import Hardware.Nvidia.SM86.Command.WholeProgramPlan
import Model.Config
import Model.Parameter
import Model.Word64
import Model.Word32
import Platform.Linux.Nvidia.Compatibility
import Platform.Linux.Nvidia.PlanHost
import Platform.Linux.Nvidia.PlanHostAdamW
import Platform.Linux.Nvidia.PlanHostRequest
import Runtime.NativePhysicalEmbeddedArtifact
import Realization.Nvidia.SM86.AdamWHalfSM86
import Runtime.NativePhysicalProgram
import Std.List
import Std.Natural

-- what the realization reports: each launch's parameter block's address,
-- the word at the AdamW step size's offset in it, and the GPFIFO table --
-- each submission's ring entry, which the host writes into the ring
-- itself when it issues the submission again
-- (Coppelius.Build.DirectArtifact passes them)
family CoppeliusRealizedTables : Type 0
constructor CoppeliusRealizedTablesValue
field unrestricted coppeliusRealizedParameterAddresses : Bytes
field unrestricted coppeliusRealizedStepWords : Bytes
field unrestricted coppeliusRealizedRingEntries : Bytes
end-family

-- The card a host drives: its plan-host ABI, and the memory its channel
-- (GPFIFO and USERD) and its arena live in -- video memory on the RTX
-- cards; on the DGX Spark's GB10, which has no frame buffer (a video
-- allocation is refused with NV_ERR_NOT_SUPPORTED), system memory, paged
-- for the arena (1.34 GB is no contiguous allocation).
family CoppeliusHostTarget : Type 0
constructor CoppeliusHostTargetValue
field unrestricted coppeliusHostTargetABI : (family NvidiaPlanHostABI)
field unrestricted coppeliusHostTargetChannelMemory : (family NvidiaPlanHostMemory)
field unrestricted coppeliusHostTargetArenaMemory : (family NvidiaPlanHostMemory)
end-family

-- a run of checkpoint gathers issued by one repeat (coppeliusGathers):
-- its first chunk and how many
family CoppeliusGatherRun : Type 0
constructor CoppeliusGatherRunValue
field unrestricted coppeliusGatherRunFirst : Nat
field unrestricted coppeliusGatherRunLength : Nat
end-family

-- Coppelius's host, derived from its plans by Platform.Linux.Nvidia.PlanHost:
-- the UVM lifecycle over the arena plan's regions, the four launch-table
-- recipes expanded at startup, and the request schedule below -- the
-- checkpoint scattered into the video arena two submissions per staged
-- chunk, the forward submission and its predictions, then the run policy
-- (Coppelius.TrainingRun): every step's context read from the token
-- stream, its AdamW scalars computed, the update issued; the checkpoint
-- gathered back one submission per chunk and published after every
-- interval; the result record.  Every number here is the arena plan's, the training run's
-- or the submission plan's; the host addresses are this pairing's placement
-- of its mappings in the process.

-- ---- host placement (Coppelius.ArenaPlan and Coppelius.Build.Placement
-- place the mappings) ----
def coppeliusHostQMDOverflow : Nat = (naturalAdd coppeliusHostQMD coppeliusSM86CompatQMDPrimaryBytesNatural)
def coppeliusQMDOverflowBytes : Nat =
  (naturalSaturatingSubtract coppeliusSM86CompatQMDBytesNatural coppeliusSM86CompatQMDPrimaryBytesNatural)

def coppeliusSystemBuffer =
  (lambda unrestricted identity : Bytes .
    (lambda unrestricted gpu : Nat .
      (lambda unrestricted host : Nat .
        (lambda unrestricted extent : Nat .
          (constructor NvidiaPlanHostBuffer NvidiaPlanHostBufferValue identity gpu host extent
            (constructor NvidiaPlanHostMemory NvidiaPlanHostSystemMemory))))))


def coppeliusHostAmpere : (family CoppeliusHostTarget) =
  (constructor CoppeliusHostTarget CoppeliusHostTargetValue nvidiaPlanHostAmpere580
    (constructor NvidiaPlanHostMemory NvidiaPlanHostVideoMemory)
    (constructor NvidiaPlanHostMemory NvidiaPlanHostVideoMemory))

-- The GB10's arena is GPU-cached system memory (the
-- allocation cudaMalloc makes there, L2-cacheable by the GPU; the paged
-- system memory it had ran L2-resident copies 4.6 times slower); it is
-- never host-mapped.
def coppeliusHostGB10 : (family CoppeliusHostTarget) =
  (constructor CoppeliusHostTarget CoppeliusHostTargetValue nvidiaPlanHostBlackwell580
    (constructor NvidiaPlanHostMemory NvidiaPlanHostSystemMemory)
    (constructor NvidiaPlanHostMemory NvidiaPlanHostGPUCachedSystemMemory))

def coppeliusHostTargetABI =
  (lambda unrestricted target : (family CoppeliusHostTarget) .
    (eliminate CoppeliusHostTarget (lambda unrestricted current : (family CoppeliusHostTarget) . (family NvidiaPlanHostABI)) target
      (branch CoppeliusHostTargetValue abi channel arena . abi)))

def coppeliusHostTargetChannelMemory =
  (lambda unrestricted target : (family CoppeliusHostTarget) .
    (eliminate CoppeliusHostTarget (lambda unrestricted current : (family CoppeliusHostTarget) . (family NvidiaPlanHostMemory)) target
      (branch CoppeliusHostTargetValue abi channel arena . channel)))

def coppeliusHostTargetArenaMemory =
  (lambda unrestricted target : (family CoppeliusHostTarget) .
    (eliminate CoppeliusHostTarget (lambda unrestricted current : (family CoppeliusHostTarget) . (family NvidiaPlanHostMemory)) target
      (branch CoppeliusHostTargetValue abi channel arena . arena)))

def coppeliusVideoBuffer =
  (lambda unrestricted memory : (family NvidiaPlanHostMemory) .
  (lambda unrestricted identity : Bytes .
    (lambda unrestricted gpu : Nat .
      (lambda unrestricted host : Nat .
        (lambda unrestricted extent : Nat .
          (constructor NvidiaPlanHostBuffer NvidiaPlanHostBufferValue identity gpu host extent memory))))))

-- The layout: the GPFIFO and its USERD page in one video buffer, the
-- pushbuffer (its extent the expansion's, `pushExpanded`), semaphores,
-- program and QMD tables in system memory at the plan's bases, the
-- checkpoint staging buffer, the QMD overflow adjacent to the QMD, and the
-- video arena itself inside the UVM external range (not host-mapped).  The
-- pairing's card name is asserted.
def coppeliusHostLayoutFor =
  (lambda unrestricted target : (family CoppeliusHostTarget) .
  (lambda unrestricted pushExpanded : Nat .
  (lambda unrestricted expectedName : Bytes .
    (constructor NvidiaPlanHostLayout NvidiaPlanHostLayoutValue
      (coppeliusVideoBuffer (coppeliusHostTargetChannelMemory target) b"gpfifo+userd" coppeliusGPFIFOBase coppeliusHostGPFIFO coppeliusSM86CompatGPFIFOAndUSERDBytesNatural)
      (coppeliusSystemBuffer b"pushbuffer" coppeliusPushbufferBase coppeliusHostPushbuffer (coppeliusPushbufferBytes pushExpanded))
      (coppeliusSystemBuffer b"semaphores" coppeliusSemaphoreBase coppeliusHostSemaphores coppeliusSemaphoreBytes)
      (coppeliusSystemBuffer b"program" coppeliusProgramBase coppeliusHostProgram coppeliusSM86CompatProgramBytesNatural)
      (coppeliusSystemBuffer b"qmd" coppeliusQMDBase coppeliusHostQMD coppeliusSM86CompatQMDPrimaryBytesNatural)
      (coppeliusSystemBuffer b"userd-alias" (naturalAdd coppeliusGPFIFOBase nvidiaPlanHostUVMGPFIFOBytes) (naturalAdd coppeliusHostGPFIFO nvidiaPlanHostUVMGPFIFOBytes) nvidiaPlanHostUSERDBytes)
      (constructor NvidiaPlanHostBuffers NvidiaPlanHostBuffersNext
        (coppeliusSystemBuffer b"checkpoint-staging" coppeliusStagingBase coppeliusHostStaging coppeliusSM86CompatCheckpointStagingBytesNatural)
        (constructor NvidiaPlanHostBuffers NvidiaPlanHostBuffersNext
          (coppeliusSystemBuffer b"qmd-overflow" (naturalAdd coppeliusQMDBase coppeliusSM86CompatQMDPrimaryBytesNatural) coppeliusHostQMDOverflow coppeliusQMDOverflowBytes)
          (constructor NvidiaPlanHostBuffers NvidiaPlanHostBuffersNext
            (coppeliusVideoBuffer (coppeliusHostTargetArenaMemory target) b"video-arena" coppeliusSM86CompatVideoBaseNatural 0 coppeliusSM86CompatVideoBytesNatural)
            (constructor NvidiaPlanHostBuffers NvidiaPlanHostBuffersEnd))))
      (coppeliusHostTargetABI target)
      (constructor NvidiaPlanHostLifecycle NvidiaPlanHostUVM coppeliusSM86CompatVideoBaseNatural coppeliusSM86CompatVideoBytesNatural)
      expectedName
      coppeliusHostErrorNotifier))))

-- the RTX pairings' layout (the arena certificate and the admissions below
-- decide it; the GB10's differs only in memory kinds and ABI)
def coppeliusHostLayout = (coppeliusHostLayoutFor coppeliusHostAmpere)

-- the four recipes and the extents their expansions must have
def coppeliusRecipeFills =
  (lambda unrestricted qmdExpanded : Nat .
    (lambda unrestricted pushExpanded : Nat .
      (constructor NvidiaPlanHostRecipeFills NvidiaPlanHostRecipeFillsNext coppeliusHostProgram coppeliusSM86CompatProgramBytesNatural
        (constructor NativePhysicalEmbeddedComponentKind NativePhysicalEmbeddedProgramTable)
        (constructor NvidiaPlanHostRecipeFills NvidiaPlanHostRecipeFillsNext coppeliusHostQMD qmdExpanded
          (constructor NativePhysicalEmbeddedComponentKind NativePhysicalEmbeddedQMDTable)
          (constructor NvidiaPlanHostRecipeFills NvidiaPlanHostRecipeFillsNext coppeliusHostPushbuffer pushExpanded
            (constructor NativePhysicalEmbeddedComponentKind NativePhysicalEmbeddedPushbuffer)
            (constructor NvidiaPlanHostRecipeFills NvidiaPlanHostRecipeFillsNext coppeliusHostGPFIFO coppeliusGPFIFOTableBytes
              (constructor NativePhysicalEmbeddedComponentKind NativePhysicalEmbeddedGPFIFO)
              (constructor NvidiaPlanHostRecipeFills NvidiaPlanHostRecipeFillsEnd)))))))

-- ---- the schedule ----
-- The checkpoint moves through the staging buffer a chunk at a time: the
-- full chunks, then the tail.
def coppeliusChunkBytes : Nat = coppeliusSM86CompatCheckpointStagingBytesNatural
def coppeliusFullChunks : Nat = (naturalDivideUnchecked coppeliusCheckpointPayloadBytesNatural coppeliusChunkBytes)
def coppeliusTailBytes : Nat = (naturalModuloUnchecked coppeliusCheckpointPayloadBytesNatural coppeliusChunkBytes)
def coppeliusTailOffset : Nat = (naturalMultiply coppeliusFullChunks coppeliusChunkBytes)

-- The load: after the first two submissions the first chunk is staged, then
-- every two submissions scatter the staged chunk and the next is staged
-- (one repeat iteration per remaining full chunk), then two submissions and
-- the tail, then one submission scatters it.  The forward submission
-- follows the run's first context; its predictions and the forward part of
-- the result record are written before the request's operation decides
-- whether training continues.
def coppeliusLoadIterations : Nat = (naturalSaturatingSubtract coppeliusFullChunks 1)
def coppeliusLoadFirst : Nat = 3
def coppeliusLoadEnd : Nat = (naturalAdd 2 (naturalMultiply 2 coppeliusLoadIterations))
def coppeliusTailLoadFirst : Nat = (naturalAdd coppeliusLoadEnd 1)
def coppeliusTailScatter : Nat = (naturalAdd coppeliusLoadEnd 3)
def coppeliusForwardSubmission : Nat = (naturalAdd coppeliusLoadEnd 4)

-- The plan's submissions by ordinal (Coppelius.Build.Graph.
-- coppeliusSubmissionSchedule): initialize; two per checkpoint chunk; the
-- forward; the update (one step: forward, backward, AdamW, the half copies,
-- the step's loss copy); the final forward; a gather per chunk; the report
-- put back.  The run issues the update, the final forward, the gathers and
-- the report again from their ordinals.
def coppeliusForwardPiece : Nat = (naturalSaturatingSubtract coppeliusForwardSubmission 1)
def coppeliusUpdatePiece : Nat = (succ coppeliusForwardPiece)
def coppeliusFinalPiece : Nat = (naturalAdd coppeliusUpdatePiece coppeliusUpdatePieces)
def coppeliusGatherPiece = (lambda unrestricted chunk : Nat . (naturalAdd (succ coppeliusFinalPiece) chunk))
def coppeliusRestorePiece : Nat = (coppeliusGatherPiece cgChunks)
-- the host's chunks are the plan's, and the report put back is its last
-- submission
def coppeliusChunksAgree :
  (equal Nat cgChunks (naturalAdd coppeliusFullChunks (naturalNonzero coppeliusTailBytes))) =
  (refl Nat cgChunks)
def coppeliusPiecesArePlans : (equal Nat (succ coppeliusRestorePiece) coppeliusSubmissionCount) =
  (refl Nat coppeliusSubmissionCount)

def coppeliusStagingAt = (lambda unrestricted offset : Nat . (naturalAdd coppeliusHostStaging offset))
def coppeliusUSERDHost : Nat = (naturalAdd coppeliusHostGPFIFO nvidiaPlanHostUVMGPFIFOBytes)

def coppeliusCheckpointIn : (family NvidiaPlanHostFile) = (constructor NvidiaPlanHostFile NvidiaPlanHostCheckpointIn)
def coppeliusCheckpointOut : (family NvidiaPlanHostFile) = (constructor NvidiaPlanHostFile NvidiaPlanHostCheckpointOut)
def coppeliusInput : (family NvidiaPlanHostFile) = (constructor NvidiaPlanHostFile NvidiaPlanHostInput)
def coppeliusPredictions : (family NvidiaPlanHostFile) = (constructor NvidiaPlanHostFile NvidiaPlanHostPredictions)
def coppeliusResult : (family NvidiaPlanHostFile) = (constructor NvidiaPlanHostFile NvidiaPlanHostResult)

def coppeliusSubmit =
  (lambda unrestricted gpPut : Nat .
    (lambda unrestricted stride : Nat .
      (lambda unrestricted tail : (family NvidiaPlanHostSteps) .
        (constructor NvidiaPlanHostSteps NvidiaPlanHostStepSubmit gpPut stride tail))))

def coppeliusStamp =
  (lambda unrestricted gpPut : Nat .
    (lambda unrestricted stride : Nat .
      (lambda unrestricted which : Nat .
        (lambda unrestricted tail : (family NvidiaPlanHostSteps) .
          (constructor NvidiaPlanHostSteps NvidiaPlanHostStepStamp gpPut stride which tail)))))

def coppeliusReadCheckpoint =
  (lambda unrestricted extent : Nat .
    (lambda unrestricted offset : Nat .
      (lambda unrestricted stride : Nat .
        (lambda unrestricted tail : (family NvidiaPlanHostSteps) .
          (constructor NvidiaPlanHostSteps NvidiaPlanHostStepRead coppeliusCheckpointIn coppeliusHostStaging extent offset stride 1 tail)))))

def coppeliusWriteCheckpoint =
  (lambda unrestricted extent : Nat .
    (lambda unrestricted tail : (family NvidiaPlanHostSteps) .
      (constructor NvidiaPlanHostSteps NvidiaPlanHostStepWrite coppeliusCheckpointOut coppeliusHostStaging extent tail)))

def coppeliusWriteResult =
  (lambda unrestricted host : Nat .
    (lambda unrestricted extent : Nat .
      (lambda unrestricted tail : (family NvidiaPlanHostSteps) .
        (constructor NvidiaPlanHostSteps NvidiaPlanHostStepWrite coppeliusResult host extent tail))))

def coppeliusRecordWord =
  (lambda unrestricted host : Nat .
    (lambda unrestricted statOffset : Nat .
      (lambda unrestricted tail : (family NvidiaPlanHostSteps) .
        (constructor NvidiaPlanHostSteps NvidiaPlanHostStepRecordStat host statOffset 4 tail))))

def coppeliusStepsEnd : (family NvidiaPlanHostSteps) = (constructor NvidiaPlanHostSteps NvidiaPlanHostStepsEnd)

def coppeliusLoadBody : (family NvidiaPlanHostSteps) =
  (coppeliusSubmit coppeliusLoadFirst 2
    (coppeliusSubmit (naturalAdd coppeliusLoadFirst 1) 2
      (coppeliusStamp (naturalAdd coppeliusLoadFirst 1) 2 32
        (coppeliusReadCheckpoint coppeliusChunkBytes coppeliusChunkBytes coppeliusChunkBytes
          (coppeliusStamp (naturalAdd coppeliusLoadFirst 1) 2 48
            coppeliusStepsEnd)))))

-- the context at the stream's cursor into the staging record: its inputs,
-- and its targets (the same window one token on)
def coppeliusReadContext =
  (lambda unrestricted tail : (family NvidiaPlanHostSteps) .
    (constructor NvidiaPlanHostSteps NvidiaPlanHostStepReadAtCursor coppeliusInput coppeliusHostStaging coppeliusContextStreamBytes 0
      (constructor NvidiaPlanHostSteps NvidiaPlanHostStepReadAtCursor coppeliusInput (coppeliusStagingAt coppeliusInputTargetsOffset)
        coppeliusContextStreamBytes coppeliusTokenBytes tail)))

def coppeliusReissue =
  (lambda unrestricted entries : Bytes .
    (lambda unrestricted gpPut : Nat .
      (lambda unrestricted piece : Nat .
        (lambda unrestricted tail : (family NvidiaPlanHostSteps) .
          (constructor NvidiaPlanHostSteps NvidiaPlanHostStepReissue gpPut 0 piece 0 (nvidiaPlanHostTableWord entries piece) 0 tail)))))

-- ---- the run ----
-- After the forward submission the ring carries the run, period by period
-- (coppeliusRunCheckpointInterval steps each): the update issued once per
-- step, then the checkpoint gathered back; the last period issues the
-- final forward before its gathers and the report put back after them.
def coppeliusPeriodSubmissions : Nat = (naturalAdd coppeliusRunCheckpointInterval cgChunks)
def coppeliusPeriodFirst =
  (lambda unrestricted period : Nat .
    (naturalAdd (succ coppeliusForwardSubmission) (naturalMultiply period coppeliusPeriodSubmissions)))
def coppeliusLastPeriod : Nat = (naturalSaturatingSubtract coppeliusRunPublications 1)
def coppeliusRunSubmissionCount : Nat =
  (naturalAdd coppeliusForwardSubmission (naturalAdd (naturalMultiply coppeliusRunPublications coppeliusPeriodSubmissions) 2))

-- One step: its context read and the cursor moved on; its AdamW scalars
-- computed into the update's AdamW launches (`adamW`); the update issued;
-- its per-row losses (the update's loss copy put them in the loss-record
-- window) appended to the result.
def coppeliusStepBody =
  (lambda unrestricted entries : Bytes .
    (lambda unrestricted adamW : (family NativePhysicalCommands) .
      (lambda unrestricted first : Nat .
        (coppeliusReadContext
          (constructor NvidiaPlanHostSteps NvidiaPlanHostStepCursorAdvance coppeliusContextStreamBytes
            (constructor NvidiaPlanHostSteps NvidiaPlanHostStepCommands adamW
              (constructor NvidiaPlanHostSteps NvidiaPlanHostStepReissue first 1 coppeliusUpdatePiece 0
                (nvidiaPlanHostTableWord entries coppeliusUpdatePiece) 0
                (coppeliusWriteResult (coppeliusStagingAt coppeliusResultLossRecordOffset) coppeliusResultWindowBytes
                  coppeliusStepsEnd))))))))

-- The checkpoint gathered back, from gpPut `first`: each chunk's gather
-- issued and the staged chunk written to the temporary.  Gathers laid end
-- to end at a common stride (equal pieces) are issued by one repeat, the
-- piece and its ring entry advancing per iteration; a chunk whose piece
-- differs (one that spans two planes takes a launch more) starts a run of
-- its own.  The runs are read off the realized table (the stride is the
-- second full chunk's step from the first), and nvidiaPlanHostReissuesAdmitted
-- holds every entry issued to it.
def coppeliusGatherEntry =
  (lambda unrestricted entries : Bytes . (lambda unrestricted chunk : Nat .
    (nvidiaPlanHostTableWord entries (coppeliusGatherPiece chunk))))

def coppeliusGatherStride =
  (lambda unrestricted entries : Bytes .
    (naturalSaturatingSubtract (coppeliusGatherEntry entries 1) (coppeliusGatherEntry entries 0)))

-- the full chunks, as runs in order (built last run first, from the last
-- chunk back: a chunk joins the run after it when that run's first entry is
-- one stride past its own)
def coppeliusGatherRuns =
  (lambda unrestricted entries : Bytes .
    (let unrestricted stride = (coppeliusGatherStride entries) in
    (nat-eliminate
      (lambda unrestricted current : Nat . (family StdList (family CoppeliusGatherRun)))
      (constructor StdList StdListEmpty (family CoppeliusGatherRun))
      (lambda unrestricted p : Nat .
        (lambda unrestricted later : (family StdList (family CoppeliusGatherRun)) .
          (let unrestricted chunk = (naturalSaturatingSubtract (naturalSaturatingSubtract coppeliusFullChunks 1) p) in
          (let unrestricted alone = (constructor StdList StdListCons (family CoppeliusGatherRun)
                                      (constructor CoppeliusGatherRun CoppeliusGatherRunValue chunk 1) later) in
          (eliminate StdList (lambda unrestricted current : (family StdList (family CoppeliusGatherRun)) . (family StdList (family CoppeliusGatherRun))) later
            (branch StdListEmpty . alone)
            (branch StdListCons run rest induction .
              (eliminate CoppeliusGatherRun (lambda unrestricted current : (family CoppeliusGatherRun) . (family StdList (family CoppeliusGatherRun))) run
                (branch CoppeliusGatherRunValue first length .
                  (nat-eliminate (lambda unrestricted current : Nat . (family StdList (family CoppeliusGatherRun)))
                    alone
                    (lambda unrestricted q : Nat . (lambda unrestricted ignored : (family StdList (family CoppeliusGatherRun)) .
                      (constructor StdList StdListCons (family CoppeliusGatherRun)
                        (constructor CoppeliusGatherRun CoppeliusGatherRunValue chunk (succ length)) rest)))
                    (naturalEqual (coppeliusGatherEntry entries first) (naturalAdd (coppeliusGatherEntry entries chunk) stride)))))))))))
      coppeliusFullChunks)))

def coppeliusGathers =
  (lambda unrestricted entries : Bytes .
    (lambda unrestricted first : Nat .
      (lambda unrestricted tail : (family NvidiaPlanHostSteps) .
        (stdListFold (family CoppeliusGatherRun) (family NvidiaPlanHostSteps)
          (lambda unrestricted run : (family CoppeliusGatherRun) .
            (lambda unrestricted after : (family NvidiaPlanHostSteps) .
              (eliminate CoppeliusGatherRun (lambda unrestricted current : (family CoppeliusGatherRun) . (family NvidiaPlanHostSteps)) run
                (branch CoppeliusGatherRunValue chunk length .
                  (constructor NvidiaPlanHostSteps NvidiaPlanHostStepRepeat length
                    (constructor NvidiaPlanHostSteps NvidiaPlanHostStepReissue (naturalAdd first chunk) 1 (coppeliusGatherPiece chunk) 1
                      (coppeliusGatherEntry entries chunk) (coppeliusGatherStride entries)
                      (coppeliusWriteCheckpoint coppeliusChunkBytes coppeliusStepsEnd))
                    after)))))
          (coppeliusReissue entries (naturalAdd first coppeliusFullChunks) (coppeliusGatherPiece coppeliusFullChunks)
            (coppeliusWriteCheckpoint coppeliusTailBytes tail))
          (coppeliusGatherRuns entries)))))

-- `chosen` when the flag is nonzero, else `otherwise`
def coppeliusStepsWhen =
  (lambda unrestricted flag : Nat .
    (lambda unrestricted chosen : (family NvidiaPlanHostSteps) .
      (lambda unrestricted otherwise : (family NvidiaPlanHostSteps) .
        (nat-eliminate (lambda unrestricted current : Nat . (family NvidiaPlanHostSteps)) otherwise
          (lambda unrestricted predecessor : Nat . (lambda unrestricted ignored : (family NvidiaPlanHostSteps) . chosen))
          flag))))

-- Period p: its steps; in the last, the final forward; the checkpoint's
-- temporary (the first period's the train operation made); the gathers;
-- in the last, the report put back; the checkpoint published, counting the
-- steps (and the contexts) taken so far.
def coppeliusPeriod =
  (lambda unrestricted entries : Bytes .
    (lambda unrestricted adamW : (family NativePhysicalCommands) .
      (lambda unrestricted period : Nat .
        (lambda unrestricted tail : (family NvidiaPlanHostSteps) .
          (let unrestricted first = (coppeliusPeriodFirst period) in
          (let unrestricted last = (naturalEqual period coppeliusLastPeriod) in
          (let unrestricted gathered = (naturalAdd (naturalAdd first coppeliusRunCheckpointInterval) last) in
          (let unrestricted taken = (naturalMultiply (succ period) coppeliusRunCheckpointInterval) in
          (let unrestricted publish =
            (constructor NvidiaPlanHostSteps NvidiaPlanHostStepCheckpointPublish taken taken tail) in
          (let unrestricted gathers =
            (coppeliusGathers entries gathered
              (coppeliusStepsWhen last
                (coppeliusReissue entries (naturalAdd gathered cgChunks) coppeliusRestorePiece publish)
                publish)) in
          (let unrestricted begun =
            (coppeliusStepsWhen period
              (constructor NvidiaPlanHostSteps NvidiaPlanHostStepCheckpointBegin gathers)
              gathers) in
          (constructor NvidiaPlanHostSteps NvidiaPlanHostStepRepeat coppeliusRunCheckpointInterval
            (coppeliusStepBody entries adamW first)
            (coppeliusStepsWhen last
              (coppeliusReissue entries (naturalAdd first coppeliusRunCheckpointInterval) coppeliusFinalPiece begun)
              begun)))))))))))))

def coppeliusRunPeriods =
  (lambda unrestricted entries : Bytes .
    (lambda unrestricted adamW : (family NativePhysicalCommands) .
      (lambda unrestricted tail : (family NvidiaPlanHostSteps) .
        (app
          (nat-eliminate
            (lambda unrestricted current : Nat . (pi unrestricted after : (family NvidiaPlanHostSteps) . (family NvidiaPlanHostSteps)))
            (lambda unrestricted after : (family NvidiaPlanHostSteps) . after)
            (lambda unrestricted p : Nat .
              (lambda unrestricted induction : (pi unrestricted after : (family NvidiaPlanHostSteps) . (family NvidiaPlanHostSteps)) .
                (lambda unrestricted after : (family NvidiaPlanHostSteps) . (induction (coppeliusPeriod entries adamW p after)))))
            coppeliusRunPublications)
          tail))))

-- The whole schedule.  The stream's cursor starts at the checkpoint's
-- position; the run's first context is staged for the forward submission
-- (read again for the first step: the cursor moves only there).
def coppeliusHostStepsFor =
  (lambda unrestricted entries : Bytes .
    (lambda unrestricted adamW : (family NativePhysicalCommands) .
  (constructor NvidiaPlanHostSteps NvidiaPlanHostStepCursorFromCheckpoint coppeliusTokenStreamHeaderBytes coppeliusContextStreamBytes
  (coppeliusSubmit 1 0
    (coppeliusSubmit 2 0
      (coppeliusReadCheckpoint coppeliusChunkBytes 0 0
        (constructor NvidiaPlanHostSteps NvidiaPlanHostStepRepeat coppeliusLoadIterations coppeliusLoadBody
          (coppeliusSubmit coppeliusTailLoadFirst 0
            (coppeliusSubmit (naturalAdd coppeliusTailLoadFirst 1) 0
              (coppeliusReadCheckpoint coppeliusTailBytes coppeliusTailOffset 0
                (coppeliusSubmit coppeliusTailScatter 0
                  (coppeliusReadContext
                    (constructor NvidiaPlanHostSteps NvidiaPlanHostStepFill (coppeliusStagingAt coppeliusResultRecordOffset) coppeliusResultRecordMarker
                      (coppeliusSubmit coppeliusForwardSubmission 0
                        (constructor NvidiaPlanHostSteps NvidiaPlanHostStepWrite coppeliusPredictions (coppeliusStagingAt coppeliusResultPredictionsOffset) coppeliusResultPredictionsBytes
                          (coppeliusWriteResult (coppeliusStagingAt coppeliusResultRecordOffset) coppeliusResultRecordForwardBytes
                            (constructor NvidiaPlanHostSteps NvidiaPlanHostStepOperation
                              (coppeliusRunPeriods entries adamW
                                (coppeliusRecordWord (naturalAdd coppeliusUSERDHost nvidiaPlanHostUSERDGPGetOffset) coppeliusResultStatGPGetOffset
                                  (coppeliusRecordWord (naturalAdd coppeliusUSERDHost nvidiaPlanHostUSERDGPPutOffset) coppeliusResultStatGPPutOffset
                                    (coppeliusRecordWord coppeliusHostSemaphores coppeliusResultStatSemaphoreOffset
                                      (constructor NvidiaPlanHostSteps NvidiaPlanHostStepWriteStat coppeliusResult
                                        (coppeliusWriteResult (coppeliusStagingAt coppeliusResultLossBlockOffset) coppeliusResultLossBlockBytes
                                          (coppeliusWriteResult (coppeliusStagingAt coppeliusResultRecordOffset) coppeliusResultRecordFinalBytes
                                            (coppeliusWriteResult coppeliusHostSemaphores coppeliusSemaphoreBytes
                                              (coppeliusWriteResult coppeliusHostErrorNotifier coppeliusResultErrorNotifierBytes
                                                (constructor NvidiaPlanHostSteps NvidiaPlanHostStepWriteStatus coppeliusResult
                                                  (constructor NvidiaPlanHostSteps NvidiaPlanHostStepWriteTimestamps coppeliusResult
                                                    coppeliusStepsEnd))))))))))))))))))))))))))))

-- ---- AdamW's step-dependent scalars ----
-- Coppelius's learner defines them exactly (Coppelius.Build.Graph: the
-- binary32 nearest lr sqrt(1 - beta2^t) / (1 - beta1^t)); the host computes
-- them in binary64 from the same ratios, continuing the step count from the
-- checkpoint's header (scripts/ci/adamw-schedule.sh holds the two to each
-- other): the constants and the powers first, then one update's scalars
-- per step, written into the update's AdamW launches.
def coppeliusHostAdamW : (family NvidiaPlanHostAdamW) =
  (constructor NvidiaPlanHostAdamW NvidiaPlanHostAdamWValue
    (constructor NvidiaPlanHostArithmetic NvidiaPlanHostBinary64)
    (constructor NvidiaPlanHostConstant NvidiaPlanHostRatio coppeliusLearningRateNumerator coppeliusLearningRateDenominator)
    (constructor NvidiaPlanHostConstant NvidiaPlanHostRatio coppeliusBeta1Numerator coppeliusBeta1Denominator)
    (constructor NvidiaPlanHostConstant NvidiaPlanHostRatio coppeliusBeta2Numerator coppeliusBeta2Denominator)
    (constructor NvidiaPlanHostConstant NvidiaPlanHostRatio (naturalMultiply coppeliusEpsilonNumerator coppeliusLossScale) coppeliusEpsilonDenominator))

def coppeliusRealizedAddresses =
  (lambda unrestricted realized : (family CoppeliusRealizedTables) .
    (eliminate CoppeliusRealizedTables (lambda unrestricted current : (family CoppeliusRealizedTables) . Bytes) realized
      (branch CoppeliusRealizedTablesValue addresses words entries . addresses)))

def coppeliusRealizedEntries =
  (lambda unrestricted realized : (family CoppeliusRealizedTables) .
    (eliminate CoppeliusRealizedTables (lambda unrestricted current : (family CoppeliusRealizedTables) . Bytes) realized
      (branch CoppeliusRealizedTablesValue addresses words entries . entries)))

-- where the host writes: the step-size word of the parameter block (the
-- epsilon beside it) of each of the update's AdamW launches
def coppeliusAdamWWordOffset : Nat = adamWHalfSM86StepSizeOffset

def coppeliusAdamWLaunches =
  (lambda unrestricted k : Nat .
    (stdListReverse Nat
      (nat-eliminate
        (lambda unrestricted current : Nat . (family StdList Nat))
        (constructor StdList StdListEmpty Nat)
        (lambda unrestricted piece : Nat .
          (lambda unrestricted rest : (family StdList Nat) .
            (constructor StdList StdListCons Nat (coppeliusAdamWLaunch k piece) rest)))
        coppeliusAdamWPieces)))

-- the word the plan bakes there: the first step's step size and epsilon
-- (Coppelius.Build.Graph)
def coppeliusAdamWBaked =
  (lambda unrestricted k : Nat .
    (let unrestricted step = (naturalAdd coppeliusFirstStep k) in
    (naturalAdd (cgAdamStepSize step) (naturalMultiply 4294967296 (cgAdamEpsilon step)))))

def coppeliusAdamWSites =
  (lambda unrestricted pushExpanded : Nat .
    (lambda unrestricted addresses : Bytes .
      (nvidiaPlanHostAdamWHostSites (coppeliusHostLayout pushExpanded b"") addresses coppeliusAdamWWordOffset coppeliusUpdatePieces coppeliusAdamWLaunches)))

-- the learner's prelude (the constants, the powers the checkpoint counts)
-- and one step's commands
def coppeliusHostLearner =
  (lambda unrestricted pushExpanded : Nat .
  (lambda unrestricted addresses : Bytes .
    (nvidiaPlanHostAdamWCommands coppeliusHostAdamW 0 (coppeliusAdamWSites pushExpanded addresses))))

def coppeliusHostStepAdamW =
  (lambda unrestricted pushExpanded : Nat .
  (lambda unrestricted addresses : Bytes .
    (nvidiaPlanHostAdamWStepCommands coppeliusHostAdamW (coppeliusAdamWSites pushExpanded addresses 0))))

def coppeliusHostStepsOf =
  (lambda unrestricted pushExpanded : Nat .
    (lambda unrestricted realized : (family CoppeliusRealizedTables) .
      (coppeliusHostStepsFor (coppeliusRealizedEntries realized)
        (coppeliusHostStepAdamW pushExpanded (coppeliusRealizedAddresses realized)))))

-- ---- contracts the pairing decides ----
-- the schedule issues the run's submissions, every step on a fresh
-- context; it moves the whole checkpoint in -- in order, every chunk once --
-- and out at every publication; every
-- submission it issues again writes the ring entry the realization gave
-- it; and the plan's device side is certified (Coppelius.Build.Graph.
-- coppeliusDeviceHazard: every pointer in a named plane, no two live
-- planes sharing a byte)
def coppeliusHostIssuesRun =
  (lambda unrestricted steps : (family NvidiaPlanHostSteps) .
    (naturalEqual (nvidiaPlanHostStepSubmissions steps) coppeliusRunSubmissionCount))
def coppeliusHostMovesCheckpoint =
  (lambda unrestricted steps : (family NvidiaPlanHostSteps) .
    (naturalAnd
      (nvidiaPlanHostFileReadsTile coppeliusCheckpointIn coppeliusCheckpointPayloadBytesNatural steps)
      (naturalEqual (nvidiaPlanHostStepBytesWritten coppeliusCheckpointOut steps)
        (naturalMultiply coppeliusRunPublications coppeliusCheckpointPayloadBytesNatural))))

-- every issue of the update trains on a context no earlier step read: the
-- cursor moves on by one context exactly once before each (it starts at
-- the checkpoint's position, the forward's context, which the first step
-- trains on and moves past), and the run issues the update once per step
def coppeliusHostStepsFresh =
  (lambda unrestricted steps : (family NvidiaPlanHostSteps) .
    (app (app (app
      (eliminate NvidiaPlanHostSteps
        (lambda unrestricted current : (family NvidiaPlanHostSteps) .
          (pi unrestricted moved : Nat . (pi unrestricted issued : Nat . (pi unrestricted ok : Nat . Nat))))
        (nvidiaPlanHostStepsUnrolled steps)
        (branch NvidiaPlanHostStepsEnd .
          (lambda unrestricted moved : Nat . (lambda unrestricted issued : Nat . (lambda unrestricted ok : Nat .
            (naturalAnd ok (naturalAnd (naturalIsZero moved) (naturalEqual issued coppeliusRunSteps)))))))
        (branch NvidiaPlanHostStepSubmit gpPut stride tail induction . induction)
        (branch NvidiaPlanHostStepRepeat count body tail bodyInduction tailInduction . tailInduction)
        (branch NvidiaPlanHostStepRead file host extent offset stride tolerateEmpty tail induction . induction)
        (branch NvidiaPlanHostStepReadCyclic host extent offset recordBytes count index indexStride tail induction . induction)
        (branch NvidiaPlanHostStepWrite file host extent tail induction . induction)
        (branch NvidiaPlanHostStepWriteAt file host extent offset tail induction . induction)
        (branch NvidiaPlanHostStepWriteStat file tail induction . induction)
        (branch NvidiaPlanHostStepWriteStatus file tail induction . induction)
        (branch NvidiaPlanHostStepWriteTimestamps file tail induction . induction)
        (branch NvidiaPlanHostStepRecordStat host statOffset extent tail induction . induction)
        (branch NvidiaPlanHostStepStamp gpPut stride which tail induction . induction)
        (branch NvidiaPlanHostStepFill host payload tail induction . induction)
        (branch NvidiaPlanHostStepAssertWord identity host expected tail induction . induction)
        (branch NvidiaPlanHostStepBindInputIdentity offset extent tail induction . induction)
        (branch NvidiaPlanHostStepWriteInputIdentity tail induction . induction)
        (branch NvidiaPlanHostStepBeginCheckpoint tail induction . induction)
        (branch NvidiaPlanHostStepOperation tail induction . induction)
        (branch NvidiaPlanHostStepSync file tail induction . induction)
        (branch NvidiaPlanHostStepDataSync file tail induction . induction)
        (branch NvidiaPlanHostStepClose file tail induction . induction)
        (branch NvidiaPlanHostStepReissue gpPut stride piece pieceStride entry entryStride tail induction .
          (lambda unrestricted moved : Nat . (lambda unrestricted issued : Nat . (lambda unrestricted ok : Nat .
            (let unrestricted update = (naturalEqual piece coppeliusUpdatePiece) in
            (induction (naturalSelect update 0 moved) (naturalAdd issued update)
              (naturalAnd ok (naturalSelect update moved 1))))))))
        (branch NvidiaPlanHostStepCommands commands tail induction . induction)
        (branch NvidiaPlanHostStepCursorFromCheckpoint base scale tail induction . induction)
        (branch NvidiaPlanHostStepReadAtCursor file host extent delta tail induction . induction)
        (branch NvidiaPlanHostStepCursorAdvance bytes tail induction .
          (lambda unrestricted moved : Nat . (lambda unrestricted issued : Nat . (lambda unrestricted ok : Nat .
            (induction 1 issued (naturalAnd ok (naturalAnd (naturalIsZero moved) (naturalEqual bytes coppeliusContextStreamBytes))))))))
        (branch NvidiaPlanHostStepCheckpointBegin tail induction . induction)
        (branch NvidiaPlanHostStepCheckpointPublish updates invocations tail induction . induction))
      0) 0) 1))

def coppeliusHostLearnerAdmitted =
  (lambda unrestricted pushExpanded : Nat .
  (lambda unrestricted realized : (family CoppeliusRealizedTables) .
    (eliminate CoppeliusRealizedTables (lambda unrestricted current : (family CoppeliusRealizedTables) . Nat) realized
      (branch CoppeliusRealizedTablesValue addresses words entries .
        (naturalAnd
          (naturalEqual adamWHalfSM86EpsilonOffset (naturalAdd coppeliusAdamWWordOffset 4))
          (nvidiaPlanHostAdamWSitesAdmitted (coppeliusHostLayout pushExpanded b"") addresses words coppeliusAdamWWordOffset
            coppeliusUpdatePieces coppeliusAdamWLaunches coppeliusAdamWBaked))))))

def coppeliusHostCertified =
  (lambda unrestricted qmdExpanded : Nat .
    (lambda unrestricted pushExpanded : Nat .
      (lambda unrestricted realized : (family CoppeliusRealizedTables) .
        (let unrestricted steps = (coppeliusHostStepsOf pushExpanded realized) in
        (naturalAnd (coppeliusHostIssuesRun steps)
          (naturalAnd (coppeliusHostStepsFresh steps)
          (naturalAnd (coppeliusHostMovesCheckpoint steps)
            (naturalAnd (nvidiaPlanHostReissuesAdmitted (coppeliusRealizedEntries realized) steps)
              (nvidiaPlanHostRequestAdmitted (coppeliusHostLayout pushExpanded b"") coppeliusRunSubmissionCount
                (coppeliusRecipeFills qmdExpanded pushExpanded) steps)))))))))

-- the host's side (its schedule, and the AdamW sites the realization
-- placed), then (only when it holds) the device's: the plan's order, then
-- (only when it holds) the run's -- a refusal stops at the first that fails
def coppeliusHostAdmitted =
  (lambda unrestricted qmdExpanded : Nat .
    (lambda unrestricted pushExpanded : Nat .
      (lambda unrestricted realized : (family CoppeliusRealizedTables) .
        (nat-eliminate (lambda unrestricted host : Nat . Nat) 0
          (lambda unrestricted p : Nat . (lambda unrestricted ignored : Nat .
            (nat-eliminate (lambda unrestricted plan : Nat . Nat) 0
              (lambda unrestricted q : Nat . (lambda unrestricted ignoredToo : Nat . (naturalIsZero (coppeliusRunDeviceHazardOf q))))
              (naturalIsZero (coppeliusDeviceHazardOf p)))))
          (naturalAnd coppeliusSchedulesAdmitted
            (naturalAnd (coppeliusHostLearnerAdmitted pushExpanded realized) (coppeliusHostCertified qmdExpanded pushExpanded realized)))))))

-- ---- the host ELF ----
-- `productIdentity` is the pairing's padded card name: the host payload's
-- identity in the envelope and the name the card must report.
-- `hostWriter` makes the host program an executable for the pairing's host:
-- Runtime.NativePhysicalNativeELF's direct writer for x86-64, Platform.
-- AArch64.HostLowering's for the DGX Spark.
def coppeliusNativeHostELF =
  (lambda unrestricted hostWriter : (pi unrestricted program : (family NativePhysicalProgram) . Bytes) .
  (lambda unrestricted hostTarget : (family CoppeliusHostTarget) .
  (lambda unrestricted productIdentity : Bytes .
    (lambda unrestricted programIdentity : Bytes .
      (lambda unrestricted qmdIdentity : Bytes .
        (lambda unrestricted pushIdentity : Bytes .
          (lambda unrestricted gpfifoIdentity : Bytes .
            (lambda unrestricted programLength : Nat .
              (lambda unrestricted qmdLength : Nat .
                (lambda unrestricted pushLength : Nat .
                  (lambda unrestricted gpfifoLength : Nat .
                    (lambda unrestricted qmdExpanded : Nat .
                      (lambda unrestricted pushExpanded : Nat .
                      (lambda unrestricted realized : (family CoppeliusRealizedTables) .
                        (nat-eliminate
                          (lambda unrestricted admitted : Nat . Bytes)
                          b""
                          (lambda unrestricted predecessor : Nat .
                            (lambda unrestricted induction : Bytes .
                              (hostWriter (nvidiaPlanHostProgram productIdentity
                                (nvidiaPlanHostRequestCommands (coppeliusHostLayoutFor hostTarget pushExpanded productIdentity)
                                  productIdentity programIdentity qmdIdentity pushIdentity gpfifoIdentity
                                  programLength qmdLength pushLength gpfifoLength
                                  coppeliusRunSubmissionCount (constructor NvidiaPlanHostInputContract NvidiaPlanHostInputStream) coppeliusCheckpointContract
                                  (coppeliusRecipeFills qmdExpanded pushExpanded)
                                  (coppeliusHostLearner pushExpanded (coppeliusRealizedAddresses realized))
                                  (coppeliusHostStepsOf pushExpanded realized))))))
                          (coppeliusHostAdmitted qmdExpanded pushExpanded realized))))))))))))))))

-- ---- what the artifact demands of a card (Platform.Linux.Nvidia.Compatibility) ----
-- the instruction set its QMDs are built for, its kernels' resources, the
-- video memory it reserves, and its host's driver branch, compute class and
-- unified memory; a pairing admits a card whose profile supplies them all.
-- The host demands read only the layout's ABI and lifecycle, so the
-- pushbuffer's extent (known only at build, from its expansion) is 0 here.
def coppeliusCompatDemands =
  (lambda unrestricted architecture : (family ModelWord32) .
    (nvidiaPlanDemands (nvidiaInstructionSetOf (modelWord32ToNatural architecture)) coppeliusWholeProgramDeviceRegions coppeliusSM86CompatVideoBytesNatural
      (nvidiaHostDemands (coppeliusHostLayout 0 b"") nvidiaNoDemands)))