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 633–642

sfLoadFragments

Full file
The tile's B fragments from shared `base`: LDSM.x4 (k, p) fills n-tiles 2p and 2p + 1 of slice k, rows 16 p.. and halves 16 k.. of the tile. Lane L names row 8 (L / 16) + L % 8 and half 8 ((L / 8) % 2) of its matrix.
633def sfLoadFragments = (lambda unrestricted base : Nat . (lambda unrestricted wait : Nat .
634  (lambda unrestricted tail : (family SM86Program) .
635    (saFor 16 (lambda unrestricted i : Nat .
636      (let unrestricted k = (naturalDivideUnchecked i 4) in (let unrestricted pair = (naturalModuloUnchecked i 4) in
637        (saOp (constructor SM86InstructionBody SM86LoadSharedMatrix (saR (saB k (naturalMultiply 2 pair))) (saR sfLane)
638          (saU (naturalAdd base (naturalAdd (naturalMultiply pair (naturalMultiply 16 sfRowBytes)) (naturalMultiply k 32))))
639          (constructor SM86SharedMatrixCount SM86SharedMatrix4)
640          (constructor SM86SharedMatrixTranspose SM86SharedMatrixNotTransposed)
641          (saCtl saSB1 (naturalSelect (naturalIsZero i) wait saWaitNone)))))))
642      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.