Add MLIR Python Bindings For Custom Dialects

A guide on how to add Python bindings for custom MLIR dialects using CMake and nanobind/pybind11.

Sun Nov 16 2025

Wed Dec 31 2025

MLIRPython BindingsIntroduction

Add MLIR Python Bindings For Custom Dialects

由于 MLIR 官方的 Python 绑定文档不是很充分, 因此这里介绍的内容不一定是最佳实践, 有的方法是通过阅读源代码和实验得来的, 由于版本不同可能行为有所差异, 请注意验证.

这里演示的Python包结构是为了将 MLIR Python绑定整体作为自定义的包的一个子模块进行构建(这样似乎更常见一点, 我们自己的包的功能是基于MLIR进行开发的, 并不单纯就是个MLIR包装器), 而不是将 MLIR Python绑定的各个子模块直接集成到自己的包中. 如果目标不同, 需要做出一些调整.

前置操作

为了启用 MLIR 的 Python 绑定, 需要在编译 MLIR 时启用 Python 绑定选项. 具体来说, 在配置 LLVM/MLIR 的 CMake 选项时, 需要添加:

-DMLIR_ENABLE_PYTHON_BINDINGS=ON
-DPython3_EXECUTABLE=$(which python3)

建议先创建一个新的 Python 虚拟环境, 例如 conda 或者 venv.

conda create -n <name> python=3.12
# or
python3.12 -m venv <path-to-venv>

激活虚拟环境之后安装一些 MLIR 需要的依赖:

pip install -r /path/to/llvm-project/mlir/python/requirements.txt

然后配置 CMake, 记得加上上面所说的两个选项. 编译完成后 MLIR 的 Python 包默认会安装到
${LLVM_BUILD_DIR}/tools/mlir/python_packages/mlir_core 目录下, 将这个路径添加到终端的 PYTHONPATH 环境变量中即可使用 MLIR 的 Python 绑定:

export PYTHONPATH=${LLVM_BUILD_DIR}/tools/mlir/python_packages/mlir_core:$PYTHONPATH

python3
>>> import mlir
>>> print(mlir.__version__)

关于 MLIR 的核心 Python 绑定, 官网文档写的还是比较详细的. 这里主要介绍怎么添加的自己方言的 Python 绑定.

工作流程

先介绍一下添加绑定的大概流程. 和绑定相关的源文件主要分成两类, 一类是和创建 Python 包相关的文件, 一类是要编译成 Python 扩展模块(实际上是一些动态库)的源文件. 首先第一类文件主要包括:

__init__.py: 用来初始化你自定义的 Python 包, 例如导入子模块等.
MyDialectOps.td: 自定义方言操作的 TableGen 文件, 用来生成操作/枚举的 Python 绑定代码.
MyDialect.py: 将 mlir-tblgen 生成的绑定代码整合到一个 Python 模块中, 并且导入编译出来的扩展模块.
MyDialect.pyi: 自己定义的 Python 包的类型存根文件, 用来给IDE进行索引. 这个可选, 不影响功能正常使用.

第二类文件主要包括:

Module.cpp: 用来定义 Python 扩展模块的初始化函数, 以及注册方言/操作等.
Types.cpp: 提供自定义类型的 Python 绑定代码.
Attributes.cpp: 提供自定义属性的 Python 绑定代码.

Python 绑定本质上是创建一个Python Wrapper, 允许你在 Python 中调用编译好的 C++ 代码. 这些 C++ 代码通常会被编译成动态库(.so 文件)并暴露出必要的符号给 Python 解释器使用.

编写Python Wrapper代码通常是通过 pybind11/nanobind 库实现的. MLIR 基于这些库编写了一些更简单的API用来简化绑定. 新版本的 MLIR 逐渐切换到 nanobind 库, 因此这里使用 nanobind 作为示例.

配置项目`CMakeLists.txt`

