*Added checks for extended semantics and exit gracefully in user passes.

*Added and udated test cases.
*Refactored verification logic.

Author: shahidact

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

@@ -3474,7 +3474,6 @@ static LogicalResult verifyExtendedMatmulSemantic(MatmulOp matmulOp,
 /// 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.");
   AffineExpr exp = bcastMap.getResult(0);
   return exp.isFunctionOfDim(0);
 }
@@ -3490,6 +3489,48 @@ static bool isValidOutputResultDim(AffineMap outputMap) {
          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))
+    return batchMatmulOp->emitOpError()
+           << "Unexpected dim expression in map result.";
+
+  // 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.";
+
+  // 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)) {
+    return batchMatmulOp->emitOpError()
+           << "Invalid batch dimension expression.";
+  }
+  return success();
+}
+
+/// This function checks if the given AffineMap for the output of a
+/// BatchMatmulOp has exactly 3 result dimensions and if the output map result
+/// dimensions are valid.
+static LogicalResult verifyOutputMap(BatchMatmulOp batchMatmulOp,
+                                     AffineMap opIndexingMap) {
+  if (opIndexingMap.getNumResults() != 3)
+    return batchMatmulOp->emitOpError()
+           << "expects 3 dims, but got (" << opIndexingMap.getNumResults()
+           << ").";
+
+  if (!isValidOutputResultDim(opIndexingMap))
+    return batchMatmulOp->emitOpError()
+           << "Invalid output map result dimension.";
+
+  return success();
+}
+
 /// Verifies the broadcast and transpose semantic specified by the explicit
 /// indexing map for the BatchMatmulOp \p op for each operand specified by \p
 /// opIndex.
@@ -3503,37 +3544,14 @@ verifyExtendedBatchMatmulSemantic(BatchMatmulOp batchMatmulOp,
 
   auto opIndexingMap = opIndexingMaps[opIndex];
   auto defaultIndexingMap = defaultIndexingMaps[opIndex];
-  // Check general validity of indexing map results.
-  if (opIndex < 2) {
-    if (!isValidResultDimExprs(opIndexingMap, defaultIndexingMap))
-      return batchMatmulOp->emitOpError()
-             << "Unexpected dim expression in map result.";
-    // Check if the requested broadcast is valid.
-    if (isBroadcasted(opIndexingMap, defaultIndexingMap)) {
-      if (!batchMatmulOp.isValidLhsRhsBroadcastMap(opIndexingMap,
-                                                   opIndex == 0)) {
-        return batchMatmulOp->emitOpError() << "Invalid broadcast requested.";
-      }
-    } else {
-      // 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.";
-
-      if (!isValidBatchDim(opIndexingMap))
-        return batchMatmulOp->emitOpError()
-               << "Invalid batch dimension expression.";
-    }
-  } else {
-    // Check for valid number of result dims of output map.
-    if (opIndexingMap.getNumResults() != 3)
-      return batchMatmulOp->emitOpError()
-             << "no. of result dim expression cannot exceed 3.";
 
-    if (!isValidOutputResultDim(opIndexingMap))
-      return batchMatmulOp->emitOpError()
-             << "Invalid output map result dimension.";
-  }
+  if (opIndex == 2 && failed(verifyOutputMap(batchMatmulOp, opIndexingMap)))
+    return failure();
+
+  if (failed(verifyInputMaps(batchMatmulOp, opIndexingMap, defaultIndexingMap,
+                             opIndex == 0)))
+    return failure();
+
   return success();
 }
 

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

