module Coppelius.Build.DeviceImages

import Coppelius.ArenaPlan
import Coppelius.SM86Capability
import Accelerator.SM86.InstructionEncoding
import Accelerator.SM86.Instruction
import Accelerator.SM86.Immediate
import Accelerator.SM86.Operands
import Accelerator.SM86.Program
import Accelerator.SM86.HalfFormat
import Accelerator.SM121.Lowering
import Coppelius.Learner
import Coppelius.Model
import Coppelius.TrainingRun
import Accelerator.SM86.RNG.RandomNormal
import Data.Bytes
import Hardware.Nvidia.SM86.Command.WholeProgramPlan
import Realization.Nvidia.SM86.AdamWHalfSM86
import Realization.Nvidia.SM86.AttentionHeadLayoutSM86
import Realization.Nvidia.SM86.Cast.SM86
import Realization.Nvidia.SM86.RotaryTableSM86
import Realization.Nvidia.SM86.Elementwise.GeneralSM86
import Realization.Nvidia.SM86.ElementwiseVectorSM86
import Realization.Nvidia.SM86.EmbeddingGatherF32SM86
import Realization.Nvidia.SM86.ExactCrossEntropyBackwardSM86
import Realization.Nvidia.SM86.ExactCrossEntropySM86
import Realization.Nvidia.SM86.GELUSM86
import Realization.Nvidia.SM86.GatedGELUSM86
import Realization.Nvidia.SM86.GEMM.HMMA.RegisterTileSM86
import Realization.Nvidia.SM86.HMMAProductionSM86
import Realization.Nvidia.SM86.GreedyArgmaxSM86
import Realization.Nvidia.SM86.IndexedRowScatterOrderedSM86
import Realization.Nvidia.SM86.IndexedRowScatterSM86
import Realization.Nvidia.SM86.LayerNormSM86
import Realization.Nvidia.SM86.SoftmaxSM86
import Realization.Nvidia.SM86.StallCompaction
import Realization.Nvidia.SM86.StreamingAttentionSM86
import Realization.Nvidia.SM86.TargetGatherSM86
import Realization.Nvidia.SM86.TiledProductSM86
import Coppelius.Build.TiledChoices
import SM86.Scoreboard
import Realization.Nvidia.SM86.Transpose.FP16SM86
import Std.List
import Std.Foundation
import Std.Natural
import Std.Word

-- A device region remains an Alpha value until the final physical artifact is
-- assembled.  The resource fields travel with the encoded program so QMD
-- construction cannot accidentally select a program address independently of
-- its register, block and shared-memory requirements.
family CoppeliusDeviceImage : Type 0
constructor CoppeliusDeviceImageValue
field unrestricted coppeliusDeviceImageIdentity : Bytes
field unrestricted coppeliusDeviceImageProgram : (family SM86Program)
field unrestricted coppeliusDeviceImageMaterial : Bytes
field unrestricted coppeliusDeviceImageRegisters : Nat
field unrestricted coppeliusDeviceImageBlockX : Nat
field unrestricted coppeliusDeviceImageSharedBytes : Nat

end-family

family CoppeliusDeviceImages : Type 0
constructor CoppeliusDeviceImagesEnd
constructor CoppeliusDeviceImagesNext
field unrestricted coppeliusDeviceImagesHead : (family CoppeliusDeviceImage)
recursive unrestricted coppeliusDeviceImagesTail

end-family

family CoppeliusDeviceImagesBuildResult : Type 0
constructor CoppeliusDeviceImagesBuildReady
field unrestricted coppeliusDeviceImagesBuildBytes : BytesBuilder
field unrestricted coppeliusDeviceImagesBuildCursor : Nat
field unrestricted coppeliusDeviceImagesBuildCount : Nat
constructor CoppeliusDeviceImagesBuildFailed
field unrestricted coppeliusDeviceImagesBuildFailureIdentity : Bytes

end-family

-- a product the model launches (the tiled products below)
family CoppeliusTiledShape : Type 0
constructor CoppeliusTiledShapeValue
field unrestricted coppeliusTiledShapeA : (family TiledProductLayout)
field unrestricted coppeliusTiledShapeB : (family TiledProductLayout)
field unrestricted coppeliusTiledShapeM : Nat
field unrestricted coppeliusTiledShapeN : Nat
field unrestricted coppeliusTiledShapeK : Nat
-- 1: C = A B + R (a residual added in the epilogue)
field unrestricted coppeliusTiledShapeResidual : Nat
end-family

def coppeliusDeviceImagesEnd : (family CoppeliusDeviceImages) =
  (constructor CoppeliusDeviceImages CoppeliusDeviceImagesEnd)

def coppeliusDeviceImagesNext =
  (lambda unrestricted image : (family CoppeliusDeviceImage) .
    (lambda unrestricted tail : (family CoppeliusDeviceImages) .
      (constructor CoppeliusDeviceImages CoppeliusDeviceImagesNext image tail)))

def coppeliusDeviceImagesAppend =
  (lambda unrestricted left : (family CoppeliusDeviceImages) .
    (lambda unrestricted right : (family CoppeliusDeviceImages) .
      (eliminate
        CoppeliusDeviceImages
        (lambda unrestricted current : (family CoppeliusDeviceImages) .
          (family CoppeliusDeviceImages))
        left
        (branch CoppeliusDeviceImagesEnd . right)
        (branch
          CoppeliusDeviceImagesNext
          image
          tail
          induction
          .
          (coppeliusDeviceImagesNext image induction)))))