为自己的项目启用Python绑定, 需要对自己的项目做一些修改. 创建一个目录专门用来存放上面所提到的第一类文件, 推荐 python/. 按照这个目录结构创建一些文件, 因为其中的子目录会影响到最终生成的Python包的结构. 假设我们的方言叫做 toy.

|-- python
    |-- _mlir_libs
    |   |-- _toy.pyi
    |-- dialects
    |   |-- toy.py
    |   |-- toyOps.td
    |-- __init__.py 
    |-- CMakeLists.txt

这么构建目录结构主要是为了和 MLIR 上游保持一致. 上游编译出来的Python包结构是:

python_package_root
|-- _mlir_libs/
    |-- _mlir<DialectName>-cpython-<platform>.so // compiled extension module
|-- dialects/
    |-- <DialectName>_ops_gen.py // auto-generated by mlir-tblgen
    |-- <DialectName>_enums_gen.py // auto-generated by mlir-tblgen (if any)
    |-- <DialectName>.py // user-defined module to import generated code and extension module
|-- include/
|-- extras/
|-- runtime/
|-- execution_engine.py
|-- ir.py
|-- passmanager.py
|-- rewrite.py

我们提供的自定义文件会按照这个结构被复制进最终生成的Python包中. 例如 toy.py 会被放到 dialects/ 目录下, _toy.pyi 会被放到 _mlir_libs/ 目录下, __init__.py 会被放到包的根目录下.

只要按照这个目录结构, 还可以添加其它自己的文件, 比如自定义的 exception 模块, 只需要把 exception.py 放到 python/ 目录下即可, 它会被放到最终生成的Python包的根目录下.

在 python/CMakeLists.txt 中引入必要的 MLIR 模块.

include(AddMLIRPython)
# 将MLIR上游的Python绑定作为我们的Python包的一个子模块
add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=toy._mlir.")

# ========== Declare the Python package ==========
declare_mlir_python_sources(ToyMLIRPythonSources)
declare_mlir_python_sources(ToyMLIRPythonExtension)

declare_mlir_python_sources(ToyMLIRPythonSources.Dialect
    ADD_TO_PARENT ToyMLIRPythonSources
)

接着添加Python包相关的文件和扩展模块的源文件.

# ========== Add source files ==========
declare_mlir_dialect_python_bindings(
    ADD_TO_PARENT ToyMLIRPythonSources.Dialect
    ROOT_DIR ${CMAKE_CURRENT_SOURCE_DIR}
    TD_FILE dialects/toyOps.td
    SOURCES
        _mlir_libs/_toy.pyi
        dialects/toy.py
        __init__.py
    DIALECT_NAME toy
)

set(TOY_PYTHON_SOURCE_DIR "${PROJECT_SOURCE_DIR}/lib/Bindings/Python")

# the extension will be compiled into _toy.cpython-<platform>.so
declare_mlir_python_extension(
    MODULE_NAME _toy
    ADD_TO_PARENT ToyMLIRPythonExtension
    ROOT_DIR "/"
    PYTHON_BINDINGS_LIBRARY nanobind
    SOURCES
        ${TOY_PYTHON_SOURCE_DIR}/Module.h # headers also need to be listed here
        ${TOY_PYTHON_SOURCE_DIR}/Types.cpp
        ${TOY_PYTHON_SOURCE_DIR}/Module.cpp
    PRIVATE_LINK_LIBS
        MLIRPass
        MLIRToy
        LLVMSupport
    EMBED_CAPI_LINK_LIBS
        MLIRCAPIIR
        MLIRCAPIDebug
        MLIRCAPIToy
)

接着创建最终的Python包.

# ========== Create the Python package ==========
set(_source_components
    ToyMLIRPythonSources
    ToyMLIRPythonExtension
    MLIRPythonSources # introduce MLIR core bindings
    MLIRPythonExtension.RegisterEverything
)

# the package will be installed under
# ${PROJECT_BINARY_DIR}/tools/toy/_mlir
set(TOY_MLIR_PYTHON_PACKAGES_DIR "${PROJECT_BINARY_DIR}/tools/toy")

