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

Adam Siemieniuk llvmlistbot at llvm.org
Mon Jan 13 08:22:42 PST 2025


================
@@ -935,7 +935,8 @@ struct RankReduceContractionOps : OpRewritePattern<FromOpTy> {
         loc, collapsedResultTy, ValueRange{collapsedLhs, collapsedRhs},
         ValueRange{collapsedInit});
     for (auto attr : contractionOp->getAttrs()) {
-      if (attr.getName() == LinalgDialect::kMemoizedIndexingMapsAttrName)
+      if (attr.getName() == LinalgDialect::kMemoizedIndexingMapsAttrName ||
+          attr.getName() == "indexing_maps")
----------------
adam-smnk wrote:

Still, consider this example:
```mlir
func.func @entry(%arg0: memref<1x5x5xf32>,
    %arg1: memref<1x5x5xf32>, %arg2: memref<1x5x5xf32>) {
  linalg.batch_matmul indexing_maps = [
      affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>,
      affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>, // Transposed B
      affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
    ]
    ins(%arg0, %arg1 : memref<1x5x5xf32>, memref<1x5x5xf32>)
    outs(%arg2: memref<1x5x5xf32>)
  return
}
```
Applying `mlir-opt -test-linalg-rank-reduce-contraction-ops` drops the maps which means transpose is lost.

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


More information about the Mlir-commits mailing list