[mlir][transform] Fix apply WithPDLPatternsOp with non-pattern op

[sotoshigoto]

Summary

Fix https://github.com/llvm/llvm-project/issues/60209

Fix crash with segmentation fault when transform::WithPDLPatternsOp is
applied with non-pattern op. Added check for existing transform ops with
pattern op.

[ftynse]

Thanks!

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

1046

Nit: expand auto.

1047–1048

Why not reporting the error when there is a parent that is possibly a top-level op? WithPDLPatterns is always expected to have one non-pattern child.

[sotoshigoto]

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

1047–1048

I think it can pass when there is a parent that may be a top-level op from the below test case.
https://github.com/llvm/llvm-project/blob/main/mlir/test/Dialect/Transform/ops.mlir#L26

But I will fix it.

[sotoshigoto]

update

[ftynse]

That would trigger the exact same error as this patch is attempting to fix. Until this patch, we would allow "empty" WithPDLPatterns at the verifier level, but it could lead to crashes when applying the transformation. We can either disallow them, or fix the application so it does nothing.

[sotoshigoto]

@ftynse I don't have commit access. Could you please commit to this patch?

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

for (Operation &op : body->getOperations()) {
     InFlightDiagnostic diag =
         emitOpError()
         << "expects only pattern and top-level transform ops in its body";
     diag.attachNote(op.getLoc()) << "offending op";
     return diag;
   }
 
   if (auto parent = getOperation()->getParentOfType<WithPDLPatternsOp>()) {
     InFlightDiagnostic diag = emitOpError() << "cannot be nested";
     diag.attachNote(parent.getLoc()) << "parent operation";
     return diag;
   }
 
+  if (!topLevelOp) {
+    InFlightDiagnostic diag = emitOpError()
+                              << "expects at least one non-pattern op";
+    return diag;
+  }
+
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // PrintOp
 //===----------------------------------------------------------------------===//
 
 void transform::PrintOp::build(OpBuilder &builder, OperationState &result,

mlir/test/Dialect/Transform/ops-invalid.mlir

^bb0(%arg0: !pdl.operation):
    // expected-error @below {{op cannot be nested}}
   transform.with_pdl_patterns %arg0 : !pdl.operation {
   ^bb1(%arg1: !pdl.operation):
   }
 }
 
 // -----
 
+// expected-error @below {{op expects at least one non-pattern op}}
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  pdl.pattern @some : benefit(1) {
+    %0 = pdl.operation "test.foo"
+    pdl.rewrite %0 with "transform.dialect"
+  }
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  // expected-error @below {{op expects at least one non-pattern op}}
+  with_pdl_patterns %arg0 : !pdl.operation {
+  ^bb1(%arg1: !pdl.operation):
+  }
+}
+
+
+// -----
+
 // expected-error @below {{expects at least one region}}
 "transform.test_transform_unrestricted_op_no_interface"() : () -> ()
 
 // -----
 
 // expected-error @below {{expects a single-block region}}
 "transform.test_transform_unrestricted_op_no_interface"() ({
 ^bb0(%arg0: !pdl.operation):

mlir/test/Dialect/Transform/ops.mlir

 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   // CHECK: sequence %[[ARG]] : !pdl.operation
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
   }
 }
 
-// CHECK: transform.sequence
-// CHECK: ^{{.+}}(%[[ARG:.+]]: !pdl.operation):
-transform.sequence failures(propagate) {
-^bb0(%arg0: !pdl.operation):
-  // CHECK: with_pdl_patterns %[[ARG]] : !pdl.operation
-  with_pdl_patterns %arg0 : !pdl.operation {
-  ^bb1(%arg1: !pdl.operation):
-  }
-}
-
 // Using the same value multiple times without consuming it is fine.
 // CHECK: transform.sequence
 // CHECK: %[[V:.+]] = sequence %{{.*}} : !pdl.operation -> !pdl.operation
 // CHECK: sequence %[[V]]
 // CHECK: sequence %[[V]]
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   %0 = transform.sequence %arg0 : !pdl.operation -> !pdl.operation failures(propagate) {