[mlir][cse] do not replace operands in previously simplified operations

[youngar]

Summary

If an operation has been inserted as a key in to the known values
hashtable, then it can not be modified in a way which changes its hash.
This change avoids modifying the operands of any previously recorded
operation, which prevents their hash from changing.

In an SSACFG region, it is impossible to visit an operation before
visiting its operands, so this is not a problem. This situation can only
happen in regions without strict dominance, such as graph regions.

[rriddle]

mlir/lib/Transforms/CSE.cpp

119–139

Can you check if this is a graph region before doing this? We should add additional checks when necessary.

[youngar]

Check if the region does not have SSA dominance before using the slow-path for replacing uses.

[youngar]

mlir/lib/Transforms/CSE.cpp

119–139

(I meant to submit this comment much earlier)
This makes sense. Since I think unregistered operations have relaxed dominance, I think I should steal the check from Dominance.cpp:

/// Return true if the region with the given index inside the operation
/// has SSA dominance.
static bool hasSSADominance(Operation *op, unsigned index) {
  auto kindInterface = dyn_cast<RegionKindInterface>(op);
  return op->isRegistered() &&
         (!kindInterface || kindInterface.hasSSADominance(index));
}

Would it make sense to make this function public?

[rriddle]

mlir/lib/Transforms/CSE.cpp

119–139

We already check for this in simplifyRegion with domInfo.hasDominanceInfo. Can you pipe through that down to here? i.e. something like:

LogicalResult CSE::simplifyOperation(ScopedMapTy &knownValues, Operation *op, bool hasSSADominance) {
  ...
}

void CSE::simplifyBlock(ScopedMapTy &knownValues, DominanceInfo &domInfo,
                        Block *bb, bool hasSSADominance) {
  for (auto &inst : *bb) {
    if (succeeded(simplifyOperation(knownValues, &inst, hasSSADominance)))
      continue;
  }
}

void CSE::simplifyRegion(ScopedMapTy &knownValues, DominanceInfo &domInfo,
                         Region &region) {
  ...

  bool hasSSADominance = domInfo.hasDominanceInfo(&region);

  // If the region only contains one block, then simplify it directly.
  if (std::next(region.begin()) == region.end()) {
    ScopedMapTy::ScopeTy scope(knownValues);
    simplifyBlock(knownValues, domInfo, &region.front(), hasSSADominance);
    return;
  }

  // If the region does not have dominanceInfo, then skip it.
  // TODO: Regions without SSA dominance should define a different
  // traversal order which is appropriate and can be used here.
  if (!hasSSADominance)
    return;

  ...
}

[youngar]

Use dominfo->hasDominanceInfo to determinine if the region has dominance. Funnel this information through function arguments.

[rriddle]

Thanks for fixing this!

mlir/lib/Transforms/CSE.cpp

129

nit: Can you use llvm::zip(existing->getResults(), op->getResults()) here?

[youngar]

Use llvm::zip

mlir/lib/Transforms/CSE.cpp

struct CFGStackNode {
     DominanceInfoNode::const_iterator childIterator;
 
     /// If this node has been fully processed yet or not.
     bool processed;
   };
 
   /// Attempt to eliminate a redundant operation. Returns success if the
   /// operation was marked for removal, failure otherwise.
-  LogicalResult simplifyOperation(ScopedMapTy &knownValues, Operation *op);
-
+  LogicalResult simplifyOperation(ScopedMapTy &knownValues, Operation *op,
+                                  bool hasSSADominance);
   void simplifyBlock(ScopedMapTy &knownValues, DominanceInfo &domInfo,
-                     Block *bb);
+                     Block *bb, bool hasSSADominance);
   void simplifyRegion(ScopedMapTy &knownValues, DominanceInfo &domInfo,
                       Region &region);
 
   void runOnOperation() override;
 
 private:
   /// Operations marked as dead and to be erased.
   std::vector<Operation *> opsToErase;
 };
 } // end anonymous namespace
 
 /// Attempt to eliminate a redundant operation.
