[mlir][CSE] Remove duplicated operations with MemRead side-effect

[clementval]

Summary

This patch enhances the CSE pass to deal with simple cases of duplicated
operations with MemoryEffects.

It allows the CSE pass to remove safely duplicate operations with the
MemoryEffects::Read that have no other side-effecting operations in
between. Other MemoryEffects::Read operation are allowed.

The use case is pretty simple so far so we can build on top of it to add
more features.

This patch is also meant to avoid a dedicated CSE pass in FIR and was
brought together afetr discussion on https://reviews.llvm.org/D112711.
It does not currently cover the full range of use cases described in
https://reviews.llvm.org/D112711 but the idea is to gradually enhance
the MLIR CSE pass to handle common use cases that can be used by
other dialects.

This patch takes advantage of the new CSE capabilities in Fir.

[clementval]

Remove unused template

[mehdi_amini]

mlir/lib/Transforms/CSE.cpp

148

Can you assert(fromOp->getParent() == toOp->getParent()); ?

151

I likely deserve a TODO: handles effect generically because there could be non-MemoryEffects, River WDYT?

160

Left over debug

204–225

[clementval]

Address review comments

[clementval]

mlir/lib/Transforms/CSE.cpp

148

Did you mean assert(fromOp->getBlock() == toOp->getBlock());? Operation does not have getParent() if I checked correctly.

204–225

Good catch! Thanks.

[mehdi_amini]

mlir/lib/Transforms/CSE.cpp

148

Right! There is actually a getParent() but it is private to Operation...

217

I wonder if there is an approach to cache some information and avoid a linear traversal of the range, I'm slightly worried we'll encountered some O(n^2) kind of patterns in large block that could make it explode here.

That may be fine to start with the naive approach though.

[rriddle]

Thanks for tackling this! Starting to build this out has been a long standing TODO.

typo in description?: remove samfely duplicates

mlir/lib/Transforms/CSE.cpp

138

Can we drop the trivial braces here now?

151

A TODO is fine, though I don't think we want other effects to have an implicit influence on MemoryEffects. Ideally any memory related things that would impact this are specified by MemoryEffects themselves (whether that includes adding additional effects or not remains unknown).

155

Do we care about Allocate and Free here?

217

+1, can we at least have some minimal caching here? LLVM EarlyCSE has a concept of "memory generation" which we could likely implement a simple form of here? (at least until we have better aliasing/analysis infra in MLIR)

[clementval]

Address some comments

mlir/lib/Transforms/CSE.cpp

151

I updated the TODO here but I guess in this particular case we really care about MemoryEffects only.

155

Not in this case. Just removed.
Maybe the function should be renamed hasWriteOpInBetween() since we just check that now. What do you think?

217

Ok I'll have a look at this and make another update for this patch. Thanks for the pointer.

[clementval]

Update toy tests

[clementval]

Add caching

[mehdi_amini]

mlir/lib/Transforms/CSE.cpp

147

I think another precondition here now is

assert( isa<MemoryEffectOpInterface>(toOp) && cast<MemoryEffectOpInterface>(toOp).hasEffect<MemoryEffects::Write>()) ?

150

You're hitting twice the map here, one lookup should be enough:

