[Mlir-commits] [mlir] 46cfdfb - [mlir][Linalg] Fix Generalize transform

[Nicolas Vasilache]

Author: Nicolas Vasilache
Date: 2023-03-31T00:38:06-07:00
New Revision: 46cfdfb5f799d01ffe902dc8f2b66b3d9c7eec4f

URL: https://github.com/llvm/llvm-project/commit/46cfdfb5f799d01ffe902dc8f2b66b3d9c7eec4f
DIFF: https://github.com/llvm/llvm-project/commit/46cfdfb5f799d01ffe902dc8f2b66b3d9c7eec4f.diff

LOG: [mlir][Linalg] Fix Generalize transform

Generalize used to fail on ops that have a null region builder.
This is incorrect, the test should be whether the op has a region or not.
In the future we may want to support 0-region ops with a region builder.

Differential Revision: https://reviews.llvm.org/D147166

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
    mlir/test/Dialect/Linalg/transform-op-generalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
index 4755fa305da31..ed1f8f6eaaee4 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
@@ -39,9 +39,12 @@ static LogicalResult generalizeNamedOpPrecondition(LinalgOp linalgOp) {
   // Check if the operation is a LinalgOp but not a GenericOp.
   if (isa<GenericOp>(linalgOp))
     return failure();
-  // Check if the operation has a region builder.
-  if (!linalgOp.getRegionBuilder())
+  // Check if the operation has exactly one region.
+  if (linalgOp->getNumRegions() != 1) {
+    assert(linalgOp->getNumRegions() == 0 && "op with multiple regions");
+    // TOD: Otherwise it needs to be built explicitly from the region builder.
     return failure();
+  }
   return success();
 }
 

diff  --git a/mlir/test/Dialect/Linalg/transform-op-generalize.mlir b/mlir/test/Dialect/Linalg/transform-op-generalize.mlir
index c83ef4c6224a9..17dd3a3508933 100644
--- a/mlir/test/Dialect/Linalg/transform-op-generalize.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-generalize.mlir
@@ -10,8 +10,23 @@ func.func @generalize_unary(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> t
   return %0 : tensor<?x?xf32>
 }
 
+// CHECK-LABEL: func @map_no_inputs(
+func.func @map_no_inputs(%input: tensor<16x32x64xf32>,
+                  %init: tensor<16x64xf32>) -> tensor<16x64xf32> {
+  // CHECK-NOT:   linalg.map
+  //     CHECK:   linalg.generic
+  %reduce = linalg.reduce
+      ins(%input:tensor<16x32x64xf32>)
+      outs(%init:tensor<16x64xf32>)
+      dimensions = [1]
+      (%in: f32, %out: f32) {
+        %0 = arith.addf %out, %in: f32
+        linalg.yield %0: f32
+      }
+  func.return %reduce : tensor<16x64xf32>
+}
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
-  %0 = transform.structured.match ops{["linalg.elemwise_unary"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  %0 = transform.structured.match interface{LinalgOp} in %arg1 : (!pdl.operation) -> !pdl.operation
   %1 = transform.structured.generalize %0
 }