module Realization.Nvidia.SM86.LinearStepPlanSM86

import Accelerator.SM86.Instruction
import Accelerator.SM86.Operands
import Accelerator.SM86.Types
import Accelerator.SM86.InstructionEncoding
import Hardware.Nvidia.SM86.Command.WholeProgramPlan
import Model.Config
import Model.Word64
import Realization.Nvidia.SM86.LinearStepSM86
import Runtime.ArenaCertificate
import Std.Natural

-- The device side of the checked linear step, declared the way Bob and
-- Coppelius declare theirs: a device arena with a placement certificate, one
-- typed code region, one launch whose parameter block names the four data
-- regions and the learning rate, and one submission.  The NVIDIA backend
-- (compiler-native-encode of the plan, at build time) derives the program
-- table, the QMD, the pushbuffer and the GPFIFO from this; nothing here
-- joins bytes or knows a wire layout.
--
-- The arena.  Every address below is arena base + resident offset, and the
-- residents' placement is decided by Runtime.ArenaCertificate in
-- Proof.CheckedLinearStep: aligned, inside the arena, pairwise disjoint.
def linearStepPlanArenaBase : Nat = 0x0000_0008_0000_0000
def linearStepPlanArenaExtent : Nat = 0x0000_0000_2000_0000

def linearStepPlanPushbufferOffset : Nat = 0x0100_0000
def linearStepPlanSemaphoreOffset : Nat = 0x0200_0000
def linearStepPlanProgramOffset : Nat = 0x0300_0000
def linearStepPlanQMDOffset : Nat = 0x0400_0000
def linearStepPlanWeightsOffset : Nat = 0x1000_0000
def linearStepPlanInputOffset : Nat = 0x1001_0000
def linearStepPlanTargetOffset : Nat = 0x1002_0000
def linearStepPlanOutputOffset : Nat = 0x1003_0000

def linearStepPlanAddress =
  (lambda unrestricted offset : Nat . (naturalAdd linearStepPlanArenaBase offset))

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

-- Code: 56 instructions of 16 bytes, padded to the region alignment, then
-- one 512-byte parameter block for the one launch.
def linearStepPlanRegionAlignment : Nat = 256
def linearStepPlanInstructionBytes : Nat = 16
def linearStepPlanParameterBytes : Nat = 512

def linearStepPlanAlign =
  (lambda unrestricted value : Nat .
    (naturalMultiply
      (naturalDivideUnchecked
        (naturalAdd value (naturalSaturatingSubtract linearStepPlanRegionAlignment 1))
        linearStepPlanRegionAlignment)
      linearStepPlanRegionAlignment))

def linearStepPlanCodeBytes : Nat =
  (naturalMultiply linearStepSM86InstructionCount linearStepPlanInstructionBytes)

def linearStepPlanCodeExtent : Nat =
  (linearStepPlanAlign linearStepPlanCodeBytes)

def linearStepPlanProgramBytes : Nat =
  (naturalAdd linearStepPlanCodeExtent linearStepPlanParameterBytes)

def linearStepPlanQMDTableBytes : Nat = 256
def linearStepPlanSemaphoreStride : Nat = 32
def linearStepPlanPushbufferBudget : Nat = 0x0100_0000

-- Byte extents of the data residents for a shape: W m x k f32, x k f32,
-- t m f32, the output record; the checked 2 x 2 has 16, 8, 8, 44.
def linearStepPlanWeightsBytesFor = (lambda unrestricted m : Nat . (lambda unrestricted k : Nat . (naturalMultiply 4 (naturalMultiply m k))))
def linearStepPlanInputBytesFor = (lambda unrestricted k : Nat . (naturalMultiply 4 k))
def linearStepPlanTargetBytesFor = (lambda unrestricted m : Nat . (naturalMultiply 4 m))
def linearStepPlanWeightsBytes : Nat = (linearStepPlanWeightsBytesFor linearStepSM86Outputs linearStepSM86Inputs)
def linearStepPlanInputBytes : Nat = (linearStepPlanInputBytesFor linearStepSM86Inputs)
def linearStepPlanTargetBytes : Nat = (linearStepPlanTargetBytesFor linearStepSM86Outputs)

