Update CSE to respect nested regions that are isolated from above. This cl also remov...

PiperOrigin-RevId: 254709704

diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index 4407ec9..e0f839d 100644 (file)
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -721,22 +721,6 @@ public:
   }
 };
 
-/// This verifies that all operands used in N-th region of the given operation
-/// are defined within that region. This is a weaker variant to the
-/// 'IsIsolatedFromAbove' trait above as it applies only to one specific region.
-template <unsigned RegionIdx> class NthRegionIsIsolatedFromAbove {
-public:
-  template <typename ConcreteType>
-  class Impl : public TraitBase<ConcreteType,
-                                NthRegionIsIsolatedFromAbove<RegionIdx>::Impl> {
-  public:
-    static LogicalResult verifyTrait(Operation *op) {
-      return success(
-          op->getRegion(RegionIdx).isIsolatedFromAbove(op->getLoc()));
-    }
-  };
-};
-
 } // end namespace OpTrait
 
 //===----------------------------------------------------------------------===//

diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index 49931cd..f39f127 100644 (file)
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -79,6 +79,10 @@ public:
     return getName().getAbstractOperation();
   }
 
+  /// Returns true if this operation has a registered operation description,
+  /// otherwise false.
+  bool isRegistered() { return getAbstractOperation(); }
+
   /// Remove this operation from its parent block and delete it.
   void erase();
 

diff --git a/mlir/lib/Transforms/CSE.cpp b/mlir/lib/Transforms/CSE.cpp
index f90e12d..188db62 100644 (file)
--- a/mlir/lib/Transforms/CSE.cpp
+++ b/mlir/lib/Transforms/CSE.cpp
@@ -110,41 +110,40 @@ struct CSE : public FunctionPass<CSE> {
     bool processed;
   };
 
-  /// Attempt to eliminate a redundant operation. Returns true if the operation
-  /// was marked for removal, false otherwise.
-  bool simplifyOperation(Operation *op);
+  /// Attempt to eliminate a redundant operation. Returns success if the
+  /// operation was marked for removal, failure otherwise.
+  LogicalResult simplifyOperation(ScopedMapTy &knownValues, Operation *op);
 
-  void simplifyBlock(DominanceInfo &domInfo, Block *bb);
-  void simplifyRegion(DominanceInfo &domInfo, Region &region);
+  void simplifyBlock(ScopedMapTy &knownValues, DominanceInfo &domInfo,
+                     Block *bb);
+  void simplifyRegion(ScopedMapTy &knownValues, DominanceInfo &domInfo,
+                      Region &region);
 
   void runOnFunction() override;
 
 private:
-  /// A scoped hash table of defining operations within a function.
-  ScopedMapTy knownValues;
-
   /// Operations marked as dead and to be erased.
   std::vector<Operation *> opsToErase;
 };
 } // end anonymous namespace
 
 /// Attempt to eliminate a redundant operation.
-bool CSE::simplifyOperation(Operation *op) {
+LogicalResult CSE::simplifyOperation(ScopedMapTy &knownValues, Operation *op) {
   // 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 false;
+    return failure();
 
   // TODO(riverriddle) We currently only eliminate non side-effecting
   // operations.
   if (!op->hasNoSideEffect())
-    return false;
+    return failure();
 
   // If the operation is already trivially dead just add it to the erase list.
   if (op->use_empty()) {
     opsToErase.push_back(op);
-    return true;
+    return success();
   }
 
   // Look for an existing definition for the operation.
@@ -162,27 +161,39 @@ bool CSE::simplifyOperation(Operation *op) {
         !op->getLoc().isa<UnknownLoc>()) {
       existing->setLoc(op->getLoc());
     }
-    return true;
+    return success();
   }
 
   // Otherwise, we add this operation to the known values map.
   knownValues.insert(op, op);
-  return false;
+  return failure();
 }
 