add_mlir_python_common_capi_library(ToyMLIRAggregateCAPI
    INSTALL_COMPONENT ToyMLIRPythonModules
    INSTALL_DESTINATION _mlir
    OUTPUT_DIRECTORY "${TOY_MLIR_PYTHON_PACKAGES_DIR}/_mlir"
    RELATIVE_INSTALL_ROOT "../../"
    DECLARED_HEADERS
        MLIRPythonCAPI.HeaderSources
    DECLARED_SOURCES
        ${_source_components}
)

# link the dialect library
# a fix under MLIR 21.1.2 due to an unknown issue
target_link_libraries(ToyMLIRAggregateCAPI
    PRIVATE
        MLIRToy 
)

add_mlir_python_package(ToyMLIRPythonModules
    ROOT_PREFIX "${TOY_MLIR_PYTHON_PACKAGES_DIR}/_mlir"
    INSTALL_PREFIX _mlir
    DECLARED_SOURCES
        ${_source_components}
    COMMON_CAPI_LINK_LIBS
        ToyMLIRAggregateCAPI
)

这样配置之后生成的完整的Python包路径就会在 ${PROJECT_BINARY_DIR}/tools/toy/_mlir 下. 将这个路径添加到 PYTHONPATH 环境变量中就能在Python中导入使用.

编写Python包文件

在 python/__init__.py 中初始化包, 基本上这样就行了. 总之按照编译出来的包的结构来写就行.

这里 MLIR 做的不太好的地方就是如果你不知道最终生成的包的结构, 很难预先写出正确的 __init__.py 文件.

from .ir import *

这是 toy.py 文件的内容, 用来导入 mlir-tblgen 生成的绑定代码, 并且导入编译出来的扩展模块.

from .toy_ops_gen import *
from .toy_enums_gen import * # if you have enums defined
from .._mlir_libs._toy import * # import the compiled extension module

这是 toyOps.td 文件的内容, 用来定义方言操作. 因为之前创建方言的时候已经写过操作的ODS文件了, 所以这里简单包含一下之前的文件即可. 比如

#ifndef TOY_PYTHON_OPS_TD
#define TOY_PYTHON_OPS_TD

include "toy/IR/ToyOps.td"

#endif // TOY_PYTHON_OPS_TD

其它的文件可以自行添加, 然后记得在 CMakeLists.txt 中添加进去即可.

编写扩展模块代码

扩展模块的代码主要是编写一些 Python Wrapper, 在进行这一步之前请确保已经阅读过上一章, 关于如何创建CAPI. 为了保证跨语言调用的稳定性, 一般是不会用C++接口直接暴露给其它语言的(主要是因为ABI问题还有Name Mangling问题), 而是通过C接口暴露出来. 因此在编写Python绑定代码时, 主要是调用之前创建的CAPI.

按照约定, 扩展模块的源代码一般放在 lib/Bindings/Python 目录下, 和上面的 CMakeLists.txt 保持一致, 这里添加三个文件: Module.cpp, Types.cpp, Module.h.

Module.cpp 作为扩展模块的入口, 用来注册自定义方言, 调用绑定类型和属性的注册代码等. 这里是一个能实现核心功能的最小示例:

#include "Module.h"
#include "mlir/IR/DialectRegistry.h"
#include "mlir/CAPI/IR.h"
#include "llvm/Support/Signals.h"
#include "llvm-c/ErrorHandling.h"
#include "toy-c/Dialect/Registration.h"
#include "toy-c/Dialect/Toy.h"

namespace nb = nanobind;
using namespace mlir;

