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 1403–1413

sqEpilogue

Full file
dQ = s (the sum): the thread's rows r and r + 8, pairs of binary32 columns
1403def sqEpilogue = (lambda unrestricted tail : (family SM86Program) .
1404  (saMovConst (sqTmp 4) (saArgument 10)
1405  (saFor 32 (lambda unrestricted i : Nat .
1406    (saFmul (naturalAdd 68 i) (naturalAdd 68 i) (sqTmp 4) (naturalSelect (naturalIsZero i) saWait3 saWaitNone)))
1407  (saWide saPointer saScratch saOne (saArgument 7)
1408  (saFor 16 (lambda unrestricted i : Nat .
1409    (let unrestricted nt = (naturalDivideUnchecked i 2) in (let unrestricted h = (naturalModuloUnchecked i 2) in
1410      (saOp (sbStore64 saPointer (saElem (saO nt) h 0)
1411        (naturalAdd (naturalMultiply nt 32) (naturalMultiply h (naturalMultiply 8 (naturalMultiply saHeadWidth 4))))
1412        saWaitNone)))))
1413    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.