4 files changed, 63 insertions, 25 deletions
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index cea5b25..9f5246d 100644
--- 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
diff --git a/mlir/python/mlir/dialects/python_test.py b/mlir/python/mlir/dialects/python_test.py
index 56d3c0f..9380896 100644
--- 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)
diff --git a/mlir/python/mlir/ir.py b/mlir/python/mlir/ir.py
index 7ddc70a..11477d0 100644
--- a/mlir/python/mlir/ir.py
+++ b/mlir/python/mlir/ir.py
@@ -12,7 +12,7 @@ from ._mlir_libs._mlir.ir import _GlobalDebug
 from ._mlir_libs._mlir import (
     register_type_caster,
     register_value_caster,
-    globals,
+    globals as _globals,
 )
 from ._mlir_libs import (
     get_dialect_registry,
@@ -32,17 +32,17 @@ def loc_tracebacks(*, max_depth: int | None = None) -> Iterable[None]:
       max_depth: Maximum number of frames to include in the location.
         If None, the default limit is used.
     """
-    old_enabled = globals.loc_tracebacks_enabled()
-    old_limit = globals.loc_tracebacks_frame_limit()
+    old_enabled = _globals.loc_tracebacks_enabled()
+    old_limit = _globals.loc_tracebacks_frame_limit()
     try:
-        globals.set_loc_tracebacks_frame_limit(max_depth)
+        _globals.set_loc_tracebacks_frame_limit(max_depth)
         if not old_enabled:
-            globals.set_loc_tracebacks_enabled(True)
+            _globals.set_loc_tracebacks_enabled(True)
             yield
     finally:
         if not old_enabled:
-            globals.set_loc_tracebacks_enabled(False)
-        globals.set_loc_tracebacks_frame_limit(old_limit)
+            _globals.set_loc_tracebacks_enabled(False)
+        _globals.set_loc_tracebacks_frame_limit(old_limit)
 
 
 # Convenience decorator for registering user-friendly Attribute builders.
diff --git a/mlir/python/requirements.txt b/mlir/python/requirements.txt
index 5ff9500..abe0925 100644
--- 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"