Source/Packages

Realization.Nvidia.SM86.StreamingAttentionSM86

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

1,203 lines221 declarations70.9 KiBSHA-256 23d4e5e2aa2a

def · lines 943–978

skTileBody

Full file
943def skTileBody = (lambda unrestricted seq : Nat . (lambda unrestricted tail : (family SM86Program) .
944  (saImad saMask saCount 4096 saMaskBase
945  -- the sub-tile's L and D (their negatives)
946  (saWide saPointer skLogOffset saOne (saArgument 6)
947  (saFor 8 (lambda unrestricted i : Nat .
948    (saLoad (skTmp i) saPointer (naturalAdd (naturalMultiply (naturalDivideUnchecked i 2) 32) (naturalMultiply (naturalModuloUnchecked i 2) 4))
949      saSB1 (naturalSelect (naturalIsZero i) saWait3 saWaitNone)))
950  (saFor 8 (lambda unrestricted i : Nat .
951    (sbNegate (skNegL (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2)) (skTmp i) (naturalSelect (naturalIsZero i) saWait1 saWaitNone)))
952  (saWide saPointer skLogOffset saOne (saArgument 7)
953  (saFor 8 (lambda unrestricted i : Nat .
954    (saLoad (skTmp i) saPointer (naturalAdd (naturalMultiply (naturalDivideUnchecked i 2) 32) (naturalMultiply (naturalModuloUnchecked i 2) 4))
955      saSB1 saWaitNone))
956  (saFor 8 (lambda unrestricted i : Nat .
957    (sbNegate (skNegD (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2)) (skTmp i) (naturalSelect (naturalIsZero i) saWait1 saWaitNone)))
958  -- S^T = K Q^T
959  (saWide saKeyPointer saKeyOffset saOne (saArgument 0)
960  (skLoadRows saKeyPointer saWait3
961  (saFor 16 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 52 i) 0))
962  (saFor 16 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 68 i) 0))
963  (skProducts skS skK 4 4 saSB2 saWait2 (naturalAdd saWait0 saWait1)
964  -- dP^T = V dO^T
965  (saWide saKeyPointer saKeyOffset saOne (saArgument 3)
966  (skLoadRows saKeyPointer saWait2
967  (skProducts skDP skV 4 4 saSB2 saWait2 saWait0
968  -- dO^T's fragments for dV, once dP^T's products are done
969  (saWide saValuePointer saValueOffset saOne (saArgument 5)
970  (skLoadColumns seq saValuePointer saWait2
971  (skMaskTile
972  (skScoreGradient
973  (skProducts skDV skPA 2 8 saSB3 saWait3 saWait0
974  -- Q^T's fragments for dK, once dV's products have read dO^T's
975  (saWide saValuePointer saValueOffset saOne (saArgument 4)
976  (skLoadColumns seq saValuePointer saWait3
977  (skProducts skDK skSA 2 8 saSB3 saWait3 saWait0
978    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.