[mlir] Fix region successor bug in forward dataflow analysis

We weren't properly visiting region successors when the terminator wasn't return like, which could create incorrect results in the analysis. This revision ensures that we properly visit region successors, to avoid optimistically assuming a value is constant when it isn't.

Differential Revision: https://reviews.llvm.org/D101783

diff --git a/mlir/lib/Analysis/DataFlowAnalysis.cpp b/mlir/lib/Analysis/DataFlowAnalysis.cpp
index ef629c1..fde9f2c 100644 (file)
--- a/mlir/lib/Analysis/DataFlowAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlowAnalysis.cpp
@@ -478,8 +478,10 @@ void ForwardDataFlowSolver::visitRegionSuccessors(
       // aren't part of the control flow.
       if (succArgs.size() != results.size()) {
         opWorklist.push_back(parentOp);
-        if (succArgs.empty())
-          return markAllPessimisticFixpoint(results);
+        if (succArgs.empty()) {
+          markAllPessimisticFixpoint(results);
+          continue;
+        }
 
         unsigned firstResIdx = succArgs[0].cast<OpResult>().getResultNumber();
         markAllPessimisticFixpoint(results.take_front(firstResIdx));
@@ -492,7 +494,7 @@ void ForwardDataFlowSolver::visitRegionSuccessors(
       for (auto it : llvm::zip(succArgs, operands))
         join(parentOp, analysis.getLatticeElement(std::get<0>(it)),
              analysis.getLatticeElement(std::get<1>(it)));
-      return;
+      continue;
     }
     assert(!region->empty() && "expected region to be non-empty");
     Block *entryBlock = &region->front();
@@ -563,8 +565,16 @@ void ForwardDataFlowSolver::visitTerminatorOperation(
     // If this terminator is not "region-like", conservatively mark all of the
     // successor values as having reached the pessimistic fixpoint.
     if (!op->hasTrait<OpTrait::ReturnLike>()) {
-      for (auto &it : regionSuccessors)
-        markAllPessimisticFixpointAndVisitUsers(it.getSuccessorInputs());
+      for (auto &it : regionSuccessors) {
+        // If the successor is a region, mark the entry block as executable so
+        // that we visit operations defined within. If the successor is the
+        // parent operation, we simply mark the control flow results as having
+        // reached the pessimistic state.
+        if (Region *region = it.getSuccessor())
+          markEntryBlockExecutable(region, /*markPessimisticFixpoint=*/true);
+        else
+          markAllPessimisticFixpointAndVisitUsers(it.getSuccessorInputs());
+      }
       return;
     }
 

diff --git a/mlir/test/Transforms/sccp-structured.mlir b/mlir/test/Transforms/sccp-structured.mlir
index 45227cd..02b2365 100644 (file)
--- a/mlir/test/Transforms/sccp-structured.mlir
+++ b/mlir/test/Transforms/sccp-structured.mlir
@@ -130,3 +130,23 @@ func @loop_inner_control_flow(%arg0 : index, %arg1 : index, %arg2 : index) -> i3
   }
   return %result : i32
 }
+
+/// Test that we can properly visit region successors when the terminator is not
+/// return-like.
+
+// CHECK-LABEL: func @overdefined_non_returnlike(
+func @overdefined_non_returnlike(%arg1 : i32) {
+  // CHECK: scf.while (%[[ARG:.*]] = %[[INPUT:.*]])
+  // CHECK-NEXT: %[[COND:.*]] = cmpi slt, %[[ARG]], %{{.*}} : i32
+  // CHECK-NEXT: scf.condition(%[[COND]]) %[[ARG]] : i32
+
+  %c2_i32 = constant 2 : i32
+   %0 = scf.while (%arg2 = %c2_i32) : (i32) -> (i32) {
+    %1 = cmpi slt, %arg2, %arg1 : i32
+    scf.condition(%1) %arg2 : i32
+  } do {
+  ^bb0(%arg2: i32):
+    scf.yield %arg2 : i32
+  }
+  return
+}