[mlir][transform] Add VerifyOp

[springerm]

Summary

This transform op runs the verifier on the targeted payload ops. It is for debugging only.

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

def SequenceOp : TransformDialectOp<"sequence",
   let extraClassDeclaration = [{
     /// Allow the dialect prefix to be omitted.
     static StringRef getDefaultDialect() { return "transform"; }
   }];
 
   let hasVerifier = 1;
 }
 
+def VerifyOp : TransformDialectOp<"verify",
+    [TransformOpInterface, TransformEachOpTrait,
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     ReportTrackingListenerFailuresOpTrait]> {
+  let summary = "Verifies the targeted ops";
+  let description = [{
+    This transform verifies the targeted ops. If at least one op fails to
+    verify, the transform fails definitely.
+
+    Note: This op was designed for debugging purposes and should be used like an
+    assertion. It is intentional that this op produces a definite failure and
+    not a silenceable one. Correctness of the program should not depend on this
+    op.
+
+    This transform reads the target handle.
+  }];
+
+  let arguments = (ins TransformHandleTypeInterface:$target);
+  let results = (outs);
+  let assemblyFormat = "$target attr-dict `:` type($target)";
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+      ::mlir::transform::TransformRewriter &rewriter,
+      ::mlir::Operation *target,
+      ::mlir::transform::ApplyToEachResultList &results,
+      ::mlir::transform::TransformState &state);
+  }];
+}
+
 def YieldOp : TransformDialectOp<"yield",
     [Terminator, DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let summary = "Yields operation handles from a transform IR region";
   let description = [{
     This terminator operation yields operation handles from regions of the
     transform IR ops back to the containing op. It is not itself associated with
     any transformation on the payload IR and is used for flow purposes only.
   }];

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

 #include "mlir/Dialect/Transform/IR/TransformAttrs.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Verifier.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Pass/PassRegistry.h"
 #include "mlir/Transforms/CSE.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"
 #include "llvm/ADT/STLExtras.h"
void transform::PrintOp::getEffects(
   onlyReadsPayload(effects);
 
   // There is no resource for stderr file descriptor, so just declare print
   // writes into the default resource.
   effects.emplace_back(MemoryEffects::Write::get());
 }
 
 //===----------------------------------------------------------------------===//
+// VerifyOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::VerifyOp::applyToOne(transform::TransformRewriter &rewriter,
+                                Operation *target,
+                                transform::ApplyToEachResultList &results,
+                                transform::TransformState &state) {
+  if (failed(::mlir::verify(target))) {
+    DiagnosedDefiniteFailure diag = emitDefiniteFailure()
+                                    << "failed to verify payload op";
+    diag.attachNote(target->getLoc()) << "payload op";
+    return diag;
+  }
+  return DiagnosedSilenceableFailure::success();
+}
+
+void transform::VerifyOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getTarget(), effects);
+}
+
+//===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
 
 void transform::YieldOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   onlyReadsHandle(getOperands(), effects);
 }

mlir/test/Dialect/Transform/test-interpreter.mlir

 // CHECK-LABEL: @affine_map_param_anchor
 func.func private @affine_map_param_anchor()
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
   // CHECK: test_produce_param(#{{.*}}) : !transform.affine_map
   transform.test_produce_param(affine_map<(d0) -> ()>) : !transform.affine_map
 }
+
+// -----
+
+func.func @verify_success(%arg0: f64) -> f64 {
+  return %arg0 : f64
+}
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+  %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  transform.verify %0 : !transform.any_op
+}
+
+// -----
+
+// expected-error @below{{fail_to_verify is set}}
+// expected-note @below{{payload op}}
+func.func @verify_failure(%arg0: f64) -> f64 {
+  return %arg0 : f64
+}
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+  %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  transform.test_produce_invalid_ir %0 : !transform.any_op
+  // expected-error @below{{failed to verify payload op}}
+  transform.verify %0 : !transform.any_op
+}

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp

 }
 
 void mlir::test::TestDummyPayloadOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   for (OpResult result : getResults())
     transform::producesHandle(result, effects);
 }
 
