[mlir][Vector] Relax condition for `splitFullAndPartialTransferPrecondition`

The `splitFullAndPartialTransferPrecondition` has a restrictive condition to
prevent the pattern to be applied recursively if it is nested under an scf.IfOp.
Relaxing the condition to the immediate parent op must not be an scf.IfOp lets
the pattern be applied more generally while still preventing recursion.

Differential Revision: https://reviews.llvm.org/D85209

diff --git a/mlir/lib/Dialect/Vector/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/VectorTransforms.cpp
index 3c23c5a..33fbed6 100644 (file)
--- a/mlir/lib/Dialect/Vector/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/VectorTransforms.cpp
@@ -2049,10 +2049,10 @@ LogicalResult mlir::vector::splitFullAndPartialTransferPrecondition(
   // Must have some masked dimension to be a candidate for splitting.
   if (!xferOp.hasMaskedDim())
     return failure();
-  // Don't split transfer operations under IfOp, this avoids applying the
-  // pattern recursively.
-  // TODO: improve the condition to make it more applicable.
-  if (xferOp.getParentOfType<scf::IfOp>())
+  // Don't split transfer operations directly under IfOp, this avoids applying
+  // the pattern recursively.
+  // TODO: improve the filtering condition to make it more applicable.
+  if (isa<scf::IfOp>(xferOp.getOperation()->getParentOp()))
     return failure();
   return success();
 }