src/tensor/batched_matmul.bend

bendtorch-core@0.1.0.0 · 395 lines · 51 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
4import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
5import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
6import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
7import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
8import bendtorch-core@0.1.0.0/src/tensor/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.
40type 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.
46type Operand is Data:
47  Operand{batch: Batch, outer: U32, inner: U32, vector: Bool}
48
49def 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
62def 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.
77def 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
88def 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
91def 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.
95def 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
106def column_dims(+columns: U32, right_vector: Bool) -> List<&2, U32>:
107  match right_vector:
108    case True{}:
109      []
110    case False{}:
111      [columns]
112
113def 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.
123type 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
127def 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
131def 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<dims>) -> 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
141def 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
154def 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
159def 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
166def inner_of(+operand: Operand) -> U32:
167  Operand{batch, outer, +inner, vector} = operand
168  inner
169
170def 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
177def 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.
188def plan_dims(left: List<&2, U32>, right: List<&2, U32>) -> Result<&2, &2, Matrix.Error, Plan>:
189  left_parsed(right, left_operand(left))
190
191def 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.
198type 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
202def 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.
208def 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.
215def 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.
222def 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.
229def 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
233def 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.
239def stretch(+index: Nat, one: Bool) -> Nat:
240  match one:
241    case True{}:
242      0n
243    case False{}:
244      index
245
246def 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.
252def run_from(fuel: Nat, +ls: Nat, +rs: Nat, +lp: Nat, +rp: Nat, la: Array<F32>, ra: Array<F32>) ->
253  F32 & (Array<F32> & Array<F32>):
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.
263def 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
269def 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.
276def dot_at(+layout: Layout, +position: Nat, la: Array<F32>, ra: Array<F32>) -> F32 & (Array<F32> & Array<F32>):
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.
282def cell_at(+layout: Layout, +position: Nat, inputs: Array<F32> & Array<F32>, inside: Bool) -> F32 & (Array<F32> & Array<F32>):
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
290def build(+depth: Nat, +splits: Nat, +given: Layout, +count: Nat, +base: Nat, la: Array<F32>, ra: Array<F32>) -> Array<F32>:
291  Parallel.build(~Layout, ~Layout, ~Array<F32>, ~Array<F32>, ~layout, ~Elementwise.split, ~cell_at, given, depth, splits, count, base, (la, ra))
292
293def build_balanced(+depth: Nat, +splits: Nat, +given: Layout, +count: Nat, +base: Nat, la: Array<F32>, ra: Array<F32>) ->
294  Model.Balanced(F32, depth, build(depth, splits, given, count, base, la, ra)):
295  Parallel.balanced(~Layout, ~Layout, ~Array<F32>, ~Array<F32>, ~layout, ~Elementwise.split, ~cell_at, given, depth, splits, count, base, (la, ra))
296
297def 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.
304def 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<F32>, right: Core.Dynamic & Array<F32>) ->
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
316def allocated(+given: Layout, +output: Shape.Shape, left: Core.Dynamic, right: Core.Dynamic,
317  result: Result<&2, &1, Core.Error, Core.Tensor<output>>) -> 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`.
330def 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
334def 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
342def planned(+left: Shape.Shape, +right: Shape.Shape, a: Core.Tensor<left>, b: Core.Tensor<right>,
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
350def 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
353def both_matrices(+left: Shape.Shape, +right: Shape.Shape) -> Bool:
354  matrices(Shape.dimensions(left), Shape.dimensions(right))
355
356def dispatched(+left: Shape.Shape, +right: Shape.Shape, a: Core.Tensor<left>, b: Core.Tensor<right>,
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.
366def routed(+left: Shape.Shape, +right: Shape.Shape, a: Core.Tensor<left>, b: Core.Tensor<right>,
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.
371def 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
376def gradient_shaped(+given: Layout, first: Core.Dynamic, second: Core.Dynamic, -dims: List<&2, U32>,
377  result: Shape.ConstructResult<dims>) -> 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.
387def 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)`.
393def 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)))