[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