*Added logic to update the indexing_map attribute for collapsed MatmulOp.

*Updated test names and comments for consistency.

Author: shahidact

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

@@ -828,8 +828,8 @@ def BatchMatmulOp : LinalgStructuredBase_Op<"batch_matmul", !listconcat([AttrSiz
     them to the same data type as the accumulator/output.
 
     Broadcast and Transpose semantics can be appiled by specifying the explicit attribute
-    'indexing_maps' as shown below.This is a list attribute, so the list must include all
-    the maps if specified.
+    'indexing_maps' as shown below. This is a list attribute, so must include maps for all
+    arguments if specified.
 
     Example Transpose:
     ```
@@ -845,15 +845,15 @@ def BatchMatmulOp : LinalgStructuredBase_Op<"batch_matmul", !listconcat([AttrSiz
     Example Broadcast:
     ```
     linalg.batch_matmul indexing_maps = [
-                       affine_map<(d0, d1, d2, d3) -> (d3)>,     // broadcast
+                       affine_map<(d0, d1, d2, d3) -> (d3)>,         // broadcast
                        affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>,
                        affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
                      ]
                      ins(%arg0, %arg1 : memref<5xf32>, memref<2x5x7xf32>)
                      outs(%arg2: memref<2x3x7xf32>)
     ```
 
-    Example Broadcast and transpose:
+    Example Broadcast and Transpose:
     ```
     linalg.batch_matmul indexing_maps = [
                        affine_map<(d0, d1, d2, d3) -> (d1, d3)>,     // broadcast
@@ -919,7 +919,7 @@ def BatchMatmulOp : LinalgStructuredBase_Op<"batch_matmul", !listconcat([AttrSiz
         return regionBuilder;
       }
 
-      /// Returns a list of AffineMap with the typical batch_matmul indexing charactristic.
+      /// Returns a list with default AffineMap(s), i.e. without broadcasts and transpositions.
       static SmallVector<AffineMap> getDefaultIndexingMaps(MLIRContext *context);
 
       /// Returns true if the given broadcast map \p bcastMap is valid for this op.

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

@@ -32,6 +32,7 @@
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
+#include <type_traits>
 
 namespace mlir {
 #define GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMSPASS
@@ -908,11 +909,11 @@ struct RankReduceContractionOps : OpRewritePattern<FromOpTy> {
                                 PatternRewriter &rewriter) const override {
     // Check to not let go the batch_matmul with extended semantic, through this
     // transform.
-    if (std::is_same<FromOpTy, BatchMatmulOp>::value) {
+    if (std::is_same<FromOpTy, BatchMatmulOp>::value ||
+        std::is_same<FromOpTy, MatmulOp>::value) {
       if (contractionOp.hasUserDefinedMaps()) {
         return rewriter.notifyMatchFailure(
-            contractionOp,
-            "only batch_matmul ops with non-extended semantics are supported");
+            contractionOp, "ops with user-defined maps are not supported");
       }
     }
 
@@ -944,10 +945,21 @@ struct RankReduceContractionOps : OpRewritePattern<FromOpTy> {
         loc, collapsedResultTy, ValueRange{collapsedLhs, collapsedRhs},
         ValueRange{collapsedInit});
     for (auto attr : contractionOp->getAttrs()) {
-      if (attr.getName() == LinalgDialect::kMemoizedIndexingMapsAttrName ||
-          attr.getName() == "indexing_maps")
+      if (attr.getName() == LinalgDialect::kMemoizedIndexingMapsAttrName)
         continue;
-      collapsedOp->setAttr(attr.getName(), attr.getValue());
+
+      // Update the indexing_maps attribute for the collapsed MatmulOp.
+      if (attr.getName() == "indexing_maps" &&
+          std::is_same<FromOpTy, BatchMatmulOp>::value &&
+          std::is_same<ToOpTy, MatmulOp>::value) {
+        SmallVector<Attribute, 3> indexingMapsAttr = llvm::map_to_vector(
+            MatmulOp::getDefaultIndexingMaps(rewriter.getContext()),
+            [](AffineMap map) -> Attribute { return AffineMapAttr::get(map); });
+        collapsedOp->setAttr(attr.getName(),
+                             rewriter.getArrayAttr(indexingMapsAttr));
+      } else {
+        collapsedOp->setAttr(attr.getName(), attr.getValue());
+      }
     }
 
     auto results = contractionOp.getResults();

mlir/lib/Dialect/Linalg/Transforms/TransposeMatmul.cpp

@@ -88,12 +88,9 @@ FailureOr<Operation *>
 mlir::linalg::transposeBatchMatmul(RewriterBase &rewriter,
                                    linalg::BatchMatmulOp batchMatmulOp,
                                    bool transposeLHS) {
-  // Check to not let go the batch_matmul with extended semantic, through this
-  // transform.
   if (batchMatmulOp.hasUserDefinedMaps()) {
     return rewriter.notifyMatchFailure(
-        batchMatmulOp,
-        "only batch_matmul ops with non-extended semantics are supported");
+        batchMatmulOp, "ops with user-defined maps are not supported");
   }
 
   if (!bufferization::hasTensorSemantics(batchMatmulOp))

mlir/test/Dialect/Linalg/named-ops.mlir

@@ -1578,11 +1578,11 @@ func.func @batch_matmul_bcast_batch_dim_B(%arg0: memref<2x3x5xf32>, %arg1: memre
 
 // -----
 
-// CHECK-LABEL: func @batch_matmul_explicit_transpose_a
+// CHECK-LABEL: func @batch_matmul_explicit_transpose_A
 //       CHECK:   linalg.batch_matmul
 //  CHECK-SAME:     ins(%{{.+}}, %{{.+}} : memref<2x5x3xf32>, memref<2x5x7xf32>)
 //  CHECK-SAME:     outs(%{{.+}} : memref<2x3x7xf32>)
-func.func @batch_matmul_explicit_transpose_a(%arg0: memref<2x5x3xf32>, %arg1: memref<2x5x7xf32>, %arg2: memref<2x3x7xf32>) {
+func.func @batch_matmul_explicit_transpose_A(%arg0: memref<2x5x3xf32>, %arg1: memref<2x5x7xf32>, %arg2: memref<2x3x7xf32>) {
   linalg.batch_matmul indexing_maps = [
                        affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>,
                        affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>,
@@ -1594,11 +1594,11 @@ func.func @batch_matmul_explicit_transpose_a(%arg0: memref<2x5x3xf32>, %arg1: me
 
 // -----
 
-// CHECK-LABEL: func @batch_matmul_explicit_transpose_b
+// CHECK-LABEL: func @batch_matmul_explicit_transpose_B
 //       CHECK:   linalg.batch_matmul
 //  CHECK-SAME:     ins(%{{.+}}, %{{.+}} : memref<2x3x5xf32>, memref<2x7x5xf32>)
 //  CHECK-SAME:     outs(%{{.+}} : memref<2x3x7xf32>)
-func.func @batch_matmul_explicit_transpose_b(%arg0: memref<2x3x5xf32>, %arg1: memref<2x7x5xf32>, %arg2: memref<2x3x7xf32>) {
+func.func @batch_matmul_explicit_transpose_B(%arg0: memref<2x3x5xf32>, %arg1: memref<2x7x5xf32>, %arg2: memref<2x3x7xf32>) {
   linalg.batch_matmul indexing_maps = [
                        affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>,
                        affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>,