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 665–675

sqEpilogue

Full file
dQ = s (the sum): the thread's rows r and r + 8, pairs of binary32 columns
665def sqEpilogue = (lambda unrestricted tail : (family SM86Program) .
666  (saMovConst (sqTmp 4) (saArgument 10)
667  (saFor 32 (lambda unrestricted i : Nat .
668    (saFmul (naturalAdd 68 i) (naturalAdd 68 i) (sqTmp 4) (naturalSelect (naturalIsZero i) saWait3 saWaitNone)))
669  (saWide saPointer saScratch saOne (saArgument 7)
670  (saFor 16 (lambda unrestricted i : Nat .
671    (let unrestricted nt = (naturalDivideUnchecked i 2) in (let unrestricted h = (naturalModuloUnchecked i 2) in
672      (saOp (sbStore64 saPointer (saElem (saO nt) h 0)
673        (naturalAdd (naturalMultiply nt 32) (naturalMultiply h (naturalMultiply 8 (naturalMultiply saHeadWidth 4))))
674        saWaitNone)))))
675    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.