src/tensor/index_map.bend

bendtorch-core@0.1.0.0 · 98 lines · 8 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/main.bend as Storage
4import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
5import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
6import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
7import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
8import bendtorch-core@0.1.0.0/src/tensor/from_values_contract.bend as FromValues
9import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
10import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise.bend as Pointwise
11import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise_proof.bend as PointwiseProof
12
13# Index-map copy: out[i] = input[map(plan, i)] for every output index i, in
14# increasing i, on a replayed input. It is the pointwise family with a scalar
15# function that keeps its first operand and the same map on both operands, so
16# the read and write histories and the final-value theorem are the pointwise
17# ones. An instance supplies `map` and runs its own loop; these lemmas carry
18# one loop step and the final state.
19
20def keep(operation: Unit, word: F32, other: F32) -> F32:
21  word
22
23def source(~P: Data, ~map: P -> Nat -> Nat, +plan: P, +input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
24  +count: Nat) -> Storage.Storage:
25  PointwiseProof.input_storage(~P, ~map, plan, input, fill, ops, count)
26
27def writes(~P: Data, ~map: P -> Nat -> Nat, +count: Nat, +plan: P, +input: Shape.Shape, +fill: F32,
28  +ops: List<&2, StorageLaws.Op>) -> List<&2, StorageLaws.Op>:
29  Pointwise.writes(~Unit, ~keep, ~P, ~map, ~map, count, Unit{}, plan, input, fill, ops, input, fill, ops)
30
31def target(~P: Data, ~map: P -> Nat -> Nat, +count: Nat, +plan: P, +input: Shape.Shape, +fill: F32,
32  +ops: List<&2, StorageLaws.Op>, +output: Shape.Shape) -> Storage.Storage:
33  Replayed.storage(output, 0.0, writes(~P, ~map, count, plan, input, fill, ops))
34
35def shift(~P: Data, ~map: P -> Nat -> Nat, +plan: P, +input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
36  +output: Shape.Shape, +cursor: Nat, -K: Storage.Storage -> Storage.Storage -> Type,
37  next: K(source(~P, ~map, plan, input, fill, ops, 1n+cursor), target(~P, ~map, 1n+cursor, plan, input, fill, ops, output))) ->
38  K(Replayed.storage(input, fill, Replayed.snoc(List.append(&2, StorageLaws.Op, ops, Pointwise.reads(~P, ~map, plan, cursor)),
39      StorageLaws.Read{U32.from_nat(map(plan, cursor))})),
40    Replayed.storage(output, 0.0, Replayed.snoc(writes(~P, ~map, cursor, plan, input, fill, ops),
41      StorageLaws.Write{U32.from_nat(cursor),
42        Replayed.value(input, fill, List.append(&2, StorageLaws.Op, ops, Pointwise.reads(~P, ~map, plan, cursor)), U32.from_nat(map(plan, cursor)))}))):
43  +read = {StorageLaws.Read{U32.from_nat(map(plan, cursor))} : StorageLaws.Op}
44  +seen = List.append(&2, StorageLaws.Op, ops, Pointwise.reads(~P, ~map, plan, cursor))
45  +word = U32.from_nat(map(plan, cursor))
46  +history = writes(~P, ~map, cursor, plan, input, fill, ops)
47  -written = Replayed.storage(output, 0.0, Replayed.snoc(history, StorageLaws.Write{U32.from_nat(cursor), Replayed.value(input, fill, seen, word)}))
48  %Equal.cong(List<&2, StorageLaws.Op>, Storage.Storage, list => Replayed.storage(input, fill, list),
49    List.append(&2, StorageLaws.Op, ops, Replayed.snoc(Pointwise.reads(~P, ~map, plan, cursor), read)), Replayed.snoc(seen, read),
50    PointwiseProof.append_snoc(ops, Pointwise.reads(~P, ~map, plan, cursor), read)) :
51    K(_, written)
52  %Equal.cong(F32, Storage.Storage, value => Replayed.storage(output, 0.0, Replayed.snoc(history, StorageLaws.Write{U32.from_nat(cursor), value})),
53    Replayed.value(input, fill, ops, word), Replayed.value(input, fill, seen, word),
54    Equal.sym(F32, Replayed.value(input, fill, seen, word), Replayed.value(input, fill, ops, word),
55      PointwiseProof.value_reads(~P, ~map, plan, cursor, input, fill, ops, word))) :
56    K(source(~P, ~map, plan, input, fill, ops, 1n+cursor), _)
57  next
58
59def copied_point(~P: Data, ~map: P -> Nat -> Nat, +plan: P, +input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
60  +output: Shape.Shape,
61  -fits: {Lifecycle.logical_length(target(~P, ~map, Shape.logical_numel(Shape.dimensions(output)), plan, input, fill, ops, output)) ==
62    Shape.logical_numel(Shape.dimensions(output)) : Nat},
63  +index: U32, +inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(output))) == True{} : Bool}) ->
64  {Core.dynamic_read_value(Core.read_dynamic(Core.Packed{output,
65    Core.TTensor{target(~P, ~map, Shape.logical_numel(Shape.dimensions(output)), plan, input, fill, ops, output), fits}}, index)) ==
66    Some{Replayed.value(input, fill, ops, U32.from_nat(map(plan, U32.to_nat(index))))} : Maybe<&2, F32>}:
67  +total = Shape.logical_numel(Shape.dimensions(output))
68  +history = writes(~P, ~map, total, plan, input, fill, ops)
69  Equal.trans(Maybe<&2, F32>,
70    Core.dynamic_read_value(Core.read_dynamic(Core.Packed{output, Core.TTensor{Replayed.storage(output, 0.0, history), fits}}, index)),
71    Some{Replayed.value(output, 0.0, history, index)},
72    Some{Replayed.value(input, fill, ops, U32.from_nat(map(plan, U32.to_nat(index))))},
73    PointwiseProof.replayed_point(output, 0.0, history, fits, index, inside),
74    Equal.cong(F32, Maybe<&2, F32>, word => Some{word}, Replayed.value(output, 0.0, history, index),
75      Replayed.value(input, fill, ops, U32.from_nat(map(plan, U32.to_nat(index)))),
76      PointwiseProof.writes_value(~Unit, ~keep, ~P, ~map, ~map, total, Unit{}, plan, input, fill, ops, input, fill, ops, output, index,
77        inside, Conversion.le_reflexive(total))))
78
79def kept(~P: Data, ~map: P -> Nat -> Nat, +plan: P, +count: Nat, +input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
80  -fits: {Lifecycle.logical_length(source(~P, ~map, plan, input, fill, ops, count)) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
81  +other: U32, inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool}) ->
82  {Core.read_value(input, Core.read_flat(input, Core.TTensor{source(~P, ~map, plan, input, fill, ops, count), fits}, other)) ==
83    Some{Replayed.value(input, fill, ops, other)} : Maybe<&2, F32>}:
84  +seen = List.append(&2, StorageLaws.Op, ops, Pointwise.reads(~P, ~map, plan, count))
85  Equal.trans(Maybe<&2, F32>, Core.read_value(input, Core.read_flat(input, Core.TTensor{Replayed.storage(input, fill, seen), fits}, other)),
86    Some{Replayed.value(input, fill, seen, other)}, Some{Replayed.value(input, fill, ops, other)},
87    Replayed.read_found(input, fill, seen, fits, other, inside),
88    Equal.cong(F32, Maybe<&2, F32>, word => Some{word}, Replayed.value(input, fill, seen, other), Replayed.value(input, fill, ops, other),
89      PointwiseProof.value_reads(~P, ~map, plan, count, input, fill, ops, other)))
90
91def stored_at(+shape: Shape.Shape, +values: List<&2, F32>,
92  +counted: {List.length(&2, F32, values) == Shape.logical_numel(Shape.dimensions(shape)) : Nat},
93  +index: U32, inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}) ->
94  {Some{Replayed.value(shape, 0.0, FromValues.writes(values, 0n), index)} == List.get(&2, F32, values, U32.to_nat(index)) : Maybe<&2, F32>}:
95  FromValues.written_at(values, 0n, Shape.count(shape), index, 0.0, U32.to_nat(index), FromValues.start_room(shape, values, counted), {==},
96    %Equal.sym(Nat, List.length(&2, F32, values), Shape.logical_numel(Shape.dimensions(shape)), counted) :
97      {Nat.is_lt(U32.to_nat(index), _) == True{} : Bool}
98    inside)