[mlir][Linalg] Fix Generalize transform

[nicolasvasilache]

Summary

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.

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

 
 using namespace mlir;
 using namespace mlir::linalg;
 
 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();
 }
 
 FailureOr<GenericOp> mlir::linalg::generalizeNamedOp(RewriterBase &rewriter,
                                                      LinalgOp linalgOp) {
   if (failed(generalizeNamedOpPrecondition(linalgOp)))
     return rewriter.notifyMatchFailure(linalgOp, "preconditions not met");
 

mlir/test/Dialect/Linalg/transform-op-generalize.mlir

 // RUN: mlir-opt --test-transform-dialect-interpreter %s | FileCheck %s
 
 // CHECK-LABEL: func.func @generalize_unary
 func.func @generalize_unary(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
 
   // CHECK-NOT:   linalg.elemwise_unary
   //     CHECK:   linalg.generic
   %0 = linalg.elemwise_unary ins(%arg0 : tensor<?x?xf32>)
                              outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
   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
 }