[mlir] [MLIR][Linalg] Introduce broadcast/transpose semantic to 'linalg.batc… (PR #122275)
Andrzej Warzyński
llvmlistbot at llvm.org
Tue Jan 21 08:04:05 PST 2025
================
@@ -1487,6 +1487,154 @@ func.func @matmul_transpose_b(%arg0: memref<3x5xf32>, %arg1: memref<7x5xf32>, %a
// -----
+// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
+// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK: #[[$ATTR_2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+
+// CHECK-LABEL: func.func @batch_matmul_bcast_batch_and_m_dim_A(
+// CHECK-SAME: %[[VAL_0:.*]]: memref<5xf32>,
+// CHECK-SAME: %[[VAL_1:.*]]: memref<2x5x7xf32>,
+// CHECK-SAME: %[[VAL_2:.*]]: memref<2x3x7xf32>) {
+// CHECK: linalg.batch_matmul ins(%[[VAL_0]], %[[VAL_1]] : memref<5xf32>, memref<2x5x7xf32>) outs(%[[VAL_2]] : memref<2x3x7xf32>) indexing_maps = [#[[$ATTR_0]], #[[$ATTR_1]], #[[$ATTR_2]]]
+// CHECK: return
+// CHECK: }
+func.func @batch_matmul_bcast_batch_and_m_dim_A(%arg0: memref<5xf32>, %arg1: memref<2x5x7xf32>, %arg2: memref<2x3x7xf32>) {
----------------
banach-space wrote:
Sorry, I thought that I have already replied to this.
> I intend batch and m gets broadcasted.
"Broadcasting" means "duplicating" something. It's not possible to duplicate something that doesn't exist :)
Here's a specific example from:
* https://mlir.llvm.org/docs/Dialects/Vector/#vectorbroadcast-vectorbroadcastop.
```mlir
%0 = arith.constant 0.0 : f32
%1 = vector.broadcast %0 : f32 to vector<16xf32>
```
Here, the existing value, `%0`, is duplicated/broadcasted across all "slots" within (presumably empty) vector `%1`. You could say that `%1` is filled with `%0` through `broadcast` operation.
Perhaps I am missing something in this context (*), but from what I can tell this is the same. And, a more accurate name for this test would be:
```
`@batch_matmul_bcast_k_to_fill_batch_and_m_dims_in_A`
```
Bit of a mouthful (suggestions welcome), but feels more accurate to me.
In any case, I want us to avoid bike-shedding this. I mostly want to make sure we are on the same page with regard to the high level logic.
(*) I appreciate that we are not broadcasting vectors here.
https://github.com/llvm/llvm-project/pull/122275
More information about the Mlir-commits
mailing list