-- Preserve the typed identity and launch-resource facts until the NVIDIA
-- backend has seen the complete launch graph.  The old byte-only fold remains
-- as an oracle helper, but release realization consumes this target-owned
-- region vocabulary so unreachable images can be removed before QMD emission.
def coppeliusNvidiaDeviceRegions =
  (lambda unrestricted images : (family CoppeliusDeviceImages) .
    (eliminate
      CoppeliusDeviceImages
      (lambda unrestricted current : (family CoppeliusDeviceImages) .
        (family NvidiaDeviceRegions))
      images
      (branch
        CoppeliusDeviceImagesEnd
        .
        (constructor NvidiaDeviceRegions NvidiaDeviceRegionsEnd))
      (branch
        CoppeliusDeviceImagesNext
        image
        tail
        induction
        .
        (eliminate
          CoppeliusDeviceImage
          (lambda unrestricted current : (family CoppeliusDeviceImage) .
            (family NvidiaDeviceRegions))
          image
          (branch
            CoppeliusDeviceImageValue
            identity
            program
            material
            registers
            blockX
            sharedBytes
            .
            (constructor
              NvidiaDeviceRegions
              NvidiaDeviceRegionsNext
              (constructor
                NvidiaDeviceRegion
                NvidiaDeviceProgramRegion
                identity
                program
                registers
                blockX
                sharedBytes
                (sm121LowerRealization program registers))
              induction))))))

def coppeliusDeviceZeroBytes =
  (lambda unrestricted count : Nat .
    (nat-eliminate
      (lambda unrestricted current : Nat . Bytes)
      b""
      (lambda unrestricted predecessor : Nat .
        (lambda unrestricted induction : Bytes .
          (bytes-cons (byte 0) induction)))
      count))

def coppeliusDeviceAlignment : Nat =
  256

def coppeliusDevicePadding =
  (lambda unrestricted cursor : Nat .
    (naturalModuloUnchecked
      (naturalSaturatingSubtract
        coppeliusDeviceAlignment
        (naturalModuloUnchecked cursor coppeliusDeviceAlignment))
      coppeliusDeviceAlignment))

def coppeliusBuildDeviceImagesFrom =
  (lambda unrestricted images : (family CoppeliusDeviceImages) .
    (eliminate
      CoppeliusDeviceImages
      (lambda unrestricted current : (family CoppeliusDeviceImages) .
        (pi unrestricted builder : BytesBuilder .
          (pi unrestricted cursor : Nat .
            (pi unrestricted count : Nat .
              (family CoppeliusDeviceImagesBuildResult)))))
      images
      (branch
        CoppeliusDeviceImagesEnd
        .
        (lambda unrestricted builder : BytesBuilder .
          (lambda unrestricted cursor : Nat .
            (lambda unrestricted count : Nat .
              (constructor
                CoppeliusDeviceImagesBuildResult
                CoppeliusDeviceImagesBuildReady
                builder
                cursor
                count)))))
      (branch
        CoppeliusDeviceImagesNext
        image
        tail
        induction
        .
        (lambda unrestricted builder : BytesBuilder .
          (lambda unrestricted cursor : Nat .
            (lambda unrestricted count : Nat .
              (eliminate
                CoppeliusDeviceImage
                (lambda unrestricted current : (family CoppeliusDeviceImage) .
                  (family CoppeliusDeviceImagesBuildResult))
                image
                (branch
                  CoppeliusDeviceImageValue
                  identity
                  program
                  material
                  registers
                  blockX
                  sharedBytes
                  .
                  (nat-eliminate
                    (lambda unrestricted nonempty : Nat .
                      (family CoppeliusDeviceImagesBuildResult))
                    (constructor
                      CoppeliusDeviceImagesBuildResult
                      CoppeliusDeviceImagesBuildFailed
                      identity)
                    (lambda unrestricted materialPredecessor : Nat .
                      (lambda unrestricted materialInduction : (family CoppeliusDeviceImagesBuildResult) .
                        (let unrestricted padding = (coppeliusDevicePadding cursor)
                        in
                          (induction
                            (bytes-builder-append
                              builder
                              (bytes-builder-append
                                (bytes-builder-chunk (coppeliusDeviceZeroBytes padding))
                                (bytes-builder-chunk material)))
                            (naturalAdd cursor (naturalAdd padding (bytes-length material)))
                            (succ count)))))
                    (naturalNonzero (bytes-length material)))))))))))

def coppeliusBuildDeviceImages =
  (lambda unrestricted images : (family CoppeliusDeviceImages) .
    (coppeliusBuildDeviceImagesFrom images (bytes-builder-empty) zero zero))

-- Build-result projections deliberately return empty bytes on refusal.  The
-- list fold above converts that sentinel into a named failed region, and the
-- final artifact projection therefore fails closed rather than publishing a
-- truncated program table.
def coppeliusProgramImage =
  (lambda unrestricted program : (family SM86Program) .
    (eliminate
      SM86ProgramEncodingResult
      (lambda unrestricted current : (family SM86ProgramEncodingResult) . Bytes)
      (sm86EncodeProgram program)
      (branch SM86ProgramEncodingSucceeded image telemetry . image)
      (branch SM86ProgramEncodingFailed ordinal failure telemetry . b"")))

def coppeliusTransposeHeadTile : Nat = 64

