[Mlir-commits] [mlir] [MLIR][Utils] Fix overflow in constantTripCount for narrow types (PR #179985)

[Matthias Springer]

================
@@ -397,11 +404,28 @@ std::optional<APInt> constantTripCount(
     return std::nullopt;
   }
 
+  // Extend stepCst to match the bitwidth of diff if needed (e.g., when diff was
+  // extended to avoid overflow). Step is always positive here, so zero-extend.
+  llvm::APInt extendedStepCst = stepCst;
+  if (extendedStepCst.getBitWidth() < diff.getBitWidth()) {
+    extendedStepCst = extendedStepCst.zext(diff.getBitWidth());
+  }
+
   // Create new APSInt instances with explicit signedness to ensure they match
-  llvm::APInt tripCount = isSigned ? diff.sdiv(stepCst) : diff.udiv(stepCst);
-  llvm::APInt remainder = isSigned ? diff.srem(stepCst) : diff.urem(stepCst);
+  llvm::APInt tripCount =
----------------
matthias-springer wrote:

At this point, we can assume that `diff` is non-negative (see above comment) and that `step` is non-negative, even if interpreted as signed integer. The reason is the above check:
```
  // For signed loops, a negative step size could indicate an infinite number of
  // iterations.
  if (isSigned && stepCst.isSignBitSet()) {
    LDBG() << "constantTripCount is infinite because step is negative";
    return std::nullopt;
  }
```

That means we don't have to distinguish between signed and unsigned semantics anymore at this point.

I.e., no need to check `isSigned`, use `udiv` and `urem` only. We need to document this reasoning here.

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