[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 5c0502150f206113ecf8dc374bcb16847f1e0f99 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      | 17 +++++++++++++++++
 2 files changed, 18 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..5de007b390c51d 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -1232,3 +1232,20 @@ 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(%arg : tensor<1xf32>) {
+  %init = arith.constant dense<0.0> : tensor<1xf32>
+  %mapped = linalg.map ins(%arg:tensor<1xf32>) outs(%init :tensor<1xf32>)
+            (%in : f32) {
+              vector.print %in : f32
+              linalg.yield %in : f32
+            }
+  func.return
+}
+
+// CHECK-LABEL: @recursive_effect
+//       CHECK: linalg.map