-- Preserve already accepted schedules byte for byte. The RTX 3090 gather
-- launch completed with its original controls but stalled when every LDG
-- was given a read barrier. Repair only programs the scoreboard refuses,
-- then admit the placed program through the same checker below.
def coppeliusGuardLateReadsIfRequired =
  (lambda unrestricted program : (family SM86Program) .
    (nat-eliminate (lambda unrestricted accepted : Nat . (family SM86Program))
      (sm86GuardLateReads program)
      (lambda unrestricted predecessor : Nat .
        (lambda unrestricted induction : (family SM86Program) . program))
      (naturalIsZero (sm86Scoreboard program))))

-- An image: its typed program, and the program's SM86 encoding, which
-- places it (Coppelius.Build.Graph names launches by that placement); the
-- backend realizes the region in the target's machine code.
-- Its program's halves are in the plan's format (Coppelius.Learner).
def coppeliusDeviceImage =
  (lambda unrestricted identity : Bytes .
    (lambda unrestricted program : (family SM86Program) .
      (lambda unrestricted registers : Nat .
        (lambda unrestricted blockX : Nat .
          (lambda unrestricted sharedBytes : Nat .
            (let unrestricted halves = (sm86ProgramWithHalfFormat coppeliusHalfFormat program) in
            (let unrestricted guarded = (coppeliusGuardLateReadsIfRequired halves) in
            (constructor
              CoppeliusDeviceImage
              CoppeliusDeviceImageValue
              identity
              guarded
              (coppeliusProgramImage guarded)
              registers
              blockX
              sharedBytes))))))))

def coppeliusElementwiseImage =
  (lambda unrestricted operation : (family Elementwise256) .
    (emitElementwise256SM86 operation))

def coppeliusArgmaxImage =
  (lambda unrestricted artifact : (family GreedyArgmaxSM86Artifact) .
    (greedyArgmaxSM86ProgramFor greedyArgmaxSM86PromotedVocabulary))

def coppeliusCastImage =
  (lambda unrestricted direction : (family CastDirection) .
    (emitCastSM86 direction))

def coppeliusVectorImage =
  (lambda unrestricted kind : (family ElementwiseVectorSM86Kind) .
    (elementwiseVectorProgram kind))

def coppeliusGatherImage : (family SM86Program) =
  (embeddingGatherF32ProgramFor
    (constructor EmbeddingGatherF32Width EmbeddingGatherF32Width512))

def coppeliusLayerNormImage =
  (lambda unrestricted variant : (family LayerNormSM86Variant) .
    (layerNormSM86Program variant 1024))

def coppeliusGELUImage =
  (lambda unrestricted kind : (family GELUSM86ProgramKind) .
    (geluSM86Program kind))

def coppeliusAttentionImage =
  (lambda unrestricted kind : (family AttentionHeadLayoutSM86Kind) .
    (attentionHeadLayoutSM86Program kind))

def coppeliusSoftmaxImage =
  (lambda unrestricted variant : (family SoftmaxSM86Variant) .
    (softmaxSM86ProgramFor variant))

def coppeliusCrossEntropyImage : (family SM86Program) =
  (exactCrossEntropySM86ProgramFor
    (constructor ExactCrossEntropySM86Variant ExactCrossEntropyRows))

def coppeliusCrossEntropyBackwardImage : (family SM86Program) =
  exactCrossEntropyBackwardProgram

-- the embedding's gradient collected in a fixed order (the rows of the
-- positions naming a token added in position order, by the block of the
-- first: Realization.Nvidia.SM86.IndexedRowScatterOrderedSM86), so training
-- takes the same bits every run; the geometry is the atomic form's
def coppeliusScatterGeometry : (family IndexedRowScatterSM86Geometry) =
  (constructor IndexedRowScatterSM86Geometry IndexedRowScatterTokenEmbeddingGradientRows512x1024)

def coppeliusScatterImage : (family SM86Program) =
  (indexedRowScatterOrderedSM86Program
    (indexedRowScatterSM86WidthNatural coppeliusScatterGeometry)
    (indexedRowScatterSM86RowsNatural coppeliusScatterGeometry))

-- AdamW over the whole flat bank in one launch, the
-- parameters' half copies written with it (Realization.Nvidia.SM86.
-- AdamWHalfSM86, its stalls compacted): a grid of coppeliusAdamWBlocks
-- blocks walks the bank's element quads with a grid stride.  The bank's
-- elements are the checkpoint's (it is P, M and V, each the parameter
-- bank); the plan holds its own bank to the same count (Coppelius.Build.
-- Graph.cgAdamWCoversTheBank).
def coppeliusAdamWBlocks : Nat = 1024
def coppeliusAdamWElements : Nat =
  (naturalDivideUnchecked coppeliusCheckpointPayloadBytesNatural 12)
def coppeliusAdamWQuads : Nat = (naturalDivideUnchecked coppeliusAdamWElements adamWHalfSM86Elements)
def coppeliusAdamWStride : Nat = (naturalMultiply coppeliusAdamWBlocks adamWHalfSM86Threads)
def coppeliusAdamWIterations : Nat =
  (adamWHalfSM86Iterations coppeliusAdamWQuads coppeliusAdamWStride)
def coppeliusAdamWImage : (family SM86Program) =
  (adamWHalfSM86 sm86CompactStallsDrained coppeliusAdamWQuads coppeliusAdamWStride)
-- what the schedule checks read (the body twice, no branch)
def coppeliusAdamWChecked : (family SM86Program) =
  (adamWHalfSM86Checked sm86CompactStallsDrained coppeliusAdamWQuads coppeliusAdamWStride)

