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 2742–2750

sksScoreGradientWord

Full file
A packed dS^T word, its read barrier SB5: the next sub-tile's row-term loads and copies wait on SB5 before anything rewrites the packed words or the address pair, and each word after the first waits for the one before it to have read the shared pair (one live SB5 token, as the copies'). Word i: slice kk = i / 4, register j = i % 4 of skSA kk -- row r + 8 (j % 2), n-tile 2 kk + j / 2, its column pair.
2742def sksScoreGradientWord = (lambda unrestricted seq : Nat . (lambda unrestricted value : (pi unrestricted i : Nat . Nat) .
2743  (lambda unrestricted i : Nat .
2744    (let unrestricted j = (naturalModuloUnchecked i 4) in
2745    (let unrestricted nt = (naturalAdd (naturalMultiply 2 (naturalDivideUnchecked i 4)) (naturalDivideUnchecked j 2)) in
2746      (constructor SM86InstructionBody SM86StoreGlobal (saR sksScoreGradientPointer) (saR (value i))
2747        (saU (naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 sksScoreGradientRowBytes))
2748          (naturalMultiply nt 16)))
2749        (constructor SM86Control SM86ControlValue (byte 15) (constructor SM86YieldMode SM86Continue)
2750          saNone saSB5 (nat-to-byte (naturalSelect (naturalIsZero i) saWaitNone saWait5)) (byte 0))))))))

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.