-- the residents and certificate for any shape: the same offsets, the
-- shape's extents, the shape's program length
def linearStepPlanProgramBytesFor =
  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
    (naturalAdd (linearStepPlanAlign (naturalMultiply (linearStepSM86InstructionCountFor m k) linearStepPlanInstructionBytes)) linearStepPlanParameterBytes)))

def linearStepPlanResident =
  (lambda unrestricted identity : Bytes .
    (lambda unrestricted offset : Nat .
      (lambda unrestricted extent : Nat .
        (lambda unrestricted alignment : Nat .
          (lambda unrestricted rest : (family ArenaResidents) .
            (constructor ArenaResidents ArenaResidentsNext
              (constructor ArenaResident ArenaResidentValue identity offset extent alignment)
              rest))))))

def linearStepPlanResidents : (family ArenaResidents) =
  (linearStepPlanResident b"pushbuffer" linearStepPlanPushbufferOffset linearStepPlanPushbufferBudget 4096
  (linearStepPlanResident b"semaphores" linearStepPlanSemaphoreOffset linearStepPlanSemaphoreStride 4096
  (linearStepPlanResident b"program" linearStepPlanProgramOffset linearStepPlanProgramBytes linearStepPlanRegionAlignment
  (linearStepPlanResident b"qmd" linearStepPlanQMDOffset linearStepPlanQMDTableBytes linearStepPlanRegionAlignment
  (linearStepPlanResident b"weights" linearStepPlanWeightsOffset linearStepPlanWeightsBytes 4
  (linearStepPlanResident b"input" linearStepPlanInputOffset linearStepPlanInputBytes 4
  (linearStepPlanResident b"target" linearStepPlanTargetOffset linearStepPlanTargetBytes 4
  (linearStepPlanResident b"output" linearStepPlanOutputOffset linearStepSM86OutputRecordBytes 4
    (constructor ArenaResidents ArenaResidentsEnd)))))))))

def linearStepPlanCertificate : Nat =
  (arenaCertificate linearStepPlanArenaExtent linearStepPlanResidents)

-- One typed region: the realization's machine bytes with its register and
-- block requirements.
def linearStepPlanRegions : (family NvidiaDeviceRegions) =
  (constructor NvidiaDeviceRegions NvidiaDeviceRegionsNext
    (constructor NvidiaDeviceRegion NvidiaDeviceRegionValue
      linearStepSM86Identity
      linearStepSM86MachineBytes
      linearStepSM86Registers
      linearStepSM86BlockWidth
      zero)
    (constructor NvidiaDeviceRegions NvidiaDeviceRegionsEnd))

-- One launch: one block of one warp, the four addresses and eta at the
-- realization's parameter offsets (64-bit words 44..47 and 50).
def linearStepPlanPatch =
  (lambda unrestricted byteOffset : Nat .
    (lambda unrestricted value : Nat .
      (lambda unrestricted rest : (family NvidiaParameterPatches) .
        (constructor NvidiaParameterPatches NvidiaParameterPatchesNext
          (constructor NvidiaParameterPatch NvidiaParameterPatchValue
            (naturalDivideUnchecked byteOffset 8)
            (linearStepPlanWord64 value))
          rest))))

-- Shape-independent parameter placement: data addresses come from the
-- admitted arena, not from the historical checked-step fixture's offsets.
def linearStepPlanParametersAt =
  (lambda unrestricted weights : Nat .
  (lambda unrestricted input : Nat .
  (lambda unrestricted target : Nat .
  (lambda unrestricted output : Nat .
  (lambda unrestricted learningRate : Nat .
    (constructor NvidiaParameterBlock NvidiaParameterBlockValue
      (linearStepPlanPatch linearStepSM86WeightsOffset weights
      (linearStepPlanPatch linearStepSM86InputOffset input
      (linearStepPlanPatch linearStepSM86TargetOffset target
      (linearStepPlanPatch linearStepSM86OutputOffset output
      (linearStepPlanPatch linearStepSM86LearningRateOffset learningRate
        (constructor NvidiaParameterPatches NvidiaParameterPatchesEnd))))))))))))

