def parallel_lift
bendtorch-core@0.1.0.0 · src/tensor/contracts/matmul_proof.bend:204
import bendtorch-core@0.1.0.0/src/tensor/contracts/matmul_proof.bend as MatmulProof
import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/tensor/contracts/matmul_parallel_proof.bend as Par
import bendtorch-core@0.1.0.0/src/storage/api.bend as Api
import bendtorch-core@0.1.0.0/src/storage/array.bend as ArrayCore
import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_storage_proof.bend as Owned
import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
204def parallel_lift(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape, +rows: U32, +inner: U32, +right_inner: U32, +columns: U32,
205 -le: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>},
206 -re: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>},
207 -oe: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>},
208 -compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat},
209 +planned: {Matrix.plan(left, right) == Done{Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}} :
210 Result<&2, &2, Matrix.Error, Matrix.Plan<left, right>>},
211 +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
212 -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat},
213 +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
214 -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat},
215 -P: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic> -> Type,
216 h: ParallelFinal(left, right, output, inner, columns, left_fill, left_ops, right_fill, right_ops, StorageLaws.created(Shape.count(output), 0.0), P)) ->
217 P(Matrix.matmul(Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}}, Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}})):
218 +plan : Matrix.Plan<left, right> = Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}
219 +facts : MatrixFacts<left, right, output, rows, inner, right_inner, columns> = plan_facts(left, right, plan, planned)
220 +le_live = left_dimensions(left, right, output, rows, inner, right_inner, columns, facts)
221 +re_live = right_dimensions(left, right, output, rows, inner, right_inner, columns, facts)
222 +oe_live = output_dimensions(left, right, output, rows, inner, right_inner, columns, facts)
223 +compatible_live = inner_compatible(left, right, output, rows, inner, right_inner, columns, facts)
224 +left_extents = Equal.cong(List<&2, U32>, List<&2, Nat>, dims => Shape.naturals(dims), Shape.dimensions(left), [rows, inner], le_live)
225 +right_extents = Equal.trans(List<&2, Nat>, Shape.extents(right), [U32.to_nat(right_inner), U32.to_nat(columns)], [U32.to_nat(inner), U32.to_nat(columns)],
226 Equal.cong(List<&2, U32>, List<&2, Nat>, dims => Shape.naturals(dims), Shape.dimensions(right), [right_inner, columns], re_live),
227 Equal.cong(Nat, List<&2, Nat>, n => [n, U32.to_nat(columns)], U32.to_nat(right_inner), U32.to_nat(inner),
228 Equal.sym(Nat, U32.to_nat(inner), U32.to_nat(right_inner), compatible_live)))
229 +output_extents = Equal.cong(List<&2, U32>, List<&2, Nat>, dims => Shape.naturals(dims), Shape.dimensions(output), [rows, columns], oe_live)
230 %Equal.sym(Result<&2, &2, Matrix.Error, Matrix.Plan<left, right>>, Matrix.plan(left, right), Done{plan}, planned) :
231 P(Matrix.parallel_planned(left, right, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}, _))
232 Replayed.new_lift(output, 0.0,
233 result => P(Matrix.parallel_allocated(left, right, output, U32.to_nat(inner), U32.to_nat(columns), Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}, result)),
234 fo => Owned.created_lift(Shape.count(output), 0.0, Conversion.shape_count_within(output),
235 st => (ParallelFinal(left, right, output, inner, columns, left_fill, left_ops, right_fill, right_ops, st, P) ->
236 @-fo: {Lifecycle.logical_length(st) == Shape.logical_numel(Shape.dimensions(output)) : Nat} ->
237 P(Matrix.parallel_allocated(left, right, output, U32.to_nat(inner), U32.to_nat(columns), Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}, Done{Core.TTensor{st, fo}}))),
238 zero => final => fo =>
239 Par.storage_shift(output, StorageLaws.replay(Par.writes(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(Shape.count(output))), Storage.Empty{}), Storage.Empty{}, Owned.replay_empty(Par.writes(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(Shape.count(output)))),
240 st => k => P(Done{(Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}}, Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}}, Core.Packed{output, Core.TTensor{st, k}})}),
241 final(left_ops)(other => {==})({==})(left_fits)(right_ops)(other => {==})({==})(right_fits))(fo),
242 d => fits => f => final => fo =>
243 Par.input_take(left, left_fill, left_ops,
244 taken => P(Matrix.parallel_output(output, U32.to_nat(inner), U32.to_nat(columns), d, U32.to_nat(Shape.count(output)), f, fo,
245 taken, Elementwise.take(Core.pack(right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits})))),
246 left_kept => left_values => left_words => ha => dl => fa => lsrc => left_kept_fits =>
247 Par.input_take(right, right_fill, right_ops,
248 taken => P(Matrix.parallel_output(output, U32.to_nat(inner), U32.to_nat(columns), d, U32.to_nat(Shape.count(output)), f, fo,
249 (Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_kept), left_kept_fits}}, Api.arr(ha, dl, fa)), taken)),
250 right_kept => right_values => right_words => hb => dr => fb => rsrc => right_kept_fits =>
251 Par.storage_shift(output,
252 StorageLaws.replay(Par.writes(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(Shape.count(output))), Storage.Owned{d, U32.to_nat(Shape.count(output)), Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}),
253 Storage.Owned{d, U32.to_nat(Shape.count(output)),
254 Matrix.build(d, Matrix.product_splits(d, U32.to_nat(inner)), U32.to_nat(inner), U32.to_nat(columns), U32.to_nat(Shape.count(output)), 0n, Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)),
255 Matrix.build_balanced(d, Matrix.product_splits(d, U32.to_nat(inner)), U32.to_nat(inner), U32.to_nat(columns), U32.to_nat(Shape.count(output)), 0n, Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)),
256 f},
257 Par.output_eq(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, ha, dl, fa, right, right_fill, right_ops, hb, dr, fb,
258 lsrc, rsrc, output, U32.to_nat(rows), left_extents, right_extents, output_extents, d, fits, f),
259 st => k => P(Done{(Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_kept), left_kept_fits}},
260 Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_kept), right_kept_fits}},
261 Core.Packed{output, Core.TTensor{st, k}})}),
262 final(left_kept)(left_values)(left_words)(left_kept_fits)(right_kept)(right_values)(right_words)(right_kept_fits))(fo))(right_fits))(left_fits))(h)(fo))