1	# Shape evidence for Lanes.map: every round reads the values array, writes a
2	# scratch tile, reads it back and stores eight values, so the array keeps its
3	# depth. Numerical loops call none of these functions.
4	import Base
5	import ../traversal_lanes.bend as Lanes
6	import ../traversal_cells.bend as Cells
7	import ../kernel_tile.bend as Tile
8	import ./storage_witness.bend as Witness
9	import ./traversal_witness.bend as TraversalWitness
10	
11	type Gathered<-values: Array<F32>,-result: Array<F32> & Tile.Tile8> is Data:
12	  Gathered{-cell: Tile.Tile8,equation: {result == (values,cell) : Array<F32> & Tile.Tile8}}
13	
14	def gathered(~Context: Data,~map: Context -> F32 -> F32,-values: Array<F32>,+position: U32,+context: Context,
15	  r0: Witness.Read<F32,values,Array.get(F32,values,position)>,
16	  r1: Witness.Read<F32,values,Array.get(F32,values,(position + 1 : U32))>,
17	  r2: Witness.Read<F32,values,Array.get(F32,values,((position + 1 : U32) + 1 : U32))>,
18	  r3: Witness.Read<F32,values,Array.get(F32,values,(((position + 1 : U32) + 1 : U32) + 1 : U32))>,
19	  r4: Witness.Read<F32,values,Array.get(F32,values,((((position + 1 : U32) + 1 : U32) + 1 : U32) + 1 : U32))>,
20	  r5: Witness.Read<F32,values,Array.get(F32,values,(((((position + 1 : U32) + 1 : U32) + 1 : U32) + 1 : U32) + 1 : U32))>,
21	  r6: Witness.Read<F32,values,Array.get(F32,values,((((((position + 1 : U32) + 1 : U32) + 1 : U32) + 1 : U32) + 1 : U32) + 1 : U32))>,
22	  r7: Witness.Read<F32,values,Array.get(F32,values,(((((((position + 1 : U32) + 1 : U32) + 1 : U32) + 1 : U32) + 1 : U32) + 1 : U32) + 1 : U32))>) ->
23	  Gathered<values,Lanes.gather(~Context,~map,values,position,context)>:
24	  match r0 r1 r2 r3 r4 r5 r6 r7:
25	    case Witness.Read{a,ea} Witness.Read{b,eb} Witness.Read{c,ec} Witness.Read{d,ed} Witness.Read{e,ee} Witness.Read{f,ef} Witness.Read{g,eg} Witness.Read{h,eh}:
26	      +p1 = (position + 1 : U32)
27	      +p2 = (p1 + 1 : U32)
28	      +p3 = (p2 + 1 : U32)
29	      +p4 = (p3 + 1 : U32)
30	      +p5 = (p4 + 1 : U32)
31	      +p6 = (p5 + 1 : U32)
32	      +p7 = (p6 + 1 : U32)
33	      %Equal.sym(Array<F32> & F32,Array.get(F32,values,position),(values,a),ea) : Gathered<values,Lanes.gather1(~Context,~map,position,context,_)>
34	      %Equal.sym(Array<F32> & F32,Array.get(F32,values,p1),(values,b),eb) : Gathered<values,Lanes.gather2(~Context,~map,p1,context,a,_)>
35	      %Equal.sym(Array<F32> & F32,Array.get(F32,values,p2),(values,c),ec) : Gathered<values,Lanes.gather3(~Context,~map,p2,context,a,b,_)>
36	      %Equal.sym(Array<F32> & F32,Array.get(F32,values,p3),(values,d),ed) : Gathered<values,Lanes.gather4(~Context,~map,p3,context,a,b,c,_)>
37	      %Equal.sym(Array<F32> & F32,Array.get(F32,values,p4),(values,e),ee) : Gathered<values,Lanes.gather5(~Context,~map,p4,context,a,b,c,d,_)>
38	      %Equal.sym(Array<F32> & F32,Array.get(F32,values,p5),(values,f),ef) : Gathered<values,Lanes.gather6(~Context,~map,p5,context,a,b,c,d,e,_)>
39	      %Equal.sym(Array<F32> & F32,Array.get(F32,values,p6),(values,g),eg) : Gathered<values,Lanes.gather7(~Context,~map,p6,context,a,b,c,d,e,f,_)>
40	      %Equal.sym(Array<F32> & F32,Array.get(F32,values,p7),(values,h),eh) : Gathered<values,Lanes.gathered(~Context,~map,context,a,b,c,d,e,f,g,_)>
41	      Gathered{Tile.Tile{map(context,a),map(context,b),map(context,c),map(context,d),map(context,e),map(context,f),map(context,g),map(context,h)},{==}}
42	
43	def gather(~Context: Data,~map: Context -> F32 -> F32,+depth: Nat,-values: Array<F32>,+position: U32,+context: Context,
44	  +witness: Witness.Witness(F32,depth,values)) -> Gathered<values,Lanes.gather(~Context,~map,values,position,context)>:
45	  +p1 = (position + 1 : U32)
46	  +p2 = (p1 + 1 : U32)
47	  +p3 = (p2 + 1 : U32)
48	  +p4 = (p3 + 1 : U32)
49	  +p5 = (p4 + 1 : U32)
50	  +p6 = (p5 + 1 : U32)
51	  +p7 = (p6 + 1 : U32)
52	  gathered(~Context,~map,values,position,context,Witness.read(F32,depth,values,position,witness),Witness.read(F32,depth,values,p1,witness),
53	    Witness.read(F32,depth,values,p2,witness),Witness.read(F32,depth,values,p3,witness),Witness.read(F32,depth,values,p4,witness),
54	    Witness.read(F32,depth,values,p5,witness),Witness.read(F32,depth,values,p6,witness),Witness.read(F32,depth,values,p7,witness))
55	
56	def scatter(+depth: Nat,-values: Array<F32>,+position: U32,-tile: Tile.Tile8,witness: Witness.Witness(F32,depth,values)) ->
57	  Witness.Witness(F32,depth,Lanes.scatter(values,position,tile)):
58	  +p1 = (position + 1 : U32)
59	  +p2 = (p1 + 1 : U32)
60	  +p3 = (p2 + 1 : U32)
61	  +p4 = (p3 + 1 : U32)
62	  +p5 = (p4 + 1 : U32)
63	  +p6 = (p5 + 1 : U32)
64	  +p7 = (p6 + 1 : U32)
65	  -o1 = Array.set(F32,values,position,Tile.lane0(tile))
66	  -o2 = Array.set(F32,o1,p1,Tile.lane1(tile))
67	  -o3 = Array.set(F32,o2,p2,Tile.lane2(tile))
68	  -o4 = Array.set(F32,o3,p3,Tile.lane3(tile))
69	  -o5 = Array.set(F32,o4,p4,Tile.lane4(tile))
70	  -o6 = Array.set(F32,o5,p5,Tile.lane5(tile))
71	  -o7 = Array.set(F32,o6,p6,Tile.lane6(tile))
72	  witness = Witness.write(F32,depth,values,position,Tile.lane0(tile),witness)
73	  witness = Witness.write(F32,depth,o1,p1,Tile.lane1(tile),witness)
74	  witness = Witness.write(F32,depth,o2,p2,Tile.lane2(tile),witness)
75	  witness = Witness.write(F32,depth,o3,p3,Tile.lane3(tile),witness)
76	  witness = Witness.write(F32,depth,o4,p4,Tile.lane4(tile),witness)
77	  witness = Witness.write(F32,depth,o5,p5,Tile.lane5(tile),witness)
78	  witness = Witness.write(F32,depth,o6,p6,Tile.lane6(tile),witness)
79	  Witness.write(F32,depth,o7,p7,Tile.lane7(tile),witness)
80	
81	type Evidence<-Context: Data,-depth: Nat,-state: Cells.State<Array<F32>,Tile.Tile8,Context>> is Data:
82	  Evidence{-values: Array<F32>,-scratch: Array<Tile.Tile8>,position: U32,slot: U32,context: Context,
83	    equation: {state == Cells.State{values,scratch,position,slot,context} : Cells.State<Array<F32>,Tile.Tile8,Context>},
84	    values_witness: Witness.Witness(F32,depth,values),scratch_witness: Witness.Witness(Tile.Tile8,6n,scratch)}
85	
86	def scattered(~Context: Data,+depth: Nat,-values: Array<F32>,-scratch: Array<Tile.Tile8>,+position: U32,+slot: U32,+context: Context,
87	  values_witness: Witness.Witness(F32,depth,values),scratch_witness: Witness.Witness(Tile.Tile8,6n,scratch),
88	  read: Witness.Read<Tile.Tile8,scratch,Array.get(Tile.Tile8,scratch,slot)>) ->
89	  Evidence<Context,depth,Cells.scattered(~Array<F32>,~Tile.Tile8,~Context,~Lanes.scatter,~8,position,slot,context,values,Array.get(Tile.Tile8,scratch,slot))>:
90	  match read:
91	    case Witness.Read{tile,equation}:
92	      %Equal.sym(Array<Tile.Tile8> & Tile.Tile8,Array.get(Tile.Tile8,scratch,slot),(scratch,tile),equation) :
93	        Evidence<Context,depth,Cells.scattered(~Array<F32>,~Tile.Tile8,~Context,~Lanes.scatter,~8,position,slot,context,values,_)>
94	      Evidence{Lanes.scatter(values,position,tile),scratch,(position + 8 : U32),U32.and((slot + 1 : U32),63),context,{==},
95	        scatter(depth,values,position,tile,values_witness),scratch_witness}
96	
97	def stored(~Context: Data,~collect: Array<F32> -> U32 -> Context -> Array<F32> & Tile.Tile8,+depth: Nat,-values: Array<F32>,-scratch: Array<Tile.Tile8>,+position: U32,+slot: U32,+context: Context,
98	  values_witness: Witness.Witness(F32,depth,values),scratch_witness: Witness.Witness(Tile.Tile8,6n,scratch),
99	  gathered: Gathered<values,collect(values,position,context)>) ->
100	  Evidence<Context,depth,Cells.stored(~Array<F32>,~Tile.Tile8,~Context,~Lanes.scatter,~8,position,slot,context,scratch,collect(values,position,context))>:
101	  match gathered:
102	    case Gathered{cell,equation}:
103	      +written = Witness.write(Tile.Tile8,6n,scratch,slot,cell,scratch_witness)
104	      %Equal.sym(Array<F32> & Tile.Tile8,collect(values,position,context),(values,cell),equation) :
105	        Evidence<Context,depth,Cells.stored(~Array<F32>,~Tile.Tile8,~Context,~Lanes.scatter,~8,position,slot,context,scratch,_)>
106	      scattered(~Context,depth,values,Array.set(Tile.Tile8,scratch,slot,cell),position,slot,context,values_witness,written,
107	        Witness.read(Tile.Tile8,6n,Array.set(Tile.Tile8,scratch,slot,cell),slot,written))
108	
109	def step(~Context: Data,~collect: Array<F32> -> U32 -> Context -> Array<F32> & Tile.Tile8,
110	  ~collected: @+depth: Nat -> @-values: Array<F32> -> @+position: U32 -> @+context: Context -> @+witness: Witness.Witness(F32,depth,values) -> Gathered<values,collect(values,position,context)>,+depth: Nat,-state: Cells.State<Array<F32>,Tile.Tile8,Context>,evidence: Evidence<Context,depth,state>) ->
111	  Evidence<Context,depth,Cells.step(~Array<F32>,~Tile.Tile8,~Context,~collect,~Lanes.scatter,~8,state)>:
112	  match evidence:
113	    case Evidence{values,scratch,+position,+slot,+context,equation,+values_witness,scratch_witness}:
114	      %Equal.sym(Cells.State<Array<F32>,Tile.Tile8,Context>,state,Cells.State{values,scratch,position,slot,context},equation) :
115	        Evidence<Context,depth,Cells.step(~Array<F32>,~Tile.Tile8,~Context,~collect,~Lanes.scatter,~8,_)>
116	      stored(~Context,~collect,depth,values,scratch,position,slot,context,values_witness,scratch_witness,
117	        collected(depth,values,position,context,values_witness))
118	
119	def finish(~Context: Data,~map: Context -> F32 -> F32,+depth: Nat,count: Nat,-state: Cells.State<Array<F32>,Tile.Tile8,Context>,evidence: Evidence<Context,depth,state>) ->
120	  Witness.Witness(F32,depth,Lanes.finish(~Context,~map,count,state)):
121	  match evidence:
122	    case Evidence{values,scratch,position,slot,context,equation,values_witness,scratch_witness}:
123	      %Equal.sym(Cells.State<Array<F32>,Tile.Tile8,Context>,state,Cells.State{values,scratch,position,slot,context},equation) :
124	        Witness.Witness(F32,depth,Lanes.finish(~Context,~map,count,_))
125	      TraversalWitness.map_in_place(~F32,~Context,~map,depth,count,position,context,values,values_witness)
126	
127	def rounds_with(~Context: Data,~map: Context -> F32 -> F32,~collect: Array<F32> -> U32 -> Context -> Array<F32> & Tile.Tile8,
128	  ~collected: @+depth: Nat -> @-values: Array<F32> -> @+position: U32 -> @+context: Context -> @+witness: Witness.Witness(F32,depth,values) -> Gathered<values,collect(values,position,context)>,+depth: Nat,count: Nat,-state: Cells.State<Array<F32>,Tile.Tile8,Context>,evidence: Evidence<Context,depth,state>) ->
129	  Witness.Witness(F32,depth,Lanes.rounds_with(~Array<F32>,~Context,~collect,~Lanes.scatter,~Lanes.finish(~Context,~map),count,state)):
130	  match count:
131	    case 0n: finish(~Context,~map,depth,0n,state,evidence)
132	    case 1n: finish(~Context,~map,depth,1n,state,evidence)
133	    case 2n: finish(~Context,~map,depth,2n,state,evidence)
134	    case 3n: finish(~Context,~map,depth,3n,state,evidence)
135	    case 4n: finish(~Context,~map,depth,4n,state,evidence)
136	    case 5n: finish(~Context,~map,depth,5n,state,evidence)
137	    case 6n: finish(~Context,~map,depth,6n,state,evidence)
138	    case 7n: finish(~Context,~map,depth,7n,state,evidence)
139	    case 1n+1n+1n+1n+1n+1n+1n+1n+rest:
140	      rounds_with(~Context,~map,~collect,~collected,depth,rest,Cells.step(~Array<F32>,~Tile.Tile8,~Context,~collect,~Lanes.scatter,~8,state),
141	        step(~Context,~collect,~collected,depth,state,evidence))
142	
143	def rounds(~Context: Data,~map: Context -> F32 -> F32,+depth: Nat,count: Nat,-state: Cells.State<Array<F32>,Tile.Tile8,Context>,evidence: Evidence<Context,depth,state>) ->
144	  Witness.Witness(F32,depth,Lanes.rounds(~Context,~map,count,state)):
145	  rounds_with(~Context,~map,~Lanes.gather(~Context,~map),~gather(~Context,~map),depth,count,state,evidence)
146	
147	def map(~Context: Data,~map: Context -> F32 -> F32,+depth: Nat,count: Nat,+position: U32,+context: Context,-values: Array<F32>,witness: Witness.Witness(F32,depth,values)) ->
148	  Witness.Witness(F32,depth,Lanes.map(~Context,~map,count,position,context,values)):
149	  rounds(~Context,~map,depth,count,Cells.State{values,Lanes.scratch(),position,0,context},
150	    Evidence{values,Lanes.scratch(),position,0,context,{==},witness,Witness.fresh(Tile.Tile8,6n,Tile.filled(0.0))})
151	
152	def map_when(~Context: Data,~operation: Context -> F32 -> F32,covered: Bool,+depth: Nat,count: Nat,+position: U32,+context: Context,-values: Array<F32>,
153	  witness: Witness.Witness(F32,depth,values)) -> Witness.Witness(F32,depth,Lanes.map_when(~Context,~operation,covered,count,position,context,values)):
154	  match covered:
155	    case True{}: map(~Context,~operation,depth,count,position,context,values,witness)
156	    case False{}: TraversalWitness.map_in_place(~F32,~Context,~operation,depth,count,position,context,values,witness)