Initial boiler-plate for python bindings.

[stellaraccident]

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.

This is a draft to back up a more formal RFC to iron out naming, location, etc.

[mehdi_amini]

Fantastic! LGTM overall :)

mlir/CMakeLists.txt

85

I don't know CMake best practices, but is this something we could put in

mlir/lib/Bindings/CMakeLists.txt

101

Can you move binding_test as test/bindings ?

mlir/lib/Bindings/Python/CMakeLists.txt

63

I am not sure this will work on MacOS?

70

Thanks for the extensive doc in the Cmake files!

[stellaraccident]

mlir/CMakeLists.txt

85

Open to it but I usually see things like this done at the lowest ancestor directory in a big project (if not in the global cmake config). In this case, setting up the dependency in a leaf directory would create a fragile ordering between the Bindings subdirectory and the test directory (which needs a couple of config items from here).

101

Afaict, the entire 'test' tree is covered by a single lit test suite. I'm not a lit config expert, but I think this can be done by:

• Having something like (in the leaf directory):

if not config.enable_bindings_python:
  config.unsupported = True

• Doing some conditional setup in the test/lit.cfg.py

I'd have to fiddle with it as I've not conditionally included parts of the test tree like this yet.

mlir/lib/Bindings/Python/CMakeLists.txt

63

Probably not now that I think of it. My experience with MacOS ended 10 years ago. I can find an incantation, but I'll need a collaborator who has a Mac to actually make it work.

Do you know if we have a CMake variable we can condition on for MacOS?

70

Yw - some of this was hard-fought in previous iterations :) We probably want to refactor some of it to an include file, but I think it is useful to write out long-hand first.

[mehdi_amini]

mlir/CMakeLists.txt

101

Yes a lit.local.cfg with the snippet you mention seems like the way to go. I can help looking into this if needed.

mlir/lib/Bindings/Python/CMakeLists.txt

63

The MacOS linker to not support linker scripts. What is the requirement on the the visibility? Why don't we build the source with -fvisibility=hidden?

The rest of LLVM CMake seems to check with:

if(APPLE)

66

Your basic flow does not depend on MLIR yet, are we gonna depend on libMLIR.so, or will it be configurable to link the static lib or the .so?

[stellaraccident]

Move test suite, add condition to linker script, and add a trivial dep on MLIR to demonstrate.

[stellaraccident]

PTAL

mlir/lib/Bindings/Python/CMakeLists.txt

63

