Bob.Build.SM86Programs

1module Bob.Build.SM86Programs
2
3import Bob.Model
4import Bob.Build.BigramLossGradientSM86
5import Realization.Nvidia.SM86.AdamWSM86
6import Realization.Nvidia.SM86.EmbeddingGatherF32SM86
7import Realization.Nvidia.SM86.GreedyArgmaxSM86
8
9-- Bob's training fixture places one example for each current-byte row.  The
10-- dedicated one-block realization walks each complete 256-logit row, emits
11-- exact full-vocabulary loss and gradient, and performs no host work.
12def bobLossGradientImage =
13  bobBigramLossGradientImage
14
15-- Prediction is a table-row gather followed by the maintained 256-way greedy
16-- argmax realization.  Tie breaking therefore remains the shared lowest-byte
17-- contract rather than a Bob-only host implementation.
18def bobPredictionRowImage : (family EmbeddingGatherF32ImageResult) =
19  (embeddingGatherF32ImageFor (constructor EmbeddingGatherF32Width EmbeddingGatherF32Width256) 1)
20
21def bobOptimizerImage : (family AdamWF32SM86ImageResult) =
22  (adamWF32BoundedImage (adamWF32Unsigned32FromNatural bobParameterCount))
23
24def bobPredictionArgmaxImage : (family GreedyArgmaxSM86Artifact) =
25  greedyArgmaxSM86BuildByteVocabulary