Source/Packages

Realization.Nvidia.SM86.StreamingAttentionSM86

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

1,129 lines205 declarations66.1 KiBSHA-256 aefcbc3b00ca

def · lines 267–277

saRowMax

Full file
the row's maximum over the quad's 64 columns, in tmp h
267def saRowMax = (lambda unrestricted h : Nat . (lambda unrestricted tail : (family SM86Program) .
268  (let unrestricted into = (saTmp h) in
269  (saFmax into (saElem (saS 0) h 0) (saElem (saS 0) h 1) saWaitNone
270    (saFor 14 (lambda unrestricted j : Nat .
271      (let unrestricted i = (naturalAdd j 2) in
272        (saFmax into into (saElem (saS (naturalDivideUnchecked i 2)) h (naturalModuloUnchecked i 2)) saWaitNone)))
273      (saShfl (saTmp 2) into 1 saWaitNone
274      (saFmax into into (saTmp 2) saWait5
275      (saShfl (saTmp 2) into 2 saWaitNone
276      (saFmax into into (saTmp 2) saWait5
277        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.