def linearStepPlanParameters =
  (linearStepPlanParametersAt
    (linearStepPlanAddress linearStepPlanWeightsOffset)
    (linearStepPlanAddress linearStepPlanInputOffset)
    (linearStepPlanAddress linearStepPlanTargetOffset)
    (linearStepPlanAddress linearStepPlanOutputOffset))

def linearStepPlanKernel : (family NvidiaLaunchKernel) =
  (constructor NvidiaLaunchKernel NvidiaLaunchKernelValue
    (linearStepPlanWord64 (linearStepPlanAddress linearStepPlanProgramOffset))
    1 1 1 1 1 linearStepPlanParameterBytes)

def linearStepPlanLaunches =
  (lambda unrestricted learningRate : Nat .
    (constructor NvidiaLaunchSchedule NvidiaLaunchScheduleOne
      (constructor NvidiaLaunchTemplate NvidiaLaunchTemplateValue
        linearStepSM86Identity
        linearStepPlanKernel
        (linearStepPlanParameters learningRate))))

def linearStepPlanLaunchCount : Nat = 1
def linearStepPlanSubmissionCount : Nat = 1

def linearStepPlanSubmissions : (family NvidiaSubmissionSchedule) =
  (constructor NvidiaSubmissionSchedule NvidiaSubmissionScheduleOne
    (constructor NvidiaSubmissionBatch NvidiaSubmissionBatchValue
      b"linear-step-submission"
      zero
      (constructor NvidiaLaunchReferences NvidiaLaunchReferencesRange zero 1)))

-- The same launch submitted `count` times, one semaphore slot per
-- submission: the step is idempotent, so every submission rewrites the same
-- record, and the host waits on the last slot.  The semaphore resident grows
-- with the count; the certificate is decided for the count in use.
def linearStepPlanSubmissionsRepeated =
  (lambda unrestricted count : Nat .
    (constructor NvidiaSubmissionSchedule NvidiaSubmissionScheduleRepeat
      count zero linearStepPlanSemaphoreStride linearStepPlanSubmissions))

def linearStepPlanResidentsFor =
  (lambda unrestricted submissions : Nat .
    (linearStepPlanResident b"pushbuffer" linearStepPlanPushbufferOffset linearStepPlanPushbufferBudget 4096
    (linearStepPlanResident b"semaphores" linearStepPlanSemaphoreOffset (naturalMultiply linearStepPlanSemaphoreStride submissions) 4096
    (linearStepPlanResident b"program" linearStepPlanProgramOffset linearStepPlanProgramBytes linearStepPlanRegionAlignment
    (linearStepPlanResident b"qmd" linearStepPlanQMDOffset linearStepPlanQMDTableBytes linearStepPlanRegionAlignment
    (linearStepPlanResident b"weights" linearStepPlanWeightsOffset linearStepPlanWeightsBytes 4
    (linearStepPlanResident b"input" linearStepPlanInputOffset linearStepPlanInputBytes 4
    (linearStepPlanResident b"target" linearStepPlanTargetOffset linearStepPlanTargetBytes 4
    (linearStepPlanResident b"output" linearStepPlanOutputOffset linearStepSM86OutputRecordBytes 4
      (constructor ArenaResidents ArenaResidentsEnd))))))))))

def linearStepPlanCertificateFor =
  (lambda unrestricted submissions : Nat .
    (arenaCertificate linearStepPlanArenaExtent (linearStepPlanResidentsFor submissions)))

