math_activation_lanes.bend
stelliferous@0.0.3.0 · 86 lines · 0 definitions
1# Arithmetic stages across each eight-value cell.
2import Base
3import stelliferous@0.0.3.0/math_fp32.bend as Arithmetic
4import stelliferous@0.0.3.0/kernel_tile.bend as Tile
5
6type Arguments is Data:
7 Arguments{original: Tile.Tile8,value: Tile.Tile8,scaled: Tile.Tile8}
8
9def 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
20def 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
24def 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
31def silu(context: Unit,tile: Tile.Tile8) -> Tile.Tile8:
32 finished(~Arithmetic.silu_derivative_finish,silu_first(tile))
33
34type GeluPrepared is Data:
35 GeluPrepared{original: Tile.Tile8,head: Tile.Tile8,low: Tile.Tile8,slope: Tile.Tile8,slope_error: Tile.Tile8}
36
37def 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
48def 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
52type 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
55def 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
66def 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
75def 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
85def gelu_tanh(context: Unit,tile: Tile.Tile8) -> Tile.Tile8:
86 finished(~Arithmetic.gelu_derivative_finish,gelu_third(gelu_second(gelu_first(tile))))