Implement Pass and Dialect plugins for mlir-opt

Implementation of Pass and Dialect Plugins that mirrors LLVM Pass Plugin
implementation from the new pass manager.

Currently the implementation only supports using the pass-pipeline option
for adding passes. This restriction is imposed by the `PassPipelineCLParser`
variable in mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:114 that loads the
parse options statically before parsing the cmd line args.

```
mlir-opt stanalone-plugin.mlir --load-dialect-plugin=lib/libStandalonePlugin.so --pass-pipeline="builtin.module(standalone-switch-bar-foo)"
```

Reviewed By: rriddle, mehdi_amini

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

12 files changed, 205 insertions, 0 deletions

diff --git a/mlir/examples/standalone/CMakeLists.txt b/mlir/examples/standalone/CMakeLists.txt
index d36a6badc16d..65461c048b3f 100644
--- a/mlir/examples/standalone/CMakeLists.txt
+++ b/mlir/examples/standalone/CMakeLists.txt
@@ -52,4 +52,5 @@ if(MLIR_ENABLE_BINDINGS_PYTHON)
 endif()
 add_subdirectory(test)
 add_subdirectory(standalone-opt)
+add_subdirectory(standalone-plugin)
 add_subdirectory(standalone-translate)
diff --git a/mlir/examples/standalone/include/Standalone/CMakeLists.txt b/mlir/examples/standalone/include/Standalone/CMakeLists.txt
index 8acf640037ee..975a4ffaa2aa 100644
--- a/mlir/examples/standalone/include/Standalone/CMakeLists.txt
+++ b/mlir/examples/standalone/include/Standalone/CMakeLists.txt
@@ -1,3 +1,7 @@
 add_mlir_dialect(StandaloneOps standalone)
 add_mlir_doc(StandaloneDialect StandaloneDialect Standalone/ -gen-dialect-doc)
 add_mlir_doc(StandaloneOps StandaloneOps Standalone/ -gen-op-doc)
+
+set(LLVM_TARGET_DEFINITIONS StandalonePasses.td)
+mlir_tablegen(StandalonePasses.h.inc --gen-pass-decls)
+add_public_tablegen_target(MLIRStandalonePassesIncGen)
diff --git a/mlir/examples/standalone/include/Standalone/StandalonePasses.h b/mlir/examples/standalone/include/Standalone/StandalonePasses.h
new file mode 100644
index 000000000000..75546d688fc1
--- /dev/null
+++ b/mlir/examples/standalone/include/Standalone/StandalonePasses.h
@@ -0,0 +1,26 @@
+//===- StandalonePasses.h - Standalone passes  ------------------*- C++ -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef STANDALONE_STANDALONEPASSES_H
+#define STANDALONE_STANDALONEPASSES_H
+
+#include "Standalone/StandaloneDialect.h"
+#include "Standalone/StandaloneOps.h"
+#include "mlir/Pass/Pass.h"
+#include <memory>
+
+namespace mlir {
+namespace standalone {
+#define GEN_PASS_DECL
+#include "Standalone/StandalonePasses.h.inc"
+
+#define GEN_PASS_REGISTRATION
+#include "Standalone/StandalonePasses.h.inc"
+} // namespace standalone
+} // namespace mlir
+
+#endif
diff --git a/mlir/examples/standalone/include/Standalone/StandalonePasses.td b/mlir/examples/standalone/include/Standalone/StandalonePasses.td
new file mode 100644
index 000000000000..4cb2be02e4a2
--- /dev/null
+++ b/mlir/examples/standalone/include/Standalone/StandalonePasses.td
@@ -0,0 +1,30 @@
+//===- StandalonePsss.td - Standalone dialect passes -------*- tablegen -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef STANDALONE_PASS
+#define STANDALONE_PASS
+
+include "mlir/Pass/PassBase.td"
+
+def StandaloneSwitchBarFoo: Pass<"standalone-switch-bar-foo", "::mlir::ModuleOp"> {
+  let summary = "Switches the name of a FuncOp named `bar` to `foo` and folds.";
+  let description = [{
+    Switches the name of a FuncOp named `bar` to `foo` and folds.
+    ```
+    func.func @bar() {
+      return
+    }
+    // Gets transformed to:
+    func.func @foo() {
+      return
+    }
+    ```
+  }];
+}
+
+#endif // STANDALONE_PASS
diff --git a/mlir/examples/standalone/lib/Standalone/CMakeLists.txt b/mlir/examples/standalone/lib/Standalone/CMakeLists.txt
index 599c1c55c792..0e2d043665c3 100644
--- a/mlir/examples/standalone/lib/Standalone/CMakeLists.txt
+++ b/mlir/examples/standalone/lib/Standalone/CMakeLists.txt
@@ -2,14 +2,17 @@ add_mlir_dialect_library(MLIRStandalone
         StandaloneTypes.cpp
         StandaloneDialect.cpp
         StandaloneOps.cpp
+        StandalonePasses.cpp
 
         ADDITIONAL_HEADER_DIRS
         ${PROJECT_SOURCE_DIR}/include/Standalone
 
         DEPENDS
         MLIRStandaloneOpsIncGen
+        MLIRStandalonePassesIncGen
 
 	LINK_LIBS PUBLIC
 	MLIRIR
         MLIRInferTypeOpInterface
+        MLIRFuncDialect
 	)
