1 import Base 2 import ../shape/core.bend as Shape 3 import ../storage/model.bend as Model 4 import ../storage/refinement.bend as Refinement 5 import ./core.bend as Core 6 import ./layout.bend as Layout 7 import ./transforms.bend as Transforms 8 import ./elementwise_totality.bend as Totality 9 import ./slice.bend as Slice 10 11 # The laws of basic slicing (slice.bend), written from PyTorch's definition 12 # of a slice: which start and stop it clamps to, how long the sliced axis 13 # is, how the row-major view [outer, extent, inner] sits around the axis, 14 # which calls it refuses, and what every cell of the forward copy and of the 15 # backward scatter holds. 16 17 type Holds<-P: Type> is Data: 18 Holds{-proof: P} 19 20 def low(+clamped: Slice.Bounds) -> Nat: 21 Slice.Bounds{start, stop} = clamped 22 start 23 24 def high(+clamped: Slice.Bounds) -> Nat: 25 Slice.Bounds{start, stop} = clamped 26 stop 27 28 # PyTorch's clamp of a slice to an axis of extent n. A start at most n is 29 # kept, a start past n becomes n. 30 law slice_start_kept: 31 for +extent: Nat 32 for +start: Nat 33 for +stop: Nat 34 for inside: {Nat.is_le(start, extent) == True{} : Bool} 35 {low(Slice.bounds(extent, start, stop)) == start : Nat} 36 37 law slice_start_clamped: 38 for +extent: Nat 39 for +start: Nat 40 for +stop: Nat 41 for past: {Nat.is_le(extent, start) == True{} : Bool} 42 {low(Slice.bounds(extent, start, stop)) == extent : Nat} 43 44 # A stop from the start to n is kept, a stop past n becomes n, and a stop 45 # at or before the start ends the slice at its start: an empty axis. 46 law slice_stop_kept: 47 for +extent: Nat 48 for +start: Nat 49 for +stop: Nat 50 for +after: {Nat.is_le(start, stop) == True{} : Bool} 51 for +inside: {Nat.is_le(stop, extent) == True{} : Bool} 52 {high(Slice.bounds(extent, start, stop)) == stop : Nat} 53 54 law slice_stop_clamped: 55 for +extent: Nat 56 for +start: Nat 57 for +stop: Nat 58 for past: {Nat.is_le(extent, stop) == True{} : Bool} 59 {high(Slice.bounds(extent, start, stop)) == extent : Nat} 60 61 law slice_stop_empty: 62 for +extent: Nat 63 for +start: Nat 64 for +stop: Nat 65 for before: {Nat.is_le(stop, start) == True{} : Bool} 66 {high(Slice.bounds(extent, start, stop)) == low(Slice.bounds(extent, start, stop)) : Nat} 67 68 # The length of the sliced axis is the number of indices start, start + 69 # step, start + 2 step, ... below the stop: each index below the length 70 # lies below the stop, and, for a positive step, the index at the length 71 # does not. 72 law slice_length_below: 73 for +start: Nat 74 for +stop: Nat 75 for +step: Nat 76 for +k: Nat 77 for below: {Nat.is_lt(k, Slice.length(Slice.Bounds{start, stop}, step)) == True{} : Bool} 78 {Nat.is_lt(Nat.add(start, Nat.mul(k, step)), stop) == True{} : Bool} 79 80 law slice_length_reaches: 81 for +start: Nat 82 for +stop: Nat 83 for +step: Nat 84 for positive: {Nat.is_lt(0n, step) == True{} : Bool} 85 {Nat.is_le(stop, Nat.add(start, Nat.mul(Slice.length(Slice.Bounds{start, stop}, step), step))) == True{} : Bool} 86 87 # The row-major view around axis `dim` of dimensions prefix ++ [extent] ++ 88 # suffix, with `dim` the rank of the prefix: the element count of the 89 # prefix, the extent, and the element count of the suffix. 90 law slice_view_around: 91 for +prefix: List<&2, U32> 92 for +extent: U32 93 for +suffix: List<&2, U32> 94 {Slice.view(List.append(&2, U32, prefix, extent <> suffix), Shape.rank_dimensions(prefix)) == 95 Slice.View{Shape.logical_numel(prefix), U32.to_nat(extent), Shape.logical_numel(suffix)} : Slice.View} 96 97 # The sliced output keeps every other axis and takes the new extent at the 98 # axis; `select` drops the axis. 99 law slice_resized_around: 100 for +prefix: List<&2, U32> 101 for +extent: U32 102 for +suffix: List<&2, U32> 103 for +other: U32 104 {Slice.resized(List.append(&2, U32, prefix, extent <> suffix), Shape.rank_dimensions(prefix), other) == 105 List.append(&2, U32, prefix, other <> suffix) : List<&2, U32>} 106 107 law select_dropped_around: 108 for +prefix: List<&2, U32> 109 for +extent: U32 110 for +suffix: List<&2, U32> 111 {Slice.dropped(List.append(&2, U32, prefix, extent <> suffix), Shape.rank_dimensions(prefix)) == 112 List.append(&2, U32, prefix, suffix) : List<&2, U32>} 113 114 # What a plan holds along its axis (the extent, the clamped start, the step 115 # and the length) and around it (the element counts before and after it). 116 def axis_part(+plan: Slice.Plan) -> Nat & Nat & Nat & Nat: 117 Slice.Plan{outer, extent, inner, start, step, length} = plan 118 (extent, start, step, length) 119 120 def view_part(+plan: Slice.Plan) -> Nat & Nat: 121 Slice.Plan{outer, extent, inner, start, step, length} = plan 122 (outer, inner) 123 124 # PyTorch's clamp of a slice position to an axis of extent n: the position 125 # if it is at most n, else n. 126 def clamp(+position: Nat, +extent: Nat) -> Nat: 127 Bool.pick(Nat, Nat.is_le(position, extent), position, extent) 128 129 # The stop after the clamp: the clamped stop, or the clamped start where the 130 # clamped stop falls before it, so the axis is empty. 131 def clamped_stop(+start: Nat, +stop: Nat, +extent: Nat) -> Nat: 132 Bool.pick(Nat, Nat.is_le(clamp(start, extent), clamp(stop, extent)), clamp(stop, extent), clamp(start, extent)) 133 134 # What a request start:stop:step plans along an axis of extent n: the 135 # extent, the clamped start, the step, and the length of the clamped range. 136 def axis_spec(+extent: Nat, +start: Nat, +stop: Nat, +step: Nat) -> Nat & Nat & Nat & Nat: 137 (extent, clamp(start, extent), step, 138 Slice.length(Slice.Bounds{clamp(start, extent), clamped_stop(start, stop, extent)}, step)) 139 140 # The plan of a request at axis `dim` of prefix ++ [extent] ++ suffix, with 141 # `dim` the rank of the prefix, holds the request's own start, stop and step, 142 # clamped to that extent, whatever the other extents, an empty input's 143 # included. 144 law slice_plan_axis: 145 for +prefix: List<&2, U32> 146 for +extent: U32 147 for +suffix: List<&2, U32> 148 for +dim: U32 149 for at: {U32.to_nat(dim) == Shape.rank_dimensions(prefix) : Nat} 150 for +start: U32 151 for +stop: U32 152 for +step: U32 153 {axis_part(Slice.plan_of(List.append(&2, U32, prefix, extent <> suffix), dim, start, stop, step)) == 154 axis_spec(U32.to_nat(extent), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step)) : Nat & Nat & Nat & Nat} 155 156 # For an input with an element, the plan's view is the element counts of the 157 # prefix and of the suffix, the row-major view around the axis. 158 law slice_plan_view: 159 for +prefix: List<&2, U32> 160 for +extent: U32 161 for +suffix: List<&2, U32> 162 for +dim: U32 163 for at: {U32.to_nat(dim) == Shape.rank_dimensions(prefix) : Nat} 164 for +start: U32 165 for +stop: U32 166 for +step: U32 167 for nonempty: {Nat.is_lt(0n, Shape.logical_numel(List.append(&2, U32, prefix, extent <> suffix))) == True{} : Bool} 168 {view_part(Slice.plan_of(List.append(&2, U32, prefix, extent <> suffix), dim, start, stop, step)) == 169 (Shape.logical_numel(prefix), Shape.logical_numel(suffix)) : Nat & Nat} 170 171 # An input with no element gets a plan with no view: neither build has a 172 # cell then, so the plan forms no product of the other extents, which can 173 # pass the largest Nat the runtime holds although no element exists. 174 law slice_plan_empty_view: 175 for +dimensions: List<&2, U32> 176 for +dim: U32 177 for +start: U32 178 for +stop: U32 179 for +step: U32 180 for empty: {Shape.logical_numel(dimensions) == 0n : Nat} 181 {view_part(Slice.plan_of(dimensions, dim, start, stop, step)) == (0n, 0n) : Nat & Nat} 182 183 # PyTorch's checks, in its order. An axis at or past the rank (every axis 184 # of a rank-0 input) is refused first, then a step of 0. 185 law slice_axis_refused: 186 for +input: Shape.Shape 187 for +dim: U32 188 for +start: U32 189 for +stop: U32 190 for +step: U32 191 for outside: {Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(input))) == False{} : Bool} 192 {Slice.plan(input, dim, start, stop, step) == Fail{Slice.LayoutError{Layout.AxisOutOfBounds{}}} : 193 Result<&2, &2, Slice.Error, Slice.Planned>} 194 195 law slice_step_refused: 196 for +input: Shape.Shape 197 for +dim: U32 198 for +start: U32 199 for +stop: U32 200 for +step: U32 201 for inside: {Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(input))) == True{} : Bool} 202 for zero: {U32.to_nat(step) == 0n : Nat} 203 {Slice.plan(input, dim, start, stop, step) == Fail{Slice.StepNotPositive{}} : Result<&2, &2, Slice.Error, Slice.Planned>} 204 205 # What an accepted slice plans: the view around the axis, the start and 206 # stop clamped to the extent, the step, the length of the sliced axis, and 207 # the output's dimensions, the input's with that length at the axis. 208 def PlannedAs(+dimensions: List<&2, U32>, +dim: U32, +start: U32, +stop: U32, +step: U32, 209 result: Result<&2, &2, Slice.Error, Slice.Planned>) -> Type: 210 match result: 211 case Fail{error}: 212 Unit 213 case Done{Slice.Planned{plan, output}}: 214 {plan == Slice.plan_of(dimensions, dim, start, stop, step) : Slice.Plan} & 215 {Shape.dimensions(output) == Slice.resized(dimensions, U32.to_nat(dim), U32.from_nat(Slice.plan_length(plan))) : List<&2, U32>} 216 217 law slice_planned: 218 for +input: Shape.Shape 219 for +dim: U32 220 for +start: U32 221 for +stop: U32 222 for +step: U32 223 PlannedAs(Shape.dimensions(input), dim, start, stop, step, Slice.plan(input, dim, start, stop, step)) 224 225 # A refused plan fails the slice with its error and hands the input back. 226 law slice_refused: 227 for input: Core.Dynamic 228 for +dim: U32 229 for +start: U32 230 for +stop: U32 231 for +step: U32 232 for error: Slice.Error 233 for refused: {Slice.plan(Totality.dynamic_shape(input), dim, start, stop, step) == Fail{error} : Result<&2, &2, Slice.Error, Slice.Planned>} 234 {Slice.slice(input, dim, start, stop, step) == (input, Fail{error}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 235 236 def rank_of(input: Core.Dynamic) -> Nat: 237 Shape.rank_dimensions(Shape.dimensions(Totality.dynamic_shape(input))) 238 239 def extent_of(input: Core.Dynamic, +dim: U32) -> Nat: 240 Slice.extent_at(Shape.dimensions(Totality.dynamic_shape(input)), dim) 241 242 # PyTorch's narrow(dim, start, length) is the slice start:start + length 243 # with step 1, once the axis is in range and start + length is at most the 244 # extent; a start + length past the extent is refused with OutOfRange, an 245 # axis past the rank as a slice refuses it. 246 law narrow_slices: 247 for input: Core.Dynamic 248 for +dim: U32 249 for +start: U32 250 for +length: U32 251 for inside: {Nat.is_lt(U32.to_nat(dim), rank_of(input)) == True{} : Bool} 252 for fits: {Nat.is_le(Nat.add(U32.to_nat(start), U32.to_nat(length)), extent_of(input, dim)) == True{} : Bool} 253 {Slice.narrow(input, dim, start, length) == Slice.slice(input, dim, start, U32.add(start, length), 1) : 254 Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 255 256 law narrow_out_of_range: 257 for input: Core.Dynamic 258 for +dim: U32 259 for +start: U32 260 for +length: U32 261 for inside: {Nat.is_lt(U32.to_nat(dim), rank_of(input)) == True{} : Bool} 262 for past: {Nat.is_le(Nat.add(U32.to_nat(start), U32.to_nat(length)), extent_of(input, dim)) == False{} : Bool} 263 {Slice.narrow(input, dim, start, length) == (input, Fail{Slice.OutOfRange{}}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 264 265 law narrow_axis_refused: 266 for input: Core.Dynamic 267 for +dim: U32 268 for +start: U32 269 for +length: U32 270 for outside: {Nat.is_lt(U32.to_nat(dim), rank_of(input)) == False{} : Bool} 271 {Slice.narrow(input, dim, start, length) == (input, Fail{Slice.LayoutError{Layout.AxisOutOfBounds{}}}) : 272 Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 273 274 # A slice's result with its output reshaped to `dimensions`. 275 def reshape_as_slice(result: Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Result<&2, &1, Slice.Error, Core.Dynamic>: 276 match result: 277 case Fail{Transforms.LayoutError{error}}: 278 Fail{Slice.LayoutError{error}} 279 case Fail{Transforms.TensorError{error}}: 280 Fail{Slice.TensorError{error}} 281 case Done{output}: 282 Done{output} 283 284 def reshaped(+dimensions: List<&2, U32>, pair: Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>) -> 285 Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>: 286 (input, result) = pair 287 match result: 288 case Fail{error}: 289 (input, Fail{error}) 290 case Done{output}: 291 (input, reshape_as_slice(Transforms.reshape_dynamic(output, dimensions))) 292 293 # PyTorch's select(dim, index) is the slice index:index + 1 with step 1, 294 # reshaped to the input's dimensions without the axis, once the axis is in 295 # range and the index below the extent; an index at or past the extent is 296 # refused with OutOfRange, an axis past the rank as a slice refuses it. 297 law select_slices: 298 for input: Core.Dynamic 299 for +dim: U32 300 for +index: U32 301 for inside: {Nat.is_lt(U32.to_nat(dim), rank_of(input)) == True{} : Bool} 302 for below: {Nat.is_lt(U32.to_nat(index), extent_of(input, dim)) == True{} : Bool} 303 {Slice.select(input, dim, index) == 304 reshaped(Slice.dropped(Shape.dimensions(Totality.dynamic_shape(input)), U32.to_nat(dim)), Slice.slice(input, dim, index, U32.add(index, 1), 1)) : 305 Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 306 307 law select_out_of_range: 308 for input: Core.Dynamic 309 for +dim: U32 310 for +index: U32 311 for inside: {Nat.is_lt(U32.to_nat(dim), rank_of(input)) == True{} : Bool} 312 for past: {Nat.is_lt(U32.to_nat(index), extent_of(input, dim)) == False{} : Bool} 313 {Slice.select(input, dim, index) == (input, Fail{Slice.OutOfRange{}}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 314 315 law select_axis_refused: 316 for input: Core.Dynamic 317 for +dim: U32 318 for +index: U32 319 for outside: {Nat.is_lt(U32.to_nat(dim), rank_of(input)) == False{} : Bool} 320 {Slice.select(input, dim, index) == (input, Fail{Slice.LayoutError{Layout.AxisOutOfBounds{}}}) : 321 Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>} 322 323 # The cells of the builds, in the coordinates of the views: output cell 324 # (o, k, r) of [outer, length, inner] sits at row-major position 325 # (o * length + k) * inner + r, and input cell (o, i, r) of [outer, extent, 326 # inner] at (o * extent + i) * inner + r. 327 def position(+rows: Nat, +inner: Nat, +o: Nat, +i: Nat, +r: Nat) -> Nat: 328 Nat.add(Nat.mul(Nat.add(Nat.mul(o, rows), i), inner), r) 329 330 # For a balanced input array A and every split, the slot of the forward 331 # build that holds output cell (o, k, r), below the count, holds input cell 332 # (o, start + k * step, r): the word PyTorch's slice puts there. 333 law slice_build_cells: 334 for +depth: Nat 335 for +splits: Nat 336 for +outer: Nat 337 for +extent: Nat 338 for +inner: Nat 339 for +start: Nat 340 for +step: Nat 341 for +length: Nat 342 for +count: Nat 343 for +base: Nat 344 for +da: Nat 345 for -A: Array 346 for a: Array 347 for same: {a == A : Array} 348 for valid: Model.Balanced(F32, da, A) 349 for +path: Model.Path(depth) 350 for +o: Nat 351 for +k: Nat 352 for +r: Nat 353 for k_below: {Nat.is_lt(k, length) == True{} : Bool} 354 for r_below: {Nat.is_lt(r, inner) == True{} : Bool} 355 for at: {Nat.add(Model.offset(depth, path), base) == position(length, inner, o, k, r) : Nat} 356 for inside: {Nat.is_lt(Nat.add(Model.offset(depth, path), base), count) == True{} : Bool} 357 Holds<{Model.observe(F32, depth, Slice.build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, a), path) == 358 Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r)))) : 359 Maybe<&2, F32>}> 360 361 # For a balanced gradient array G and every split, the slot of the backward 362 # build that holds input cell (o, start + k * step, r), with k below the 363 # length and a positive step, holds gradient cell (o, k, r) ... 364 law slice_scatter_hit: 365 for +depth: Nat 366 for +splits: Nat 367 for +outer: Nat 368 for +extent: Nat 369 for +inner: Nat 370 for +start: Nat 371 for +step: Nat 372 for +length: Nat 373 for +count: Nat 374 for +base: Nat 375 for +dg: Nat 376 for -G: Array 377 for g: Array 378 for same: {g == G : Array} 379 for valid: Model.Balanced(F32, dg, G) 380 for +path: Model.Path(depth) 381 for +o: Nat 382 for +k: Nat 383 for +r: Nat 384 for positive: {Nat.is_lt(0n, step) == True{} : Bool} 385 for k_below: {Nat.is_lt(k, length) == True{} : Bool} 386 for i_below: {Nat.is_lt(Nat.add(start, Nat.mul(k, step)), extent) == True{} : Bool} 387 for r_below: {Nat.is_lt(r, inner) == True{} : Bool} 388 for at: {Nat.add(Model.offset(depth, path), base) == position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r) : Nat} 389 for inside: {Nat.is_lt(Nat.add(Model.offset(depth, path), base), count) == True{} : Bool} 390 Holds<{Model.observe(F32, depth, Slice.scatter_build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, g), path) == 391 Model.observe(F32, dg, G, Refinement.route(dg, U32.from_nat(position(length, inner, o, k, r)))) : Maybe<&2, F32>}> 392 393 # ... and the slot of input cell (o, i, r) where no k below the length has 394 # i = start + k * step holds +0.0, as PyTorch's slice_backward leaves it. 395 law slice_scatter_miss: 396 for +depth: Nat 397 for +splits: Nat 398 for +outer: Nat 399 for +extent: Nat 400 for +inner: Nat 401 for +start: Nat 402 for +step: Nat 403 for +length: Nat 404 for +count: Nat 405 for +base: Nat 406 for +dg: Nat 407 for -G: Array 408 for g: Array 409 for same: {g == G : Array} 410 for valid: Model.Balanced(F32, dg, G) 411 for +path: Model.Path(depth) 412 for +o: Nat 413 for +i: Nat 414 for +r: Nat 415 for i_below: {Nat.is_lt(i, extent) == True{} : Bool} 416 for r_below: {Nat.is_lt(r, inner) == True{} : Bool} 417 for missed: @+k: Nat -> @below: {Nat.is_lt(k, length) == True{} : Bool} -> 418 {Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i) == False{} : Bool} 419 for at: {Nat.add(Model.offset(depth, path), base) == position(extent, inner, o, i, r) : Nat} 420 for inside: {Nat.is_lt(Nat.add(Model.offset(depth, path), base), count) == True{} : Bool} 421 Holds<{Model.observe(F32, depth, Slice.scatter_build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, g), path) == 422 Some{0.0} : Maybe<&2, F32>}>