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 649–658

sfLoadFragmentsTransposed

Full file
The same B fragments from a tile staged as rows of the other operand (keys down, columns across): LDSM.T (k, p) reads keys 16 k.., columns 16 p.. through `lane` (lane L names key 8 ((L / 8) % 2) + L % 8, column 8 (L / 16)), filling the registers sfLoadFragments fills from the transposed tile.
649def sfLoadFragmentsTransposed = (lambda unrestricted lane : Nat . (lambda unrestricted base : Nat . (lambda unrestricted wait : Nat .
650  (lambda unrestricted tail : (family SM86Program) .
651    (saFor 16 (lambda unrestricted i : Nat .
652      (let unrestricted k = (naturalDivideUnchecked i 4) in (let unrestricted pair = (naturalModuloUnchecked i 4) in
653        (saOp (constructor SM86InstructionBody SM86LoadSharedMatrix (saR (saB k (naturalMultiply 2 pair))) (saR lane)
654          (saU (naturalAdd base (naturalAdd (naturalMultiply k (naturalMultiply 16 sfRowBytes)) (naturalMultiply pair 32))))
655          (constructor SM86SharedMatrixCount SM86SharedMatrix4)
656          (constructor SM86SharedMatrixTranspose SM86SharedMatrixTransposed)
657          (saCtl saSB1 (naturalSelect (naturalIsZero i) wait saWaitNone)))))))
658      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.