[Mlir-commits] [llvm] [mlir] [Linalg] Fix crash in vectorizeScalableVectorPrecondition with undersized vector sizes (PR #205493)

[Alessandro Potenza]

https://github.com/alepot55 commented:

@banach-space asked for this to move "somewhere higher up", and @egebeysel asked whether strict equality is worth enforcing. There is a fact in the tree that I think answers both, and I did not see it raised here.

**Equality is already the contract.** `vectorizeLinalgOpPrecondition` opens with:

```cpp
// Check API contract for input vector sizes.
if (!inputVectorSizes.empty() &&
    failed(vector::isValidMaskedInputVector(linalgOp.getStaticLoopRanges(),
                                            inputVectorSizes)))
  return failure();
```

and the first thing `vector::isValidMaskedInputVector` does is:

```cpp
if (inputVectorSizes.size() != shape.size()) {
  LDBG() << "Input vector sizes don't match the number of loops";
  return failure();
}
```

`shape` there is `getStaticLoopRanges()`, so its length is the loop count, and `VectorUtils.h` documents this as requirement 1 of the masking contract. For any LinalgOp with a non-empty `inputVectorSizes`, a mismatch in either direction is already rejected today.

The only reason the crash is reachable is ordering. In `vectorizeOpPrecondition`, `vectorizeScalableVectorPrecondition` runs before the `TypeSwitch` that holds that check, so the assertion fires before the contract is ever consulted.

That makes the `>` test added here a weaker restatement of a rule the tree already has, in a third place. Two ways to end up with one rule rather than two:

- move the `vectorizeScalableVectorPrecondition` call below the `TypeSwitch`. Both have to pass, so only the order in which they fail changes.
- or hoist the `isValidMaskedInputVector` call into `vectorizeOpPrecondition` for the LinalgOp case. For that case only: `vectorizePackOpPrecondition`, `vectorizeUnPackOpPrecondition` and the pad and insert_slice hooks each compare against a rank of their own, not against the loop count.

**On whether `<` is worth enforcing.** It does not crash, but it is not harmless either:

```cpp
auto iterators = linalgOp.getIteratorTypesArray();
SmallVector<bool> scalableFlags(inputScalableVecDims);
int64_t idx = scalableFlags.size() - 1;
while (!scalableFlags[idx]) { ...; iterators.pop_back(); scalableFlags.pop_back(); --idx; }
```

`iterators` starts at the loop count and `scalableFlags` at the length of the user's list, and the loop pops them together. With three loops and flags `[true, false]`, the loop runs once, and `iterators.back()` afterwards is the iterator of loop 1 while the flag being analysed is index 0. The hook then answers against the wrong iterator instead of declining. So the equality @egebeysel is in favour of is not only tidier, it is what keeps this hook honest, and it is already in the tree one call away.

One smaller thing: the PR description says the crash happens "when the arrays are shorter than the number of loops", but the reproducer passes three sizes to a two-loop `linalg.add`, which is the opposite. Worth fixing before this is squashed into a commit message.

All of the above is from reading. I have an MLIR build coming up and can confirm the misalignment on it if that is useful.

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