[Mlir-commits] [mlir] [mlir][Vector] Fold vector.step compared to constant (PR #161615)

[Jakub Kuderski]

================
@@ -7602,6 +7602,106 @@ void StepOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
   setResultRanges(getResult(), result);
 }
 
+namespace {
+
+/// Fold `vector.step -> arith.cmpi` when the step value is compared to a
+/// constant large enough such that the result is the same at all indices.
+///
+/// For example, rewrite the 'greater than' comparison below,
+///
+/// ```mlir
+/// %cst = arith.constant dense<7> : vector<3xindex>
+/// %stp = vector.step : vector<3xindex>
+/// %out = arith.cmpi ugt, %stp, %cst : vector<3xindex>
+/// ```
+///
+/// as,
+///
+/// ```mlir
+/// %out = arith.constant dense<false> : vector<3xi1>.
+/// ```
+///
+/// Above [0, 1, 2] > [7, 7, 7] => [false, false, false]. Because the result is
+/// false at ALL indices we fold. If the constant was 1, then
+/// [0, 1, 2] > [1, 1, 1] => [false, false, true] and we do fold, conservatively
----------------
kuhar wrote:

```suggestion
/// Above `[0, 1, 2] > [7, 7, 7]` => `[false, false, false]`. Because the result is
/// false at ALL indices we fold. If the constant was 1, then
/// `[0, 1, 2]` > [1, 1, 1]` => `[false, false, true]` and we do fold, conservatively
```

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