[mlir][transform] Add ApplyRegisteredPassOp transform op

[springerm]

Summary

This transform op runs a pass on the target op.

[ftynse]

I would consider evolving this into constructors of the entire pass pipeline rather than eager runners. So it can benefit from, e.g., parallelization and analysis preservation.

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

384–398

Can this be factored out? Something similar exists for "apply patterns".

[springerm]

rebase

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

def ApplyCanonicalizationPatternsOp
   let summary = "Populates canonicalization patterns";
   let description = [{
     This op populates all canonicalization patterns of all loaded dialects in
     an `apply_patterns` transform.
   }];
   let assemblyFormat = "attr-dict";
 }
 
+def ApplyRegisteredPassOp : TransformDialectOp<"apply_registered_pass",
+    [TransformOpInterface, TransformEachOpTrait,
+     FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface]> {
+  let summary = "Applies the specified registered pass";
+  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 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.
+  }];
+
+  let arguments = (ins TransformHandleTypeInterface:$target,
+                       StrAttr:$pass_name,
+                       DefaultValuedAttr<StrAttr, "\"\"">:$options);
+  let results = (outs TransformHandleTypeInterface:$result);
+  let assemblyFormat = [{
+    $pass_name `to` $target attr-dict `:` functional-type(operands, results)
+  }];
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+      ::mlir::Operation *target,
+      ::mlir::transform::ApplyToEachResultList &results,
+      ::mlir::transform::TransformState &state);
+  }];
+}
+
 def CastOp : TransformDialectOp<"cast",
     [TransformOpInterface, TransformEachOpTrait,
      DeclareOpInterfaceMethods<CastOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let arguments = (ins TransformHandleTypeInterface:$input);
   let results = (outs TransformHandleTypeInterface:$output);
   let assemblyFormat = "$input attr-dict `:` type($input) `to` type($output)";
 

mlir/lib/Dialect/Transform/IR/CMakeLists.txt

add_mlir_dialect_library(MLIRTransformDialect
   MLIRMatchInterfacesIncGen
   MLIRTransformDialectIncGen
   MLIRTransformInterfacesIncGen
 
   LINK_LIBS PUBLIC
   MLIRCastInterfaces
   MLIRIR
   MLIRParser
+  MLIRPass
   MLIRRewrite
   MLIRSideEffectInterfaces
   MLIRTransforms
   MLIRTransformDialectUtils
   )

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/FunctionImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Pass/PassRegistry.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
 #include <optional>
 
 #define DEBUG_TYPE "transform-dialect"
static void printSequenceOpOperands(OpAsmPrinter &printer, Operation *op,
                                     ValueRange extraBindings,
                                     TypeRange extraBindingTypes);
 static void printForeachMatchSymbols(OpAsmPrinter &printer, Operation *op,
                                      ArrayAttr matchers, ArrayAttr actions);
 static ParseResult parseForeachMatchSymbols(OpAsmParser &parser,
                                             ArrayAttr &matchers,
                                             ArrayAttr &actions);
 
+/// Helper function to check if the given transform op is contained in (or
+/// equal to) the given payload target op. In that case, an error is returned.
+/// Transforming transform IR that is currently executing is generally unsafe.
+static DiagnosedSilenceableFailure
+ensurePayloadIsSeparateFromTransform(transform::TransformOpInterface transform,
+                                     Operation *payload) {
+  Operation *transformAncestor = transform.getOperation();
+  while (transformAncestor) {
+    if (transformAncestor == payload) {
+      DiagnosedDefiniteFailure diag =
+          transform.emitDefiniteFailure()
+          << "cannot apply transform to itself (or one of its ancestors)";
+      diag.attachNote(payload->getLoc()) << "target payload op";
+      return diag;
+    }
+    transformAncestor = transformAncestor->getParentOp();
+  }
+  return DiagnosedSilenceableFailure::success();
+}
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Transform/IR/TransformOps.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // AlternativesOp
 //===----------------------------------------------------------------------===//
 
 OperandRange transform::AlternativesOp::getSuccessorEntryOperands(

 DiagnosedSilenceableFailure
 transform::ApplyPatternsOp::applyToOne(Operation *target,
                                        ApplyToEachResultList &results,
                                        transform::TransformState &state) {
   // Make sure that this transform is not applied to itself. Modifying the
   // transform IR while it is being interpreted is generally dangerous. Even
   // more so for the ApplyPatternsOp because the GreedyPatternRewriteDriver
   // performs many additional simplifications such as dead code elimination.
-  Operation *transformAncestor = getOperation();
-  while (transformAncestor) {
-    if (transformAncestor == target) {
-      DiagnosedDefiniteFailure diag =
-          emitDefiniteFailure()
-          << "cannot apply transform to itself (or one of its ancestors)";
-      diag.attachNote(target->getLoc()) << "target payload op";
-      return diag;
-    }
-    transformAncestor = transformAncestor->getParentOp();
-  }
+  DiagnosedSilenceableFailure payloadCheck =
+      ensurePayloadIsSeparateFromTransform(*this, target);
+  if (!payloadCheck.succeeded())
+    return payloadCheck;
 
   // Gather all specified patterns.
   MLIRContext *ctx = target->getContext();
   RewritePatternSet patterns(ctx);
   if (!getRegion().empty()) {
     for (Operation &op : getRegion().front()) {
       cast<transform::PatternDescriptorOpInterface>(&op).populatePatterns(
           patterns);
void transform::ApplyCanonicalizationPatternsOp::populatePatterns(
   MLIRContext *ctx = patterns.getContext();
   for (Dialect *dialect : ctx->getLoadedDialects())
     dialect->getCanonicalizationPatterns(patterns);
   for (RegisteredOperationName op : ctx->getRegisteredOperations())
     op.getCanonicalizationPatterns(patterns, ctx);
 }
 
 //===----------------------------------------------------------------------===//
+// ApplyRegisteredPassOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::ApplyRegisteredPassOp::applyToOne(Operation *target,
+                                             ApplyToEachResultList &results,
+                                             transform::TransformState &state) {
+  // Make sure that this transform is not applied to itself. Modifying the
+  // 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();
+  }
+
+  // Create pass manager with a single pass and run it.
+  PassManager pm(getContext());
+  if (failed(passInfo->addToPipeline(pm, getOptions(), [&](const Twine &msg) {
+        emitError(msg);
+        return failure();
+      }))) {
+    return emitDefiniteFailure()
+           << "failed to add pass 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);
+  return DiagnosedSilenceableFailure::success();
+}
+
+//===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::CastOp::applyToOne(Operation *target, ApplyToEachResultList &results,
                               transform::TransformState &state) {
   results.push_back(target);
   return DiagnosedSilenceableFailure::success();

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

+// RUN: mlir-opt %s --test-transform-dialect-interpreter -allow-unregistered-dialect --split-input-file --verify-diagnostics | FileCheck %s
+
+// CHECK-LABEL: func @successful_pass_application(
+//       CHECK:   %[[c5:.*]] = arith.constant 5 : index
+//       CHECK:   return %[[c5]]
+func.func @successful_pass_application(%t: tensor<5xf32>) -> index {
+  %c0 = arith.constant 0 : index
+  %dim = tensor.dim %t, %c0 : tensor<5xf32>
+  return %dim : index
+}
+
+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
+}
+
+// -----
+
+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}}
+  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}}
+  // expected-error @below {{trying to schedule a pass on an operation not marked as 'IsolatedFromAbove'}}
+  %dim = tensor.dim %t, %c0 : tensor<5xf32>
+  return %dim : index
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %1 = transform.structured.match ops{["tensor.dim"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  // expected-error @below {{pass pipeline failed}}
+  transform.apply_registered_pass "canonicalize" to %1 : (!transform.any_op) -> !transform.any_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}}
+  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
+}
+
+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 {options = "top-down=false"} : (!transform.any_op) -> !transform.any_op
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     srcs = glob(["lib/Dialect/Transform/IR/*.cpp"]),
     hdrs = glob(["include/mlir/Dialect/Transform/IR/*.h"]),
     deps = [
         ":Analysis",
         ":CallOpInterfaces",
         ":CastInterfaces",
         ":ControlFlowInterfaces",
         ":IR",
+        ":Pass",
         ":Rewrite",
         ":SideEffectInterfaces",
         ":Support",
         ":TransformDialectEnumsIncGen",
         ":TransformDialectIncGen",
         ":TransformDialectInterfacesIncGen",
         ":TransformDialectMatchInterfacesIncGen",
         ":TransformDialectUtils",