[Mlir-commits] [mlir] [mlir] [linalg] fix side effect of linalg op (PR #114045)

[donald chen]

https://github.com/cxy-1993 updated https://github.com/llvm/llvm-project/pull/114045

>From 3d4f708e2cfe261c2ddfb654abcf23cce6ea95a4 Mon Sep 17 00:00:00 2001
From: donald chen <chenxunyu1993 at gmail.com>
Date: Tue, 29 Oct 2024 12:57:12 +0000
Subject: [PATCH] [mlir] [linalg] fix side effect of linalg op

Linalg op need to take into account memory side effects happening
inside the region when determining their own side effects.

This patch fixed issue https://github.com/llvm/llvm-project/issues/112881
---
 .../Dialect/Linalg/IR/LinalgStructuredOps.td   |  1 +
 mlir/test/Dialect/Linalg/canonicalize.mlir     | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index bfc609bd708164..c2fee8ea55c960 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -30,6 +30,7 @@ class LinalgStructuredBase_Op<string mnemonic, list<Trait> props>
        SingleBlockImplicitTerminator<"YieldOp">,
        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
        DeclareOpInterfaceMethods<ConditionallySpeculatable>,
+       RecursiveMemoryEffects,
        DestinationStyleOpInterface,
        LinalgStructuredInterface,
        ReifyRankedShapedTypeOpInterface], props)> {
diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index 4bc2ed140da91a..d15d8f6a475ec3 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -1232,3 +1232,21 @@ func.func @transpose_buffer(%input: memref<?xf32>,
 //  CHECK-SAME:            %[[VAL_1:.*]]: memref<?xf32>) {
 //       CHECK:     linalg.transpose ins(%[[VAL_0]] : memref<?xf32>)
 //  CHECK-SAME:       outs(%[[VAL_1]] : memref<?xf32>) permutation = [0]
+
+// -----
+
+// This test checks linalg op has a recursive memory effect. Otherwise
+// linalg.map without user would not be deleted in canonicalize.
+func.func @recursive_effect(%arg0: memref<1xf32>, %arg1 : tensor<1xf32>) {
+  %c1 = arith.constant 1 : index
+  %init = arith.constant dense<0.0> : tensor<1xf32>
+  %mapped = linalg.map ins(%arg1:tensor<1xf32>) outs(%init :tensor<1xf32>)
+            (%in : f32) {
+              memref.store %in, %arg0[%c1] : memref<1xf32>
+              linalg.yield %in : f32
+            }
+  func.return
+}
+
+// CHECK-LABEL: @recursive_effect
+//       CHECK: linalg.map