module Bob.Build.PhysicalComponents

import Accelerator.SM86.InstructionEncoding
import Bob.ArenaPlan
import Bob.Build.SM86Programs
import Bob.Learner
import Bob.Model
import Bob.TrainingRun
import Compiler.Planning.VA
import Data.Bytes
import Float32Exact
import BinaryFloatModel
import Float32Model
import FloatLiteralSpec
import Hardware.Nvidia.SM86.Command.QMD
import Hardware.Nvidia.SM86.Command.WholeProgramPlan
import Model.Config
import Model.Parameter
import Model.Word64
import Runtime.ArenaCertificate
import Runtime.DeviceArenaCertificate
import Platform.Linux.Nvidia.PlanHost
import Platform.Linux.Nvidia.PlanHostRequest
import Realization.Nvidia.SM86.AdamWSM86
import Realization.Nvidia.SM86.EmbeddingGatherF32SM86
import Realization.Nvidia.SM86.GreedyArgmaxSM86
import Std.Natural

-- Bob declares what runs on the device: four typed regions, the launches that
-- use them with their kernel parameters, and the order they are submitted in.
-- The NVIDIA backend places and prunes the regions, places the parameter
-- blocks, and encodes the QMD table, pushbuffer and GPFIFO from this plan.
-- Nothing here joins bytes or knows where anything lands.

-- Bob's four device kernels, each encoded by its maintained owner.
def bobGatherMachineBytes : Bytes =
  (eliminate
    EmbeddingGatherF32ImageResult
    (lambda unrestricted current : (family EmbeddingGatherF32ImageResult) . Bytes)
    bobPredictionRowImage
    (branch EmbeddingGatherF32ImageReady image digest telemetry . image)
    (branch EmbeddingGatherF32ImageFailed error ordinal detail telemetry . b""))

def bobLossMachineBytes : Bytes =
  (eliminate
    SM86ProgramEncodingResult
    (lambda unrestricted current : (family SM86ProgramEncodingResult) . Bytes)
    bobLossGradientImage
    (branch SM86ProgramEncodingSucceeded image telemetry . image)
    (branch SM86ProgramEncodingFailed ordinal failure telemetry . b""))

def bobAdamMachineBytes : Bytes =
  (eliminate
    AdamWF32SM86ImageResult
    (lambda unrestricted current : (family AdamWF32SM86ImageResult) . Bytes)
    bobOptimizerImage
    (branch AdamWF32SM86ImageReady image digest encoding identity telemetry . image)
    (branch AdamWF32SM86ImageFailed error ordinal detail telemetry . b""))

def bobArgmaxMachineBytes : Bytes =
  (eliminate
    SM86ProgramEncodingResult
    (lambda unrestricted current : (family SM86ProgramEncodingResult) . Bytes)
    (sm86EncodeProgram greedyArgmaxSM86ByteVocabularyProgram)
    (branch SM86ProgramEncodingSucceeded image telemetry . image)
    (branch SM86ProgramEncodingFailed ordinal failure telemetry . b""))

def bobWord64 =
  (lambda unrestricted value : Nat . (modelWord64FromNaturalTruncated value))

-- Device regions.  Registers, block width and shared memory are properties of
-- each kernel; the backend derives the QMD resource fields from them.
def bobRegion =
  (lambda unrestricted identity : Bytes .
    (lambda unrestricted material : Bytes .
      (lambda unrestricted registers : Nat .
        (lambda unrestricted blockX : Nat .
          (constructor NvidiaDeviceRegion NvidiaDeviceRegionValue
            identity material registers blockX zero)))))

def bobRegions : (family NvidiaDeviceRegions) =
  (constructor NvidiaDeviceRegions NvidiaDeviceRegionsNext
    (bobRegion b"bob-row-gather" bobGatherMachineBytes 24 64)
    (constructor NvidiaDeviceRegions NvidiaDeviceRegionsNext
      (bobRegion b"bob-loss-gradient" bobLossMachineBytes 32 256)
      (constructor NvidiaDeviceRegions NvidiaDeviceRegionsNext
        (bobRegion b"bob-adamw" bobAdamMachineBytes 32 256)
        (constructor NvidiaDeviceRegions NvidiaDeviceRegionsNext
          (bobRegion b"bob-greedy-argmax" bobArgmaxMachineBytes 32 1)
          (constructor NvidiaDeviceRegions NvidiaDeviceRegionsEnd)))))

