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 2157–2184

sqdPrologueAhead

Full file
2157def sqdPrologueAhead = (lambda unrestricted ahead : Nat . (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
2158  (lambda unrestricted rest : (family SM86Program) .
2159  (let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
2160  (let unrestricted headBytes = (naturalMultiply seq (naturalMultiply saHeadWidth saHalfBytes)) in
2161  (saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
2162  (saS2R saTileIndex (saTileAxis order)
2163  (saS2R saKeyValueHead (saKeyValueAxis order)
2164  (saS2R saHead (saGroupAxis order)
2165  (saMovImm saOne 1
2166  (saMovImmAfter saScratch 31 saWait5
2167  (saImad saHead saKeyValueHead groups saHead
2168  (saCountDownTile order (naturalDivideUnchecked seq saTile)
2169  (saAnd saLane saTid saScratch
2170  (saShr saWarp saTid 5
2171  (saShr saQuad saLane 2
2172  (saMovImm saScratch 3
2173  (saAnd saQuadLane saLane saScratch
2174  (saImad saRow saWarp 16 saQuad
2175  (saImad saRow saTileIndex saTile saRow
2176  -- dQ's offset (binary32 rows), as the query pass's, for its epilogue
2177  (saMovImm saScratch 0
2178  (saImad saScratch saRow (naturalMultiply saHeadWidth 4) saScratch
2179  (saImad saScratch saQuadLane 8 saScratch
2180  (saImad saScratch saHead (naturalMultiply seq (naturalMultiply saHeadWidth 4)) saScratch
2181  -- key tiles 0 .. i, dQ from +0
2182  (saAddImm saCount saTileIndex 1
2183  (saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 68 i) 0))
2184  (sqdSetupAhead ahead seq headBytes rest))))))))))))))))))))))))))))))

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.