[mlir] Fix missing check on nested op values in LICM

LICM checks that nested ops depend only on values defined outside
before performing hoisting.
However, it specifically omits to check for terminators which can
lead to SSA violations.
This revision fixes the incorrect behavior.

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

mlir/lib/Transforms/LoopInvariantCodeMotion.cpp

@@ -66,15 +66,14 @@ static bool canBeHoisted(Operation *op,
   // can be hoisted.
   for (auto &region : op->getRegions()) {
     for (auto &block : region) {
-      for (auto &innerOp : block.without_terminator())
+      for (auto &innerOp : block)
         if (!canBeHoisted(&innerOp, definedOutside))
           return false;
     }
   }
   return true;
 }
 
-
 LogicalResult mlir::moveLoopInvariantCode(LoopLikeOpInterface looplike) {
   auto &loopBody = looplike.getLoopBody();
 

mlir/test/Transforms/loop-invariant-code-motion.mlir

@@ -292,3 +292,33 @@ func @parallel_loop_with_invariant() {
   return
 }
 
+// -----
+
+func private @make_val() -> (index)
+
+// CHECK-LABEL: func @nested_uses_inside
+func @nested_uses_inside(%lb: index, %ub: index, %step: index) {
+  %true = arith.constant true
+
+  // Check that ops that contain nested uses to values not defiend outside 
+  // remain in the loop.
+  // CHECK-NEXT: arith.constant
+  // CHECK-NEXT: scf.for
+  // CHECK-NEXT:   call @
+  // CHECK-NEXT:   call @
+  // CHECK-NEXT:   scf.if
+  // CHECK-NEXT:     scf.yield
+  // CHECK-NEXT:   else
+  // CHECK-NEXT:     scf.yield
+  scf.for %i = %lb to %ub step %step {
+    %val = call @make_val() : () -> (index)
+    %val2 = call @make_val() : () -> (index)
+    %r = scf.if %true -> (index) {
+      scf.yield %val: index
+    } else {
+      scf.yield %val2: index
+    }
+  }
+  return 
+}
+