diff --git a/mlir/examples/standalone/lib/Standalone/StandalonePasses.cpp b/mlir/examples/standalone/lib/Standalone/StandalonePasses.cpp
new file mode 100644
index 000000000000..6af45c9769a8
--- /dev/null
+++ b/mlir/examples/standalone/lib/Standalone/StandalonePasses.cpp
@@ -0,0 +1,48 @@
+//===- StandalonePasses.cpp - Standalone passes -----------------*- C++ -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+#include "Standalone/StandalonePasses.h"
+
+namespace mlir::standalone {
+#define GEN_PASS_DEF_STANDALONESWITCHBARFOO
+#include "Standalone/StandalonePasses.h.inc"
+
+namespace {
+class StandaloneSwitchBarFooRewriter : public OpRewritePattern<func::FuncOp> {
+public:
+  using OpRewritePattern<func::FuncOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(func::FuncOp op,
+                                PatternRewriter &rewriter) const final {
+    if (op.getSymName() == "bar") {
+      rewriter.updateRootInPlace(op, [&op]() { op.setSymName("foo"); });
+      return success();
+    }
+    return failure();
+  }
+};
+
+class StandaloneSwitchBarFoo
+    : public impl::StandaloneSwitchBarFooBase<StandaloneSwitchBarFoo> {
+public:
+  using impl::StandaloneSwitchBarFooBase<
+      StandaloneSwitchBarFoo>::StandaloneSwitchBarFooBase;
+  void runOnOperation() final {
+    RewritePatternSet patterns(&getContext());
+    patterns.add<StandaloneSwitchBarFooRewriter>(&getContext());
+    FrozenRewritePatternSet patternSet(std::move(patterns));
+    if (failed(applyPatternsAndFoldGreedily(getOperation(), patternSet)))
+      signalPassFailure();
+  }
+};
+} // namespace
+} // namespace mlir::standalone
diff --git a/mlir/examples/standalone/standalone-opt/standalone-opt.cpp b/mlir/examples/standalone/standalone-opt/standalone-opt.cpp
index 4cf7e15de593..e75db35e5260 100644
--- a/mlir/examples/standalone/standalone-opt/standalone-opt.cpp
+++ b/mlir/examples/standalone/standalone-opt/standalone-opt.cpp
@@ -21,9 +21,11 @@
 #include "llvm/Support/ToolOutputFile.h"
 
 #include "Standalone/StandaloneDialect.h"
