1	# Measured flat FP32 tile representation; logical views stay out of hot loops.
2	import Base
3	import ./kernel_block.bend as Block
4	import ./math_sequence.bend as Sequence
5	import ./math_fp32.bend as Fp32
6	
7	type Tile8 is Data:
8	  Tile{a0: F32, a1: F32, a2: F32, a3: F32, a4: F32, a5: F32, a6: F32, a7: F32}
9	
10	def filled(+value: F32) -> Tile8:
11	  Tile{value,value,value,value,value,value,value,value}
12	
13	def lane0(tile: Tile8) -> F32:
14	  Tile{a,b,c,d,e,f,g,h} = tile
15	  a
16	
17	def lane1(tile: Tile8) -> F32:
18	  Tile{a,b,c,d,e,f,g,h} = tile
19	  b
20	
21	def lane2(tile: Tile8) -> F32:
22	  Tile{a,b,c,d,e,f,g,h} = tile
23	  c
24	
25	def lane3(tile: Tile8) -> F32:
26	  Tile{a,b,c,d,e,f,g,h} = tile
27	  d
28	
29	def lane4(tile: Tile8) -> F32:
30	  Tile{a,b,c,d,e,f,g,h} = tile
31	  e
32	
33	def lane5(tile: Tile8) -> F32:
34	  Tile{a,b,c,d,e,f,g,h} = tile
35	  f
36	
37	def lane6(tile: Tile8) -> F32:
38	  Tile{a,b,c,d,e,f,g,h} = tile
39	  g
40	
41	def lane7(tile: Tile8) -> F32:
42	  Tile{a,b,c,d,e,f,g,h} = tile
43	  h
44	
45	# One ordered multiply-add per lane: eight output positions, one weight.
46	def update(accumulator: Tile8,input: Tile8,weight: F32) -> Tile8:
47	  Tile{a,b,c,d,e,f,g,h} = accumulator
48	  Tile{x0,x1,x2,x3,x4,x5,x6,x7} = input
49	  +weight = weight
50	  Tile{Fp32.multiply_add(a,x0,weight),Fp32.multiply_add(b,x1,weight),Fp32.multiply_add(c,x2,weight),Fp32.multiply_add(d,x3,weight),
51	    Fp32.multiply_add(e,x4,weight),Fp32.multiply_add(f,x5,weight),Fp32.multiply_add(g,x6,weight),Fp32.multiply_add(h,x7,weight)}
52	
53	def from_sequence(vector: Sequence.Values(F32,8n)) -> Tile8:
54	  Sequence.Elements{a0,Sequence.Elements{a1,Sequence.Elements{a2,Sequence.Elements{a3,Sequence.Elements{a4,Sequence.Elements{a5,Sequence.Elements{a6,
55	          Sequence.Elements{a7,Unit{}}}}}}}}} = vector
56	  Tile{a0,a1,a2,a3,a4,a5,a6,a7}
57	
58	type HeadContext<-Context: Data> is Data:
59	  HeadContext{head: F32,context: Context}
60	
61	def read_with_head(
62	  ~State: Type,~Context: Data,~read: State -> Context -> Nat -> State & F32,
63	  state: State,combined: HeadContext<Context>,index: Nat
64	) -> State & F32:
65	  HeadContext{head,context} = combined
66	  match index:
67	    case 0n: (state,head)
68	    case 1n+rest: read(state,context,1n+rest)
69	
70	def pairs_value(vector: Block.Twin<Block.Twin<Block.Twin<F32>>>) -> Tile8:
71	  Block.Twin{Block.Twin{Block.Twin{a0,a1},Block.Twin{a2,a3}},Block.Twin{Block.Twin{a4,a5},
72	      Block.Twin{a6,a7}}} = vector
73	  Tile{a0,a1,a2,a3,a4,a5,a6,a7}
74	
75	def finish_pairs(~State: Type,result: State & Block.Twin<Block.Twin<Block.Twin<F32>>>) -> State & Tile8:
76	  (state,vector) = result
77	  (state,pairs_value(vector))
78	
79	def read_tail(~State: Type,~Context: Data,~read: State -> Context -> Nat -> State & F32,first: State & F32,context: Context) -> State & Tile8:
80	  (state,head) = first
81	  finish_pairs(~State,
82	    Block.read_pair(~State,~HeadContext<Context>,~Block.Twin<Block.Twin<F32>>,
83	      ~Block.read_pair(~State,~HeadContext<Context>,~Block.Twin<F32>,
84	        ~Block.read_pair(~State,~HeadContext<Context>,~F32,
85	          ~read_with_head(~State,~Context,~read),~1n),~2n),~4n,
86	      state,HeadContext{head,context},0n))
87	
88	# Logical view only; runtime kernels never traverse it.
89	
90	def view(tile: Tile8) -> Sequence.Values(F32,8n):
91	  Tile{a0,a1,a2,a3,a4,a5,a6,a7} = tile
92	  Sequence.Elements{a0,Sequence.Elements{a1,Sequence.Elements{a2,Sequence.Elements{a3,Sequence.Elements{a4,Sequence.Elements{a5,Sequence.Elements{a6,
93	          Sequence.Elements{a7,Unit{}}}}}}}}}
94	
95	def element_at(index: Sequence.Index(8n),tile: Tile8) -> F32:
96	  Sequence.get(F32,8n,index,view(tile))