Realization.Nvidia.SM86.StreamingAttentionSM86

1483def sqRotatedEpilogue = (lambda unrestricted seq : Nat . (lambda unrestricted width : Nat .
1484  (lambda unrestricted tail : (family SM86Program) .
1485  (saMovConst (sqTmp 4) (saArgument 10)
1486  (saFor 32 (lambda unrestricted i : Nat .
1487    (saFmul (naturalAdd 68 i) (naturalAdd 68 i) (sqTmp 4) (naturalSelect (naturalIsZero i) saWait3 saWaitNone)))
1488  -- the column, position and head from the binary32 store's offset (bytes
1489  -- of [head][seq][64]): element e = offset / 4, column e % 64, row e / 64
1490  -- = head seq + position
1491  (saShr sqrElement saScratch 2
1492  (saMovImm sqrColumn 63
1493  (saAnd sqrColumn sqrElement sqrColumn
1494  (saShr sqrRow sqrElement 6
1495  (saMovImm sqrPosition (naturalSaturatingSubtract seq 1)
1496  (saAnd sqrPosition sqrRow sqrPosition
1497  (saShr sqrHead sqrRow (sqrLog2 seq)
1498  -- the output: (position width + 64 head + column) 2 bytes
1499  (saImad sqrOutput sqrPosition width sqrColumn
1500  (saImad sqrOutput sqrHead saHeadWidth sqrOutput
1501  (saImad sqrOutput sqrOutput saHalfBytes 255
1502  (saWide sqrPointer sqrOutput saOne (saArgument 7)
1503  -- the tables: (position 32 + column) 4 bytes
1504  (saImad sqrTable sqrPosition 32 sqrColumn
1505  (saMovImm sqrFour 4
1506  (saWide sqrCosine sqrTable sqrFour (saArgument streamingAttentionQueryRotatedCosineArgument)
1507  (saWide sqrSine sqrTable sqrFour (saArgument streamingAttentionQueryRotatedSineArgument)
1508  -- every table word: rows r and r + 8, columns 8 j + c and + 1 (j < 4)
1509  (saFor 16 (lambda unrestricted i : Nat .
1510    (let unrestricted h = (naturalDivideUnchecked i 8) in
1511    (let unrestricted j = (naturalDivideUnchecked (naturalModuloUnchecked i 8) 2) in
1512    (let unrestricted e = (naturalModuloUnchecked i 2) in
1513    (let unrestricted offset = (naturalMultiply 4 (naturalAdd (naturalMultiply 256 h) (naturalAdd (naturalMultiply 8 j) e))) in
1514    (lambda unrestricted after : (family SM86Program) .
1515      (saLoad (sqrWord h 0 j e) sqrCosine offset saSB4 saWaitNone
1516      (saLoad (sqrWord h 1 j e) sqrSine offset saSB4 saWaitNone after))))))))
1517  -- the sines negated (the inverse rotation), once every word is in
1518  (saFor 16 (lambda unrestricted i : Nat .
1519    (let unrestricted h = (naturalDivideUnchecked i 8) in
1520    (let unrestricted j = (naturalDivideUnchecked (naturalModuloUnchecked i 8) 2) in
1521    (let unrestricted e = (naturalModuloUnchecked i 2) in
1522      (sbNegate (sqrWord h 1 j e) (sqrWord h 1 j e) (naturalSelect i saWaitNone saWait4))))))
1523  (saFor 2 (lambda unrestricted h : Nat . (lambda unrestricted rest : (family SM86Program) .
1524    (saFor 4 (lambda unrestricted j : Nat . (sqrPairs width h j)) rest)))
1525    tail)))))))))))))))))))))))