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 2063–2072

sqdLoadKeys

Full file
K^T's B fragments of slice k (keys 16 k..), n-tile pair p (columns 16 p..), LDSM.T through pair p's lane: the registers the query pass's sqsLoadFragments fills from its padded tile
2063def sqdLoadKeys = (lambda unrestricted base : Nat . (lambda unrestricted wait : Nat .
2064  (lambda unrestricted tail : (family SM86Program) .
2065    (saFor 16 (lambda unrestricted i : Nat .
2066      (let unrestricted k = (naturalDivideUnchecked i 4) in (let unrestricted pair = (naturalModuloUnchecked i 4) in
2067        (saOp (constructor SM86InstructionBody SM86LoadSharedMatrix (saR (saB k (naturalMultiply 2 pair))) (saR (sqdLaneKeys pair))
2068          (saU (naturalAdd base (naturalMultiply k (naturalMultiply 16 sqdRowBytes))))
2069          (constructor SM86SharedMatrixCount SM86SharedMatrix4)
2070          (constructor SM86SharedMatrixTranspose SM86SharedMatrixTransposed)
2071          (saCtl saSB1 (naturalSelect (naturalIsZero i) wait saWaitNone)))))))
2072      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.