[mlir][transform] Add PrintOp to transform dialect

This op dumps the associated payload IR to stderr. It has proven useful for printf-style debugging.

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

3 files changed, 66 insertions, 0 deletions

diff --git a/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td b/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
index bc20edfed77d..4b1bb02ee757 100644
--- a/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
+++ b/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
@@ -290,6 +290,25 @@ def PDLMatchOp : TransformDialectOp<"pdl_match",
                        "functional-type(operands, results)";
 }
 
+def PrintOp : TransformDialectOp<"print",
+    [DeclareOpInterfaceMethods<TransformOpInterface>,
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let summary = "Dump each payload op";
+  let description = [{
+    This op dumps each payload op that is associated with the `target` operand
+    to stderr. It also prints the `name` string attribute. If no target is
+    specified, the top-level op is dumped.
+
+    This op is useful for printf-style debugging.
+  }];
+
+  let arguments = (ins Optional<TransformTypeInterface>:$target,
+                       OptionalAttr<StrAttr>:$name);
+  let results = (outs);
+  let assemblyFormat = "$target attr-dict (`:` type($target)^)?";
+}
+
+
 def ReplicateOp : TransformDialectOp<"replicate",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
diff --git a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
index 9d9f1ea67adf..2be1bea91fbe 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
+++ b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
@@ -807,3 +807,37 @@ LogicalResult transform::WithPDLPatternsOp::verify() {
 
   return success();
 }
+
+//===----------------------------------------------------------------------===//
+// PrintOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::PrintOp::apply(transform::TransformResults &results,
+                          transform::TransformState &state) {
+  llvm::errs() << "[[[ IR printer: ";
+  if (getName().has_value())
+    llvm::errs() << *getName() << " ";
+
+  if (!getTarget()) {
+    llvm::errs() << "top-level ]]]\n" << *state.getTopLevel() << "\n";
+    return DiagnosedSilenceableFailure::success();
+  }
+
+  llvm::errs() << "]]]\n";
+  ArrayRef<Operation *> targets = state.getPayloadOps(getTarget());
+  for (Operation *target : targets)
+    llvm::errs() << *target << "\n";
+
+  return DiagnosedSilenceableFailure::success();
+}
+
+void transform::PrintOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  onlyReadsHandle(getTarget(), effects);
+  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());
+}
diff --git a/mlir/test/Dialect/Transform/ops.mlir b/mlir/test/Dialect/Transform/ops.mlir
index 506fce063d61..0d27f9248556 100644
--- a/mlir/test/Dialect/Transform/ops.mlir
+++ b/mlir/test/Dialect/Transform/ops.mlir
@@ -67,3 +67,16 @@ transform.sequence failures(propagate) {
   // CHECK: cast %{{.*}} : !transform.any_op to !transform.op<"builtin.module">
   %1 = cast %0: !transform.any_op to !transform.op<"builtin.module">
 }
+
+// CHECK: transform.sequence
+// CHECK: print
+// CHECK: print
+// CHECK: print
+// CHECK: print
+transform.sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  transform.print %arg0 : !pdl.operation
+  transform.print
+  transform.print %arg0 {name = "test"} : !pdl.operation
+  transform.print {name = "test"}
+}