module Hardware.Nvidia.SM86.Command.AffineLaunchSchedule

import Hardware.Nvidia.SM86.Command.WholeProgramPlan
import Model.Parameter
import Model.Word64
import Std.Natural
import Std.Word

-- A repeated body may change only its parameter words between iterations.
-- Compare adjacent body instances by launch ordinal and word. Differences
-- are modulo 2^64 because device pointers are encoded as 64-bit words. The
-- body is retained once; WholeProgramPlan expands its affine parameters
-- when it emits the launch-table recipe. Callers prove the body shape and
-- pointer progression for the geometry they instantiate.
def nvidiaAffineNoLaunches : (family NvidiaLaunchSchedule) =
  (constructor NvidiaLaunchSchedule NvidiaLaunchScheduleEmpty)
def nvidiaAffineScheduleIf =
  (lambda unrestricted condition : Nat .
    (lambda unrestricted whenTrue : (family NvidiaLaunchSchedule) .
      (lambda unrestricted whenFalse : (family NvidiaLaunchSchedule) .
        (nat-eliminate
          (lambda unrestricted current : Nat . (family NvidiaLaunchSchedule))
          whenFalse
          (lambda unrestricted p : Nat .
            (lambda unrestricted induction : (family NvidiaLaunchSchedule) . whenTrue))
          condition))))

def nvidiaAffineWordMaximum : Nat =
  (modelWord64Natural stdU64AllOnes)

def nvidiaAffineWordDifference =
  (lambda unrestricted a : Nat .
    (lambda unrestricted b : Nat .
      (naturalSelect
        (naturalLessOrEqual a b)
        (naturalSaturatingSubtract b a)
        (naturalSaturatingSubtract
          nvidiaAffineWordMaximum
          (naturalSaturatingSubtract (naturalSaturatingSubtract a b) 1)))))

def nvidiaAffinePatchValue =
  (lambda unrestricted patches : (family NvidiaParameterPatches) .
    (lambda unrestricted word : Nat .
      (eliminate
        NvidiaParameterPatches
        (lambda unrestricted current : (family NvidiaParameterPatches) . Nat)
        patches
        (branch NvidiaParameterPatchesEnd . 0)
        (branch
          NvidiaParameterPatchesNext
          head
          tail
          induction
          .
          (eliminate
            NvidiaParameterPatch
            (lambda unrestricted current : (family NvidiaParameterPatch) . Nat)
            head
            (branch
              NvidiaParameterPatchValue
              at
              value
              .
              (naturalSelect (naturalEqual at word) (modelWord64Natural value) induction)))))))

def nvidiaAffinePatchPresent =
  (lambda unrestricted patches : (family NvidiaParameterPatches) .
    (lambda unrestricted word : Nat .
      (eliminate
        NvidiaParameterPatches
        (lambda unrestricted current : (family NvidiaParameterPatches) . Nat)
        patches
        (branch NvidiaParameterPatchesEnd . 0)
        (branch
          NvidiaParameterPatchesNext
          head
          tail
          induction
          .
          (eliminate
            NvidiaParameterPatch
            (lambda unrestricted current : (family NvidiaParameterPatch) . Nat)
            head
            (branch
              NvidiaParameterPatchValue
              at
              value
              .
              (naturalOr (naturalEqual at word) induction)))))))

def nvidiaAffinePatchesOf =
  (lambda unrestricted template : (family NvidiaLaunchTemplate) .
    (eliminate
      NvidiaLaunchTemplate
      (lambda unrestricted current : (family NvidiaLaunchTemplate) .
        (family NvidiaParameterPatches))
      template
      (branch
        NvidiaLaunchTemplateValue
        identity
        kernel
        block
        .
        (eliminate
          NvidiaParameterBlock
          (lambda unrestricted current : (family NvidiaParameterBlock) .
            (family NvidiaParameterPatches))
          block
          (branch NvidiaParameterBlockValue patches . patches)))))

