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 2920–2949

sksSetup

Full file
lane addresses, the copy role, Q/dO and L/D offsets at the last sub-tile, then that sub-tile into buffer 0
2920def sksSetup = (lambda unrestricted seq : Nat . (lambda unrestricted headBytes : Nat .
2921  (lambda unrestricted tail : (family SM86Program) .
2922  (let unrestricted last = (naturalSaturatingSubtract (naturalDivideUnchecked seq skSubTile) 1) in
2923  (saShr sksT0 saLane 4
2924  (saMovImm sksT1 7
2925  (saAnd sksT1 saLane sksT1
2926  (saImad sksT1 sksT0 8 sksT1
2927  (saMovImm sksLane 0
2928  (saImad sksLane sksT1 sfRowBytes sksLane
2929  (saShr sksT1 saLane 3
2930  (saAnd sksT1 sksT1 saOne
2931  (saImad sksLane sksT1 16 sksLane
2932  (saMovImm sksLaneTransposed 7
2933  (saAnd sksLaneTransposed saLane sksLaneTransposed
2934  (saImad sksT1 sksT1 8 sksLaneTransposed
2935  (saMovImm sksLaneTransposed 0
2936  (saImad sksLaneTransposed sksT1 sfRowBytes sksLaneTransposed
2937  (saImad sksLaneTransposed sksT0 16 sksLaneTransposed
2938  (saShr sksT0 saTid 3
2939  (saMovImm sksT1 7
2940  (saAnd sksT1 saTid sksT1
2941  (saMovImm sksChunkBytes 0
2942  (saImad sksChunkBytes sksT1 16 sksChunkBytes
2943  (saImad sksStore sksT0 sfRowBytes sksChunkBytes
2944  (saImad (sksOffset 0) sksT0 (naturalMultiply saHeadWidth saHalfBytes) sksChunkBytes
2945  (saImad (sksOffset 0) skHead headBytes (sksOffset 0)
2946  (saAddImm (sksOffset 0) (sksOffset 0) (naturalMultiply last (naturalMultiply skSubTile (naturalMultiply saHeadWidth saHalfBytes)))
2947  (saFor 3 (lambda unrestricted i : Nat .
2948    (saAddImm (sksOffset (succ i)) (sksOffset 0) (naturalMultiply (succ i) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))))
2949  (sksCopyTile 0 0 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.