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

[nicolasvasilache]

Summary

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.

[herhut]

Great find. Thanks for the fix.

mlir/lib/Transforms/LoopInvariantCodeMotion.cpp

static bool canBeHoisted(Operation *op,
   } else if (!op->hasTrait<OpTrait::HasRecursiveSideEffects>()) {
     return false;
   }
 
   // Recurse into the regions for this op and check whether the contained ops
   // 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();
 
   // We use two collections here as we need to preserve the order for insertion
   // and this is easiest.
   SmallPtrSet<Operation *, 8> willBeMovedSet;
   SmallVector<Operation *, 8> opsToMove;
 

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

func @parallel_loop_with_invariant() {
   // CHECK-NEXT:   arith.addi %arg0, %arg1 : index
   // CHECK-NEXT:   yield
   // CHECK-NEXT: }
   // CHECK-NEXT: return
 
   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 
+}
+