Source/Packages

Realization.Nvidia.SM86.GatedGELUSM86

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

170 lines18 declarations7.4 KiBSHA-256 fe6a7ec05912

def · lines 108–136

ggBackwardElement

Full file
---- backward ---- R4:R5 dGated, R6:R7 gelu, R8:R9 up, R10:R11 gate, R12:R13 upGradient, R14:R15 gateGradient; R22, R23 dGated; R24, R25 gelu; R26, R27 up; R28, R29 gate; R30 s, R31 c, R32 h, R33 1, R34 3c; element e's from 36 + 16e; R70, R71 the half pairs
108def ggBackwardElement = (lambda unrestricted e : Nat . (lambda unrestricted tail : (family SM86Program) .
109  (let unrestricted b = (naturalAdd 36 (naturalMultiply 16 e)) in
110  (let unrestricted x = (naturalAdd 28 e) in
111  (let unrestricted r = (lambda unrestricted i : Nat . (naturalAdd b i)) in
112  -- upGradient = dGated gelu; g = dGated up
113  (saFmul (r 0) (naturalAdd 22 e) (naturalAdd 24 e) (naturalSelect e saWaitNone saWait0)
114  (saFmul (r 1) (naturalAdd 22 e) (naturalAdd 26 e) saWaitNone
115  -- t = tanh(s (x^3 c + x))
116  (saFmul (r 2) x x saWaitNone
117  (saFmul (r 3) (r 2) x saWaitNone
118  (ggFfma (r 4) (r 3) 31 x
119  (saFmul (r 4) (r 4) 30 saWaitNone
120  (ggTanh (r 5) (r 4)
121  -- 1 - t^2
122  (saFmul (r 6) (r 5) (r 5) saWait1
123  (saFneg (r 6) (r 6)
124  (saFadd (r 6) 33 (r 6) saWaitNone
125  -- (x^2 3c + 1) s
126  (ggFfma (r 7) (r 2) 34 33
127  (saFmul (r 7) (r 7) 30 saWaitNone
128  -- (1 + t) h
129  (saFadd (r 8) 33 (r 5) saWaitNone
130  (saFmul (r 8) (r 8) 32 saWaitNone
131  -- x (1 - t^2) (x^2 3c + 1) s h
132  (saFmul (r 9) x (r 6) saWaitNone
133  (saFmul (r 9) (r 9) (r 7) saWaitNone
134  (saFmul (r 9) (r 9) 32 saWaitNone
135  (saFadd (r 10) (r 8) (r 9) saWaitNone
136  (saFmul (r 11) (r 1) (r 10) saWaitNone 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.