1	# Two arithmetic stages for each eight-value cell.
2	import Base
3	import ./tensor_access.bend as Access
4	import ./math_activation.bend as Act
5	import ./math_fp32.bend as G
6	import ./kernel_tile.bend as Tile
7	import ./traversal_lanes.bend as Lanes
8	import ./traversal_cells.bend as Cells
9	import ./traversal_array.bend as Traversal
10	
11	type Arguments is Data:
12	  Arguments{error: Tile.Tile8,head: Tile.Tile8,original: Tile.Tile8}
13	
14	def 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	
25	def 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	
29	def 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	
36	def gathered(read: Array<F32> & Tile.Tile8) -> Array<F32> & Tile.Tile8:
37	  (values,tile) = read
38	  (values,second(first(tile)))
39	
40	def identity(context: Unit,value: F32) -> F32:
41	  value
42	
43	def gather(values: Array<F32>,position: U32,context: Unit) -> Array<F32> & Tile.Tile8:
44	  gathered(Lanes.gather(~Unit,~identity,values,position,context))
45	
46	def 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	
49	def 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)