[Mlir-commits] [mlir] 09d09fc - [MLIR][Linalg] Avoid generalizing `linalg.map`

[Lorenzo Chelini]

Author: Lorenzo Chelini
Date: 2023-06-26T11:55:56+02:00
New Revision: 09d09fc335ac8255d2c3c80048922ee5c68c06c3

URL: https://github.com/llvm/llvm-project/commit/09d09fc335ac8255d2c3c80048922ee5c68c06c3
DIFF: https://github.com/llvm/llvm-project/commit/09d09fc335ac8255d2c3c80048922ee5c68c06c3.diff

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

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.

Reviewed By: nicolasvasilache

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

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
    mlir/test/Dialect/Linalg/generalize-named-ops.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
index ed1f8f6eaaee4..146217d55b124 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
@@ -36,8 +36,10 @@ 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) {

diff  --git a/mlir/test/Dialect/Linalg/generalize-named-ops.mlir b/mlir/test/Dialect/Linalg/generalize-named-ops.mlir
index ad7980ba4ff01..f276953b20616 100644
--- a/mlir/test/Dialect/Linalg/generalize-named-ops.mlir
+++ b/mlir/test/Dialect/Linalg/generalize-named-ops.mlir
@@ -272,3 +272,17 @@ func.func @batch_reduce_gemm(%lhs: memref<7x8x9xf32>, %rhs: memref<7x9x8xf32>, %
 // 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
+}