[MLIR][Linalg] Avoid generalizing `linalg.map`

[chelini]

Summary

We cannot trivially generalize linalg.map, as it does not use the
output as a region argument in the block, while linalg.generic expects
many region arguments as the input/output operands.

[nicolasvasilache]

Interesting, do you have a recommendation here ?

[chelini]

In D153442#4440213, @nicolasvasilache wrote:

Interesting, do you have a recommendation here ?

I would re-evaluate https://reviews.llvm.org/D134444 as it may create inconsistency between linalg.generic and named ops. I agree with the direction, but it is not just an ergonomics change. We could think of making linalg.generic consistent with this behaviour by dropping unused bb args or making linalg.named ops compatible with generics. Moving from named to generic (and vice-versa) is an essential trait of linalg and should be maintained, worth having an offline discussion.

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

 } // namespace mlir
 
 #define DEBUG_TYPE "linalg-generalization"
 
 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))
+  // Bailout if `linalgOp` is already a generic or a linalg.map. We cannot
+  // trivially generalize a `linalg.map`, as it does not use the output as
+  // region arguments in the block.
+  if (isa<GenericOp>(linalgOp) || isa<MapOp>(linalgOp))
     return failure();
   // 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();

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

 // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
 // CHECK-SAME: iterator_types = ["reduction", "parallel", "parallel", "reduction"]}
 // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<7x8x9xf32>, memref<7x9x8xf32>)
 // CHECK-SAME: outs(%{{.+}} : memref<8x8xf32>
 // CHECK:         ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
 // CHECK:         %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32
 // CHECK:         %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32
 // CHECK:         linalg.yield %[[ADD]] : f32
+
+// -----
+
+// CHECK-LABEL: generalize_linalg_map
+func.func @generalize_linalg_map(%arg0: memref<1x8x8x8xf32>) {
+  %cst = arith.constant 0.000000e+00 : f32
+  // CHECK: linalg.map
+  // CHECK-NOT: linalg.generic
+  linalg.map outs(%arg0 : memref<1x8x8x8xf32>)
+    () {
+      linalg.yield %cst : f32
+    }
+  return
+}