[Mlir-commits] [mlir] [MLIR][Linalg] Introduce broadcast/transpose semantic to 'linalg.batc… (PR #122275)

[Adam Siemieniuk]

================
@@ -3450,6 +3467,46 @@ static LogicalResult verifyExtendedMatmulSemantic(MatmulOp matmulOp,
   return success();
 }
 
+/// Checks if the given AffineMap represents a valid batch dimension.
+/// It checks if the first result dimension is a function of the first
+/// dimension.
+static bool isValidBatchDim(AffineMap bcastMap) {
+  assert(bcastMap.getNumResults() == 3 && "Expected three result dim expr.");
----------------
adam-smnk wrote:

Using the example you provided trips the assertion:
```mlir
linalg.batch_matmul indexing_maps = [
                       affine_map<(d0, d1, d2, d3) -> (d0, d1, d3, d3)>,
                       affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>,
                       affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
                      ]
                      ins(%arg0, %arg1 : memref<?x?x?xf32>, memref<?x?x?xf32>)
                      outs(%arg2 :memref<?x?x?xf32>)
```

https://github.com/llvm/llvm-project/pull/122275