-- the bank is whole quads, the grid stride reaches every quad (the
-- image's bound leaves the rest of the last pass alone), and the bound is
-- a 32-bit immediate
def coppeliusAdamWWholeQuads :
  (equal Nat (naturalMultiply adamWHalfSM86Elements coppeliusAdamWQuads) coppeliusAdamWElements) =
  (refl Nat coppeliusAdamWElements)
def coppeliusAdamWReachesEveryElement :
  (equal Nat (naturalLessOrEqual coppeliusAdamWQuads (naturalMultiply coppeliusAdamWIterations coppeliusAdamWStride)) 1) =
  (refl Nat 1)
def coppeliusAdamWBoundFits :
  (equal Nat (naturalLess coppeliusAdamWQuads 4294967296) 1) =
  (refl Nat 1)

-- the gated activation's forward and backward
-- (Realization.Nvidia.SM86.GatedGELUSM86), their stalls compacted
def coppeliusGatedGELUForwardImage : (family SM86Program) = (sm86CompactStalls gatedGELUForwardSM86)
def coppeliusGatedGELUBackwardImage : (family SM86Program) = (sm86CompactStalls gatedGELUBackwardSM86)

def coppeliusCoreDeviceImages : (family CoppeliusDeviceImages) =
  (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"random-normal" randomNormalSM86Program 24 256 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"rotary-table" rotaryTableSM86Program rotaryTableSM86RegisterCount 256 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"fill" (coppeliusVectorImage (constructor ElementwiseVectorSM86Kind ElementwiseVectorSM86Fill)) 16 256 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"copy" (coppeliusVectorImage (constructor ElementwiseVectorSM86Kind ElementwiseVectorSM86Copy)) 24 256 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"cast-down" (coppeliusCastImage (constructor CastDirection Float32ToFloat16)) 24 256 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"cast-up" (coppeliusCastImage (constructor CastDirection Float16ToFloat32)) 24 256 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"add" (coppeliusElementwiseImage (constructor Elementwise256 Add256)) 24 256 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"scale" (coppeliusElementwiseImage (constructor Elementwise256 Scale256)) 24 256 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"argmax" (coppeliusArgmaxImage greedyArgmaxSM86BuildPromoted) 64 256 96)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"gather" coppeliusGatherImage 40 128 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"layernorm-forward" (coppeliusLayerNormImage (constructor LayerNormSM86Variant LayerNormSM86Forward512)) 40 (layerNormSM86BlockX (constructor LayerNormSM86Variant LayerNormSM86Forward512)) 128)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"qk-rope" (coppeliusAttentionImage (constructor AttentionHeadLayoutSM86Kind AttentionHeadLayoutQKRotary)) 64 128 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"value-transpose" (coppeliusAttentionImage (constructor AttentionHeadLayoutSM86Kind AttentionHeadLayoutValueTranspose)) 24 512 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"head-to-token" (coppeliusAttentionImage (constructor AttentionHeadLayoutSM86Kind AttentionHeadLayoutHeadToToken)) 32 256 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"token-to-head" (coppeliusAttentionImage (constructor AttentionHeadLayoutSM86Kind AttentionHeadLayoutTokenToHead)) 32 256 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"inverse-rope-qkv" (coppeliusAttentionImage (constructor AttentionHeadLayoutSM86Kind AttentionHeadLayoutInverseRoPEQKVMerge)) 64 128 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"softmax-forward" (coppeliusSoftmaxImage (constructor SoftmaxSM86Variant SoftmaxSM86CausalForward1024)) 32 256 128)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"softmax-backward" (coppeliusSoftmaxImage (constructor SoftmaxSM86Variant SoftmaxSM86CausalBackward1024)) 32 256 128)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"layernorm-backward" (coppeliusLayerNormImage (constructor LayerNormSM86Variant LayerNormSM86InputBackward512)) 48 (layerNormSM86BlockX (constructor LayerNormSM86Variant LayerNormSM86InputBackward512)) 128)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"layernorm-parameters" (coppeliusLayerNormImage (constructor LayerNormSM86Variant LayerNormSM86ParameterGradient512)) 40 (layerNormSM86BlockX (constructor LayerNormSM86Variant LayerNormSM86ParameterGradient512)) 128)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"cross-entropy" coppeliusCrossEntropyImage 48 256 128)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"cross-entropy-backward" coppeliusCrossEntropyBackwardImage 40 256 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"embedding-scatter" coppeliusScatterImage 40 512 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"adamw" coppeliusAdamWImage adamWHalfSM86Registers adamWHalfSM86Threads 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"gated-gelu-forward" coppeliusGatedGELUForwardImage gatedGELUForwardSM86Registers gatedGELUSM86Threads 0)
    (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"gated-gelu-backward" coppeliusGatedGELUBackwardImage gatedGELUBackwardSM86Registers gatedGELUSM86Threads 0)
    coppeliusDeviceImagesEnd))))))))))))))))))))))))))

-- the attention's head layouts (the products read their operands in place)
def coppeliusTransposeDeviceImages : (family CoppeliusDeviceImages) =
  (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"transpose-g8-r1024-c64" (coppeliusFP16TransposeTiledProgram coppeliusTransposeHeadTile 8 1024 64) 40 (coppeliusFP16TransposeBlockThreads coppeliusTransposeHeadTile) 0)
    (coppeliusDeviceImagesNext
      (coppeliusDeviceImage b"transpose-g8-r64-c1024" (coppeliusFP16TransposeTiledProgram coppeliusTransposeHeadTile 8 64 1024) 40 (coppeliusFP16TransposeBlockThreads coppeliusTransposeHeadTile) 0)
      coppeliusDeviceImagesEnd))

