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 858–891

skScoreGradient

Full file
P^T = 2^(t - L), dS^T = P^T (dP^T - D), in S^T's registers
858def skScoreGradient = (lambda unrestricted tail : (family SM86Program) .
859  (saFor 16 (lambda unrestricted i : Nat .
860    (let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
861      (let unrestricted b = (naturalModuloUnchecked i 2) in
862      (let unrestricted e = (saElem (skS nt) h b) in
863      (lambda unrestricted rest : (family SM86Program) .
864        (saFfma e e saScale (skNegL nt b)
865          (saEx2 e e saWaitNone rest))))))))
866  -- P^T's fragments, before dS^T replaces it
867  (saFor 2 (lambda unrestricted kk : Nat .
868    (let unrestricted left = (skS (naturalMultiply 2 kk)) in (let unrestricted right = (skS (succ (naturalMultiply 2 kk))) in
869      (lambda unrestricted rest : (family SM86Program) .
870        (saOp (constructor SM86InstructionBody SM86FloatPairToPackedHalfPair (saR (skPA kk)) (saR (succ left)) (saR left)
871          (saAfter (naturalSelect (naturalIsZero kk) saWait4 saWaitNone)))
872        (saPack (succ (skPA kk)) (naturalAdd left 3) (naturalAdd left 2)
873        (saPack (naturalAdd (skPA kk) 2) (succ right) right
874        (saPack (naturalAdd (skPA kk) 3) (naturalAdd right 3) (naturalAdd right 2)
875          rest))))))))
876  (saFor 16 (lambda unrestricted i : Nat .
877    (let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
878      (let unrestricted b = (naturalModuloUnchecked i 2) in
879      (let unrestricted e = (saElem (skS nt) h b) in (let unrestricted g = (saElem (skDP nt) h b) in
880      (lambda unrestricted rest : (family SM86Program) .
881        (saFadd g g (skNegD nt b) saWaitNone
882          (saFmul e e g saWaitNone rest)))))))))
883  (saFor 2 (lambda unrestricted kk : Nat .
884    (let unrestricted left = (skS (naturalMultiply 2 kk)) in (let unrestricted right = (skS (succ (naturalMultiply 2 kk))) in
885      (lambda unrestricted rest : (family SM86Program) .
886        (saPack (skSA kk) (succ left) left
887        (saPack (succ (skSA kk)) (naturalAdd left 3) (naturalAdd left 2)
888        (saPack (naturalAdd (skSA kk) 2) (succ right) right
889        (saPack (naturalAdd (skSA kk) 3) (naturalAdd right 3) (naturalAdd right 2)
890          rest))))))))
891    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.