[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

1013

Nit: expand auto.

1014–1015

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

1014–1015

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 (!topLevelOp) {
+    if (auto *parentOp = getOperation()->getParentOp()) {
+      if (!parentOp->hasTrait<
+              ::mlir::transform::PossibleTopLevelTransformOpTrait>()) {
+        InFlightDiagnostic diag = emitOpError()
+                                  << "expects at least one non-pattern op";
+        return diag;
+      }
+    }
+  }
+
   if (auto parent = getOperation()->getParentOfType<WithPDLPatternsOp>()) {
     InFlightDiagnostic diag = emitOpError() << "cannot be nested";
     diag.attachNote(parent.getLoc()) << "parent operation";
     return diag;
   }
 
   return success();
 }

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"
+  }
+}
+
+// -----
+
 // 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):