[mlir][vector] Error when running `--canonicalize` with out of bounds positions in `vector.extract`

[math-fehr]

The following code

func.func @extract_poison_idx(%arg0: vector<16xf32>) -> f32 {
  %x = arith.constant -2 : index
  %0 = vector.extract %arg0[%x]: f32 from vector<16xf32>
  return %0 : f32
}

raises the following error when running --canonicalize

<source>:3:8: error: 'vector.extract' op expected position attribute #1 to be a non-negative integer smaller than the corresponding vector dimension or poison (-1)
  %0 = vector.extract %arg0[%x]: f32 from vector<16xf32>
       ^

This is due to -1 being the only accepted static position that is out of bounds in a vector.extract or vector.insert.

There should be two ways to fix this:

• Either we allow out of bounds static positions that are not -1 (which are also poison).
• Or the vector.extract builders should map out of bounds static positions to -1.
• Or it is the role of the users to make sure they do not pass non -1 out of bounds static positions.

I feel both options come with a cost, and would personally prefer solution 1, as it follows LLVM semantics (extractelement <4 x i32> %vec, i32 -2 is poison as well), and is roughly more modular. I feel solution 3 is exactly why we have this bug in the first place, and would be quite error-prone overall, because it is really easy to miss. Solution 2 is okay for me, but I feel following LLVM would make more sense here.

What do you think? I can write the patch, I just want to make sure we agree on a solution first before going into it. It also applies for vector.insert and vector.shuffle from what I see.

[math-fehr]

@kuhar @banach-space @dcaballe @Groverkss, I'm pinging you explicitly as you were the ones reviewing/writing #123488 which introduced the -1 constraint.

[kuhar]

The combination of the verifier & canonicalizer is too strict -- we shouldn't be checking (non-attribute) operands like this.

With the current op specs, this is immediate UB, and the fix would be either:

• Teach canon patterns to leave invalid constant operands as-is.
• Turn invalid constant operands into valid attribute operands (-1 or any other valid index). This is fine because lowering can make things more defined (refinement).

I'd be leaning towards not for defining extract/insert with OOB indices as poison -- I'm not convinced yet it improves the overall consistency/ergonomic. I said 'leaning' because in terms of runtime overhead in the dynamic case, I don't think it will cause any extra overhead during lowering:

• SPIR-V does treat any OOB index (including -1) as immediate UB, and we already have to account for that in the lowering by clamping. Adding more OOB cases won't make the situation worse for SPIR-V.
• LLVM already treats OOB indices as poison: https://llvm.org/docs/LangRef.html#id184.

If others would prefer consistency with the LLVM instructions, I won't object. If we do go this way, I think this should be a separate PR from the one that makes the current semantics work properly.

[math-fehr]

That makes sense to me!
Technically here it is not that we check non-attribute operands, but it's the fact that the folder wouldn't map a wrong access to -1.

I think I have a more general question which is why do we require OOB accesses to be mapped to -1? Looking at the origin of it #123488, it seems that there was a misunderstanding on how LLVM works, as -1 is not a special case they are using AFAIK, and any OOB value is allowed (and canonicalize the operation to poison). I am wondering if we should not instead allow all values in the verifier, which should simplify a bit the semantics and avoid these kind of errors. Unless maybe there are other considerations I'm missing?

And yes, for chosing wether extract/insert should be OOB or poison, the only datapoint I have is that LLVM uses poison instead of UB so that the operations can be safely hoisted from loops (because of ConditionallySpeculatable), which is way harder to do when immediate UB is involved.

[kuhar]

I think I have a more general question which is why do we require OOB accesses to be mapped to -1? Looking at the origin of it #123488, it seems that there was a misunderstanding on how LLVM works, as -1 is not a special case they are using AFAIK, and any OOB value is allowed (and canonicalize the operation to poison)

I think this follows the shufflevector semantics, where there's dedicated poison index: https://llvm.org/docs/LangRef.html#id193. @dcaballe should know if this is intended or not.

I am wondering if we should not instead allow all values in the verifier, which should simplify a bit the semantics and avoid these kind of errors. Unless maybe there are other considerations I'm missing?

I think this may make sense but I think it'd be better to stage this and first fix the current semantics before allowing more OOB indices.

If we canonicalize constant OOB indices towards -1, this should fix the error, and be compatible with allowing more OOB values later on.

[banach-space]

Thanks for raising this issue, @math-fehr !

The combination of the verifier & canonicalizer is too strict -- we shouldn't be checking (non-attribute) operands like this.

Agreed. This violates the (undocumented?) rule that a verifier should only rely on the "local" info. This is not "local":

  %x = arith.constant -2 : index

is that what you had in mind, Jakub?

I think I have a more general question which is why do we require OOB accesses to be mapped to -1?

Just to clarify, -1 was introduced to model "poison" values specifically. I don't recall discussing OOB.

I think this follows the shufflevector semantics, where there's dedicated poison index: https://llvm.org/docs/LangRef.html#id193.

I actually don't see anything specific about -1 🤔

If others would prefer consistency with the LLVM instructions, I won't object.

I don't really have any data-points to refer to, so for now I'd keep -1 as the only valid option and to use that to model poison specifically. Mostly to keep things simple.

[kuhar]

is that what you had in mind, Jakub?

Yes, exactly. This is documented here: https://mlir.llvm.org/getting_started/DeveloperGuide/#ir-verifier

[math-fehr]

Agreed. This violates the (undocumented?) rule that a verifier should only rely on the "local" info. This is not "local":

I think there is a misunderstanding on this, the error comes from the folder that folds an OOB access, which triggers a verifier error after the operation is folded.

Just to clarify, -1 was introduced to model "poison" values specifically. I don't recall discussing OOB.

Ohhh, this was a misunderstanding on my end, my bad.

I think this may make sense but I think it'd be better to stage this and first fix the current semantics before allowing more OOB indices.

Yes, I'll write a patch to fix the folder, so then at least this should now work.

[kuhar]

I think there is a misunderstanding on this, the error comes from the folder that folds an OOB access, which triggers a verifier error after the operation is folded.

I think we are on the same page: it's a combination of the folder and the verifier that's at fault. Folders must not turn valid IR into invalid IR. The fix is to stop folding such inputs.

[dcaballe]

Using -1 to represent a poison index comes from the shuffle vector implementation in LLVM: https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/IR/Instructions.h#L1901-L1909. We extended it to vector.extract and vector.insert ops in MLIR for consistency and simplicity since they are all operations indexing a position within a vector and their indices are often propagated (with some minor transformations) across these ops in canonicalizations and folders.

I don't see a problem with allowing more OOB indices. This was initially constrained to test the waters. I would canonicalize OOB to -1, though, to avoid having arbitrary odd numbers in the IR.