@@ -138,6 +138,16 @@ transposePackedMatmul(RewriterBase &rewriter, linalg::LinalgOp linalgOp,
 FailureOr<PackResult>
 linalg::blockPackMatmul(RewriterBase &rewriter, linalg::LinalgOp linalgOp,
                         const ControlBlockPackMatmulFn &controlPackMatmul) {
+  // Check to not let go the batch_matmul with extended semantic, through this
+  // transform.
+  if (auto *batchMatmulOp = dyn_cast<linalg::BatchMatmulOp>(&linalgOp)) {
+    if (batchMatmulOp->hasUserDefinedMaps()) {
+      return rewriter.notifyMatchFailure(
+          *batchMatmulOp,
+          "only batch_matmul ops with non-extended semantics are supported");
+    }
+  }
+
   if (linalgOp.hasPureBufferSemantics())
     return rewriter.notifyMatchFailure(linalgOp, "require tensor semantics");
 

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

@@ -906,6 +906,15 @@ 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) {
+      if (contractionOp.hasUserDefinedMaps()) {
+        return rewriter.notifyMatchFailure(
+            contractionOp,
+            "only batch_matmul ops with non-extended semantics are supported");
+      }
+    }
 
     auto loc = contractionOp.getLoc();
     auto inputs = contractionOp.getDpsInputs();

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

@@ -88,6 +88,14 @@ 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");
+  }
+
   if (!bufferization::hasTensorSemantics(batchMatmulOp))
     return rewriter.notifyMatchFailure(
         batchMatmulOp, "only matmul ops with tensors are supported");

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

@@ -1004,21 +1004,22 @@ func.func @matmul_transpose_a_b_explicit(%arg0: memref<5x3xf32>, %arg1: memref<7
 // CHECK: #[[$ATTR_2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
 
 // CHECK-LABEL:   func.func @batch_matmul(
-// CHECK-SAME:                            %[[VAL_0:.*]]: tensor<?x?x?xf32>, %[[VAL_1:.*]]: tensor<?x?x?xf32>,
-// CHECK-SAME:                            %[[VAL_2:.*]]: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
-// CHECK:           linalg.generic {indexing_maps = [#[[$ATTR_0]], #[[$ATTR_1]], #[[$ATTR_2]]], iterator_types = ["parallel", "parallel", "parallel", "reduction"]} ins(%[[VAL_0]], %[[VAL_1]] : tensor<?x?x?xf32>, tensor<?x?x?xf32>) outs(%[[VAL_2]] : tensor<?x?x?xf32>) {
+// CHECK-SAME:                            %[[VAL_0:.*]]: tensor<2x3x5xf32>,
+// CHECK-SAME:                            %[[VAL_1:.*]]: tensor<2x5x7xf32>,
+// CHECK-SAME:                            %[[VAL_2:.*]]: tensor<2x3x7xf32>) -> tensor<2x3x7xf32> {
+// CHECK:           %[[VAL_3:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_0]], #[[$ATTR_1]], #[[$ATTR_2]]], iterator_types = ["parallel", "parallel", "parallel", "reduction"]} ins(%[[VAL_0]], %[[VAL_1]] : tensor<2x3x5xf32>, tensor<2x5x7xf32>) outs(%[[VAL_2]] : tensor<2x3x7xf32>) {
 // CHECK:           arith.mulf
 // CHECK:           arith.addf
 
-func.func @batch_matmul(%arg0: tensor<?x?x?xf32>, %arg1: tensor<?x?x?xf32>, %arg2: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
+func.func @batch_matmul(%arg0: tensor<2x3x5xf32>, %arg1: tensor<2x5x7xf32>, %arg2: tensor<2x3x7xf32>) -> tensor<2x3x7xf32> {
   %0 = linalg.batch_matmul indexing_maps = [
                             affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>,
                             affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>,
                             affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
                            ]
-    ins(%arg0, %arg1: tensor<?x?x?xf32>, tensor<?x?x?xf32>)
-    outs(%arg2: tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-  return %0 : tensor<?x?x?xf32>
+    ins(%arg0, %arg1: tensor<2x3x5xf32>, tensor<2x5x7xf32>)
+    outs(%arg2: tensor<2x3x7xf32>) -> tensor<2x3x7xf32>
+  return %0 : tensor<2x3x7xf32>
 }
 
 // -----

mlir/test/Dialect/Linalg/invalid.mlir

@@ -1408,7 +1408,7 @@ func.func @invalid_B_map_result_num_batch_matmul(%arg0: memref<?x?x?xf32>, %arg1
 // -----
 
 func.func @invalid_C_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 expects 3 dims, but got (2).}}
   linalg.batch_matmul indexing_maps = [
                             affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>,
                             affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>,