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

[Andrzej Warzyński]

================
@@ -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.
----------------
banach-space wrote:

[nit] To me this is a fairly significant (and not immediately intuitive) part of the design. Perhaps move above the signature?

Also, is this based on some prior-art? Just curious, this does make sense. 

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