[Mlir-commits] [mlir] [MLIR] Add a getStaticTripCount method to LoopLikeOpInterface (PR #158679)

[Mehdi Amini]

================
@@ -264,22 +274,102 @@ getValuesSortedByKey(ArrayRef<Attribute> keys, ArrayRef<int64_t> values,
 
 /// Return the number of iterations for a loop with a lower bound `lb`, upper
 /// bound `ub` and step `step`.
-std::optional<int64_t> constantTripCount(OpFoldResult lb, OpFoldResult ub,
-                                         OpFoldResult step) {
+std::optional<APInt> constantTripCount(
+    OpFoldResult lb, OpFoldResult ub, OpFoldResult step, bool isSigned,
+    llvm::function_ref<std::optional<llvm::APSInt>(Value, Value, bool)>
+        computeUbMinusLb) {
+  // This is the bitwidth used to return 0 when loop does not execute.
+  // We infer it from the type of the bound if it isn't an index type.
+  int bitwidth = IndexType::kInternalStorageBitWidth;
+  bool isIndex = true;
+  if (auto lbAttr = dyn_cast<Attribute>(lb)) {
+    if (auto intAttr = dyn_cast<IntegerAttr>(lbAttr)) {
+      if (auto intType = dyn_cast<IntegerType>(intAttr.getType())) {
+        bitwidth = intType.getWidth();
+        isIndex = intType.isIndex();
+      }
+    }
+  } else {
+    auto val = cast<Value>(lb);
+    if (auto intType = dyn_cast<IntegerType>(val.getType())) {
+      bitwidth = intType.getWidth();
+      isIndex = intType.isIndex();
+    }
+  }
+
   if (lb == ub)
-    return 0;
+    return APInt(bitwidth, 0);
 
-  std::optional<int64_t> lbConstant = getConstantIntValue(lb);
-  if (!lbConstant)
-    return std::nullopt;
-  std::optional<int64_t> ubConstant = getConstantIntValue(ub);
-  if (!ubConstant)
-    return std::nullopt;
-  std::optional<int64_t> stepConstant = getConstantIntValue(step);
-  if (!stepConstant || *stepConstant == 0)
-    return std::nullopt;
+  std::optional<std::pair<APInt, bool>> maybeStepCst =
+      getConstantAPIntValue(step);
+
+  if (maybeStepCst) {
+    auto &stepCst = maybeStepCst->first;
+    if (stepCst.isZero())
+      return stepCst;
+    if (stepCst.isNegative())
+      return APInt(bitwidth, 0);
+  }
+
+  if (isIndex) {
+    LDBG()
+        << "Computing loop trip count for index type may break with overflow";
+    // TODO: we can't compute the trip count for index type. We should fix this
----------------
joker-eph wrote:

You may have missed this thread I opened: https://discourse.llvm.org/t/index-type-and-assumption-about-bitwidth/88287 :)

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