Source/Packages

Realization.Nvidia.SM86.StreamingAttentionSM86

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

1,129 lines205 declarations66.1 KiBSHA-256 aefcbc3b00ca

def · lines 893–928

skTileBody

Full file
893def skTileBody = (lambda unrestricted seq : Nat . (lambda unrestricted tail : (family SM86Program) .
894  (saImad saMask saCount 4096 saMaskBase
895  -- the sub-tile's L and D (their negatives)
896  (saWide saPointer skLogOffset saOne (saArgument 6)
897  (saFor 8 (lambda unrestricted i : Nat .
898    (saLoad (skTmp i) saPointer (naturalAdd (naturalMultiply (naturalDivideUnchecked i 2) 32) (naturalMultiply (naturalModuloUnchecked i 2) 4))
899      saSB1 (naturalSelect (naturalIsZero i) saWait3 saWaitNone)))
900  (saFor 8 (lambda unrestricted i : Nat .
901    (sbNegate (skNegL (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2)) (skTmp i) (naturalSelect (naturalIsZero i) saWait1 saWaitNone)))
902  (saWide saPointer skLogOffset saOne (saArgument 7)
903  (saFor 8 (lambda unrestricted i : Nat .
904    (saLoad (skTmp i) saPointer (naturalAdd (naturalMultiply (naturalDivideUnchecked i 2) 32) (naturalMultiply (naturalModuloUnchecked i 2) 4))
905      saSB1 saWaitNone))
906  (saFor 8 (lambda unrestricted i : Nat .
907    (sbNegate (skNegD (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2)) (skTmp i) (naturalSelect (naturalIsZero i) saWait1 saWaitNone)))
908  -- S^T = K Q^T
909  (saWide saKeyPointer saKeyOffset saOne (saArgument 0)
910  (skLoadRows saKeyPointer saWait3
911  (saFor 16 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 52 i) 0))
912  (saFor 16 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 68 i) 0))
913  (skProducts skS skK 4 4 saSB2 saWait2 (naturalAdd saWait0 saWait1)
914  -- dP^T = V dO^T
915  (saWide saKeyPointer saKeyOffset saOne (saArgument 3)
916  (skLoadRows saKeyPointer saWait2
917  (skProducts skDP skV 4 4 saSB2 saWait2 saWait0
918  -- dO^T's fragments for dV, once dP^T's products are done
919  (saWide saValuePointer saValueOffset saOne (saArgument 5)
920  (skLoadColumns seq saValuePointer saWait2
921  (skMaskTile
922  (skScoreGradient
923  (skProducts skDV skPA 2 8 saSB3 saWait3 saWait0
924  -- Q^T's fragments for dK, once dV's products have read dO^T's
925  (saWide saValuePointer saValueOffset saOne (saArgument 4)
926  (skLoadColumns seq saValuePointer saWait3
927  (skProducts skDK skSA 2 8 saSB3 saWait3 saWait0
928    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.