Initial boiler-plate for python bindings.

Summary:
* Native '_mlir' extension module.
* Python mlir/__init__.py trampoline module.
* Lit test that checks a message.
* Uses some cmake configurations that have worked for me in the past but likely needs further elaboration.

Subscribers: mgorny, mehdi_amini, rriddle, jpienaar, shauheen, antiagainst, nicolasvasilache, arpith-jacob, mgester, lucyrfox, aartbik, liufengdb, stephenneuendorffer, Joonsoo, grosul1, Kayjukh, jurahul, msifontes

Tags: #mlir

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

Author: stellaraccident

mlir/CMakeLists.txt

@@ -51,6 +51,38 @@ option(MLIR_INCLUDE_TESTS
 option(MLIR_INCLUDE_INTEGRATION_TESTS
        "Generate build targets for the MLIR integration tests.")
 
+#-------------------------------------------------------------------------------
+# Python Bindings Configuration
+# Requires:
+#   The pybind11 library can be found (set with -DPYBIND_DIR=...)
+#   The python executable is correct (set with -DPYTHON_EXECUTABLE=...)
+#
+# Version locking
+# ---------------
+# By default, python extensions are version locked to specific Python libraries.
+# This linking mode is somewhat more consistent across platforms and surfaces
+# undefined symbols at link time (vs runtime). It is suitable for development
+# workflows but can be disabled for more flexible deployment by
+# setting -DMLIR_PYTHON_BINDINGS_VERSION_LOCKED=OFF
+#-------------------------------------------------------------------------------
+
+set(MLIR_BINDINGS_PYTHON_ENABLED 0 CACHE BOOL
+       "Enables building of Python bindings.")
+set(MLIR_PYTHON_BINDINGS_VERSION_LOCKED 1 CACHE BOOL
+       "Links to specific python libraries, resolving all symbols.")
+
+if(MLIR_BINDINGS_PYTHON_ENABLED)
+  find_package(PythonInterp REQUIRED)
+  find_package(PythonLibs REQUIRED)
+  message(STATUS "Found python include dirs: ${PYTHON_INCLUDE_DIRS}")
+  message(STATUS "Found ppython libraries: ${PYTHON_LIBRARIES}")
+  find_package(pybind11 CONFIG REQUIRED)
+  message(STATUS "Found pybind11 v${pybind11_VERSION}: ${pybind11_INCLUDE_DIRS}")
+  message(STATUS "Python prefix = '${PYTHON_MODULE_PREFIX}', "
+                 "suffix = '${PYTHON_MODULE_SUFFIX}', "
+                 "extension = '${PYTHON_MODULE_EXTENSION}")
+endif()
+
 include_directories( "include")
 include_directories( ${MLIR_INCLUDE_DIR})
 

mlir/lib/Bindings/CMakeLists.txt

@@ -0,0 +1,3 @@
+if(MLIR_BINDINGS_PYTHON_ENABLED)
+  add_subdirectory(Python)
+endif()

mlir/lib/Bindings/Python/CMakeLists.txt

@@ -0,0 +1,70 @@
+# Normally on unix-like platforms, extensions are built as "MODULE" libraries
+# and do not explicitly link to the python shared object. This allows for
+# some greater deployment flexibility since the extension will bind to
+# symbols in the python interpreter on load. However, it also keeps the
+# linker from erroring on undefined symbols, leaving this to (usually obtuse)
+# runtime errors. Building in "SHARED" mode with an explicit link to the
+# python libraries allows us to build with the expectation of no undefined
+# symbols, which is better for development.
+if(MLIR_PYTHON_BINDINGS_VERSION_LOCKED)
+  set(PYEXT_LINK_MODE SHARED)
+  set(PYEXT_LIBADD ${PYTHON_LIBRARIES})
+else()
+  set(PYEXT_LINK_MODE MODULE)
+  set(PYEXT_LIBADD)
+endif()
+
+# The actual extension library produces a shared-object or DLL and has
+# sources that must be compiled in accordance with pybind11 needs (RTTI and
+# exceptions).
+add_library(MLIRBindingsPythonExtension ${PYEXT_LINK_MODE}
+  MainModule.cpp
+)
+
+target_include_directories(MLIRBindingsPythonExtension PRIVATE
+  "${PYTHON_INCLUDE_DIRS}"
+  "${pybind11_INCLUDE_DIRS}")
+
+# The extension itself must be compiled with RTTI and exceptions enabled.
+# Also, some warning classes triggered by pybind11 are disabled.
+target_compile_options(MLIRBindingsPythonExtension PRIVATE
+  $<$<OR:$<CXX_COMPILER_ID:Clang>,$<CXX_COMPILER_ID:AppleClang>,$<CXX_COMPILER_ID:GNU>>:
+    # Enable RTTI and exceptions.
+    -frtti -fexceptions
+    # Noisy pybind warnings
+    -Wno-unused-value
+    -Wno-covered-switch-default
+  >
+  $<$<CXX_COMPILER_ID:MSVC>:
+    # Enable RTTI and exceptions.
+    /EHsc /GR>
+)
+
+# Configure the output to match python expectations.
+set_target_properties(
+  MLIRBindingsPythonExtension PROPERTIES
+    LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}"
+    OUTPUT_NAME "_mlir"
+    PREFIX "${PYTHON_MODULE_PREFIX}"
+    SUFFIX "${PYTHON_MODULE_SUFFIX}${PYTHON_MODULE_EXTENSION}"
+)
+
+# pybind11 requires binding code to be compiled with -fvisibility=hidden
+# For static linkage, better code can be generated if the entire project
+# compiles that way, but that is not enforced here. Instead, include a linker
+# script that explicitly hides anything but the PyInit_* symbols, allowing gc
+# to take place.
+# TODO: Add a Windows .def file and figure out the right thing to do on MacOS.
+set_target_properties(
+  MLIRBindingsPythonExtension PROPERTIES CXX_VISIBILITY_PRESET "hidden")
+if(NOT MSVC AND NOT APPLE)
+  set_target_properties(MLIRBindingsPythonExtension
+    PROPERTIES LINK_FLAGS
+    "-Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/unix_version.lds")
+endif()
+
+target_link_libraries(MLIRBindingsPythonExtension
+  PRIVATE
+  MLIRIR
+  ${PYEXT_LIBADD}
+)

