[mlir][transform] Add PrintOp to transform dialect

[springerm]

Summary

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

[ftynse]

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

824

You can do ->print(llvm::errs()) and avoid this comment.

839–842

We have more readable helpers onlyReadsPayload, modifiesPayload, ...

[springerm]

address comments

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

let arguments = (ins
     SymbolRefAttr:$pattern_name);
   let results = (outs
     Res<TransformTypeInterface, "Handle to the matched Payload IR ops">:$matched);
 
   let assemblyFormat = "$pattern_name `in` $root attr-dict `:` "
                        "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>,
      AllTypesMatch<["handles", "replicated"]>]> {
   let summary = "Lists payload ops multiple times in the new handle";
   let description = [{
     Produces a new handle associated with a list of payload IR ops that is
     computed by repeating the list of payload IR ops associated with the

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

LogicalResult transform::WithPDLPatternsOp::verify() {
   if (auto parent = getOperation()->getParentOfType<WithPDLPatternsOp>()) {
     InFlightDiagnostic diag = emitOpError() << "cannot be nested";
     diag.attachNote(parent.getLoc()) << "parent operation";
     return diag;
   }
 
   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());
+}

mlir/test/Dialect/Transform/ops.mlir

 // CHECK: transform.sequence
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   // CHECK: cast %{{.*}} : !pdl.operation to !transform.any_op
   %0 = cast %arg0: !pdl.operation to !transform.any_op
   // 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"}
+}