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 2778–2798

sksCopyTile

Full file
chunk j: Q's (j < 4) or dO's row r + 8 (j % 4), chunk c, through pointer pair j % 2, each copy after the previous one's reads
2778def sksCopyTile = (lambda unrestricted buffer : Nat . (lambda unrestricted guarded : Nat .
2779  (lambda unrestricted tail : (family SM86Program) .
2780    (saFor 8 (lambda unrestricted j : Nat .
2781      (let unrestricted body =
2782            (constructor SM86InstructionBody SM86LoadGlobalToShared (saR sksStore)
2783              (saU (naturalAdd buffer
2784                (naturalAdd (naturalSelect (naturalLess j 4) 0 sksQueryTile)
2785                  (naturalMultiply (naturalModuloUnchecked j 4) (naturalMultiply 8 sfRowBytes)))))
2786              (saR (sksPair (naturalModuloUnchecked j 2))) (saU 0)
2787              (constructor SM86Control SM86ControlValue (byte 15) (constructor SM86YieldMode SM86Continue)
2788                saNone saSB5 (nat-to-byte saWait5) (byte 0))) in
2789      (lambda unrestricted rest : (family SM86Program) .
2790        (saOp (constructor SM86InstructionBody SM86IntegerMultiplyAddWideConstant
2791            (saR (sksPair (naturalModuloUnchecked j 2))) (saR (sksOffset (naturalModuloUnchecked j 4))) (saR saOne)
2792            (byte 0) (saU (naturalSelect (naturalLess j 4) (saArgument 0) (saArgument 3))) saPlain)
2793          (nat-eliminate (lambda unrestricted mode : Nat . (family SM86Program))
2794            (saOp body rest)
2795            (lambda unrestricted predecessor : Nat . (lambda unrestricted induction : (family SM86Program) .
2796              (saWhen saP1 body rest)))
2797            guarded)))))
2798    (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.