mlir/lib/Bindings/Python/MainModule.cpp

@@ -0,0 +1,27 @@
+//===- MainModule.cpp - Main 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 <tuple>
+
+#include <pybind11/pybind11.h>
+
+#include "mlir/IR/MLIRContext.h"
+
+using namespace mlir;
+
+PYBIND11_MODULE(_mlir, m) {
+  m.doc() = "MLIR Python Native Extension";
+
+  m.def("get_test_value", []() {
+    // This is just calling a method on the MLIRContext as a smoketest
+    // for linkage.
+    MLIRContext context;
+    return std::make_tuple(std::string("From the native module"),
+                           context.isMultithreadingEnabled());
+  });
+}

mlir/lib/Bindings/Python/mlir/__init__.py

@@ -0,0 +1,11 @@
+#  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
+
+# Note that the only function of this module is currently to load the
+# native module and re-export its symbols. In the future, this file is
+# reserved as a trampoline to handle environment specific loading needs
+# and arbitrate any one-time initialization needed in various shared-library
+# scenarios.
+
+from _mlir import *

mlir/lib/Bindings/Python/unix_version.lds

@@ -0,0 +1,4 @@
+{
+  global: PyInit__mlir;
+  local: *;
+};

mlir/lib/CMakeLists.txt

@@ -2,6 +2,7 @@
 add_flag_if_supported("-Werror=global-constructors" WERROR_GLOBAL_CONSTRUCTOR)
 
 add_subdirectory(Analysis)
+add_subdirectory(Bindings)
 add_subdirectory(Conversion)
 add_subdirectory(Dialect)
 add_subdirectory(EDSC)
@@ -15,4 +16,4 @@ add_subdirectory(Support)
 add_subdirectory(TableGen)
 add_subdirectory(Target)
 add_subdirectory(Transforms)
-add_subdirectory(Translation)
+add_subdirectory(Translation)

mlir/test/Bindings/Python/lit.local.cfg

@@ -0,0 +1,2 @@
+if not config.enable_bindings_python:
+  config.unsupported = True

mlir/test/Bindings/Python/smoke_test.py

@@ -0,0 +1,6 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+import mlir
+
+# CHECK: From the native module
+print(mlir.get_test_value())

mlir/test/CMakeLists.txt