+LogicalResult mlir::test::TestDummyPayloadOp::verify() {
+  if (getFailToVerify())
+    return emitOpError() << "fail_to_verify is set";
+  return success();
+}
+
 DiagnosedSilenceableFailure
 mlir::test::TestTrackedRewriteOp::apply(transform::TransformRewriter &rewriter,
                                         transform::TransformResults &results,
                                         transform::TransformState &state) {
   int64_t numIterations = 0;
 
   // `getPayloadOps` returns an iterator that skips ops that are erased in the
   // loop body. Replacement ops are not enumerated.

 }
 
 void mlir::test::TestNotifyPayloadOpReplacedOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getOriginal(), effects);
   transform::onlyReadsHandle(getReplacement(), effects);
 }
 
+DiagnosedSilenceableFailure mlir::test::TestProduceInvalidIR::applyToOne(
+    transform::TransformRewriter &rewriter, Operation *target,
+    transform::ApplyToEachResultList &results,
+    transform::TransformState &state) {
+  // Provide some IR that does not verify.
+  rewriter.setInsertionPointToStart(&target->getRegion(0).front());
+  rewriter.create<TestDummyPayloadOp>(target->getLoc(), TypeRange(),
+                                      ValueRange(), /*failToVerify=*/true);
+  return DiagnosedSilenceableFailure::success();
+}
+
+void mlir::test::TestProduceInvalidIR::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getTarget(), effects);
+  transform::modifiesPayload(effects);
+}
+
 namespace {
 /// Test extension of the Transform dialect. Registers additional ops and
 /// declares PDL as dependent dialect since the additional ops are using PDL
 /// types for operands and results.
 class TestTransformDialectExtension
     : public transform::TransformDialectExtension<
           TestTransformDialectExtension> {
 public:

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.td

 // be created with "RewriterBase::replaceOpWithNewOp" (needed for a test case).
 // Since only TransformOpInterface can be injected into the transform dialect,
 // this op implements the interface, even though it is not used as a transform
 // op.
 def TestDummyPayloadOp
   : Op<Transform_Dialect, "test_dummy_payload_op",
       [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
        TransformOpInterface]> {
-  let arguments = (ins Variadic<AnyType>:$args);
+  let arguments = (ins Variadic<AnyType>:$args,
+                       UnitAttr:$fail_to_verify);
   let results = (outs Variadic<AnyType>:$outs);
   let assemblyFormat = "$args attr-dict `:` functional-type(operands, results)";
   let cppNamespace = "::mlir::test";
+  let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     DiagnosedSilenceableFailure apply(transform::TransformRewriter &rewriter,
                                       transform::TransformResults &results,
                                       transform::TransformState &state) {
       llvm_unreachable("op should not be used as a transform");
       return DiagnosedSilenceableFailure::definiteFailure();
     }
: Op<Transform_Dialect, "test_notify_payload_op_replaced",
         DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let arguments = (ins TransformHandleTypeInterface:$original,
                        TransformHandleTypeInterface:$replacement);
   let results = (outs);
   let assemblyFormat = "$original `,` $replacement attr-dict `:` functional-type(operands, results)";
   let cppNamespace = "::mlir::test";
 }
 
+def TestProduceInvalidIR
+  : Op<Transform_Dialect, "test_produce_invalid_ir",
+       [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+        TransformOpInterface, TransformEachOpTrait]> {
+  let arguments = (ins TransformHandleTypeInterface:$target);
+  let results = (outs);
+  let assemblyFormat = "$target attr-dict `:` type($target)";
+  let cppNamespace = "::mlir::test";
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
+        ::mlir::Operation *target,
+        ::mlir::transform::ApplyToEachResultList &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
 #endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_TD