module Coppelius.ArenaPlan

import Compiler.Planning.VA
import Platform.Linux.Nvidia.PlanHost
import Coppelius.SM86Capability
import Model.Config
import Model.Parameter
import Model.Word64
import Std.Natural
import Std.Physical

-- Fixed addresses are part of this artifact ABI.  The generic address-plan
-- and UVM/RM owners remain shared; this value records only Coppelius placement.
family CoppeliusArenaPlan : Type 0
constructor CoppeliusArenaPlanValue
field unrestricted coppeliusArenaProfileIdentity : Bytes
field unrestricted coppeliusArenaVideoBase : DeviceAddress
field unrestricted coppeliusArenaVideoBytes : ByteCount
field unrestricted coppeliusArenaVideoEnd : DeviceAddress
field unrestricted coppeliusArenaAlignmentBytes : ByteAlignment
field unrestricted coppeliusArenaPersistentEndOffset : ByteCount
field unrestricted coppeliusArenaWorkspaceEndOffset : ByteCount
field unrestricted coppeliusArenaProgramBytes : ByteCount
field unrestricted coppeliusArenaQMDBytes : ByteCount
field unrestricted coppeliusArenaOutputStagingBytes : ByteCount
field unrestricted coppeliusArenaGPFIFOAndUSERDBytes : ByteCount
field unrestricted coppeliusArenaCheckpointStagingBytes : ByteCount
field unrestricted coppeliusArenaAllowedHostFallbacks : Nat

end-family

-- The naturals are the values the host derivation (Coppelius.Build.NativeHost)
-- places by; the typed values are the same numbers for the plan.
def coppeliusSM86CompatVideoBaseNatural : Nat =
  0x0000_0009_0000_0000

def coppeliusSM86CompatVideoBase : DeviceAddress =
  (stdDeviceAddress (modelWord64FromNaturalTruncated coppeliusSM86CompatVideoBaseNatural))

def coppeliusSM86CompatVideoBytesNatural : Nat =
  2_147_483_648

def coppeliusSM86CompatVideoBytes : ByteCount =
  (stdByteCount (modelWord64FromNaturalTruncated coppeliusSM86CompatVideoBytesNatural))

def coppeliusSM86CompatVideoEnd : DeviceAddress =
  (stdDeviceAddress 0x0000_0009_8000_0000)

def coppeliusSM86CompatArenaAlignment : ByteAlignment =
  (stdByteAlignment 2_097_152)

def coppeliusSM86CompatPersistentEndOffset : ByteCount =
  (stdByteCount 1_038_630_912)

def coppeliusSM86CompatWorkspaceEndOffset : ByteCount =
  (stdByteCount 1_207_205_888)

def coppeliusSM86CompatProgramBytesNatural : Nat =
  4_194_304

def coppeliusSM86CompatProgramBytes : ByteCount =
  (stdByteCount (modelWord64FromNaturalTruncated coppeliusSM86CompatProgramBytesNatural))

def coppeliusSM86CompatQMDBytesNatural : Nat =
  8_388_608

def coppeliusSM86CompatQMDBytes : ByteCount =
  (stdByteCount (modelWord64FromNaturalTruncated coppeliusSM86CompatQMDBytesNatural))

-- The host maps the first half of the QMD arena directly; any compiler-owned
-- overflow is mapped as the adjacent typed region.
def coppeliusSM86CompatQMDPrimaryBytesNatural : Nat =
  4_194_304

def coppeliusSM86CompatOutputStagingBytes : ByteCount =
  (stdByteCount 4_194_304)

def coppeliusSM86CompatGPFIFOAndUSERDBytesNatural : Nat =
  266_240

def coppeliusSM86CompatGPFIFOAndUSERDBytes : ByteCount =
  (stdByteCount (modelWord64FromNaturalTruncated coppeliusSM86CompatGPFIFOAndUSERDBytesNatural))

def coppeliusSM86CompatCheckpointStagingBytesNatural : Nat =
  4_194_304

def coppeliusSM86CompatCheckpointStagingBytes : ByteCount =
  (stdByteCount (modelWord64FromNaturalTruncated coppeliusSM86CompatCheckpointStagingBytesNatural))

def coppeliusSM86CompatArenaPlan : (family CoppeliusArenaPlan) =
  (record
    CoppeliusArenaPlan
    (coppeliusArenaProfileIdentity = coppeliusSM86CompatProfileIdentity)
    (coppeliusArenaVideoBase = coppeliusSM86CompatVideoBase)
    (coppeliusArenaVideoBytes = coppeliusSM86CompatVideoBytes)
    (coppeliusArenaVideoEnd = coppeliusSM86CompatVideoEnd)
    (coppeliusArenaAlignmentBytes = coppeliusSM86CompatArenaAlignment)
    (coppeliusArenaPersistentEndOffset = coppeliusSM86CompatPersistentEndOffset)
    (coppeliusArenaWorkspaceEndOffset = coppeliusSM86CompatWorkspaceEndOffset)
    (coppeliusArenaProgramBytes = coppeliusSM86CompatProgramBytes)
    (coppeliusArenaQMDBytes = coppeliusSM86CompatQMDBytes)
    (coppeliusArenaOutputStagingBytes = coppeliusSM86CompatOutputStagingBytes)
    (coppeliusArenaGPFIFOAndUSERDBytes = coppeliusSM86CompatGPFIFOAndUSERDBytes)
    (coppeliusArenaCheckpointStagingBytes = coppeliusSM86CompatCheckpointStagingBytes)
    (coppeliusArenaAllowedHostFallbacks = 0))

-- ---- placement ----
-- The buffers whose place the plan itself depends on, in the card's address
-- space from the device origin (Compiler.Planning.VA) and in the process from
-- the plan host's mapping origin (the arena itself is the UVM external range
-- above): the channel's GPFIFO; the staging window the kernels read their
-- input from and write their results to (Coppelius.TrainingRun lays it
-- out); the program region the device images are placed in
-- (Coppelius.Build.DeviceImages).  The buffers the plan sizes follow them
-- (Coppelius.Build.Placement).  The host side of the arena certificate
-- decides that all of them are disjoint and clear of the host's own
-- mappings.
def coppeliusDeviceOrigin : Nat = vaBaseNatural

def coppeliusDeviceSpaceAlignment : Nat =
  0x0100_0000

def coppeliusGPFIFOBase : Nat =
  coppeliusDeviceOrigin

def coppeliusStagingBase : Nat =
  (vaPlaceAfter
    coppeliusGPFIFOBase
    coppeliusSM86CompatGPFIFOAndUSERDBytesNatural
    coppeliusDeviceSpaceAlignment)

def coppeliusProgramBase : Nat =
  (vaPlaceAfter
    coppeliusStagingBase
    coppeliusSM86CompatCheckpointStagingBytesNatural
    coppeliusDeviceSpaceAlignment)

def coppeliusHostGPFIFO : Nat =
  nvidiaPlanHostMappingOrigin

def coppeliusHostErrorNotifier : Nat =
  (vaPlaceAfter
    coppeliusHostGPFIFO
    coppeliusSM86CompatGPFIFOAndUSERDBytesNatural
    nvidiaPlanHostMappingAlignment)

def coppeliusHostStaging : Nat =
  (vaPlaceAfter
    coppeliusHostErrorNotifier
    nvidiaPlanHostErrorPageBytes
    nvidiaPlanHostMappingAlignment)

def coppeliusHostProgram : Nat =
  (vaPlaceAfter
    coppeliusHostStaging
    coppeliusSM86CompatCheckpointStagingBytesNatural
    nvidiaPlanHostMappingAlignment)