NB_MODULE(_toy, m) {
    m.doc() = "Toy MLIR Python Native Extension";
    llvm::sys::PrintStackTraceOnErrorSignal("");
    LLVMEnablePrettyStackTrace();

    // register dialect
    m.def("register_dialect", [](MlirContext context) {
        mlirDialectHandle toy = mlirGetDialectHandle__toy__();
        DialectRegistry registry;
        unwrap(context)->appendDialectRegistry(registry);
        mlirDialectHandleRegisterDialect(toy, context);
        mlirDialectHandleLoadDialect(toy, context);
    }, nb::arg("context") = nb::none())

    // define other functions if needed
    // m.def ...

    // populate types and attributes
    python::populateToyIRTypes(m);
    // populateToyIRAttributes(m); same as above if you have attributes
    toyMlirRegisterAllPasses(); // defined in CAPI
}

这样就定义了一个名为 _toy 的扩展模块, 并在这个模块中提供了一个函数 register_dialect 用来把方言注册到给定的上下文中, 例如可以这样使用:

# _toy extension has been introduced in toy.py before
from nexus._mlir.dialects import toy
from nexus._mlir.ir import Context

with Context() as ctx:
    toy.register_dialect(ctx)

在 Module.h 中声明一些公共的函数和头文件

#ifndef TOY_PYTHON_MODULE_H
#define TOY_PYTHON_MODULE_H

#include "nanobind/nanobind.h"
#include "mlir/Bindings/Python/NanobindAdaptors.h"

namespace mlir::python {
void populateToyIRTypes(nanobind::module_ &m);
// void populateToyIRAttributes(nanobind::module_ &m); same as above
}
#endif // TOY_PYTHON_MODULE_H

在 Types.cpp 中实现自定义类型的绑定代码. 它们基本都遵循相同的模式, 这里给出一个示例:

#include "Module.h"
#include "toy-c/Dialect/ToyTypes.h"

namespace nb = nanobind;
using namespace mlir;
using namespace mlir::python;
using namespace mlir::python::nanobind_adaptors;

void python::populateToyIRTypes(nb::module_ &m) {
    mlir_type_subclass(m, "FixedType", toyMlirIsAFixedType, toyMlirFixedTypeGetTypeID)
        .def_staticmethod("get", [](size_t width, size_t frac, MlirContext context) {
            return toyMlirFixedTypeGet(context, width, frac);
        }, nb::arg("width"), nb::arg("frac"), nb::arg("context") = nb::none(), "Create a FixedType instance")
        .def_property_readonly("width", [](MlirType type) {
            return toyMlirFixedTypeGetWidth(type);
        }, "Get the width of the FixedType")
        .def_property_readonly("frac", [](MlirType type) {
            return toyMlirFixedTypeGetFrac(type);
        }, "Get the frac of the FixedType");
}

这样就定义了一个名为 FixedType 的定点数类型. 里面涉及到的所有CAPI函数都是在之前创建CAPI时定义好的.

mlir_type_subclass 类的构造函数接受几个参数, 第一个是 nanobind::module_ 对象的引用, 第二个是要定义的类型名称, 第三个是一个函数指针, 用来判断某个 MlirType 是否是该类型的实例, 第四个是一个函数指针, 用来获取该类型的类型ID. 这些函数都是在CAPI中定义好的.

这个类提供一些API用来定义Python类型的方法,

def_staticmethod 用来定义静态方法
def_property_readonly 用来定义只读属性,
def_classmethod 用来定义类方法
def 定义实例方法

这些方法都接受几个参数, 第一个是方法名称, 第二个是一个 lambda 函数, 用来实现方法的功能, 后面可以接一些可选参数, 例如参数名称和文档字符串等.

nb::arg("") 用来指定参数名称, nb::arg("") = nb::none() 用来指定参数的默认值为 None, 也可以指定其它默认值, 例如 nb::int_(1), nb::float_(0.0) 等. nb:: 命名空间里定义了几乎所有Python内置类型.

对于使用 pybind11 的用户来讲, 将命名空间从 nanobind 换成 pybind11, 其余的代码几乎完全一样.

另外, MLIR自动会给你定义一些类的方法(只要你创建 mlir_type_subclass), 包括 isinstance 方法和 __repr__ 方法. 如果提供了获取类型ID的函数, 则会生成一个 get_static_typeid 静态方法.

