[Mlir-commits] [mlir] [mlir][VectorOps] Fold extract on constant_mask (PR #183780)

[Lukas Sommer]

================
@@ -2298,6 +2298,59 @@ class ExtractOpFromCreateMask final : public OpRewritePattern<ExtractOp> {
   }
 };
 
+// Pattern to rewrite a ExtractOp(ConstantMask) -> ConstantMask.
+class ExtractOpFromConstantMask final : public OpRewritePattern<ExtractOp> {
+public:
+  using Base::Base;
+
+  LogicalResult matchAndRewrite(ExtractOp extractOp,
+                                PatternRewriter &rewriter) const override {
+    auto constantMaskOp =
+        extractOp.getSource().getDefiningOp<vector::ConstantMaskOp>();
+    if (!constantMaskOp)
+      return failure();
+
+    auto extractedMaskType =
+        dyn_cast<VectorType>(extractOp.getResult().getType());
+    if (!extractedMaskType)
+      return failure();
----------------
sommerlukas wrote:

No strong reason, I mainly followed `ExtractOpFromCreateMask`, which also bails out in that case. I've now added support for the scalar case in the latest commit. 

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