[Mlir-commits] [mlir] [mlir][spirv] Support poison index when converting vector.insert/extract (PR #125560)

[Jakub Kuderski]

================
@@ -137,6 +137,26 @@ struct VectorBroadcastConvert final
   }
 };
 
+// SPIR-V does not have a concept of a poison index for certain instructions,
+// which creates a UB hazard when lowering from otherwise equivalent Vector
+// dialect instructions, because this index will be considered out-of-bounds.
+// To avoid this, this function implements a dynamic sanitization, arbitrarily
+// choosing to replace the poison index with index 0 (always in-bounds).
+static Value sanitizeDynamicIndex(ConversionPatternRewriter &rewriter,
+                                  Location loc, Value dynamicIndex,
+                                  int64_t kPoisonIndex) {
+  Value poisonIndex = rewriter.create<spirv::ConstantOp>(
+      loc, dynamicIndex.getType(),
+      rewriter.getIntegerAttr(dynamicIndex.getType(), kPoisonIndex));
+  Value cmpResult =
+      rewriter.create<spirv::IEqualOp>(loc, dynamicIndex, poisonIndex);
+  Value sanitizedIndex = rewriter.create<spirv::SelectOp>(
+      loc, cmpResult,
+      spirv::ConstantOp::getZero(dynamicIndex.getType(), loc, rewriter),
----------------
kuhar wrote:

Selects are fairly expensive on many GPU targets. Instead, Can we clamp power-of-two vectors using an `and` with a mask? I'd expect that any performance-oriented code uses pow2 vectors. 

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