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 1630–1650

sqsCopyTile

Full file
chunk j: K's (j < 4) or V's row r + 16 (j % 4), chunk c, through pointer pair j % 2; each copy waits for the previous one's register reads, which also frees the pair the next pointer re-forms
1630def sqsCopyTile = (lambda unrestricted buffer : Nat . (lambda unrestricted guarded : Nat .
1631  (lambda unrestricted tail : (family SM86Program) .
1632    (saFor 8 (lambda unrestricted j : Nat .
1633      (let unrestricted body =
1634            (constructor SM86InstructionBody SM86LoadGlobalToShared (saR sqsStore)
1635              (saU (naturalAdd buffer
1636                (naturalAdd (naturalSelect (naturalLess j 4) 0 sfValueTile)
1637                  (naturalMultiply (naturalModuloUnchecked j 4) (naturalMultiply 16 sfRowBytes)))))
1638              (saR (sqsPair (naturalModuloUnchecked j 2))) (saU 0)
1639              (constructor SM86Control SM86ControlValue (byte 15) (constructor SM86YieldMode SM86Continue)
1640                saNone saSB5 (nat-to-byte saWait5) (byte 0))) in
1641      (lambda unrestricted rest : (family SM86Program) .
1642        (saOp (constructor SM86InstructionBody SM86IntegerMultiplyAddWideConstant
1643            (saR (sqsPair (naturalModuloUnchecked j 2))) (saR (sqsOffset (naturalModuloUnchecked j 4))) (saR saOne)
1644            (byte 0) (saU (naturalSelect (naturalLess j 4) (saArgument 1) (saArgument 2))) saPlain)
1645          (nat-eliminate (lambda unrestricted mode : Nat . (family SM86Program))
1646            (saOp body rest)
1647            (lambda unrestricted predecessor : Nat . (lambda unrestricted induction : (family SM86Program) .
1648              (saWhen saP1 body rest)))
1649            guarded)))))
1650    (sfCommit 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.