+#include "Standalone/StandalonePasses.h"
 
 int main(int argc, char **argv) {
   mlir::registerAllPasses();
+  mlir::standalone::registerPasses();
   // TODO: Register standalone passes here.
 
   mlir::DialectRegistry registry;
diff --git a/mlir/examples/standalone/standalone-plugin/CMakeLists.txt b/mlir/examples/standalone/standalone-plugin/CMakeLists.txt
new file mode 100644
index 000000000000..961a3ea2c906
--- /dev/null
+++ b/mlir/examples/standalone/standalone-plugin/CMakeLists.txt
@@ -0,0 +1,22 @@
+get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
+set(LIBS
+        MLIRIR
+        MLIRPass
+        MLIRPluginsLib
+        MLIRStandalone
+        MLIRTransformUtils
+        )
+
+add_mlir_dialect_library(StandalonePlugin
+        SHARED
+        standalone-plugin.cpp
+
+        DEPENDS
+        MLIRStandalone
+        )
+
+llvm_update_compile_flags(StandalonePlugin)
+target_link_libraries(StandalonePlugin PRIVATE ${LIBS})
+
+mlir_check_all_link_libraries(StandalonePlugin)
diff --git a/mlir/examples/standalone/standalone-plugin/standalone-plugin.cpp b/mlir/examples/standalone/standalone-plugin/standalone-plugin.cpp
new file mode 100644
index 000000000000..129b86b9db55
--- /dev/null
+++ b/mlir/examples/standalone/standalone-plugin/standalone-plugin.cpp
@@ -0,0 +1,39 @@
+//===- standalone-plugin.cpp ------------------------------------*- C++ -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/InitAllDialects.h"
+#include "mlir/InitAllPasses.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Tools/Plugins/DialectPlugin.h"
+
+#include "Standalone/StandaloneDialect.h"
+#include "Standalone/StandalonePasses.h"
+
+using namespace mlir;
+
+/// Dialect plugin registration mechanism.
+/// Observe that it also allows to register passes.
+/// Necessary symbol to register the dialect plugin.
+extern "C" LLVM_ATTRIBUTE_WEAK DialectPluginLibraryInfo
+mlirGetDialectPluginInfo() {
+  return {MLIR_PLUGIN_API_VERSION, "Standalone", LLVM_VERSION_STRING,
+          [](DialectRegistry *registry) {
+            registry->insert<mlir::standalone::StandaloneDialect>();
+            mlir::standalone::registerPasses();
+          }};
+}
+
+/// Pass plugin registration mechanism.
+/// Necessary symbol to register the pass plugin.
+extern "C" LLVM_ATTRIBUTE_WEAK PassPluginLibraryInfo mlirGetPassPluginInfo() {
+  return {MLIR_PLUGIN_API_VERSION, "StandalonePasses", LLVM_VERSION_STRING,
+          []() { mlir::standalone::registerPasses(); }};
+}
diff --git a/mlir/examples/standalone/test/Standalone/standalone-pass-plugin.mlir b/mlir/examples/standalone/test/Standalone/standalone-pass-plugin.mlir
new file mode 100644
index 000000000000..5af4b3d92e34
--- /dev/null
+++ b/mlir/examples/standalone/test/Standalone/standalone-pass-plugin.mlir
@@ -0,0 +1,13 @@
+// RUN: mlir-opt %s --load-pass-plugin=%standalone_libs/libStandalonePlugin.so --pass-pipeline="builtin.module(standalone-switch-bar-foo)" | FileCheck %s
+
+module {
+  // CHECK-LABEL: func @foo()
+  func.func @bar() {
+    return
+  }
+
+  // CHECK-LABEL: func @abar()
+  func.func @abar() {
+    return
+  }
+}
diff --git a/mlir/examples/standalone/test/Standalone/standalone-plugin.mlir b/mlir/examples/standalone/test/Standalone/standalone-plugin.mlir
new file mode 100644
index 000000000000..3f935db8910c
--- /dev/null
+++ b/mlir/examples/standalone/test/Standalone/standalone-plugin.mlir
@@ -0,0 +1,13 @@
+// RUN: mlir-opt %s --load-dialect-plugin=%standalone_libs/libStandalonePlugin.so --pass-pipeline="builtin.module(standalone-switch-bar-foo)" | FileCheck %s
+
+module {
+  // CHECK-LABEL: func @foo()
+  func.func @bar() {
+    return
+  }
+
+  // CHECK-LABEL: func @standalone_types(%arg0: !standalone.custom<"10">)
+  func.func @standalone_types(%arg0: !standalone.custom<"10">) {
+    return
+  }
+}
diff --git a/mlir/examples/standalone/test/lit.cfg.py b/mlir/examples/standalone/test/lit.cfg.py
index a6a3d2444d0c..3e4ceee3865c 100644
--- a/mlir/examples/standalone/test/lit.cfg.py
+++ b/mlir/examples/standalone/test/lit.cfg.py
@@ -44,12 +44,16 @@ config.excludes = ['Inputs', 'Examples', 'CMakeLists.txt', 'README.txt', 'LICENS
 # test_exec_root: The root path where tests should be run.
 config.test_exec_root = os.path.join(config.standalone_obj_root, 'test')
 config.standalone_tools_dir = os.path.join(config.standalone_obj_root, 'bin')
+config.standalone_libs_dir = os.path.join(config.standalone_obj_root, 'lib')
+
+config.substitutions.append(('%standalone_libs', config.standalone_libs_dir))
 
 # Tweak the PATH to include the tools dir.
 llvm_config.with_environment('PATH', config.llvm_tools_dir, append_path=True)
 
 tool_dirs = [config.standalone_tools_dir, config.llvm_tools_dir]
 tools = [
+    'mlir-opt',
     'standalone-capi-test',
     'standalone-opt',
     'standalone-translate',