-LogicalResult CSE::simplifyOperation(ScopedMapTy &knownValues, Operation *op) {
+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.
+    // 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());
+            });
+      }
+
+      // 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());
     }
 
     ++numCSE;
     return success();
   }
 
   // Otherwise, we add this operation to the known values map.
   knownValues.insert(op, op);
   return failure();
 }
 
 void CSE::simplifyBlock(ScopedMapTy &knownValues, DominanceInfo &domInfo,
-                        Block *bb) {
+                        Block *bb, bool hasSSADominance) {
   for (auto &inst : *bb) {
     // If the operation is simplified, we don't process any held regions.
-    if (succeeded(simplifyOperation(knownValues, &inst)))
+    if (succeeded(simplifyOperation(knownValues, &inst, hasSSADominance)))
       continue;
 
     // If this operation is isolated above, we can't process nested regions with
     // the given 'knownValues' map. This would cause the insertion of implicit
     // captures in explicit capture only regions.
     if (inst.mightHaveTrait<OpTrait::IsIsolatedFromAbove>()) {
       ScopedMapTy nestedKnownValues;
       for (auto &region : inst.getRegions())
         simplifyRegion(nestedKnownValues, domInfo, region);
       continue;
     }
 
     // Otherwise, process nested regions normally.
     for (auto &region : inst.getRegions())
       simplifyRegion(knownValues, domInfo, region);
   }
 }
 
 void CSE::simplifyRegion(ScopedMapTy &knownValues, DominanceInfo &domInfo,
                          Region &region) {
   // If the region is empty there is nothing to do.
   if (region.empty())
     return;
 
+  bool hasSSADominance = domInfo.hasDominanceInfo(&region);
+
   // If the region only contains one block, then simplify it directly.
   if (std::next(region.begin()) == region.end()) {
     ScopedMapTy::ScopeTy scope(knownValues);
-    simplifyBlock(knownValues, domInfo, &region.front());
+    simplifyBlock(knownValues, domInfo, &region.front(), hasSSADominance);
     return;
   }
 
   // If the region does not have dominanceInfo, then skip it.
   // TODO: Regions without SSA dominance should define a different
   // traversal order which is appropriate and can be used here.
-  if (!domInfo.hasDominanceInfo(&region))
+  if (!hasSSADominance)
     return;
 
   // Note, deque is being used here because there was significant performance
   // gains over vector when the container becomes very large due to the
   // specific access patterns. If/when these performance issues are no
   // longer a problem we can change this to vector. For more information see
   // the llvm mailing list discussion on this:
   // http://lists.llvm.org/pipermail/llvm-commits/Week-of-Mon-20120116/135228.html
   std::deque<std::unique_ptr<CFGStackNode>> stack;
 
   // Process the nodes of the dom tree for this region.
   stack.emplace_back(std::make_unique<CFGStackNode>(
       knownValues, domInfo.getRootNode(&region)));
 
   while (!stack.empty()) {
     auto &currentNode = stack.back();
 
     // Check to see if we need to process this node.
     if (!currentNode->processed) {
       currentNode->processed = true;
-      simplifyBlock(knownValues, domInfo, currentNode->node->getBlock());
+      simplifyBlock(knownValues, domInfo, currentNode->node->getBlock(),
+                    hasSSADominance);
     }
 
     // Otherwise, check to see if we need to process a child node.
     if (currentNode->childIterator != currentNode->node->end()) {
       auto *childNode = *(currentNode->childIterator++);
       stack.emplace_back(
           std::make_unique<CFGStackNode>(knownValues, childNode));
     } else {

mlir/test/Transforms/cse.mlir

func @nested_isolated() -> i32 {
   "foo.region"() ({
     // CHECK-NEXT: constant 1
     %foo = constant 1 : i32
     "foo.yield"(%foo) : (i32) -> ()
   }) : () -> ()
 
   return %0 : i32
 }
+
+/// This test is checking that CSE gracefully handles values in graph regions
+/// where the use occurs before the def, and one of the defs could be CSE'd with
+/// the other.
+// CHECK-LABEL: @use_before_def
+func @use_before_def() {
+  // CHECK-NEXT: test.graph_region
+  test.graph_region {
+    // CHECK-NEXT: addi %c1_i32, %c1_i32_0
+    %0 = addi %1, %2 : i32
+
+    // CHECK-NEXT: constant 1
+    // CHECK-NEXT: constant 1
+    %1 = constant 1 : i32
+    %2 = constant 1 : i32
+
+    // CHECK-NEXT: "foo.yield"(%0) : (i32) -> ()
+    "foo.yield"(%0) : (i32) -> ()
+  }
+  return
+}