-- the images every launch but the loss's uses, in the order the backend
-- places them
def coppeliusImportedDeviceImages : (family CoppeliusDeviceImages) =
  (coppeliusDeviceImagesAppend coppeliusCoreDeviceImages coppeliusTransposeDeviceImages)

-- the training loss's target-logit gather (Realization.Nvidia.SM86.
-- TargetGatherSM86), placed after them: it fills `correct` with each row's
-- target logit (docs/observability PRD item 1)
def coppeliusTargetGatherImage : (family SM86Program) = targetGatherSM86Coppelius
def coppeliusLossDeviceImages : (family CoppeliusDeviceImages) =
  (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"loss-target-gather" coppeliusTargetGatherImage (sm86RegisterDemand targetGatherSM86Coppelius) 256 0)
    coppeliusDeviceImagesEnd)

-- the heads' attention as streaming launches
-- (Realization.Nvidia.SM86.StreamingAttentionSM86) at the model's sequence
-- and heads, the forward and the backward's two, their stalls compacted
-- (Realization.Nvidia.SM86.StallCompaction)
def coppeliusStreamingImage =
  (lambda unrestricted kernel : (pi unrestricted seq : Nat . (pi unrestricted heads : Nat . (family SM86Program))) .
    (sm86CompactStalls (kernel (stdU32ToNatural modelCoppeliusBlock) (stdU32ToNatural modelCoppeliusHeads))))
def coppeliusStreamingDeviceImages : (family CoppeliusDeviceImages) =
  (coppeliusDeviceImagesNext
    (coppeliusDeviceImage b"attention-forward" (coppeliusStreamingImage streamingAttentionForwardSM86)
      streamingAttentionForwardSM86Registers streamingAttentionForwardSM86Threads 0)
    (coppeliusDeviceImagesNext
      (coppeliusDeviceImage b"attention-query" (coppeliusStreamingImage streamingAttentionQuerySM86)
        streamingAttentionQuerySM86Registers streamingAttentionQuerySM86Threads 0)
      (coppeliusDeviceImagesNext
        (coppeliusDeviceImage b"attention-key" (coppeliusStreamingImage streamingAttentionKeySM86)
          streamingAttentionKeySM86Registers streamingAttentionKeySM86Threads 0)
        coppeliusDeviceImagesEnd)))

-- ---- the tiled products ----
def stdListFromShapes = (lambda unrestricted head : (family CoppeliusTiledShape) .
  (lambda unrestricted tail : (family StdList (family CoppeliusTiledShape)) .
    (constructor StdList StdListCons (family CoppeliusTiledShape) head tail)))

-- C (m x n) = A B (Realization.Nvidia.SM86.TiledProductSM86), each operand
-- read where it is stored: A as rows (m x k) or as columns (A^T, k x m), B
-- as rows (B^T, n x k) or as columns (k x n).  The model's products are
-- three patterns: the forward's X W^T (rows, rows), a data gradient's dY W
-- (rows, columns) and a weight gradient's dY^T X (columns, columns).  Each
-- product's tile and stages are a candidate measured on the card and
-- recorded (Coppelius.Build.TiledChoices, scripts/alpha/tiled-choices.py)
-- among the candidates below.

def coppeliusTiledRows : (family TiledProductLayout) = (constructor TiledProductLayout TiledProductRows)
def coppeliusTiledColumns : (family TiledProductLayout) = (constructor TiledProductLayout TiledProductColumns)
def coppeliusTiledIsColumns = (lambda unrestricted layout : (family TiledProductLayout) .
  (eliminate TiledProductLayout (lambda unrestricted c : (family TiledProductLayout) . Nat) layout
    (branch TiledProductRows . 0) (branch TiledProductColumns . 1)))