Added the conditional. I suspect that this will work as-is on MacOS but without some additional linker config may error at runtime if loaded into a process that exposes a different version of LLVM symbols (but I'm running off of some very old knowledge about the vagaries of linkage issues on that platform).

The requirement on visibility imposed by pybind is that any sources that use pybind must be compiled with -fvisibility=hidden, which is what the CXX_VISIBILITY_PRESET property is setting (in a platform agnostic way).

The additional thing I'm doing here is using a linker script to force the non-pybind *dependencies* to be hidden, which could otherwise only be controlled by setting an LLVM level CXX flag. This makes it so that "static" shared-library extension modules are self contained and do not leak LLVM exports to/from other shared-libraries that may be resident in the same process.

If building LLVM to only produce a python extension (i.e. as part of publishing to a pip archive or something), we could get somewhat better/smaller binaries by forcing all of LLVM to build with -fvisibility=hidden, but forcing it with linker scripts/flags is a nice middle ground that gives us the isolation while not leaking necessary compiler flags to the whole project.

66

I just added a trivial dep on MLIR to make sure it works.

I'm anticipating needing an option to control this behavior as all combinations are necessary in the limit. Currently this is building a self-contained ("static") shared library, which seems like a good starting point (and is what I generally prefer for development). I'd like to follow-up with the knobs for a true shared build (probably in conjunction with rules to copy the sources and built artifact to a dedicated directory in the build tree). Getting the shared linkage right can be a pain across platforms, so I'd like to make sure the simple static case works first/by default.

[mehdi_amini]

mlir/lib/Bindings/Python/CMakeLists.txt

63

OK now I get it, originally I thought you'd like against libMLIR.so and so the binding shared lib would be a "thin" object with naturally only the binding code.

It is likely not an issue with the two-levels namespace on MacOS in general (even though you can still hide things in the linker using an export list, which LLVM does for the LTO plugin: https://github.com/llvm/llvm-project/blob/master/llvm/tools/lto/CMakeLists.txt#L22 )

mlir/test/lit.cfg.py

45

Do you need these additions?

[stellaraccident]

mlir/lib/Bindings/Python/CMakeLists.txt

63

We definitely want to support the thin case but I'd like to iterate to get there and make sure we support the fat case, since that is important for various distribution scenarios (and has fewer moving parts to cause sharp edges -- which is why I prefer it for dev and cross builds).

mlir/test/lit.cfg.py

45

Yes - otherwise, adding the .py extension to the test suite advice causes lit to match its own config file and try to run it as a test :/

[stellaraccident]

Rebase to head.

[stellaraccident]

Squash.

[stephenneuendorffer]

mlir/test/CMakeLists.txt

9–13

I assume these were just never needed in the past? might be nice to have tests that depend on them? Is the rest dependent on this or should it maybe be separate?

[stellaraccident]

mlir/test/CMakeLists.txt

9–13

Yes, I found them missing when trying to figure out why -DMY_FEATURE_ENABLED=ON caused parse errors on the site.cfg. It looks like instead of adding them here, someone just started using the convention of defining them as 0/1 in the main file (and requiring a numeric flag on the command line). Adding them, makes it robust to saying -DMLIR_VULKAN_RUNNER_ENABLED=ON (versus '1') -- so I think people were just getting by before: there are tests for each of these and if you pass a non 0/1 bool on the command line you'll get a parse error without this addition.

This was fun to find. Let me bust this out and send separately.

[stellaraccident]

Rebase.

mlir/CMakeLists.txt

 
 option(MLIR_INCLUDE_TESTS
        "Generate build targets for the MLIR unit tests."
        ${LLVM_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})
 
 # Adding tools/mlir-tblgen here as calling add_tablegen sets some variables like
 # MLIR_TABLEGEN_EXE in PARENT_SCOPE which gets lost if that folder is included
 # from another directory like tools
 add_subdirectory(tools/mlir-tblgen)
 
 add_subdirectory(include/mlir)
 add_subdirectory(lib)
 if (MLIR_INCLUDE_TESTS)
   add_definitions(-DMLIR_INCLUDE_TESTS)
   add_subdirectory(unittests)
   add_subdirectory(test)
 endif()
 if (MLIR_INCLUDE_INTEGRATION_TESTS)
   add_definitions(-DMLIR_INCLUDE_INTEGRATION_TESTS)
   add_subdirectory(integration_test)
 endif()
 # Tools needs to come late to ensure that MLIR_ALL_LIBS is populated.
 # Generally things after this point may depend on MLIR_ALL_LIBS or libMLIR.so.
 add_subdirectory(tools)
 
 if( LLVM_INCLUDE_EXAMPLES )

mlir/lib/Bindings/CMakeLists.txt

+if(MLIR_BINDINGS_PYTHON_ENABLED)
+  add_subdirectory(Python)
+endif()

mlir/lib/Bindings/Python/CMakeLists.txt

+# 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

+//===- 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

+#  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

+{
+  global: PyInit__mlir;
+  local: *;
+};

mlir/lib/CMakeLists.txt

 # Enable errors for any global constructors.
 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)
 add_subdirectory(ExecutionEngine)
 add_subdirectory(IR)
 add_subdirectory(Interfaces)
 add_subdirectory(Parser)
 add_subdirectory(Pass)
 add_subdirectory(Reducer)
 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

+if not config.enable_bindings_python:
+  config.unsupported = True

mlir/test/Bindings/Python/smoke_test.py

+# RUN: %PYTHON %s | FileCheck %s
+
+import mlir
+
+# CHECK: From the native module
+print(mlir.get_test_value())

mlir/test/CMakeLists.txt

 add_subdirectory(EDSC)
 add_subdirectory(mlir-cpu-runner)
 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
   MLIR_ROCM_CONVERSIONS_ENABLED
   MLIR_ROCM_RUNNER_ENABLED
   MLIR_VULKAN_RUNNER_ENABLED
   )
 
 # Passed to lit.site.cfg.py.in to set up the path where to find the libraries
 # for linalg integration tests.
 set(MLIR_DIALECT_LINALG_INTEGRATION_TEST_LIB_DIR ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
 set(MLIR_RUNNER_UTILS_DIR ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
 
 # Passed to lit.site.cfg.py.in to set up the path where to find the libraries

 endif()
 
 if(MLIR_VULKAN_RUNNER_ENABLED)
   list(APPEND MLIR_TEST_DEPENDS
     mlir-vulkan-runner
   )
 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}
   )
 set_target_properties(check-mlir PROPERTIES FOLDER "Tests")
 
 add_lit_testsuites(MLIR ${CMAKE_CURRENT_SOURCE_DIR}
   DEPENDS ${MLIR_TEST_DEPS}
 )

