src/tensor/contracts/broadcast_proof.bend
bendtorch-core@0.1.0.0 · 106 lines · 6 definitions
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/shape/flat.bend as Flat
4import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
5import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
6import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
7import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
8import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
9import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
10import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
11import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise.bend as Pointwise
12import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise_proof.bend as Proof
13import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise_records.bend as Records
14import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast.bend as Broadcast
15
16def output_inside(+plan: Shape.BroadcastPlan, +index: Nat,
17 below: {Nat.is_lt(index, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) ->
18 {True{} == Nat.is_lt(index, Shape.elements(Shape.output_shape(plan))) : Bool}:
19 %Equal.sym(Nat, Shape.elements(Shape.output_shape(plan)), Shape.numel(Shape.extents(Shape.output_shape(plan))),
20 Shape.elements_exact(Shape.output_shape(plan))) : {True{} == Nat.is_lt(index, _) : Bool}
21 Equal.sym(Bool, Nat.is_lt(index, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))), True{}, below)
22
23def left_walk_bound(+plan: Shape.BroadcastPlan,
24 map: Shape.Index<Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))>,
25 +index: Nat, below: {Nat.is_lt(index, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) ->
26 {Nat.is_lt(Broadcast.left_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), map, index),
27 Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan)))) == True{} : Bool}:
28 match map:
29 case Shape.Indexed{rows, exact}:
30 Equal.sym(Bool, True{}, Nat.is_lt(Shape.walk_left(rows, index), Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan)))),
31 %Shape.elements_exact(Shape.left_shape(plan)) : {True{} == Nat.is_lt(Shape.walk_left(rows, index), _) : Bool}
32 Flat.left_bound(plan, rows, exact, index, output_inside(plan, index, below)))
33 case Shape.Unindexed{zero}:
34 Empty.absurd({Nat.is_lt(0n, Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan)))) == True{} : Bool},
35 Arith.lt_zero(index,
36 %Equal.sym(Nat, 0n, Shape.numel(Shape.extents(Shape.output_shape(plan))), zero) : {Nat.is_lt(index, _) == True{} : Bool}
37 below))
38
39def right_walk_bound(+plan: Shape.BroadcastPlan,
40 map: Shape.Index<Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))>,
41 +index: Nat, below: {Nat.is_lt(index, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) ->
42 {Nat.is_lt(Broadcast.right_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), map, index),
43 Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))) == True{} : Bool}:
44 match map:
45 case Shape.Indexed{rows, exact}:
46 Equal.sym(Bool, True{}, Nat.is_lt(Shape.walk_right(rows, index), Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))),
47 %Shape.elements_exact(Shape.right_shape(plan)) : {True{} == Nat.is_lt(Shape.walk_right(rows, index), _) : Bool}
48 Flat.right_bound(plan, rows, exact, index, output_inside(plan, index, below)))
49 case Shape.Unindexed{zero}:
50 Empty.absurd({Nat.is_lt(0n, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))) == True{} : Bool},
51 Arith.lt_zero(index,
52 %Equal.sym(Nat, 0n, Shape.numel(Shape.extents(Shape.output_shape(plan))), zero) : {Nat.is_lt(index, _) == True{} : Bool}
53 below))
54
55def left_inside(+plan: Shape.BroadcastPlan, +count: Nat,
56 +room: {Nat.is_le(count, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) ->
57 {Pointwise.inside(~Shape.BroadcastPlan, ~Broadcast.left_map, plan, count, Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan)))) == True{} : Bool}:
58 match count:
59 case 0n:
60 {==}
61 case 1n+previous:
62 +total = Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))
63 +bound = Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan)))
64 %Equal.sym(Bool, Pointwise.inside(~Shape.BroadcastPlan, ~Broadcast.left_map, plan, previous, bound), True{},
65 left_inside(plan, previous, Conversion.le_predecessor(previous, total, room))) :
66 {Bool.and(_, Nat.is_lt(Broadcast.left_map(plan, previous), bound)) == True{} : Bool}
67 %Equal.sym(Bool, Nat.is_lt(Broadcast.left_map(plan, previous), bound), True{},
68 left_walk_bound(plan, Shape.index_map(plan), previous, Conversion.le_successor_lt(previous, total, room))) :
69 {Bool.and(True{}, _) == True{} : Bool}
70 {==}
71
72def right_inside(+plan: Shape.BroadcastPlan, +count: Nat,
73 +room: {Nat.is_le(count, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) ->
74 {Pointwise.inside(~Shape.BroadcastPlan, ~Broadcast.right_map, plan, count, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))) == True{} : Bool}:
75 match count:
76 case 0n:
77 {==}
78 case 1n+previous:
79 +total = Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))
80 +bound = Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))
81 %Equal.sym(Bool, Pointwise.inside(~Shape.BroadcastPlan, ~Broadcast.right_map, plan, previous, bound), True{},
82 right_inside(plan, previous, Conversion.le_predecessor(previous, total, room))) :
83 {Bool.and(_, Nat.is_lt(Broadcast.right_map(plan, previous), bound)) == True{} : Bool}
84 %Equal.sym(Bool, Nat.is_lt(Broadcast.right_map(plan, previous), bound), True{},
85 right_walk_bound(plan, Shape.index_map(plan), previous, Conversion.le_successor_lt(previous, total, room))) :
86 {Bool.and(True{}, _) == True{} : Bool}
87 {==}
88
89def broadcast_pointwise(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan,
90 +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
91 -left_fits: {Lifecycle.logical_length(Replayed.storage(Shape.left_shape(plan), left_fill, left_ops)) ==
92 Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan))) : Nat},
93 +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
94 -right_fits: {Lifecycle.logical_length(Replayed.storage(Shape.right_shape(plan), right_fill, right_ops)) ==
95 Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}) ->
96 Records.PointwiseContract<Shape.output_shape(plan), Shape.left_shape(plan), Shape.right_shape(plan),
97 index => Broadcast.cell(operation, plan, left_fill, left_ops, right_fill, right_ops, index),
98 other => Replayed.value(Shape.left_shape(plan), left_fill, left_ops, other),
99 other => Replayed.value(Shape.right_shape(plan), right_fill, right_ops, other),
100 Broadcast.binary(operation, plan,
101 Core.Packed{Shape.left_shape(plan), Core.TTensor{Replayed.storage(Shape.left_shape(plan), left_fill, left_ops), left_fits}},
102 Core.Packed{Shape.right_shape(plan), Core.TTensor{Replayed.storage(Shape.right_shape(plan), right_fill, right_ops), right_fits}})>:
103 +total = Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))
104 Proof.contract(~Elementwise.Operation, ~Elementwise.scalar, ~Shape.BroadcastPlan, ~Broadcast.left_map, ~Broadcast.right_map,
105 operation, plan, Shape.left_shape(plan), left_fill, left_ops, Shape.right_shape(plan), right_fill, right_ops, Shape.output_shape(plan),
106 left_inside(plan, total, Conversion.le_reflexive(total)), right_inside(plan, total, Conversion.le_reflexive(total)), left_fits, right_fits)