def nvidiaAffineAdjust =
  (lambda unrestricted ordinal : Nat .
    (lambda unrestricted word : Nat .
      (lambda unrestricted delta : Nat .
        (lambda unrestricted tail : (family NvidiaParameterAdjustments) .
          (nat-eliminate
            (lambda unrestricted current : Nat . (family NvidiaParameterAdjustments))
            tail
            (lambda unrestricted q : Nat .
              (lambda unrestricted ignored : (family NvidiaParameterAdjustments) .
                (constructor
                  NvidiaParameterAdjustments
                  NvidiaParameterAdjustmentsNext
                  (constructor
                    NvidiaParameterAdjustment
                    NvidiaParameterAdjustmentValue
                    (constructor
                      NvidiaParameterAdjustmentScope
                      NvidiaParameterAdjustmentBlock
                      ordinal)
                    word
                    (modelWord64FromNaturalTruncated delta))
                  tail)))
            (naturalNonzero delta))))))

-- Compare sparse words for one launch. An absent word is zero in the launch
-- table, so additions and removals are deltas from or to zero. The ordinal
-- scope prevents a pointer in one kernel from relocating the same ABI slot
-- when that slot is a scalar in another kernel.
def nvidiaAffineBlockDeltas =
  (lambda unrestricted ordinal : Nat .
    (lambda unrestricted first : (family NvidiaParameterPatches) .
      (lambda unrestricted second : (family NvidiaParameterPatches) .
        (lambda unrestricted tail : (family NvidiaParameterAdjustments) .
          (eliminate
            NvidiaParameterPatches
            (lambda unrestricted current : (family NvidiaParameterPatches) .
              (family NvidiaParameterAdjustments))
            first
            (branch
              NvidiaParameterPatchesEnd
              .
              (eliminate
                NvidiaParameterPatches
                (lambda unrestricted current : (family NvidiaParameterPatches) .
                  (family NvidiaParameterAdjustments))
                second
                (branch NvidiaParameterPatchesEnd . tail)
                (branch
                  NvidiaParameterPatchesNext
                  head
                  rest
                  induction
                  .
                  (eliminate
                    NvidiaParameterPatch
                    (lambda unrestricted current : (family NvidiaParameterPatch) .
                      (family NvidiaParameterAdjustments))
                    head
                    (branch
                      NvidiaParameterPatchValue
                      word
                      value
                      .
                      (nat-eliminate
                        (lambda unrestricted current : Nat . (family NvidiaParameterAdjustments))
                        (nvidiaAffineAdjust ordinal word (modelWord64Natural value) induction)
                        (lambda unrestricted q : Nat .
                          (lambda unrestricted ignored : (family NvidiaParameterAdjustments) .
                            induction))
                        (nvidiaAffinePatchPresent first word)))))))
            (branch
              NvidiaParameterPatchesNext
              head
              rest
              induction
              .
              (eliminate
                NvidiaParameterPatch
                (lambda unrestricted current : (family NvidiaParameterPatch) .
                  (family NvidiaParameterAdjustments))
                head
                (branch
                  NvidiaParameterPatchValue
                  word
                  value
                  .
                  (nvidiaAffineAdjust
                    ordinal
                    word
                    (nvidiaAffineWordDifference (modelWord64Natural value) (nvidiaAffinePatchValue second word))
                    induction)))))))))

def nvidiaAffineOneOf =
  (lambda unrestricted schedule : (family NvidiaLaunchSchedule) .
    (eliminate
      NvidiaLaunchSchedule
      (lambda unrestricted current : (family NvidiaLaunchSchedule) .
        (family NvidiaParameterPatches))
      schedule
      (branch
        NvidiaLaunchScheduleEmpty
        .
        (constructor NvidiaParameterPatches NvidiaParameterPatchesEnd))
      (branch NvidiaLaunchScheduleOne template . (nvidiaAffinePatchesOf template))
      (branch
        NvidiaLaunchScheduleAppend
        left
        right
        il
        ir
        .
        (constructor NvidiaParameterPatches NvidiaParameterPatchesEnd))
      (branch
        NvidiaLaunchScheduleRepeat
        count
        iteration
        body
        ib
        .
        (constructor NvidiaParameterPatches NvidiaParameterPatchesEnd))))

def nvidiaAffineLeftOf =
  (lambda unrestricted schedule : (family NvidiaLaunchSchedule) .
    (eliminate
      NvidiaLaunchSchedule
      (lambda unrestricted current : (family NvidiaLaunchSchedule) . (family NvidiaLaunchSchedule))
      schedule
      (branch NvidiaLaunchScheduleEmpty . schedule)
      (branch NvidiaLaunchScheduleOne template . schedule)
      (branch NvidiaLaunchScheduleAppend left right il ir . left)
      (branch NvidiaLaunchScheduleRepeat count iteration body ib . body)))

