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 2874–2914

sksTileBodyMasked

Full file
2874def sksTileBodyMasked = (lambda unrestricted mask : (pi unrestricted tail : (family SM86Program) . (family SM86Program)) .
2875  (lambda unrestricted seq : Nat . (lambda unrestricted buffer : Nat . (lambda unrestricted other : Nat .
2876  (lambda unrestricted tail : (family SM86Program) .
2877  (sfBarrier
2878  (saGreater saP1 saCount 1
2879  (saFor 4 (lambda unrestricted i : Nat .
2880    (saAddImm (sksOffset i) (sksOffset i) (naturalSaturatingSubtract 4294967296 (naturalMultiply skSubTile (naturalMultiply saHeadWidth saHalfBytes)))))
2881  (sksCopyTile other 1
2882  (sfWaitGroups 1 saWait5
2883  (sfBarrier
2884  (saImad saMask saCount 4096 saMaskBase
2885  -- the sub-tile's L and D, in flight under the products below
2886  (saWide saPointer skLogOffset saOne (saArgument 6)
2887  (sksLoadRowTermsInto skNegL saPointer
2888  (saWide sksRowTermPointer skLogOffset saOne (saArgument 7)
2889  (sksLoadRowTermsInto skNegD sksRowTermPointer
2890  (saAddImm skLogOffset skLogOffset (naturalSaturatingSubtract 4294967296 (naturalMultiply skSubTile 4))
2891  -- S^T = K Q^T
2892  (sksLoadFragments 0 buffer saWait3
2893  (skProducts 1 skS skK 4 4 saSB2 saWait2 saWait1
2894  -- dP^T = V dO^T
2895  (sksLoadFragments 0 (naturalAdd buffer sksQueryTile) saWait2
2896  (skProducts 1 skDP skV 4 4 saSB2 saWait2 saWait1
2897  -- dO^T's fragments for dV, once dP^T's products are done
2898  (sksLoadFragments 1 (naturalAdd buffer sksQueryTile) saWait2
2899  (mask
2900  -- their negatives, once they have arrived (exact, as before)
2901  (saFor 8 (lambda unrestricted i : Nat .
2902    (sbNegate (skNegL (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2))
2903      (skNegL (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2))
2904      (naturalSelect (naturalIsZero i) (naturalAdd saWait4 saWait5) saWaitNone)))
2905  (saFor 8 (lambda unrestricted i : Nat .
2906    (sbNegate (skNegD (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2))
2907      (skNegD (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2)) saWaitNone))
2908  (sksScoreGradient
2909  (skProducts 0 skDV skPA 2 8 saSB3 saWait3 saWait1
2910  -- Q^T's fragments for dK, once dV's products have read dO^T's
2911  (sksLoadFragments 1 buffer saWait3
2912  (skProducts 0 skDK skSA 2 8 saSB3 saWait3 saWait1
2913  (sksStoreScoreGradients seq
2914    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.