-void CSE::simplifyBlock(DominanceInfo &domInfo, Block *bb) {
-  for (auto &i : *bb) {
+void CSE::simplifyBlock(ScopedMapTy &knownValues, DominanceInfo &domInfo,
+                        Block *bb) {
+  for (auto &inst : *bb) {
     // If the operation is simplified, we don't process any held regions.
-    if (simplifyOperation(&i))
+    if (succeeded(simplifyOperation(knownValues, &inst)))
+      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.isRegistered() || inst.isKnownIsolatedFromAbove()) {
+      ScopedMapTy nestedKnownValues;
+      for (auto &region : inst.getRegions())
+        simplifyRegion(nestedKnownValues, domInfo, region);
       continue;
+    }
 
-    // Simplify any held blocks.
-    for (auto &region : i.getRegions())
-      simplifyRegion(domInfo, region);
+    // Otherwise, process nested regions normally.
+    for (auto &region : inst.getRegions())
+      simplifyRegion(knownValues, domInfo, region);
   }
 }
 
-void CSE::simplifyRegion(DominanceInfo &domInfo, Region &region) {
+void CSE::simplifyRegion(ScopedMapTy &knownValues, DominanceInfo &domInfo,
+                         Region &region) {
   // If the region is empty there is nothing to do.
   if (region.empty())
     return;
@@ -190,7 +201,7 @@ void CSE::simplifyRegion(DominanceInfo &domInfo, Region &region) {
   // If the region only contains one block, then simplify it directly.
   if (std::next(region.begin()) == region.end()) {
     ScopedMapTy::ScopeTy scope(knownValues);
-    simplifyBlock(domInfo, &region.front());
+    simplifyBlock(knownValues, domInfo, &region.front());
     return;
   }
 
@@ -212,7 +223,7 @@ void CSE::simplifyRegion(DominanceInfo &domInfo, Region &region) {
     // Check to see if we need to process this node.
     if (!currentNode->processed) {
       currentNode->processed = true;
-      simplifyBlock(domInfo, currentNode->node->getBlock());
+      simplifyBlock(knownValues, domInfo, currentNode->node->getBlock());
     }
 
     // Otherwise, check to see if we need to process a child node.
@@ -229,13 +240,14 @@ void CSE::simplifyRegion(DominanceInfo &domInfo, Region &region) {
 }
 
 void CSE::runOnFunction() {
-  simplifyRegion(getAnalysis<DominanceInfo>(), getFunction().getBody());
+  /// A scoped hash table of defining operations within a function.
+  ScopedMapTy knownValues;
+  simplifyRegion(knownValues, getAnalysis<DominanceInfo>(),
+                 getFunction().getBody());
 
   // If no operations were erased, then we mark all analyses as preserved.
-  if (opsToErase.empty()) {
-    markAllAnalysesPreserved();
-    return;
-  }
+  if (opsToErase.empty())
+    return markAllAnalysesPreserved();
 
   /// Erase any operations that were marked as dead during simplification.
   for (auto *op : opsToErase)

diff --git a/mlir/test/Transforms/cse.mlir b/mlir/test/Transforms/cse.mlir
index b0ba7b9..e398ecc 100644 (file)
--- a/mlir/test/Transforms/cse.mlir
+++ b/mlir/test/Transforms/cse.mlir
@@ -220,3 +220,27 @@ func @up_propagate_region() -> i32 {
   }) : () -> (i32)
   return %0 : i32
 }
+
+/// This test checks that nested regions that are isolated from above are
+/// properly handled.
+// CHECK-LABEL: @nested_isolated
+func @nested_isolated() -> i32 {
+  // CHECK-NEXT: constant 1
+  %0 = constant 1 : i32
+
+  // CHECK-NEXT: @nested_func
+  func @nested_func() {
+    // CHECK-NEXT: constant 1
+    %foo = constant 1 : i32
+    "foo.yield"(%foo) : (i32) -> ()
+  }
+
+  // CHECK: "foo.region"
+  "foo.region"() ({
+    // CHECK-NEXT: constant 1
+    %foo = constant 1 : i32
+    "foo.yield"(%foo) : (i32) -> ()
+  }) : () -> ()
+
+  return %0 : i32
+}