def linearStepPlanResidentsForShape =
  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
    (lambda unrestricted submissions : Nat .
      (linearStepPlanResident b"pushbuffer" linearStepPlanPushbufferOffset linearStepPlanPushbufferBudget 4096
      (linearStepPlanResident b"semaphores" linearStepPlanSemaphoreOffset (naturalMultiply linearStepPlanSemaphoreStride submissions) 4096
      (linearStepPlanResident b"program" linearStepPlanProgramOffset (linearStepPlanProgramBytesFor m k) linearStepPlanRegionAlignment
      (linearStepPlanResident b"qmd" linearStepPlanQMDOffset linearStepPlanQMDTableBytes linearStepPlanRegionAlignment
      (linearStepPlanResident b"weights" linearStepPlanWeightsOffset (linearStepPlanWeightsBytesFor m k) 4
      (linearStepPlanResident b"input" linearStepPlanInputOffset (linearStepPlanInputBytesFor k) 4
      (linearStepPlanResident b"target" linearStepPlanTargetOffset (linearStepPlanTargetBytesFor m) 4
      (linearStepPlanResident b"output" linearStepPlanOutputOffset (linearStepSM86OutputRecordBytesFor m k) 4
        (constructor ArenaResidents ArenaResidentsEnd))))))))))))

def linearStepPlanCertificateForShape =
  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
    (lambda unrestricted submissions : Nat .
      (arenaCertificate linearStepPlanArenaExtent (linearStepPlanResidentsForShape m k submissions)))))

-- a program's machine bytes by the compiler's encoder (none when it refuses)
def linearStepPlanMachineBytesOf =
  (lambda unrestricted program : (family SM86Program) .
    (eliminate SM86ProgramEncodingResult
      (lambda unrestricted current : (family SM86ProgramEncodingResult) . Bytes)
      (sm86EncodeProgram program)
      (branch SM86ProgramEncodingSucceeded bytes telemetry . bytes)
      (branch SM86ProgramEncodingFailed index failure telemetry . b"")))

-- a program's instructions, and the program region it needs: its code
-- aligned, then the parameter block -- for the generator's program, the
-- shape's (Proof.LinearStepGeneratorForall: its length is the count)
def linearStepPlanInstructionsOf =
  (lambda unrestricted program : (family SM86Program) .
    (eliminate SM86Program (lambda unrestricted current : (family SM86Program) . Nat) program
      (branch SM86ProgramEnd . zero)
      (branch SM86ProgramNext head tail induction . (succ induction))))

def linearStepPlanProgramBytesOf =
  (lambda unrestricted program : (family SM86Program) .
    (naturalAdd (linearStepPlanAlign (naturalMultiply (linearStepPlanInstructionsOf program) linearStepPlanInstructionBytes)) linearStepPlanParameterBytes))

-- the residents and certificate for a shape's data and a program's region
def linearStepPlanResidentsForProgram =
  (lambda unrestricted program : (family SM86Program) .
  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
    (lambda unrestricted submissions : Nat .
      (linearStepPlanResident b"pushbuffer" linearStepPlanPushbufferOffset linearStepPlanPushbufferBudget 4096
      (linearStepPlanResident b"semaphores" linearStepPlanSemaphoreOffset (naturalMultiply linearStepPlanSemaphoreStride submissions) 4096
      (linearStepPlanResident b"program" linearStepPlanProgramOffset (linearStepPlanProgramBytesOf program) linearStepPlanRegionAlignment
      (linearStepPlanResident b"qmd" linearStepPlanQMDOffset linearStepPlanQMDTableBytes linearStepPlanRegionAlignment
      (linearStepPlanResident b"weights" linearStepPlanWeightsOffset (linearStepPlanWeightsBytesFor m k) 4
      (linearStepPlanResident b"input" linearStepPlanInputOffset (linearStepPlanInputBytesFor k) 4
      (linearStepPlanResident b"target" linearStepPlanTargetOffset (linearStepPlanTargetBytesFor m) 4
      (linearStepPlanResident b"output" linearStepPlanOutputOffset (linearStepSM86OutputRecordBytesFor m k) 4
        (constructor ArenaResidents ArenaResidentsEnd)))))))))))))

def linearStepPlanCertificateForProgram =
  (lambda unrestricted program : (family SM86Program) .
  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
    (lambda unrestricted submissions : Nat .
      (arenaCertificate linearStepPlanArenaExtent (linearStepPlanResidentsForProgram program m k submissions))))))

