[mlir][vector] Update how OOB accesses are handled by vector.extract/vector.insert

[banach-space]

At the moment, only -1 is used to "model" poison values - i.e., accesses that produce poison:

  %c_neg_1 = arith.constant -1 : index
  %0 = vector.insert %value_to_store, %dest[%c_neg_1] : vector<5xf32> into vector<4x5xf32>
  %1 = vector.extract %src[%c_neg_1, 0] : f32 from vector<4x5xf32>

From the discussion so far, it seems that:

• We should treat any out-of-bounds (OOB) index — not just -1 — as producing a poison value (see @dcaballe’s supportive comment).
• We should allow using %poison = ub.poison : index directly as an index that generates poison.

This direction would bring the dialect more in line with LLVM semantics. From LLVM LangRef:

If idx exceeds the length of val for a fixed-length vector, the result is a poison value

VOICES AGAINST / CONCERNS

We should be mindful of @kuhar’s comment. While not opposed to the proposed direction:

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

... we should keep him in the loop and ensure any concerns he raises are addressed.

PRIOR DISCUSSION

Relevant context and discussion can be found in:

• [mlir][vector] Error when running --canonicalize with out of bounds positions in vector.extract #134516
• [mlir][vector] Prevent folding of OOB values in insert/extract #135498
• [mlir][vector] Update the folder for vector.{insert|extract} #136579

Also, this issue is to track NEXT STEPS from this comment.

CC @dcaballe @math-fehr @kuhar

[llvmbot]

@llvm/issue-subscribers-mlir-vectorops

Author: Andrzej Warzyński (banach-space)

At the moment, only `-1` is used to "model" poison values - i.e., accesses that produce poison:

  %c_neg_1 = arith.constant -1 : index
  %0 = vector.insert %value_to_store, %dest[%c_neg_1] : vector<5xf32> into vector<4x5xf32>
  %1 = vector.extract %src[%c_neg_1, 0] : f32 from vector<4x5xf32>

From the discussion so far, it seems that:

• We should treat any out-of-bounds (OOB) index — not just -1 — as producing a poison value (see @dcaballe’s supportive comment).
• We should allow using %poison = ub.poison : index directly as an index that generates poison.

This direction would bring the dialect more in line with LLVM semantics. From LLVM LangRef:

> If idx exceeds the length of val for a fixed-length vector, the result is a poison value

VOICES AGAINST / CONCERNS

We should be mindful of @kuhar’s comment. While not opposed to the proposed direction:

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

... we should keep him in the loop and ensure any concerns he raises are addressed.

PRIOR DISCUSSION

Relevant context and discussion can be found in:

• [mlir][vector] Error when running --canonicalize with out of bounds positions in vector.extract #134516
• [mlir][vector] Prevent folding of OOB values in insert/extract #135498
• [mlir][vector] Update the folder for vector.{insert|extract} #136579

Also, this issue is to track NEXT STEPS from this comment.

[joker-eph]

I commented elsewhere on a command mentioning

With the current op specs, this is immediate UB,

The op is marked as "Pure", either it isn't immediate UB or it's a bug in ODS.
You can't speculate an operation that is "immediate UB" for example.

It was also mentioned that SPIRV has immediate UB behavior (?), which would be a problem: speculating an extract and then converting to SPIRV would be incorrect.
I wonder how it works in LLVM? Does the SPIRV backend need to guard the extract in the lowering to prevent UB?

[kuhar]

With the current op specs, this is immediate UB,

The op is marked as "Pure", either it isn't immediate UB or it's a bug in ODS. You can't speculate an operation that is "immediate UB" for example.

This is a very good point! I think we'd either have to make it conditionally speculatable with the existing semantics, or allow for OOB accesses. I'm fine with either.

It was also mentioned that SPIRV has immediate UB behavior (?), which would be a problem: speculating an extract and then converting to SPIRV would be incorrect.
I wonder how it works in LLVM? Does the SPIRV backend need to guard the extract in the lowering to prevent UB?

Unlike the vector dialect, SPIR-V has two different sets of ops: one for static indices, one for dynamic indices. My understanding is that the former are speculatable and the latter are not. During vector to spir-v lowering, we have to clamp the dynamic case to remain inbounds -- this adds runtime overhead, but dynamic inserts/extracts are not expected to be performant on GPU anyway...

[joker-eph]

This is a very good point! I think we'd either have to make it conditionally speculatable with the existing semantics, or allow for OOB accesses.

I'm missing something: what about poison here?

[kuhar]

I'm not sure what you mean?
allow for OOB accesses ==> make that produce poison

[joker-eph]

Oh gotcha! (I understood "allow for OOB accesses" -> "make it legally supported by the operation somehow")

[dcaballe]

Going back to how to represent poison in a constant index array:

I thought that the reason -1 was used was that it is not possible to have an #ub.poison attribute in the constant index array

This is an example where we need to represent constant indices within the range of an integer type + some extra values with specific semantics (i.e., poison). Representing dynamic dimensions is another case where this problem occurs (we are currently using std::numeric_limits<int64_t>>::min, IIRC). It looks like we need a more powerful internal representation for cases where we want to represent the full range of an integer type + these special values without stepping into the integer type range for that.

I briefly looked into this some time ago and tried to create a new flavor of DynamicIndexList but run into some parsing/printing issues (I can't remember the details but at least we'd need a way to specify how these special values are parsed/printed and represented internally). I'm bringing this up to see if @joker-eph, @jpienaar or anyone would have any suggestion or idea to accomplish this as it goes well into the MLIR internals and it looks like a recurring issue that we may want to solve properly?

Thoughts?