[Mlir-commits] [mlir] [mlir][vector] Replace unused shuffle operands / results with poison (PR #190763)

[Diego Caballero]

================
@@ -3472,12 +3472,55 @@ class ShuffleInterleave : public OpRewritePattern<ShuffleOp> {
   }
 };
 
+/// Pattern to replace usused shuffle operands / results with poison.
+class FoldUnusedShuffleOperand final : public OpRewritePattern<ShuffleOp> {
+public:
+  using Base::Base;
+
+  LogicalResult matchAndRewrite(ShuffleOp op,
+                                PatternRewriter &rewriter) const override {
+    // Replace with poison if all mask elements are poison.
+    if (llvm::all_of(op.getMask(), [](int64_t mask) {
+          return mask == ShuffleOp::kPoisonIndex;
+        })) {
+      rewriter.replaceOpWithNewOp<ub::PoisonOp>(op, op.getType());
+      return success();
+    }
+
+    // Replace V1 with poison if it is not used.
+    int64_t leadingV1Size = op.getV1VectorType().getRank() > 0
+                                ? op.getV1VectorType().getDimSize(0)
+                                : 1;
+    bool isV1Used = llvm::any_of(op.getMask(), [&](int64_t mask) {
+      return mask != ShuffleOp::kPoisonIndex && mask < leadingV1Size;
+    });
+    if (!isV1Used && !matchPattern(op.getV1(), ub::m_Poison())) {
+      Value poison =
+          ub::PoisonOp::create(rewriter, op.getLoc(), op.getV1VectorType());
+      rewriter.modifyOpInPlace(op, [&]() { op.getV1Mutable().assign(poison); });
+      return success();
+    }
+
+    // Replace V2 with poison if it is not used.
+    bool isV2Used = llvm::any_of(op.getMask(), [&](int64_t mask) {
+      return mask != ShuffleOp::kPoisonIndex && mask >= leadingV1Size;
+    });
+    if (!isV2Used && !matchPattern(op.getV2(), ub::m_Poison())) {
+      Value poison =
+          ub::PoisonOp::create(rewriter, op.getLoc(), op.getV2VectorType());
+      rewriter.modifyOpInPlace(op, [&]() { op.getV2Mutable().assign(poison); });
+      return success();
+    }
----------------
dcaballe wrote:

Refactor these two blocks into only a parametrized one?

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