module Bob.Build.SM86Programs

import Bob.Model
import Bob.Build.BigramLossGradientSM86
import Realization.Nvidia.SM86.AdamWSM86
import Realization.Nvidia.SM86.EmbeddingGatherF32SM86
import Realization.Nvidia.SM86.GreedyArgmaxSM86

-- Bob's training fixture places one example for each current-byte row.  The
-- dedicated one-block realization walks each complete 256-logit row, emits
-- exact full-vocabulary loss and gradient, and performs no host work.
def bobLossGradientImage =
  bobBigramLossGradientImage

-- Prediction is a table-row gather followed by the maintained 256-way greedy
-- argmax realization.  Tie breaking therefore remains the shared lowest-byte
-- contract rather than a Bob-only host implementation.
def bobPredictionRowImage : (family EmbeddingGatherF32ImageResult) =
  (embeddingGatherF32ImageFor (constructor EmbeddingGatherF32Width EmbeddingGatherF32Width256) 1)

def bobOptimizerImage : (family AdamWF32SM86ImageResult) =
  (adamWF32BoundedImage (adamWF32Unsigned32FromNatural bobParameterCount))

def bobPredictionArgmaxImage : (family GreedyArgmaxSM86Artifact) =
  greedyArgmaxSM86BuildByteVocabulary