[Mlir-commits] [mlir] [mlir][Vector] Add `vector.shuffle` fold for poison inputs (PR #125608)

[Diego Caballero]

================
@@ -2684,25 +2682,45 @@ OpFoldResult vector::ShuffleOp::fold(FoldAdaptor adaptor) {
   if (!v1Attr || !v2Attr)
     return {};
 
+  // Fold shuffle poison, poison -> poison.
+  bool isV1Poison = isa<ub::PoisonAttr>(v1Attr);
+  bool isV2Poison = isa<ub::PoisonAttr>(v2Attr);
+  if (isV1Poison && isV2Poison)
+    return ub::PoisonAttr::get(getContext());
+
   // Only support 1-D for now to avoid complicated n-D DenseElementsAttr
   // manipulation.
   if (v1Type.getRank() != 1)
     return {};
 
-  int64_t v1Size = v1Type.getDimSize(0);
+  // Poison input attributes need special handling as they are not
+  // DenseElementsAttr. If an index is poison, we select the first element of
+  // the first non-poison input.
----------------
dcaballe wrote:

Not sure I follow the prior-art part. Do you mean why we pick the first element of the first non-poison input? Poison is basically UB so given that we can't represent a partially poison vector we just make a random decision, which is ok as part of the UB behavior.

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