[mlir][transform] ApplyRegisteredPassOp: Support pass pipelines

[springerm]

Summary

The same transform op can now be used to apply registered pass pipelines.

This revision also adds a helper function for querying PassPipelineInfo objects and moves the corresponding lookup function for PassInfo objects to the PassInfo class.

[springerm]

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

722–724

Note: This is same as resolvePipelineElement: first look for a pass pipeline with the given name, if none was found look for a pass.

[springerm]

Closed by 2f8690b1e2bf09174660cd66ef8a27066848555e.

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

::mlir::DiagnosedSilenceableFailure applyToOne(
       ::mlir::transform::ApplyToEachResultList &results,
       ::mlir::transform::TransformState &state);
   }];
 }
 
 def ApplyRegisteredPassOp : TransformDialectOp<"apply_registered_pass",
     [TransformOpInterface, TransformEachOpTrait,
      FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface]> {
-  let summary = "Applies the specified registered pass";
+  let summary = "Applies the specified registered pass or pass pipeline";
   let description = [{
-    This transform applies the specified pass to the targeted ops. The name of
-    the pass is specified as a string attribute, as set during pass
-    registration. Optionally, pass options may be specified as a string
-    attribute. The pass options syntax is identical to the one used with
-    "mlir-opt".
+    This transform applies the specified pass or pass pipeline to the targeted
+    ops. The name of the pass/pipeline is specified as a string attribute, as
+    set during pass/pipeline registration. Optionally, pass options may be
+    specified as a string attribute. The pass options syntax is identical to the
+    one used with "mlir-opt".
+
+    This op first looks for a pass pipeline with the specified name. If no such
+    pipeline exists, it looks for a pass with the specified name. If no such
+    pass exists either, this op fails definitely.
 
     This transform consumes the target handle and produces a new handle that is
     mapped to the same op. Passes are not allowed to remove/modify the operation
     that they operate on, so the target op is guaranteed to still exist. The
     target handle is invalidated because a pass may arbitrarily modify the body
     of targeted ops.
   }];
 

mlir/include/mlir/Pass/Pass.h

 /// derived pass object, e.g its kind and abstract TypeID.
 class Pass {
 public:
   virtual ~Pass() = default;
 
   /// Returns the unique identifier that corresponds to this pass.
   TypeID getTypeID() const { return passID; }
 
-  /// Returns the pass info for the specified pass class or null if unknown.
-  static const PassInfo *lookupPassInfo(StringRef passArg);
-
   /// Returns the pass info for this pass, or null if unknown.
   const PassInfo *lookupPassInfo() const {
-    return lookupPassInfo(getArgument());
+    return PassInfo::lookup(getArgument());
   }
 
   /// Returns the derived pass name.
   virtual StringRef getName() const = 0;
 
   /// Register dependent dialects for the current pass.
   /// A pass is expected to register the dialects it will create entities for
   /// (Operations, Types, Attributes), other than dialect that exists in the

mlir/include/mlir/Pass/PassRegistry.h

 /// A structure to represent the information of a registered pass pipeline.
 class PassPipelineInfo : public PassRegistryEntry {
 public:
   PassPipelineInfo(
       StringRef arg, StringRef description, const PassRegistryFunction &builder,
       std::function<void(function_ref<void(const detail::PassOptions &)>)>
           optHandler)
       : PassRegistryEntry(arg, description, builder, std::move(optHandler)) {}
+
+  /// Returns the pass pipeline info for the specified pass pipeline or null if
+  /// unknown.
+  static const PassPipelineInfo *lookup(StringRef pipelineArg);
 };
 
 /// A structure to represent the information for a derived pass class.
 class PassInfo : public PassRegistryEntry {
 public:
   /// PassInfo constructor should not be invoked directly, instead use
   /// PassRegistration or registerPass.
   PassInfo(StringRef arg, StringRef description,
            const PassAllocatorFunction &allocator);
+
+  /// Returns the pass info for the specified pass class or null if unknown.
+  static const PassInfo *lookup(StringRef passArg);
 };
 
 //===----------------------------------------------------------------------===//
 // PassRegistration
 //===----------------------------------------------------------------------===//
 
 /// Register a specific dialect pipeline registry function with the system,
 /// typically used through the PassPipelineRegistration template.

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

DiagnosedSilenceableFailure transform::ApplyRegisteredPassOp::applyToOne(
   // transform IR while it is being interpreted is generally dangerous. Even
   // more so when applying passes because they may perform a wide range of IR
   // modifications.
   DiagnosedSilenceableFailure payloadCheck =
       ensurePayloadIsSeparateFromTransform(*this, target);
   if (!payloadCheck.succeeded())
     return payloadCheck;
 
-  // Get pass from registry.
-  const PassInfo *passInfo = Pass::lookupPassInfo(getPassName());
-  if (!passInfo) {
-    return emitDefiniteFailure() << "unknown pass: " << getPassName();
-  }
+  // Get pass or pass pipeline from registry.
+  const PassRegistryEntry *info = PassPipelineInfo::lookup(getPassName());
+  if (!info)
+    info = PassInfo::lookup(getPassName());
+  if (!info)
+    return emitDefiniteFailure()
+           << "unknown pass or pass pipeline: " << getPassName();
 
-  // Create pass manager with a single pass and run it.
+  // Create pass manager and run the pass or pass pipeline.
   PassManager pm(getContext());
-  if (failed(passInfo->addToPipeline(pm, getOptions(), [&](const Twine &msg) {
+  if (failed(info->addToPipeline(pm, getOptions(), [&](const Twine &msg) {
         emitError(msg);
         return failure();
       }))) {
     return emitDefiniteFailure()
-           << "failed to add pass to pipeline: " << getPassName();
+           << "failed to add pass or pass pipeline to pipeline: "
+           << getPassName();
   }
   if (failed(pm.run(target))) {
     auto diag = emitSilenceableError() << "pass pipeline failed";
     diag.attachNote(target->getLoc()) << "target op";
     return diag;
   }
 
   results.push_back(target);

mlir/lib/Pass/PassRegistry.cpp

void mlir::registerPass(const PassAllocatorFunction &function) {
   if (it->second != entryTypeID)
     llvm::report_fatal_error(
         "pass allocator creates a different pass than previously "
         "registered for pass " +
         arg);
 }
 
 /// Returns the pass info for the specified pass argument or null if unknown.
-const PassInfo *mlir::Pass::lookupPassInfo(StringRef passArg) {
+const PassInfo *mlir::PassInfo::lookup(StringRef passArg) {
   auto it = passRegistry->find(passArg);
   return it == passRegistry->end() ? nullptr : &it->second;
 }
 
+/// Returns the pass pipeline info for the specified pass pipeline argument or
+/// null if unknown.
+const PassPipelineInfo *mlir::PassPipelineInfo::lookup(StringRef pipelineArg) {
+  auto it = passPipelineRegistry->find(pipelineArg);
+  return it == passPipelineRegistry->end() ? nullptr : &it->second;
+}
+
 //===----------------------------------------------------------------------===//
 // PassOptions
 //===----------------------------------------------------------------------===//
 
 LogicalResult detail::pass_options::parseCommaSeparatedList(
     llvm::cl::Option &opt, StringRef argName, StringRef optionStr,
     function_ref<LogicalResult(StringRef)> elementParseFn) {
   // Functor used for finding a character in a string, and skipping over

 /// Resolve a single element of the pipeline.
 LogicalResult
 TextualPipeline::resolvePipelineElement(PipelineElement &element,
                                         ErrorHandlerT errorHandler) {
   // If the inner pipeline of this element is not empty, this is an operation
   // pipeline.
   if (!element.innerPipeline.empty())
     return resolvePipelineElements(element.innerPipeline, errorHandler);
+
   // Otherwise, this must be a pass or pass pipeline.
   // Check to see if a pipeline was registered with this name.
-  auto pipelineRegistryIt = passPipelineRegistry->find(element.name);
-  if (pipelineRegistryIt != passPipelineRegistry->end()) {
-    element.registryEntry = &pipelineRegistryIt->second;
+  if ((element.registryEntry = PassPipelineInfo::lookup(element.name)))
     return success();
-  }
 
   // If not, then this must be a specific pass name.
-  if ((element.registryEntry = Pass::lookupPassInfo(element.name)))
+  if ((element.registryEntry = PassInfo::lookup(element.name)))
     return success();
 
   // Emit an error for the unknown pass.
   auto *rawLoc = element.name.data();
   return errorHandler(rawLoc, "'" + element.name +
                                   "' does not refer to a "
                                   "registered pass or pass pipeline");
 }

mlir/test/Dialect/Transform/test-pass-application.mlir

 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %1 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   transform.apply_registered_pass "canonicalize" to %1 : (!transform.any_op) -> !transform.any_op
 }
 
 // -----
 
+// CHECK-LABEL: func @pass_pipeline(
+func.func @pass_pipeline() {
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %1 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  // This pipeline does not do anything. Just make sure that the pipeline is
+  // found and no error is produced.
+  transform.apply_registered_pass "test-options-pass-pipeline" to %1 : (!transform.any_op) -> !transform.any_op
+}
+
+// -----
+
 func.func @invalid_pass_name() {
   return
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %1 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  // expected-error @below {{unknown pass: non-existing-pass}}
+  // expected-error @below {{unknown pass or pass pipeline: non-existing-pass}}
   transform.apply_registered_pass "non-existing-pass" to %1 : (!transform.any_op) -> !transform.any_op
 }
 
 // -----
 
 func.func @not_isolated_from_above(%t: tensor<5xf32>) -> index {
   %c0 = arith.constant 0 : index
   // expected-note @below {{target op}}

 
 func.func @invalid_pass_option() {
   return
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %1 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  // expected-error @below {{failed to add pass to pipeline: canonicalize}}
+  // expected-error @below {{failed to add pass or pass pipeline to pipeline: canonicalize}}
   transform.apply_registered_pass "canonicalize" to %1 {options = "invalid-option=1"} : (!transform.any_op) -> !transform.any_op
 }
 
 // -----
 
 // CHECK-LABEL: func @valid_pass_option()
 func.func @valid_pass_option() {
   return