Source/Packages

Realization.Nvidia.SM86.StreamingAttentionSM86

packages/realizations/cooperative/nvidia-sm86/src/Realization/Nvidia/SM86/StreamingAttentionSM86.alpha

1,203 lines221 declarations70.9 KiBSHA-256 23d4e5e2aa2a

def · lines 279–289

saRowMax

Full file
the row's maximum over the quad's 64 columns, in tmp h
279def saRowMax = (lambda unrestricted h : Nat . (lambda unrestricted tail : (family SM86Program) .
280  (let unrestricted into = (saTmp h) in
281  (saFmax into (saElem (saS 0) h 0) (saElem (saS 0) h 1) saWaitNone
282    (saFor 14 (lambda unrestricted j : Nat .
283      (let unrestricted i = (naturalAdd j 2) in
284        (saFmax into into (saElem (saS (naturalDivideUnchecked i 2)) h (naturalModuloUnchecked i 2)) saWaitNone)))
285      (saShfl (saTmp 2) into 1 saWaitNone
286      (saFmax into into (saTmp 2) saWait5
287      (saShfl (saTmp 2) into 2 saWaitNone
288      (saFmax into into (saTmp 2) saWait5
289        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.