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 2803–2821

sksLoadFragments

Full file
B fragments from shared `base`: plain, 4 slices x 4 n-tiles (32 rows) -- LDSM.x4 (k, p) n-tiles 2p, 2p + 1 of slice k -- or transposed, 2 slices x 8 n-tiles (64 columns) -- LDSM.T (kk, p) rows 16 kk.., columns 16 p..
2803def sksLoadFragments = (lambda unrestricted transposed : Nat .
2804  (lambda unrestricted base : Nat . (lambda unrestricted wait : Nat .
2805  (lambda unrestricted tail : (family SM86Program) .
2806    (saFor 8 (lambda unrestricted i : Nat .
2807      (let unrestricted k = (naturalSelect transposed (naturalDivideUnchecked i 4) (naturalDivideUnchecked i 2)) in
2808      (let unrestricted pair = (naturalSelect transposed (naturalModuloUnchecked i 4) (naturalModuloUnchecked i 2)) in
2809        (saOp (constructor SM86InstructionBody SM86LoadSharedMatrix
2810          (saR (skB (naturalSelect transposed 8 4) k (naturalMultiply 2 pair)))
2811          (saR (naturalSelect transposed sksLaneTransposed sksLane))
2812          (saU (naturalAdd base
2813            (naturalSelect transposed
2814              (naturalAdd (naturalMultiply k (naturalMultiply 16 sfRowBytes)) (naturalMultiply pair 32))
2815              (naturalAdd (naturalMultiply pair (naturalMultiply 16 sfRowBytes)) (naturalMultiply k 32)))))
2816          (constructor SM86SharedMatrixCount SM86SharedMatrix4)
2817          (eliminate StdBool (lambda unrestricted c : (family StdBool) . (family SM86SharedMatrixTranspose)) (stdBoolFromNatural transposed)
2818            (branch StdTrue . (constructor SM86SharedMatrixTranspose SM86SharedMatrixTransposed))
2819            (branch StdFalse . (constructor SM86SharedMatrixTranspose SM86SharedMatrixNotTransposed)))
2820          (saCtl saSB1 (naturalSelect (naturalIsZero i) wait saWaitNone)))))))
2821      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.