对于属性的绑定代码, 要使用的类型是 mlir_attribute_subclass, 其余过程和上面完全一样.

void python::populateToyIRAttributes(nb::module_ &m) {
    mlir_attribute_subclass(...)
        .def_staticmethod(...)
        .def_property_readonly(...)
        .def(...)
}

最后生成出来的Python Wrapper类似于这样的效果(这只是效果示例, 执行功能的机器码还是被打包进动态库里面了):

from nexus._mlir.ir import Type

class FixedType(Type):
    @staticmethod
    def get_static_typeid() -> int: ...
    @staticmethod
    def isinstance(type: Type) -> bool: ...
    @staticmethod
    def get(width: int, frac: int, context: Context = None) -> FixedType: ...
    @property
    def width(self) -> int: ...
    @property
    def frac(self) -> int: ...

不过这些绑定代码最终都会被编译并打包到扩展模块里面, 也就是动态库 _toy.cpython-<platform>.so 里面. IDE并不能直接读取动态库中的代码获得各种类型信息, 因此建议为自定义的Python包编写类型存根文件(.pyi 文件), 以便IDE能够正确索引类型信息和方法签名. 例如上面就是一个为 FixedType 编写的类型存根文件片段. 将存根文件放到 python/_mlir_libs/_toy.pyi 就可以一起被打包进最终的Python包中并发挥作用.

调试技巧

有的时候可能在导入自定义包的时候会遇到问题, 这里记录一些方法来帮助调试.

`ImportError`说了和没说一样

有的时候在导入自定义包的时候会遇到类似下面的错误:

ImportError: Error loading extension module _mlir_libs._toy

或者类似的, 总之是就给你报错 ImportError 但是什么有用的信息都不告诉你. 这种时候是因为调用的动态库中的C++函数抛异常了, 从Python的角度看它也不知道抛了什么异常, 只知道导入失败了.

这种时候可以请出GDB大法, 终端中打开GDB, 然后运行Python解释器, 在GDB中设置断点捕获C++异常抛出的地方, 例如:

gdb --args python
(gdb) catch throw
(gdb) run /error/producing/script.py
... some info ...
catchpoint 1 (throw)
(gdb) bt
# ... backtrace info ...

这样可以看到异常抛出的堆栈信息, 进而定位问题. 前提是你在构建C++部分的时候开启了调试信息, 否则堆栈信息可能不完整.

Back to Archives

Add MLIR Python Bindings For Custom Dialects

A guide on how to add Python bindings for custom MLIR dialects using CMake and nanobind/pybind11.

Sun Nov 16 2025

Wed Dec 31 2025

MLIRPython BindingsIntroduction

Add MLIR Python Bindings For Custom Dialects

前置操作

为了启用 MLIR 的 Python 绑定, 需要在编译 MLIR 时启用 Python 绑定选项. 具体来说, 在配置 LLVM/MLIR 的 CMake 选项时, 需要添加:

-DMLIR_ENABLE_PYTHON_BINDINGS=ON
-DPython3_EXECUTABLE=$(which python3)

建议先创建一个新的 Python 虚拟环境, 例如 conda 或者 venv.

conda create -n <name> python=3.12
# or
python3.12 -m venv <path-to-venv>

激活虚拟环境之后安装一些 MLIR 需要的依赖:

pip install -r /path/to/llvm-project/mlir/python/requirements.txt

export PYTHONPATH=${LLVM_BUILD_DIR}/tools/mlir/python_packages/mlir_core:$PYTHONPATH

python3
>>> import mlir
>>> print(mlir.__version__)

关于 MLIR 的核心 Python 绑定, 官网文档写的还是比较详细的. 这里主要介绍怎么添加的自己方言的 Python 绑定.

工作流程

__init__.py: 用来初始化你自定义的 Python 包, 例如导入子模块等.
MyDialectOps.td: 自定义方言操作的 TableGen 文件, 用来生成操作/枚举的 Python 绑定代码.
MyDialect.py: 将 mlir-tblgen 生成的绑定代码整合到一个 Python 模块中, 并且导入编译出来的扩展模块.
MyDialect.pyi: 自己定义的 Python 包的类型存根文件, 用来给IDE进行索引. 这个可选, 不影响功能正常使用.

