src/tensor/contracts/pointwise.bend
bendtorch-core@0.1.0.0 · 138 lines · 17 definitions
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
4import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
5import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
6
7
8
9
10
11type State<-shape: Shape.Shape> is Type:
12 State{left: Core.Dynamic, right: Core.Dynamic, output: Core.Tensor<shape>}
13
14def 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
25def 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
36def 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
47def 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
55def 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
66def 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
76def 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
81def 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
88def 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
94def 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
101def 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
112def 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
120def 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
128def 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
132def 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
136def 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))