[mlir] Ignore effects on allocated results when checking whether the op is trivially dead.

In the current state, this is only special cased for Allocation effects, but any effects on results allocated by the operation may be ignored when checking whether the op may be removed, as none of them are possible to be observed if the result is unused.

A use case for this is for IRs for languages which always initialize on allocation. To correctly model such operations, a Write as well as an Allocation effect should be placed on the result. This would prevent the Op from being deleted if unused however. This patch fixes that issue.

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

mlir/lib/Interfaces/SideEffectInterfaces.cpp

@@ -8,6 +8,8 @@
 
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
+#include "llvm/ADT/SmallPtrSet.h"
+
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
@@ -62,11 +64,20 @@ static bool wouldOpBeTriviallyDeadImpl(Operation *rootOp) {
       // memory.
       SmallVector<MemoryEffects::EffectInstance, 1> effects;
       effectInterface.getEffects(effects);
-      if (!llvm::all_of(effects, [op](const MemoryEffects::EffectInstance &it) {
-            // We can drop allocations if the value is a result of the
-            // operation.
-            if (isa<MemoryEffects::Allocate>(it.getEffect()))
-              return it.getValue() && it.getValue().getDefiningOp() == op;
+
+      // Gather all results of this op that are allocated.
+      SmallPtrSet<Value, 4> allocResults;
+      for (const MemoryEffects::EffectInstance &it : effects)
+        if (isa<MemoryEffects::Allocate>(it.getEffect()) && it.getValue() &&
+            it.getValue().getDefiningOp() == op)
+          allocResults.insert(it.getValue());
+
+      if (!llvm::all_of(effects, [&allocResults](
+                                     const MemoryEffects::EffectInstance &it) {
+            // We can drop effects if the value is an allocation and is a result
+            // of the operation.
+            if (allocResults.contains(it.getValue()))
+              return true;
             // Otherwise, the effect must be a read.
             return isa<MemoryEffects::Read>(it.getEffect());
           })) {

mlir/test/Transforms/canonicalize-dce.mlir

@@ -173,3 +173,17 @@ func.func @f() {
   }
   return
 }
+
+// -----
+
+
+// Test case: Delete ops that only have side-effects on an allocated result.
+
+// CHECK:      func @f()
+// CHECK-NOT:    test_effects_result
+// CHECK-NEXT:   return
+
+func.func @f() {
+  %0 = "test.test_effects_result"() : () -> i32
+  return
+}

mlir/test/lib/Dialect/Test/TestOps.td

@@ -2867,6 +2867,10 @@ def TestEffectsRead : TEST_Op<"op_with_memread",
 def TestEffectsWrite : TEST_Op<"op_with_memwrite",
     [MemoryEffects<[MemWrite]>]>;
 
+def TestEffectsResult : TEST_Op<"test_effects_result"> {
+  let results = (outs Res<I32, "", [MemAlloc, MemWrite]>);
+}
+
 //===----------------------------------------------------------------------===//
 // Test Ops with verifiers
 //===----------------------------------------------------------------------===//