[Mlir-commits] [mlir] [mlir][vector] Fix crash on untraceable masks in getCompressedMaskOp (PR #207299)

[Andrzej Warzyński]

================
@@ -98,10 +98,16 @@ static FailureOr<Operation *> getCompressedMaskOp(OpBuilder &rewriter,
     if (auto extractOp = dyn_cast<vector::ExtractOp>(maskOp)) {
       maskOp = extractOp.getSource().getDefiningOp();
       extractOps.push_back(extractOp);
+    } else {
+      // Unsupported mask-defining op (e.g. a block argument, which has no
+      // defining op, or an op we cannot trace through). Bail out rather than
+      // looping forever or dereferencing a null op below.
+      break;
     }
   }
----------------
banach-space wrote:

I prefer early and clear bail out. Also, I would update the comments.
```suggestion
// Chain of extract Ops that lead to the original Op that created the mask.
// TODO: add support to `vector.broadcast`.
SmallVector<vector::ExtractOp, 2> extractOps;

// Find the mask creation operation.
while (maskOp &&
!isa<arith::ConstantOp, vector::CreateMaskOp, vector::ConstantMaskOp>(
maskOp)) {

  auto extractOp = dyn_cast<vector::ExtractOp>(maskOp));
  if (!extractOp) {
        // Unsupported mask-defining op (e.g. a block argument, which has no
        // defining op, or an op we cannot trace through). Bail out rather than
        // looping forever or dereferencing a null op below.
     return failure()
  }
  maskOp = extractOp.getSource().getDefiningOp();
  extractOps.push_back(extractOp);
}
```


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