First step of adding compiler plugins to the build system. (#12598)

* CMake plumbing to support static, linked-in plugins.
* Example plugin added.
* PluginManager wired into compiler driver and API.

Per RFC #12520

CMakeLists.txt

@@ -445,6 +445,7 @@ include(iree_copts)
 include(iree_cc_binary)
 include(iree_cc_library)
 include(iree_cc_test)
+include(iree_compiler_plugin)
 include(iree_import_binary)
 include(iree_external_cmake_options)
 include(iree_tablegen_library)
@@ -733,6 +734,7 @@ else()
   endif()
 
   set(BUILD_SHARED_LIBS ${_IREE_ORIG_BUILD_SHARED_LIBS} CACHE BOOL "" FORCE)
+  iree_compiler_configure_plugins()
 endif()
 
 #-------------------------------------------------------------------------------

build_tools/bazel/build_defs.oss.bzl

@@ -73,6 +73,13 @@ def iree_compiler_cc_library(deps = [], **kwargs):
         **kwargs
     )
 
+def iree_compiler_register_plugin(plugin_id, target):
+    """Mirror of the CMake iree_compiler_register_plugin function.
+
+    Does nothing in bazel currently.
+    """
+    pass
+
 def iree_compiler_cc_test(deps = [], **kwargs):
     """Used for cc_test targets within the //compiler tree.
 

build_tools/bazel_to_cmake/bazel_to_cmake_converter.py

@@ -367,6 +367,16 @@ def cc_library(self,
   def iree_compiler_cc_library(self, deps=[], **kwargs):
     self.cc_library(deps=deps + ["//compiler/src:defs"], **kwargs)
 
+  def iree_compiler_register_plugin(self, plugin_id, target):
+    plugin_id_block = _convert_string_arg_block("PLUGIN_ID",
+                                                plugin_id,
+                                                quote=False)
+    target_block = _convert_single_target_block("TARGET", target)
+    self.converter.body += (f"iree_compiler_register_plugin(\n"
+                            f"{plugin_id_block}"
+                            f"{target_block}"
+                            f")\n\n")
+
   def iree_runtime_cc_library(self, deps=[], **kwargs):
     self.cc_library(deps=deps + ["//runtime/src:runtime_defines"], **kwargs)
 

build_tools/cmake/iree_compiler_plugin.cmake

@@ -0,0 +1,86 @@
+# Copyright 2023 The IREE Authors
+#
+# 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
+
+set(IREE_COMPILER_PLUGINS "" CACHE STRING "List of named in-tree plugins (under compiler/plugins) to statically compile")
+
+# Ids of all plugins that have been included in the configure step. This
+# may include plugins that we do not statically link but we do build.
+set_property(GLOBAL PROPERTY IREE_COMPILER_INCLUDED_PLUGIN_IDS "")
+# Ids of all plugins that we have statically linked.
+set_property(GLOBAL PROPERTY IREE_COMPILER_LINKED_PLUGIN_IDS "")
+# Libraries to add to the plugin registry for all statically linked plugins.
+set_property(GLOBAL PROPERTY IREE_COMPILER_LINKED_PLUGIN_LIBS "")
+
+# iree_compiler_register_plugin()
+# Within a plugin package, registers the plugin by id with the build system,
+# associating it with a registration target.
+function(iree_compiler_register_plugin)
+  cmake_parse_arguments(
+    _RULE
+    ""
+    "PACKAGE;PLUGIN_ID;TARGET"
+    ""
+    ${ARGN}
+  )
+
+  # Prefix the library with the package name, so we get: iree_package_name.
+  if(_RULE_PACKAGE)
+    set(_PACKAGE_NS "${_RULE_PACKAGE}")
+    string(REPLACE "::" "_" _PACKAGE_NAME ${_RULE_PACKAGE})
+  else()
+    iree_package_ns(_PACKAGE_NS)
+    iree_package_name(_PACKAGE_NAME)
+  endif()
+
+  # Replace target passed by ::name with iree::package::name
+  list(TRANSFORM _RULE_TARGET REPLACE "^::" "${_PACKAGE_NS}::")
+
+  # Do nothing if plugin registration is taking place standalone.
+  # See the IREE_COMPILER_IN_ADD_PLUGIN variable set in
+  # iree_compiler_add_plugin.
+  if(NOT IREE_COMPILER_IN_ADD_PLUGIN)
+    message(STATUS "Not registering plugin (out of tree)")
+    return()
+  endif()
+
+  # TODO: Can have more control on what gets linked.
+  message(STATUS "Registering static linked compiler plugin ${_RULE_PLUGIN_ID}")
+  set_property(GLOBAL APPEND PROPERTY IREE_COMPILER_LINKED_PLUGIN_IDS "${_RULE_PLUGIN_ID}")
+  set_property(GLOBAL APPEND PROPERTY IREE_COMPILER_LINKED_PLUGIN_LIBS "${_RULE_TARGET}")
+endfunction()
+
+# iree_compiler_configure_plugins()
+# Configures all in-tree and out-of-tree plugins based on global settings.
+function(iree_compiler_configure_plugins)
+  # Process in-tree plugins.
+  foreach(_plugin_id ${IREE_COMPILER_PLUGINS})
+    set(_plugin_src_dir "${IREE_SOURCE_DIR}/compiler/plugins/${_plugin_id}")
+    iree_compiler_add_plugin("${_plugin_id}" "${_plugin_src_dir}")
+  endforeach()
+endfunction()
+
+# iree_compiler_add_plugin(src bin)
+# Adds a compiler plugin based on id and source directory.
+function(iree_compiler_add_plugin plugin_id plugin_src_dir)
+  # Add a guard so that we know if we are adding an in-tree plugin.
+  if(IREE_COMPILER_IN_ADD_PLUGIN)
+    message(FATAL_ERROR "Cannot recursively add plugins")
+  endif()
+
+  message(STATUS "Adding static compiler plugin ${plugin_id} (from ${plugin_src_dir})")
+  get_property(_existing_plugin_ids GLOBAL PROPERTY IREE_COMPILER_INCLUDED_PLUGIN_IDS)
+  if("${plugin_id}" IN_LIST _existing_plugin_ids)
+    message(SEND_ERROR "Plugin already registered: ${_plugin_id}")
+    return()
+  endif()
+
+  # Include it.
+  set(IREE_COMPILER_IN_ADD_PLUGIN "${plugin_id}")
+  set_property(GLOBAL APPEND PROPERTY IREE_COMPILER_INCLUDED_PLUGIN_IDS "${plugin_id}")
+  set(_binary_dir "${IREE_BINARY_DIR}/compiler/plugins/${_plugin_id}")
+  add_subdirectory("${plugin_src_dir}" "${_binary_dir}")
+  unset(IREE_COMPILER_IN_ADD_PLUGIN)
+endfunction()

compiler/plugins/example/BUILD.bazel

@@ -0,0 +1,35 @@
+# Copyright 2023 The IREE Authors
+#
+# 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
+
+load("//build_tools/bazel:build_defs.oss.bzl", "iree_compiler_register_plugin")
+
+package(
+    default_visibility = ["//visibility:public"],
+    features = ["layering_check"],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+cc_library(
+    name = "defs",
+    includes = ["src"],
+)
+
+cc_library(
+    name = "registration",
+    srcs = [
+        "src/PluginRegistration.cpp",
+    ],
+    deps = [
+        ":defs",
+        "//compiler/src/iree/compiler/PluginAPI",
+        "@llvm-project//mlir:IR",
+    ],
+)
+
+iree_compiler_register_plugin(
+    plugin_id = "example",
+    target = ":registration",
+)

compiler/plugins/example/CMakeLists.txt

@@ -0,0 +1,41 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from           #
+# compiler/plugins/example/BUILD.bazel                                         #
+#                                                                              #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary   #
+# CMake-only content.                                                          #
+#                                                                              #
+# To disable autogeneration for this file entirely, delete this header.        #
+################################################################################
+
+iree_add_all_subdirs()
+
+iree_cc_library(
+  NAME
+    defs
+  INCLUDES
+    "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/src>"
+    "$<BUILD_INTERFACE:${CMAKE_CURRENT_BINARY_DIR}/src>"
+  PUBLIC
+)
+
+iree_cc_library(
+  NAME
+    registration
+  SRCS
+    "src/PluginRegistration.cpp"
+  DEPS
+    ::defs
+    MLIRIR
+    iree::compiler::PluginAPI
+  PUBLIC
+)
+
+iree_compiler_register_plugin(
+  PLUGIN_ID
+    example
+  TARGET
+    ::registration
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###

compiler/plugins/example/README.md

@@ -0,0 +1,4 @@
+# Example plugin which exercises the plugin API.
+
+TODO: This probably doesn't belong here long term but including it lets me
+build out the hierarchy and get the code layout right.

compiler/plugins/example/src/PluginRegistration.cpp

@@ -0,0 +1,45 @@
+// Copyright 2023 The IREE Authors
+//
+// 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 "iree/compiler/PluginAPI/Client.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/MLIRContext.h"
+
+using namespace mlir;
+using namespace mlir::iree_compiler;
+
+namespace {
+
+struct MyOptions {
+  bool flag = false;
+
+  void bindOptions(OptionsBinder &binder) {
+    static llvm::cl::OptionCategory category("IREE Example Plugin");
+    binder.opt<bool>("iree-example-flag", flag,
+                     llvm::cl::desc("Dummy flag for the example plugin"),
+                     llvm::cl::cat(category));
+  }
+};
+
+struct MySession : public PluginSession<MySession, MyOptions> {
+  LogicalResult onActivate() override {
+    mlir::emitRemark(mlir::UnknownLoc::get(context))
+        << "This remark is from the example plugin activation (flag="
+        << options.flag << ")";
+    return success();
+  }
+};
+
+}  // namespace
+
+IREE_DEFINE_COMPILER_OPTION_FLAGS(MyOptions);
+
+extern "C" bool iree_register_compiler_plugin_example(
+    mlir::iree_compiler::PluginRegistrar *registrar) {
+  registrar->registerPlugin<MySession>("example");
+  return true;
+}

compiler/src/iree/compiler/API/Internal/BUILD.bazel

@@ -28,6 +28,8 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Dialect/VM/Target/Bytecode",
         "//compiler/src/iree/compiler/Dialect/VM/Target/C",
         "//compiler/src/iree/compiler/Pipelines",
+        "//compiler/src/iree/compiler/PluginAPI",
+        "//compiler/src/iree/compiler/PluginAPI:PluginManager",
         "//compiler/src/iree/compiler/Tools:init_llvmir_translations",
         "//compiler/src/iree/compiler/Tools:init_passes_and_dialects",
         "//compiler/src/iree/compiler/Tools:init_targets",

compiler/src/iree/compiler/API/Internal/CMakeLists.txt

@@ -28,6 +28,8 @@ iree_cc_library(
     iree::compiler::Dialect::VM::Target::C
     iree::compiler::Dialect::VM::Target::init_targets
     iree::compiler::Pipelines
+    iree::compiler::PluginAPI
+    iree::compiler::PluginAPI::PluginManager
     iree::compiler::Tools::init_llvmir_translations
     iree::compiler::Tools::init_passes_and_dialects
     iree::compiler::Tools::init_targets