@@ -4,6 +4,7 @@ add_subdirectory(SDBM)
 add_subdirectory(lib)
 
 llvm_canonicalize_cmake_booleans(
+  MLIR_BINDINGS_PYTHON_ENABLED
   LLVM_BUILD_EXAMPLES
   MLIR_CUDA_CONVERSIONS_ENABLED
   MLIR_CUDA_RUNNER_ENABLED
@@ -83,6 +84,12 @@ if(MLIR_VULKAN_RUNNER_ENABLED)
   )
 endif()
 
+if(MLIR_BINDINGS_PYTHON_ENABLED)
+  list(APPEND MLIR_TEST_DEPENDS
+    MLIRBindingsPythonExtension
+  )
+endif()
+
 add_lit_testsuite(check-mlir "Running the MLIR regression tests"
   ${CMAKE_CURRENT_BINARY_DIR}
   DEPENDS ${MLIR_TEST_DEPENDS}

mlir/test/lit.cfg.py

@@ -21,7 +21,7 @@
 config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)
 
 # suffixes: A list of file extensions to treat as test files.
-config.suffixes = ['.td', '.mlir', '.toy', '.ll', '.tc']
+config.suffixes = ['.td', '.mlir', '.toy', '.ll', '.tc', '.py']
 
 # test_source_root: The root path where tests are located.
 config.test_source_root = os.path.dirname(__file__)
@@ -41,7 +41,8 @@
 # excludes: A list of directories to exclude from the testsuite. The 'Inputs'
 # subdirectories contain auxiliary inputs for various tests in their parent
 # directories.
-config.excludes = ['Inputs', 'CMakeLists.txt', 'README.txt', 'LICENSE.txt']
+config.excludes = ['Inputs', 'CMakeLists.txt', 'README.txt', 'LICENSE.txt',
+                   'lit.cfg.py', 'lit.site.cfg.py']
 
 # test_source_root: The root path where tests are located.
 config.test_source_root = os.path.dirname(__file__)
@@ -62,6 +63,7 @@
 
 # The following tools are optional
 tools.extend([
+    ToolSubst('%PYTHON', config.python_executable),
     ToolSubst('toy-ch1', unresolved='ignore'),
     ToolSubst('toy-ch2', unresolved='ignore'),
     ToolSubst('toy-ch3', unresolved='ignore'),
@@ -71,7 +73,7 @@
     ToolSubst('%linalg_test_lib_dir', config.linalg_test_lib_dir, unresolved='ignore'),
     ToolSubst('%mlir_runner_utils_dir', config.mlir_runner_utils_dir, unresolved='ignore'),
     ToolSubst('%rocm_wrapper_library_dir', config.rocm_wrapper_library_dir, unresolved='ignore'),
-    ToolSubst('%vulkan_wrapper_library_dir', config.vulkan_wrapper_library_dir, unresolved='ignore')
+    ToolSubst('%vulkan_wrapper_library_dir', config.vulkan_wrapper_library_dir, unresolved='ignore'),
 ])
 
 llvm_config.add_tool_substitutions(tools, tool_dirs)
@@ -89,3 +91,13 @@
 # to be available for JIT tests.
 if config.target_triple:
     config.available_features.add('default_triple')
+
+# Add the python path for both the source and binary tree.
+# Note that presently, the python sources come from the source tree and the
+# binaries come from the build tree. This should be unified to the build tree
+# by copying/linking sources to build.
+if config.enable_bindings_python:
+    llvm_config.with_environment('PYTHONPATH', [
+        os.path.join(config.mlir_src_root, "lib", "Bindings", "Python"),
+        os.path.join(config.mlir_obj_root, "lib", "Bindings", "Python"),
+    ], append_path=True)

mlir/test/lit.site.cfg.py.in

@@ -43,6 +43,7 @@ config.rocm_wrapper_library_dir = "@MLIR_ROCM_WRAPPER_LIBRARY_DIR@"
 config.enable_rocm_runner = @MLIR_ROCM_RUNNER_ENABLED@
 config.vulkan_wrapper_library_dir = "@MLIR_VULKAN_WRAPPER_LIBRARY_DIR@"
 config.enable_vulkan_runner = @MLIR_VULKAN_RUNNER_ENABLED@
+config.enable_bindings_python = @MLIR_BINDINGS_PYTHON_ENABLED@
 
 # Support substitution of the tools_dir with user parameters. This is
 # used when we can't determine the tool dir at configuration time.