[mlir][spirv] Add support for poison index values to `spirv.CompositeExtract/InsertOp`

[dcaballe]

#123488 adds support for poison indices to vector.insert/extract ops. Similarly, we should define how poison indices are represented for spirv.CompositeExtract/InsertOp and enable the lowering from Vector to SPIR-V for these indices.

We can use the lit tests introduced here as reproducers (they currently error out):

• @extract_poison_idx: https://github.com/llvm/llvm-project/pull/123488/files#diff-38974e5f35197a568a112dcab4e1261396394b503c26c19a9454aa1625a84454R178-R185
• @insert_poison_idx: https://github.com/llvm/llvm-project/pull/123488/files#diff-38974e5f35197a568a112dcab4e1261396394b503c26c19a9454aa1625a84454R267-R274

[llvmbot]

@llvm/issue-subscribers-mlir

Author: Diego Caballero (dcaballe)

https://github.com//pull/123488 adds support for poison indices to `vector.insert/extract` ops. Similarly, we should define how poison indices are represented for `spirv.CompositeExtract/InsertOp` and enable the lowering from Vector to SPIR-V for these indices.

We can use the lit tests introduced here as reproducers (they currently error out):

• @<!-- -->extract_poison_idx: https://github.com/llvm/llvm-project/pull/123488/files#diff-38974e5f35197a568a112dcab4e1261396394b503c26c19a9454aa1625a84454R178-R185
• @<!-- -->insert_poison_idx: https://github.com/llvm/llvm-project/pull/123488/files#diff-38974e5f35197a568a112dcab4e1261396394b503c26c19a9454aa1625a84454R267-R274

[llvmbot]

@llvm/issue-subscribers-mlir-spirv

Author: Diego Caballero (dcaballe)

https://github.com//pull/123488 adds support for poison indices to `vector.insert/extract` ops. Similarly, we should define how poison indices are represented for `spirv.CompositeExtract/InsertOp` and enable the lowering from Vector to SPIR-V for these indices.

We can use the lit tests introduced here as reproducers (they currently error out):

• @<!-- -->extract_poison_idx: https://github.com/llvm/llvm-project/pull/123488/files#diff-38974e5f35197a568a112dcab4e1261396394b503c26c19a9454aa1625a84454R178-R185
• @<!-- -->insert_poison_idx: https://github.com/llvm/llvm-project/pull/123488/files#diff-38974e5f35197a568a112dcab4e1261396394b503c26c19a9454aa1625a84454R267-R274

[kuhar]

SPIR-V treats OOB (including -1) indices in extract/insert as immediate undefined behavior. For this reason, we need to handle two cases:

1. Static: replace -1 with some arbitrary index (say, 0).
2. Dynamic: clamp the range of indices to [0, n-1]. For power-of-two vectors, this can be an and with a mask. This will incur runtime overhead but should be acceptable since dynamic inserts/extracts are not expected to be that performant anyway. If the overhead is non-negligible, we may need to rely on ValueBounds analysis to remove these or introduce an inbounds attribute at the level of the vector dialect.

See https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_composite_instructions

[andfau-amd]

Just to be absolutely clear, this is a lowering task (vector -> spirv) and does not involve changing the behaviour of the spirv ops directly?

[kuhar]

Correct.

[andfau-amd]

The constant indices case is trivial, but the dynamic indices case is scary, because there isn't a definitive target for SPIR-V, so I'm not how overhead should be estimated, or who should make that decision. Would simply always clamping be acceptable for now, and then it can be refined later if someone complains? Dynamic indices are already a rare and slow path anyway, I think.

[kuhar]

Would simply always clamping be acceptable for now, and then it can be refined later if someone complains? Dynamic indices are already a rare and slow path anyway, I think.

I think we can start with that. Dynamic indexing is not fast in general anyway. If we learn that the performance regresses too much, we can use the ValueBounds analysis to see which clamps can be dropped, and if that's not enough, consider adding an inbounds attribute to these vector ops.

because there isn't a definitive target for SPIR-V, so I'm not how overhead should be estimated, or who should make that decision

We can start with some real-world model from IREE on any random GPU (say, rdna3) to show that there's not meaningful perf regression.