第二类文件主要包括:

Module.cpp: 用来定义 Python 扩展模块的初始化函数, 以及注册方言/操作等.
Types.cpp: 提供自定义类型的 Python 绑定代码.
Attributes.cpp: 提供自定义属性的 Python 绑定代码.

配置项目`CMakeLists.txt`

|-- python
    |-- _mlir_libs
    |   |-- _toy.pyi
    |-- dialects
    |   |-- toy.py
    |   |-- toyOps.td
    |-- __init__.py 
    |-- CMakeLists.txt

这么构建目录结构主要是为了和 MLIR 上游保持一致. 上游编译出来的Python包结构是:

python_package_root
|-- _mlir_libs/
    |-- _mlir<DialectName>-cpython-<platform>.so // compiled extension module
|-- dialects/
    |-- <DialectName>_ops_gen.py // auto-generated by mlir-tblgen
    |-- <DialectName>_enums_gen.py // auto-generated by mlir-tblgen (if any)
    |-- <DialectName>.py // user-defined module to import generated code and extension module
|-- include/
|-- extras/
|-- runtime/
|-- execution_engine.py
|-- ir.py
|-- passmanager.py
|-- rewrite.py

在 python/CMakeLists.txt 中引入必要的 MLIR 模块.

include(AddMLIRPython)
# 将MLIR上游的Python绑定作为我们的Python包的一个子模块
add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=toy._mlir.")

# ========== Declare the Python package ==========
declare_mlir_python_sources(ToyMLIRPythonSources)
declare_mlir_python_sources(ToyMLIRPythonExtension)

declare_mlir_python_sources(ToyMLIRPythonSources.Dialect
    ADD_TO_PARENT ToyMLIRPythonSources
)

接着添加Python包相关的文件和扩展模块的源文件.

# ========== Add source files ==========
declare_mlir_dialect_python_bindings(
    ADD_TO_PARENT ToyMLIRPythonSources.Dialect
    ROOT_DIR ${CMAKE_CURRENT_SOURCE_DIR}
    TD_FILE dialects/toyOps.td
    SOURCES
        _mlir_libs/_toy.pyi
        dialects/toy.py
        __init__.py
    DIALECT_NAME toy
)

set(TOY_PYTHON_SOURCE_DIR "${PROJECT_SOURCE_DIR}/lib/Bindings/Python")

# the extension will be compiled into _toy.cpython-<platform>.so
declare_mlir_python_extension(
    MODULE_NAME _toy
    ADD_TO_PARENT ToyMLIRPythonExtension
    ROOT_DIR "/"
    PYTHON_BINDINGS_LIBRARY nanobind
    SOURCES
        ${TOY_PYTHON_SOURCE_DIR}/Module.h # headers also need to be listed here
        ${TOY_PYTHON_SOURCE_DIR}/Types.cpp
        ${TOY_PYTHON_SOURCE_DIR}/Module.cpp
    PRIVATE_LINK_LIBS
        MLIRPass
        MLIRToy
        LLVMSupport
    EMBED_CAPI_LINK_LIBS
        MLIRCAPIIR
        MLIRCAPIDebug
        MLIRCAPIToy
)

接着创建最终的Python包.

# ========== Create the Python package ==========
set(_source_components
    ToyMLIRPythonSources
    ToyMLIRPythonExtension
    MLIRPythonSources # introduce MLIR core bindings
    MLIRPythonExtension.RegisterEverything
)

# the package will be installed under
# ${PROJECT_BINARY_DIR}/tools/toy/_mlir
set(TOY_MLIR_PYTHON_PACKAGES_DIR "${PROJECT_BINARY_DIR}/tools/toy")