-- the plan for a shape and a program for it -- the generator's, one a
-- proposer rewrote, or another realization's (the looped one): the program
-- region is the program's; the same offsets and launch shape. Keep this typed
-- until whole-program encoding: arena placement must be checkable without
-- evaluating a build-only machine-image intrinsic.
def linearStepPlanRegionsForProgram =
  (lambda unrestricted program : (family SM86Program) .
    (constructor NvidiaDeviceRegions NvidiaDeviceRegionsNext
      (constructor NvidiaDeviceRegion NvidiaDeviceProgramRegion
        linearStepSM86Identity
        program
        (sm86RegisterDemand program)
        linearStepSM86BlockWidth
        zero
        (constructor NvidiaDeviceRealization NvidiaDeviceRealizationRefused
          b"linear step requires an established SM86 realization"))
      (constructor NvidiaDeviceRegions NvidiaDeviceRegionsEnd)))

def linearStepPlanForProgram =
  (lambda unrestricted program : (family SM86Program) .
  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
  (lambda unrestricted component : (family NvidiaWholeProgramComponent) .
    (lambda unrestricted architecture : (family ModelWord32) .
      (lambda unrestricted computeClass : (family ModelWord32) .
        (lambda unrestricted spaVersion : (family ModelWord32) .
          (lambda unrestricted learningRate : Nat .
          (lambda unrestricted submissions : Nat .
            (constructor NvidiaWholeProgramPlan NvidiaWholeProgramPlanValue
              component
              architecture
              computeClass
              spaVersion
              (linearStepPlanWord64 (linearStepPlanAddress linearStepPlanProgramOffset))
              (linearStepPlanWord64 (linearStepPlanAddress linearStepPlanQMDOffset))
              (linearStepPlanWord64 (linearStepPlanAddress linearStepPlanPushbufferOffset))
              (linearStepPlanWord64 (linearStepPlanAddress linearStepPlanSemaphoreOffset))
              linearStepPlanRegionAlignment
              (linearStepPlanProgramBytesOf program)
              (linearStepPlanRegionsForProgram program)
              (linearStepPlanLaunches learningRate)
              linearStepPlanLaunchCount
              linearStepPlanQMDTableBytes
              (linearStepPlanSubmissionsRepeated submissions)
              submissions
              submissions))))))))))

def linearStepPlanComponentForProgram =
  (lambda unrestricted program : (family SM86Program) .
  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
  (lambda unrestricted component : (family NvidiaWholeProgramComponent) .
    (lambda unrestricted architecture : (family ModelWord32) .
      (lambda unrestricted computeClass : (family ModelWord32) .
        (lambda unrestricted spaVersion : (family ModelWord32) .
          (lambda unrestricted learningRate : Nat .
          (lambda unrestricted submissions : Nat .
            (compiler-native-encode
              (family NvidiaWholeProgramPlan)
              (linearStepPlanForProgram program m k component architecture computeClass spaVersion learningRate submissions)))))))))))

-- the plan for any shape: the generator's program for it
def linearStepPlanRegionsFor =
  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
    (linearStepPlanRegionsForProgram (linearStepSM86ProgramFor m k))))
def linearStepPlanFor =
  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
    (linearStepPlanForProgram (linearStepSM86ProgramFor m k) m k)))
def linearStepPlanComponentFor =
  (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
    (linearStepPlanComponentForProgram (linearStepSM86ProgramFor m k) m k)))

-- the checked 2 x 2: the generator's plan for that shape
def linearStepPlan = (linearStepPlanFor linearStepSM86Outputs linearStepSM86Inputs)

-- The backend boundary: the compiler realizes one component of the plan.
def linearStepPlanComponent =
  (lambda unrestricted component : (family NvidiaWholeProgramComponent) .
    (lambda unrestricted architecture : (family ModelWord32) .
      (lambda unrestricted computeClass : (family ModelWord32) .
        (lambda unrestricted spaVersion : (family ModelWord32) .
          (lambda unrestricted learningRate : Nat .
          (lambda unrestricted submissions : Nat .
            (compiler-native-encode
              (family NvidiaWholeProgramPlan)
              (linearStepPlan component architecture computeClass spaVersion learningRate submissions))))))))