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 2413–2426

skPackRowFragments

Full file
Pack a 16-element score plane into two eight-half A fragments. Source tiles are visited before later packed destinations can overwrite them; destination may therefore start at the source plane's first register. The checked packing root covers every destination word, signed zero and BF16 ties. This is a register transformation, not a change of precision.
2413def skPackRowFragments = (lambda unrestricted source : (pi unrestricted nt : Nat . Nat) .
2414  (lambda unrestricted destination : (pi unrestricted kk : Nat . Nat) .
2415    (lambda unrestricted wait : Nat . (lambda unrestricted tail : (family SM86Program) .
2416      (saFor 2 (lambda unrestricted kk : Nat .
2417        (let unrestricted left = (source (naturalMultiply 2 kk)) in
2418        (let unrestricted right = (source (succ (naturalMultiply 2 kk))) in
2419          (lambda unrestricted rest : (family SM86Program) .
2420            (saOp (constructor SM86InstructionBody SM86FloatPairToPackedHalfPair
2421              (saR (destination kk)) (saR (succ left)) (saR left)
2422              (saAfter (naturalSelect (naturalIsZero kk) wait saWaitNone)))
2423            (saPack (succ (destination kk)) (naturalAdd left 3) (naturalAdd left 2)
2424            (saPack (naturalAdd (destination kk) 2) (succ right) right
2425            (saPack (naturalAdd (destination kk) 3) (naturalAdd right 3) (naturalAdd right 2)
2426              rest)))))))) 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.