*Added logic to ensure the indexing_map attribute can be dropped for collapsed contraction op.

*Refactored some tests and methods for better naming, comments and readability.

Author: shahidact

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

@@ -710,7 +710,10 @@ def LinalgStructuredInterface
     >,
     InterfaceMethod<
       /*desc=*/[{
-        Return true if the user has supplied an explicit indexing maps for this op.
+        Returns true if the user has supplied explicit indexing maps that are
+        different from default indexing maps for this op. Returns `false` otherwise.
+        Note, if the user define maps that are identical to the default maps,
+        this method returns `false`.
       }],
       /*retTy=*/"bool",
       /*methodName=*/"hasUserDefinedMaps",

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

@@ -674,8 +674,7 @@ def MatmulOp : LinalgStructuredBase_Op<"matmul", [
       static unsigned getNumRegionArgs();
       std::string getLibraryCallName();
       bool hasDynamicIndexingMaps();
-      /// Check if the op has broadcast and/or transpose semantic. Returns true if the
-      /// user defined indexing maps are not equal to default map.
+      /// Returns true if the user defined indexing maps are not equal to default maps.
       bool hasUserDefinedMaps();
     }];
 }
@@ -933,8 +932,7 @@ def BatchMatmulOp : LinalgStructuredBase_Op<"batch_matmul", !listconcat([AttrSiz
       static unsigned getNumRegionArgs();
       bool hasDynamicIndexingMaps() { return true; }
       std::string getLibraryCallName();
-      /// Check if the op has broadcast and/or transpose semantic. Returns true if the
-      /// user defined indexing maps are not equal to default map.
+      /// Returns true if the user defined indexing maps are not equal to default maps.
       bool hasUserDefinedMaps();
     }];
 }

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

@@ -3426,11 +3426,10 @@ Operation *LinalgDialect::materializeConstant(OpBuilder &builder,
   return arith::ConstantOp::materialize(builder, value, type, loc);
 }
 
