Source/Packages

Realization.Nvidia.SM86.StreamingAttentionSM86

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

1,203 lines221 declarations70.9 KiBSHA-256 23d4e5e2aa2a

def · lines 684–694

sqEpilogue

Full file
dQ = s (the sum): the thread's rows r and r + 8, pairs of binary32 columns
684def sqEpilogue = (lambda unrestricted tail : (family SM86Program) .
685  (saMovConst (sqTmp 4) (saArgument 10)
686  (saFor 32 (lambda unrestricted i : Nat .
687    (saFmul (naturalAdd 68 i) (naturalAdd 68 i) (sqTmp 4) (naturalSelect (naturalIsZero i) saWait3 saWaitNone)))
688  (saWide saPointer saScratch saOne (saArgument 7)
689  (saFor 16 (lambda unrestricted i : Nat .
690    (let unrestricted nt = (naturalDivideUnchecked i 2) in (let unrestricted h = (naturalModuloUnchecked i 2) in
691      (saOp (sbStore64 saPointer (saElem (saO nt) h 0)
692        (naturalAdd (naturalMultiply nt 32) (naturalMultiply h (naturalMultiply 8 (naturalMultiply saHeadWidth 4))))
693        saWaitNone)))))
694    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.