address Jakub's comments

Signed-off-by: James Newling <[email protected]>

Author: newling

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

@@ -7628,62 +7628,64 @@ struct StepCompareFolder : public OpRewritePattern<StepOp> {
     const int64_t stepSize = stepOp.getResult().getType().getNumElements();
 
     for (auto &use : stepOp.getResult().getUses()) {
-      if (auto cmpiOp = dyn_cast<arith::CmpIOp>(use.getOwner())) {
-        const unsigned stepOperandNumber = use.getOperandNumber();
-
-        // arith.cmpi canonicalizer makes constants final operands.
-        if (stepOperandNumber != 0)
-          continue;
-
-        // Check that operand 1 is a constant.
-        unsigned constOperandNumber = 1;
-        Value otherOperand = cmpiOp.getOperand(constOperandNumber);
-        auto maybeConstValue = getConstantIntValue(otherOperand);
-        if (!maybeConstValue.has_value())
-          continue;
+      auto cmpiOp = dyn_cast<arith::CmpIOp>(use.getOwner());
+      if (!cmpiOp)
+        continue;
 
-        int64_t constValue = maybeConstValue.value();
-        arith::CmpIPredicate pred = cmpiOp.getPredicate();
+      // arith.cmpi canonicalizer makes constants final operands.
+      const unsigned stepOperandNumber = use.getOperandNumber();
+      if (stepOperandNumber != 0)
+        continue;
 
-        auto maybeSplat = [&]() -> std::optional<bool> {
-          // Handle ult (unsigned less than) and uge (unsigned greater equal).
-          if ((pred == arith::CmpIPredicate::ult ||
-               pred == arith::CmpIPredicate::uge) &&
-              stepSize <= constValue)
-            return pred == arith::CmpIPredicate::ult;
+      // Check that operand 1 is a constant.
+      unsigned constOperandNumber = 1;
+      Value otherOperand = cmpiOp.getOperand(constOperandNumber);
+      auto maybeConstValue = getConstantIntValue(otherOperand);
+      if (!maybeConstValue.has_value())
+        continue;
 
-          // Handle ule and ugt.
-          if ((pred == arith::CmpIPredicate::ule ||
-               pred == arith::CmpIPredicate::ugt) &&
-              stepSize <= constValue + 1)
-            return pred == arith::CmpIPredicate::ule;
+      int64_t constValue = maybeConstValue.value();
+      arith::CmpIPredicate pred = cmpiOp.getPredicate();
+
+      auto maybeSplat = [&]() -> std::optional<bool> {
+        // Handle ult (unsigned less than) and uge (unsigned greater equal).
+        if ((pred == arith::CmpIPredicate::ult ||
+             pred == arith::CmpIPredicate::uge) &&
+            stepSize <= constValue)
+          return pred == arith::CmpIPredicate::ult;
+
+        // Handle ule and ugt.
+        if ((pred == arith::CmpIPredicate::ule ||
+             pred == arith::CmpIPredicate::ugt) &&
+            stepSize - 1 <= constValue) {
+          return pred == arith::CmpIPredicate::ule;
+        }
 
-          // Handle eq and ne.
-          if ((pred == arith::CmpIPredicate::eq ||
-               pred == arith::CmpIPredicate::ne) &&
-              stepSize <= constValue)
-            return pred == arith::CmpIPredicate::ne;
+        // Handle eq and ne.
+        if ((pred == arith::CmpIPredicate::eq ||
+             pred == arith::CmpIPredicate::ne) &&
+            stepSize <= constValue)
+          return pred == arith::CmpIPredicate::ne;
 
-          return std::optional<bool>();
-        }();
+        return std::nullopt;
+      }();
 
-        if (!maybeSplat.has_value())
-          continue;
+      if (!maybeSplat.has_value())
+        continue;
 
-        rewriter.setInsertionPointAfter(cmpiOp);
+      rewriter.setInsertionPointAfter(cmpiOp);
 
-        auto type = dyn_cast<VectorType>(cmpiOp.getResult().getType());
-        if (!type)
-          continue;
+      auto type = dyn_cast<VectorType>(cmpiOp.getResult().getType());
+      if (!type)
+        continue;
 
-        DenseElementsAttr boolAttr =
-            DenseElementsAttr::get(type, maybeSplat.value());
-        Value splat = mlir::arith::ConstantOp::create(rewriter, cmpiOp.getLoc(),
-                                                      type, boolAttr);
+      DenseElementsAttr boolAttr =
+          DenseElementsAttr::get(type, maybeSplat.value());
+      Value splat = mlir::arith::ConstantOp::create(rewriter, cmpiOp.getLoc(),
+                                                    type, boolAttr);
 
-        rewriter.replaceOp(cmpiOp, splat);
-        return success();
-      }
+      rewriter.replaceOp(cmpiOp, splat);
+      return success();
     }
 
     return failure();

mlir/test/Dialect/Vector/canonicalize/vector-step.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -canonicalize="test-convergence" -split-input-file -allow-unregistered-dialect | FileCheck %s
+// RUN: mlir-opt %s -canonicalize="test-convergence" -split-input-file | FileCheck %s
 
 ///===----------------------------------------------===//
 ///  Tests of `StepCompareFolder`
@@ -59,6 +59,20 @@ func.func @check_ugt_constant_3_rhs() -> vector<3xi1> {
   return %1 : vector<3xi1>
 }
 
+// -----
+
+// CHECK-LABEL: @check_ugt_constant_max_rhs
+//       CHECK: %[[CST:.*]] = arith.constant dense<false> : vector<3xi1>
+//       CHECK: return %[[CST]] : vector<3xi1>
+func.func @check_ugt_constant_max_rhs() -> vector<3xi1> {
+  // The largest i64 possible:
+  %cst = arith.constant dense<0x7fffffffffffffff> : vector<3xindex>
+  %0 = vector.step : vector<3xindex>
+  %1 = arith.cmpi ugt, %0, %cst: vector<3xindex>
+  return %1 : vector<3xi1>
+}
+
+
 // -----
 
 // CHECK-LABEL: @check_ugt_constant_2_rhs