1 import Base 2 import ../../shape/core.bend as Shape 3 import ../../storage/main.bend as Storage 4 import ../../storage/LAWS.bend as StorageLaws 5 import ../../storage/arith.bend as Arith 6 import ../core.bend as Core 7 import ../lifecycle.bend as Lifecycle 8 import ../word_conversion.bend as Conversion 9 import ../read_contract.bend as Reads 10 import ../from_values_contract.bend as FromValues 11 12 # A replayed tensor holds the storage that `Storage.new(count, fill)` creates, 13 # changed by a list of storage operations in order. `Core.new(shape, fill)` is 14 # the empty history, `Core.from_values` writes each value once, and every 15 # contract-family output is a write history over a `+0.0` fill. 16 def storage(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>) -> Storage.Storage: 17 StorageLaws.replay(ops, StorageLaws.created(Shape.count(shape), fill)) 18 19 def value(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +index: U32) -> F32: 20 StorageLaws.expected(ops, U32.to_nat(Shape.count(shape)), index, fill) 21 22 def snoc(+ops: List<&2, StorageLaws.Op>, +op: StorageLaws.Op) -> List<&2, StorageLaws.Op>: 23 List.append(&2, StorageLaws.Op, ops, [op]) 24 25 def stored(-shape: Shape.Shape, tensor: Core.Tensor) -> Storage.Storage: 26 Core.TTensor{storage, valid} = tensor 27 storage 28 29 def read_stored(-shape: Shape.Shape, pair: Core.Tensor & Result<&2, &2, Core.Error, F32>) -> Storage.Storage: 30 (tensor, result) = pair 31 stored(shape, tensor) 32 33 def write_stored(-shape: Shape.Shape, pair: Core.Tensor & Result<&2, &2, Core.Error, Unit>) -> Storage.Storage: 34 (tensor, result) = pair 35 stored(shape, tensor) 36 37 def write_status(-shape: Shape.Shape, pair: Core.Tensor & Result<&2, &2, Core.Error, Unit>) -> Bool: 38 (tensor, result) = pair 39 Result.is_done(&2, &2, Core.Error, Unit, result) 40 41 def some_value(+word: F32, +other: F32, same: {Some{word} == Some{other} : Maybe<&2, F32>}) -> {word == other : F32}: 42 Equal.cong(Maybe<&2, F32>, F32, found => Maybe.default(&2, F32, found, word), Some{word}, Some{other}, same) 43 44 def fits(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>) -> 45 {Lifecycle.logical_length(storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}: 46 Equal.trans(Nat, Lifecycle.logical_length(storage(shape, fill, ops)), StorageLaws.size(storage(shape, fill, ops)), 47 Shape.logical_numel(Shape.dimensions(shape)), 48 Equal.sym(Nat, StorageLaws.size(storage(shape, fill, ops)), Lifecycle.logical_length(storage(shape, fill, ops)), 49 Reads.size_is_length(storage(shape, fill, ops))), 50 Equal.trans(Nat, StorageLaws.size(storage(shape, fill, ops)), U32.to_nat(Shape.count(shape)), 51 Shape.logical_numel(Shape.dimensions(shape)), 52 StorageLaws.replay_size(Shape.count(shape), fill, ops, Conversion.shape_count_within(shape)), Reads.count_exact(shape))) 53 54 def replay_snoc(+ops: List<&2, StorageLaws.Op>, +op: StorageLaws.Op, +shape: Shape.Shape, +fill: F32) -> 55 {storage(shape, fill, snoc(ops, op)) == StorageLaws.step(storage(shape, fill, ops), op) : Storage.Storage}: 56 FromValues.replay_snoc(ops, op, StorageLaws.created(Shape.count(shape), fill)) 57 58 def read_owner_projection(-shape: Shape.Shape, -original: Storage.Storage, 59 -valid: {Lifecycle.logical_length(original) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}, +index: U32, 60 pair: Storage.Storage & Result<&2, &2, Storage.Error, F32>, 61 -same: {pair == Storage.get(original, index) : Storage.Storage & Result<&2, &2, Storage.Error, F32>}) -> 62 {read_stored(shape, Core.read_result(shape, original, valid, index, pair, same)) == 63 StorageLaws.owner(Result<&2, &2, Storage.Error, F32>, pair) : Storage.Storage}: 64 (storage, result) = pair 65 {==} 66 67 def read_moved(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, 68 -valid: {Lifecycle.logical_length(storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}, +index: U32) -> 69 {read_stored(shape, Core.read_flat(shape, Core.TTensor{storage(shape, fill, ops), valid}, index)) == 70 storage(shape, fill, snoc(ops, StorageLaws.Read{index})) : Storage.Storage}: 71 Equal.trans(Storage.Storage, read_stored(shape, Core.read_flat(shape, Core.TTensor{storage(shape, fill, ops), valid}, index)), 72 StorageLaws.kept(storage(shape, fill, ops), index), storage(shape, fill, snoc(ops, StorageLaws.Read{index})), 73 read_owner_projection(shape, storage(shape, fill, ops), valid, index, Storage.get(storage(shape, fill, ops), index), {==}), 74 Equal.sym(Storage.Storage, storage(shape, fill, snoc(ops, StorageLaws.Read{index})), StorageLaws.kept(storage(shape, fill, ops), index), 75 replay_snoc(ops, StorageLaws.Read{index}, shape, fill))) 76 77 def read_found(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, 78 -valid: {Lifecycle.logical_length(storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}, +index: U32, 79 inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}) -> 80 {Core.read_value(shape, Core.read_flat(shape, Core.TTensor{storage(shape, fill, ops), valid}, index)) == 81 Some{value(shape, fill, ops, index)} : Maybe<&2, F32>}: 82 Equal.trans(Maybe<&2, F32>, Core.read_value(shape, Core.read_flat(shape, Core.TTensor{storage(shape, fill, ops), valid}, index)), 83 Lifecycle.read_value(Storage.get(storage(shape, fill, ops), index)), Some{value(shape, fill, ops, index)}, 84 Core.read_result_value(shape, storage(shape, fill, ops), valid, index, Storage.get(storage(shape, fill, ops), index), {==}), 85 Reads.stored_value(storage(shape, fill, ops), index, value(shape, fill, ops, index), 86 StorageLaws.replay_get(Shape.count(shape), fill, ops, index, Conversion.shape_count_within(shape), 87 Reads.word_inside(shape, index, inside)))) 88 89 def read_shift(+shape: Shape.Shape, -source: Storage.Storage, -target: Storage.Storage, -word: F32, -found: F32, 90 moved: {target == source : Storage.Storage}, same: {found == word : F32}, 91 -P: Core.Tensor & Result<&2, &2, Core.Error, F32> -> Type, 92 h: @-fits: {Lifecycle.logical_length(target) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} -> 93 P((Core.TTensor{target, fits}, Done{found}))) -> 94 @-fits: {Lifecycle.logical_length(source) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} -> 95 P((Core.TTensor{source, fits}, Done{word})): 96 %moved : (@-fits: {Lifecycle.logical_length(_) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} -> 97 P((Core.TTensor{_, fits}, Done{word}))) 98 %same : (@-fits: {Lifecycle.logical_length(target) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} -> 99 P((Core.TTensor{target, fits}, Done{_}))) 100 h 101 102 def read_pair_lift(+shape: Shape.Shape, -kept: Storage.Storage, +found: F32, 103 pair: Core.Tensor & Result<&2, &2, Core.Error, F32>, 104 moved: {read_stored(shape, pair) == kept : Storage.Storage}, 105 read: {Core.read_value(shape, pair) == Some{found} : Maybe<&2, F32>}, 106 -P: Core.Tensor & Result<&2, &2, Core.Error, F32> -> Type, 107 h: @-fits: {Lifecycle.logical_length(kept) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} -> 108 P((Core.TTensor{kept, fits}, Done{found}))) -> P(pair): 109 match pair: 110 case (Core.TTensor{storage, valid}, Fail{error}): 111 Empty.absurd(P((Core.TTensor{storage, valid}, Fail{error})), 112 Arith.clash(Equal.cong(Maybe<&2, F32>, Bool, word => Maybe.is_some(&2, F32, word), None{}, Some{found}, read))) 113 case (Core.TTensor{storage, valid}, Done{+word}): 114 read_shift(shape, storage, kept, word, found, Equal.sym(Storage.Storage, storage, kept, moved), some_value(found, word, Equal.sym(Maybe<&2, F32>, Some{word}, Some{found}, read)), 115 P, h)(valid) 116 117 def read_lift(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, 118 -valid: {Lifecycle.logical_length(storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}, +index: U32, 119 inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}, 120 -P: Core.Tensor & Result<&2, &2, Core.Error, F32> -> Type, 121 h: @-fits: {Lifecycle.logical_length(storage(shape, fill, snoc(ops, StorageLaws.Read{index}))) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} -> 122 P((Core.TTensor{storage(shape, fill, snoc(ops, StorageLaws.Read{index})), fits}, Done{value(shape, fill, ops, index)}))) -> 123 P(Core.read_flat(shape, Core.TTensor{storage(shape, fill, ops), valid}, index)): 124 read_pair_lift(shape, storage(shape, fill, snoc(ops, StorageLaws.Read{index})), value(shape, fill, ops, index), 125 Core.read_flat(shape, Core.TTensor{storage(shape, fill, ops), valid}, index), 126 read_moved(shape, fill, ops, valid, index), read_found(shape, fill, ops, valid, index, inside), P, h) 127 128 def write_owner_projection(-shape: Shape.Shape, -original: Storage.Storage, 129 -valid: {Lifecycle.logical_length(original) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}, +index: U32, +word: F32, 130 pair: Storage.Storage & Result<&2, &2, Storage.Error, Unit>, 131 -same: {pair == Storage.set(original, index, word) : Storage.Storage & Result<&2, &2, Storage.Error, Unit>}) -> 132 {write_stored(shape, Core.write_result(shape, original, valid, index, word, pair, same)) == 133 StorageLaws.owner(Result<&2, &2, Storage.Error, Unit>, pair) : Storage.Storage}: 134 (storage, result) = pair 135 {==} 136 137 def write_status_projection(-shape: Shape.Shape, -original: Storage.Storage, 138 -valid: {Lifecycle.logical_length(original) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}, +index: U32, +word: F32, 139 pair: Storage.Storage & Result<&2, &2, Storage.Error, Unit>, 140 -same: {pair == Storage.set(original, index, word) : Storage.Storage & Result<&2, &2, Storage.Error, Unit>}) -> 141 {write_status(shape, Core.write_result(shape, original, valid, index, word, pair, same)) == 142 Result.is_done(&2, &2, Storage.Error, Unit, StorageLaws.outcome(Result<&2, &2, Storage.Error, Unit>, pair)) : Bool}: 143 (storage, result) = pair 144 match result: 145 case Fail{error}: 146 {==} 147 case Done{unit}: 148 {==} 149 150 def write_moved(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, 151 -valid: {Lifecycle.logical_length(storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}, 152 +index: U32, +word: F32) -> 153 {write_stored(shape, Core.write_flat(shape, Core.TTensor{storage(shape, fill, ops), valid}, index, word)) == 154 storage(shape, fill, snoc(ops, StorageLaws.Write{index, word})) : Storage.Storage}: 155 Equal.trans(Storage.Storage, write_stored(shape, Core.write_flat(shape, Core.TTensor{storage(shape, fill, ops), valid}, index, word)), 156 StorageLaws.written(storage(shape, fill, ops), index, word), storage(shape, fill, snoc(ops, StorageLaws.Write{index, word})), 157 write_owner_projection(shape, storage(shape, fill, ops), valid, index, word, Storage.set(storage(shape, fill, ops), index, word), {==}), 158 Equal.sym(Storage.Storage, storage(shape, fill, snoc(ops, StorageLaws.Write{index, word})), StorageLaws.written(storage(shape, fill, ops), index, word), 159 replay_snoc(ops, StorageLaws.Write{index, word}, shape, fill))) 160 161 def write_done(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, 162 -valid: {Lifecycle.logical_length(storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}, 163 +index: U32, +word: F32, 164 inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}) -> 165 {write_status(shape, Core.write_flat(shape, Core.TTensor{storage(shape, fill, ops), valid}, index, word)) == True{} : Bool}: 166 Equal.trans(Bool, write_status(shape, Core.write_flat(shape, Core.TTensor{storage(shape, fill, ops), valid}, index, word)), 167 Result.is_done(&2, &2, Storage.Error, Unit, StorageLaws.status(storage(shape, fill, ops), index, word)), True{}, 168 write_status_projection(shape, storage(shape, fill, ops), valid, index, word, Storage.set(storage(shape, fill, ops), index, word), {==}), 169 Equal.cong(Result<&2, &2, Storage.Error, Unit>, Bool, Result.is_done(&2, &2, Storage.Error, Unit), 170 StorageLaws.status(storage(shape, fill, ops), index, word), Done{Unit{}}, 171 StorageLaws.replay_set(Shape.count(shape), fill, ops, index, word, Conversion.shape_count_within(shape), 172 Reads.word_inside(shape, index, inside)))) 173 174 def write_shift(+shape: Shape.Shape, -source: Storage.Storage, -target: Storage.Storage, 175 moved: {target == source : Storage.Storage}, 176 -P: Core.Tensor & Result<&2, &2, Core.Error, Unit> -> Type, 177 h: @-fits: {Lifecycle.logical_length(target) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} -> 178 P((Core.TTensor{target, fits}, Done{Unit{}}))) -> 179 @-fits: {Lifecycle.logical_length(source) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} -> 180 P((Core.TTensor{source, fits}, Done{Unit{}})): 181 %moved : (@-fits: {Lifecycle.logical_length(_) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} -> 182 P((Core.TTensor{_, fits}, Done{Unit{}}))) 183 h 184 185 def write_pair_lift(+shape: Shape.Shape, -kept: Storage.Storage, 186 pair: Core.Tensor & Result<&2, &2, Core.Error, Unit>, 187 moved: {write_stored(shape, pair) == kept : Storage.Storage}, 188 done: {write_status(shape, pair) == True{} : Bool}, 189 -P: Core.Tensor & Result<&2, &2, Core.Error, Unit> -> Type, 190 h: @-fits: {Lifecycle.logical_length(kept) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} -> 191 P((Core.TTensor{kept, fits}, Done{Unit{}}))) -> P(pair): 192 match pair: 193 case (Core.TTensor{storage, valid}, Fail{error}): 194 Empty.absurd(P((Core.TTensor{storage, valid}, Fail{error})), Arith.clash(done)) 195 case (Core.TTensor{storage, valid}, Done{Unit{}}): 196 write_shift(shape, storage, kept, Equal.sym(Storage.Storage, storage, kept, moved), P, h)(valid) 197 198 def write_lift(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, 199 -valid: {Lifecycle.logical_length(storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}, 200 +index: U32, +word: F32, 201 inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}, 202 -P: Core.Tensor & Result<&2, &2, Core.Error, Unit> -> Type, 203 h: @-fits: {Lifecycle.logical_length(storage(shape, fill, snoc(ops, StorageLaws.Write{index, word}))) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} -> 204 P((Core.TTensor{storage(shape, fill, snoc(ops, StorageLaws.Write{index, word})), fits}, Done{Unit{}}))) -> 205 P(Core.write_flat(shape, Core.TTensor{storage(shape, fill, ops), valid}, index, word)): 206 write_pair_lift(shape, storage(shape, fill, snoc(ops, StorageLaws.Write{index, word})), 207 Core.write_flat(shape, Core.TTensor{storage(shape, fill, ops), valid}, index, word), 208 write_moved(shape, fill, ops, valid, index, word), write_done(shape, fill, ops, valid, index, word, inside), P, h) 209 210 def new_lift(+shape: Shape.Shape, +fill: F32, -P: Result<&2, &1, Core.Error, Core.Tensor> -> Type, 211 h: @-fits: {Lifecycle.logical_length(storage(shape, fill, Nil{})) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} -> 212 P(Done{Core.TTensor{storage(shape, fill, Nil{}), fits}})) -> P(Core.new(shape, fill)): 213 Reads.new_lift(shape, fill, P, h) 214 215 def expected_snoc(+ops: List<&2, StorageLaws.Op>, +op: StorageLaws.Op, +n: Nat, +index: U32, +word: F32) -> 216 {StorageLaws.expected(snoc(ops, op), n, index, word) == StorageLaws.latest(n, op, index, StorageLaws.expected(ops, n, index, word)) : F32}: 217 match ops: 218 case Nil{}: 219 {==} 220 case +first <> +rest: 221 expected_snoc(rest, op, n, index, StorageLaws.latest(n, first, index, word)) 222 223 def value_after_read(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +read: U32, +index: U32) -> 224 {value(shape, fill, snoc(ops, StorageLaws.Read{read}), index) == value(shape, fill, ops, index) : F32}: 225 expected_snoc(ops, StorageLaws.Read{read}, U32.to_nat(Shape.count(shape)), index, fill)