Revert "[MLIR][Python] use FetchContent_Declare for nanobind and remove pybind (#161230)" (#162309)

This reverts commit 84a214856a. This gives us more time to work out the alternative and also people to migrate
2026-01-16 05:32:28 +08:00 · 2025-10-07 12:30:10 -04:00
parent 278a99e8e9
commit 93097b2d47
18 changed files with 1082 additions and 121 deletions
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -440,11 +440,11 @@ declare_mlir_dialect_python_bindings(
  DIALECT_NAME smt)

 declare_mlir_dialect_python_bindings(
-  ADD_TO_PARENT MLIRPythonSources.Dialects
-  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
-  TD_FILE dialects/SPIRVOps.td
-  SOURCES dialects/spirv.py
-  DIALECT_NAME spirv)
+    ADD_TO_PARENT MLIRPythonSources.Dialects
+    ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+    TD_FILE dialects/SPIRVOps.td
+    SOURCES dialects/spirv.py
+    DIALECT_NAME spirv)

 declare_mlir_dialect_python_bindings(
  ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -501,6 +501,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
  MODULE_NAME _mlir
  ADD_TO_PARENT MLIRPythonSources.Core
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    MainModule.cpp
    IRAffine.cpp
@@ -539,6 +540,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
 declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
  MODULE_NAME _mlirRegisterEverything
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    RegisterEverything.cpp
  PRIVATE_LINK_LIBS
@@ -549,10 +551,11 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
    MLIRCAPIRegisterEverything
 )

-declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
  MODULE_NAME _mlirDialectsLinalg
  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectLinalg.cpp
  PRIVATE_LINK_LIBS
@@ -562,10 +565,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Nanobind
    MLIRCAPILinalg
 )

-declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
  MODULE_NAME _mlirDialectsGPU
  ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectGPU.cpp
  PRIVATE_LINK_LIBS
@@ -575,10 +579,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Nanobind
    MLIRCAPIGPU
 )

-declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
  MODULE_NAME _mlirDialectsLLVM
  ADD_TO_PARENT MLIRPythonSources.Dialects.llvm
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectLLVM.cpp
  PRIVATE_LINK_LIBS
@@ -588,10 +593,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Nanobind
    MLIRCAPILLVM
 )

-declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
  MODULE_NAME _mlirDialectsQuant
  ADD_TO_PARENT MLIRPythonSources.Dialects.quant
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectQuant.cpp
  PRIVATE_LINK_LIBS
@@ -601,10 +607,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Nanobind
    MLIRCAPIQuant
 )

-declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
  MODULE_NAME _mlirDialectsNVGPU
  ADD_TO_PARENT MLIRPythonSources.Dialects.nvgpu
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectNVGPU.cpp
  PRIVATE_LINK_LIBS
@@ -614,10 +621,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Nanobind
    MLIRCAPINVGPU
 )

-declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
  MODULE_NAME _mlirDialectsPDL
  ADD_TO_PARENT MLIRPythonSources.Dialects.pdl
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectPDL.cpp
  PRIVATE_LINK_LIBS
@@ -627,10 +635,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Nanobind
    MLIRCAPIPDL
 )

-declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
  MODULE_NAME _mlirDialectsSparseTensor
  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectSparseTensor.cpp
  PRIVATE_LINK_LIBS
@@ -640,10 +649,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Nanobind
    MLIRCAPISparseTensor
 )

-declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
  MODULE_NAME _mlirDialectsTransform
  ADD_TO_PARENT MLIRPythonSources.Dialects.transform
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectTransform.cpp
  PRIVATE_LINK_LIBS
@@ -653,10 +663,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Nanobind
    MLIRCAPITransformDialect
 )

