[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 dabfcdb0689d31ada16d7e137006eb802e51fd69 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..4cd45bd62abd1f 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 a user would be DCEd.
+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