add_mlir_python_common_capi_library(ToyMLIRAggregateCAPI
    INSTALL_COMPONENT ToyMLIRPythonModules
    INSTALL_DESTINATION _mlir
    OUTPUT_DIRECTORY "${TOY_MLIR_PYTHON_PACKAGES_DIR}/_mlir"
    RELATIVE_INSTALL_ROOT "../../"
    DECLARED_HEADERS
        MLIRPythonCAPI.HeaderSources
    DECLARED_SOURCES
        ${_source_components}
)

# link the dialect library
# a fix under MLIR 21.1.2 due to an unknown issue
target_link_libraries(ToyMLIRAggregateCAPI
    PRIVATE
        MLIRToy 
)

add_mlir_python_package(ToyMLIRPythonModules
    ROOT_PREFIX "${TOY_MLIR_PYTHON_PACKAGES_DIR}/_mlir"
    INSTALL_PREFIX _mlir
    DECLARED_SOURCES
        ${_source_components}
    COMMON_CAPI_LINK_LIBS
        ToyMLIRAggregateCAPI
)

这样配置之后生成的完整的Python包路径就会在 ${PROJECT_BINARY_DIR}/tools/toy/_mlir 下. 将这个路径添加到 PYTHONPATH 环境变量中就能在Python中导入使用.

编写Python包文件

在 python/__init__.py 中初始化包, 基本上这样就行了. 总之按照编译出来的包的结构来写就行.

这里 MLIR 做的不太好的地方就是如果你不知道最终生成的包的结构, 很难预先写出正确的 __init__.py 文件.

from .ir import *

这是 toy.py 文件的内容, 用来导入 mlir-tblgen 生成的绑定代码, 并且导入编译出来的扩展模块.

from .toy_ops_gen import *
from .toy_enums_gen import * # if you have enums defined
from .._mlir_libs._toy import * # import the compiled extension module

这是 toyOps.td 文件的内容, 用来定义方言操作. 因为之前创建方言的时候已经写过操作的ODS文件了, 所以这里简单包含一下之前的文件即可. 比如

#ifndef TOY_PYTHON_OPS_TD
#define TOY_PYTHON_OPS_TD

include "toy/IR/ToyOps.td"

#endif // TOY_PYTHON_OPS_TD

其它的文件可以自行添加, 然后记得在 CMakeLists.txt 中添加进去即可.

编写扩展模块代码

按照约定, 扩展模块的源代码一般放在 lib/Bindings/Python 目录下, 和上面的 CMakeLists.txt 保持一致, 这里添加三个文件: Module.cpp, Types.cpp, Module.h.

Module.cpp 作为扩展模块的入口, 用来注册自定义方言, 调用绑定类型和属性的注册代码等. 这里是一个能实现核心功能的最小示例:

#include "Module.h"
#include "mlir/IR/DialectRegistry.h"
#include "mlir/CAPI/IR.h"
#include "llvm/Support/Signals.h"
#include "llvm-c/ErrorHandling.h"
#include "toy-c/Dialect/Registration.h"
#include "toy-c/Dialect/Toy.h"

namespace nb = nanobind;
using namespace mlir;

NB_MODULE(_toy, m) {
    m.doc() = "Toy MLIR Python Native Extension";
    llvm::sys::PrintStackTraceOnErrorSignal("");
    LLVMEnablePrettyStackTrace();

    // register dialect
    m.def("register_dialect", [](MlirContext context) {
        mlirDialectHandle toy = mlirGetDialectHandle__toy__();
        DialectRegistry registry;
        unwrap(context)->appendDialectRegistry(registry);
        mlirDialectHandleRegisterDialect(toy, context);
        mlirDialectHandleLoadDialect(toy, context);
    }, nb::arg("context") = nb::none())

    // define other functions if needed
    // m.def ...

    // populate types and attributes
    python::populateToyIRTypes(m);
    // populateToyIRAttributes(m); same as above if you have attributes
    toyMlirRegisterAllPasses(); // defined in CAPI
}

