[mlir] Fix region successor bug in forward dataflow analysis

[rriddle]

Summary

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.

[mehdi_amini]

mlir/lib/Analysis/DataFlowAnalysis.cpp

570

This deserves a comment, I don't quite get it immediately

[rriddle]

update

mlir/lib/Analysis/DataFlowAnalysis.cpp

if (!region) {
 
       // Mark the results outside of the input range as having reached the
       // pessimistic fixpoint.
       // TODO: This isn't exactly ideal. There may be situations in which a
       // region operation can provide information for certain results that
       // 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));
         markAllPessimisticFixpoint(
             results.drop_front(firstResIdx + succArgs.size()));
       }
 
       // Update the lattice for any operation results.
       OperandRange operands = getInputsForRegion(/*index=*/llvm::None);
       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();
     markBlockExecutable(entryBlock);
 
     // If all of the arguments have already reached a fixpoint, the arguments
     // have already been fully resolved.
     Block::BlockArgListType arguments = entryBlock->getArguments();
analysis.getSuccessorsForOperands(regionInterface,
                                       parentRegion->getRegionNumber(),
                                       operandLattices, regionSuccessors);
     if (regionSuccessors.empty())
       return;
 
     // 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)
+      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;
     }
 
     // Otherwise, propagate the operand lattice states to the successors.
     OperandRange operands = op->getOperands();
     return visitRegionSuccessors(parentOp, regionSuccessors,
                                  [&](Optional<unsigned>) { return operands; });
   }

mlir/test/Transforms/sccp-structured.mlir

%result = scf.for %i0 = %arg0 to %arg1 step %arg2 iter_args(%si = %cst_1) -> (i32) {
     } else {
       %si_inc = addi %si, %cst_1 : i32
       scf.yield %si_inc : i32
     }
     scf.yield %inner_res : i32
   }
   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
+}