auto memEffectsCacheInfo = memEffectsCache.find(fromOp);
if (memEffectsCacheInfo != memEffectsCache.end()) {

[rriddle]

mlir/lib/Transforms/CSE.cpp

63

Can you document this?

65

Can we drop the mlir:: here?

144

Can you move the documentation to the declaration of the function?

150

try_emplace would avoid the additional lookup(via operator[]) further below.

164

If the operation doesn't implement the interface, we need to treat it conservatively (i.e. assume it writes). Can you add a test for this?

266

This is likely worth a comment.

[clementval]

Address review comments

[clementval]

mlir/lib/Transforms/CSE.cpp

147

fromOp and toOp have the MemoryEffects::Read and we are looking for write in between.

[rriddle]

mlir/lib/Transforms/CSE.cpp

64–65

operation until MemoryEffects information is valid

Can you tweak this a bit? It's a bit difficult for my head to parse.

211

I think the second part of this is off, the operation could also allocate/write to memory. Should this be !memEffects.onlyHasEffect<...>()`?

[clementval]

Address review comments

mlir/lib/Transforms/CSE.cpp

211

You are right. We want onlyHasEffect here.

[rriddle]

This seems like a good start, we can beef up the caching/analysis as necessary. Thanks for pushing this! Really awesome to see this being developed.

mlir/test/Transforms/cse.mlir

291–304

Can we trim this a bit? Are all of these necessary?

[clementval]

Thanks for the review @mehdi_amini @rriddle

mlir/test/Transforms/cse.mlir

291–304

I can trim it a bit before landing.

[mehdi_amini]

mlir/lib/Transforms/CSE.cpp

147

Mmmm... OK but then there is something else: if toOp does not have MemoryEffects::Write, it may not be in the map.
But that means that we may jump ahead of it using the map, and then wouldn't it be possible for the loop line 165 to go past the end of the block and crash?
With something like:

read // A - fromOp
...
read // B - toOp
...
write // C

If the first read is in the map, then memEffectsCache.find(fromOp); could return C and so nextOp will be already beyond toOp?

Should we do while (nextOp && nextOp.isBeforeInBlock(toOp)) instead?

[clementval]

mlir/lib/Transforms/CSE.cpp

147

We never have write op in the map (at least not with the current code). We just have the information if there is write op between two read operations. fromOp and toOp are always the same kind of operation with the read memory effect.

In the example below in the first call we will look between fromOp 1 and toOp 1. In the second call, without cache we would have to look between fromOp 2 and toOp 2. With the cache we can start directly from toOp 1. So we check between toOp 1 and toOp 2.

read // A - fromOp 1 / fromOp2
...
read // B - toOp 1
...
write // C
...
read // D - toOp 2

[mehdi_amini]

mlir/lib/Transforms/CSE.cpp

147

OK, I missed the final update line 179!
There is also an implicit assumption that you call this function with <fromOp 1, topOp1> first and <fromOp 1, topOp2> second right? (I was reading this function in isolation)

150

I don't think you implemented it? River is right that try_emplace here would get you the iterator to use for the updates line 173 and line 179.

[clementval]

mlir/lib/Transforms/CSE.cpp

147

My bad. We can have write op in the map but in this case we exit directly.

1. We check between fromOp1 and toOp 1. We store the write in cache.
2. We check between fromOp2 and toOp 2. We get the information from the cache that there is write op in between.

read // A - fromOp 1 / fromOp2
...
write // C
...
read // B - toOp 1
...
read // D - toOp 2

[clementval]

mlir/lib/Transforms/CSE.cpp

150

Right. Let me update that.

[schweitz]

Good to see these improvements going in to MLIR's CSE, Valentin.

[clementval]

• Use try_emplace
• Trim test with multiple load

[mehdi_amini]

Still LGTM :)

flang/include/flang/Optimizer/Dialect/FIROps.td

let description = [{
     ```
   }];
 
   let arguments = (ins Arg<fir_HeapType, "", [MemFree]>:$heapref);
 
   let assemblyFormat = "$heapref attr-dict `:` qualified(type($heapref))";
 }
 