这样就定义了一个名为 _toy 的扩展模块, 并在这个模块中提供了一个函数 register_dialect 用来把方言注册到给定的上下文中, 例如可以这样使用:

# _toy extension has been introduced in toy.py before
from nexus._mlir.dialects import toy
from nexus._mlir.ir import Context

with Context() as ctx:
    toy.register_dialect(ctx)

在 Module.h 中声明一些公共的函数和头文件

#ifndef TOY_PYTHON_MODULE_H
#define TOY_PYTHON_MODULE_H

#include "nanobind/nanobind.h"
#include "mlir/Bindings/Python/NanobindAdaptors.h"

namespace mlir::python {
void populateToyIRTypes(nanobind::module_ &m);
// void populateToyIRAttributes(nanobind::module_ &m); same as above
}
#endif // TOY_PYTHON_MODULE_H

在 Types.cpp 中实现自定义类型的绑定代码. 它们基本都遵循相同的模式, 这里给出一个示例:

#include "Module.h"
#include "toy-c/Dialect/ToyTypes.h"

namespace nb = nanobind;
using namespace mlir;
using namespace mlir::python;
using namespace mlir::python::nanobind_adaptors;

void python::populateToyIRTypes(nb::module_ &m) {
    mlir_type_subclass(m, "FixedType", toyMlirIsAFixedType, toyMlirFixedTypeGetTypeID)
        .def_staticmethod("get", [](size_t width, size_t frac, MlirContext context) {
            return toyMlirFixedTypeGet(context, width, frac);
        }, nb::arg("width"), nb::arg("frac"), nb::arg("context") = nb::none(), "Create a FixedType instance")
        .def_property_readonly("width", [](MlirType type) {
            return toyMlirFixedTypeGetWidth(type);
        }, "Get the width of the FixedType")
        .def_property_readonly("frac", [](MlirType type) {
            return toyMlirFixedTypeGetFrac(type);
        }, "Get the frac of the FixedType");
}

这样就定义了一个名为 FixedType 的定点数类型. 里面涉及到的所有CAPI函数都是在之前创建CAPI时定义好的.

这个类提供一些API用来定义Python类型的方法,

def_staticmethod 用来定义静态方法
def_property_readonly 用来定义只读属性,
def_classmethod 用来定义类方法
def 定义实例方法

对于使用 pybind11 的用户来讲, 将命名空间从 nanobind 换成 pybind11, 其余的代码几乎完全一样.

对于属性的绑定代码, 要使用的类型是 mlir_attribute_subclass, 其余过程和上面完全一样.

void python::populateToyIRAttributes(nb::module_ &m) {
    mlir_attribute_subclass(...)
        .def_staticmethod(...)
        .def_property_readonly(...)
        .def(...)
}

最后生成出来的Python Wrapper类似于这样的效果(这只是效果示例, 执行功能的机器码还是被打包进动态库里面了):

from nexus._mlir.ir import Type

class FixedType(Type):
    @staticmethod
    def get_static_typeid() -> int: ...
    @staticmethod
    def isinstance(type: Type) -> bool: ...
    @staticmethod
    def get(width: int, frac: int, context: Context = None) -> FixedType: ...
    @property
    def width(self) -> int: ...
    @property
    def frac(self) -> int: ...

调试技巧

有的时候可能在导入自定义包的时候会遇到问题, 这里记录一些方法来帮助调试.

`ImportError`说了和没说一样

有的时候在导入自定义包的时候会遇到类似下面的错误:

ImportError: Error loading extension module _mlir_libs._toy

这种时候可以请出GDB大法, 终端中打开GDB, 然后运行Python解释器, 在GDB中设置断点捕获C++异常抛出的地方, 例如:

gdb --args python
(gdb) catch throw
(gdb) run /error/producing/script.py
... some info ...
catchpoint 1 (throw)
(gdb) bt
# ... backtrace info ...

这样可以看到异常抛出的堆栈信息, 进而定位问题. 前提是你在构建C++部分的时候开启了调试信息, 否则堆栈信息可能不完整.