def nvidiaAffineRightOf =
  (lambda unrestricted schedule : (family NvidiaLaunchSchedule) .
    (eliminate
      NvidiaLaunchSchedule
      (lambda unrestricted current : (family NvidiaLaunchSchedule) . (family NvidiaLaunchSchedule))
      schedule
      (branch NvidiaLaunchScheduleEmpty . schedule)
      (branch NvidiaLaunchScheduleOne template . schedule)
      (branch NvidiaLaunchScheduleAppend left right il ir . right)
      (branch NvidiaLaunchScheduleRepeat count iteration body ib . body)))

-- the deltas between two bodies of the same shape, blocks numbered in the
-- expanded body from `offset`; a nested repeat's difference applies to
-- every copy of its body
def nvidiaAffineDeltas =
  (lambda unrestricted first : (family NvidiaLaunchSchedule) .
    (eliminate
      NvidiaLaunchSchedule
      (lambda unrestricted current : (family NvidiaLaunchSchedule) .
        (pi unrestricted second : (family NvidiaLaunchSchedule) .
          (pi unrestricted offset : Nat .
            (pi unrestricted tail : (family NvidiaParameterAdjustments) .
              (family NvidiaParameterAdjustments)))))
      first
      (branch
        NvidiaLaunchScheduleEmpty
        .
        (lambda unrestricted second : (family NvidiaLaunchSchedule) .
          (lambda unrestricted offset : Nat .
            (lambda unrestricted tail : (family NvidiaParameterAdjustments) . tail))))
      (branch
        NvidiaLaunchScheduleOne
        template
        .
        (lambda unrestricted second : (family NvidiaLaunchSchedule) .
          (lambda unrestricted offset : Nat .
            (lambda unrestricted tail : (family NvidiaParameterAdjustments) .
              (nvidiaAffineBlockDeltas offset (nvidiaAffinePatchesOf template) (nvidiaAffineOneOf second) tail)))))
      (branch
        NvidiaLaunchScheduleAppend
        left
        right
        il
        ir
        .
        (lambda unrestricted second : (family NvidiaLaunchSchedule) .
          (lambda unrestricted offset : Nat .
            (lambda unrestricted tail : (family NvidiaParameterAdjustments) .
              (il
                (nvidiaAffineLeftOf second)
                offset
                (ir (nvidiaAffineRightOf second) (naturalAdd offset (nvidiaLaunchScheduleCount left)) tail))))))
      (branch
        NvidiaLaunchScheduleRepeat
        count
        iteration
        body
        ib
        .
        (lambda unrestricted second : (family NvidiaLaunchSchedule) .
          (lambda unrestricted offset : Nat .
            (lambda unrestricted tail : (family NvidiaParameterAdjustments) .
              (nat-eliminate
                (lambda unrestricted current : Nat . (family NvidiaParameterAdjustments))
                tail
                (lambda unrestricted p : Nat .
                  (lambda unrestricted induction : (family NvidiaParameterAdjustments) .
                    (ib
                      (nvidiaAffineRightOf second)
                      (naturalAdd offset (naturalMultiply p (nvidiaLaunchScheduleCount body)))
                      induction)))
                count)))))))

-- `count` bodies of an affine loop, as one Repeat
def nvidiaAffineLaunchRepeat =
  (lambda unrestricted count : Nat .
    (lambda unrestricted body : (pi unrestricted index : Nat . (family NvidiaLaunchSchedule)) .
      (nvidiaAffineScheduleIf
        (naturalLess count 2)
        (nvidiaAffineScheduleIf (naturalIsZero count) nvidiaAffineNoLaunches (body 0))
        (let unrestricted first =
          (body 0)
          in
          (constructor
            NvidiaLaunchSchedule
            NvidiaLaunchScheduleRepeat
            count
            (constructor
              NvidiaParameterIteration
              NvidiaParameterIterationAffine
              (nvidiaAffineDeltas
                first
                (body 1)
                0
                (constructor NvidiaParameterAdjustments NvidiaParameterAdjustmentsEnd)))
            first)))))