-- the candidates, as scripts/alpha/tiled-choices.py numbers them (block M x
-- block N x stages, warps along M x N, each warp's rows x columns):
--   0  128 x 128 x 4   2 x 4   64 x 32
--   1  128 x 128 x 2   2 x 4   64 x 32
--   2   64 x  64 x 4   2 x 2   32 x 32
--   3   64 x  64 x 2   2 x 2   32 x 32
--   4  128 x  64 x 4   4 x 2   32 x 32   (48 KiB: two blocks an SM)
--   5   64 x 128 x 4   2 x 4   32 x 32
--   6  256 x 128 x 2   4 x 2   64 x 64   (the most registers the SM86 file allows)
--   7  128 x 256 x 2   2 x 4   64 x 64
-- Candidates over a card's shared memory or registers exist but never serve
-- (coppeliusTiledTileServes).  The cuBLAS tiles measured on the GB10
-- (research gemm/cublas-shapes.md: 128 x 96 .. 256 x 128, three stages) are
-- where 4 .. 7 came from.
def coppeliusTiledCandidateCount : Nat = 8
-- candidate i's entry of a column of eight
def coppeliusTiledPick = (lambda unrestricted i : Nat .
  (lambda unrestricted v0 : Nat . (lambda unrestricted v1 : Nat . (lambda unrestricted v2 : Nat . (lambda unrestricted v3 : Nat .
  (lambda unrestricted v4 : Nat . (lambda unrestricted v5 : Nat . (lambda unrestricted v6 : Nat . (lambda unrestricted v7 : Nat .
    (naturalSelect (naturalEqual i 0) v0 (naturalSelect (naturalEqual i 1) v1 (naturalSelect (naturalEqual i 2) v2
    (naturalSelect (naturalEqual i 3) v3 (naturalSelect (naturalEqual i 4) v4 (naturalSelect (naturalEqual i 5) v5
    (naturalSelect (naturalEqual i 6) v6 v7))))))))))))))))
def coppeliusTiledCandidate = (lambda unrestricted i : Nat .
  (tiledProductTile
    (coppeliusTiledPick i 2 2 2 2 4 2 4 2)
    (coppeliusTiledPick i 4 4 2 2 2 4 2 4)
    (coppeliusTiledPick i 4 4 2 2 2 2 4 4)
    (coppeliusTiledPick i 4 4 4 4 4 4 8 8)
    (coppeliusTiledPick i 4 2 4 2 4 4 2 2)))
-- a product's recorded candidate
def coppeliusTiledChoiceFor = (lambda unrestricted residual : Nat .
  (lambda unrestricted a : (family TiledProductLayout) . (lambda unrestricted b : (family TiledProductLayout) .
  (lambda unrestricted m : Nat . (lambda unrestricted n : Nat . (lambda unrestricted k : Nat .
    (coppeliusTiledChoice residual (coppeliusTiledIsColumns a) (coppeliusTiledIsColumns b) m n k)))))))
def coppeliusTiledTile = (lambda unrestricted residual : Nat .
  (lambda unrestricted a : (family TiledProductLayout) . (lambda unrestricted b : (family TiledProductLayout) .
  (lambda unrestricted m : Nat . (lambda unrestricted n : Nat . (lambda unrestricted k : Nat .
    (coppeliusTiledCandidate (coppeliusTiledChoiceFor residual a b m n k))))))))
-- 1 when a tile serves the product: it divides C, its stages divide K's
-- steps, its threads copy its operand tiles evenly, its registers fit the
-- SM86 file, and its stages fit the shared memory -- every compatible
-- card's (Coppelius.SM86Capability), not only the GB10's
def coppeliusTiledTileServes = (lambda unrestricted t : (family TiledProductTile) .
  (lambda unrestricted m : Nat . (lambda unrestricted n : Nat . (lambda unrestricted k : Nat .
    (tiledProductSM86ShapeAdmitted t m n k coppeliusSM86CompatSharedBytesPerBlock)))))
def coppeliusTiledLayoutName = (lambda unrestricted layout : (family TiledProductLayout) .
  (eliminate TiledProductLayout (lambda unrestricted c : (family TiledProductLayout) . Bytes) layout
    (branch TiledProductRows . b"r") (branch TiledProductColumns . b"c")))
-- the stored rows' lengths: A's (k as rows, m as columns), B's (k, n)
def coppeliusTiledStrideA = (lambda unrestricted layout : (family TiledProductLayout) . (lambda unrestricted m : Nat . (lambda unrestricted k : Nat .
  (naturalSelect (coppeliusTiledIsColumns layout) m k))))
def coppeliusTiledStrideB = (lambda unrestricted layout : (family TiledProductLayout) . (lambda unrestricted n : Nat . (lambda unrestricted k : Nat .
  (naturalSelect (coppeliusTiledIsColumns layout) n k))))
def coppeliusTiledName = (lambda unrestricted residual : Nat .
  (lambda unrestricted a : (family TiledProductLayout) . (lambda unrestricted b : (family TiledProductLayout) .
  (lambda unrestricted m : Nat . (lambda unrestricted n : Nat . (lambda unrestricted k : Nat .
    (bytes-append b"tiled-" (bytes-append (coppeliusTiledLayoutName a) (bytes-append (coppeliusTiledLayoutName b)
      (bytes-append b"-m" (bytes-append (naturalDecimalBytesWithin 8 m) (bytes-append b"-n" (bytes-append (naturalDecimalBytesWithin 8 n)
        (bytes-append b"-k" (bytes-append (naturalDecimalBytesWithin 8 k)
          (nat-eliminate (lambda unrestricted c : Nat . Bytes) b"" (lambda unrestricted p : Nat . (lambda unrestricted ignored : Bytes . b"-add")) residual))))))))))))))))

-- a product's program with tile t: looped 1 the loop, 0 the checked
-- unrolling (Proof.GB10TiledChoiceProbe times every candidate this way)
def coppeliusTiledProgramWith = (lambda unrestricted t : (family TiledProductTile) .
  (lambda unrestricted looped : Nat . (lambda unrestricted residual : Nat .
  (lambda unrestricted a : (family TiledProductLayout) . (lambda unrestricted b : (family TiledProductLayout) .
  (lambda unrestricted m : Nat . (lambda unrestricted n : Nat . (lambda unrestricted k : Nat .
    (let unrestricted sa = (coppeliusTiledStrideA a m k) in (let unrestricted sb = (coppeliusTiledStrideB b n k) in
    (eliminate StdBool (lambda unrestricted c : (family StdBool) . (family SM86Program)) (stdBoolFromNatural looped)
      (branch StdTrue .
        (eliminate StdBool (lambda unrestricted c : (family StdBool) . (family SM86Program)) (stdBoolFromNatural residual)
          (branch StdTrue . (tiledProductSM86Residual sm86CompactStallsDrained t a b k sa sb n))
          (branch StdFalse . (tiledProductSM86 sm86CompactStallsDrained t a b k sa sb n))))
      (branch StdFalse .
        (eliminate StdBool (lambda unrestricted c : (family StdBool) . (family SM86Program)) (stdBoolFromNatural residual)
          (branch StdTrue . (tiledProductSM86ResidualChecked sm86CompactStallsDrained t a b k sa sb n))
          (branch StdFalse . (tiledProductSM86Checked sm86CompactStallsDrained t a b k sa sb n)))))))))))))))

-- a product's program with its recorded tile
def coppeliusTiledProgramAs = (lambda unrestricted looped : Nat . (lambda unrestricted residual : Nat .
  (lambda unrestricted a : (family TiledProductLayout) . (lambda unrestricted b : (family TiledProductLayout) .
  (lambda unrestricted m : Nat . (lambda unrestricted n : Nat . (lambda unrestricted k : Nat .
    (coppeliusTiledProgramWith (coppeliusTiledTile residual a b m n k) looped residual a b m n k))))))))

def coppeliusTiledShapeImage = (lambda unrestricted shape : (family CoppeliusTiledShape) .
  (eliminate CoppeliusTiledShape (lambda unrestricted c : (family CoppeliusTiledShape) . (family CoppeliusDeviceImage)) shape
    (branch CoppeliusTiledShapeValue a b m n k residual .
      (let unrestricted tile = (coppeliusTiledTile residual a b m n k) in
      (coppeliusDeviceImage (coppeliusTiledName residual a b m n k) (coppeliusTiledProgramAs 1 residual a b m n k)
        (tiledProductSM86Registers tile) (tiledProductSM86Threads tile) (tiledProductSM86SharedBytes tile))))))
-- 1 when the product has a recorded choice, its tile serves it, and the
-- checked schedule holds on both of SM86.Scoreboard's passes
def coppeliusTiledShapeAdmitted = (lambda unrestricted shape : (family CoppeliusTiledShape) .
  (eliminate CoppeliusTiledShape (lambda unrestricted c : (family CoppeliusTiledShape) . Nat) shape
    (branch CoppeliusTiledShapeValue a b m n k residual .
      (let unrestricted checked = (coppeliusGuardLateReadsIfRequired (coppeliusTiledProgramAs 0 residual a b m n k)) in
      -- a recorded choice, a candidate that serves the product
      (naturalAnd (naturalLess (coppeliusTiledChoiceFor residual a b m n k) coppeliusTiledCandidateCount)
      (naturalAnd (coppeliusTiledTileServes (coppeliusTiledTile residual a b m n k) m n k)
        (naturalAnd (naturalIsZero (sm86Scoreboard checked)) (naturalIsZero (sm86FixedLatencyHazard checked)))))))))
def coppeliusTiledShapeWith = (lambda unrestricted residual : Nat . (lambda unrestricted a : (family TiledProductLayout) . (lambda unrestricted b : (family TiledProductLayout) .
  (lambda unrestricted m : Nat . (lambda unrestricted n : Nat . (lambda unrestricted k : Nat .
    (constructor CoppeliusTiledShape CoppeliusTiledShapeValue a b m n k residual)))))))
def coppeliusTiledShape = (coppeliusTiledShapeWith 0)
-- C = A B + R
def coppeliusTiledShapeAdding = (coppeliusTiledShapeWith 1)
-- every product Coppelius.Build.Graph launches (sequence 1024, width 512,
-- the query/key/value 1536, the feed-forward 1408, the vocabulary 12288)
def coppeliusTiledShapes : (family StdList (family CoppeliusTiledShape)) =
  (stdListFromShapes
    (coppeliusTiledShape coppeliusTiledRows coppeliusTiledRows 1024 1536 512)
    (stdListFromShapes (coppeliusTiledShapeAdding coppeliusTiledRows coppeliusTiledRows 1024 512 512)
    (stdListFromShapes (coppeliusTiledShape coppeliusTiledRows coppeliusTiledRows 1024 1408 512)
    (stdListFromShapes (coppeliusTiledShapeAdding coppeliusTiledRows coppeliusTiledRows 1024 512 1408)
    (stdListFromShapes (coppeliusTiledShape coppeliusTiledRows coppeliusTiledRows 1024 12288 512)
    (stdListFromShapes (coppeliusTiledShape coppeliusTiledRows coppeliusTiledColumns 1024 512 1536)
    (stdListFromShapes (coppeliusTiledShape coppeliusTiledRows coppeliusTiledColumns 1024 512 512)
    (stdListFromShapes (coppeliusTiledShape coppeliusTiledRows coppeliusTiledColumns 1024 512 1408)
    (stdListFromShapes (coppeliusTiledShapeAdding coppeliusTiledRows coppeliusTiledColumns 1024 512 1408)
    (stdListFromShapes (coppeliusTiledShape coppeliusTiledRows coppeliusTiledColumns 1024 1408 512)
    (stdListFromShapes (coppeliusTiledShape coppeliusTiledRows coppeliusTiledColumns 1024 512 12288)
    (stdListFromShapes (coppeliusTiledShape coppeliusTiledColumns coppeliusTiledColumns 1536 512 1024)
    (stdListFromShapes (coppeliusTiledShape coppeliusTiledColumns coppeliusTiledColumns 512 512 1024)
    (stdListFromShapes (coppeliusTiledShape coppeliusTiledColumns coppeliusTiledColumns 1408 512 1024)
    (stdListFromShapes (coppeliusTiledShape coppeliusTiledColumns coppeliusTiledColumns 512 1408 1024)
    (stdListFromShapes (coppeliusTiledShape coppeliusTiledColumns coppeliusTiledColumns 12288 512 1024)
      (constructor StdList StdListEmpty (family CoppeliusTiledShape))))))))))))))))))
-- every product Coppelius launches has a recorded choice, and it serves the
-- product (decided before any program is built)
def coppeliusTiledChoicesAdmitted : Nat =
  (eliminate StdList (lambda unrestricted c : (family StdList (family CoppeliusTiledShape)) . Nat) coppeliusTiledShapes
    (branch StdListEmpty . 1)
    (branch StdListCons head tail induction .
      (naturalAnd induction
        (eliminate CoppeliusTiledShape (lambda unrestricted c : (family CoppeliusTiledShape) . Nat) head
          (branch CoppeliusTiledShapeValue a b m n k residual .
            (naturalAnd (naturalLess (coppeliusTiledChoiceFor residual a b m n k) coppeliusTiledCandidateCount)
              (coppeliusTiledTileServes (coppeliusTiledTile residual a b m n k) m n k)))))))
def coppeliusTiledChoicesServe : (equal Nat coppeliusTiledChoicesAdmitted 1) = (refl Nat 1)
def coppeliusTiledDeviceImages : (family CoppeliusDeviceImages) =
  (eliminate StdList (lambda unrestricted c : (family StdList (family CoppeliusTiledShape)) . (family CoppeliusDeviceImages)) coppeliusTiledShapes
    (branch StdListEmpty . coppeliusDeviceImagesEnd)
    (branch StdListCons head tail induction . (coppeliusDeviceImagesNext (coppeliusTiledShapeImage head) induction)))
def coppeliusTiledSchedulesAdmitted : Nat =
  (eliminate StdList (lambda unrestricted c : (family StdList (family CoppeliusTiledShape)) . Nat) coppeliusTiledShapes
    (branch StdListEmpty . 1)
    (branch StdListCons head tail induction . (naturalAnd (coppeliusTiledShapeAdmitted head) induction)))

-- 1 when a program's schedule holds on both of SM86.Scoreboard's passes
def coppeliusScheduleHolds = (lambda unrestricted program : (family SM86Program) .
  (let unrestricted guarded = (coppeliusGuardLateReadsIfRequired program) in
    (naturalAnd (naturalIsZero (sm86Scoreboard guarded)) (naturalIsZero (sm86FixedLatencyHazard guarded)))))
def coppeliusDeviceImages : (family CoppeliusDeviceImages) =
  (coppeliusDeviceImagesAppend coppeliusImportedDeviceImages
    (coppeliusDeviceImagesAppend coppeliusLossDeviceImages
      (coppeliusDeviceImagesAppend coppeliusStreamingDeviceImages coppeliusTiledDeviceImages)))

-- Every placed SM86 image is checked after half-format selection and the
-- late-read guard. This covers the streaming and imported images that the
-- older sampled admission did not inspect. The fixed-latency pass remains on
-- the previously qualified unrolled paths: older imported kernels, beginning
-- with random-normal, have pre-existing fixed-latency model mismatches and
-- require a separate rescheduling migration before that pass can gate all.
def coppeliusDeviceImageScheduleRefusal : Bytes =
  (eliminate CoppeliusDeviceImages
    (lambda unrestricted current : (family CoppeliusDeviceImages) . Bytes)
    coppeliusDeviceImages
    (branch CoppeliusDeviceImagesEnd . b"")
    (branch CoppeliusDeviceImagesNext image tail induction .
      (eliminate CoppeliusDeviceImage
        (lambda unrestricted current : (family CoppeliusDeviceImage) . Bytes) image
        (branch CoppeliusDeviceImageValue identity program material registers blockX sharedBytes .
          (nat-eliminate (lambda unrestricted accepted : Nat . Bytes)
            identity
            (lambda unrestricted predecessor : Nat . (lambda unrestricted ignored : Bytes . induction))
            (naturalIsZero (sm86Scoreboard program)))))))

def coppeliusDeviceImageSchedulesAdmitted : Nat =
  (naturalIsZero (bytes-length coppeliusDeviceImageScheduleRefusal))

def coppeliusSchedulesAdmitted : Nat =
  (naturalAnd coppeliusDeviceImageSchedulesAdmitted
    (naturalAnd coppeliusTiledSchedulesAdmitted
      (naturalAnd (coppeliusScheduleHolds coppeliusAdamWChecked)
        (naturalAnd (coppeliusScheduleHolds coppeliusGatedGELUForwardImage)
          (coppeliusScheduleHolds coppeliusGatedGELUBackwardImage)))))

-- the device address of the image named `identity`: where the backend
-- places its region, by the same fold (0, which no launch can name, when
-- there is none)
def coppeliusDeviceImageAddress =
  (lambda unrestricted identity : Bytes .
    (app
      (eliminate CoppeliusDeviceImages
        (lambda unrestricted current : (family CoppeliusDeviceImages) . (pi unrestricted cursor : Nat . Nat))
        coppeliusDeviceImages
        (branch CoppeliusDeviceImagesEnd . (lambda unrestricted cursor : Nat . 0))
        (branch CoppeliusDeviceImagesNext image tail induction .
          (lambda unrestricted cursor : Nat .
            (eliminate CoppeliusDeviceImage (lambda unrestricted current : (family CoppeliusDeviceImage) . Nat) image
              (branch CoppeliusDeviceImageValue name program material registers blockX sharedBytes .
                (let unrestricted offset = (naturalAdd cursor (coppeliusDevicePadding cursor)) in
                  (naturalSelect (bytes-equal name identity)
                    (naturalAdd coppeliusProgramBase offset)
                    (induction (naturalAdd offset (bytes-length material))))))))))
      0))

def coppeliusWholeProgramDeviceRegions : (family NvidiaDeviceRegions) =
  (coppeliusNvidiaDeviceRegions coppeliusDeviceImages)

def coppeliusDeviceImagesBuild : (family CoppeliusDeviceImagesBuildResult) =
  (coppeliusBuildDeviceImages coppeliusDeviceImages)