1	# Arithmetic stages across each eight-value cell.
2	import Base
3	import ./math_fp32.bend as Arithmetic
4	import ./kernel_tile.bend as Tile
5	
6	type Arguments is Data:
7	  Arguments{original: Tile.Tile8,value: Tile.Tile8,scaled: Tile.Tile8}
8	
9	def arguments(a: Arithmetic.DerivativeArgument,b: Arithmetic.DerivativeArgument,c: Arithmetic.DerivativeArgument,d: Arithmetic.DerivativeArgument,e: Arithmetic.DerivativeArgument,f: Arithmetic.DerivativeArgument,g: Arithmetic.DerivativeArgument,h: Arithmetic.DerivativeArgument) -> Arguments:
10	  Arithmetic.DerivativeArgument{ax,av,ascale} = a
11	  Arithmetic.DerivativeArgument{bx,bv,bs} = b
12	  Arithmetic.DerivativeArgument{cx,cv,cs} = c
13	  Arithmetic.DerivativeArgument{dx,dv,ds} = d
14	  Arithmetic.DerivativeArgument{ex,ev,es} = e
15	  Arithmetic.DerivativeArgument{fx,fv,fs} = f
16	  Arithmetic.DerivativeArgument{gx,gv,gs} = g
17	  Arithmetic.DerivativeArgument{hx,hv,hs} = h
18	  Arguments{Tile.Tile{ax,bx,cx,dx,ex,fx,gx,hx},Tile.Tile{av,bv,cv,dv,ev,fv,gv,hv},Tile.Tile{ascale,bs,cs,ds,es,fs,gs,hs}}
19	
20	def silu_first(tile: Tile.Tile8) -> Arguments:
21	  Tile.Tile{a,b,c,d,e,f,g,h} = tile
22	  arguments(Arithmetic.silu_derivative_prepare(a),Arithmetic.silu_derivative_prepare(b),Arithmetic.silu_derivative_prepare(c),Arithmetic.silu_derivative_prepare(d),Arithmetic.silu_derivative_prepare(e),Arithmetic.silu_derivative_prepare(f),Arithmetic.silu_derivative_prepare(g),Arithmetic.silu_derivative_prepare(h))
23	
24	def finished(~finish: Arithmetic.DerivativeArgument -> F32,args: Arguments) -> Tile.Tile8:
25	  Arguments{original,value,scaled} = args
26	  Tile.Tile{ax,bx,cx,dx,ex,fx,gx,hx} = original
27	  Tile.Tile{av,bv,cv,dv,ev,fv,gv,hv} = value
28	  Tile.Tile{ascale,bs,cs,ds,es,fs,gs,hs} = scaled
29	  Tile.Tile{finish(Arithmetic.DerivativeArgument{ax,av,ascale}),finish(Arithmetic.DerivativeArgument{bx,bv,bs}),finish(Arithmetic.DerivativeArgument{cx,cv,cs}),finish(Arithmetic.DerivativeArgument{dx,dv,ds}),finish(Arithmetic.DerivativeArgument{ex,ev,es}),finish(Arithmetic.DerivativeArgument{fx,fv,fs}),finish(Arithmetic.DerivativeArgument{gx,gv,gs}),finish(Arithmetic.DerivativeArgument{hx,hv,hs})}
30	
31	def silu(context: Unit,tile: Tile.Tile8) -> Tile.Tile8:
32	  finished(~Arithmetic.silu_derivative_finish,silu_first(tile))
33	
34	type GeluPrepared is Data:
35	  GeluPrepared{original: Tile.Tile8,head: Tile.Tile8,low: Tile.Tile8,slope: Tile.Tile8,slope_error: Tile.Tile8}
36	
37	def gelu_prepared(a: Arithmetic.GeluPrepared,b: Arithmetic.GeluPrepared,c: Arithmetic.GeluPrepared,d: Arithmetic.GeluPrepared,e: Arithmetic.GeluPrepared,f: Arithmetic.GeluPrepared,g: Arithmetic.GeluPrepared,h: Arithmetic.GeluPrepared) -> GeluPrepared:
38	  Arithmetic.GeluPrepared{a0,a1,a2,a3,a4} = a
39	  Arithmetic.GeluPrepared{b0,b1,b2,b3,b4} = b
40	  Arithmetic.GeluPrepared{c0,c1,c2,c3,c4} = c
41	  Arithmetic.GeluPrepared{d0,d1,d2,d3,d4} = d
42	  Arithmetic.GeluPrepared{e0,e1,e2,e3,e4} = e
43	  Arithmetic.GeluPrepared{f0,f1,f2,f3,f4} = f
44	  Arithmetic.GeluPrepared{g0,g1,g2,g3,g4} = g
45	  Arithmetic.GeluPrepared{h0,h1,h2,h3,h4} = h
46	  GeluPrepared{Tile.Tile{a0,b0,c0,d0,e0,f0,g0,h0},Tile.Tile{a1,b1,c1,d1,e1,f1,g1,h1},Tile.Tile{a2,b2,c2,d2,e2,f2,g2,h2},Tile.Tile{a3,b3,c3,d3,e3,f3,g3,h3},Tile.Tile{a4,b4,c4,d4,e4,f4,g4,h4}}
47	
48	def gelu_first(tile: Tile.Tile8) -> GeluPrepared:
49	  Tile.Tile{a,b,c,d,e,f,g,h} = tile
50	  gelu_prepared(Arithmetic.gelu_derivative_prepare(a),Arithmetic.gelu_derivative_prepare(b),Arithmetic.gelu_derivative_prepare(c),Arithmetic.gelu_derivative_prepare(d),Arithmetic.gelu_derivative_prepare(e),Arithmetic.gelu_derivative_prepare(f),Arithmetic.gelu_derivative_prepare(g),Arithmetic.gelu_derivative_prepare(h))
51	
52	type GeluReduced is Data:
53	  GeluReduced{original: Tile.Tile8,power: Tile.Tile8,denominator: Tile.Tile8,error: Tile.Tile8,slope: Tile.Tile8,slope_error: Tile.Tile8}
54	
55	def gelu_reduced(a: Arithmetic.GeluReduced,b: Arithmetic.GeluReduced,c: Arithmetic.GeluReduced,d: Arithmetic.GeluReduced,e: Arithmetic.GeluReduced,f: Arithmetic.GeluReduced,g: Arithmetic.GeluReduced,h: Arithmetic.GeluReduced) -> GeluReduced:
56	  Arithmetic.GeluReduced{a0,a1,a2,a3,a4,a5} = a
57	  Arithmetic.GeluReduced{b0,b1,b2,b3,b4,b5} = b
58	  Arithmetic.GeluReduced{c0,c1,c2,c3,c4,c5} = c
59	  Arithmetic.GeluReduced{d0,d1,d2,d3,d4,d5} = d
60	  Arithmetic.GeluReduced{e0,e1,e2,e3,e4,e5} = e
61	  Arithmetic.GeluReduced{f0,f1,f2,f3,f4,f5} = f
62	  Arithmetic.GeluReduced{g0,g1,g2,g3,g4,g5} = g
63	  Arithmetic.GeluReduced{h0,h1,h2,h3,h4,h5} = h
64	  GeluReduced{Tile.Tile{a0,b0,c0,d0,e0,f0,g0,h0},Tile.Tile{a1,b1,c1,d1,e1,f1,g1,h1},Tile.Tile{a2,b2,c2,d2,e2,f2,g2,h2},Tile.Tile{a3,b3,c3,d3,e3,f3,g3,h3},Tile.Tile{a4,b4,c4,d4,e4,f4,g4,h4},Tile.Tile{a5,b5,c5,d5,e5,f5,g5,h5}}
65	
66	def gelu_second(args: GeluPrepared) -> GeluReduced:
67	  GeluPrepared{original,head,low,slope,slope_error} = args
68	  Tile.Tile{a0,b0,c0,d0,e0,f0,g0,h0} = original
69	  Tile.Tile{a1,b1,c1,d1,e1,f1,g1,h1} = head
70	  Tile.Tile{a2,b2,c2,d2,e2,f2,g2,h2} = low
71	  Tile.Tile{a3,b3,c3,d3,e3,f3,g3,h3} = slope
72	  Tile.Tile{a4,b4,c4,d4,e4,f4,g4,h4} = slope_error
73	  gelu_reduced(Arithmetic.gelu_derivative_exponential(Arithmetic.GeluPrepared{a0,a1,a2,a3,a4}),Arithmetic.gelu_derivative_exponential(Arithmetic.GeluPrepared{b0,b1,b2,b3,b4}),Arithmetic.gelu_derivative_exponential(Arithmetic.GeluPrepared{c0,c1,c2,c3,c4}),Arithmetic.gelu_derivative_exponential(Arithmetic.GeluPrepared{d0,d1,d2,d3,d4}),Arithmetic.gelu_derivative_exponential(Arithmetic.GeluPrepared{e0,e1,e2,e3,e4}),Arithmetic.gelu_derivative_exponential(Arithmetic.GeluPrepared{f0,f1,f2,f3,f4}),Arithmetic.gelu_derivative_exponential(Arithmetic.GeluPrepared{g0,g1,g2,g3,g4}),Arithmetic.gelu_derivative_exponential(Arithmetic.GeluPrepared{h0,h1,h2,h3,h4}))
74	
75	def gelu_third(args: GeluReduced) -> Arguments:
76	  GeluReduced{original,power,denominator,error,slope,slope_error} = args
77	  Tile.Tile{a0,b0,c0,d0,e0,f0,g0,h0} = original
78	  Tile.Tile{a1,b1,c1,d1,e1,f1,g1,h1} = power
79	  Tile.Tile{a2,b2,c2,d2,e2,f2,g2,h2} = denominator
80	  Tile.Tile{a3,b3,c3,d3,e3,f3,g3,h3} = error
81	  Tile.Tile{a4,b4,c4,d4,e4,f4,g4,h4} = slope
82	  Tile.Tile{a5,b5,c5,d5,e5,f5,g5,h5} = slope_error
83	  arguments(Arithmetic.gelu_derivative_evaluated(Arithmetic.GeluReduced{a0,a1,a2,a3,a4,a5}),Arithmetic.gelu_derivative_evaluated(Arithmetic.GeluReduced{b0,b1,b2,b3,b4,b5}),Arithmetic.gelu_derivative_evaluated(Arithmetic.GeluReduced{c0,c1,c2,c3,c4,c5}),Arithmetic.gelu_derivative_evaluated(Arithmetic.GeluReduced{d0,d1,d2,d3,d4,d5}),Arithmetic.gelu_derivative_evaluated(Arithmetic.GeluReduced{e0,e1,e2,e3,e4,e5}),Arithmetic.gelu_derivative_evaluated(Arithmetic.GeluReduced{f0,f1,f2,f3,f4,f5}),Arithmetic.gelu_derivative_evaluated(Arithmetic.GeluReduced{g0,g1,g2,g3,g4,g5}),Arithmetic.gelu_derivative_evaluated(Arithmetic.GeluReduced{h0,h1,h2,h3,h4,h5}))
84	
85	def gelu_tanh(context: Unit,tile: Tile.Tile8) -> Tile.Tile8:
86	  finished(~Arithmetic.gelu_derivative_finish,gelu_third(gelu_second(gelu_first(tile))))