Source/Packages

Realization.Nvidia.SM86.StreamingAttentionSM86

packages/realizations/cooperative/nvidia-sm86/src/Realization/Nvidia/SM86/StreamingAttentionSM86.alpha

3,230 lines493 declarations193.3 KiBSHA-256 cc0366b3d857

def · lines 1353–1391

sqRowTermsWith

Full file
D for the thread's rows: its sixteen dO values of row r (and of r + 8) against O's, summed over the quad; -D and -L kept, D written by the quad's first lane D for the thread's rows, stored at 4 (h seq + row) by the quad's first lane through `storePointer`
1353def sqRowTermsWith = (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
1354  (lambda unrestricted storePointer : Nat . (lambda unrestricted tail : (family SM86Program) .
1355  (let unrestricted rowBytes = (naturalMultiply heads (naturalMultiply saHeadWidth saHalfBytes)) in
1356  -- O in dO's fragment pattern, into the B registers (free before the loop).
1357  -- Q's loads name saPointer as their address and may read it after issue:
1358  -- re-form it only once they (SB1) have retired, else an O address can
1359  -- replace a Q address before its late read (the c28712e8 hazard; the
1360  -- checker found it at ordinal 64).
1361  (saOp (constructor SM86InstructionBody SM86IntegerMultiplyAddWideConstant (saR saPointer) (saR saScratch) (saR saOne)
1362    (byte 0) (saU (saArgument 5)) (saAfter saWait1))
1363  (saFor 16 (lambda unrestricted i : Nat .
1364    (let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
1365      (saLoad (naturalAdd 100 i) saPointer
1366        (naturalAdd (naturalMultiply kt 32)
1367          (naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 rowBytes))
1368            (naturalMultiply (naturalDivideUnchecked j 2) 16)))
1369        saSB0 saWaitNone))))
1370  (saMovImm (sqTmp 0) 0 (saMovImm (sqTmp 1) 0
1371  -- register j of slice kt is row r (j even) or r + 8 (j odd); both halves
1372  (saFor 32 (lambda unrestricted i : Nat .
1373    (let unrestricted r = (naturalDivideUnchecked i 2) in (let unrestricted high = (naturalModuloUnchecked i 2) in
1374      (let unrestricted h = (naturalModuloUnchecked (naturalModuloUnchecked r 4) 2) in
1375      (lambda unrestricted rest : (family SM86Program) .
1376        (sbWiden (sqTmp 2) (naturalAdd 100 r) high (naturalSelect (naturalIsZero i) (naturalAdd saWait0 saWait1) saWaitNone)
1377        (sbWiden (sqTmp 3) (naturalAdd (sqDOut 0) r) high saWaitNone
1378        (saFfma (sqTmp h) (sqTmp 2) (sqTmp 3) (sqTmp h)
1379          rest))))))))
1380  (saFor 2 (lambda unrestricted h : Nat . (lambda unrestricted rest : (family SM86Program) .
1381    (saShfl (sqTmp 2) (sqTmp h) 1 saWaitNone
1382    (saFadd (sqTmp h) (sqTmp h) (sqTmp 2) saWait5
1383    (saShfl (sqTmp 2) (sqTmp h) 2 saWaitNone
1384    (saFadd (sqTmp h) (sqTmp h) (sqTmp 2) saWait5
1385    (saFneg (sqNegD h) (sqTmp h)
1386      rest)))))))
1387  (saWide storePointer saRow saOne (saArgument 8)
1388  (saGreater saP0 saQuadLane 0
1389  (saUnless saP0 (saStore storePointer (sqTmp 0) 0 saWaitNone)
1390  (saUnless saP0 (saStore storePointer (sqTmp 1) 32 saWaitNone)
1391    tail)))))))))))))))

The compiler supplied declaration spans and resolved links from this source snapshot. This page does not assert that this file belongs to a checked closure.