[mlir][scf] Interpret trip counts as unsigned integers (#178060)

Trip counts represent iteration counts and are always non-negative. This
PR fixes all call sites to correctly use `getZExtValue()` instead of
`getSExtValue()` when extracting trip count values from `APInt`. Also
documents to clarify results are unsigned.

Author: jhalakpatel

mlir/include/mlir/Dialect/SCF/Utils/Utils.h

@@ -251,6 +251,7 @@ FailureOr<scf::ParallelOp> parallelLoopUnrollByFactors(
 /// Get constant trip counts for each of the induction variables of the given
 /// loop operation. If any of the loop's trip counts is not constant, return an
 /// empty vector.
+/// TODO(#178506): Should return SmallVector<uint64_t> for correct signedness.
 llvm::SmallVector<int64_t>
 getConstLoopTripCounts(mlir::LoopLikeOpInterface loopOp);
 

mlir/include/mlir/Dialect/Utils/StaticValueUtils.h

@@ -211,10 +211,11 @@ foldDynamicOffsetSizeList(SmallVectorImpl<OpFoldResult> &offsetsOrSizes);
 LogicalResult foldDynamicStrideList(SmallVectorImpl<OpFoldResult> &strides);
 
 /// Return the number of iterations for a loop with a lower bound `lb`, upper
-/// bound `ub` and step `step`. The `isSigned` flag indicates whether the loop
-/// comparison between lb and ub is signed or unsigned. A negative step or a
-/// lower bound greater than the upper bound are considered invalid and will
-/// yield a zero trip count.
+/// bound `ub` and step `step`, as an unsigned integer. The `isSigned` flag
+/// indicates whether the loop comparison between lb and ub is signed or
+/// unsigned. (The result of this function must be interpreted as an unsigned
+/// integer.) A lower bound greater than the upper bound is considered invalid
+/// and will yield a zero trip count.
 /// The `computeUbMinusLb` callback is invoked to compute the difference between
 /// the upper and lower bound when not constant. It can be used by the client
 /// to compute a static difference when the bounds are not constant.

mlir/lib/Dialect/SCF/IR/SCF.cpp

@@ -443,7 +443,7 @@ LogicalResult ForOp::promoteIfSingleIteration(RewriterBase &rewriter) {
   LDBG() << "promoteIfSingleIteration tripCount is " << tripCount
          << " for loop "
          << OpWithFlags(getOperation(), OpPrintingFlags().skipRegions());
-  if (!tripCount.has_value() || tripCount->getSExtValue() > 1)
+  if (!tripCount.has_value() || tripCount->getZExtValue() > 1)
     return failure();
 
   if (*tripCount == 0) {

mlir/lib/Dialect/SCF/Utils/Utils.cpp

@@ -397,9 +397,12 @@ FailureOr<UnrolledLoopInfo> mlir::loopUnrollByFactor(
       return UnrolledLoopInfo{forOp, std::nullopt};
     }
 
+    // TODO(#178506): This may overflow for large trip counts. Should use
+    // uint64_t.
     int64_t tripCountEvenMultiple =
-        constTripCount->getSExtValue() -
-        (constTripCount->getSExtValue() % unrollFactor);
+        constTripCount->getZExtValue() -
+        (constTripCount->getZExtValue() % unrollFactor);
+    // TODO(#178506): This may overflow when computing upperBoundUnrolledCst.
     int64_t upperBoundUnrolledCst = lbCst + tripCountEvenMultiple * stepCst;
     int64_t stepUnrolledCst = stepCst * unrollFactor;
 
@@ -501,9 +504,9 @@ LogicalResult mlir::loopUnrollFull(scf::ForOp forOp) {
   const APInt &tripCount = *mayBeConstantTripCount;
   if (tripCount.isZero())
     return success();
-  if (tripCount.getSExtValue() == 1)
+  if (tripCount.getZExtValue() == 1)
     return forOp.promoteIfSingleIteration(rewriter);
-  return loopUnrollByFactor(forOp, tripCount.getSExtValue());
+  return loopUnrollByFactor(forOp, tripCount.getZExtValue());
 }
 
 /// Check if bounds of all inner loops are defined outside of `forOp`
@@ -554,7 +557,7 @@ LogicalResult mlir::loopUnrollJamByFactor(scf::ForOp forOp,
               "trip "
               "count";
     unrollJamFactor = tripCount->getZExtValue();
-  } else if (tripCount->getSExtValue() % unrollJamFactor != 0) {
+  } else if (tripCount->getZExtValue() % unrollJamFactor != 0) {
     LDBG() << "failed to unroll and jam: unsupported trip count that is not a "
               "multiple of unroll jam factor";
     return failure();
@@ -1567,13 +1570,16 @@ mlir::getConstLoopTripCounts(mlir::LoopLikeOpInterface loopOp) {
   std::optional<SmallVector<OpFoldResult>> steps = loopOp.getLoopSteps();
   if (!loBnds || !upBnds || !steps)
     return {};
+  // TODO(#178506): The result should be SmallVector<uint64_t> and use uint64_t
+  // for trip counts.
   llvm::SmallVector<int64_t> tripCounts;
   for (auto [lb, ub, step] : llvm::zip(*loBnds, *upBnds, *steps)) {
+    // TODO(#178506): Signedness is not handled correctly here.
     std::optional<llvm::APInt> numIter = constantTripCount(
         lb, ub, step, /*isSigned=*/true, scf::computeUbMinusLb);
     if (!numIter)
       return {};
-    tripCounts.push_back(numIter->getSExtValue());
+    tripCounts.push_back(numIter->getZExtValue());
   }
   return tripCounts;
 }

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

@@ -337,8 +337,6 @@ std::optional<APInt> constantTripCount(
       // case applies, so the static trip count is unknown.
       return std::nullopt;
     }
-    if (stepCst.isNegative())
-      return APInt(bitwidth, 0);
   }
 
   if (isIndex) {
@@ -392,6 +390,14 @@ std::optional<APInt> constantTripCount(
     return std::nullopt;
   }
   auto &stepCst = maybeStepCst->first;
+  // 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;
+  }
+
+  // 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);
   if (!remainder.isZero())

mlir/test/Dialect/SCF/trip_count.mlir

@@ -699,4 +699,105 @@ func.func @trip_count_arith_add_nuw_loop_unsigned_invalid(%lb : i32, %other : i3
     scf.yield %arg0 : i32
   }
   return %1 : i32
+}
+
+// -----
+
+// CHECK-LABEL:func.func @trip_count_i8_unsigned_full_range(
+func.func @trip_count_i8_unsigned_full_range(%a : i32, %b : i32) -> i32 {
+  %c0 = arith.constant 0 : i8
+  %c255 = arith.constant 255 : i8
+  %c1 = arith.constant 1 : i8
+
+  // Unsigned i8 from 0 to 255: trip count is 255
+  // Trip counts are returned in their natural bitwidth and printed as signed.
+  // 255 in i8 is represented as -1 when printed in signed format.
+  // CHECK: "test.trip-count" = -1 : i8
+  %r = scf.for unsigned %i = %c0 to %c255 step %c1 iter_args(%0 = %a) -> i32 : i8 {
+    scf.yield %b : i32
+  }
+  return %r : i32
+}
+
+// -----
+
+// CHECK-LABEL:func.func @trip_count_i8_unsigned_partial_range(
+func.func @trip_count_i8_unsigned_partial_range(%a : i32, %b : i32) -> i32 {
+  %c0 = arith.constant 0 : i8
+  %c200 = arith.constant 200 : i8
+  %c1 = arith.constant 1 : i8
+
+  // Unsigned i8 from 0 to 200: trip count is 200
+  // 200 in i8 is represented as -56 when printed in signed format.
+  // CHECK: "test.trip-count" = -56 : i8
+  %r = scf.for unsigned %i = %c0 to %c200 step %c1 iter_args(%0 = %a) -> i32 : i8 {
+    scf.yield %b : i32
+  }
+  return %r : i32
+}
+
+// -----
+
+// CHECK-LABEL:func.func @trip_count_i8_unsigned_high_range(
+func.func @trip_count_i8_unsigned_high_range(%a : i32, %b : i32) -> i32 {
+  %c200 = arith.constant 200 : i8
+  %c255 = arith.constant 255 : i8
+  %c1 = arith.constant 1 : i8
+
+  // Unsigned i8 from 200 to 255: trip count is 55
+  // CHECK: "test.trip-count" = 55 : i8
+  %r = scf.for unsigned %i = %c200 to %c255 step %c1 iter_args(%0 = %a) -> i32 : i8 {
+    scf.yield %b : i32
+  }
+  return %r : i32
+}
+
+// -----
+
+// CHECK-LABEL:func.func @trip_count_i8_signed_crossing_zero(
+func.func @trip_count_i8_signed_crossing_zero(%a : i32, %b : i32) -> i32 {
+  %c-128 = arith.constant -128 : i32
+  %c127 = arith.constant 127 : i32
+  %c1 = arith.constant 1 : i32
+
+  // Signed i32 from -128 to 127, crossing zero
+  // CHECK: "test.trip-count" = 255
+  %r = scf.for %i = %c-128 to %c127 step %c1 iter_args(%0 = %a) -> i32 : i32 {
+    scf.yield %b : i32
+  }
+  return %r : i32
+}
+
+// -----
+
+// CHECK-LABEL:func.func @trip_count_i16_unsigned_full_range(
+func.func @trip_count_i16_unsigned_full_range(%a : i32, %b : i32) -> i32 {
+  %c0 = arith.constant 0 : i16
+  %c65535 = arith.constant 65535 : i16
+  %c1 = arith.constant 1 : i16
+
+  // Unsigned i16 from 0 to 65535: trip count is 65535
+  // 65535 in i16 is represented as -1 when printed in signed format.
+  // CHECK: "test.trip-count" = -1 : i16
+  %r = scf.for unsigned %i = %c0 to %c65535 step %c1 iter_args(%0 = %a) -> i32 : i16 {
+    scf.yield %b : i32
+  }
+  return %r : i32
+}
+
+// -----
+
+// CHECK-LABEL:func.func @trip_count_i8_unsigned_step_2(
+func.func @trip_count_i8_unsigned_step_2(%a : i32, %b : i32) -> i32 {
+  %c0 = arith.constant 0 : i8
+  %c255 = arith.constant 255 : i8
+  %c2 = arith.constant 2 : i8
+
+  // Unsigned i8 from 0 to 255 step 2: trip count is 128 (255/2 rounded up)
+  // 128 in i8 is represented as -128 when printed in signed format.
+  // CHECK: "test.trip-count" = -128 : i8
+  %r = scf.for unsigned %i = %c0 to %c255 step %c2 iter_args(%0 = %a) -> i32 : i8 {
+    scf.yield %b : i32
+  }
+  return %r : i32
 }

mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp

@@ -50,7 +50,7 @@ struct TestSCFForUtilsPass
             "test.trip-count",
             IntegerAttr::get(IntegerType::get(&getContext(),
                                               tripCount.value().getBitWidth()),
-                             tripCount.value().getSExtValue()));
+                             tripCount.value().getZExtValue()));
       else
         loopOp->setDiscardableAttr("test.trip-count",
                                    StringAttr::get(&getContext(), "none"));