1 import Base 2 import ../shape/core.bend as Shape 3 import ../shape/nat_proofs.bend as Nats 4 import ../shape/index_proofs.bend as Index 5 import ../storage/arith.bend as Arith 6 import ../storage/api.bend as Api 7 import ../storage/model.bend as Model 8 import ../storage/refinement.bend as Refinement 9 import ./core.bend as Core 10 import ./layout.bend as Layout 11 import ./transforms.bend as Transforms 12 import ./slice.bend as Slice 13 import ./slice_LAWS.bend as Laws 14 import ./contracts/build_facts.bend as Facts 15 import ./elementwise.bend as Elementwise 16 import ./elementwise_parallel_proof.bend as Parallel 17 import ./parallel_proofs.bend as Generic 18 import ./batched_matmul_PROOF.bend as BatchedProof 19 import ../storage/array.bend as ArrayCore 20 import ./elementwise_walk_proof.bend as Walk 21 import ../autograd/reduce_fiber_proof.bend as Fiber 22 23 # The clamp. 24 25 def min_left(+a: Nat, +b: Nat, h: {Nat.is_le(a, b) == True{} : Bool}) -> {Nat.min(a, b) == a : Nat}: 26 match a b: 27 case 0n _: 28 {==} 29 case 1n+p 0n: 30 Empty.absurd({Nat.min(1n+p, 0n) == 1n+p : Nat}, Arith.clash(h)) 31 case 1n+p 1n+q: 32 Equal.cong(Nat, Nat, x => 1n+x, Nat.min(p, q), p, min_left(p, q, h)) 33 34 def min_right(+a: Nat, +b: Nat, h: {Nat.is_le(b, a) == True{} : Bool}) -> {Nat.min(a, b) == b : Nat}: 35 match a b: 36 case 0n 0n: 37 {==} 38 case 0n 1n+q: 39 Empty.absurd({Nat.min(0n, 1n+q) == 1n+q : Nat}, Arith.clash(h)) 40 case 1n+p 0n: 41 {==} 42 case 1n+p 1n+q: 43 Equal.cong(Nat, Nat, x => 1n+x, Nat.min(p, q), q, min_right(p, q, h)) 44 45 def min_le_right(+a: Nat, +b: Nat) -> {Nat.is_le(Nat.min(a, b), b) == True{} : Bool}: 46 match a b: 47 case 0n _: 48 Arith.le_zero(b) 49 case 1n+p 0n: 50 {==} 51 case 1n+p 1n+q: 52 min_le_right(p, q) 53 54 # min(a, b) <= min(c, b) for a <= c. 55 def min_mono(+a: Nat, +c: Nat, +b: Nat, h: {Nat.is_le(a, c) == True{} : Bool}) -> {Nat.is_le(Nat.min(a, b), Nat.min(c, b)) == True{} : Bool}: 56 match a c b: 57 case 0n _ _: 58 Arith.le_zero(Nat.min(c, b)) 59 case 1n+p 0n _: 60 Empty.absurd({Nat.is_le(Nat.min(1n+p, b), Nat.min(0n, b)) == True{} : Bool}, Arith.clash(h)) 61 case 1n+p 1n+q 0n: 62 {==} 63 case 1n+p 1n+q 1n+r: 64 min_mono(p, q, r, h) 65 66 def max_right(+a: Nat, +b: Nat, h: {Nat.is_le(a, b) == True{} : Bool}) -> {Nat.max(a, b) == b : Nat}: 67 match a b: 68 case 0n _: 69 {==} 70 case 1n+p 0n: 71 Empty.absurd({Nat.max(1n+p, 0n) == 0n : Nat}, Arith.clash(h)) 72 case 1n+p 1n+q: 73 Equal.cong(Nat, Nat, x => 1n+x, Nat.max(p, q), q, max_right(p, q, h)) 74 75 def max_left(+a: Nat, +b: Nat, h: {Nat.is_le(b, a) == True{} : Bool}) -> {Nat.max(a, b) == a : Nat}: 76 match a b: 77 case 0n 0n: 78 {==} 79 case 0n 1n+q: 80 Empty.absurd({Nat.max(0n, 1n+q) == 0n : Nat}, Arith.clash(h)) 81 case 1n+p 0n: 82 {==} 83 case 1n+p 1n+q: 84 Equal.cong(Nat, Nat, x => 1n+x, Nat.max(p, q), p, max_left(p, q, h)) 85 86 def Laws.slice_start_clamped(extent, start, stop, past): 87 min_right(start, extent, past) 88 89 def Laws.slice_start_kept(extent, start, stop, inside): 90 min_left(start, extent, inside) 91 92 def Laws.slice_stop_clamped(extent, start, stop, past): 93 %Equal.sym(Nat, Nat.min(stop, extent), extent, min_right(stop, extent, past)) : 94 {Nat.max(Nat.min(start, extent), _) == extent : Nat} 95 max_right(Nat.min(start, extent), extent, min_le_right(start, extent)) 96 97 def Laws.slice_stop_kept(extent, start, stop, after, inside): 98 %Equal.sym(Nat, Nat.min(start, extent), start, min_left(start, extent, Facts.le_trans(start, stop, extent, after, inside))) : 99 {Nat.max(_, Nat.min(stop, extent)) == stop : Nat} 100 %Equal.sym(Nat, Nat.min(stop, extent), stop, min_left(stop, extent, inside)) : 101 {Nat.max(start, _) == stop : Nat} 102 max_right(start, stop, after) 103 104 def Laws.slice_stop_empty(extent, start, stop, before): 105 max_left(Nat.min(start, extent), Nat.min(stop, extent), min_mono(stop, start, extent, before)) 106 107 # The length. 108 109 def add_zero_right(+number: Nat) -> {Nat.add(number, 0n) == number : Nat}: 110 match number: 111 case 0n: 112 {==} 113 case 1n+p: 114 Equal.cong(Nat, Nat, x => 1n+x, Nat.add(p, 0n), p, add_zero_right(p)) 115 116 # Nat facts the closed form of the length needs. 117 118 def nat_clash(+x: Nat, e: {1n+x == 0n : Nat}) -> Empty: 119 Arith.clash(Equal.cong(Nat, Bool, v => Nat.is_eq(v, 0n), 1n+x, 0n, e)) 120 121 # stop - low = 1 + rest gives stop = low + (1 + rest). 122 def sub_back(+stop: Nat, +low: Nat, +rest: Nat, e: {Nat.sub(stop, low) == 1n+rest : Nat}) -> {Nat.add(low, 1n+rest) == stop : Nat}: 123 match stop low: 124 case 0n _: 125 Empty.absurd({Nat.add(low, 1n+rest) == 0n : Nat}, nat_clash(rest, Equal.sym(Nat, 0n, 1n+rest, e))) 126 case 1n+ap 0n: 127 Equal.sym(Nat, 1n+ap, 1n+rest, e) 128 case 1n+ap 1n+lp: 129 Equal.cong(Nat, Nat, x => 1n+x, Nat.add(lp, 1n+rest), ap, sub_back(ap, lp, rest, e)) 130 131 # stop - low = 0 gives stop <= low. 132 def sub_zero_le(+stop: Nat, +low: Nat, e: {Nat.sub(stop, low) == 0n : Nat}) -> {Nat.is_le(stop, low) == True{} : Bool}: 133 match stop low: 134 case 0n _: 135 Arith.le_zero(low) 136 case 1n+ap 0n: 137 Empty.absurd({Nat.is_le(1n+ap, 0n) == True{} : Bool}, nat_clash(ap, e)) 138 case 1n+ap 1n+lp: 139 sub_zero_le(ap, lp, e) 140 141 def lt_succ_le(+k: Nat, +q: Nat, h: {Nat.is_lt(k, 1n+q) == True{} : Bool}) -> {Nat.is_le(k, q) == True{} : Bool}: 142 match k q: 143 case 0n _: 144 Arith.le_zero(q) 145 case 1n+p 0n: 146 Empty.absurd({Nat.is_le(1n+p, 0n) == True{} : Bool}, Arith.lt_zero(p, h)) 147 case 1n+p 1n+r: 148 lt_succ_le(p, r, h) 149 150 def lt_le_succ(+r: Nat, +s: Nat, h: {Nat.is_lt(r, s) == True{} : Bool}) -> {Nat.is_le(1n+r, s) == True{} : Bool}: 151 match r s: 152 case 0n 0n: 153 Empty.absurd({Nat.is_le(1n, 0n) == True{} : Bool}, Arith.clash(h)) 154 case 0n 1n+q: 155 Arith.le_zero(q) 156 case 1n+p 0n: 157 Empty.absurd({Nat.is_le(2n+p, 0n) == True{} : Bool}, Arith.clash(h)) 158 case 1n+p 1n+q: 159 lt_le_succ(p, q, h) 160 161 # k <= rest / step gives k * step <= rest; at step 0 the product is 0. 162 def scaled_within(+k: Nat, +rest: Nat, +s: Nat, +k_le: {Nat.is_le(k, Nat.div(rest, 1n+s)) == True{} : Bool}, 163 given: Fiber.Decomposed) -> {Nat.is_le(Nat.mul(k, 1n+s), rest) == True{} : Bool}: 164 Fiber.Decomposed{below, sum} = given 165 Fiber.le_trans(Nat.mul(k, 1n+s), Nat.mul(Nat.div(rest, 1n+s), 1n+s), rest, 166 Fiber.mul_le_mono(k, Nat.div(rest, 1n+s), 1n+s, k_le), 167 %sum : {Nat.is_le(Nat.mul(Nat.div(rest, 1n+s), 1n+s), _) == True{} : Bool} 168 Api.le_add_right(Nat.mul(Nat.div(rest, 1n+s), 1n+s), Nat.mod(rest, 1n+s))) 169 170 def scaled_le(+k: Nat, +rest: Nat, +step: Nat, below: {Nat.is_lt(k, 1n+Nat.div(rest, step)) == True{} : Bool}) -> 171 {Nat.is_le(Nat.mul(k, step), rest) == True{} : Bool}: 172 match step: 173 case 0n: 174 %Equal.sym(Nat, Nat.mul(k, 0n), 0n, Nats.mul_zero_right(k)) : {Nat.is_le(_, rest) == True{} : Bool} 175 Arith.le_zero(rest) 176 case 1n+s: 177 scaled_within(k, rest, s, lt_succ_le(k, Nat.div(rest, 1n+s), below), Fiber.decompose(rest, 1n+s, {==})) 178 179 def below_span(+start: Nat, +stop: Nat, +step: Nat, +k: Nat, +span: Nat, e: {Nat.sub(stop, start) == span : Nat}, 180 below: {Nat.is_lt(k, Slice.counted(span, step)) == True{} : Bool}) -> 181 {Nat.is_lt(Nat.add(start, Nat.mul(k, step)), stop) == True{} : Bool}: 182 match span: 183 case 0n: 184 Empty.absurd({Nat.is_lt(Nat.add(start, Nat.mul(k, step)), stop) == True{} : Bool}, Arith.lt_zero(k, below)) 185 case 1n+rest: 186 %sub_back(stop, start, rest, e) : {Nat.is_lt(Nat.add(start, Nat.mul(k, step)), _) == True{} : Bool} 187 %Equal.sym(Nat, Nat.add(start, 1n+rest), 1n+Nat.add(start, rest), Arith.add_succ(start, rest)) : 188 {Nat.is_lt(Nat.add(start, Nat.mul(k, step)), _) == True{} : Bool} 189 Arith.lt_succ(Nat.add(start, Nat.mul(k, step)), Nat.add(start, rest), 190 Facts.le_add_cong(start, Nat.mul(k, step), rest, scaled_le(k, rest, step, below))) 191 192 def Laws.slice_length_below(start, stop, step, k, below): 193 below_span(start, stop, step, k, Nat.sub(stop, start), {==}, below) 194 195 # 1 + rest <= step + (rest / step) * step for a positive step: rest's 196 # remainder is below the step. 197 def covered(+rest: Nat, +s: Nat, given: Fiber.Decomposed) -> 198 {Nat.is_le(1n+rest, Nat.add(1n+s, Nat.mul(Nat.div(rest, 1n+s), 1n+s))) == True{} : Bool}: 199 Fiber.Decomposed{below, sum} = given 200 +q = Nat.mul(Nat.div(rest, 1n+s), 1n+s) 201 %sum : {Nat.is_le(1n+_, Nat.add(1n+s, q)) == True{} : Bool} 202 %Walk.add_comm(q, 1n+s) : {Nat.is_le(1n+Nat.add(q, Nat.mod(rest, 1n+s)), _) == True{} : Bool} 203 %Arith.add_succ(q, Nat.mod(rest, 1n+s)) : {Nat.is_le(_, Nat.add(q, 1n+s)) == True{} : Bool} 204 Facts.le_add_cong(q, 1n+Nat.mod(rest, 1n+s), 1n+s, lt_le_succ(Nat.mod(rest, 1n+s), 1n+s, below)) 205 206 def cover(+rest: Nat, +step: Nat, positive: {Nat.is_lt(0n, step) == True{} : Bool}) -> 207 {Nat.is_le(1n+rest, Nat.add(step, Nat.mul(Nat.div(rest, step), step))) == True{} : Bool}: 208 match step: 209 case 0n: 210 Empty.absurd({Nat.is_le(1n+rest, Nat.add(0n, Nat.mul(Nat.div(rest, 0n), 0n))) == True{} : Bool}, Arith.clash(positive)) 211 case 1n+s: 212 covered(rest, s, Fiber.decompose(rest, 1n+s, {==})) 213 214 def reaches_span(+start: Nat, +stop: Nat, +step: Nat, +span: Nat, e: {Nat.sub(stop, start) == span : Nat}, 215 positive: {Nat.is_lt(0n, step) == True{} : Bool}) -> 216 {Nat.is_le(stop, Nat.add(start, Nat.mul(Slice.counted(span, step), step))) == True{} : Bool}: 217 match span: 218 case 0n: 219 %Equal.sym(Nat, Nat.add(start, 0n), start, add_zero_right(start)) : {Nat.is_le(stop, _) == True{} : Bool} 220 sub_zero_le(stop, start, e) 221 case 1n+rest: 222 %sub_back(stop, start, rest, e) : {Nat.is_le(_, Nat.add(start, Nat.mul(1n+Nat.div(rest, step), step))) == True{} : Bool} 223 Facts.le_add_cong(start, 1n+rest, Nat.add(step, Nat.mul(Nat.div(rest, step), step)), cover(rest, step, positive)) 224 225 def Laws.slice_length_reaches(start, stop, step, positive): 226 reaches_span(start, stop, step, Nat.sub(stop, start), {==}, positive) 227 228 # The view, the output dimensions and select's dimensions. 229 230 def view_around(+prefix: List<&2, U32>, +extent: U32, +suffix: List<&2, U32>, +outer: Nat) -> 231 {Slice.view_at(List.append(&2, U32, prefix, extent <> suffix), Shape.rank_dimensions(prefix), outer) == 232 Slice.View{Nat.mul(outer, Shape.logical_numel(prefix)), U32.to_nat(extent), Shape.logical_numel(suffix)} : Slice.View}: 233 match prefix: 234 case Nil{}: 235 %Equal.sym(Nat, Nat.mul(outer, 1n), outer, Nats.mul_one_right(outer)) : 236 {Slice.View{outer, U32.to_nat(extent), Shape.logical_numel(suffix)} == 237 Slice.View{_, U32.to_nat(extent), Shape.logical_numel(suffix)} : Slice.View} 238 {==} 239 case +head <> +tail: 240 %Nats.mul_assoc(outer, U32.to_nat(head), Shape.logical_numel(tail)) : 241 {Slice.view_at(List.append(&2, U32, tail, extent <> suffix), Shape.rank_dimensions(tail), Nat.mul(outer, U32.to_nat(head))) == 242 Slice.View{_, U32.to_nat(extent), Shape.logical_numel(suffix)} : Slice.View} 243 view_around(tail, extent, suffix, Nat.mul(outer, U32.to_nat(head))) 244 245 def Laws.slice_view_around(prefix, extent, suffix): 246 Equal.trans(Slice.View, Slice.view(List.append(&2, U32, prefix, extent <> suffix), Shape.rank_dimensions(prefix)), 247 Slice.View{Nat.mul(1n, Shape.logical_numel(prefix)), U32.to_nat(extent), Shape.logical_numel(suffix)}, 248 Slice.View{Shape.logical_numel(prefix), U32.to_nat(extent), Shape.logical_numel(suffix)}, 249 view_around(prefix, extent, suffix, 1n), 250 Equal.cong(Nat, Slice.View, x => Slice.View{x, U32.to_nat(extent), Shape.logical_numel(suffix)}, 251 Nat.mul(1n, Shape.logical_numel(prefix)), Shape.logical_numel(prefix), add_zero_right(Shape.logical_numel(prefix)))) 252 253 def Laws.slice_resized_around(prefix, extent, suffix, other): 254 match prefix: 255 case Nil{}: 256 {==} 257 case +head <> +tail: 258 Equal.cong(List<&2, U32>, List<&2, U32>, rest => head <> rest, 259 Slice.resized(List.append(&2, U32, tail, extent <> suffix), Shape.rank_dimensions(tail), other), 260 List.append(&2, U32, tail, other <> suffix), Laws.slice_resized_around(tail, extent, suffix, other)) 261 262 def Laws.select_dropped_around(prefix, extent, suffix): 263 match prefix: 264 case Nil{}: 265 {==} 266 case +head <> +tail: 267 Equal.cong(List<&2, U32>, List<&2, U32>, rest => head <> rest, 268 Slice.dropped(List.append(&2, U32, tail, extent <> suffix), Shape.rank_dimensions(tail)), 269 List.append(&2, U32, tail, suffix), Laws.select_dropped_around(tail, extent, suffix)) 270 271 # The plan from the request. 272 273 def pick_succ(+c: Bool, +p: Nat, +q: Nat) -> {Bool.pick(Nat, c, 1n+p, 1n+q) == 1n+Bool.pick(Nat, c, p, q) : Nat}: 274 match c: 275 case True{}: 276 {==} 277 case False{}: 278 {==} 279 280 def min_pick(+a: Nat, +b: Nat) -> {Nat.min(a, b) == Bool.pick(Nat, Nat.is_le(a, b), a, b) : Nat}: 281 match a b: 282 case 0n 0n: 283 {==} 284 case 0n 1n+q: 285 {==} 286 case 1n+p 0n: 287 {==} 288 case 1n+p 1n+q: 289 Equal.trans(Nat, 1n+Nat.min(p, q), 1n+Bool.pick(Nat, Nat.is_le(p, q), p, q), Bool.pick(Nat, Nat.is_le(p, q), 1n+p, 1n+q), 290 Equal.cong(Nat, Nat, x => 1n+x, Nat.min(p, q), Bool.pick(Nat, Nat.is_le(p, q), p, q), min_pick(p, q)), 291 Equal.sym(Nat, Bool.pick(Nat, Nat.is_le(p, q), 1n+p, 1n+q), 1n+Bool.pick(Nat, Nat.is_le(p, q), p, q), pick_succ(Nat.is_le(p, q), p, q))) 292 293 def max_pick(+a: Nat, +b: Nat) -> {Nat.max(a, b) == Bool.pick(Nat, Nat.is_le(a, b), b, a) : Nat}: 294 match a b: 295 case 0n 0n: 296 {==} 297 case 0n 1n+q: 298 {==} 299 case 1n+p 0n: 300 {==} 301 case 1n+p 1n+q: 302 Equal.trans(Nat, 1n+Nat.max(p, q), 1n+Bool.pick(Nat, Nat.is_le(p, q), q, p), Bool.pick(Nat, Nat.is_le(p, q), 1n+q, 1n+p), 303 Equal.cong(Nat, Nat, x => 1n+x, Nat.max(p, q), Bool.pick(Nat, Nat.is_le(p, q), q, p), max_pick(p, q)), 304 Equal.sym(Nat, Bool.pick(Nat, Nat.is_le(p, q), 1n+q, 1n+p), 1n+Bool.pick(Nat, Nat.is_le(p, q), q, p), pick_succ(Nat.is_le(p, q), q, p))) 305 306 def view_extent(+given: Slice.View) -> Nat: 307 Slice.View{outer, extent, inner} = given 308 extent 309 310 def view_pair(+given: Slice.View) -> Nat & Nat: 311 Slice.View{outer, extent, inner} = given 312 (outer, inner) 313 314 # The clamp of bounds is PyTorch's, written with min and max. 315 def clamp_axis(+n: Nat, +s: Nat, +t: Nat, +p: Nat) -> 316 {(n, Nat.min(s, n), p, Slice.length(Slice.Bounds{Nat.min(s, n), Nat.max(Nat.min(s, n), Nat.min(t, n))}, p)) == 317 Laws.axis_spec(n, s, t, p) : Nat & Nat & Nat & Nat}: 318 +cs = Bool.pick(Nat, Nat.is_le(s, n), s, n) 319 +ct = Bool.pick(Nat, Nat.is_le(t, n), t, n) 320 %Equal.sym(Nat, Nat.min(s, n), cs, min_pick(s, n)) : 321 {(n, _, p, Slice.length(Slice.Bounds{_, Nat.max(_, Nat.min(t, n))}, p)) == Laws.axis_spec(n, s, t, p) : Nat & Nat & Nat & Nat} 322 %Equal.sym(Nat, Nat.min(t, n), ct, min_pick(t, n)) : 323 {(n, cs, p, Slice.length(Slice.Bounds{cs, Nat.max(cs, _)}, p)) == Laws.axis_spec(n, s, t, p) : Nat & Nat & Nat & Nat} 324 %Equal.sym(Nat, Nat.max(cs, ct), Bool.pick(Nat, Nat.is_le(cs, ct), ct, cs), max_pick(cs, ct)) : 325 {(n, cs, p, Slice.length(Slice.Bounds{cs, _}, p)) == Laws.axis_spec(n, s, t, p) : Nat & Nat & Nat & Nat} 326 {==} 327 328 def viewed_axis(+given: Slice.View, +s: Nat, +t: Nat, +p: Nat) -> 329 {Laws.axis_part(Slice.viewed(given, s, t, p)) == Laws.axis_spec(view_extent(given), s, t, p) : Nat & Nat & Nat & Nat}: 330 match given: 331 case Slice.View{+o, +n, +i}: 332 clamp_axis(n, s, t, p) 333 334 def viewed_view(+given: Slice.View, +s: Nat, +t: Nat, +p: Nat) -> 335 {Laws.view_part(Slice.viewed(given, s, t, p)) == view_pair(given) : Nat & Nat}: 336 match given: 337 case Slice.View{+o, +n, +i}: 338 {==} 339 340 def axis_extent_around(+prefix: List<&2, U32>, +extent: U32, +suffix: List<&2, U32>) -> 341 {Slice.axis_extent(List.append(&2, U32, prefix, extent <> suffix), Shape.rank_dimensions(prefix)) == U32.to_nat(extent) : Nat}: 342 match prefix: 343 case Nil{}: 344 {==} 345 case +head <> +tail: 346 axis_extent_around(tail, extent, suffix) 347 348 def sized_extent(+prefix: List<&2, U32>, +extent: U32, +suffix: List<&2, U32>, +empty: Bool) -> 349 {view_extent(Slice.sized(List.append(&2, U32, prefix, extent <> suffix), Shape.rank_dimensions(prefix), empty)) == U32.to_nat(extent) : Nat}: 350 match empty: 351 case True{}: 352 axis_extent_around(prefix, extent, suffix) 353 case False{}: 354 Equal.cong(Slice.View, Nat, v => view_extent(v), 355 Slice.view(List.append(&2, U32, prefix, extent <> suffix), Shape.rank_dimensions(prefix)), 356 Slice.View{Shape.logical_numel(prefix), U32.to_nat(extent), Shape.logical_numel(suffix)}, 357 Laws.slice_view_around(prefix, extent, suffix)) 358 359 def Laws.slice_plan_axis(prefix, extent, suffix, dim, at, start, stop, step): 360 +dims = List.append(&2, U32, prefix, extent <> suffix) 361 %Equal.sym(Nat, U32.to_nat(dim), Shape.rank_dimensions(prefix), at) : 362 {Laws.axis_part(Slice.viewed(Slice.sized(dims, _, Slice.holds_zero(dims)), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step))) == 363 Laws.axis_spec(U32.to_nat(extent), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step)) : Nat & Nat & Nat & Nat} 364 Equal.trans(Nat & Nat & Nat & Nat, 365 Laws.axis_part(Slice.viewed(Slice.sized(dims, Shape.rank_dimensions(prefix), Slice.holds_zero(dims)), 366 U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step))), 367 Laws.axis_spec(view_extent(Slice.sized(dims, Shape.rank_dimensions(prefix), Slice.holds_zero(dims))), 368 U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step)), 369 Laws.axis_spec(U32.to_nat(extent), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step)), 370 viewed_axis(Slice.sized(dims, Shape.rank_dimensions(prefix), Slice.holds_zero(dims)), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step)), 371 Equal.cong(Nat, Nat & Nat & Nat & Nat, n => Laws.axis_spec(n, U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step)), 372 view_extent(Slice.sized(dims, Shape.rank_dimensions(prefix), Slice.holds_zero(dims))), U32.to_nat(extent), 373 sized_extent(prefix, extent, suffix, Slice.holds_zero(dims)))) 374 375 # Some dimension 0 makes the element count 0. 376 def zero_head(+extent: U32, +rest: List<&2, U32>, h: {Slice.holds_zero(extent <> rest) == True{} : Bool}, 377 rec: @h: {Slice.holds_zero(rest) == True{} : Bool} -> {Shape.logical_numel(rest) == 0n : Nat}, 378 decision: Shape.BoolCase) -> {Shape.logical_numel(extent <> rest) == 0n : Nat}: 379 match decision: 380 case Shape.BTrue{equal}: 381 %Equal.sym(Nat, U32.to_nat(extent), 0n, Arith.eq_true(U32.to_nat(extent), 0n, Equal.sym(Bool, True{}, Nat.is_eq(U32.to_nat(extent), 0n), equal))) : 382 {Nat.mul(_, Shape.logical_numel(rest)) == 0n : Nat} 383 {==} 384 case Shape.BFalse{equal}: 385 %Equal.sym(Nat, Shape.logical_numel(rest), 0n, 386 rec(Equal.trans(Bool, Slice.holds_zero(rest), Bool.or(Nat.is_eq(U32.to_nat(extent), 0n), Slice.holds_zero(rest)), True{}, 387 Equal.cong(Bool, Bool, x => Bool.or(x, Slice.holds_zero(rest)), False{}, Nat.is_eq(U32.to_nat(extent), 0n), equal), h))) : 388 {Nat.mul(U32.to_nat(extent), _) == 0n : Nat} 389 Nats.mul_zero_right(U32.to_nat(extent)) 390 391 def zero_numel(+dimensions: List<&2, U32>, h: {Slice.holds_zero(dimensions) == True{} : Bool}) -> 392 {Shape.logical_numel(dimensions) == 0n : Nat}: 393 match dimensions: 394 case Nil{}: 395 Empty.absurd({Shape.logical_numel(Nil{}) == 0n : Nat}, Arith.clash(h)) 396 case +extent <> +rest: 397 zero_head(extent, rest, h, hr => zero_numel(rest, hr), Shape.bool_case(Nat.is_eq(U32.to_nat(extent), 0n))) 398 399 # A product a * b of 0 with a not 0 has b = 0. 400 def add_zero_left_part(+b: Nat, +x: Nat, e: {Nat.add(b, x) == 0n : Nat}) -> {b == 0n : Nat}: 401 match b: 402 case 0n: 403 {==} 404 case 1n+q: 405 Empty.absurd({1n+q == 0n : Nat}, nat_clash(Nat.add(q, x), e)) 406 407 def mul_zero_factor(+a: Nat, +b: Nat, e: {Nat.mul(a, b) == 0n : Nat}, nonzero: {False{} == Nat.is_eq(a, 0n) : Bool}) -> {b == 0n : Nat}: 408 match a: 409 case 0n: 410 Empty.absurd({b == 0n : Nat}, Arith.clash(nonzero)) 411 case 1n+p: 412 add_zero_left_part(b, Nat.mul(p, b), e) 413 414 # An element count of 0 has some dimension 0. 415 def zero_found(+extent: U32, +rest: List<&2, U32>, e: {Shape.logical_numel(extent <> rest) == 0n : Nat}, 416 rec: @e: {Shape.logical_numel(rest) == 0n : Nat} -> {Slice.holds_zero(rest) == True{} : Bool}, 417 decision: Shape.BoolCase) -> {Slice.holds_zero(extent <> rest) == True{} : Bool}: 418 match decision: 419 case Shape.BTrue{equal}: 420 %equal : {Bool.or(_, Slice.holds_zero(rest)) == True{} : Bool} 421 {==} 422 case Shape.BFalse{+equal}: 423 %equal : {Bool.or(_, Slice.holds_zero(rest)) == True{} : Bool} 424 rec(mul_zero_factor(U32.to_nat(extent), Shape.logical_numel(rest), e, equal)) 425 426 def holds_true(+dimensions: List<&2, U32>, e: {Shape.logical_numel(dimensions) == 0n : Nat}) -> {Slice.holds_zero(dimensions) == True{} : Bool}: 427 match dimensions: 428 case Nil{}: 429 Empty.absurd({Slice.holds_zero(Nil{}) == True{} : Bool}, nat_clash(0n, e)) 430 case +extent <> +rest: 431 zero_found(extent, rest, e, er => holds_true(rest, er), Shape.bool_case(Nat.is_eq(U32.to_nat(extent), 0n))) 432 433 def holds_false(+dimensions: List<&2, U32>, +nonempty: {Nat.is_lt(0n, Shape.logical_numel(dimensions)) == True{} : Bool}, 434 decision: Shape.BoolCase) -> {Slice.holds_zero(dimensions) == False{} : Bool}: 435 match decision: 436 case Shape.BFalse{equal}: 437 Equal.sym(Bool, False{}, Slice.holds_zero(dimensions), equal) 438 case Shape.BTrue{equal}: 439 Empty.absurd({Slice.holds_zero(dimensions) == False{} : Bool}, 440 Arith.clash(Equal.trans(Bool, False{}, Nat.is_lt(0n, Shape.logical_numel(dimensions)), True{}, 441 Equal.cong(Nat, Bool, v => Nat.is_lt(0n, v), 0n, Shape.logical_numel(dimensions), 442 Equal.sym(Nat, Shape.logical_numel(dimensions), 0n, zero_numel(dimensions, Equal.sym(Bool, True{}, Slice.holds_zero(dimensions), equal)))), 443 nonempty))) 444 445 def Laws.slice_plan_view(prefix, extent, suffix, dim, at, start, stop, step, nonempty): 446 +dims = List.append(&2, U32, prefix, extent <> suffix) 447 %Equal.sym(Nat, U32.to_nat(dim), Shape.rank_dimensions(prefix), at) : 448 {Laws.view_part(Slice.viewed(Slice.sized(dims, _, Slice.holds_zero(dims)), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step))) == 449 (Shape.logical_numel(prefix), Shape.logical_numel(suffix)) : Nat & Nat} 450 %Equal.sym(Bool, Slice.holds_zero(dims), False{}, holds_false(dims, nonempty, Shape.bool_case(Slice.holds_zero(dims)))) : 451 {Laws.view_part(Slice.viewed(Slice.sized(dims, Shape.rank_dimensions(prefix), _), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step))) == 452 (Shape.logical_numel(prefix), Shape.logical_numel(suffix)) : Nat & Nat} 453 %Equal.sym(Slice.View, Slice.view(dims, Shape.rank_dimensions(prefix)), 454 Slice.View{Shape.logical_numel(prefix), U32.to_nat(extent), Shape.logical_numel(suffix)}, Laws.slice_view_around(prefix, extent, suffix)) : 455 {Laws.view_part(Slice.viewed(_, U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step))) == 456 (Shape.logical_numel(prefix), Shape.logical_numel(suffix)) : Nat & Nat} 457 {==} 458 459 def Laws.slice_plan_empty_view(dimensions, dim, start, stop, step, empty): 460 %Equal.sym(Bool, Slice.holds_zero(dimensions), True{}, holds_true(dimensions, empty)) : 461 {Laws.view_part(Slice.viewed(Slice.sized(dimensions, U32.to_nat(dim), _), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step))) == 462 (0n, 0n) : Nat & Nat} 463 {==} 464 465 # The checks and the plan. 466 467 def Laws.slice_axis_refused(input, dim, start, stop, step, outside): 468 %Equal.sym(Bool, Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(input))), False{}, outside) : 469 {Slice.checked(Shape.dimensions(input), dim, start, stop, step, Slice.refused(_, Nat.is_lt(0n, U32.to_nat(step)))) == 470 Fail{Slice.LayoutError{Layout.AxisOutOfBounds{}}} : Result<&2, &2, Slice.Error, Slice.Planned>} 471 {==} 472 473 def Laws.slice_step_refused(input, dim, start, stop, step, inside, zero): 474 %Equal.sym(Bool, Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(input))), True{}, inside) : 475 {Slice.checked(Shape.dimensions(input), dim, start, stop, step, Slice.refused(_, Nat.is_lt(0n, U32.to_nat(step)))) == 476 Fail{Slice.StepNotPositive{}} : Result<&2, &2, Slice.Error, Slice.Planned>} 477 %Equal.sym(Nat, U32.to_nat(step), 0n, zero) : 478 {Slice.checked(Shape.dimensions(input), dim, start, stop, step, Slice.refused(True{}, Nat.is_lt(0n, _))) == 479 Fail{Slice.StepNotPositive{}} : Result<&2, &2, Slice.Error, Slice.Planned>} 480 {==} 481 482 def planned_shaped(+dimensions: List<&2, U32>, +dim: U32, +start: U32, +stop: U32, +step: U32, 483 built: Shape.ConstructResult) -> 484 Laws.PlannedAs(dimensions, dim, start, stop, step, 485 Slice.shaped(Slice.plan_of(dimensions, dim, start, stop, step), 486 Slice.resized(dimensions, U32.to_nat(dim), U32.from_nat(Slice.plan_length(Slice.plan_of(dimensions, dim, start, stop, step)))), built)): 487 match built: 488 case Shape.SError{error}: 489 Unit{} 490 case Shape.SBuilt{+shape, preserved}: 491 ({==}, preserved) 492 493 def planned_checked(+dimensions: List<&2, U32>, +dim: U32, +start: U32, +stop: U32, +step: U32, refused: Maybe<&2, Slice.Error>) -> 494 Laws.PlannedAs(dimensions, dim, start, stop, step, Slice.checked(dimensions, dim, start, stop, step, refused)): 495 match refused: 496 case Some{error}: 497 Unit{} 498 case None{}: 499 planned_shaped(dimensions, dim, start, stop, step, 500 Shape.construct(Slice.resized(dimensions, U32.to_nat(dim), U32.from_nat(Slice.plan_length(Slice.plan_of(dimensions, dim, start, stop, step)))))) 501 502 def Laws.slice_planned(input, dim, start, stop, step): 503 planned_checked(Shape.dimensions(input), dim, start, stop, step, Slice.refusal(Shape.dimensions(input), dim, step)) 504 505 def Laws.slice_refused(input, dim, start, stop, step, error, refused): 506 Core.Packed{+shape, tensor} = input 507 %Equal.sym(Result<&2, &2, Slice.Error, Slice.Planned>, Slice.plan(shape, dim, start, stop, step), Fail{error}, refused) : 508 {Slice.sliced(shape, tensor, _) == (Core.Packed{shape, tensor}, Fail{error}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 509 {==} 510 511 # narrow and select. 512 513 def Laws.narrow_out_of_range(input, dim, start, length, inside, past): 514 Core.Packed{+shape, tensor} = input 515 %Equal.sym(Bool, Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(shape))), True{}, inside) : 516 {Slice.narrow_checked(shape, tensor, dim, start, length, Slice.refused(_, Nat.is_lt(0n, U32.to_nat(1)))) == 517 (Core.Packed{shape, tensor}, Fail{Slice.OutOfRange{}}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 518 %Equal.sym(Bool, Nat.is_le(Nat.add(U32.to_nat(start), U32.to_nat(length)), Slice.extent_at(Shape.dimensions(shape), dim)), False{}, past) : 519 {Slice.narrow_ranged(shape, tensor, dim, start, length, _) == 520 (Core.Packed{shape, tensor}, Fail{Slice.OutOfRange{}}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 521 {==} 522 523 def Laws.narrow_slices(input, dim, start, length, inside, fits): 524 Core.Packed{+shape, tensor} = input 525 %Equal.sym(Bool, Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(shape))), True{}, inside) : 526 {Slice.narrow_checked(shape, tensor, dim, start, length, Slice.refused(_, Nat.is_lt(0n, U32.to_nat(1)))) == 527 Slice.slice(Core.Packed{shape, tensor}, dim, start, U32.add(start, length), 1) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 528 %Equal.sym(Bool, Nat.is_le(Nat.add(U32.to_nat(start), U32.to_nat(length)), Slice.extent_at(Shape.dimensions(shape), dim)), True{}, fits) : 529 {Slice.narrow_ranged(shape, tensor, dim, start, length, _) == 530 Slice.slice(Core.Packed{shape, tensor}, dim, start, U32.add(start, length), 1) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 531 {==} 532 533 def Laws.narrow_axis_refused(input, dim, start, length, outside): 534 Core.Packed{+shape, tensor} = input 535 %Equal.sym(Bool, Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(shape))), False{}, outside) : 536 {Slice.narrow_checked(shape, tensor, dim, start, length, Slice.refused(_, Nat.is_lt(0n, U32.to_nat(1)))) == 537 (Core.Packed{shape, tensor}, Fail{Slice.LayoutError{Layout.AxisOutOfBounds{}}}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 538 {==} 539 540 def transformed_as(result: Result<&2, &1, Transforms.Error, Core.Dynamic>) -> 541 {Slice.transformed(result) == Laws.reshape_as_slice(result) : Result<&2, &1, Slice.Error, Core.Dynamic>}: 542 match result: 543 case Fail{Transforms.LayoutError{error}}: 544 {==} 545 case Fail{Transforms.TensorError{error}}: 546 {==} 547 case Done{output}: 548 {==} 549 550 def reshaped_as(+dimensions: List<&2, U32>, pair: Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>) -> 551 {Slice.reshaped(dimensions, pair) == Laws.reshaped(dimensions, pair) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}: 552 match pair: 553 case Tuple{input, Fail{error}}: 554 {==} 555 case Tuple{input, Done{output}}: 556 Equal.cong(Result<&2, &1, Slice.Error, Core.Dynamic>, Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>, 557 result => (input, result), Slice.transformed(Transforms.reshape_dynamic(output, dimensions)), 558 Laws.reshape_as_slice(Transforms.reshape_dynamic(output, dimensions)), transformed_as(Transforms.reshape_dynamic(output, dimensions))) 559 560 def Laws.select_out_of_range(input, dim, index, inside, past): 561 Core.Packed{+shape, tensor} = input 562 %Equal.sym(Bool, Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(shape))), True{}, inside) : 563 {Slice.select_checked(shape, tensor, dim, index, Slice.refused(_, Nat.is_lt(0n, U32.to_nat(1)))) == 564 (Core.Packed{shape, tensor}, Fail{Slice.OutOfRange{}}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 565 %Equal.sym(Bool, Nat.is_lt(U32.to_nat(index), Slice.extent_at(Shape.dimensions(shape), dim)), False{}, past) : 566 {Slice.select_ranged(shape, tensor, dim, index, _) == 567 (Core.Packed{shape, tensor}, Fail{Slice.OutOfRange{}}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 568 {==} 569 570 def Laws.select_slices(input, dim, index, inside, below): 571 Core.Packed{+shape, tensor} = input 572 %Equal.sym(Bool, Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(shape))), True{}, inside) : 573 {Slice.select_checked(shape, tensor, dim, index, Slice.refused(_, Nat.is_lt(0n, U32.to_nat(1)))) == 574 Laws.reshaped(Slice.dropped(Shape.dimensions(shape), U32.to_nat(dim)), Slice.slice(Core.Packed{shape, tensor}, dim, index, U32.add(index, 1), 1)) : 575 Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 576 %Equal.sym(Bool, Nat.is_lt(U32.to_nat(index), Slice.extent_at(Shape.dimensions(shape), dim)), True{}, below) : 577 {Slice.select_ranged(shape, tensor, dim, index, _) == 578 Laws.reshaped(Slice.dropped(Shape.dimensions(shape), U32.to_nat(dim)), Slice.slice(Core.Packed{shape, tensor}, dim, index, U32.add(index, 1), 1)) : 579 Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 580 reshaped_as(Slice.dropped(Shape.dimensions(shape), U32.to_nat(dim)), Slice.slice(Core.Packed{shape, tensor}, dim, index, U32.add(index, 1), 1)) 581 582 def Laws.select_axis_refused(input, dim, index, outside): 583 Core.Packed{+shape, tensor} = input 584 %Equal.sym(Bool, Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(shape))), False{}, outside) : 585 {Slice.select_checked(shape, tensor, dim, index, Slice.refused(_, Nat.is_lt(0n, U32.to_nat(1)))) == 586 (Core.Packed{shape, tensor}, Fail{Slice.LayoutError{Layout.AxisOutOfBounds{}}}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 587 {==} 588 589 # The cells of the builds. 590 591 def take_value(-b: Array, pair: Array & F32) -> 592 {Pair.fst(F32, Array & Array, Slice.take(b, pair)) == Refinement.second(F32, pair) : F32}: 593 (a, x) = pair 594 {==} 595 596 def take_state(-b: Array, pair: Array & F32) -> 597 {Pair.snd(F32, Array & Array, Slice.take(b, pair)) == (Refinement.first(F32, pair), b) : Array & Array}: 598 (a, x) = pair 599 {==} 600 601 # A read of A through `take` gives the word A holds there ... 602 def read_word(+da: Nat, -A: Array, -b: Array, +index: Nat, -valid: Model.Balanced(F32, da, A)) -> 603 Parallel.Fact<{Some{Pair.fst(F32, Array & Array, Slice.take(b, Array.get(F32, A, U32.from_nat(index))))} == 604 Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(index))) : Maybe<&2, F32>}>: 605 Parallel.Fact{ 606 %Equal.sym(F32, Pair.fst(F32, Array & Array, Slice.take(b, Array.get(F32, A, U32.from_nat(index)))), 607 Refinement.second(F32, Array.get(F32, A, U32.from_nat(index))), take_value(b, Array.get(F32, A, U32.from_nat(index)))) : 608 {Some{_} == Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(index))) : Maybe<&2, F32>} 609 {\{Parallel.Fact: q => q} : 610 Parallel.Fact<{Some{Refinement.second(F32, Array.get(F32, A, U32.from_nat(index)))} == 611 Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(index))) : Maybe<&2, F32>}> -> 612 {Some{Refinement.second(F32, Array.get(F32, A, U32.from_nat(index)))} == 613 Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(index))) : Maybe<&2, F32>}}( 614 Parallel.read_value(da, A, A, U32.from_nat(index), Array.get(F32, A, U32.from_nat(index)), {==}, {==}, valid))} 615 616 # ... and hands both arrays back. 617 def read_keeps(+da: Nat, -A: Array, -b: Array, +index: Nat, -valid: Model.Balanced(F32, da, A)) -> 618 Parallel.Fact<{Pair.snd(F32, Array & Array, Slice.take(b, Array.get(F32, A, U32.from_nat(index)))) == (A, b) : 619 Array & Array}>: 620 Parallel.Fact{Equal.trans(Array & Array, Pair.snd(F32, Array & Array, Slice.take(b, Array.get(F32, A, U32.from_nat(index)))), 621 (Refinement.first(F32, Array.get(F32, A, U32.from_nat(index))), b), (A, b), 622 take_state(b, Array.get(F32, A, U32.from_nat(index))), 623 Equal.cong(Array, Array & Array, x => (x, b), Refinement.first(F32, Array.get(F32, A, U32.from_nat(index))), A, 624 {\{Parallel.Fact: q => q} : 625 Parallel.Fact<{Refinement.first(F32, Array.get(F32, A, U32.from_nat(index))) == A : Array}> -> 626 {Refinement.first(F32, Array.get(F32, A, U32.from_nat(index))) == A : Array}}( 627 Parallel.read_keeps(da, A, A, U32.from_nat(index), Array.get(F32, A, U32.from_nat(index)), {==}, {==}, valid))))} 628 629 def empty() -> Array: 630 Array.new(F32, 0n, 0.0) 631 632 def cell_keeps(+plan: Slice.Plan, +at: Nat, +da: Nat, -A: Array, inside: Bool, -valid: Model.Balanced(F32, da, A)) -> 633 Parallel.Fact<{Pair.snd(F32, Array & Array, Slice.cell(plan, at, (A, empty()), inside)) == (A, empty()) : Array & Array}>: 634 match inside: 635 case False{}: 636 Parallel.Fact{{==}} 637 case True{}: 638 read_keeps(da, A, empty(), Slice.source(plan, at), valid) 639 640 def scattered_keeps(+da: Nat, -A: Array, at: Maybe<&2, Nat>, -valid: Model.Balanced(F32, da, A)) -> 641 Parallel.Fact<{Pair.snd(F32, Array & Array, Slice.scattered((A, empty()), at)) == (A, empty()) : Array & Array}>: 642 match at: 643 case None{}: 644 Parallel.Fact{{==}} 645 case Some{+index}: 646 read_keeps(da, A, empty(), index, valid) 647 648 def scatter_keeps(+plan: Slice.Plan, +at: Nat, +da: Nat, -A: Array, inside: Bool, -valid: Model.Balanced(F32, da, A)) -> 649 Parallel.Fact<{Pair.snd(F32, Array & Array, Slice.scatter_cell(plan, at, (A, empty()), inside)) == (A, empty()) : Array & Array}>: 650 match inside: 651 case False{}: 652 Parallel.Fact{{==}} 653 case True{}: 654 scattered_keeps(da, A, Slice.target(plan, at), valid) 655 656 def SplitFacts(-A: Array) -> Type: 657 {Pair.fst(Array & Array, Array & Array, Elementwise.split((A, empty()))) == (A, empty()) : Array & Array} & 658 {Pair.snd(Array & Array, Array & Array, Elementwise.split((A, empty()))) == (A, empty()) : Array & Array} 659 660 def split_facts(+da: Nat, -A: Array, -valid: Model.Balanced(F32, da, A)) -> Parallel.Fact: 661 Parallel.Fact{{\{Parallel.Fact: q => q} : Parallel.Fact -> SplitFacts(A)}( 662 Parallel.split_equal(da, 0n, A, empty(), valid, ArrayCore.new_balanced(F32, 0n, 0.0)))} 663 664 # Every slot of the forward build holds what its cell computes from (A, []). 665 # The Fact holds the cells eta-expanded: `--verdict` rejects the unexpanded 666 # function in this position although `--check-only` accepts it. 667 def build_cells(+depth: Nat, +splits: Nat, +plan: Slice.Plan, +count: Nat, +base: Nat, +da: Nat, -A: Array, a: Array, 668 -same: {a == A : Array}, -valid: Model.Balanced(F32, da, A)) -> 669 Parallel.Fact, ~Array, ~Slice.cell, plan, (A, empty()), depth, count, base, 670 Slice.build(depth, splits, plan, count, base, a))>: 671 Parallel.Fact{path => 672 {\{Generic.Fact: q => q} : 673 Generic.Fact, ~Array, ~Slice.cell, plan, (A, empty()), depth, count, base, 674 Slice.build(depth, splits, plan, count, base, a))> -> 675 Generic.Cells(~Slice.Plan, ~Array, ~Array, ~Slice.cell, plan, (A, empty()), depth, count, base, 676 Slice.build(depth, splits, plan, count, base, a))}( 677 Generic.build_cells(~Slice.Plan, ~Slice.Plan, ~Array, ~Array, ~Slice.same, ~Elementwise.split, ~Slice.cell, 678 plan, (A, empty()), depth, splits, count, base, (a, empty()), Parallel.pair_equal(a, empty(), A, empty(), same, {==}), 679 at => inside => {\{Parallel.Fact: q => q} : 680 Parallel.Fact<{Pair.snd(F32, Array & Array, Slice.cell(plan, at, (A, empty()), inside)) == (A, empty()) : Array & Array}> -> 681 {Pair.snd(F32, Array & Array, Slice.cell(plan, at, (A, empty()), inside)) == (A, empty()) : Array & Array}}( 682 cell_keeps(plan, at, da, A, inside, valid)), 683 Pair.fst({Pair.fst(Array & Array, Array & Array, Elementwise.split((A, empty()))) == (A, empty()) : Array & Array}, 684 {Pair.snd(Array & Array, Array & Array, Elementwise.split((A, empty()))) == (A, empty()) : Array & Array}, 685 {\{Parallel.Fact: q => q} : Parallel.Fact -> SplitFacts(A)}(split_facts(da, A, valid))), 686 Pair.snd({Pair.fst(Array & Array, Array & Array, Elementwise.split((A, empty()))) == (A, empty()) : Array & Array}, 687 {Pair.snd(Array & Array, Array & Array, Elementwise.split((A, empty()))) == (A, empty()) : Array & Array}, 688 {\{Parallel.Fact: q => q} : Parallel.Fact -> SplitFacts(A)}(split_facts(da, A, valid)))))(path)} 689 690 def scatter_cells(+depth: Nat, +splits: Nat, +plan: Slice.Plan, +count: Nat, +base: Nat, +da: Nat, -A: Array, a: Array, 691 -same: {a == A : Array}, -valid: Model.Balanced(F32, da, A)) -> 692 Parallel.Fact, ~Array, ~Slice.scatter_cell, plan, (A, empty()), depth, count, base, 693 Slice.scatter_build(depth, splits, plan, count, base, a))>: 694 Parallel.Fact{path => 695 {\{Generic.Fact: q => q} : 696 Generic.Fact, ~Array, ~Slice.scatter_cell, plan, (A, empty()), depth, count, base, 697 Slice.scatter_build(depth, splits, plan, count, base, a))> -> 698 Generic.Cells(~Slice.Plan, ~Array, ~Array, ~Slice.scatter_cell, plan, (A, empty()), depth, count, base, 699 Slice.scatter_build(depth, splits, plan, count, base, a))}( 700 Generic.build_cells(~Slice.Plan, ~Slice.Plan, ~Array, ~Array, ~Slice.same, ~Elementwise.split, ~Slice.scatter_cell, 701 plan, (A, empty()), depth, splits, count, base, (a, empty()), Parallel.pair_equal(a, empty(), A, empty(), same, {==}), 702 at => inside => {\{Parallel.Fact: q => q} : 703 Parallel.Fact<{Pair.snd(F32, Array & Array, Slice.scatter_cell(plan, at, (A, empty()), inside)) == (A, empty()) : Array & Array}> -> 704 {Pair.snd(F32, Array & Array, Slice.scatter_cell(plan, at, (A, empty()), inside)) == (A, empty()) : Array & Array}}( 705 scatter_keeps(plan, at, da, A, inside, valid)), 706 Pair.fst({Pair.fst(Array & Array, Array & Array, Elementwise.split((A, empty()))) == (A, empty()) : Array & Array}, 707 {Pair.snd(Array & Array, Array & Array, Elementwise.split((A, empty()))) == (A, empty()) : Array & Array}, 708 {\{Parallel.Fact: q => q} : Parallel.Fact -> SplitFacts(A)}(split_facts(da, A, valid))), 709 Pair.snd({Pair.fst(Array & Array, Array & Array, Elementwise.split((A, empty()))) == (A, empty()) : Array & Array}, 710 {Pair.snd(Array & Array, Array & Array, Elementwise.split((A, empty()))) == (A, empty()) : Array & Array}, 711 {\{Parallel.Fact: q => q} : Parallel.Fact -> SplitFacts(A)}(split_facts(da, A, valid)))))(path)} 712 713 # The arithmetic of the positions. 714 715 def sub_zero(+number: Nat) -> {Nat.sub(number, 0n) == number : Nat}: 716 match number: 717 case 0n: 718 {==} 719 case 1n+p: 720 {==} 721 722 def sub_cancel(+a: Nat, +b: Nat) -> {Nat.sub(Nat.add(a, b), a) == b : Nat}: 723 match a: 724 case 0n: 725 sub_zero(b) 726 case 1n+p: 727 sub_cancel(p, b) 728 729 # (k * step) / step = k for a positive step. 730 def div_step(+k: Nat, +step: Nat, +positive: {Nat.is_lt(0n, step) == True{} : Bool}) -> {Nat.div(Nat.mul(k, step), step) == k : Nat}: 731 Equal.trans(Nat, Nat.div(Nat.mul(k, step), step), Nat.div(Nat.add(Nat.mul(k, step), 0n), step), k, 732 Equal.cong(Nat, Nat, x => Nat.div(x, step), Nat.mul(k, step), Nat.add(Nat.mul(k, step), 0n), 733 Equal.sym(Nat, Nat.add(Nat.mul(k, step), 0n), Nat.mul(k, step), add_zero_right(Nat.mul(k, step)))), 734 BatchedProof.div_exact(k, step, 0n, positive)) 735 736 # The forward cell (o, k, r) reads input cell (o, start + k * step, r). 737 def source_at(+outer: Nat, +extent: Nat, +inner: Nat, +start: Nat, +step: Nat, +length: Nat, +o: Nat, +k: Nat, +r: Nat, 738 +k_below: {Nat.is_lt(k, length) == True{} : Bool}, +r_below: {Nat.is_lt(r, inner) == True{} : Bool}) -> 739 {Slice.source(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(length, inner, o, k, r)) == 740 Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r) : Nat}: 741 +row = Nat.add(Nat.mul(o, length), k) 742 +pos = Laws.position(length, inner, o, k, r) 743 %Equal.sym(Nat, Nat.div(pos, inner), row, BatchedProof.div_exact(row, inner, r, r_below)) : 744 {Nat.add(Nat.mul(Nat.add(Nat.mul(Nat.div(_, length), extent), Nat.add(start, Nat.mul(Nat.mod(_, length), step))), inner), Nat.mod(pos, inner)) == 745 Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r) : Nat} 746 %Equal.sym(Nat, Nat.mod(pos, inner), r, BatchedProof.mod_exact(row, inner, r, r_below)) : 747 {Nat.add(Nat.mul(Nat.add(Nat.mul(Nat.div(row, length), extent), Nat.add(start, Nat.mul(Nat.mod(row, length), step))), inner), _) == 748 Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r) : Nat} 749 %Equal.sym(Nat, Nat.div(row, length), o, BatchedProof.div_exact(o, length, k, k_below)) : 750 {Nat.add(Nat.mul(Nat.add(Nat.mul(_, extent), Nat.add(start, Nat.mul(Nat.mod(row, length), step))), inner), r) == 751 Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r) : Nat} 752 %Equal.sym(Nat, Nat.mod(row, length), k, BatchedProof.mod_exact(o, length, k, k_below)) : 753 {Nat.add(Nat.mul(Nat.add(Nat.mul(o, extent), Nat.add(start, Nat.mul(_, step))), inner), r) == 754 Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r) : Nat} 755 {==} 756 757 def Laws.slice_build_cells(depth, splits, outer, extent, inner, start, step, length, count, base, da, A, a, same, valid, path, o, k, r, 758 k_below, r_below, at, inside): 759 Laws.Holds{ 760 %Equal.trans(Nat, Slice.source(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base)), 761 Slice.source(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(length, inner, o, k, r)), 762 Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r), 763 Equal.cong(Nat, Nat, x => Slice.source(Slice.Plan{outer, extent, inner, start, step, length}, x), 764 Nat.add(Model.offset(depth, path), base), Laws.position(length, inner, o, k, r), at), 765 source_at(outer, extent, inner, start, step, length, o, k, r, k_below, r_below)) : 766 {Model.observe(F32, depth, Slice.build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, a), path) == 767 Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(_))) : Maybe<&2, F32>} 768 Equal.trans(Maybe<&2, F32>, 769 Model.observe(F32, depth, Slice.build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, a), path), 770 Some{Pair.fst(F32, Array & Array, Slice.cell(Slice.Plan{outer, extent, inner, start, step, length}, 771 Nat.add(Model.offset(depth, path), base), (A, empty()), True{}))}, 772 Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(Slice.source(Slice.Plan{outer, extent, inner, start, step, length}, 773 Nat.add(Model.offset(depth, path), base))))), 774 %inside : 775 {Model.observe(F32, depth, Slice.build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, a), path) == 776 Some{Pair.fst(F32, Array & Array, Slice.cell(Slice.Plan{outer, extent, inner, start, step, length}, 777 Nat.add(Model.offset(depth, path), base), (A, empty()), _))} : Maybe<&2, F32>} 778 {\{Parallel.Fact: q => q} : 779 Parallel.Fact, ~Array, ~Slice.cell, Slice.Plan{outer, extent, inner, start, step, length}, 780 (A, empty()), depth, count, base, Slice.build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, a))> -> 781 Generic.Cells(~Slice.Plan, ~Array, ~Array, ~Slice.cell, Slice.Plan{outer, extent, inner, start, step, length}, 782 (A, empty()), depth, count, base, Slice.build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, a))}( 783 build_cells(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, da, A, a, same, valid))(path), 784 {\{Parallel.Fact: q => q} : 785 Parallel.Fact<{Some{Pair.fst(F32, Array & Array, Slice.take(empty(), Array.get(F32, A, U32.from_nat(Slice.source(Slice.Plan{outer, extent, inner, start, step, length}, 786 Nat.add(Model.offset(depth, path), base))))))} == 787 Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(Slice.source(Slice.Plan{outer, extent, inner, start, step, length}, 788 Nat.add(Model.offset(depth, path), base))))) : Maybe<&2, F32>}> -> 789 {Some{Pair.fst(F32, Array & Array, Slice.take(empty(), Array.get(F32, A, U32.from_nat(Slice.source(Slice.Plan{outer, extent, inner, start, step, length}, 790 Nat.add(Model.offset(depth, path), base))))))} == 791 Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(Slice.source(Slice.Plan{outer, extent, inner, start, step, length}, 792 Nat.add(Model.offset(depth, path), base))))) : Maybe<&2, F32>}}( 793 read_word(da, A, empty(), Slice.source(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base)), valid)))} 794 795 # The backward target of input cell (o, i, r). 796 797 def picked_k(+outer: Nat, +extent: Nat, +inner: Nat, +start: Nat, +step: Nat, +length: Nat, +row: Nat, +within: Nat, +i: Nat, +k: Nat) -> 798 Maybe<&2, Nat>: 799 Slice.picked(Slice.Plan{outer, extent, inner, start, step, length}, row, within, i, k, Slice.hits(start, step, length, i, k)) 800 801 def picked_at(+outer: Nat, +extent: Nat, +inner: Nat, +start: Nat, +step: Nat, +length: Nat, +row: Nat, +within: Nat, +i: Nat) -> 802 Maybe<&2, Nat>: 803 picked_k(outer, extent, inner, start, step, length, row, within, i, Nat.div(Nat.sub(i, start), step)) 804 805 # `Slice.target` from the row and the remainder of the position. 806 def target_of(+outer: Nat, +extent: Nat, +inner: Nat, +start: Nat, +step: Nat, +length: Nat, +row: Nat, +within: Nat) -> Maybe<&2, Nat>: 807 picked_at(outer, extent, inner, start, step, length, row, within, Nat.mod(row, extent)) 808 809 # Position (o * extent + i) * inner + r splits into row o * extent + i and 810 # remainder r, and the row into o and i. 811 def target_split(+outer: Nat, +extent: Nat, +inner: Nat, +start: Nat, +step: Nat, +length: Nat, +o: Nat, +i: Nat, +r: Nat, 812 +i_below: {Nat.is_lt(i, extent) == True{} : Bool}, +r_below: {Nat.is_lt(r, inner) == True{} : Bool}) -> 813 {Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(extent, inner, o, i, r)) == 814 picked_at(outer, extent, inner, start, step, length, Nat.add(Nat.mul(o, extent), i), r, i) : Maybe<&2, Nat>}: 815 +row = Nat.add(Nat.mul(o, extent), i) 816 +pos = Laws.position(extent, inner, o, i, r) 817 %Equal.sym(Nat, Nat.div(pos, inner), row, BatchedProof.div_exact(row, inner, r, r_below)) : 818 {target_of(outer, extent, inner, start, step, length, _, Nat.mod(pos, inner)) == 819 picked_at(outer, extent, inner, start, step, length, row, r, i) : Maybe<&2, Nat>} 820 %Equal.sym(Nat, Nat.mod(pos, inner), r, BatchedProof.mod_exact(row, inner, r, r_below)) : 821 {target_of(outer, extent, inner, start, step, length, row, _) == picked_at(outer, extent, inner, start, step, length, row, r, i) : Maybe<&2, Nat>} 822 %Equal.sym(Nat, Nat.mod(row, extent), i, BatchedProof.mod_exact(o, extent, i, i_below)) : 823 {picked_at(outer, extent, inner, start, step, length, row, r, _) == picked_at(outer, extent, inner, start, step, length, row, r, i) : Maybe<&2, Nat>} 824 {==} 825 826 # Input cell (o, i, r) that no k below the length reaches reads nothing. 827 def missed_k(+outer: Nat, +extent: Nat, +inner: Nat, +start: Nat, +step: Nat, +length: Nat, +row: Nat, +r: Nat, +i: Nat, +k: Nat, 828 missed: @+k: Nat -> @below: {Nat.is_lt(k, length) == True{} : Bool} -> {Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i) == False{} : Bool}, 829 decision: Shape.BoolCase) -> 830 {picked_k(outer, extent, inner, start, step, length, row, r, i, k) == None{} : Maybe<&2, Nat>}: 831 match decision: 832 case Shape.BFalse{equal}: 833 %equal : {Slice.picked(Slice.Plan{outer, extent, inner, start, step, length}, row, r, i, k, _) == None{} : Maybe<&2, Nat>} 834 {==} 835 case Shape.BTrue{equal}: 836 +rest = Arith.and_right(Nat.is_le(start, i), Bool.and(Nat.is_lt(k, length), Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i)), 837 Equal.sym(Bool, True{}, Slice.hits(start, step, length, i, k), equal)) 838 Empty.absurd({picked_k(outer, extent, inner, start, step, length, row, r, i, k) == None{} : Maybe<&2, Nat>}, 839 Arith.clash(Equal.trans(Bool, False{}, Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i), True{}, 840 Equal.sym(Bool, Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i), False{}, 841 missed(k, Arith.and_left(Nat.is_lt(k, length), Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i), rest))), 842 Arith.and_right(Nat.is_lt(k, length), Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i), rest)))) 843 844 def target_miss(+outer: Nat, +extent: Nat, +inner: Nat, +start: Nat, +step: Nat, +length: Nat, +o: Nat, +i: Nat, +r: Nat, 845 +i_below: {Nat.is_lt(i, extent) == True{} : Bool}, +r_below: {Nat.is_lt(r, inner) == True{} : Bool}, 846 missed: @+k: Nat -> @below: {Nat.is_lt(k, length) == True{} : Bool} -> {Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i) == False{} : Bool}) -> 847 {Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(extent, inner, o, i, r)) == None{} : Maybe<&2, Nat>}: 848 Equal.trans(Maybe<&2, Nat>, Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(extent, inner, o, i, r)), 849 picked_at(outer, extent, inner, start, step, length, Nat.add(Nat.mul(o, extent), i), r, i), None{}, 850 target_split(outer, extent, inner, start, step, length, o, i, r, i_below, r_below), 851 missed_k(outer, extent, inner, start, step, length, Nat.add(Nat.mul(o, extent), i), r, i, Nat.div(Nat.sub(i, start), step), missed, 852 Shape.bool_case(Slice.hits(start, step, length, i, Nat.div(Nat.sub(i, start), step))))) 853 854 def ScatterCells(+depth: Nat, +splits: Nat, +plan: Slice.Plan, +count: Nat, +base: Nat, -G: Array, -g: Array) -> Type: 855 Generic.Cells(~Slice.Plan, ~Array, ~Array, ~Slice.scatter_cell, plan, (G, empty()), depth, count, base, 856 Slice.scatter_build(depth, splits, plan, count, base, g)) 857 858 # The slot of input cell (o, i, r), below the count, holds what the gradient 859 # target of the cell names. 860 def scatter_slot(+depth: Nat, +splits: Nat, +plan: Slice.Plan, +count: Nat, +base: Nat, +dg: Nat, -G: Array, g: Array, 861 -same: {g == G : Array}, -valid: Model.Balanced(F32, dg, G), +path: Model.Path(depth), 862 -inside: {Nat.is_lt(Nat.add(Model.offset(depth, path), base), count) == True{} : Bool}) -> 863 Parallel.Fact<{Model.observe(F32, depth, Slice.scatter_build(depth, splits, plan, count, base, g), path) == 864 Some{Pair.fst(F32, Array & Array, Slice.scattered((G, empty()), Slice.target(plan, Nat.add(Model.offset(depth, path), base))))} : 865 Maybe<&2, F32>}>: 866 Parallel.Fact{ 867 %inside : 868 {Model.observe(F32, depth, Slice.scatter_build(depth, splits, plan, count, base, g), path) == 869 Some{Pair.fst(F32, Array & Array, Slice.scatter_cell(plan, Nat.add(Model.offset(depth, path), base), (G, empty()), _))} : 870 Maybe<&2, F32>} 871 {\{Parallel.Fact: q => q} : Parallel.Fact -> ScatterCells(depth, splits, plan, count, base, G, g)}( 872 scatter_cells(depth, splits, plan, count, base, dg, G, g, same, valid))(path)} 873 874 def Laws.slice_scatter_miss(depth, splits, outer, extent, inner, start, step, length, count, base, dg, G, g, same, valid, path, o, i, r, 875 i_below, r_below, missed, at, inside): 876 Laws.Holds{ 877 Equal.trans(Maybe<&2, F32>, 878 Model.observe(F32, depth, Slice.scatter_build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, g), path), 879 Some{Pair.fst(F32, Array & Array, Slice.scattered((G, empty()), 880 Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base))))}, 881 Some{0.0}, 882 {\{Parallel.Fact: q => q} : 883 Parallel.Fact<{Model.observe(F32, depth, Slice.scatter_build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, g), path) == 884 Some{Pair.fst(F32, Array & Array, Slice.scattered((G, empty()), 885 Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base))))} : Maybe<&2, F32>}> -> 886 {Model.observe(F32, depth, Slice.scatter_build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, g), path) == 887 Some{Pair.fst(F32, Array & Array, Slice.scattered((G, empty()), 888 Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base))))} : Maybe<&2, F32>}}( 889 scatter_slot(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, dg, G, g, same, valid, path, inside)), 890 %Equal.sym(Maybe<&2, Nat>, Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base)), None{}, 891 Equal.trans(Maybe<&2, Nat>, Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base)), 892 Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(extent, inner, o, i, r)), None{}, 893 Equal.cong(Nat, Maybe<&2, Nat>, x => Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, x), 894 Nat.add(Model.offset(depth, path), base), Laws.position(extent, inner, o, i, r), at), 895 target_miss(outer, extent, inner, start, step, length, o, i, r, i_below, r_below, missed))) : 896 {Some{Pair.fst(F32, Array & Array, Slice.scattered((G, empty()), _))} == Some{0.0} : Maybe<&2, F32>} 897 {==})} 898 899 # start <= start + k * step, k < length and start + k * step is itself. 900 def hits_true(+start: Nat, +step: Nat, +length: Nat, +k: Nat, +k_below: {Nat.is_lt(k, length) == True{} : Bool}) -> 901 {Slice.hits(start, step, length, Nat.add(start, Nat.mul(k, step)), k) == True{} : Bool}: 902 %Equal.sym(Bool, Nat.is_le(start, Nat.add(start, Nat.mul(k, step))), True{}, Api.le_add_right(start, Nat.mul(k, step))) : 903 {Bool.and(_, Bool.and(Nat.is_lt(k, length), Nat.is_eq(Nat.add(start, Nat.mul(k, step)), Nat.add(start, Nat.mul(k, step))))) == True{} : Bool} 904 %Equal.sym(Bool, Nat.is_lt(k, length), True{}, k_below) : 905 {Bool.and(True{}, Bool.and(_, Nat.is_eq(Nat.add(start, Nat.mul(k, step)), Nat.add(start, Nat.mul(k, step))))) == True{} : Bool} 906 %Equal.sym(Bool, Nat.is_eq(Nat.add(start, Nat.mul(k, step)), Nat.add(start, Nat.mul(k, step))), True{}, 907 Arith.eq_refl(Nat.add(start, Nat.mul(k, step)))) : 908 {Bool.and(True{}, Bool.and(True{}, _)) == True{} : Bool} 909 {==} 910 911 # Input cell (o, start + k * step, r) reads gradient cell (o, k, r). 912 def target_hit(+outer: Nat, +extent: Nat, +inner: Nat, +start: Nat, +step: Nat, +length: Nat, +o: Nat, +k: Nat, +r: Nat, 913 +positive: {Nat.is_lt(0n, step) == True{} : Bool}, +k_below: {Nat.is_lt(k, length) == True{} : Bool}, 914 +i_below: {Nat.is_lt(Nat.add(start, Nat.mul(k, step)), extent) == True{} : Bool}, +r_below: {Nat.is_lt(r, inner) == True{} : Bool}) -> 915 {Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r)) == 916 Some{Laws.position(length, inner, o, k, r)} : Maybe<&2, Nat>}: 917 +i = Nat.add(start, Nat.mul(k, step)) 918 +row = Nat.add(Nat.mul(o, extent), i) 919 Equal.trans(Maybe<&2, Nat>, Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(extent, inner, o, i, r)), 920 picked_at(outer, extent, inner, start, step, length, row, r, i), Some{Laws.position(length, inner, o, k, r)}, 921 target_split(outer, extent, inner, start, step, length, o, i, r, i_below, r_below), 922 %Equal.sym(Nat, Nat.div(Nat.sub(i, start), step), k, 923 Equal.trans(Nat, Nat.div(Nat.sub(i, start), step), Nat.div(Nat.mul(k, step), step), k, 924 Equal.cong(Nat, Nat, x => Nat.div(x, step), Nat.sub(i, start), Nat.mul(k, step), sub_cancel(start, Nat.mul(k, step))), 925 div_step(k, step, positive))) : 926 {picked_k(outer, extent, inner, start, step, length, row, r, i, _) == Some{Laws.position(length, inner, o, k, r)} : Maybe<&2, Nat>} 927 %Equal.sym(Bool, Slice.hits(start, step, length, i, k), True{}, hits_true(start, step, length, k, k_below)) : 928 {Slice.picked(Slice.Plan{outer, extent, inner, start, step, length}, row, r, i, k, _) == Some{Laws.position(length, inner, o, k, r)} : 929 Maybe<&2, Nat>} 930 %Equal.sym(Nat, Nat.div(row, extent), o, BatchedProof.div_exact(o, extent, i, i_below)) : 931 {Some{Nat.add(Nat.mul(Nat.add(Nat.mul(_, length), k), inner), r)} == Some{Laws.position(length, inner, o, k, r)} : Maybe<&2, Nat>} 932 {==}) 933 934 def Laws.slice_scatter_hit(depth, splits, outer, extent, inner, start, step, length, count, base, dg, G, g, same, valid, path, o, k, r, 935 positive, k_below, i_below, r_below, at, inside): 936 Laws.Holds{ 937 Equal.trans(Maybe<&2, F32>, 938 Model.observe(F32, depth, Slice.scatter_build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, g), path), 939 Some{Pair.fst(F32, Array & Array, Slice.scattered((G, empty()), 940 Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base))))}, 941 Model.observe(F32, dg, G, Refinement.route(dg, U32.from_nat(Laws.position(length, inner, o, k, r)))), 942 {\{Parallel.Fact: q => q} : 943 Parallel.Fact<{Model.observe(F32, depth, Slice.scatter_build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, g), path) == 944 Some{Pair.fst(F32, Array & Array, Slice.scattered((G, empty()), 945 Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base))))} : Maybe<&2, F32>}> -> 946 {Model.observe(F32, depth, Slice.scatter_build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, g), path) == 947 Some{Pair.fst(F32, Array & Array, Slice.scattered((G, empty()), 948 Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base))))} : Maybe<&2, F32>}}( 949 scatter_slot(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, dg, G, g, same, valid, path, inside)), 950 %Equal.sym(Maybe<&2, Nat>, Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base)), 951 Some{Laws.position(length, inner, o, k, r)}, 952 Equal.trans(Maybe<&2, Nat>, Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base)), 953 Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r)), 954 Some{Laws.position(length, inner, o, k, r)}, 955 Equal.cong(Nat, Maybe<&2, Nat>, x => Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, x), 956 Nat.add(Model.offset(depth, path), base), Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r), at), 957 target_hit(outer, extent, inner, start, step, length, o, k, r, positive, k_below, i_below, r_below))) : 958 {Some{Pair.fst(F32, Array & Array, Slice.scattered((G, empty()), _))} == 959 Model.observe(F32, dg, G, Refinement.route(dg, U32.from_nat(Laws.position(length, inner, o, k, r)))) : Maybe<&2, F32>} 960 {\{Parallel.Fact: q => q} : 961 Parallel.Fact<{Some{Pair.fst(F32, Array & Array, Slice.take(empty(), Array.get(F32, G, U32.from_nat(Laws.position(length, inner, o, k, r)))))} == 962 Model.observe(F32, dg, G, Refinement.route(dg, U32.from_nat(Laws.position(length, inner, o, k, r)))) : Maybe<&2, F32>}> -> 963 {Some{Pair.fst(F32, Array & Array, Slice.take(empty(), Array.get(F32, G, U32.from_nat(Laws.position(length, inner, o, k, r)))))} == 964 Model.observe(F32, dg, G, Refinement.route(dg, U32.from_nat(Laws.position(length, inner, o, k, r)))) : Maybe<&2, F32>}}( 965 read_word(dg, G, empty(), Laws.position(length, inner, o, k, r), valid)))}