Add Python bindings (#1064)

Garra1980 · tkarna · web-flow · commit d2e405ef6c92 · 2025-04-15T09:32:45.000-05:00
Co-authored-by: Tuomas Karna &lt;tuomas.karna@intel.com&gt;
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -272,6 +272,13 @@ else ()
     set(IMEX_ENABLE_PVC_TARGET 0)
 endif()
 
+set(IMEX_ENABLE_BINDINGS_PYTHON 0 CACHE BOOL "Enable IMEX Python Binding")
+if (IMEX_ENABLE_BINDINGS_PYTHON)
+    set(CMAKE_CXX_FLAGS " ${CMAKE_CXX_FLAGS} -Wno-error=address -Wno-error=pragmas -Wno-error=pedantic -Wno-error=parentheses")
+    set(IMEX_ENABLE_BINDINGS_PYTHON 1)
+else ()
+    set(IMEX_ENABLE_BINDINGS_PYTHON 0)
+endif()
 
 if (NOT (CMAKE_BUILD_TYPE MATCHES Debug))
     if (IMEX_ENABLE_DEBUG_BUILD)
@@ -326,6 +333,9 @@ option(IMEX_INCLUDE_DOCS "Generate build targets for the IMEX docs." ON)
 if (IMEX_INCLUDE_DOCS)
     add_subdirectory(docs)
 endif()
+if (IMEX_ENABLE_BINDINGS_PYTHON)
+    add_subdirectory(python)
+endif()
 
 # Custom target to install all imex headers
 add_custom_target(imex-headers)
diff --git a/README.md b/README.md
@@ -146,6 +146,17 @@ cd build
 cmake --build build --target doc_doxygen
 ```
 
+### Building Python bindings
+
+IMEX can be built with Python bindings. The `imex_mlir` Python package functions similarly to the [MLIR Python bindings](https://mlir.llvm.org/docs/Bindings/Python/), providing access to the MLIR objects as well as IMEX dialects and passes. To enable the bindings:
+
+- LLVM must be built with `-DMLIR_ENABLE_BINDINGS_PYTHON=1`.
+- IMEX must be built with `-DIMEX_ENABLE_BINDINGS_PYTHON=1`.
+- On systems with multiple Python implementations, you may need to specify the desired Python version, for example, `-DPython3_EXECUTABLE=$(which python3)`.
+- pybind11 and nanobind are additional build dependencies.
+
+The `imex_mlir` package is located in the `<IMEX_ROOT>/python_packages/` directory, where `<IMEX_ROOT>` refers to the IMEX build or installation directory. To use the Python bindings, include `<IMEX_ROOT>/python_packages/` in your `PYTHONPATH` environment variable. Usage examples can be found in the [Python test suite](test/python/).
+
 ## Adding a new dialect
 ```sh
 # enter root directory of mlir-extension
diff --git a/include/imex-c/Dialects.h b/include/imex-c/Dialects.h
@@ -0,0 +1,28 @@
+//===- Dialects.h - IMEX CAPI Registration ----------------------*- C++ -*-===//
+//
+// Copyright 2025 Intel Corporation
+// Part of the IMEX 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the IMEX C API dialects.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef IMEX_C_DIALECTS_H
+#define IMEX_C_DIALECTS_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Region, region);
+
+#ifdef __cplusplus
+}
+#endif
+#endif // IMEX_C_DIALECTS_H
diff --git a/include/imex-c/Passes.h b/include/imex-c/Passes.h
@@ -0,0 +1,28 @@
+//===- Passes.h - IMEX CAPI Registration ------------------------*- C++ -*-===//
+//
+// Copyright 2025 Intel Corporation
+// Part of the IMEX 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the IMEX C API passes.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef IMEX_MLIR_C_PASSES_H
+#define IMEX_MLIR_C_PASSES_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_CAPI_EXPORTED void mlirRegisterAllIMEXPasses(void);
+
+#ifdef __cplusplus
+}
+#endif
+#endif // IMEX_MLIR_C_PASSES_H
diff --git a/lib/CAPI/CMakeLists.txt b/lib/CAPI/CMakeLists.txt
@@ -0,0 +1,44 @@
+set(IMEX_ALL_LIBS
+  IMEXTransforms
+  IMEXXeTileTransforms
+  IMEXXeTileDialect
+  IMEXNDArrayTransforms
+  IMEXNDArrayDialect
+  IMEXGPUXDialect
+  IMEXDistRuntimeTransforms
+  IMEXDistRuntimeDialect
+  IMEXRegionTransforms
+  IMEXRegionDialect
+  IMEXArithToVC
+  IMEXDropRegions
+  IMEXGPUToGPUX
+  IMEXGPUToSPIRV
+  IMEXGPUXToLLVM
+  IMEXMathToVC
+  IMEXNDArrayToLinalg
+  IMEXSCFToGPU
+  IMEXXeGPUToVC
+  IMEXXeTileToXeGPU
+  MLIRXeGPUToXeVM
+  MLIRXeVMToLLVM
+  IMEXUtil
+  imex_runner_utils
+)
+
+if(IMEX_ENABLE_L0_RUNTIME)
+  list(APPEND IMEX_ALL_LIBS level-zero-runtime)
+endif()
+
+if(IMEX_ENABLE_SYCL_RUNTIME)
+  list(APPEND IMEX_ALL_LIBS sycl-runtime)
+endif()
+
+add_mlir_public_c_api_library(ImexCAPI
+  Dialects.cpp
+  Passes.cpp
+  LINK_LIBS PUBLIC
+  IMEXRegionDialect
+  ${IMEX_ALL_LIBS}
+)
+target_link_libraries(obj.ImexCAPI PUBLIC ImexInterface)
+set_property(GLOBAL APPEND PROPERTY IMEX_MLIR_LIBS ImexCAPI)
diff --git a/lib/CAPI/Dialects.cpp b/lib/CAPI/Dialects.cpp
@@ -0,0 +1,16 @@
+//===- Dialects.cpp - IMEX CAPI Registration --------------------*- C++ -*-===//
+//
+// Copyright 2025 Intel Corporation
+// Part of the IMEX 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 "imex-c/Dialects.h"
+#include "imex/Dialect/Region/IR/RegionOps.h"
+#include "mlir/CAPI/Registration.h"
+
+// name, namespace, classname
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Region, region,
+                                      imex::region::RegionDialect)
diff --git a/lib/CAPI/Passes.cpp b/lib/CAPI/Passes.cpp
@@ -0,0 +1,23 @@
+//===- Passes.cpp - IMEX CAPI Registration ----------------------*- C++ -*-===//
+//
+// Copyright 2025 Intel Corporation
+// Part of the IMEX 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 "imex/InitIMEXPasses.h"
+#include "mlir-c/Pass.h"
+#include "mlir/CAPI/Pass.h"
+
+using namespace imex;
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_CAPI_EXPORTED void mlirRegisterAllIMEXPasses() { registerAllPasses(); }
+#ifdef __cplusplus
+}
+#endif
diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt
@@ -1,4 +1,5 @@
 add_subdirectory(Dialect)