mlir/test/lit.cfg.py

 # Configuration file for the 'lit' test runner.
 
 # name: The name of this test suite.
 config.name = 'MLIR'
 
 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__)
 
 # test_exec_root: The root path where tests should be run.
 config.test_exec_root = os.path.join(config.mlir_obj_root, 'test')
 
 config.substitutions.append(('%PATH%', config.environment['PATH']))
 config.substitutions.append(('%shlibext', config.llvm_shlib_ext))
 config.substitutions.append(("%mlir_src_root", config.mlir_src_root))
 
 llvm_config.with_system_environment(
     ['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP'])
 
 llvm_config.use_default_substitutions()
 
 # 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__)
 
 # test_exec_root: The root path where tests should be run.
 config.test_exec_root = os.path.join(config.mlir_obj_root, 'test')
 
 # Tweak the PATH to include the tools dir.
 llvm_config.with_environment('PATH', config.llvm_tools_dir, append_path=True)
 
 tool_dirs = [config.mlir_tools_dir, config.llvm_tools_dir]
 tools = [
     'mlir-opt',
     'mlir-tblgen',
     'mlir-translate',
     'mlir-edsc-builder-api-test',
 ]
 
 # 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'),
     ToolSubst('toy-ch4', unresolved='ignore'),
     ToolSubst('toy-ch5', unresolved='ignore'),
     ToolSubst('%cuda_wrapper_library_dir', config.cuda_wrapper_library_dir, unresolved='ignore'),
     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)
 
 
 # FileCheck -enable-var-scope is enabled by default in MLIR test
 # This option avoids to accidentally reuse variable across -LABEL match,
 # it can be explicitly opted-in by prefixing the variable name with $
 config.environment['FILECHECK_OPTS'] = "-enable-var-scope"
 
 
 # LLVM can be configured with an empty default triple
 # by passing ` -DLLVM_DEFAULT_TARGET_TRIPLE="" `.
 # This is how LLVM filters tests that require the host target
 # 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

 config.run_cuda_tests = @MLIR_CUDA_CONVERSIONS_ENABLED@
 config.cuda_wrapper_library_dir = "@MLIR_CUDA_WRAPPER_LIBRARY_DIR@"
 config.enable_cuda_runner = @MLIR_CUDA_RUNNER_ENABLED@
 config.run_rocm_tests = @MLIR_ROCM_CONVERSIONS_ENABLED@
 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.
 try:
     config.llvm_tools_dir = config.llvm_tools_dir % lit_config.params
     config.llvm_shlib_dir = config.llvm_shlib_dir % lit_config.params
 except KeyError:
     e = sys.exc_info()[1]
     key, = e.args
     lit_config.fatal("unable to find %r parameter, use '--param=%s=VALUE'" % (key,key))
 
 import lit.llvm
 lit.llvm.initialize(lit_config, config)
 
 # Let the main config do the real work.
 lit_config.load_config(config, "@MLIR_SOURCE_DIR@/test/lit.cfg.py")