[Mlir-commits] [mlir] [mlir][linalg] Improve contraction specialization (PR #201300)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-linalg

Author: Longsheng Mou (CoTinker)

<details>
<summary>Changes</summary>

Extend the matcher to recognize additional supported boolean contraction form(and+or). Fixes #<!-- -->198235.

---
Full diff: https://github.com/llvm/llvm-project/pull/201300.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp (+20-9) 
- (modified) mlir/test/Dialect/Linalg/roundtrip-morphism-linalg-named-ops.mlir (+18) 


``````````diff
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
index cc8a69cef8e5d..dcc46116e77bc 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
@@ -450,6 +450,20 @@ static FailureOr<LinalgOp> specializeLinalgMmt4D(RewriterBase &rewriter,
                                            namedOpMaps);
 }
 
+static bool isSupportedContractionPair(Operation *first, Operation *second) {
+  if (isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second))
+    return true;
+  if (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second))
+    return true;
+  if (isa<complex::MulOp>(first) && isa<complex::AddOp>(second))
+    return true;
+  if (isa<arith::AndIOp>(first) && isa<arith::OrIOp>(second) &&
+      first->getResult(0).getType().isInteger(1))
+    return true;
+
+  return false;
+}
+
 // Converts linalg.generic to named linalg.*matmul* where possible.
 static FailureOr<LinalgOp> specializeLinalgContractions(RewriterBase &rewriter,
                                                         GenericOp genericOp,
@@ -463,15 +477,12 @@ static FailureOr<LinalgOp> specializeLinalgContractions(RewriterBase &rewriter,
                    [](AffineMap m) { return !m.isProjectedPermutation(); }))
     return failure();
 
-  // Only mul+add contraction is supported.
-  // Currently, there is no way to control the contraction body type in named
-  // and category ops which all default to mul+add only.
-  if (!mlir::linalg::detail::isContractionBody(
-          *genericOp.getBlock(), [](Operation *first, Operation *second) {
-            return (isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
-                   (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)) ||
-                   (isa<complex::MulOp>(first) && isa<complex::AddOp>(second));
-          }))
+  // Only contractions that can be represented by named linalg ops are
+  // eligible for specialization:
+  //   - mul + add    (floating-point, integer, complex)
+  //   - and + or     (i1)
+  if (!mlir::linalg::detail::isContractionBody(*genericOp.getBlock(),
+                                               isSupportedContractionPair))
     return failure();
 
   // Determine the cast type for the named matmul op, or bail out if casts
diff --git a/mlir/test/Dialect/Linalg/roundtrip-morphism-linalg-named-ops.mlir b/mlir/test/Dialect/Linalg/roundtrip-morphism-linalg-named-ops.mlir
index 604077f6f7834..3ea6f99827484 100644
--- a/mlir/test/Dialect/Linalg/roundtrip-morphism-linalg-named-ops.mlir
+++ b/mlir/test/Dialect/Linalg/roundtrip-morphism-linalg-named-ops.mlir
@@ -222,6 +222,24 @@ func.func @matmul(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
 
 // -----
 
+func.func @matmul_bool(%A: tensor<?x?xi1>, %B: tensor<?x?xi1>,
+    %Out: tensor<?x?xi1>) -> tensor<?x?xi1> {
+  %0 = linalg.matmul
+    ins(%A, %B : tensor<?x?xi1>, tensor<?x?xi1>)
+    outs(%Out : tensor<?x?xi1>) -> tensor<?x?xi1>
+  return %0 : tensor<?x?xi1>
+}
+
+// CHECK-LABEL: @matmul_bool
+// CHECK-SAME: %[[A:.+]]: tensor<?x?xi1>, %[[B:.+]]: tensor<?x?xi1>,
+// CHECK-SAME: %[[OUT:.+]]: tensor<?x?xi1>) -> tensor<?x?xi1>
+// CHECK-NOT: linalg.generic
+// CHECK: linalg.matmul
+// CHECK-SAME: ins(%[[A]], %[[B]] : tensor<?x?xi1>, tensor<?x?xi1>)
+// CHECK-SAME: outs(%[[OUT]] : tensor<?x?xi1>) -> tensor<?x?xi1>
+
+// -----
+
 // Check matmul with unsigned cast is correctly raised back to named op.
 func.func @matmul_unsigned_cast(%A: tensor<16x8xi16>, %B: tensor<8x32xi64>,
     %Out: tensor<16x32xi32>) -> tensor<16x32xi32> {

``````````

</details>


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