+add_subdirectory(CAPI)
 add_subdirectory(Conversion)
 add_subdirectory(Transforms)
 add_subdirectory(ExecutionEngine)
diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
@@ -0,0 +1,84 @@
+if(IMEX_ENABLE_BINDINGS_PYTHON)
+  if(NOT MLIR_ENABLE_BINDINGS_PYTHON)
+    message(STATUS "Failed to enable Python API due to the 'MLIR_ENABLE_BINDINGS_PYTHON' for LLVM is not ON.")
+    set(IMEX_ENABLE_BINDINGS_PYTHON OFF CACHE BOOL "" FORCE)
+    return()
+  endif()
+  message(STATUS "Enabling Python API")
+else()
+  return()
+endif()
+
+include(MLIRDetectPythonEnv)
+mlir_configure_python_dev_packages()
+include(AddMLIRPython)
+
+# Specifies that all MLIR packages are co-located under the `imex_mlir`
+# top level package (the API has been embedded in a relocatable way).
+# TODO: Add an upstream cmake param for this vs having a global here.
+add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=imex_mlir.")
+
+
+################################################################################
+# Sources
+################################################################################
+
+declare_mlir_python_sources(ImexPythonSources)
+
+declare_mlir_python_sources(ImexPythonSources.Common
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/imex_mlir"
+  ADD_TO_PARENT ImexPythonSources
+  SOURCES
+    _mlir_libs/_site_initialize_0.py
+)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT ImexPythonSources
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/imex_mlir"
+  TD_FILE dialects/RegionOps.td
+  SOURCES
+    dialects/region.py
+  DIALECT_NAME region)
+
+declare_mlir_python_extension(ImexPythonSources.Extension
+  MODULE_NAME _imex_mlir
+  ADD_TO_PARENT ImexPythonSources
+  SOURCES
+    ImexExtension.cpp
+  EMBED_CAPI_LINK_LIBS
+    ImexCAPI
+  PYTHON_BINDINGS_LIBRARY nanobind
+)
+
+
+################################################################################
+# Common CAPI
+################################################################################
+
+add_mlir_python_common_capi_library(ImexPythonCAPI
+  INSTALL_COMPONENT ImexPythonModules
+  INSTALL_DESTINATION python_packages/imex_mlir/_mlir_libs
+  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/imex_mlir/_mlir_libs"
+  RELATIVE_INSTALL_ROOT "../../../.."
+  DECLARED_SOURCES
+    ImexPythonSources
+    MLIRPythonExtension.RegisterEverything
+    MLIRPythonSources.Core
+    MLIRPythonSources.ExecutionEngine
+)
+
+################################################################################
+# Instantiation of all Python modules
+################################################################################
+
+add_mlir_python_modules(ImexPythonModules
+  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/imex_mlir"
+  INSTALL_PREFIX "python_packages/imex_mlir"
+  DECLARED_SOURCES
+    ImexPythonSources
+    MLIRPythonExtension.RegisterEverything
+    MLIRPythonSources
+    MLIRPythonSources.ExecutionEngine
+  COMMON_CAPI_LINK_LIBS
+    ImexPythonCAPI
+  )
diff --git a/python/ImexExtension.cpp b/python/ImexExtension.cpp
@@ -0,0 +1,37 @@
+//===- ImexExtension.cpp - Extension module -------------------------------===//
+//
+// Copyright 2025 Intel Corporation
+// Part of the IMEX 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 "imex-c/Dialects.h"
+#include "imex-c/Passes.h"
+#include "mlir/Bindings/Python/Nanobind.h"
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
+
+namespace nb = nanobind;
+
+NB_MODULE(_imex_mlir, m) {
+  m.doc() = "Intel Extension for MLIR (IMEX) Python binding";
+
+  mlirRegisterAllIMEXPasses();
+
+  //===--------------------------------------------------------------------===//
+  // region dialect
+  //===--------------------------------------------------------------------===//
+  auto regionM = m.def_submodule("region");
+
+  regionM.def(
+      "register_dialect",
+      [](MlirContext context, bool load) {
+        MlirDialectHandle handle = mlirGetDialectHandle__region__();
+        mlirDialectHandleRegisterDialect(handle, context);
+        if (load) {
+          mlirDialectHandleLoadDialect(handle, context);
+        }
+      },
+      nb::arg("context").none() = nb::none(), nb::arg("load") = true);
+}
diff --git a/python/imex_mlir/_mlir_libs/_site_initialize_0.py b/python/imex_mlir/_mlir_libs/_site_initialize_0.py
@@ -0,0 +1,13 @@
+# ===-- _site_initialize_0.py - For site init -----------------*- Python -*-===#
+#
+# Copyright 2025 Intel Corporation
+# Part of the IMEX 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
+#
+# ===-----------------------------------------------------------------------===#
+
+
+def context_init_hook(context):
+    # this will invoke imex binding registrations
+    from imex_mlir.dialects.region import EnvironmentRegionOp  # noqa
diff --git a/python/imex_mlir/dialects/RegionOps.td b/python/imex_mlir/dialects/RegionOps.td
@@ -0,0 +1,15 @@
+//===-- RegionOps.td - Python bindings for region dialect --*- tablegen -*-===//
+//
+// Copyright 2025 Intel Corporation
+// Part of the IMEX 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
+//
+//===---------------------------------------------------------------------===//
+
+#ifndef PYTHON_BINDINGS_REGION_OPS
+#define PYTHON_BINDINGS_REGION_OPS
+
+include "imex/Dialect/Region/IR/RegionOps.td"
+
+#endif
diff --git a/python/imex_mlir/dialects/region.py b/python/imex_mlir/dialects/region.py
@@ -0,0 +1,11 @@
+# ===-- region.py - Region dialect bindings -------------------*- Python -*-===#
+#
+# Copyright 2025 Intel Corporation
+# Part of the IMEX 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
+#
+# ===-----------------------------------------------------------------------===#
+
+from ._region_ops_gen import *
+from .._mlir_libs._imex_mlir.region import *
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
@@ -64,6 +64,12 @@ if(IMEX_ENABLE_L0_RUNTIME OR IMEX_ENABLE_SYCL_RUNTIME)
         )
 endif()
 