-declare_mlir_python_extension(MLIRPythonExtension.Dialects.IRDL.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.IRDL.Pybind
  MODULE_NAME _mlirDialectsIRDL
  ADD_TO_PARENT MLIRPythonSources.Dialects.irdl
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectIRDL.cpp
  PRIVATE_LINK_LIBS
@@ -670,6 +681,7 @@ declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
  MODULE_NAME _mlirAsyncPasses
  ADD_TO_PARENT MLIRPythonSources.Dialects.async
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    AsyncPasses.cpp
  PRIVATE_LINK_LIBS
@@ -683,6 +695,7 @@ if(MLIR_ENABLE_EXECUTION_ENGINE)
    MODULE_NAME _mlirExecutionEngine
    ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
    ROOT_DIR "${PYTHON_SOURCE_DIR}"
+    PYTHON_BINDINGS_LIBRARY nanobind
    SOURCES
      ExecutionEngineModule.cpp
    PRIVATE_LINK_LIBS
@@ -696,6 +709,7 @@ declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
  MODULE_NAME _mlirGPUPasses
  ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    GPUPasses.cpp
  PRIVATE_LINK_LIBS
@@ -708,6 +722,7 @@ declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
  MODULE_NAME _mlirLinalgPasses
  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    LinalgPasses.cpp
  PRIVATE_LINK_LIBS
@@ -716,10 +731,11 @@ declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
    MLIRCAPILinalg
 )

-declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
  MODULE_NAME _mlirDialectsSMT
  ADD_TO_PARENT MLIRPythonSources.Dialects.smt
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectSMT.cpp
    # Headers must be included explicitly so they are installed.
@@ -736,6 +752,7 @@ declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
  MODULE_NAME _mlirSparseTensorPasses
  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    SparseTensorPasses.cpp
  PRIVATE_LINK_LIBS
@@ -748,6 +765,7 @@ declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
  MODULE_NAME _mlirTransformInterpreter
  ADD_TO_PARENT MLIRPythonSources.Dialects.transform
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    TransformInterpreter.cpp
  PRIVATE_LINK_LIBS
@@ -789,10 +807,23 @@ if(MLIR_INCLUDE_TESTS)
    ADD_TO_PARENT MLIRPythonTestSources.Dialects.PythonTest
    SOURCES "dialects/_python_test_ops_gen.py")

+  declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtensionPybind11
+    MODULE_NAME _mlirPythonTestPybind11
+    ADD_TO_PARENT MLIRPythonTestSources.Dialects
+    ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
+    PYTHON_BINDINGS_LIBRARY pybind11
+    SOURCES
+      PythonTestModulePybind11.cpp
+    PRIVATE_LINK_LIBS
+      LLVMSupport
+    EMBED_CAPI_LINK_LIBS
+      MLIRCAPIPythonTestDialect
+  )
  declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtensionNanobind
    MODULE_NAME _mlirPythonTestNanobind
    ADD_TO_PARENT MLIRPythonTestSources.Dialects
    ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
+    PYTHON_BINDINGS_LIBRARY nanobind
    SOURCES
      PythonTestModuleNanobind.cpp
    PRIVATE_LINK_LIBS
--- a/mlir/python/mlir/dialects/python_test.py
+++ b/mlir/python/mlir/dialects/python_test.py
@@ -5,7 +5,12 @@
 from ._python_test_ops_gen import *


-def register_python_test_dialect(registry):
-    from .._mlir_libs import _mlirPythonTestNanobind
+def register_python_test_dialect(registry, use_nanobind):
+    if use_nanobind:
+        from .._mlir_libs import _mlirPythonTestNanobind

-    _mlirPythonTestNanobind.register_dialect(registry)
+        _mlirPythonTestNanobind.register_dialect(registry)
+    else:
+        from .._mlir_libs import _mlirPythonTestPybind11
+
+        _mlirPythonTestPybind11.register_dialect(registry)
--- a/mlir/python/requirements.txt
+++ b/mlir/python/requirements.txt
@@ -1,4 +1,6 @@
+nanobind>=2.9, <3.0
 numpy>=1.19.5, <=2.1.2
+pybind11>=2.10.0, <=2.13.6
 PyYAML>=5.4.0, <=6.0.1
 ml_dtypes>=0.1.0, <=0.6.0; python_version<"3.13"   # provides several NumPy dtype extensions, including the bf16
 ml_dtypes>=0.5.0, <=0.6.0; python_version>="3.13"