1 import Base 2 import ../shape/core.bend as Shape 3 import ../storage/main.bend as Storage 4 import ../storage/model.bend as Model 5 import ./core.bend as Core 6 import ./elementwise.bend as Elementwise 7 import ./matmul.bend as Matrix 8 import ./parallel.bend as Parallel 9 10 # PyTorch's `torch.matmul`: batched products with 1-D promotion. 11 # 12 # Two rank-2 operands are `Matrix.matmul` itself, the same call, so their 13 # output words, errors and laws are those of `matmul`. Every other pair of 14 # ranks from 1 to 4 is planned here, as `torch.matmul` plans it: 15 # 16 # - A 1-D left operand [k] is the row [1, k], and the output drops that row; 17 # a 1-D right operand [k] is the column [k, 1], and the output drops that 18 # column. So [k] @ [k] is a rank-0 dot product. 19 # - Every dimension but the last two of an operand is a batch dimension. 20 # The batch dimensions of the two operands broadcast as elementwise 21 # operations do: right-aligned, an extent of 1 stretches. The output is 22 # the broadcast batch dimensions, then [m, n] with a dropped row or column 23 # left out. 24 # - A rank-0 operand fails with RankMismatch (PyTorch: "both arguments to 25 # matmul need to be at least 1D"), unequal inner extents with 26 # InnerDimensionMismatch, batch dimensions that do not broadcast with 27 # ShapeError{IncompatibleBroadcast}, and an output past the capacity with 28 # ShapeError{CapacityExceeded}. 29 # 30 # Output cell (b, row, column) adds left[b, row, k] * right[b, k, column] 31 # for k = 0, 1, ..., inner - 1, starting from +0.0: the terms and the order 32 # of `Matrix.matmul`'s cell (row, column) on the two matrices of batch b, so 33 # each batch gets the words `matmul` writes for its matrices. A broadcast 34 # operand is read in place at its own batch, never copied. The cells are 35 # built by `Parallel.build`, so no word depends on the split or the thread 36 # count. 37 38 # The batch extents of an operand, padded on the left with ones to two, and 39 # how many it has. 40 type Batch is Data: 41 Batch{rank: Nat, first: U32, second: U32} 42 43 # An operand split into its batch and its matrix: the outer extent (rows of 44 # the left operand, columns of the right one) and the inner extent. A 1-D 45 # operand is a vector: an outer extent of 1 that the output drops. 46 type Operand is Data: 47 Operand{batch: Batch, outer: U32, inner: U32, vector: Bool} 48 49 def left_operand(dims: List<&2, U32>) -> Result<&2, &2, Matrix.Error, Operand>: 50 match dims: 51 case +k <> Nil{}: 52 Done{Operand{Batch{0n, 1, 1}, 1, k, True{}}} 53 case +m <> +k <> Nil{}: 54 Done{Operand{Batch{0n, 1, 1}, m, k, False{}}} 55 case +b <> +m <> +k <> Nil{}: 56 Done{Operand{Batch{1n, 1, b}, m, k, False{}}} 57 case +a <> +b <> +m <> +k <> Nil{}: 58 Done{Operand{Batch{2n, a, b}, m, k, False{}}} 59 case rest: 60 Fail{Matrix.RankMismatch{}} 61 62 def right_operand(dims: List<&2, U32>) -> Result<&2, &2, Matrix.Error, Operand>: 63 match dims: 64 case +k <> Nil{}: 65 Done{Operand{Batch{0n, 1, 1}, 1, k, True{}}} 66 case +k <> +n <> Nil{}: 67 Done{Operand{Batch{0n, 1, 1}, n, k, False{}}} 68 case +b <> +k <> +n <> Nil{}: 69 Done{Operand{Batch{1n, 1, b}, n, k, False{}}} 70 case +a <> +b <> +k <> +n <> Nil{}: 71 Done{Operand{Batch{2n, a, b}, n, k, False{}}} 72 case rest: 73 Fail{Matrix.RankMismatch{}} 74 75 # One batch axis: equal extents, or an extent of 1 that stretches to the 76 # other one. 77 def stretched(+l: U32, +r: U32, same: Bool, left_one: Bool, right_one: Bool) -> Result<&2, &2, Matrix.Error, U32>: 78 match same left_one right_one: 79 case True{} _ _: 80 Done{l} 81 case False{} True{} _: 82 Done{r} 83 case False{} False{} True{}: 84 Done{l} 85 case False{} False{} False{}: 86 Fail{Matrix.ShapeError{Shape.IncompatibleBroadcast{}}} 87 88 def axis(+l: U32, +r: U32) -> Result<&2, &2, Matrix.Error, U32>: 89 stretched(l, r, Nat.is_eq(U32.to_nat(l), U32.to_nat(r)), Nat.is_eq(U32.to_nat(l), 1n), Nat.is_eq(U32.to_nat(r), 1n)) 90 91 def larger(+a: Nat, +b: Nat) -> Nat: 92 Matrix.least(b, a, Nat.is_lt(a, b)) 93 94 # The output's batch dimensions: as many as the operand with more of them. 95 def batch_dims(+rank: Nat, +first: U32, +second: U32) -> List<&2, U32>: 96 match rank: 97 case 0n: 98 [] 99 case 1n+rest: 100 match rest: 101 case 0n: 102 [second] 103 case 1n+more: 104 [first, second] 105 106 def column_dims(+columns: U32, right_vector: Bool) -> List<&2, U32>: 107 match right_vector: 108 case True{}: 109 [] 110 case False{}: 111 [columns] 112 113 def matrix_dims(+rows: U32, left_vector: Bool, +columns: U32, right_vector: Bool) -> List<&2, U32>: 114 match left_vector: 115 case True{}: 116 column_dims(columns, right_vector) 117 case False{}: 118 rows <> column_dims(columns, right_vector) 119 120 # Everything a product and its two gradients need: the output shape, the 121 # output's batch dimensions, the matrix extents, both operands' padded 122 # batches and the output's, and which operands are vectors. 123 type Plan is Data: 124 Plan{output: Shape.Shape, batch: List<&2, U32>, rows: Nat, inner: Nat, columns: Nat, 125 l1: Nat, l2: Nat, r1: Nat, r2: Nat, o1: Nat, o2: Nat, left_vector: Bool, right_vector: Bool} 126 127 def output_of(+plan: Plan) -> Shape.Shape: 128 Plan{+output, batch, rows, inner, columns, l1, l2, r1, r2, o1, o2, lv, rv} = plan 129 output 130 131 def planned_shape(+batch: List<&2, U32>, +rows: U32, +inner: U32, +columns: U32, 132 +l1: U32, +l2: U32, +r1: U32, +r2: U32, +o1: U32, +o2: U32, +lv: Bool, +rv: Bool, 133 -dims: List<&2, U32>, result: Shape.ConstructResult) -> Result<&2, &2, Matrix.Error, Plan>: 134 match result: 135 case Shape.SError{error}: 136 Fail{Matrix.ShapeError{error}} 137 case Shape.SBuilt{+shape, preserved}: 138 Done{Plan{shape, batch, U32.to_nat(rows), U32.to_nat(inner), U32.to_nat(columns), 139 U32.to_nat(l1), U32.to_nat(l2), U32.to_nat(r1), U32.to_nat(r2), U32.to_nat(o1), U32.to_nat(o2), lv, rv}} 140 141 def batched(+rank: Nat, +rows: U32, +inner: U32, +columns: U32, +l1: U32, +l2: U32, +r1: U32, +r2: U32, 142 +lv: Bool, +rv: Bool, first: Result<&2, &2, Matrix.Error, U32>, 143 second: Result<&2, &2, Matrix.Error, U32>) -> Result<&2, &2, Matrix.Error, Plan>: 144 match first second: 145 case Fail{error} second: 146 Fail{error} 147 case Done{o1} Fail{error}: 148 Fail{error} 149 case Done{+o1} Done{+o2}: 150 planned_shape(batch_dims(rank, o1, o2), rows, inner, columns, l1, l2, r1, r2, o1, o2, lv, rv, 151 List.append(&2, U32, batch_dims(rank, o1, o2), matrix_dims(rows, lv, columns, rv)), 152 Shape.construct(List.append(&2, U32, batch_dims(rank, o1, o2), matrix_dims(rows, lv, columns, rv)))) 153 154 def shaped(+left: Operand, +right: Operand) -> Result<&2, &2, Matrix.Error, Plan>: 155 Operand{Batch{+lr, +l1, +l2}, +rows, +inner, +lv} = left 156 Operand{Batch{+rr, +r1, +r2}, +columns, right_inner, +rv} = right 157 batched(larger(lr, rr), rows, inner, columns, l1, l2, r1, r2, lv, rv, axis(l1, r1), axis(l2, r2)) 158 159 def compatible(+left: Operand, +right: Operand, same: Bool) -> Result<&2, &2, Matrix.Error, Plan>: 160 match same: 161 case False{}: 162 Fail{Matrix.InnerDimensionMismatch{}} 163 case True{}: 164 shaped(left, right) 165 166 def inner_of(+operand: Operand) -> U32: 167 Operand{batch, outer, +inner, vector} = operand 168 inner 169 170 def both_parsed(+left: Operand, result: Result<&2, &2, Matrix.Error, Operand>) -> Result<&2, &2, Matrix.Error, Plan>: 171 match result: 172 case Fail{error}: 173 Fail{error} 174 case Done{+right}: 175 compatible(left, right, Nat.is_eq(U32.to_nat(inner_of(left)), U32.to_nat(inner_of(right)))) 176 177 def left_parsed(right: List<&2, U32>, result: Result<&2, &2, Matrix.Error, Operand>) -> Result<&2, &2, Matrix.Error, Plan>: 178 match result: 179 case Fail{error}: 180 Fail{error} 181 case Done{+left}: 182 both_parsed(left, right_operand(right)) 183 184 # The plan of left @ right for operands of dimensions `left` and `right`, 185 # each of rank 1 to 4: the left operand's rank is checked first, then the 186 # right one's, the inner extents, the batch dimensions and the output's 187 # capacity. 188 def plan_dims(left: List<&2, U32>, right: List<&2, U32>) -> Result<&2, &2, Matrix.Error, Plan>: 189 left_parsed(right, left_operand(left)) 190 191 def plan(+left: Shape.Shape, +right: Shape.Shape) -> Result<&2, &2, Matrix.Error, Plan>: 192 plan_dims(Shape.dimensions(left), Shape.dimensions(right)) 193 194 # What every cell of a build reads: the orders and extents of the product 195 # (as in `Matrix.Layout`), and each operand's padded batch extents with the 196 # output's second batch extent, which splits a batch number into its two 197 # indices. 198 type Layout is Data: 199 Layout{left: Matrix.Order, right: Matrix.Order, inner: Nat, rows: Nat, columns: Nat, 200 l1: Nat, l2: Nat, r1: Nat, r2: Nat, o2: Nat} 201 202 def layout(+given: Layout) -> Layout: 203 Layout{+lo, +ro, +inner, +rows, +columns, +l1, +l2, +r1, +r2, +o2} = given 204 Layout{lo, ro, inner, rows, columns, l1, l2, r1, r2, o2} 205 206 # left @ right: left [.., rows, inner] read by rows, right [.., inner, 207 # columns] read by columns, as `matmul` reads them. 208 def forward_layout(+plan: Plan) -> Layout: 209 Plan{output, batch, +rows, +inner, +columns, +l1, +l2, +r1, +r2, o1, +o2, lv, rv} = plan 210 Layout{Matrix.Kept{}, Matrix.Swapped{}, inner, rows, columns, l1, l2, r1, r2, o2} 211 212 # g @ right^T, the left operand's gradient over the output's batch: g 213 # [.., rows, columns] by rows, right [.., inner, columns] by its rows, the 214 # columns being the terms. 215 def left_layout(+plan: Plan) -> Layout: 216 Plan{output, batch, +rows, +inner, +columns, l1, l2, +r1, +r2, +o1, +o2, lv, rv} = plan 217 Layout{Matrix.Kept{}, Matrix.Kept{}, columns, rows, inner, o1, o2, r1, r2, o2} 218 219 # left^T @ g, the right operand's gradient over the output's batch: left 220 # [.., rows, inner] by its columns, g [.., rows, columns] by columns, the 221 # rows being the terms. 222 def right_layout(+plan: Plan) -> Layout: 223 Plan{output, batch, +rows, +inner, +columns, +l1, +l2, r1, r2, +o1, +o2, lv, rv} = plan 224 Layout{Matrix.Swapped{}, Matrix.Swapped{}, rows, inner, columns, l1, l2, o1, o2, o2} 225 226 # The dimensions of each operand's gradient before it is reduced to the 227 # operand's shape: the output's batch, then the operand's own matrix 228 # dimensions. 229 def left_dims(+plan: Plan) -> List<&2, U32>: 230 Plan{output, +batch, +rows, +inner, columns, l1, l2, r1, r2, o1, o2, +lv, rv} = plan 231 List.append(&2, U32, batch, matrix_dims(U32.from_nat(rows), lv, U32.from_nat(inner), False{})) 232 233 def right_dims(+plan: Plan) -> List<&2, U32>: 234 Plan{output, +batch, rows, +inner, +columns, l1, l2, r1, r2, o1, o2, lv, +rv} = plan 235 List.append(&2, U32, batch, matrix_dims(U32.from_nat(inner), False{}, U32.from_nat(columns), rv)) 236 237 # Index `index` of a batch axis of `extent`: the index itself, or 0 where the 238 # extent is 1 and stretches. 239 def stretch(+index: Nat, one: Bool) -> Nat: 240 match one: 241 case True{}: 242 0n 243 case False{}: 244 index 245 246 def batch_number(+first: Nat, +second: Nat, +b1: Nat, +b2: Nat) -> Nat: 247 Nat.add(Nat.mul(stretch(b1, Nat.is_eq(first, 1n)), second), stretch(b2, Nat.is_eq(second, 1n))) 248 249 # The `fuel` terms of a cell from offsets `lp` and `rp`, from +0.0, each 250 # next term `ls` and `rs` further on: `Matrix.product_run` after the first 251 # reads. 252 def run_from(fuel: Nat, +ls: Nat, +rs: Nat, +lp: Nat, +rp: Nat, la: Array, ra: Array) -> 253 F32 & (Array & Array): 254 match fuel: 255 case 0n: 256 (0.0, (la, ra)) 257 case 1n+rest: 258 Matrix.product_run(rest, ls, rs, lp, rp, 0.0, Array.get(F32, la, U32.from_nat(lp)), Array.get(F32, ra, U32.from_nat(rp))) 259 260 # The first term's offsets of cell `position`: batch b = position / (rows * 261 # columns) splits into (b / o2, b % o2), and each operand reads its own 262 # batch, a stretched axis at 0, then the cell's row and column at term 0. 263 def left_start(+layout: Layout, +position: Nat) -> Nat: 264 Layout{+lo, ro, +inner, +rows, +columns, +l1, +l2, r1, r2, +o2} = layout 265 +batch = Nat.div(position, Nat.mul(rows, columns)) 266 Nat.add(Nat.mul(batch_number(l1, l2, Nat.div(batch, o2), Nat.mod(batch, o2)), Nat.mul(rows, inner)), 267 Matrix.read_offset(lo, inner, rows, Nat.div(Nat.mod(position, Nat.mul(rows, columns)), columns), 0n)) 268 269 def right_start(+layout: Layout, +position: Nat) -> Nat: 270 Layout{lo, +ro, +inner, +rows, +columns, l1, l2, +r1, +r2, +o2} = layout 271 +batch = Nat.div(position, Nat.mul(rows, columns)) 272 Nat.add(Nat.mul(batch_number(r1, r2, Nat.div(batch, o2), Nat.mod(batch, o2)), Nat.mul(inner, columns)), 273 Matrix.read_offset(ro, inner, columns, Nat.mod(Nat.mod(position, Nat.mul(rows, columns)), columns), 0n)) 274 275 # Cell `position`: all `inner` terms, from +0.0. 276 def dot_at(+layout: Layout, +position: Nat, la: Array, ra: Array) -> F32 & (Array & Array): 277 Layout{+lo, +ro, +inner, +rows, +columns, l1, l2, r1, r2, o2} = layout 278 run_from(inner, Matrix.stride(lo, rows), Matrix.stride(ro, columns), left_start(layout, position), right_start(layout, position), la, ra) 279 280 # The cell function `Parallel.build` runs: a cell at or past the count is 281 # +0.0 padding. 282 def cell_at(+layout: Layout, +position: Nat, inputs: Array & Array, inside: Bool) -> F32 & (Array & Array): 283 (la, ra) = inputs 284 match inside: 285 case False{}: 286 (0.0, (la, ra)) 287 case True{}: 288 dot_at(layout, position, la, ra) 289 290 def build(+depth: Nat, +splits: Nat, +given: Layout, +count: Nat, +base: Nat, la: Array, ra: Array) -> Array: 291 Parallel.build(~Layout, ~Layout, ~Array, ~Array, ~layout, ~Elementwise.split, ~cell_at, given, depth, splits, count, base, (la, ra)) 292 293 def build_balanced(+depth: Nat, +splits: Nat, +given: Layout, +count: Nat, +base: Nat, la: Array, ra: Array) -> 294 Model.Balanced(F32, depth, build(depth, splits, given, count, base, la, ra)): 295 Parallel.balanced(~Layout, ~Layout, ~Array, ~Array, ~layout, ~Elementwise.split, ~cell_at, given, depth, splits, count, base, (la, ra)) 296 297 def terms_of(+given: Layout) -> Nat: 298 Layout{lo, ro, +inner, rows, columns, l1, l2, r1, r2, o2} = given 299 inner 300 301 # Both inputs come back to the caller as the kept halves of their clones, 302 # and the build reads the other halves. It forks by work as `Matrix.splits` 303 # says. 304 def filled(+given: Layout, +output: Shape.Shape, +depth: Nat, +length: Nat, 305 -valid: {Storage.bounds(depth, length) == True{} : Bool}, 306 -exact: {length == Shape.logical_numel(Shape.dimensions(output)) : Nat}, 307 left: Core.Dynamic & Array, right: Core.Dynamic & Array) -> 308 Result<&2, &1, Core.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>: 309 (kept_left, la) = left 310 (kept_right, ra) = right 311 Done{(kept_left, kept_right, Core.pack(output, Core.TTensor{Storage.Owned{depth, length, 312 build(depth, Matrix.splits(depth, terms_of(given)), given, length, 0n, la, ra), 313 build_balanced(depth, Matrix.splits(depth, terms_of(given)), given, length, 0n, la, ra), 314 valid}, exact}))} 315 316 def allocated(+given: Layout, +output: Shape.Shape, left: Core.Dynamic, right: Core.Dynamic, 317 result: Result<&2, &1, Core.Error, Core.Tensor>) -> Result<&2, &1, Core.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>: 318 match result: 319 case Fail{error}: 320 Fail{error} 321 case Done{Core.TTensor{storage, exact}}: 322 match storage: 323 case Storage.Empty{}: 324 Done{(left, right, Core.pack(output, Core.TTensor{Storage.Empty{}, exact}))} 325 case Storage.Owned{+depth, +length, array, balanced, valid}: 326 filled(given, output, depth, length, valid, exact, Elementwise.take(left), Elementwise.take(right)) 327 328 # The product of the storages of `left` and `right` read as `given` says, 329 # into a new tensor of shape `output`. 330 def product(+given: Layout, +output: Shape.Shape, left: Core.Dynamic, right: Core.Dynamic) -> 331 Result<&2, &1, Core.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>: 332 allocated(given, output, left, right, Core.new(output, 0.0)) 333 334 def tensor_failed(result: Result<&2, &1, Core.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>) -> 335 Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>: 336 match result: 337 case Fail{error}: 338 Fail{Matrix.TensorError{error}} 339 case Done{triple}: 340 Done{triple} 341 342 def planned(+left: Shape.Shape, +right: Shape.Shape, a: Core.Tensor, b: Core.Tensor, 343 result: Result<&2, &2, Matrix.Error, Plan>) -> Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>: 344 match result: 345 case Fail{error}: 346 Fail{error} 347 case Done{+plan}: 348 tensor_failed(product(forward_layout(plan), output_of(plan), Core.Packed{left, a}, Core.Packed{right, b})) 349 350 def matrices(dims: List<&2, U32>, other: List<&2, U32>) -> Bool: 351 Bool.and(Nat.is_eq(Shape.rank_dimensions(dims), 2n), Nat.is_eq(Shape.rank_dimensions(other), 2n)) 352 353 def both_matrices(+left: Shape.Shape, +right: Shape.Shape) -> Bool: 354 matrices(Shape.dimensions(left), Shape.dimensions(right)) 355 356 def dispatched(+left: Shape.Shape, +right: Shape.Shape, a: Core.Tensor, b: Core.Tensor, 357 +dims: List<&2, U32>, +other: List<&2, U32>, both: Bool) -> Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>: 358 match both: 359 case True{}: 360 Matrix.matmul(Core.Packed{left, a}, Core.Packed{right, b}) 361 case False{}: 362 planned(left, right, a, b, plan_dims(dims, other)) 363 364 # The operands' dimensions `dims` and `other` decide: two rank-2 operands go 365 # to `Matrix.matmul`, every other pair to the plan. 366 def routed(+left: Shape.Shape, +right: Shape.Shape, a: Core.Tensor, b: Core.Tensor, 367 +dims: List<&2, U32>, +other: List<&2, U32>) -> Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>: 368 dispatched(left, right, a, b, dims, other, matrices(dims, other)) 369 370 # torch.matmul(left, right). Both inputs come back to the caller. 371 def matmul(left: Core.Dynamic, right: Core.Dynamic) -> Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>: 372 Core.Packed{+left_shape, a} = left 373 Core.Packed{+right_shape, b} = right 374 routed(left_shape, right_shape, a, b, Shape.dimensions(left_shape), Shape.dimensions(right_shape)) 375 376 def gradient_shaped(+given: Layout, first: Core.Dynamic, second: Core.Dynamic, -dims: List<&2, U32>, 377 result: Shape.ConstructResult) -> Result<&2, &1, Core.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>: 378 match result: 379 case Shape.SError{error}: 380 Fail{Core.ShapeError{error}} 381 case Shape.SBuilt{+shape, preserved}: 382 product(given, shape, first, second) 383 384 # The left operand's gradient over the output's batch, g @ right^T, of 385 # dimensions `left_dims(plan)`: g has the output's shape and right the right 386 # operand's. A broadcast left operand still has to be summed to its shape. 387 def left_gradient(+plan: Plan, g: Core.Dynamic, right: Core.Dynamic) -> 388 Result<&2, &1, Core.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>: 389 gradient_shaped(left_layout(plan), g, right, left_dims(plan), Shape.construct(left_dims(plan))) 390 391 # The right operand's gradient over the output's batch, left^T @ g, of 392 # dimensions `right_dims(plan)`. 393 def right_gradient(+plan: Plan, left: Core.Dynamic, g: Core.Dynamic) -> 394 Result<&2, &1, Core.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>: 395 gradient_shaped(right_layout(plan), left, g, right_dims(plan), Shape.construct(right_dims(plan)))