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 2567–2582

skpStorePairs

Full file
the 16 pairs of `accumulator` (n-tile nt, row half h) from `pointer` (binary32) and `halfPointer` (bfloat16), each packed into `packed` i
2567def skpStorePairs = (lambda unrestricted width : Nat .
2568  (lambda unrestricted accumulator : (pi unrestricted nt : Nat . Nat) . (lambda unrestricted packed : (pi unrestricted i : Nat . Nat) .
2569  (lambda unrestricted pointer : Nat . (lambda unrestricted halfPointer : Nat . (lambda unrestricted tail : (family SM86Program) .
2570  (saFor 16 (lambda unrestricted i : Nat .
2571    (let unrestricted nt = (naturalDivideUnchecked i 2) in (let unrestricted h = (naturalModuloUnchecked i 2) in
2572    (lambda unrestricted after : (family SM86Program) .
2573      (saOp (sbStore64 pointer (saElem (accumulator nt) h 0)
2574        (naturalAdd (naturalMultiply nt 32) (naturalMultiply h (naturalMultiply 8 (naturalMultiply width 4))))
2575        saWaitNone)
2576      (saOp (constructor SM86InstructionBody SM86FloatPairToPackedHalfPair
2577          (saR (packed i)) (saR (saElem (accumulator nt) h 1)) (saR (saElem (accumulator nt) h 0)) saPlain)
2578      (saOp (saStore halfPointer (packed i)
2579        (naturalAdd (naturalMultiply nt 16) (naturalMultiply h (naturalMultiply 8 (naturalMultiply width saHalfBytes))))
2580        saWaitNone)
2581        after)))))))
2582    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.