[Mlir-commits] [mlir] [mlir] Implement inferResultRanges for vector.step op (PR #151536)

[Krzysztof Drewniak]

================
@@ -7197,6 +7197,31 @@ Value mlir::vector::makeArithReduction(OpBuilder &b, Location loc,
   return selectPassthru(b, mask, result, acc);
 }
 
+//===----------------------------------------------------------------------===//
+// StepOp
+//===----------------------------------------------------------------------===//
+
+void StepOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                               SetIntRangeFn setResultRanges) {
+  auto resultType = cast<VectorType>(getType());
+  if (resultType.isScalable()) {
+    return;
+  }
+  std::optional<ConstantIntRanges> result;
+  Type elementType = resultType.getElementType();
+  unsigned bitwidth = elementType.isIndex()
+                          ? IndexType::kInternalStorageBitWidth
+                          : elementType.getIntOrFloatBitWidth();
+  int64_t size = resultType.getShape()[0];
+  for (int64_t val : llvm::seq<int64_t>(size)) {
----------------
krzysz00 wrote:

... This is O(N) for no good reason - I'd just create the range from 0 to `size - 1`, and `fromSigned` it

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