1	import Base
2	import ../../shape/core.bend as Shape
3	import ../../storage/LAWS.bend as StorageLaws
4	import ../core.bend as Core
5	import ./replayed.bend as Replayed
6	
7	# Pointwise family: out[i] = f(operation, left[ia(plan, i)], right[ib(plan, i)])
8	# for every flat output index i, written in increasing i. `f` is the scalar
9	# function and `ia`, `ib` the index maps.
10	
11	type State<-shape: Shape.Shape> is Type:
12	  State{left: Core.Dynamic, right: Core.Dynamic, output: Core.Tensor<shape>}
13	
14	def write_step(-shape: Shape.Shape,
15	  rec: State<shape> -> (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>,
16	  left: Core.Dynamic, right: Core.Dynamic, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) ->
17	  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
18	  (output, result) = pair
19	  match result:
20	    case Fail{error}:
21	      ((left, right), Fail{error})
22	    case Done{unit}:
23	      rec(State{left, right, output})
24	
25	def right_step(~O: Data, ~f: O -> F32 -> F32 -> F32, +operation: O, -shape: Shape.Shape, +index: U32,
26	  rec: State<shape> -> (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>,
27	  left: Core.Dynamic, output: Core.Tensor<shape>, x: F32, pair: Core.Dynamic & Result<&2, &2, Core.Error, F32>) ->
28	  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
29	  (right, result) = pair
30	  match result:
31	    case Fail{error}:
32	      ((left, right), Fail{error})
33	    case Done{y}:
34	      write_step(shape, rec, left, right, Core.write_flat(shape, output, index, f(operation, x, y)))
35	
36	def left_step(~O: Data, ~f: O -> F32 -> F32 -> F32, +operation: O, -shape: Shape.Shape, +index: U32, +right_index: U32,
37	  rec: State<shape> -> (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>,
38	  right: Core.Dynamic, output: Core.Tensor<shape>, pair: Core.Dynamic & Result<&2, &2, Core.Error, F32>) ->
39	  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
40	  (left, result) = pair
41	  match result:
42	    case Fail{error}:
43	      ((left, right), Fail{error})
44	    case Done{x}:
45	      right_step(~O, ~f, operation, shape, index, rec, left, output, x, Core.read_dynamic(right, right_index))
46	
47	def step(~O: Data, ~f: O -> F32 -> F32 -> F32, ~P: Data, ~ia: P -> Nat -> Nat, ~ib: P -> Nat -> Nat,
48	  +operation: O, +plan: P, -shape: Shape.Shape, +cursor: Nat,
49	  rec: State<shape> -> (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>,
50	  state: State<shape>) -> (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
51	  State{left, right, output} = state
52	  left_step(~O, ~f, operation, shape, U32.from_nat(cursor), U32.from_nat(ib(plan, cursor)), rec, right, output,
53	    Core.read_dynamic(left, U32.from_nat(ia(plan, cursor))))
54	
55	def loop(~O: Data, ~f: O -> F32 -> F32 -> F32, ~P: Data, ~ia: P -> Nat -> Nat, ~ib: P -> Nat -> Nat,
56	  fuel: Nat, +operation: O, +plan: P, +shape: Shape.Shape, +cursor: Nat, state: State<shape>) ->
57	  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
58	  match fuel:
59	    case 0n:
60	      State{left, right, output} = state
61	      ((left, right), Done{Core.pack(shape, output)})
62	    case 1n+rest:
63	      step(~O, ~f, ~P, ~ia, ~ib, operation, plan, shape, cursor,
64	        state => loop(~O, ~f, ~P, ~ia, ~ib, rest, operation, plan, shape, 1n+cursor, state), state)
65	
66	def allocated(~O: Data, ~f: O -> F32 -> F32 -> F32, ~P: Data, ~ia: P -> Nat -> Nat, ~ib: P -> Nat -> Nat,
67	  +operation: O, +plan: P, +shape: Shape.Shape, left: Core.Dynamic, right: Core.Dynamic,
68	  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) ->
69	  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
70	  match result:
71	    case Fail{error}:
72	      ((left, right), Fail{error})
73	    case Done{output}:
74	      loop(~O, ~f, ~P, ~ia, ~ib, Shape.elements(shape), operation, plan, shape, 0n, State{left, right, output})
75	
76	def run(~O: Data, ~f: O -> F32 -> F32 -> F32, ~P: Data, ~ia: P -> Nat -> Nat, ~ib: P -> Nat -> Nat,
77	  +operation: O, +plan: P, +shape: Shape.Shape, left: Core.Dynamic, right: Core.Dynamic) ->
78	  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
79	  allocated(~O, ~f, ~P, ~ia, ~ib, operation, plan, shape, left, right, Core.new(shape, 0.0))
80	
81	def inside(~P: Data, ~map: P -> Nat -> Nat, +plan: P, +count: Nat, +bound: Nat) -> Bool:
82	  match count:
83	    case 0n:
84	      True{}
85	    case 1n+previous:
86	      Bool.and(inside(~P, ~map, plan, previous, bound), Nat.is_lt(map(plan, previous), bound))
87	
88	def cell(~O: Data, ~f: O -> F32 -> F32 -> F32, ~P: Data, ~ia: P -> Nat -> Nat, ~ib: P -> Nat -> Nat,
89	  +operation: O, +plan: P, +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
90	  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +index: Nat) -> F32:
91	  f(operation, Replayed.value(left, left_fill, left_ops, U32.from_nat(ia(plan, index))),
92	    Replayed.value(right, right_fill, right_ops, U32.from_nat(ib(plan, index))))
93	
94	def reads(~P: Data, ~map: P -> Nat -> Nat, +plan: P, +count: Nat) -> List<&2, StorageLaws.Op>:
95	  match count:
96	    case 0n:
97	      Nil{}
98	    case 1n+previous:
99	      Replayed.snoc(reads(~P, ~map, plan, previous), StorageLaws.Read{U32.from_nat(map(plan, previous))})
100	
101	def writes(~O: Data, ~f: O -> F32 -> F32 -> F32, ~P: Data, ~ia: P -> Nat -> Nat, ~ib: P -> Nat -> Nat,
102	  +count: Nat, +operation: O, +plan: P, +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
103	  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>) -> List<&2, StorageLaws.Op>:
104	  match count:
105	    case 0n:
106	      Nil{}
107	    case 1n+previous:
108	      Replayed.snoc(writes(~O, ~f, ~P, ~ia, ~ib, previous, operation, plan, left, left_fill, left_ops, right, right_fill, right_ops),
109	        StorageLaws.Write{U32.from_nat(previous),
110	          cell(~O, ~f, ~P, ~ia, ~ib, operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, previous)})
111	
112	def output_point(result: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>, +index: U32) -> Maybe<&2, F32>:
113	  (inputs, outcome) = result
114	  match outcome:
115	    case Fail{error}:
116	      None{}
117	    case Done{output}:
118	      Core.dynamic_read_value(Core.read_dynamic(output, index))
119	
120	def output_dimensions(result: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) -> Maybe<&2, List<&2, U32>>:
121	  (inputs, outcome) = result
122	  match outcome:
123	    case Fail{error}:
124	      None{}
125	    case Done{Core.Packed{shape, tensor}}:
126	      Some{Shape.dimensions(shape)}
127	
128	def succeeded(result: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) -> Bool:
129	  (inputs, outcome) = result
130	  Result.is_done(&2, &1, Core.Error, Core.Dynamic, outcome)
131	
132	def left_point(result: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>, +index: U32) -> Maybe<&2, F32>:
133	  ((left, right), outcome) = result
134	  Core.dynamic_read_value(Core.read_dynamic(left, index))
135	
136	def right_point(result: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>, +index: U32) -> Maybe<&2, F32>:
137	  ((left, right), outcome) = result
138	  Core.dynamic_read_value(Core.read_dynamic(right, index))