[Mlir-commits] [mlir] ac0a70f - [mlir] Split out Python bindings entry point into a separate file
Alex Zinenko
llvmlistbot at llvm.org
Thu Apr 29 02:18:33 PDT 2021
Author: Alex Zinenko
Date: 2021-04-29T11:18:25+02:00
New Revision: ac0a70f3737ecb2c0586c00240d14e46ff00644e
URL: https://github.com/llvm/llvm-project/commit/ac0a70f3737ecb2c0586c00240d14e46ff00644e
DIFF: https://github.com/llvm/llvm-project/commit/ac0a70f3737ecb2c0586c00240d14e46ff00644e.diff
LOG: [mlir] Split out Python bindings entry point into a separate file
This will allow the bindings to be built as a library and reused in out-of-tree
projects that want to provide bindings on top of MLIR bindings.
Reviewed By: stellaraccident, mikeurbach
Differential Revision: https://reviews.llvm.org/D101075
Added:
mlir/lib/Bindings/Python/IRModule.cpp
Modified:
mlir/lib/Bindings/Python/CMakeLists.txt
mlir/lib/Bindings/Python/MainModule.cpp
Removed:
################################################################################
diff --git a/mlir/lib/Bindings/Python/CMakeLists.txt b/mlir/lib/Bindings/Python/CMakeLists.txt
index bbccea63c461..580405f09619 100644
--- a/mlir/lib/Bindings/Python/CMakeLists.txt
+++ b/mlir/lib/Bindings/Python/CMakeLists.txt
@@ -84,6 +84,7 @@ add_mlir_python_extension(MLIRCoreBindingsPythonExtension _mlir
IRAffine.cpp
IRAttributes.cpp
IRCore.cpp
+ IRModule.cpp
IRTypes.cpp
PybindUtils.cpp
Pass.cpp
diff --git a/mlir/lib/Bindings/Python/IRModule.cpp b/mlir/lib/Bindings/Python/IRModule.cpp
new file mode 100644
index 000000000000..08ce06da8783
--- /dev/null
+++ b/mlir/lib/Bindings/Python/IRModule.cpp
@@ -0,0 +1,146 @@
+//===- IRModule.cpp - IR pybind module ------------------------------------===//
+//
+// Part of the LLVM Project, 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 "IRModule.h"
+#include "Globals.h"
+#include "PybindUtils.h"
+
+#include <vector>
+
+namespace py = pybind11;
+using namespace mlir;
+using namespace mlir::python;
+
+// -----------------------------------------------------------------------------
+// PyGlobals
+// -----------------------------------------------------------------------------
+
+PyGlobals *PyGlobals::instance = nullptr;
+
+PyGlobals::PyGlobals() {
+ assert(!instance && "PyGlobals already constructed");
+ instance = this;
+}
+
+PyGlobals::~PyGlobals() { instance = nullptr; }
+
+void PyGlobals::loadDialectModule(llvm::StringRef dialectNamespace) {
+ py::gil_scoped_acquire();
+ if (loadedDialectModulesCache.contains(dialectNamespace))
+ return;
+ // Since re-entrancy is possible, make a copy of the search prefixes.
+ std::vector<std::string> localSearchPrefixes = dialectSearchPrefixes;
+ py::object loaded;
+ for (std::string moduleName : localSearchPrefixes) {
+ moduleName.push_back('.');
+ moduleName.append(dialectNamespace.data(), dialectNamespace.size());
+
+ try {
+ py::gil_scoped_release();
+ loaded = py::module::import(moduleName.c_str());
+ } catch (py::error_already_set &e) {
+ if (e.matches(PyExc_ModuleNotFoundError)) {
+ continue;
+ } else {
+ throw;
+ }
+ }
+ break;
+ }
+
+ // Note: Iterator cannot be shared from prior to loading, since re-entrancy
+ // may have occurred, which may do anything.
+ loadedDialectModulesCache.insert(dialectNamespace);
+}
+
+void PyGlobals::registerDialectImpl(const std::string &dialectNamespace,
+ py::object pyClass) {
+ py::gil_scoped_acquire();
+ py::object &found = dialectClassMap[dialectNamespace];
+ if (found) {
+ throw SetPyError(PyExc_RuntimeError, llvm::Twine("Dialect namespace '") +
+ dialectNamespace +
+ "' is already registered.");
+ }
+ found = std::move(pyClass);
+}
+
+void PyGlobals::registerOperationImpl(const std::string &operationName,
+ py::object pyClass,
+ py::object rawOpViewClass) {
+ py::gil_scoped_acquire();
+ py::object &found = operationClassMap[operationName];
+ if (found) {
+ throw SetPyError(PyExc_RuntimeError, llvm::Twine("Operation '") +
+ operationName +
+ "' is already registered.");
+ }
+ found = std::move(pyClass);
+ rawOpViewClassMap[operationName] = std::move(rawOpViewClass);
+}
+
+llvm::Optional<py::object>
+PyGlobals::lookupDialectClass(const std::string &dialectNamespace) {
+ py::gil_scoped_acquire();
+ loadDialectModule(dialectNamespace);
+ // Fast match against the class map first (common case).
+ const auto foundIt = dialectClassMap.find(dialectNamespace);
+ if (foundIt != dialectClassMap.end()) {
+ if (foundIt->second.is_none())
+ return llvm::None;
+ assert(foundIt->second && "py::object is defined");
+ return foundIt->second;
+ }
+
+ // Not found and loading did not yield a registration. Negative cache.
+ dialectClassMap[dialectNamespace] = py::none();
+ return llvm::None;
+}
+
+llvm::Optional<pybind11::object>
+PyGlobals::lookupRawOpViewClass(llvm::StringRef operationName) {
+ {
+ py::gil_scoped_acquire();
+ auto foundIt = rawOpViewClassMapCache.find(operationName);
+ if (foundIt != rawOpViewClassMapCache.end()) {
+ if (foundIt->second.is_none())
+ return llvm::None;
+ assert(foundIt->second && "py::object is defined");
+ return foundIt->second;
+ }
+ }
+
+ // Not found. Load the dialect namespace.
+ auto split = operationName.split('.');
+ llvm::StringRef dialectNamespace = split.first;
+ loadDialectModule(dialectNamespace);
+
+ // Attempt to find from the canonical map and cache.
+ {
+ py::gil_scoped_acquire();
+ auto foundIt = rawOpViewClassMap.find(operationName);
+ if (foundIt != rawOpViewClassMap.end()) {
+ if (foundIt->second.is_none())
+ return llvm::None;
+ assert(foundIt->second && "py::object is defined");
+ // Positive cache.
+ rawOpViewClassMapCache[operationName] = foundIt->second;
+ return foundIt->second;
+ } else {
+ // Negative cache.
+ rawOpViewClassMap[operationName] = py::none();
+ return llvm::None;
+ }
+ }
+}
+
+void PyGlobals::clearImportCache() {
+ py::gil_scoped_acquire();
+ loadedDialectModulesCache.clear();
+ rawOpViewClassMapCache.clear();
+}
diff --git a/mlir/lib/Bindings/Python/MainModule.cpp b/mlir/lib/Bindings/Python/MainModule.cpp
index 79128f2677a5..60c282d1d9d3 100644
--- a/mlir/lib/Bindings/Python/MainModule.cpp
+++ b/mlir/lib/Bindings/Python/MainModule.cpp
@@ -20,135 +20,6 @@ namespace py = pybind11;
using namespace mlir;
using namespace mlir::python;
-// -----------------------------------------------------------------------------
-// PyGlobals
-// -----------------------------------------------------------------------------
-
-PyGlobals *PyGlobals::instance = nullptr;
-
-PyGlobals::PyGlobals() {
- assert(!instance && "PyGlobals already constructed");
- instance = this;
-}
-
-PyGlobals::~PyGlobals() { instance = nullptr; }
-
-void PyGlobals::loadDialectModule(llvm::StringRef dialectNamespace) {
- py::gil_scoped_acquire();
- if (loadedDialectModulesCache.contains(dialectNamespace))
- return;
- // Since re-entrancy is possible, make a copy of the search prefixes.
- std::vector<std::string> localSearchPrefixes = dialectSearchPrefixes;
- py::object loaded;
- for (std::string moduleName : localSearchPrefixes) {
- moduleName.push_back('.');
- moduleName.append(dialectNamespace.data(), dialectNamespace.size());
-
- try {
- py::gil_scoped_release();
- loaded = py::module::import(moduleName.c_str());
- } catch (py::error_already_set &e) {
- if (e.matches(PyExc_ModuleNotFoundError)) {
- continue;
- } else {
- throw;
- }
- }
- break;
- }
-
- // Note: Iterator cannot be shared from prior to loading, since re-entrancy
- // may have occurred, which may do anything.
- loadedDialectModulesCache.insert(dialectNamespace);
-}
-
-void PyGlobals::registerDialectImpl(const std::string &dialectNamespace,
- py::object pyClass) {
- py::gil_scoped_acquire();
- py::object &found = dialectClassMap[dialectNamespace];
- if (found) {
- throw SetPyError(PyExc_RuntimeError, llvm::Twine("Dialect namespace '") +
- dialectNamespace +
- "' is already registered.");
- }
- found = std::move(pyClass);
-}
-
-void PyGlobals::registerOperationImpl(const std::string &operationName,
- py::object pyClass,
- py::object rawOpViewClass) {
- py::gil_scoped_acquire();
- py::object &found = operationClassMap[operationName];
- if (found) {
- throw SetPyError(PyExc_RuntimeError, llvm::Twine("Operation '") +
- operationName +
- "' is already registered.");
- }
- found = std::move(pyClass);
- rawOpViewClassMap[operationName] = std::move(rawOpViewClass);
-}
-
-llvm::Optional<py::object>
-PyGlobals::lookupDialectClass(const std::string &dialectNamespace) {
- py::gil_scoped_acquire();
- loadDialectModule(dialectNamespace);
- // Fast match against the class map first (common case).
- const auto foundIt = dialectClassMap.find(dialectNamespace);
- if (foundIt != dialectClassMap.end()) {
- if (foundIt->second.is_none())
- return llvm::None;
- assert(foundIt->second && "py::object is defined");
- return foundIt->second;
- }
-
- // Not found and loading did not yield a registration. Negative cache.
- dialectClassMap[dialectNamespace] = py::none();
- return llvm::None;
-}
-
-llvm::Optional<pybind11::object>
-PyGlobals::lookupRawOpViewClass(llvm::StringRef operationName) {
- {
- py::gil_scoped_acquire();
- auto foundIt = rawOpViewClassMapCache.find(operationName);
- if (foundIt != rawOpViewClassMapCache.end()) {
- if (foundIt->second.is_none())
- return llvm::None;
- assert(foundIt->second && "py::object is defined");
- return foundIt->second;
- }
- }
-
- // Not found. Load the dialect namespace.
- auto split = operationName.split('.');
- llvm::StringRef dialectNamespace = split.first;
- loadDialectModule(dialectNamespace);
-
- // Attempt to find from the canonical map and cache.
- {
- py::gil_scoped_acquire();
- auto foundIt = rawOpViewClassMap.find(operationName);
- if (foundIt != rawOpViewClassMap.end()) {
- if (foundIt->second.is_none())
- return llvm::None;
- assert(foundIt->second && "py::object is defined");
- // Positive cache.
- rawOpViewClassMapCache[operationName] = foundIt->second;
- return foundIt->second;
- } else {
- // Negative cache.
- rawOpViewClassMap[operationName] = py::none();
- return llvm::None;
- }
- }
-}
-
-void PyGlobals::clearImportCache() {
- py::gil_scoped_acquire();
- loadedDialectModulesCache.clear();
- rawOpViewClassMapCache.clear();
-}
-
// -----------------------------------------------------------------------------
// Module initialization.
// -----------------------------------------------------------------------------
More information about the Mlir-commits
mailing list