[Mlir-commits] [mlir] [mlir][Affine] Fix LICM incorrectly hoisting stores from zero-trip-count loops (PR #189165)

[Mehdi Amini]

================
@@ -174,13 +174,22 @@ void LoopInvariantCodeMotion::runOnAffineForOp(AffineForOp forOp) {
   SmallVector<Operation *, 8> opsToMove;
   SmallPtrSet<Operation *, 8> opsWithUsers;
 
+  // If the trip count is statically known to be zero, do not hoist ops with
+  // side effects: the loop body never executes, so hoisting would cause them
+  // to run unconditionally, changing program semantics. Pure (side-effect-free
+  // and speculatable) ops are still eligible for hoisting.
+  std::optional<uint64_t> tripCount = getConstantTripCount(forOp);
+  bool zeroTripCount = tripCount.has_value() && *tripCount == 0;
----------------
joker-eph wrote:

Completely, let me make it actually conservative (done: see the test in @unknown_trip_count_store_not_hoisted)

I'm wondering if we could plug some integer range analysis here to be more aggressive... It's a bit unfortunate that we only have a query for "getConstantTripCount" when we just need "executeAtLeastOnce". This would make it for a nice extra API on the LoopOpInterface I think.

That said the challenge then is to make it inter-operate with more complex IR analyses (like integer range) or other symbolic relationships.

https://github.com/llvm/llvm-project/pull/189165