diff --git a/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp b/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp --- a/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp +++ b/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp @@ -66,7 +66,7 @@ // can be hoisted. for (auto ®ion : op->getRegions()) { for (auto &block : region) { - for (auto &innerOp : block.without_terminator()) + for (auto &innerOp : block) if (!canBeHoisted(&innerOp, definedOutside)) return false; } @@ -74,7 +74,6 @@ return true; } - LogicalResult mlir::moveLoopInvariantCode(LoopLikeOpInterface looplike) { auto &loopBody = looplike.getLoopBody(); diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir --- a/mlir/test/Transforms/loop-invariant-code-motion.mlir +++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir @@ -292,3 +292,33 @@ 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 +} +