+if(IMEX_ENABLE_BINDINGS_PYTHON)
+    list(APPEND IMEX_TEST_DEPENDS
+        ImexPythonModules
+        )
+endif()
+
 # "Gen" is the root of all generated test cases
 add_subdirectory(Gen)
 
diff --git a/test/lit.cfg.py b/test/lit.cfg.py
@@ -72,3 +72,11 @@
 ]
 
 llvm_config.add_tool_substitutions(tools, tool_dirs)
+
+llvm_config.with_environment(
+    "PYTHONPATH",
+    [
+        os.path.join(config.imex_obj_root, "python_packages"),
+    ],
+    append_path=True,
+)
diff --git a/test/lit.site.cfg.py.in b/test/lit.site.cfg.py.in
@@ -41,6 +41,7 @@ config.imex_enable_bf16_tests = @IMEX_ENABLE_BF16_TESTS@
 config.imex_enable_excluded_tests = @IMEX_ENABLE_EXCLUDED_TESTS@
 config.imex_enable_ats_target = @IMEX_ENABLE_ATS_TARGET@
 config.imex_enable_pvc_target = @IMEX_ENABLE_PVC_TARGET@
+config.imex_enable_bindings_python = @IMEX_ENABLE_BINDINGS_PYTHON@
 
 if sys.platform.startswith('win32'):
     config.shlib_prefix = ""
diff --git a/test/python/builtin.py b/test/python/builtin.py
diff --git a/test/python/lit.local.cfg b/test/python/lit.local.cfg
diff --git a/test/python/pass.py b/test/python/pass.py
diff --git a/test/python/region.py b/test/python/region.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`add_subdirectory(Dialect)`
	`2`	`+add_subdirectory(CAPI)`
`2`	`3`	`add_subdirectory(Conversion)`
`3`	`4`	`add_subdirectory(Transforms)`
`4`	`5`	`add_subdirectory(ExecutionEngine)`