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 945–968

skEpilogue

Full file
dK = s (the sum), rows; dV^T, columns: element (key r + 8h, column 8 nt + 2 (lane % 4) + b) at (8 nt + b) seq + 8 h words past the thread's
945def skEpilogue = (lambda unrestricted seq : Nat . (lambda unrestricted tail : (family SM86Program) .
946  (saMovConst (skTmp 0) (saArgument 11)
947  (saFor 32 (lambda unrestricted i : Nat .
948    (saFmul (naturalAdd 116 i) (naturalAdd 116 i) (skTmp 0) (naturalSelect (naturalIsZero i) saWait3 saWaitNone)))
949  (saWide saPointer saScratch saOne (saArgument 8)
950  (saFor 16 (lambda unrestricted i : Nat .
951    (let unrestricted nt = (naturalDivideUnchecked i 2) in (let unrestricted h = (naturalModuloUnchecked i 2) in
952      (saOp (sbStore64 saPointer (saElem (skDK nt) h 0)
953        (naturalAdd (naturalMultiply nt 32) (naturalMultiply h (naturalMultiply 8 (naturalMultiply saHeadWidth 4))))
954        saWaitNone)))))
955  -- dV^T's offset: head h's 64 x seq plane, the quad lane's first column,
956  -- the key
957  (saMovImm (skTmp 1) 0
958  (saImad (skTmp 1) saQuadLane (naturalMultiply 8 seq) (skTmp 1)
959  (saImad (skTmp 1) skHead (naturalMultiply saHeadWidth (naturalMultiply seq 4)) (skTmp 1)
960  (saImad (skTmp 1) skRow 4 (skTmp 1)
961  (saWide saPointer (skTmp 1) saOne (saArgument 9)
962  (saFor 32 (lambda unrestricted i : Nat .
963    (let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
964      (let unrestricted b = (naturalModuloUnchecked i 2) in
965      (saOp (saStore saPointer (saElem (skDV nt) h b)
966        (naturalAdd (naturalMultiply (naturalAdd (naturalMultiply 8 nt) b) (naturalMultiply seq 4)) (naturalMultiply h 32))
967        saWaitNone))))))
968    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.