-- A launch names its region by the region's place in the declared order.  The
-- places follow from the region sizes and the plan's alignment; the backend
-- uses them only as keys and repacks whatever survives.
def bobRegionAlignment : Nat =
  256

def bobAlign =
  (lambda unrestricted value : Nat .
    (naturalMultiply
      (naturalDivideUnchecked
        (naturalAdd value (naturalSaturatingSubtract bobRegionAlignment 1))
        bobRegionAlignment)
      bobRegionAlignment))

def bobGatherPlace : Nat =
  zero

def bobLossPlace : Nat =
  (bobAlign (naturalAdd bobGatherPlace (bytes-length bobGatherMachineBytes)))

def bobAdamPlace : Nat =
  (bobAlign (naturalAdd bobLossPlace (bytes-length bobLossMachineBytes)))

def bobArgmaxPlace : Nat =
  (bobAlign (naturalAdd bobAdamPlace (bytes-length bobAdamMachineBytes)))

def bobCodeExtent : Nat =
  (bobAlign (naturalAdd bobArgmaxPlace (bytes-length bobArgmaxMachineBytes)))

def bobParameterBytes : Nat =
  512

def bobKernel =
  (lambda unrestricted place : Nat .
    (lambda unrestricted gridX : Nat .
      (constructor NvidiaLaunchKernel NvidiaLaunchKernelValue
        (bobWord64 (naturalAdd bobProgramBase place))
        gridX 1 1 1 1 bobParameterBytes)))

