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

[zero9178]

Summary

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.

[Mogball]

This makes sense to me, but you might want someone else's opinion.

mlir/lib/Interfaces/SideEffectInterfaces.cpp

66–67

[rriddle]

mlir/lib/Interfaces/SideEffectInterfaces.cpp

67–81

This by itself isn't safe, because it's also making an assumption about aliasing. It's assuming that the op has properly deduced what effects should be placed on inputs/etc., which may not necessarily be the case. I think the driving example should be fine though, i.e. if the value is also known to be an alloc, it should be fine regardless of what other effects were placed on it (given that it's a newly constructed value).

[zero9178]

Change the logic to only apply to results that are also allocated

[mehdi_amini]

mlir/lib/Interfaces/SideEffectInterfaces.cpp

72

Why are you precomputing this in allocResults instead of doing it directly at the point of checking?

[mehdi_amini]

mlir/lib/Interfaces/SideEffectInterfaces.cpp

72

Forget it, I figured: when we have an effect on a result, we don't know if this particular result is also allocated otherwise. The issue is that effects aren't grouped...

mlir/test/Transforms/canonicalize-dce.mlir

180

The comment is outdated

[zero9178]

mlir/test/Transforms/canonicalize-dce.mlir

180

In the previous diff the comment read "Delete ops that only have side-effects on an result.".

I have since updated it and added "allocated result", which is the current state of the patch and accurate I believe.

mlir/lib/Interfaces/SideEffectInterfaces.cpp

 //===- SideEffectInterfaces.cpp - SideEffects in MLIR ---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
+#include "llvm/ADT/SmallPtrSet.h"
+
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // SideEffect Interfaces
 //===----------------------------------------------------------------------===//
 
 /// Include the definitions of the side effect interfaces.
 #include "mlir/Interfaces/SideEffectInterfaces.cpp.inc"
while (!effectingOps.empty()) {
     }
 
     // If the op has memory effects, try to characterize them to see if the op
     // is trivially dead here.
     if (auto effectInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
       // Check to see if this op either has no effects, or only allocates/reads
       // 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());
           })) {
         return false;
       }
       continue;
 
       // Otherwise, if the op has recursive side effects we can treat the
       // operation itself as having no effects.

mlir/test/Transforms/canonicalize-dce.mlir

 func.func @f() {
   test.graph_region {
     %0 = "math.exp"(%1) : (f32) -> f32
     %1 = "math.exp"(%0) : (f32) -> f32
     "test.terminator"() : ()->()
   }
   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

 def TestEffectsRead : TEST_Op<"op_with_memread",
     [MemoryEffects<[MemRead]>]> {
   let results = (outs AnyInteger);
 }
 
 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
 //===----------------------------------------------------------------------===//
 
 def TestVerifiersOp : TEST_Op<"verifiers",
                               [SingleBlock, NoTerminator, IsolatedFromAbove]> {
   let arguments = (ins I32:$input);
   let regions = (region SizedRegion<1>:$region);