math_gelu_lanes.bend
stelliferous@0.0.3.0 · 52 lines · 0 definitions
1# Two arithmetic stages for each eight-value cell.
2import Base
3import stelliferous@0.0.3.0/tensor_access.bend as Access
4import stelliferous@0.0.3.0/math_activation.bend as Act
5import stelliferous@0.0.3.0/math_fp32.bend as G
6import stelliferous@0.0.3.0/kernel_tile.bend as Tile
7import stelliferous@0.0.3.0/traversal_lanes.bend as Lanes
8import stelliferous@0.0.3.0/traversal_cells.bend as Cells
9import stelliferous@0.0.3.0/traversal_array.bend as Traversal
10
11type Arguments is Data:
12 Arguments{error: Tile.Tile8,head: Tile.Tile8,original: Tile.Tile8}
13
14def arguments(a: G.Argument,b: G.Argument,c: G.Argument,d: G.Argument,e: G.Argument,f: G.Argument,g: G.Argument,h: G.Argument) -> Arguments:
15 G.Argument{ax,aw,ae} = a
16 G.Argument{bx,bw,be} = b
17 G.Argument{cx,cw,ce} = c
18 G.Argument{dx,dw,de} = d
19 G.Argument{ex,ew,ee} = e
20 G.Argument{fx,fw,fe} = f
21 G.Argument{gx,gw,ge} = g
22 G.Argument{hx,hw,he} = h
23 Arguments{Tile.Tile{ae,be,ce,de,ee,fe,ge,he},Tile.Tile{aw,bw,cw,dw,ew,fw,gw,hw},Tile.Tile{ax,bx,cx,dx,ex,fx,gx,hx}}
24
25def first(tile: Tile.Tile8) -> Arguments:
26 Tile.Tile{a,b,c,d,e,f,g,h} = tile
27 arguments(G.gelu_prepared(a),G.gelu_prepared(b),G.gelu_prepared(c),G.gelu_prepared(d),G.gelu_prepared(e),G.gelu_prepared(f),G.gelu_prepared(g),G.gelu_prepared(h))
28
29def second(args: Arguments) -> Tile.Tile8:
30 Arguments{error,head,original} = args
31 Tile.Tile{ae,be,ce,de,ee,fe,ge,he} = error
32 Tile.Tile{aw,bw,cw,dw,ew,fw,gw,hw} = head
33 Tile.Tile{ax,bx,cx,dx,ex,fx,gx,hx} = original
34 Tile.Tile{G.gelu_finish(G.Argument{ax,aw,ae}),G.gelu_finish(G.Argument{bx,bw,be}),G.gelu_finish(G.Argument{cx,cw,ce}),G.gelu_finish(G.Argument{dx,dw,de}),G.gelu_finish(G.Argument{ex,ew,ee}),G.gelu_finish(G.Argument{fx,fw,fe}),G.gelu_finish(G.Argument{gx,gw,ge}),G.gelu_finish(G.Argument{hx,hw,he})}
35
36def gathered(read: Array<F32> & Tile.Tile8) -> Array<F32> & Tile.Tile8:
37 (values,tile) = read
38 (values,second(first(tile)))
39
40def identity(context: Unit,value: F32) -> F32:
41 value
42
43def gather(values: Array<F32>,position: U32,context: Unit) -> Array<F32> & Tile.Tile8:
44 gathered(Lanes.gather(~Unit,~identity,values,position,context))
45
46def map(count: Nat,position: U32,values: Array<F32>) -> Array<F32>:
47 Lanes.rounds_with(~Array<F32>,~Unit,~gather,~Lanes.scatter,~Lanes.finish(~Unit,~Access.apply_element(~F32,~Act.gelu_tanh_value)),count,Cells.State{values,Lanes.scratch(),position,0,Unit{}})
48
49def map_when(covered: Bool,count: Nat,position: U32,values: Array<F32>) -> Array<F32>:
50 match covered:
51 case True{}: map(count,position,values)
52 case False{}: Traversal.map_in_place(~F32,~Unit,~Access.apply_element(~F32,~Act.gelu_tanh_value),count,position,Unit{},values)