-/// Returns true if the result AffineExpr of the \p explicitMap is same as \p
-/// defaultMap.
-static bool isValidResultDimExprs(AffineMap explictMap, AffineMap defaultMap) {
-  auto explicitRange = explictMap.getResults();
-  auto defaultRange = defaultMap.getResults();
+// Returns true if the result expression of `subMap` are a subset of `fullMap`.
+static bool areResultExprsSubsetOf(AffineMap subMap, AffineMap fullMap) {
+  auto explicitRange = subMap.getResults();
+  auto defaultRange = fullMap.getResults();
   DenseSet<AffineExpr> explicitSet(explicitRange.begin(), explicitRange.end());
   DenseSet<AffineExpr> defaultSet(defaultRange.begin(), defaultRange.end());
   llvm::set_union(explicitSet, defaultSet);
@@ -3455,7 +3454,7 @@ static LogicalResult verifyExtendedMatmulSemantic(MatmulOp matmulOp,
   auto opIndexingMap = opIndexingMaps[opIndex];
   auto defaultIndexingMap = defaultIndexingMaps[opIndex];
   // Check general validity of indexing map results.
-  if (!isValidResultDimExprs(opIndexingMap, defaultIndexingMap))
+  if (!areResultExprsSubsetOf(opIndexingMap, defaultIndexingMap))
     return matmulOp->emitOpError()
            << "Unexpected dim expression in map result.";
 
@@ -3470,44 +3469,31 @@ 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) {
-  AffineExpr exp = bcastMap.getResult(0);
-  return exp.isFunctionOfDim(0);
-}
-
-/// Checks if the given AffineMap's result dimensions are valid output result
-/// dimensions.
-static bool isValidOutputResultDim(AffineMap outputMap) {
-  enum Indices { batchPos, mPos, nPos };
-  AffineExpr exp0 = outputMap.getResult(batchPos);
-  AffineExpr exp1 = outputMap.getResult(mPos);
-  AffineExpr exp2 = outputMap.getResult(nPos);
-  return exp0.isFunctionOfDim(batchPos) && exp1.isFunctionOfDim(mPos) &&
-         exp2.isFunctionOfDim(nPos);
-}
-
 // Check general validity of input indexing map.
 static LogicalResult verifyInputMaps(BatchMatmulOp batchMatmulOp,
                                      AffineMap opIndexingMap,
                                      AffineMap defaultIndexingMap, bool isLHS) {
   // Check the result dims are valid.
-  if (!isValidResultDimExprs(opIndexingMap, defaultIndexingMap))
+  if (!areResultExprsSubsetOf(opIndexingMap, defaultIndexingMap))
     return batchMatmulOp->emitOpError()
-           << "Unexpected dim expression in map result.";
+           << "Unexpected result dim expression (outside the set of default "
+              "result dims).";
 
   // Check for valid number of result dims of input maps.
   if (opIndexingMap.getNumResults() > 3)
     return batchMatmulOp->emitOpError()
-           << "no. of result dim expression cannot exceed 3.";
+           << "no. of result dim expressions exceeds 3.";
+
+  auto hasValidBatchDim = [](AffineMap map) {
+    AffineExpr batchDim = map.getResult(0);
+    return batchDim.isFunctionOfDim(0);
+  };
 
   // Check if the requested broadcast is valid.
   if (isBroadcasted(opIndexingMap, defaultIndexingMap)) {
     if (!batchMatmulOp.isValidLhsRhsBroadcastMap(opIndexingMap, isLHS))
       return batchMatmulOp->emitOpError() << "Invalid broadcast requested.";
-  } else if (!isValidBatchDim(opIndexingMap)) {
+  } else if (!hasValidBatchDim(opIndexingMap)) {
     return batchMatmulOp->emitOpError()
            << "Invalid batch dimension expression.";
   }
@@ -3524,7 +3510,13 @@ static LogicalResult verifyOutputMap(BatchMatmulOp batchMatmulOp,
            << "expects 3 dims, but got (" << opIndexingMap.getNumResults()
            << ").";
 
-  if (!isValidOutputResultDim(opIndexingMap))
+  auto areValidOutputResultDim = [](AffineMap outputMap) {
+    return outputMap.getResult(0).isFunctionOfDim(0) &&
+           outputMap.getResult(1).isFunctionOfDim(1) &&
+           outputMap.getResult(2).isFunctionOfDim(2);
+  };
+
+  if (!areValidOutputResultDim(opIndexingMap))
     return batchMatmulOp->emitOpError()
            << "Invalid output map result dimension.";
 
@@ -3941,7 +3933,8 @@ bool BatchMatmulOp::hasUserDefinedMaps() {
 
 /// Returns true if the given broadcast map \p bcastMap is valid for this op.
 bool BatchMatmulOp::isValidLhsRhsBroadcastMap(AffineMap bcastMap, bool isLHS) {
-  assert(bcastMap.getNumResults() < 3 && "Expected single result dim expr.");
+  assert(bcastMap.getNumResults() < 3 &&
+         "Expected less than 3 result dim expr.");
   bool isValid = false;
   enum Indices { batchPos, mPos, nPos, kPos };
   if (bcastMap.getNumResults() == 1) {

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

@@ -907,14 +907,9 @@ struct RankReduceContractionOps : OpRewritePattern<FromOpTy> {
 
   LogicalResult matchAndRewrite(FromOpTy contractionOp,
                                 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 ||
-        std::is_same<FromOpTy, MatmulOp>::value) {
-      if (contractionOp.hasUserDefinedMaps()) {
-        return rewriter.notifyMatchFailure(
-            contractionOp, "ops with user-defined maps are not supported");
-      }
+    if (contractionOp.hasUserDefinedMaps()) {
+      return rewriter.notifyMatchFailure(
+          contractionOp, "ops with user-defined maps are not supported");
     }
 
     auto loc = contractionOp.getLoc();
@@ -945,21 +940,10 @@ 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")
         continue;
-
-      // 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());
-      }
+      collapsedOp->setAttr(attr.getName(), attr.getValue());
     }
 
     auto results = contractionOp.getResults();

mlir/test/Dialect/Linalg/invalid.mlir

@@ -1303,7 +1303,7 @@ func.func @missing_indexing_map_batch_matmul(%arg0: memref<?x?x?xf32>, %arg1: me
 // -----
 
 func.func @invalid_dim_expr_batch_matmul_a(%arg0: memref<?x?x?xf32>, %arg1: memref<?x?x?xf32>, %arg2: memref<?x?x?xf32>) {
-  // expected-error @+1 {{Unexpected dim expression in map result}}
+  // expected-error @+1 {{Unexpected result dim expression (outside the set of default result dims)}}
   linalg.batch_matmul indexing_maps = [
                        affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>,
                        affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>,
@@ -1316,7 +1316,7 @@ func.func @invalid_dim_expr_batch_matmul_a(%arg0: memref<?x?x?xf32>, %arg1: memr
 // -----
 
 func.func @invalid_dim_expr_batch_matmul_b(%arg0: memref<?x?x?xf32>, %arg1: memref<?x?x?xf32>, %arg2: memref<?x?x?xf32>) {
-  // expected-error @+1 {{Unexpected dim expression in map result}}
+  // expected-error @+1 {{Unexpected result dim expression (outside the set of default result dims)}}
   linalg.batch_matmul indexing_maps = [
                        affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>,
                        affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>,
@@ -1407,7 +1407,7 @@ func.func @invalid_batch_dim_batch_matmul_b(%arg0: memref<?x?x?xf32>, %arg1: mem
 // -----
 
 func.func @invalid_A_map_result_num_batch_matmul(%arg0: memref<?x?x?xf32>, %arg1: memref<?x?x?xf32>, %arg2: memref<?x?xf32>) {
-  // expected-error @+1 {{'linalg.batch_matmul' op no. of result dim expression cannot exceed 3.}}
+  // expected-error @+1 {{'linalg.batch_matmul' op no. of result dim expressions exceeds 3.}}
   linalg.batch_matmul indexing_maps = [
                             affine_map<(d0, d1, d2, d3) -> (d0, d1, d3, d3)>,
                             affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>,
@@ -1421,7 +1421,7 @@ func.func @invalid_A_map_result_num_batch_matmul(%arg0: memref<?x?x?xf32>, %arg1
 // -----
 
 func.func @invalid_B_map_result_num_batch_matmul(%arg0: memref<?x?x?xf32>, %arg1: memref<?x?x?xf32>, %arg2: memref<?x?xf32>) {
-  // expected-error @+1 {{'linalg.batch_matmul' op no. of result dim expression cannot exceed 3.}}
+  // expected-error @+1 {{'linalg.batch_matmul' op no. of result dim expressions exceeds 3.}}
   linalg.batch_matmul indexing_maps = [
                             affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>,
                             affine_map<(d0, d1, d2, d3) -> (d0, d3, d2, d3)>,

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

@@ -1489,14 +1489,14 @@ func.func @matmul_transpose_b(%arg0: memref<3x5xf32>, %arg1: memref<7x5xf32>, %a
 // 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-LABEL:   func.func @batch_matmul_bcast_k_to_fill_missing_dims_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>) {
+func.func @batch_matmul_bcast_k_to_fill_missing_dims_A(%arg0: memref<5xf32>, %arg1: memref<2x5x7xf32>, %arg2: memref<2x3x7xf32>) {
   linalg.batch_matmul indexing_maps = [
                        affine_map<(d0, d1, d2, d3) -> (d3)>,
                        affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>,