[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 =
+      isSigned ? diff.sdiv(extendedStepCst) : diff.udiv(extendedStepCst);
+  llvm::APInt remainder =
+      isSigned ? diff.srem(extendedStepCst) : diff.urem(extendedStepCst);
   if (!remainder.isZero())
     tripCount = tripCount + 1;
+
+  // Truncate back to original bitwidth if we extended for overflow prevention.
+  // This is safe because ceil(diff/step) ≤ 2^bitwidth - 1, which always fits
+  // in bitwidth bits when interpreted as unsigned (trip counts are inherently
+  // non-negative regardless of loop comparison signedness).
+  if (tripCount.getBitWidth() > bitwidth)
----------------
matthias-springer wrote:

No truncation needed anymore here.

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