diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp --- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp +++ b/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 @@ // memory. SmallVector 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(it.getEffect())) - return it.getValue() && it.getValue().getDefiningOp() == op; + + // Gather all results of this op that are allocated. + SmallPtrSet allocResults; + for (const MemoryEffects::EffectInstance &it : effects) + if (isa(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(it.getEffect()); })) { diff --git a/mlir/test/Transforms/canonicalize-dce.mlir b/mlir/test/Transforms/canonicalize-dce.mlir --- a/mlir/test/Transforms/canonicalize-dce.mlir +++ b/mlir/test/Transforms/canonicalize-dce.mlir @@ -173,3 +173,17 @@ } 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 +} diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td --- a/mlir/test/lib/Dialect/Test/TestOps.td +++ b/mlir/test/lib/Dialect/Test/TestOps.td @@ -2867,6 +2867,10 @@ def TestEffectsWrite : TEST_Op<"op_with_memwrite", [MemoryEffects<[MemWrite]>]>; +def TestEffectsResult : TEST_Op<"test_effects_result"> { + let results = (outs Res); +} + //===----------------------------------------------------------------------===// // Test Ops with verifiers //===----------------------------------------------------------------------===//