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 714–739

sfTileBodyMasked

Full file
714def sfTileBodyMasked = (lambda unrestricted mask : (pi unrestricted tail : (family SM86Program) . (family SM86Program)) .
715  (lambda unrestricted rows : Nat . (lambda unrestricted buffer : Nat . (lambda unrestricted other : Nat .
716  (lambda unrestricted tail : (family SM86Program) .
717  (sfBarrier
718  (saGreater saP1 saCount 1
719  (saFor 4 (lambda unrestricted i : Nat . (saAddImm (sfKeyOffset i) (sfKeyOffset i) (naturalMultiply saTile (naturalMultiply saHeadWidth saHalfBytes))))
720  (saFor 4 (lambda unrestricted i : Nat . (saAddImm (sfValueOffset i) (sfValueOffset i)
721    (naturalSelect rows (naturalMultiply saTile (naturalMultiply saHeadWidth saHalfBytes)) (naturalMultiply saTile saHalfBytes))))
722  (sfCopyTile other 1
723  (sfWaitGroups 1 saWait5
724  (sfBarrier
725  (saImad saMask saCount 4096 saMaskBase
726  (sfLoadFragments buffer saWaitNone
727  (sfScores
728  (sfChooseWith rows
729    (sfLoadFragmentsTransposed sfT1 (naturalAdd buffer sfValueTile) saWait2)
730    (sfLoadFragments (naturalAdd buffer sfValueTile) saWait2)
731  (mask
732  (saScaleRow 0 (saScaleRow 1
733  (saRowMax 0 (saRowMax 1
734  (saRescaleFactor 0 (saRescaleFactor 1
735  (saExponentials 0 (saExponentials 1
736  (saRescaleOutput 0 (saRescaleOutput 1
737  (saPackP
738  (sfValues
739    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.