-def fir_LoadOp : fir_OneResultOp<"load"> {
+def fir_LoadOp : fir_OneResultOp<"load", [MemoryEffects<[MemRead]>]> {
   let summary = "load a value from a memory reference";
   let description = [{
     Load a value from a memory reference into an ssa-value (virtual register).
     Produces an immutable ssa-value of the referent type. A memory reference
     has type `!fir.ref<T>`, `!fir.heap<T>`, or `!fir.ptr<T>`.
 
     ```mlir
       %a = fir.alloca i32
def fir_CharConvertOp : fir_Op<"char_convert", []> {
 
   let assemblyFormat = [{
     $from `for` $count `to` $to attr-dict `:` type(operands)
   }];
 
   let hasVerifier = 1;
 }
 
-def fir_StoreOp : fir_Op<"store", []> {
+def fir_StoreOp : fir_Op<"store", [MemoryEffects<[MemWrite]>]> {
   let summary = "store an SSA-value to a memory location";
 
   let description = [{
     Store an ssa-value (virtual register) to a memory reference.  The stored
     value must be of the same type as the referent type of the memory
     reference.
 
     ```mlir

flang/test/Fir/cse.fir

+// RUN: fir-opt --cse -split-input-file %s | FileCheck %s
+
+// Check that the redundant fir.load is removed.
+func @fun(%arg0: !fir.ref<i64>) -> i64 {
+    %0 = fir.load %arg0 : !fir.ref<i64>
+    %1 = fir.load %arg0 : !fir.ref<i64>
+    %2 = arith.addi %0, %1 : i64
+    return %2 : i64
+}
+
+// CHECK-LABEL: func @fun
+// CHECK-NEXT:    %[[LOAD:.*]] = fir.load %{{.*}} : !fir.ref<i64>
+// CHECK-NEXT:    %{{.*}} = arith.addi %[[LOAD]], %[[LOAD]] : i64
+
+// -----
+
+// CHECK-LABEL: func @fun(
+// CHECK-SAME:            %[[A:.*]]: !fir.ref<i64>
+func @fun(%a : !fir.ref<i64>) -> i64 {
+  // CHECK: %[[LOAD:.*]] = fir.load %[[A]] : !fir.ref<i64>
+  %1 = fir.load %a : !fir.ref<i64>
+  %2 = fir.load %a : !fir.ref<i64>
+  // CHECK-NEXT: %{{.*}} = arith.addi %[[LOAD]], %[[LOAD]] : i64
+  %3 = arith.addi %1, %2 : i64
+  %4 = fir.load %a : !fir.ref<i64>
+  // CHECK-NEXT: %{{.*}} = arith.addi
+  %5 = arith.addi %3, %4 : i64
+  %6 = fir.load %a : !fir.ref<i64>
+  // CHECK-NEXT: %{{.*}} = arith.addi
+  %7 = arith.addi %5, %6 : i64
+  %8 = fir.load %a : !fir.ref<i64>
+  // CHECK-NEXT: %{{.*}} = arith.addi
+  %9 = arith.addi %7, %8 : i64
+  %10 = fir.load %a : !fir.ref<i64>
+  // CHECK-NEXT: %{{.*}} = arith.addi
+  %11 = arith.addi %10, %9 : i64
+  %12 = fir.load %a : !fir.ref<i64>
+  // CHECK-NEXT: %{{.*}} = arith.addi
+  %13 = arith.addi %11, %12 : i64
+  // CHECK-NEXT: return %{{.*}} : i64
+  return %13 : i64
+}
+
+// -----
+
+func @fun(%a : !fir.ref<i64>) -> i64 {
+  cf.br ^bb1
+^bb1:
+  %1 = fir.load %a : !fir.ref<i64>
+  %2 = fir.load %a : !fir.ref<i64>
+  %3 = arith.addi %1, %2 : i64
+  cf.br ^bb2
+^bb2:
+  %4 = fir.load %a : !fir.ref<i64>
+  %5 = arith.subi %4, %4 : i64
+  return %5 : i64
+}

mlir/lib/Transforms/CSE.cpp

 struct CSE : public CSEBase<CSE> {
   /// Shared implementation of operation elimination and scoped map definitions.
   using AllocatorTy = llvm::RecyclingAllocator<
       llvm::BumpPtrAllocator,
       llvm::ScopedHashTableVal<Operation *, Operation *>>;
   using ScopedMapTy = llvm::ScopedHashTable<Operation *, Operation *,
                                             SimpleOperationInfo, AllocatorTy>;
 
+  /// Cache holding MemoryEffects information between two operations. The first
+  /// operation is stored has the key. The second operation is stored inside a
+  /// pair in the value. The pair also hold the MemoryEffects between those
+  /// two operations. If the MemoryEffects is nullptr then we assume there is
+  /// no operation with MemoryEffects::Write between the two operations.
+  using MemEffectsCache =
+      DenseMap<Operation *, std::pair<Operation *, MemoryEffects::Effect *>>;
+
   /// Represents a single entry in the depth first traversal of a CFG.
   struct CFGStackNode {
     CFGStackNode(ScopedMapTy &knownValues, DominanceInfoNode *node)
         : scope(knownValues), node(node), childIterator(node->begin()) {}
 
     /// Scope for the known values.
     ScopedMapTy::ScopeTy scope;
 
struct CSE : public CSEBase<CSE> {
   LogicalResult simplifyOperation(ScopedMapTy &knownValues, Operation *op,
                                   bool hasSSADominance);
   void simplifyBlock(ScopedMapTy &knownValues, Block *bb, bool hasSSADominance);
   void simplifyRegion(ScopedMapTy &knownValues, Region &region);
 
   void runOnOperation() override;
 
 private:
+  void replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
+                            Operation *existing, bool hasSSADominance);
+
+  /// Check if there is side-effecting operations other than the given effect
+  /// between the two operations.
+  bool hasOtherSideEffectingOpInBetween(Operation *fromOp, Operation *toOp);
+
   /// Operations marked as dead and to be erased.
   std::vector<Operation *> opsToErase;
   DominanceInfo *domInfo = nullptr;
+  MemEffectsCache memEffectsCache;
 };
 } // namespace
 
-/// Attempt to eliminate a redundant operation.
-LogicalResult CSE::simplifyOperation(ScopedMapTy &knownValues, Operation *op,
-                                     bool hasSSADominance) {
-  // Don't simplify terminator operations.
-  if (op->hasTrait<OpTrait::IsTerminator>())
-    return failure();
-
-  // If the operation is already trivially dead just add it to the erase list.
-  if (isOpTriviallyDead(op)) {
-    opsToErase.push_back(op);
-    ++numDCE;
-    return success();
-  }
-
-  // Don't simplify operations with nested blocks. We don't currently model
-  // equality comparisons correctly among other things. It is also unclear
-  // whether we would want to CSE such operations.
-  if (op->getNumRegions() != 0)
-    return failure();
-
-  // TODO: We currently only eliminate non side-effecting
-  // operations.
-  if (!MemoryEffectOpInterface::hasNoEffect(op))
-    return failure();
-
-  // Look for an existing definition for the operation.
-  if (auto *existing = knownValues.lookup(op)) {
-
-    // If we find one then replace all uses of the current operation with the
-    // existing one and mark it for deletion. We can only replace an operand in
-    // an operation if it has not been visited yet.
-    if (hasSSADominance) {
-      // If the region has SSA dominance, then we are guaranteed to have not
-      // visited any use of the current operation.
-      op->replaceAllUsesWith(existing);
-      opsToErase.push_back(op);
-    } else {
-      // When the region does not have SSA dominance, we need to check if we
-      // have visited a use before replacing any use.
-      for (auto it : llvm::zip(op->getResults(), existing->getResults())) {
-        std::get<0>(it).replaceUsesWithIf(
-            std::get<1>(it), [&](OpOperand &operand) {
-              return !knownValues.count(operand.getOwner());
-            });
-      }
+void CSE::replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
+                               Operation *existing, bool hasSSADominance) {
 
-      // There may be some remaining uses of the operation.
-      if (op->use_empty())
-        opsToErase.push_back(op);
-    }
 
-    // If the existing operation has an unknown location and the current
-    // operation doesn't, then set the existing op's location to that of the
-    // current op.
-    if (existing->getLoc().isa<UnknownLoc>() &&
-        !op->getLoc().isa<UnknownLoc>()) {
-      existing->setLoc(op->getLoc());
-    }
+  if (existing->getLoc().isa<UnknownLoc>() && !op->getLoc().isa<UnknownLoc>())
+
+  ++numCSE;
 
+bool CSE::hasOtherSideEffectingOpInBetween(Operation *fromOp, Operation *toOp) {
+  assert(fromOp->getBlock() == toOp->getBlock());
+  assert(
+      isa<MemoryEffectOpInterface>(fromOp) &&
+      cast<MemoryEffectOpInterface>(fromOp).hasEffect<MemoryEffects::Read>() &&
+      isa<MemoryEffectOpInterface>(toOp) &&
+      cast<MemoryEffectOpInterface>(toOp).hasEffect<MemoryEffects::Read>());
+  Operation *nextOp = fromOp->getNextNode();
+  auto result =
+      memEffectsCache.try_emplace(fromOp, std::make_pair(fromOp, nullptr));
+  if (result.second) {
+    auto memEffectsCachePair = result.first->second;
+    if (memEffectsCachePair.second == nullptr) {
+      // No MemoryEffects::Write has been detected until the cached operation.
+      // Continue looking from the cached operation to toOp.
+      nextOp = memEffectsCachePair.first;
+    } else {
+      // MemoryEffects::Write has been detected before so there is no need to
+      // check further.
+      return true;
+    }
+  }
+  while (nextOp && nextOp != toOp) {
+    auto nextOpMemEffects = dyn_cast<MemoryEffectOpInterface>(nextOp);
+    // TODO: Do we need to handle other effects generically?
+    // If the operation does not implement the MemoryEffectOpInterface we
+    // conservatively assumes it writes.
+    if ((nextOpMemEffects &&
+         nextOpMemEffects.hasEffect<MemoryEffects::Write>()) ||
+        !nextOpMemEffects) {
+      result.first->second =
+          std::make_pair(nextOp, MemoryEffects::Write::get());
+      return true;
+    }
+    nextOp = nextOp->getNextNode();
+  }
+  result.first->second = std::make_pair(toOp, nullptr);
+  return false;
+}
+
+/// Attempt to eliminate a redundant operation.
+LogicalResult CSE::simplifyOperation(ScopedMapTy &knownValues, Operation *op,
+                                     bool hasSSADominance) {
+  // Don't simplify terminator operations.
+  if (op->hasTrait<OpTrait::IsTerminator>())
+    return failure();
+
+  // If the operation is already trivially dead just add it to the erase list.
+  if (isOpTriviallyDead(op)) {
+    opsToErase.push_back(op);
+    ++numDCE;
+    return success();
+  }
+
+  // Don't simplify operations with nested blocks. We don't currently model
+  // equality comparisons correctly among other things. It is also unclear
+  // whether we would want to CSE such operations.
+  if (op->getNumRegions() != 0)
+    return failure();
+
+  // Some simple use case of operation with memory side-effect are dealt with
+  // here. Operations with no side-effect are done after.
+  if (!MemoryEffectOpInterface::hasNoEffect(op)) {
+    auto memEffects = dyn_cast<MemoryEffectOpInterface>(op);
+    // TODO: Only basic use case for operations with MemoryEffects::Read can be
+    // eleminated now. More work needs to be done for more complicated patterns
+    // and other side-effects.
+    if (!memEffects || !memEffects.onlyHasEffect<MemoryEffects::Read>())
+      return failure();
+
+    // Look for an existing definition for the operation.
+    if (auto *existing = knownValues.lookup(op)) {
+      if (existing->getBlock() == op->getBlock() &&
+          !hasOtherSideEffectingOpInBetween(existing, op)) {
+        // The operation that can be deleted has been reach with no
+        // side-effecting operations in between the existing operation and
+        // this one so we can remove the duplicate.
+        replaceUsesAndDelete(knownValues, op, existing, hasSSADominance);
+        return success();
+      }
+    }
+    knownValues.insert(op, op);
+    return failure();
+  }
+
+  // Look for an existing definition for the operation.
+  if (auto *existing = knownValues.lookup(op)) {
+    replaceUsesAndDelete(knownValues, op, existing, hasSSADominance);
     ++numCSE;
     return success();
   }
 
   // Otherwise, we add this operation to the known values map.
   knownValues.insert(op, op);
   return failure();
 }
if (op.mightHaveTrait<OpTrait::IsIsolatedFromAbove>()) {
         simplifyRegion(nestedKnownValues, region);
       continue;
     }
 
     // Otherwise, process nested regions normally.
     for (auto &region : op.getRegions())
       simplifyRegion(knownValues, region);
   }
+  // Clear the MemoryEffects cache since its usage is by block only.
+  memEffectsCache.clear();
 }
 
 void CSE::simplifyRegion(ScopedMapTy &knownValues, Region &region) {
   // If the region is empty there is nothing to do.
   if (region.empty())
     return;
 
   bool hasSSADominance = domInfo->hasSSADominance(&region);

mlir/test/Examples/Toy/Ch5/affine-lowering.mlir

 // CHECK:         affine.store [[VAL_5]], [[VAL_8]][1, 2] : memref<2x3xf64>
 // CHECK:         affine.for [[VAL_9:%.*]] = 0 to 3 {
 // CHECK:           affine.for [[VAL_10:%.*]] = 0 to 2 {
 // CHECK:             [[VAL_11:%.*]] = affine.load [[VAL_8]]{{\[}}[[VAL_10]], [[VAL_9]]] : memref<2x3xf64>
 // CHECK:             affine.store [[VAL_11]], [[VAL_7]]{{\[}}[[VAL_9]], [[VAL_10]]] : memref<3x2xf64>
 // CHECK:         affine.for [[VAL_12:%.*]] = 0 to 3 {
 // CHECK:           affine.for [[VAL_13:%.*]] = 0 to 2 {
 // CHECK:             [[VAL_14:%.*]] = affine.load [[VAL_7]]{{\[}}[[VAL_12]], [[VAL_13]]] : memref<3x2xf64>
-// CHECK:             [[VAL_15:%.*]] = affine.load [[VAL_7]]{{\[}}[[VAL_12]], [[VAL_13]]] : memref<3x2xf64>
-// CHECK:             [[VAL_16:%.*]] = arith.mulf [[VAL_14]], [[VAL_15]] : f64
+// CHECK:             [[VAL_16:%.*]] = arith.mulf [[VAL_14]], [[VAL_14]] : f64
 // CHECK:             affine.store [[VAL_16]], [[VAL_6]]{{\[}}[[VAL_12]], [[VAL_13]]] : memref<3x2xf64>
 // CHECK:         toy.print [[VAL_6]] : memref<3x2xf64>
 // CHECK:         memref.dealloc [[VAL_8]] : memref<2x3xf64>
 // CHECK:         memref.dealloc [[VAL_7]] : memref<3x2xf64>
 // CHECK:         memref.dealloc [[VAL_6]] : memref<3x2xf64>
 
 // OPT-LABEL: func @main()
 // OPT-DAG:     [[VAL_0:%.*]] = arith.constant 1.000000e+00 : f64

mlir/test/Examples/Toy/Ch6/affine-lowering.mlir

 // CHECK:         affine.store [[VAL_5]], [[VAL_8]][1, 2] : memref<2x3xf64>
 // CHECK:         affine.for [[VAL_9:%.*]] = 0 to 3 {
 // CHECK:           affine.for [[VAL_10:%.*]] = 0 to 2 {
 // CHECK:             [[VAL_11:%.*]] = affine.load [[VAL_8]]{{\[}}[[VAL_10]], [[VAL_9]]] : memref<2x3xf64>
 // CHECK:             affine.store [[VAL_11]], [[VAL_7]]{{\[}}[[VAL_9]], [[VAL_10]]] : memref<3x2xf64>
 // CHECK:         affine.for [[VAL_12:%.*]] = 0 to 3 {
 // CHECK:           affine.for [[VAL_13:%.*]] = 0 to 2 {
 // CHECK:             [[VAL_14:%.*]] = affine.load [[VAL_7]]{{\[}}[[VAL_12]], [[VAL_13]]] : memref<3x2xf64>
-// CHECK:             [[VAL_15:%.*]] = affine.load [[VAL_7]]{{\[}}[[VAL_12]], [[VAL_13]]] : memref<3x2xf64>
-// CHECK:             [[VAL_16:%.*]] = arith.mulf [[VAL_14]], [[VAL_15]] : f64
+// CHECK:             [[VAL_16:%.*]] = arith.mulf [[VAL_14]], [[VAL_14]] : f64
 // CHECK:             affine.store [[VAL_16]], [[VAL_6]]{{\[}}[[VAL_12]], [[VAL_13]]] : memref<3x2xf64>
 // CHECK:         toy.print [[VAL_6]] : memref<3x2xf64>
 // CHECK:         memref.dealloc [[VAL_8]] : memref<2x3xf64>
 // CHECK:         memref.dealloc [[VAL_7]] : memref<3x2xf64>
 // CHECK:         memref.dealloc [[VAL_6]] : memref<3x2xf64>
 
 // OPT-LABEL: func @main()
 // OPT-DAG:     [[VAL_0:%.*]] = arith.constant 1.000000e+00 : f64

mlir/test/Examples/Toy/Ch7/affine-lowering.mlir

 // CHECK:         affine.store [[VAL_5]], [[VAL_8]][1, 2] : memref<2x3xf64>
 // CHECK:         affine.for [[VAL_9:%.*]] = 0 to 3 {
 // CHECK:           affine.for [[VAL_10:%.*]] = 0 to 2 {
 // CHECK:             [[VAL_11:%.*]] = affine.load [[VAL_8]]{{\[}}[[VAL_10]], [[VAL_9]]] : memref<2x3xf64>
 // CHECK:             affine.store [[VAL_11]], [[VAL_7]]{{\[}}[[VAL_9]], [[VAL_10]]] : memref<3x2xf64>
 // CHECK:         affine.for [[VAL_12:%.*]] = 0 to 3 {
 // CHECK:           affine.for [[VAL_13:%.*]] = 0 to 2 {
 // CHECK:             [[VAL_14:%.*]] = affine.load [[VAL_7]]{{\[}}[[VAL_12]], [[VAL_13]]] : memref<3x2xf64>
-// CHECK:             [[VAL_15:%.*]] = affine.load [[VAL_7]]{{\[}}[[VAL_12]], [[VAL_13]]] : memref<3x2xf64>
-// CHECK:             [[VAL_16:%.*]] = arith.mulf [[VAL_14]], [[VAL_15]] : f64
+// CHECK:             [[VAL_16:%.*]] = arith.mulf [[VAL_14]], [[VAL_14]] : f64
 // CHECK:             affine.store [[VAL_16]], [[VAL_6]]{{\[}}[[VAL_12]], [[VAL_13]]] : memref<3x2xf64>
 // CHECK:         toy.print [[VAL_6]] : memref<3x2xf64>
 // CHECK:         memref.dealloc [[VAL_8]] : memref<2x3xf64>
 // CHECK:         memref.dealloc [[VAL_7]] : memref<3x2xf64>
 // CHECK:         memref.dealloc [[VAL_6]] : memref<3x2xf64>
 
 // OPT-LABEL: func @main()
 // OPT-DAG:     [[VAL_0:%.*]] = arith.constant 1.000000e+00 : f64

mlir/test/Transforms/cse.mlir

test.graph_region {
     %1 = arith.constant 1 : i32
     %2 = arith.constant 1 : i32
 
     // CHECK-NEXT: "foo.yield"(%0) : (i32) -> ()
     "foo.yield"(%0) : (i32) -> ()
   }
   return
 } 
+
+/// This test is checking that CSE is removing duplicated read op that follow
+/// other.
+// CHECK-LABEL: @remove_direct_duplicated_read_op
+func @remove_direct_duplicated_read_op() -> i32 {
+  // CHECK-NEXT: %[[READ_VALUE:.*]] = "test.op_with_memread"() : () -> i32
+  %0 = "test.op_with_memread"() : () -> (i32)
+  %1 = "test.op_with_memread"() : () -> (i32)
+  // CHECK-NEXT: %{{.*}} = arith.addi %[[READ_VALUE]], %[[READ_VALUE]] : i32
+  %2 = arith.addi %0, %1 : i32
+  return %2 : i32
+}
+
+/// This test is checking that CSE is removing duplicated read op that follow
+/// other.
+// CHECK-LABEL: @remove_multiple_duplicated_read_op
+func @remove_multiple_duplicated_read_op() -> i64 {
+  // CHECK: %[[READ_VALUE:.*]] = "test.op_with_memread"() : () -> i64
+  %0 = "test.op_with_memread"() : () -> (i64)
+  %1 = "test.op_with_memread"() : () -> (i64)
+  // CHECK-NEXT: %{{.*}} = arith.addi %{{.*}}, %[[READ_VALUE]] : i64
+  %2 = arith.addi %0, %1 : i64
+  %3 = "test.op_with_memread"() : () -> (i64)
+  // CHECK-NEXT: %{{.*}} = arith.addi %{{.*}}, %{{.*}} : i64
+  %4 = arith.addi %2, %3 : i64
+  %5 = "test.op_with_memread"() : () -> (i64)
+  // CHECK-NEXT: %{{.*}} = arith.addi %{{.*}}, %{{.*}} : i64
+  %6 = arith.addi %4, %5 : i64
+  // CHECK-NEXT: return %{{.*}} : i64
+  return %6 : i64
+}
+
+/// This test is checking that CSE is not removing duplicated read op that
+/// have write op in between.
+// CHECK-LABEL: @dont_remove_duplicated_read_op_with_sideeffecting
+func @dont_remove_duplicated_read_op_with_sideeffecting() -> i32 {
+  // CHECK-NEXT: %[[READ_VALUE0:.*]] = "test.op_with_memread"() : () -> i32
+  %0 = "test.op_with_memread"() : () -> (i32)
+  "test.op_with_memwrite"() : () -> ()
+  // CHECK: %[[READ_VALUE1:.*]] = "test.op_with_memread"() : () -> i32
+  %1 = "test.op_with_memread"() : () -> (i32)
+  // CHECK-NEXT: %{{.*}} = arith.addi %[[READ_VALUE0]], %[[READ_VALUE1]] : i32
+  %2 = arith.addi %0, %1 : i32
+  return %2 : i32
+}

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

let arguments = (ins
         );
 
     let results = (outs Res<AnyMemRef, "", [MemAlloc<TestResource>]>);
 }
 
 def TestEffectsOpB : TEST_Op<"op_with_effects_b",
     [MemoryEffects<[MemWrite<TestResource>]>]>;
 
+def TestEffectsRead : TEST_Op<"op_with_memread",
+    [MemoryEffects<[MemRead]>]> {
+  let results = (outs AnyInteger);
+}
+
+def TestEffectsWrite : TEST_Op<"op_with_memwrite",
+    [MemoryEffects<[MemWrite]>]>;
+
 #endif // TEST_OPS