Add simple_io_sample custom compiler plugin. (#12745)
Includes updates to the plugin mechanism:
* Enables building plugins at a specific path (in or out of repo)
* Re-organizes initialization sequence to separate dialect registration
from activation
* Wires into iree-opt
* Adds a mechanism for extending pipelines with additional passes and
wires that into the preprocessing pipeline
Progress on #12520
diff --git a/build_tools/cmake/iree_compiler_plugin.cmake b/build_tools/cmake/iree_compiler_plugin.cmake
index 2a56be9..5d492a8 100644
--- a/build_tools/cmake/iree_compiler_plugin.cmake
+++ b/build_tools/cmake/iree_compiler_plugin.cmake
@@ -5,6 +5,7 @@
# 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")
+set(IREE_COMPILER_PLUGIN_PATHS "" CACHE STRING "Paths to external compiler plugins")
# 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.
@@ -60,6 +61,17 @@
set(_plugin_src_dir "${IREE_SOURCE_DIR}/compiler/plugins/${_plugin_id}")
iree_compiler_add_plugin("${_plugin_id}" "${_plugin_src_dir}")
endforeach()
+ unset(_plugin_id)
+ unset(_plugin_src_dir)
+
+ # Process out of tree plugins.
+ foreach(_plugin_src_dir ${IREE_COMPILER_PLUGIN_PATHS})
+ # TODO: Support some path mangling to allow overriding the plugin id
+ # if it is not literally the last path component.
+ cmake_path(ABSOLUTE_PATH _plugin_src_dir BASE_DIRECTORY "${IREE_SOURCE_DIR}" NORMALIZE)
+ cmake_path(GET _plugin_src_dir FILENAME _plugin_id)
+ iree_compiler_add_plugin("${_plugin_id}" "${_plugin_src_dir}")
+ endforeach()
endfunction()
# iree_compiler_add_plugin(src bin)
@@ -81,6 +93,17 @@
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}")
+
+ # Force enable BUILD_SHARED_LIBS for the compiler if instructed.
+ set(_IREE_ORIG_BUILD_SHARED_LIBS ${BUILD_SHARED_LIBS})
+ if(IREE_COMPILER_BUILD_SHARED_LIBS)
+ set(BUILD_SHARED_LIBS ON CACHE BOOL "" FORCE)
+ endif()
+
add_subdirectory("${plugin_src_dir}" "${_binary_dir}")
+
+ # Reset BUILD_SHARED_LIBS.
+ set(BUILD_SHARED_LIBS ${_IREE_ORIG_BUILD_SHARED_LIBS} CACHE BOOL "" FORCE)
+
unset(IREE_COMPILER_IN_ADD_PLUGIN)
-endfunction()
\ No newline at end of file
+endfunction()
diff --git a/compiler/src/iree/compiler/API/Internal/BUILD.bazel b/compiler/src/iree/compiler/API/Internal/BUILD.bazel
index 8a2addf..bf8294f 100644
--- a/compiler/src/iree/compiler/API/Internal/BUILD.bazel
+++ b/compiler/src/iree/compiler/API/Internal/BUILD.bazel
@@ -74,11 +74,14 @@
],
deps = [
"//compiler/bindings/c:headers",
+ "//compiler/src/iree/compiler/PluginAPI:PluginManager",
"//compiler/src/iree/compiler/Tools:init_passes_and_dialects",
"//compiler/src/iree/compiler/Tools:init_targets",
"@llvm-project//llvm:Support",
+ "@llvm-project//mlir:Debug",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:MlirOptLib",
+ "@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
],
)
diff --git a/compiler/src/iree/compiler/API/Internal/CMakeLists.txt b/compiler/src/iree/compiler/API/Internal/CMakeLists.txt
index 80322a8..a9e3130 100644
--- a/compiler/src/iree/compiler/API/Internal/CMakeLists.txt
+++ b/compiler/src/iree/compiler/API/Internal/CMakeLists.txt
@@ -71,9 +71,12 @@
"IREEOptToolEntryPoint.cpp"
DEPS
LLVMSupport
+ MLIRDebug
MLIRIR
MLIROptLib
+ MLIRPass
MLIRSupport
+ iree::compiler::PluginAPI::PluginManager
iree::compiler::Tools::init_passes_and_dialects
iree::compiler::Tools::init_targets
iree::compiler::bindings::c::headers
diff --git a/compiler/src/iree/compiler/API/Internal/Embed.cpp b/compiler/src/iree/compiler/API/Internal/Embed.cpp
index 3a43fea..aabd131 100644
--- a/compiler/src/iree/compiler/API/Internal/Embed.cpp
+++ b/compiler/src/iree/compiler/API/Internal/Embed.cpp
@@ -138,7 +138,7 @@
}
pluginManager.globalInitialize();
pluginManager.registerPasses();
- pluginManager.registerDialects(registry);
+ pluginManager.registerGlobalDialects(registry);
}
void GlobalInit::registerCommandLineOptions() {
@@ -198,7 +198,14 @@
LogicalResult activatePluginsOnce() {
if (!pluginsActivated) {
pluginsActivated = true;
- pluginActivationStatus = pluginSession.activatePlugins(&context);
+ if (failed(pluginSession.initializePlugins())) {
+ pluginActivationStatus = failure();
+ } else {
+ DialectRegistry registry;
+ pluginSession.registerDialects(registry);
+ context.appendDialectRegistry(registry);
+ pluginActivationStatus = pluginSession.activatePlugins(&context);
+ }
}
return pluginActivationStatus;
}
@@ -503,17 +510,9 @@
Error *outputVMCSource(Output &output);
Error *outputHALExecutable(Output &output);
- IREEVMPipelineHooks &getHooks() {
- static IREEVMPipelineHooks hooks = {
- // buildConstEvalPassPipelineCallback =
- [](OpPassManager &pm) {
- pm.addPass(ConstEval::createJitGlobalsPass());
- }};
- return hooks;
- }
-
Session &session;
PassManager passManager;
+ IREEVMPipelineHooks pipelineHooks;
// Diagnostic handlers are instantiated upon parsing the source (when we
// have the SrcMgr) and held for the duration of the invocation. Each will
@@ -543,6 +542,16 @@
mlir::applyDefaultTimingPassManagerCLOptions(passManager);
}
passManager.addInstrumentation(std::make_unique<PassTracing>());
+
+ // Since the jitter invokes much of the top-level compiler recursively,
+ // it must be injected at the top-level here vs in the pass pipeline
+ // (or else the circular dependency cannot be resolved).
+ pipelineHooks.buildConstEvalPassPipelineCallback = [](OpPassManager &pm) {
+ pm.addPass(ConstEval::createJitGlobalsPass());
+ };
+ // The PluginSession implements PipelineExtensions and delegates it to
+ // activated plugins.
+ pipelineHooks.pipelineExtensions = &session.pluginSession;
}
bool Invocation::parseSource(Source &source) {
@@ -610,7 +619,7 @@
session.bindingOptions, session.inputOptions,
session.preprocessingOptions, session.highLevelOptimizationOptions,
session.schedulingOptions, session.halTargetOptions,
- session.vmTargetOptions, getHooks(), passManager, *compileToPhase);
+ session.vmTargetOptions, pipelineHooks, passManager, *compileToPhase);
break;
}
case IREE_COMPILER_PIPELINE_HAL_EXECUTABLE: {
diff --git a/compiler/src/iree/compiler/API/Internal/IREEOptToolEntryPoint.cpp b/compiler/src/iree/compiler/API/Internal/IREEOptToolEntryPoint.cpp
index adf0a35..6fcd2cf 100644
--- a/compiler/src/iree/compiler/API/Internal/IREEOptToolEntryPoint.cpp
+++ b/compiler/src/iree/compiler/API/Internal/IREEOptToolEntryPoint.cpp
@@ -8,20 +8,114 @@
//
// Based on mlir-opt but registers the passes and dialects we care about.
+#include "iree/compiler/PluginAPI/PluginManager.h"
#include "iree/compiler/Tools/init_dialects.h"
#include "iree/compiler/Tools/init_passes.h"
#include "iree/compiler/Tools/init_targets.h"
#include "iree/compiler/tool_entry_points_api.h"
#include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/Process.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/ToolOutputFile.h"
+#include "mlir/Debug/Counter.h"
+#include "mlir/IR/AsmState.h"
#include "mlir/IR/Dialect.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Support/FileUtilities.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Tools/mlir-opt/MlirOptMain.h"
+using namespace llvm;
+using namespace mlir;
+
+static LogicalResult ireeOptMainFromCL(int argc, char **argv,
+ llvm::StringRef toolName,
+ DialectRegistry ®istry) {
+ static cl::opt<std::string> inputFilename(
+ cl::Positional, cl::desc("<input file>"), cl::init("-"));
+
+ static cl::opt<std::string> outputFilename("o", cl::desc("Output filename"),
+ cl::value_desc("filename"),
+ cl::init("-"));
+
+ InitLLVM y(argc, argv);
+
+ // Register any command line options.
+ MlirOptMainConfig::registerCLOptions();
+ registerAsmPrinterCLOptions();
+ registerMLIRContextCLOptions();
+ registerPassManagerCLOptions();
+ registerDefaultTimingManagerCLOptions();
+ tracing::DebugCounter::registerCLOptions();
+ auto &pluginManagerOptions =
+ mlir::iree_compiler::PluginManagerOptions::FromFlags::get();
+
+ // Build the list of dialects as a header for the --help message.
+ std::string helpHeader = (toolName + "\nAvailable Dialects: ").str();
+ {
+ llvm::raw_string_ostream os(helpHeader);
+ interleaveComma(registry.getDialectNames(), os,
+ [&](auto name) { os << name; });
+ }
+
+ // We support a limited form of the PluginManager, allowing it to perform
+ // global initialization and dialect registration.
+ mlir::iree_compiler::PluginManager pluginManager;
+ if (!pluginManager.loadAvailablePlugins()) {
+ llvm::errs() << "error: Failed to initialize IREE compiler plugins\n";
+ return failure();
+ }
+ pluginManager.initializeCLI();
+
+ // Parse pass names in main to ensure static initialization completed.
+ cl::ParseCommandLineOptions(argc, argv, helpHeader);
+ MlirOptMainConfig config = MlirOptMainConfig::createFromCLOptions();
+
+ // The local binder is meant for overriding session-level options, but for
+ // tools like this it is unused.
+ auto localBinder = mlir::iree_compiler::OptionsBinder::local();
+ pluginManager.globalInitialize();
+ pluginManager.registerPasses();
+ pluginManager.registerGlobalDialects(registry);
+ mlir::iree_compiler::PluginManagerSession pluginSession(
+ pluginManager, localBinder, pluginManagerOptions);
+ if (failed(pluginSession.initializePlugins())) return failure();
+ pluginSession.registerDialects(registry);
+
+ // When reading from stdin and the input is a tty, it is often a user mistake
+ // and the process "appears to be stuck". Print a message to let the user know
+ // about it!
+ if (inputFilename == "-" &&
+ sys::Process::FileDescriptorIsDisplayed(fileno(stdin)))
+ llvm::errs() << "(processing input from stdin now, hit ctrl-c/ctrl-d to "
+ "interrupt)\n";
+
+ // Set up the input file.
+ std::string errorMessage;
+ auto file = openInputFile(inputFilename, &errorMessage);
+ if (!file) {
+ llvm::errs() << errorMessage << "\n";
+ return failure();
+ }
+
+ auto output = openOutputFile(outputFilename, &errorMessage);
+ if (!output) {
+ llvm::errs() << errorMessage << "\n";
+ return failure();
+ }
+ if (failed(MlirOptMain(output->os(), std::move(file), registry, config)))
+ return failure();
+
+ // Keep the output file if the invocation of MlirOptMain was successful.
+ output->keep();
+ return success();
+}
+
int ireeOptRunMain(int argc, char **argv) {
llvm::setBugReportMsg(
"Please report issues to https://github.com/openxla/iree/issues and "
"include the crash backtrace.\n");
- llvm::InitLLVM y(argc, argv);
mlir::DialectRegistry registry;
mlir::iree_compiler::registerAllDialects(registry);
@@ -32,9 +126,8 @@
// TODO: this should be upstreamed.
mlir::linalg::transform::registerDropSchedulePass();
- if (failed(MlirOptMain(argc, argv, "IREE modular optimizer driver\n",
- registry,
- /*preloadDialectsInContext=*/false))) {
+ if (failed(ireeOptMainFromCL(argc, argv, "IREE modular optimizer driver\n",
+ registry))) {
return 1;
}
return 0;
diff --git a/compiler/src/iree/compiler/Pipelines/Pipelines.cpp b/compiler/src/iree/compiler/Pipelines/Pipelines.cpp
index cb1c611..e7f0736 100644
--- a/compiler/src/iree/compiler/Pipelines/Pipelines.cpp
+++ b/compiler/src/iree/compiler/Pipelines/Pipelines.cpp
@@ -127,7 +127,8 @@
break;
default:
IREE_TRACE_ADD_BEGIN_FRAME_PASS(passManager, "Preprocessing");
- IREE::buildPreprocessingPassPipeline(passManager, preprocessingOptions);
+ IREE::buildPreprocessingPassPipeline(passManager, preprocessingOptions,
+ hooks.pipelineExtensions);
IREE_TRACE_ADD_END_FRAME_PASS(passManager, "Preprocessing");
if (compileTo == IREEVMPipelinePhase::Preprocessing)
return; // early-exit
diff --git a/compiler/src/iree/compiler/Pipelines/Pipelines.h b/compiler/src/iree/compiler/Pipelines/Pipelines.h
index c135f47..acf7776 100644
--- a/compiler/src/iree/compiler/Pipelines/Pipelines.h
+++ b/compiler/src/iree/compiler/Pipelines/Pipelines.h
@@ -16,6 +16,8 @@
namespace mlir {
namespace iree_compiler {
+class PipelineExtensions;
+
// Hooks for injecting behavior into the IREEVM pipeline. Since these are not
// derived from CLI options, we maintain them as a separate struct.
struct IREEVMPipelineHooks {
@@ -26,6 +28,9 @@
// the constant evaluator, which needs to recursively invoke these
// pipelines.
std::function<void(OpPassManager &)> buildConstEvalPassPipelineCallback;
+
+ // Applies pipeline extensions to the built pipeline if not nullptr.
+ PipelineExtensions *pipelineExtensions = nullptr;
};
enum class IREEVMPipelinePhase {
diff --git a/compiler/src/iree/compiler/PluginAPI/Client.cpp b/compiler/src/iree/compiler/PluginAPI/Client.cpp
index 1162a4b..341fac8 100644
--- a/compiler/src/iree/compiler/PluginAPI/Client.cpp
+++ b/compiler/src/iree/compiler/PluginAPI/Client.cpp
@@ -14,6 +14,8 @@
namespace mlir::iree_compiler {
+PipelineExtensions::~PipelineExtensions() = default;
+
AbstractPluginRegistration::~AbstractPluginRegistration() = default;
AbstractPluginSession::~AbstractPluginSession() = default;
@@ -31,7 +33,9 @@
void PluginRegistrar::registerPlugin(
std::unique_ptr<AbstractPluginRegistration> registration) {
- std::string_view id = registration->getPluginId();
+ // Need to copy the id since in the error case, the registration will be
+ // deleted before reporting the error message.
+ std::string id = std::string(registration->getPluginId());
auto foundIt = registrations.insert(
std::make_pair(llvm::StringRef(id), std::move(registration)));
if (!foundIt.second) {
diff --git a/compiler/src/iree/compiler/PluginAPI/Client.h b/compiler/src/iree/compiler/PluginAPI/Client.h
index 94ede00..b22f473 100644
--- a/compiler/src/iree/compiler/PluginAPI/Client.h
+++ b/compiler/src/iree/compiler/PluginAPI/Client.h
@@ -4,6 +4,9 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#ifndef IREE_COMPILER_PLUGINAPI_CLIENT_H_
+#define IREE_COMPILER_PLUGINAPI_CLIENT_H_
+
#include <optional>
#include <string_view>
@@ -13,6 +16,7 @@
namespace mlir {
class DialectRegistry;
class MLIRContext;
+class OpPassManager;
} // namespace mlir
namespace mlir::iree_compiler {
@@ -29,6 +33,16 @@
static void bindOptions(OptionsBinder &binder) {}
};
+// Entrypoints for extending IREE's pass pipelines at various stages.
+// Override what is needed.
+class PipelineExtensions {
+ public:
+ virtual ~PipelineExtensions();
+
+ // Adds passes to the |buildPreprocessingPassPipeline| pipeline at the end.
+ virtual void extendPreprocessingPassPipeline(OpPassManager &passManager) {}
+};
+
// Abstract class representing a plugin registration. It is responsible for
// various global initialization and creation of plugin sessions that mirror
// the lifetime of and |iree_compiler_session_t| for when the plugin is
@@ -72,7 +86,7 @@
// change behavior. It is safer to customize the context on a per-session
// basis in a plugin session's activate() method (i.e. if registering
// interfaces or behavior changes extensions).
- virtual void registerDialects(DialectRegistry ®istry) {}
+ virtual void registerGlobalDialects(DialectRegistry ®istry) {}
// Creates an uninitialized session. If the CLI was initialized, then this
// should also ensure that any command line options were managed properly into
@@ -96,15 +110,24 @@
// Most users will inherit from this class via the PluginSession CRTP helper,
// which adds some niceties and support for global command line option
// registration.
-class AbstractPluginSession {
+class AbstractPluginSession : public PipelineExtensions {
public:
virtual ~AbstractPluginSession();
+ // Called prior to context initialization in order to register dialects.
+ void registerDialects(DialectRegistry ®istry) {
+ onRegisterDialects(registry);
+ }
+
// Called after the session has been fully constructed. If it fails, then
// it should emit an appropriate diagnostic.
LogicalResult activate(MLIRContext *context);
protected:
+ // Called from registerDialects() prior to initializing the context and
+ // prior to onActivate().
+ virtual void onRegisterDialects(DialectRegistry ®istry) {}
+
// Called from the activate() method once pre-conditions are verified and the
// context is set.
virtual LogicalResult onActivate() { return success(); };
@@ -122,7 +145,7 @@
// AbstractPluginRegistration.
static void globalInitialize() {}
static void registerPasses() {}
- static void registerDialects(DialectRegistry ®istry) {}
+ static void registerGlobalDialects(DialectRegistry ®istry) {}
struct Registration : public AbstractPluginRegistration {
using AbstractPluginRegistration::AbstractPluginRegistration;
@@ -135,9 +158,9 @@
// Actually need to capture the reference, not a copy. So get a pointer.
globalCLIOptions = &OptionsFromFlags<OptionsTy>::get();
}
- void registerDialects(DialectRegistry ®istry) override {
+ void registerGlobalDialects(DialectRegistry ®istry) override {
// Forward to the CRTP derived type.
- DerivedTy::registerDialects(registry);
+ DerivedTy::registerGlobalDialects(registry);
}
std::unique_ptr<AbstractPluginSession> createUninitializedSession(
OptionsBinder &localOptionsBinder) override {
@@ -181,3 +204,5 @@
};
} // namespace mlir::iree_compiler
+
+#endif // IREE_COMPILER_PLUGINAPI_CLIENT_H_
diff --git a/compiler/src/iree/compiler/PluginAPI/PluginManager.cpp b/compiler/src/iree/compiler/PluginAPI/PluginManager.cpp
index a920181..3c9fda9 100644
--- a/compiler/src/iree/compiler/PluginAPI/PluginManager.cpp
+++ b/compiler/src/iree/compiler/PluginAPI/PluginManager.cpp
@@ -66,9 +66,9 @@
}
}
-void PluginManager::registerDialects(DialectRegistry ®istry) {
+void PluginManager::registerGlobalDialects(DialectRegistry ®istry) {
for (auto &kv : registrations) {
- kv.second->registerDialects(registry);
+ kv.second->registerGlobalDialects(registry);
}
}
@@ -82,7 +82,7 @@
}
}
-LogicalResult PluginManagerSession::activatePlugins(MLIRContext *context) {
+LogicalResult PluginManagerSession::initializePlugins() {
auto getAvailableIds = [&]() -> llvm::SmallVector<llvm::StringRef> {
llvm::SmallVector<llvm::StringRef> availableIds;
for (auto &kv : allPluginSessions) {
@@ -105,25 +105,36 @@
// sorting accordingly. For now, what you say is what you get.
for (auto &pluginId : options.plugins) {
if (options.printPluginInfo) {
- llvm::errs() << "[IREE plugins]: Activating plugin '" << pluginId
+ llvm::errs() << "[IREE plugins]: Initializing plugin '" << pluginId
<< "'\n";
}
auto foundIt = allPluginSessions.find(pluginId);
if (foundIt == allPluginSessions.end()) {
- auto diag = mlir::emitError(mlir::UnknownLoc::get(context))
- << "could not activate requested IREE plugin '" << pluginId
- << "' because it is not registered (available plugins: ";
- llvm::interleaveComma(getAvailableIds(), diag);
- diag << ")";
+ llvm::errs()
+ << "[IREE plugins error]: could not activate requested IREE plugin '"
+ << pluginId << "' because it is not registered (available plugins: ";
+ llvm::interleaveComma(getAvailableIds(), llvm::errs());
+ llvm::errs() << ")\n";
return failure();
}
- AbstractPluginSession *instance = foundIt->second.get();
- if (failed(instance->activate(context))) return failure();
- activatedSessions.push_back(instance);
+ initializedSessions.push_back(foundIt->second.get());
}
return success();
}
+void PluginManagerSession::registerDialects(DialectRegistry ®istry) {
+ for (auto *s : initializedSessions) {
+ s->registerDialects(registry);
+ }
+}
+
+LogicalResult PluginManagerSession::activatePlugins(MLIRContext *context) {
+ for (auto *s : initializedSessions) {
+ if (failed(s->activate(context))) return failure();
+ }
+ return success();
+}
+
} // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/PluginAPI/PluginManager.h b/compiler/src/iree/compiler/PluginAPI/PluginManager.h
index b39a802..c76e5f7 100644
--- a/compiler/src/iree/compiler/PluginAPI/PluginManager.h
+++ b/compiler/src/iree/compiler/PluginAPI/PluginManager.h
@@ -4,6 +4,9 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#ifndef IREE_COMPILER_PLUGINAPI_PLUGINMANAGER_H_
+#define IREE_COMPILER_PLUGINAPI_PLUGINMANAGER_H_
+
#include <optional>
#include <string_view>
#include <vector>
@@ -66,29 +69,44 @@
// Calls through to AbstractPluginRegistration::registerDialects for all
// available plugins.
- void registerDialects(DialectRegistry ®istry);
+ void registerGlobalDialects(DialectRegistry ®istry);
private:
friend class PluginManagerSession;
};
// Holds activated plugins for an |iree_compiler_session_t|.
-class PluginManagerSession {
+class PluginManagerSession : public PipelineExtensions {
public:
PluginManagerSession(PluginManager &pluginManager, OptionsBinder &binder,
PluginManagerOptions &options);
+ // Initializes all plugins that should be activated by default.
+ LogicalResult initializePlugins();
+
+ // Invokes registerDialects() on all initialized plugins.
+ void registerDialects(DialectRegistry ®istry);
+
// Activates plugins as configured.
LogicalResult activatePlugins(MLIRContext *context);
+ // Forward pipeline extensions.
+ void extendPreprocessingPassPipeline(OpPassManager &passManager) override {
+ for (auto *s : initializedSessions) {
+ s->extendPreprocessingPassPipeline(passManager);
+ }
+ }
+
private:
PluginManagerOptions &options;
// At construction, uninitialized plugin sessions are created for all
// registered plugins so that CLI options can be set properly.
llvm::StringMap<std::unique_ptr<AbstractPluginSession>> allPluginSessions;
- // Activation state.
- llvm::SmallVector<AbstractPluginSession *> activatedSessions;
+ // Initialized list of plugins.
+ llvm::SmallVector<AbstractPluginSession *> initializedSessions;
};
} // namespace mlir::iree_compiler
+
+#endif // IREE_COMPILER_PLUGINAPI_PLUGINMANAGER_H_
diff --git a/compiler/src/iree/compiler/Preprocessing/BUILD.bazel b/compiler/src/iree/compiler/Preprocessing/BUILD.bazel
index 11993e9..13d3feb 100644
--- a/compiler/src/iree/compiler/Preprocessing/BUILD.bazel
+++ b/compiler/src/iree/compiler/Preprocessing/BUILD.bazel
@@ -22,6 +22,7 @@
],
deps = [
"//compiler/src/iree/compiler/Pipelines:Options",
+ "//compiler/src/iree/compiler/PluginAPI",
"//compiler/src/iree/compiler/Preprocessing/Common:Transforms",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
diff --git a/compiler/src/iree/compiler/Preprocessing/CMakeLists.txt b/compiler/src/iree/compiler/Preprocessing/CMakeLists.txt
index 67a6ed2..dcef2c2 100644
--- a/compiler/src/iree/compiler/Preprocessing/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Preprocessing/CMakeLists.txt
@@ -22,6 +22,7 @@
MLIRArithDialect
MLIRPass
iree::compiler::Pipelines::Options
+ iree::compiler::PluginAPI
iree::compiler::Preprocessing::Common::Transforms
PUBLIC
)
diff --git a/compiler/src/iree/compiler/Preprocessing/Passes.cpp b/compiler/src/iree/compiler/Preprocessing/Passes.cpp
index f30688e..27e0132 100644
--- a/compiler/src/iree/compiler/Preprocessing/Passes.cpp
+++ b/compiler/src/iree/compiler/Preprocessing/Passes.cpp
@@ -15,41 +15,37 @@
namespace iree_compiler {
namespace IREE {
-void buildPreprocessingPassPipeline(
- OpPassManager &passManager,
- const PreprocessingOptions &preprocessingOptions) {
- auto pipelineStr = preprocessingOptions.preprocessingPassPipeline;
- if (pipelineStr.empty()) {
- return;
- }
+namespace {
+void extendWithTextPipeline(OpPassManager &passManager,
+ StringRef textPipeline) {
+ StringRef orig = textPipeline;
// Strip the `builtin.module(...)` that surrounds the pass pipeline
// description. On failure an assertion is triggered, but in release builds
// it just will silently return and not raise an error. There is no
// way to handle the error in caller currently.
- StringRef text(pipelineStr);
- size_t pos = text.find_first_of("(");
+ size_t pos = textPipeline.find_first_of("(");
if (pos == StringRef::npos) {
llvm::errs() << "ERROR: expected preprocessing pass pipeline string to be "
"nested within `builtin.module(..)`; got `"
- << pipelineStr << "`\n";
+ << orig << "`\n";
return;
}
- if (text.substr(0, pos) != "builtin.module") {
+ if (textPipeline.substr(0, pos) != "builtin.module") {
llvm::errs() << "ERROR: expected preprocessing pass pipeline string to be "
"nested within `builtin.module(..)`; got `"
- << pipelineStr << "`\n";
+ << orig << "`\n";
return;
}
- if (text.back() != ')') {
- llvm::errs() << "ERROR: mismatched parenthesis in pass pipeline `"
- << pipelineStr << "`\n";
+ if (textPipeline.back() != ')') {
+ llvm::errs() << "ERROR: mismatched parenthesis in pass pipeline `" << orig
+ << "`\n";
return;
}
- text = text.substr(pos + 1);
- if (failed(parsePassPipeline(text.drop_back(), passManager))) {
- llvm::errs() << "ERROR: mismatched parenthesis in pass pipeline `"
- << pipelineStr << "`\n";
+ textPipeline = textPipeline.substr(pos + 1);
+ if (failed(parsePassPipeline(textPipeline.drop_back(), passManager))) {
+ llvm::errs() << "ERROR: mismatched parenthesis in pass pipeline `" << orig
+ << "`\n";
return;
}
LLVM_DEBUG({
@@ -58,6 +54,23 @@
});
}
+} // namespace
+
+void buildPreprocessingPassPipeline(
+ OpPassManager &passManager,
+ const PreprocessingOptions &preprocessingOptions,
+ PipelineExtensions *pipelineExtensions) {
+ auto pipelineStr = preprocessingOptions.preprocessingPassPipeline;
+ if (!preprocessingOptions.preprocessingPassPipeline.empty()) {
+ extendWithTextPipeline(passManager,
+ preprocessingOptions.preprocessingPassPipeline);
+ }
+
+ if (pipelineExtensions) {
+ pipelineExtensions->extendPreprocessingPassPipeline(passManager);
+ }
+}
+
void registerPreprocessingPasses() { registerCommonPreprocessingPasses(); }
} // namespace IREE
diff --git a/compiler/src/iree/compiler/Preprocessing/Passes.h b/compiler/src/iree/compiler/Preprocessing/Passes.h
index 0eae95c..7768e82 100644
--- a/compiler/src/iree/compiler/Preprocessing/Passes.h
+++ b/compiler/src/iree/compiler/Preprocessing/Passes.h
@@ -10,6 +10,7 @@
#include <functional>
#include "iree/compiler/Pipelines/Options.h"
+#include "iree/compiler/PluginAPI/Client.h"
#include "mlir/Pass/PassManager.h"
namespace mlir {
@@ -22,8 +23,9 @@
/// on the sequence of preprocessing passes to run after conversion from input
/// dialects like `mhlo`/`tosa` before running the core IREE compilation
/// pipelines (starting with the flow pipeline).
-void buildPreprocessingPassPipeline(OpPassManager &passManager,
- const PreprocessingOptions &options);
+void buildPreprocessingPassPipeline(
+ OpPassManager &passManager, const PreprocessingOptions &options,
+ PipelineExtensions *pipelineExtensions = nullptr);
void registerPreprocessingPasses();
@@ -31,4 +33,4 @@
} // namespace iree_compiler
} // namespace mlir
-#endif // IREE_COMPILER_PREPROCESSING_PASSES_H_
\ No newline at end of file
+#endif // IREE_COMPILER_PREPROCESSING_PASSES_H_
diff --git a/samples/CMakeLists.txt b/samples/CMakeLists.txt
index 1e70c9d..057c48d 100644
--- a/samples/CMakeLists.txt
+++ b/samples/CMakeLists.txt
@@ -4,4 +4,19 @@
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-iree_add_all_subdirs()
+# Not CMake directories:
+# colab
+# models
+# vision_inference
+# Samples that are actually external projects do not get directly
+# included:
+# compiler_plugins
+
+add_subdirectory(custom_dispatch)
+add_subdirectory(custom_module)
+add_subdirectory(dynamic_shapes)
+add_subdirectory(emitc_modules)
+add_subdirectory(py_custom_module)
+add_subdirectory(simple_embedding)
+add_subdirectory(static_library)
+add_subdirectory(variables_and_state)
diff --git a/samples/compiler_plugins/simple_io_sample/BUILD.bazel b/samples/compiler_plugins/simple_io_sample/BUILD.bazel
new file mode 100644
index 0000000..32556d8
--- /dev/null
+++ b/samples/compiler_plugins/simple_io_sample/BUILD.bazel
@@ -0,0 +1,152 @@
+# 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", "iree_gentbl_cc_library", "iree_tablegen_doc", "iree_td_library")
+
+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 = [
+ ":IR",
+ ":Transforms",
+ ":defs",
+ "//compiler/src/iree/compiler/PluginAPI",
+ "@llvm-project//mlir:IR",
+ "@llvm-project//mlir:Pass",
+ ],
+)
+
+iree_compiler_register_plugin(
+ plugin_id = "simple_io_sample",
+ target = ":registration",
+)
+
+iree_td_library(
+ name = "td_files",
+ srcs = [
+ "src/simple_io_sample/IR/SimpleIOOps.td",
+ "src/simple_io_sample/Transforms/Passes.td",
+ ],
+ deps = [
+ "@llvm-project//mlir:ControlFlowInterfacesTdFiles",
+ "@llvm-project//mlir:FuncTdFiles",
+ "@llvm-project//mlir:InferTypeOpInterfaceTdFiles",
+ "@llvm-project//mlir:OpBaseTdFiles",
+ "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+ "@llvm-project//mlir:ViewLikeInterfaceTdFiles",
+ ],
+)
+
+cc_library(
+ name = "IR",
+ srcs = [
+ "src/simple_io_sample/IR/SimpleIODialect.cpp",
+ "src/simple_io_sample/IR/SimpleIOOps.cpp",
+ "src/simple_io_sample/IR/SimpleIOOps.cpp.inc",
+ ],
+ hdrs = [
+ "src/simple_io_sample/IR/SimpleIODialect.h",
+ "src/simple_io_sample/IR/SimpleIOOps.h",
+ "src/simple_io_sample/IR/SimpleIOOps.h.inc",
+ ],
+ deps = [
+ ":SimpleIOOpsGen",
+ ":defs",
+ "@llvm-project//llvm:Support",
+ "@llvm-project//mlir:FuncDialect",
+ "@llvm-project//mlir:IR",
+ "@llvm-project//mlir:Support",
+ ],
+)
+
+iree_gentbl_cc_library(
+ name = "SimpleIOOpsGen",
+ tbl_outs = [
+ (
+ ["--gen-dialect-decls"],
+ "src/simple_io_sample/IR/SimpleIODialect.h.inc",
+ ),
+ (
+ ["--gen-dialect-defs"],
+ "src/simple_io_sample/IR/SimpleIODialect.cpp.inc",
+ ),
+ (
+ ["--gen-op-decls"],
+ "src/simple_io_sample/IR/SimpleIOOps.h.inc",
+ ),
+ (
+ ["--gen-op-defs"],
+ "src/simple_io_sample/IR/SimpleIOOps.cpp.inc",
+ ),
+ ],
+ tblgen = "@llvm-project//mlir:mlir-tblgen",
+ td_file = "src/simple_io_sample/IR/SimpleIOOps.td",
+ deps = [":td_files"],
+)
+
+cc_library(
+ name = "Transforms",
+ srcs = [
+ "src/simple_io_sample/Transforms/LegalizeSimpleIO.cpp",
+ ],
+ hdrs = [
+ "src/simple_io_sample/Transforms/Passes.h",
+ "src/simple_io_sample/Transforms/Passes.h.inc",
+ ],
+ deps = [
+ ":IR",
+ ":PassesIncGen",
+ ":defs",
+ "@llvm-project//mlir:FuncDialect",
+ "@llvm-project//mlir:IR",
+ "@llvm-project//mlir:Pass",
+ ],
+)
+
+iree_gentbl_cc_library(
+ name = "PassesIncGen",
+ tbl_outs = [
+ (
+ ["--gen-pass-decls"],
+ "src/simple_io_sample/Transforms/Passes.h.inc",
+ ),
+ ],
+ tblgen = "@llvm-project//mlir:mlir-tblgen",
+ td_file = "src/simple_io_sample/Transforms/Passes.td",
+ deps = [
+ ":td_files",
+ "@llvm-project//mlir:PassBaseTdFiles",
+ ],
+)
+
+iree_tablegen_doc(
+ name = "SimpleIODialectDocGen",
+ tbl_outs = [
+ (
+ [
+ "--gen-dialect-doc",
+ "-dialect=simple_io",
+ ],
+ "src/simple_io_sample/IR/SimpleIODialect.md",
+ ),
+ ],
+ tblgen = "@llvm-project//mlir:mlir-tblgen",
+ td_file = "src/simple_io_sample/IR/SimpleIOOps.td",
+ deps = [":td_files"],
+)
diff --git a/samples/compiler_plugins/simple_io_sample/CMakeLists.txt b/samples/compiler_plugins/simple_io_sample/CMakeLists.txt
new file mode 100644
index 0000000..12ba8a4
--- /dev/null
+++ b/samples/compiler_plugins/simple_io_sample/CMakeLists.txt
@@ -0,0 +1,113 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from #
+# samples/compiler_plugins/simple_io_sample/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
+ ::IR
+ ::Transforms
+ ::defs
+ MLIRIR
+ MLIRPass
+ iree::compiler::PluginAPI
+ PUBLIC
+)
+
+iree_compiler_register_plugin(
+ PLUGIN_ID
+ simple_io_sample
+ TARGET
+ ::registration
+)
+
+iree_cc_library(
+ NAME
+ IR
+ HDRS
+ "src/simple_io_sample/IR/SimpleIODialect.h"
+ "src/simple_io_sample/IR/SimpleIOOps.h"
+ "src/simple_io_sample/IR/SimpleIOOps.h.inc"
+ SRCS
+ "src/simple_io_sample/IR/SimpleIODialect.cpp"
+ "src/simple_io_sample/IR/SimpleIOOps.cpp"
+ "src/simple_io_sample/IR/SimpleIOOps.cpp.inc"
+ DEPS
+ ::SimpleIOOpsGen
+ ::defs
+ LLVMSupport
+ MLIRFuncDialect
+ MLIRIR
+ MLIRSupport
+ PUBLIC
+)
+
+iree_tablegen_library(
+ NAME
+ SimpleIOOpsGen
+ TD_FILE
+ "src/simple_io_sample/IR/SimpleIOOps.td"
+ OUTS
+ --gen-dialect-decls src/simple_io_sample/IR/SimpleIODialect.h.inc
+ --gen-dialect-defs src/simple_io_sample/IR/SimpleIODialect.cpp.inc
+ --gen-op-decls src/simple_io_sample/IR/SimpleIOOps.h.inc
+ --gen-op-defs src/simple_io_sample/IR/SimpleIOOps.cpp.inc
+)
+
+iree_cc_library(
+ NAME
+ Transforms
+ HDRS
+ "src/simple_io_sample/Transforms/Passes.h"
+ "src/simple_io_sample/Transforms/Passes.h.inc"
+ SRCS
+ "src/simple_io_sample/Transforms/LegalizeSimpleIO.cpp"
+ DEPS
+ ::IR
+ ::PassesIncGen
+ ::defs
+ MLIRFuncDialect
+ MLIRIR
+ MLIRPass
+ PUBLIC
+)
+
+iree_tablegen_library(
+ NAME
+ PassesIncGen
+ TD_FILE
+ "src/simple_io_sample/Transforms/Passes.td"
+ OUTS
+ --gen-pass-decls src/simple_io_sample/Transforms/Passes.h.inc
+)
+
+iree_tablegen_doc(
+ NAME
+ SimpleIODialectDocGen
+ TD_FILE
+ "src/simple_io_sample/IR/SimpleIOOps.td"
+ OUTS
+ --gen-dialect-doc -dialect=simple_io src/simple_io_sample/IR/SimpleIODialect.md
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/samples/compiler_plugins/simple_io_sample/README.md b/samples/compiler_plugins/simple_io_sample/README.md
new file mode 100644
index 0000000..d123902
--- /dev/null
+++ b/samples/compiler_plugins/simple_io_sample/README.md
@@ -0,0 +1,37 @@
+# SimpleIO Compiler Plugin Sample
+
+WARNING: This sample is under construction.
+
+This sample demonstrates a compiler plugin which:
+
+* Adds a new dialect to IREE
+* Implements pre-processor lowerings to transform ops to internal
+ implementations (TODO)
+* Has a python-based runner that implements the IO ops in pure python (TODO)
+* Illustrates some advanced features of the way such things can be
+ constructed (custom types, async, etc) (TODO)
+* Show how to test such a plugin (TODO)
+
+To use this, the plugin must be built into the compiler via:
+
+```
+-DIREE_COMPILER_PLUGIN_PATHS=samples/compiler_plugins/simple_io_sample
+```
+
+It can then be activated in either `iree-opt` or `iree-compile` via the
+option `--iree-plugin=simple_io_sample`.
+
+To compile a sample:
+
+```
+iree-compile --iree-plugin=simple_io_sample test/print.mlir -o /tmp/print.vmfb
+python run_mock.py /tmp/print.vmfb
+```
+
+Should print:
+
+```
+--- Loading /tmp/print.vmfb
+--- Running main()
++++ HELLO FROM SIMPLE_IO
+```
diff --git a/samples/compiler_plugins/simple_io_sample/src/PluginRegistration.cpp b/samples/compiler_plugins/simple_io_sample/src/PluginRegistration.cpp
new file mode 100644
index 0000000..5ddd8b7
--- /dev/null
+++ b/samples/compiler_plugins/simple_io_sample/src/PluginRegistration.cpp
@@ -0,0 +1,55 @@
+// 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"
+#include "mlir/Pass/Pass.h"
+#include "simple_io_sample/IR/SimpleIODialect.h"
+#include "simple_io_sample/Transforms/Passes.h"
+
+using namespace mlir;
+using namespace mlir::iree_compiler;
+
+namespace detail {
+namespace {
+
+#define GEN_PASS_REGISTRATION
+#include "simple_io_sample/Transforms/Passes.h.inc"
+
+} // namespace
+} // namespace detail
+
+namespace {
+
+struct MyOptions {
+ void bindOptions(OptionsBinder &binder) {}
+};
+
+struct MySession : public PluginSession<MySession, MyOptions> {
+ static void registerPasses() { ::detail::registerPasses(); }
+
+ void onRegisterDialects(DialectRegistry ®istry) override {
+ registry.insert<IREE::SimpleIO::SimpleIODialect>();
+ }
+
+ LogicalResult onActivate() override { return success(); }
+
+ void extendPreprocessingPassPipeline(OpPassManager &pm) override {
+ pm.addPass(IREE::SimpleIO::createLegalizeSimpleIOPass());
+ }
+};
+
+} // namespace
+
+IREE_DEFINE_COMPILER_OPTION_FLAGS(MyOptions);
+
+extern "C" bool iree_register_compiler_plugin_simple_io_sample(
+ mlir::iree_compiler::PluginRegistrar *registrar) {
+ registrar->registerPlugin<MySession>("simple_io_sample");
+ return true;
+}
diff --git a/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/IR/SimpleIODialect.cpp b/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/IR/SimpleIODialect.cpp
new file mode 100644
index 0000000..313a015
--- /dev/null
+++ b/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/IR/SimpleIODialect.cpp
@@ -0,0 +1,22 @@
+// 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 "simple_io_sample/IR/SimpleIODialect.h"
+
+#include "simple_io_sample/IR/SimpleIOOps.h"
+
+namespace mlir::iree_compiler::IREE::SimpleIO {
+
+void SimpleIODialect::initialize() {
+ addOperations<
+#define GET_OP_LIST
+#include "simple_io_sample/IR/SimpleIOOps.cpp.inc"
+ >();
+}
+
+} // namespace mlir::iree_compiler::IREE::SimpleIO
+
+#include "simple_io_sample/IR/SimpleIODialect.cpp.inc"
diff --git a/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/IR/SimpleIODialect.h b/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/IR/SimpleIODialect.h
new file mode 100644
index 0000000..06fc3bb
--- /dev/null
+++ b/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/IR/SimpleIODialect.h
@@ -0,0 +1,16 @@
+// 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
+
+#ifndef IREE_SAMPLES_COMPILER_PLUGINS_SIMPLE_IO_SAMPLE_IR_SIMPLEIODIALECT_H_
+#define IREE_SAMPLES_COMPILER_PLUGINS_SIMPLE_IO_SAMPLE_IR_SIMPLEIODIALECT_H_
+
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+
+// Include generated.
+#include "simple_io_sample/IR/SimpleIODialect.h.inc" // IWYU pragma: keep
+
+#endif // IREE_SAMPLES_COMPILER_PLUGINS_SIMPLE_IO_SAMPLE_IR_SIMPLEIODIALECT_H_
diff --git a/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/IR/SimpleIOOps.cpp b/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/IR/SimpleIOOps.cpp
new file mode 100644
index 0000000..1acebd4
--- /dev/null
+++ b/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/IR/SimpleIOOps.cpp
@@ -0,0 +1,14 @@
+// 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 "simple_io_sample/IR/SimpleIOOps.h"
+
+#include "mlir/IR/OpImplementation.h"
+
+// clang-format off
+#define GET_OP_CLASSES
+#include "simple_io_sample/IR/SimpleIOOps.cpp.inc" // IWYU pragma: keep
+// clang-format on
diff --git a/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/IR/SimpleIOOps.h b/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/IR/SimpleIOOps.h
new file mode 100644
index 0000000..f9602db
--- /dev/null
+++ b/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/IR/SimpleIOOps.h
@@ -0,0 +1,17 @@
+// 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
+
+#ifndef IREE_SAMPLES_COMPILER_PLUGINS_SIMPLE_IO_SAMPLE_IR_SIMPLEIOOPS_H_
+#define IREE_SAMPLES_COMPILER_PLUGINS_SIMPLE_IO_SAMPLE_IR_SIMPLEIOOPS_H_
+
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Operation.h"
+
+// Include generated.
+#define GET_OP_CLASSES
+#include "simple_io_sample/IR/SimpleIOOps.h.inc" // IWYU pragma: keep
+
+#endif // IREE_SAMPLES_COMPILER_PLUGINS_SIMPLE_IO_SAMPLE_IR_SIMPLEIOOPS_H_
diff --git a/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/IR/SimpleIOOps.td b/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/IR/SimpleIOOps.td
new file mode 100644
index 0000000..b5324f0
--- /dev/null
+++ b/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/IR/SimpleIOOps.td
@@ -0,0 +1,31 @@
+// 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
+
+#ifndef SIMPLE_IO_SAMPLE
+#define SIMPLE_IO_SAMPLE
+
+include "mlir/IR/OpBase.td"
+
+def SimpleIO_Dialect : Dialect {
+ let name = "simple_io";
+ let cppNamespace = "::mlir::iree_compiler::IREE::SimpleIO";
+}
+
+class SimpleIO_Op<string mnemonic, list<Trait> traits = []> :
+ Op<SimpleIO_Dialect, mnemonic, traits> {
+}
+
+def SimpleIO_PrintOp : SimpleIO_Op<"print", []> {
+ let summary = [{Print}];
+ let arguments = (ins);
+ let results = (outs);
+
+ let assemblyFormat = [{
+ attr-dict
+ }];
+}
+
+#endif
diff --git a/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/Transforms/LegalizeSimpleIO.cpp b/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/Transforms/LegalizeSimpleIO.cpp
new file mode 100644
index 0000000..4e60da2
--- /dev/null
+++ b/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/Transforms/LegalizeSimpleIO.cpp
@@ -0,0 +1,54 @@
+// 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 "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "simple_io_sample/IR/SimpleIOOps.h"
+#include "simple_io_sample/Transforms/Passes.h"
+
+#define GEN_PASS_DEF_LEGALIZESIMPLEIO
+#include "simple_io_sample/Transforms/Passes.h.inc"
+
+namespace mlir::iree_compiler::IREE::SimpleIO {
+namespace {
+
+class LegalizeSimpleIOPass
+ : public ::impl::LegalizeSimpleIOBase<LegalizeSimpleIOPass> {
+ public:
+ void runOnOperation() override {
+ auto *context = &getContext();
+ // TODO: This is all just a placeholder. To make it real, we should be
+ // checking if the import already exists and likely doing some more fancy
+ // lowering.
+ // Add imports.
+ auto m = getOperation();
+ auto importBuilder = OpBuilder::atBlockBegin(m.getBody());
+ importBuilder
+ .create<func::FuncOp>(m.getLoc(), "simple_io.print",
+ FunctionType::get(context, {}, {}))
+ .setPrivate();
+
+ // Legalize operations.
+ m.walk([&](Operation *op) {
+ if (auto printOp = dyn_cast<IREE::SimpleIO::PrintOp>(op)) {
+ OpBuilder b(op);
+ b.create<func::CallOp>(printOp.getLoc(), "simple_io.print",
+ TypeRange{});
+ printOp.erase();
+ }
+ });
+ }
+};
+
+} // namespace
+
+std::unique_ptr<OperationPass<ModuleOp>> createLegalizeSimpleIOPass() {
+ return std::make_unique<LegalizeSimpleIOPass>();
+}
+
+} // namespace mlir::iree_compiler::IREE::SimpleIO
diff --git a/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/Transforms/Passes.h b/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/Transforms/Passes.h
new file mode 100644
index 0000000..58c01d8
--- /dev/null
+++ b/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/Transforms/Passes.h
@@ -0,0 +1,20 @@
+// 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
+
+#ifndef IREE_SAMPLES_COMPILER_PLUGINS_SIMPLE_IO_SAMPLE_TRANSFORMS_PASSES_H_
+#define IREE_SAMPLES_COMPILER_PLUGINS_SIMPLE_IO_SAMPLE_TRANSFORMS_PASSES_H_
+
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+
+namespace mlir::iree_compiler::IREE::SimpleIO {
+
+std::unique_ptr<OperationPass<ModuleOp>> createLegalizeSimpleIOPass();
+
+} // namespace mlir::iree_compiler::IREE::SimpleIO
+
+#endif // IREE_SAMPLES_COMPILER_PLUGINS_SIMPLE_IO_SAMPLE_TRANSFORMS_PASSES_H_
diff --git a/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/Transforms/Passes.td b/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/Transforms/Passes.td
new file mode 100644
index 0000000..38f7f20
--- /dev/null
+++ b/samples/compiler_plugins/simple_io_sample/src/simple_io_sample/Transforms/Passes.td
@@ -0,0 +1,19 @@
+// 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 "mlir/Pass/PassBase.td"
+
+#ifndef IREE_SIMPLEIO_PASSES
+#define IREE_SIMPLEIO_PASSES
+
+def LegalizeSimpleIO : Pass<"iree-simpleio-legalize", "mlir::ModuleOp"> {
+ let summary = "Legalizes the simpleio sample ops";
+ let constructor = [{
+ ::mlir::iree_compiler::IREE::SimpleIO::createLegalizeSimpleIOPass()
+ }];
+}
+
+#endif // IREE_SIMPLEIO_PASSES
diff --git a/samples/compiler_plugins/simple_io_sample/test/print.mlir b/samples/compiler_plugins/simple_io_sample/test/print.mlir
new file mode 100644
index 0000000..b27b7b8
--- /dev/null
+++ b/samples/compiler_plugins/simple_io_sample/test/print.mlir
@@ -0,0 +1,7 @@
+// RUN: iree-opt --iree-plugin=simple_io_sample --iree-print-plugin-info %s
+
+func.func @main() {
+ // CHECK: call @simple_io.print
+ simple_io.print
+ func.return
+}
diff --git a/samples/compiler_plugins/simple_io_sample/test/run_mock.py b/samples/compiler_plugins/simple_io_sample/test/run_mock.py
new file mode 100644
index 0000000..db9454a
--- /dev/null
+++ b/samples/compiler_plugins/simple_io_sample/test/run_mock.py
@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+# 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
+
+# TODO: Turn this into a real test and wire it up.
+# Usage:
+# iree-compile --iree-plugin=simple_io_sample print.mlir -o print.vmfb
+# run_mock.py print.vmfb
+
+import iree.runtime as rt
+import sys
+
+input_file = sys.argv[1]
+print(f"--- Loading {input_file}")
+
+with open(input_file, "rb") as f:
+ vmfb_contents = f.read()
+
+
+def create_simple_io_module():
+
+ class SimpleIO:
+
+ def __init__(self, iface):
+ ...
+
+ def print_impl(self):
+ print("+++ HELLO FROM SIMPLE_IO")
+
+ iface = rt.PyModuleInterface("simple_io", SimpleIO)
+ iface.export("print", "0v_v", SimpleIO.print_impl)
+ return iface.create()
+
+
+config = rt.Config("local-sync")
+main_module = rt.VmModule.from_flatbuffer(config.vm_instance, vmfb_contents)
+modules = config.default_vm_modules + (
+ create_simple_io_module(),
+ main_module,
+)
+context = rt.SystemContext(vm_modules=modules, config=config)
+
+print("--- Running main()")
+context.modules.module.main()
