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 294–304

saRowMax

Full file
the row's maximum over the quad's 64 columns, in tmp h
294def saRowMax = (lambda unrestricted h : Nat . (lambda unrestricted tail : (family SM86Program) .
295  (let unrestricted into = (saTmp h) in
296  (saFmax into (saElem (saS 0) h 0) (saElem (saS 0) h 1) saWaitNone
297    (saFor 14 (lambda unrestricted j : Nat .
298      (let unrestricted i = (naturalAdd j 2) in
299        (saFmax into into (saElem (saS (naturalDivideUnchecked i 2)) h (naturalModuloUnchecked i 2)) saWaitNone)))
300      (saShfl (saTmp 2) into 1 saWaitNone
301      (saFmax into into (saTmp 2) saWait5
302      (saShfl (saTmp 2) into 2 saWaitNone
303      (saFmax into into (saTmp 2) saWait5
304        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.