-- Kernel parameters are sparse 8-byte words in a 512-byte constant buffer.
-- Every Bob kernel reads its arguments from the SM86 parameter base on: the
-- i-th argument is word `bobArgument i`; a kernel's binary32 scalars follow
-- its pointers, two to a word (the loss kernel's after six pointer slots,
-- AdamW's after its four).
def bobArgument =
  (lambda unrestricted i : Nat . (naturalAdd (naturalDivideUnchecked qmdKernelParameterBase 8) i))

def bobNoPatches : (family NvidiaParameterPatches) =
  (constructor NvidiaParameterPatches NvidiaParameterPatchesEnd)

def bobPatch =
  (lambda unrestricted word : Nat .
    (lambda unrestricted value : (family ModelWord64) .
      (lambda unrestricted rest : (family NvidiaParameterPatches) .
        (constructor NvidiaParameterPatches NvidiaParameterPatchesNext
          (constructor NvidiaParameterPatch NvidiaParameterPatchValue word value)
          rest))))

-- Consecutive words taken from little-endian bytes, used for the f32 scalars a
-- kernel reads after its addresses.
def bobWordPatches =
  (lambda unrestricted words : Nat .
    (nat-eliminate
      (lambda unrestricted current : Nat .
        (pi unrestricted word : Nat .
          (pi unrestricted input : Bytes .
            (pi unrestricted rest : (family NvidiaParameterPatches) .
              (family NvidiaParameterPatches)))))
      (lambda unrestricted word : Nat .
        (lambda unrestricted input : Bytes .
          (lambda unrestricted rest : (family NvidiaParameterPatches) . rest)))
      (lambda unrestricted predecessor : Nat .
        (lambda unrestricted induction :
          (pi unrestricted word : Nat .
            (pi unrestricted input : Bytes .
              (pi unrestricted rest : (family NvidiaParameterPatches) .
                (family NvidiaParameterPatches)))) .
          (lambda unrestricted word : Nat .
            (lambda unrestricted input : Bytes .
              (lambda unrestricted rest : (family NvidiaParameterPatches) .
                (eliminate
                  DataBytesWord64DecodeResult
                  (lambda unrestricted current : (family DataBytesWord64DecodeResult) .
                    (family NvidiaParameterPatches))
                  (dataBytesReadWord64LE input)
                  (branch DataBytesWord64Decoded value remaining .
                    (bobPatch word value (induction (succ word) remaining rest)))
                  (branch DataBytesWord64DecodeFailed error . rest)))))))
      words))

def bobBlock =
  (lambda unrestricted patches : (family NvidiaParameterPatches) .
    (constructor NvidiaParameterBlock NvidiaParameterBlockValue patches))

-- Bob's device arena: one system-memory staging buffer the host maps, each
-- tensor after the previous, page-aligned: the parameters and the two
-- optimizer moments (the checkpoint, contiguous, in that order), the
-- gradients, the input record (Bob.TrainingRun), the two loss blocks, the
-- prompts' logit rows and their predictions.  The host (Bob.Build.NativeHost)
-- reads the input record and the checkpoint into it and writes the results
-- out of it.
def bobArenaNext =
  (lambda unrestricted address : Nat .
    (lambda unrestricted bytes : Nat . (vaPlaceAfter address bytes nvidiaPlanHostPageBytes)))
def bobTableBytes : Nat = (naturalMultiply bobCheckpointBytesPerScalar bobParameterCount)
def bobRowBytes : Nat = (naturalMultiply bobCheckpointBytesPerScalar bobVocabularySize)
-- a loss block: one binary32 per row
def bobLossBytes : Nat = bobRowBytes
def bobLogitRowsBytes : Nat = (naturalMultiply bobPromptCount bobRowBytes)
def bobPredictionsBytes : Nat = (naturalMultiply bobPromptCount bobPredictionBytes)
def bobParametersAddress : Nat = bobArenaBase
def bobFirstMomentAddress : Nat = (bobArenaNext bobParametersAddress bobTableBytes)
def bobSecondMomentAddress : Nat = (bobArenaNext bobFirstMomentAddress bobTableBytes)
def bobGradientsAddress : Nat = (bobArenaNext bobSecondMomentAddress bobTableBytes)
def bobInputAddress : Nat = (bobArenaNext bobGradientsAddress bobTableBytes)
def bobTargetsAddress : Nat = (naturalAdd bobInputAddress bobInputTargetsOffset)
def bobPromptBytesAddress : Nat = (naturalAdd bobInputAddress bobInputPromptsOffset)
def bobInitialLossAddress : Nat = (bobArenaNext bobInputAddress bobInputRecordBytes)
def bobTrainingLossAddress : Nat = (bobArenaNext bobInitialLossAddress bobLossBytes)
def bobLogitRowsAddress : Nat = (bobArenaNext bobTrainingLossAddress bobLossBytes)
def bobPredictionsAddress : Nat = (bobArenaNext bobLogitRowsAddress bobLogitRowsBytes)
def bobArenaBytes : Nat =
  (naturalSaturatingSubtract (naturalAdd bobPredictionsAddress bobPredictionsBytes) bobParametersAddress)
-- the result record: the initial loss block, the training loss block, the
-- logit rows
def bobResultInitialLossOffset : Nat = 0
def bobResultTrainingLossOffset : Nat = bobLossBytes
def bobResultLogitRowsOffset : Nat = (naturalMultiply 2 bobLossBytes)

-- log2 e, ln 2 and 1 / 256 as binary32. BigramLossGradientSM86 converts
-- the log2 reduction back to natural logarithms for each row's cross entropy;
-- inverse rows scales the gradients, not the stored row losses. The row
-- block must be reduced separately before reporting a mean objective. Each
-- constant is held to its exact value (Float32Exact).
def bobInverseRows : Nat = 0x3b800000

def bobInverseRowsDerived :
  (equal Nat bobInverseRows (float32ExactRational 1 bobVocabularySize)) =
  (refl Nat bobInverseRows)

def bobWordPair =
  (lambda unrestricted low : Nat .
    (lambda unrestricted high : Nat . (naturalAdd low (naturalMultiply high 4294967296))))

def bobLossBlock =
  (lambda unrestricted lossAddress : Nat .
    (bobBlock
      (bobPatch (bobArgument 0) (bobWord64 bobParametersAddress)
        (bobPatch (bobArgument 1) (bobWord64 bobGradientsAddress)
          (bobPatch (bobArgument 2) (bobWord64 bobTargetsAddress)
            (bobPatch (bobArgument 3) (bobWord64 lossAddress)
              (bobPatch (bobArgument 6) (bobWord64 (bobWordPair float32Log2E float32Ln2))
                (bobPatch (bobArgument 7) (bobWord64 (bobWordPair bobInverseRows 0)) bobNoPatches))))))))

def bobGatherBlock =
  (lambda unrestricted index : Nat .
    (bobBlock
      (bobPatch (bobArgument 0) (bobWord64 (naturalAdd bobLogitRowsAddress (naturalMultiply index bobRowBytes)))
        (bobPatch (bobArgument 1) (bobWord64 (naturalAdd bobPromptBytesAddress (naturalMultiply index bobInputWordBytes)))
          (bobPatch (bobArgument 2) (bobWord64 bobParametersAddress) bobNoPatches)))))

def bobArgmaxBlock =
  (lambda unrestricted index : Nat .
    (bobBlock
      (bobPatch (bobArgument 0) (bobWord64 (naturalAdd bobPredictionsAddress (naturalMultiply index bobPredictionBytes)))
        (bobPatch (bobArgument 1) (bobWord64 (naturalAdd bobLogitRowsAddress (naturalMultiply index bobRowBytes)))
          bobNoPatches))))

-- AdamW's nine f32 scalars for one step, read after its four addresses.
def bobAdamBlock =
  (lambda unrestricted scalars : Bytes .
    (bobBlock
      (bobPatch (bobArgument 0) (bobWord64 bobParametersAddress)
        (bobPatch (bobArgument 1) (bobWord64 bobGradientsAddress)
          (bobPatch (bobArgument 2) (bobWord64 bobFirstMomentAddress)
            (bobPatch (bobArgument 3) (bobWord64 bobSecondMomentAddress)
              (bobWordPatches 5 (bobArgument 4) (bytes-append scalars b"\x00\x00\x00\x00") bobNoPatches)))))))

def bobLaunch =
  (lambda unrestricted identity : Bytes .
    (lambda unrestricted kernel : (family NvidiaLaunchKernel) .
      (lambda unrestricted block : (family NvidiaParameterBlock) .
        (constructor NvidiaLaunchSchedule NvidiaLaunchScheduleOne
          (constructor NvidiaLaunchTemplate NvidiaLaunchTemplateValue identity kernel block)))))

def bobThen =
  (lambda unrestricted left : (family NvidiaLaunchSchedule) .
    (lambda unrestricted right : (family NvidiaLaunchSchedule) .
      (constructor NvidiaLaunchSchedule NvidiaLaunchScheduleAppend left right)))

def bobLossLaunch =
  (lambda unrestricted lossAddress : Nat .
    (bobLaunch b"bob-loss-gradient" (bobKernel bobLossPlace 1) (bobLossBlock lossAddress)))

-- Predicting one prompt gathers its logits row and takes the greedy argmax.
def bobPredictLaunches =
  (lambda unrestricted index : Nat .
    (bobThen
      (bobLaunch b"bob-row-gather" (bobKernel bobGatherPlace 1) (bobGatherBlock index))
      (bobLaunch b"bob-greedy-argmax" (bobKernel bobArgmaxPlace 1) (bobArgmaxBlock index))))

-- One training step updates the table from the current gradients and then
-- evaluates the loss and the next gradients.
def bobTrainingStep =
  (lambda unrestricted scalars : Bytes .
    (bobThen
      (bobLaunch b"bob-adamw" (bobKernel bobAdamPlace 256) (bobAdamBlock scalars))
      (bobLossLaunch bobTrainingLossAddress)))

def bobTrainingSteps : Nat =
  bobUpdatesPerInvocation

-- Steps 1 .. count, in order, each with its own scalars.
def bobTrainingSchedule =
  (lambda unrestricted count : Nat .
    (nat-eliminate
      (lambda unrestricted current : Nat . (family NvidiaLaunchSchedule))
      (constructor NvidiaLaunchSchedule NvidiaLaunchScheduleEmpty)
      (lambda unrestricted predecessor : Nat .
        (lambda unrestricted induction : (family NvidiaLaunchSchedule) .
          (bobThen induction (bobTrainingStep (bobAdamStepScalars (succ predecessor))))))
      count))

-- Prompts 0 .. count - 1, in order.
def bobPredictionSchedule : (family NvidiaLaunchSchedule) =
  (nat-eliminate
    (lambda unrestricted current : Nat . (family NvidiaLaunchSchedule))
    (constructor NvidiaLaunchSchedule NvidiaLaunchScheduleEmpty)
    (lambda unrestricted predecessor : Nat .
      (lambda unrestricted induction : (family NvidiaLaunchSchedule) .
        (bobThen induction (bobPredictLaunches predecessor))))
    bobPromptCount)

def bobInitialLoss : (family NvidiaLaunchSchedule) =
  (bobLossLaunch bobInitialLossAddress)

-- the forward launches: the initial loss and the prompts' predictions,
-- before the request's operation decides whether training follows
def bobForwardSchedule : (family NvidiaLaunchSchedule) =
  (bobThen bobInitialLoss bobPredictionSchedule)

def bobLaunchSchedule : (family NvidiaLaunchSchedule) =
  (bobThen bobForwardSchedule (bobTrainingSchedule bobTrainingSteps))

def bobForwardLaunches : Nat = (nvidiaLaunchScheduleCount bobForwardSchedule)

-- the launch of update k's AdamW (from 0; first in its step): after the
-- forward launches, the earlier steps'.  The host writes each update's
-- step size and epsilon into its parameter block
-- (Platform.Linux.Nvidia.PlanHostAdamW); the values baked here are the
-- first invocation's.
def bobStepLaunches : Nat =
  (nvidiaLaunchScheduleCount (bobTrainingStep (bobAdamStepScalars 1)))

def bobAdamWLaunch =
  (lambda unrestricted k : Nat . (naturalAdd bobForwardLaunches (naturalMultiply k bobStepLaunches)))
def bobLaunchCount : Nat = (nvidiaLaunchScheduleCount bobLaunchSchedule)

-- Each launch is submitted alone.  The run submits every launch once, in
-- order, and then predicts again from the trained table.
def bobSubmitEach =
  (lambda unrestricted count : Nat .
    (lambda unrestricted firstLaunch : Nat .
      (lambda unrestricted firstSemaphore : Nat .
        (constructor NvidiaSubmissionSchedule NvidiaSubmissionScheduleRepeat
          count 1 prSemaphoreStride
          (constructor NvidiaSubmissionSchedule NvidiaSubmissionScheduleOne
            (constructor NvidiaSubmissionBatch NvidiaSubmissionBatchValue
              b"bob-submission"
              firstSemaphore
              (constructor NvidiaLaunchReferences NvidiaLaunchReferencesRange firstLaunch 1)))))))

def bobPredictionLaunches : Nat =
  (nvidiaLaunchScheduleCount bobPredictionSchedule)

def bobSubmissionSchedule : (family NvidiaSubmissionSchedule) =
  (constructor NvidiaSubmissionSchedule NvidiaSubmissionScheduleAppend
    (bobSubmitEach bobLaunchCount zero zero)
    (bobSubmitEach bobPredictionLaunches (nvidiaLaunchScheduleCount bobInitialLoss) (naturalMultiply bobLaunchCount prSemaphoreStride)))

def bobSubmissionCount : Nat =
  (naturalAdd bobLaunchCount bobPredictionLaunches)

-- The program component is exactly the retained code and one parameter block
-- per launch; the QMD table is one 256-byte record per launch.
def bobProgramBytes : Nat =
  (naturalAdd bobCodeExtent (naturalMultiply bobLaunchCount bobParameterBytes))

def bobQMDTableBytes : Nat =
  (naturalMultiply bobLaunchCount qmdRecordBytes)

def bobWholeProgramPlan =
  (lambda unrestricted component : (family NvidiaWholeProgramComponent) .
    (lambda unrestricted architecture : (family ModelWord32) .
      (lambda unrestricted computeClass : (family ModelWord32) .
        (lambda unrestricted spaVersion : (family ModelWord32) .
          (constructor
            NvidiaWholeProgramPlan
            NvidiaWholeProgramPlanValue
            component
            architecture
            computeClass
            spaVersion
            (bobWord64 bobProgramBase)
            (bobWord64 bobQMDBase)
            (bobWord64 bobPushbufferBase)
            (bobWord64 bobSemaphoreBase)
            bobRegionAlignment
            bobProgramBytes
            bobRegions
            bobLaunchSchedule
            bobLaunchCount
            bobQMDTableBytes
            bobSubmissionSchedule
            bobSubmissionCount
            bobSubmissionCount)))))

def bobPhysicalComponent =
  (lambda unrestricted component : (family NvidiaWholeProgramComponent) .
    (lambda unrestricted architecture : (family ModelWord32) .
      (lambda unrestricted computeClass : (family ModelWord32) .
        (lambda unrestricted spaVersion : (family ModelWord32) .
          (compiler-native-encode
            (family NvidiaWholeProgramPlan)
            (bobWholeProgramPlan component architecture computeClass spaVersion))))))

-- ---- the device side of the arena ----
-- Every tensor Bob places lives for the whole program (the arena above);
-- each launch's identity is its phase.  Runtime.DeviceArenaCertificate
-- decides that the tensors are disjoint and that every pointer a launch
-- passes lies in one of them; the host is admitted only then
-- (Bob.Build.NativeHost).
def bobPlane =
  (lambda unrestricted name : Bytes .
    (lambda unrestricted address : Nat .
      (lambda unrestricted bytes : Nat .
        (lambda unrestricted rest : (family ArenaResidents) .
          (constructor ArenaResidents ArenaResidentsNext
            (constructor ArenaResident ArenaResidentValue name address bytes bobRegionAlignment)
            rest)))))

def bobNoPlanes : (family ArenaResidents) =
  (constructor ArenaResidents ArenaResidentsEnd)

def bobDeviceWindows : (family ArenaResidents) =
  (bobPlane b"arena" bobParametersAddress bobArenaBytes bobNoPlanes)

def bobDeviceGlobals : (family ArenaResidents) =
  (bobPlane b"parameters" bobParametersAddress bobTableBytes
    (bobPlane b"first-moments" bobFirstMomentAddress bobTableBytes
      (bobPlane b"second-moments" bobSecondMomentAddress bobTableBytes
        (bobPlane b"gradients" bobGradientsAddress bobTableBytes
          (bobPlane b"input-record" bobInputAddress bobInputRecordBytes
            (bobPlane b"initial-loss" bobInitialLossAddress bobLossBytes
              (bobPlane b"training-loss" bobTrainingLossAddress bobLossBytes
                (bobPlane b"logit-rows" bobLogitRowsAddress bobLogitRowsBytes
                  (bobPlane b"predictions" bobPredictionsAddress bobPredictionsBytes bobNoPlanes)))))))))

def bobPhase =
  (lambda unrestricted identity : Bytes .
    (lambda unrestricted rest : (family DeviceArenaPhases) .
      (constructor DeviceArenaPhases DeviceArenaPhasesNext
        (constructor DeviceArenaPhase DeviceArenaPhaseValue identity bobNoPlanes bobNoPlanes)
        rest)))

def bobDevicePhases : (family DeviceArenaPhases) =
  (bobPhase b"bob-loss-gradient"
    (bobPhase b"bob-row-gather"
      (bobPhase b"bob-greedy-argmax"
        (bobPhase b"bob-adamw" (constructor DeviceArenaPhases DeviceArenaPhasesEnd)))))

-- 0 when the device side is certified; a function, so a build computes it
-- only where it is applied
def bobDeviceHazardOf =
  (lambda unrestricted ignored : Nat .
    (deviceArenaHazard bobDeviceWindows bobDeviceGlobals bobDevicePhases bobLaunchSchedule bobSubmissionSchedule))

def bobDeviceHazard : Nat =
  (bobDeviceHazardOf zero)