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 1655–1670

sqsLoadFragments

Full file
a tile's B fragments from shared `base` through lane address `lane`, plain (rows as the forward's) or transposed (LDSM.T (kk, p): keys 16 kk.., columns 16 p..; lane L names key 8 ((L / 8) % 2) + L % 8, column 8 (L / 16))
1655def sqsLoadFragments = (lambda unrestricted lane : Nat . (lambda unrestricted transposed : Nat .
1656  (lambda unrestricted base : Nat . (lambda unrestricted wait : Nat .
1657  (lambda unrestricted tail : (family SM86Program) .
1658    (saFor 16 (lambda unrestricted i : Nat .
1659      (let unrestricted k = (naturalDivideUnchecked i 4) in (let unrestricted pair = (naturalModuloUnchecked i 4) in
1660        (saOp (constructor SM86InstructionBody SM86LoadSharedMatrix (saR (saB k (naturalMultiply 2 pair))) (saR lane)
1661          (saU (naturalAdd base
1662            (naturalSelect transposed
1663              (naturalAdd (naturalMultiply k (naturalMultiply 16 sfRowBytes)) (naturalMultiply pair 32))
1664              (naturalAdd (naturalMultiply pair (naturalMultiply 16 sfRowBytes)) (naturalMultiply k 32)))))
1665          (constructor SM86SharedMatrixCount SM86SharedMatrix4)
1666          (eliminate StdBool (lambda unrestricted c : (family StdBool) . (family SM86SharedMatrixTranspose)) (stdBoolFromNatural transposed)
1667            (branch StdTrue . (constructor SM86SharedMatrixTranspose SM86SharedMatrixTransposed))
1668            (branch StdFalse . (constructor SM86SharedMatrixTranspose SM86SharedMatrixNotTransposed)))
1669          (saCtl saSB1 (naturalSelect (naturalIsZero i) wait saWaitNone)))))))
1670      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.