[Mlir-commits] [mlir] [mlir][vector] Allow signless integer element types in `vector.step` (PR #205142)

[Victor Perez]

victor-eds wrote:

Thanks for the careful look. Yep, this is subtler than it first appears. Let's discuss the options we have.

_One framing point first._ I'd separate _what `vector.step` means_ for a narrow element type from _how we lower it_. The wrap-around behaviour is a property of the op, not of any particular lowering: both when lowering to a constant and `llvm.stepvector`, we honour it. 

I am not a big fan of bailing out on "large vectors" when lowering to `arith.constant` because any threshold is arbitrary, and the only non-arbitrary one (bail right at the wrap-around boundary) doesn't help either, since the constant lowering already encodes the correct semantics.

I see three paths, with different trade-offs:

1. **Keep current lowering** (modulo the fix making the pattern parametric in the `index` bitwidth, defaulting to 64 for backward compatibility). In this case, I'd rather pass the bitwidth in as an argument than read it from the data layout, which could change behaviour for existing pipelines.

```mlir
%0 = vector.step : vector<8xindex>
%1 = vector.step : vector<8xi32>
%2 = vector.step : vector<[8]xindex>
%3 = vector.step : vector<[8]xi32>
```

lowers to (`index` bitwidth = 64):

```mlir
%0 = llvm.mlir.constant(dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xindex>) : vector<4xi64>
%1 = llvm.mlir.constant(dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xi32>) : vector<8xi32>
%2 = llvm.intr.stepvector : vector<[8]xi64>
%3 = llvm.intr.stepvector : vector<[8]xi32>
```
  
2. **Always lower to `llvm.stepvector`** (cleaner than above; drop the `arith.constant` path; constants only for SPIR-V). The op stays in the module until target lowering.

```mlir
%0 = llvm.intr.stepvector : vector<4xi64>
%1 = llvm.intr.stepvector : vector<8xi32>
%2 = llvm.intr.stepvector : vector<[8]xi64>
%3 = llvm.intr.stepvector : vector<[8]xi32>
```

This is the option that depends on your question to @kuhar, as, if anything relies on `StepToArithConstantOpRewrite`, dropping it outright is the risky one.

3. **Fold + lower to `llvm.stepvector`** (my preference). Fixed-width integer element type + fixed-length vector folds to a constant (encoding the wrap-around semantics); everything else stays as `vector.step` and lowers to `llvm.stepvector`. The previous code would fold to:

```mlir
%0 = vector.step : vector<8xindex>
%1 = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xi32>
%2 = vector.step : vector<[8]xindex>
%3 = vector.step : vector<8xi32>
```

and  then lower to:

```mlir
%0 = llvm.intr.stepvector : vector<4xi64>
%1 = llvm.mlir.constant(dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xi32>) : vector<8xi32>
%2 = llvm.intr.stepvector : vector<[8]xi64>
%3 = llvm.intr.stepvector : vector<8xi32>
```

---

I lean towards 3. Happy to go with 2 if nothing depends on `StepToArithConstantOpRewrite`.

 

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