[Mlir-commits] [mlir] [mlir][vector] Adds pattern rewrite for maskable Ops (PR #83827)

[Diego Caballero]

================
@@ -212,6 +211,64 @@ static Value createMul(Location loc, Value x, Value y, bool isInt,
 
 namespace {
 
+/// A pattern for ops that implement `MaskableOpInterface` and that _might_ be
+/// masked (i.e. inside `vector.mask` Op region). In particular:
+///   1. It matches `SourceOp` operation, Op.
+///   2. If Op is masked, retrieves the mask and updates the insertion point to
+///   avoid inserting new ops into `vector.mask` Op region (which only allows
+///   one Op). If the Op is not masked, this step is a nop.
+///   3. Invokes `matchAndRewriteMaskableOp` on Op that might be nested (not
+///   required) in the matched `vector.mask` operation from step 2.
+///
+/// It frees the patterns implementing this class from worrying about the
+/// logic to update the insertion point. However, those patterns are still
+/// responsible for providing an updated version of:
+///   * the source Op when mask _is not_ present,
+///   * the source Op *and* the mask Op when mask _is_ present.
+template <class SourceOp>
+struct MaybeMaskedOpRewritePattern : OpRewritePattern<SourceOp> {
----------------
dcaballe wrote:

This is the main concern I've been struggling with for a while: this pattern matches the nested op and the rewrite would replace its parent. I'm not quite sure pattern rewrite allows that even though it works for this cases. We can make this match against the vector mask op but then we would need another pattern to match against the source op for the unmasked cases. I couldn't find an elegant way to provide this functionality but maybe you can :)

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