Skip to content

Commit ce14c95

Browse files
committed
Expand docs on transpose and broadcast
1 parent cd04568 commit ce14c95

File tree

1 file changed

+15
-4
lines changed

1 file changed

+15
-4
lines changed

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -727,10 +727,21 @@ def ContractOp : LinalgStructuredBase_Op<"contract", [
727727
outs(%C: tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
728728
```
729729

730-
Note that by permuting the dims in the co-domains of the `affine_map`s, we
731-
can apply arbitrary transposes to the inputs and output. Similarly,
732-
arbitrary broadcasts can be achieved through leaving out dims on either
733-
input operand - these dims' inferred iter type will be parallel.
730+
Note that by permuting dims in the co-domains of the `affine_map`s arbitrary
731+
transposes can be applied to the inputs and output. Similarly, arbitrary
732+
broadcasts can be achieved through leaving out dims on either input operand
733+
(these dims' inferred iter type will be parallel). For example, the
734+
following is a variant of batch-matmul where a transposition is applied to
735+
`A` while matrix `B` gets broadcasted along the batch dimension:
736+
737+
```
738+
linalg.contract
739+
indexing_maps = [affine_map<(batch, m, n, k) -> (batch, k, m)>,
740+
affine_map<(batch, m, n, k) -> (k, n)>,
741+
affine_map<(batch, m, n, k) -> (batch, m, n)>]
742+
ins(%A, %B: memref<?x?x?xf32>, memref<?x?xf32>)
743+
outs(%C: memref<?x?x?xf32>)
744+
```
734745

735746
Numeric casting is performed on the operands to the inner multiplication,
736747
promoting them to the same data type as the accumulator/output.

0 commit comments

Comments
 (0)