[Mlir-commits] [mlir] [mlir][python] fix symbol resolution on MacOS (PR #174057)

[Maksim Levental]

https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/174057

>From ce95ec59d592bc88f76731302f2fb21d139e585f Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Tue, 30 Dec 2025 22:09:33 -0800
Subject: [PATCH] [mlir][python] fix flatnamespace

---
 mlir/cmake/modules/AddMLIRPython.cmake        |  6 +++++
 mlir/examples/standalone/CMakeLists.txt       |  4 +++-
 mlir/examples/standalone/test/lit.cfg.py      | 12 +++++-----
 .../standalone/test/python/smoketest.py       | 22 +++++++++++++++++--
 mlir/test/Examples/standalone/test.toy        |  2 ++
 mlir/test/Examples/standalone/test.wheel.toy  |  8 +++++--
 6 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index ca90151e76268..8c301faf0941a 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -766,6 +766,12 @@ function(add_mlir_python_extension libname extname)
     FREE_THREADED
     ${ARG_SOURCES}
   )
+  if(APPLE)
+    # In llvm/cmake/modules/HandleLLVMOptions.cmake:268 we set -Wl,-flat_namespace which breaks
+    # the default name spacing on MacOS and causes "cross-wired" symbol resolution when multiple
+    # bindings packages are loaded.
+    target_link_options(${libname} PRIVATE "LINKER:-twolevel_namespace")
+  endif()
 
   if (NOT MLIR_DISABLE_CONFIGURE_PYTHON_DEV_PACKAGES
       AND (LLVM_COMPILER_IS_GCC_COMPATIBLE OR CLANG_CL))
diff --git a/mlir/examples/standalone/CMakeLists.txt b/mlir/examples/standalone/CMakeLists.txt
index c6c49fde12d2e..955c9ec7a7b4c 100644
--- a/mlir/examples/standalone/CMakeLists.txt
+++ b/mlir/examples/standalone/CMakeLists.txt
@@ -71,7 +71,9 @@ if(MLIR_ENABLE_BINDINGS_PYTHON)
   endif()
   add_subdirectory(python)
 endif()
-add_subdirectory(test)
+if(MLIR_INCLUDE_TESTS)
+  add_subdirectory(test)
+endif()
 add_subdirectory(standalone-opt)
 if(NOT WIN32)
   add_subdirectory(standalone-plugin)
diff --git a/mlir/examples/standalone/test/lit.cfg.py b/mlir/examples/standalone/test/lit.cfg.py
index e27dddd7fb0b9..89cdd6889a1f2 100644
--- a/mlir/examples/standalone/test/lit.cfg.py
+++ b/mlir/examples/standalone/test/lit.cfg.py
@@ -61,10 +61,8 @@
 
 llvm_config.add_tool_substitutions(tools, tool_dirs)
 
-llvm_config.with_environment(
-    "PYTHONPATH",
-    [
-        os.path.join(config.mlir_obj_dir, "python_packages", "standalone"),
-    ],
-    append_path=True,
-)
+python_path = [os.path.join(config.mlir_obj_dir, "python_packages", "standalone")]
+if "PYTHONPATH" in os.environ:
+    python_path += [os.environ["PYTHONPATH"]]
+
+llvm_config.with_environment("PYTHONPATH", python_path, append_path=True)
diff --git a/mlir/examples/standalone/test/python/smoketest.py b/mlir/examples/standalone/test/python/smoketest.py
index f8819841fac45..7526feb7fcd8a 100644
--- a/mlir/examples/standalone/test/python/smoketest.py
+++ b/mlir/examples/standalone/test/python/smoketest.py
@@ -1,4 +1,5 @@
-# RUN: %python %s nanobind | FileCheck %s
+# RUN: %python %s 2>&1 | FileCheck %s
+import sys
 
 from mlir_standalone.ir import *
 from mlir_standalone.dialects import standalone_nanobind as standalone_d
@@ -13,4 +14,21 @@
     )
     # CHECK: %[[C:.*]] = arith.constant 2 : i32
     # CHECK: standalone.foo %[[C]] : i32
-    print(str(module))
+    print(str(module), file=sys.stderr)
+
+# CHECK-NOT: RuntimeWarning: nanobind: type '{{.*}}' was already registered!
+
+from mlir.ir import *
+from mlir.dialects import (
+    amdgpu,
+    gpu,
+    irdl,
+    linalg,
+    llvm,
+    nvgpu,
+    pdl,
+    quant,
+    smt,
+    sparse_tensor,
+    transform,
+)
diff --git a/mlir/test/Examples/standalone/test.toy b/mlir/test/Examples/standalone/test.toy
index a88c115ebf197..dc3c17f3da3d9 100644
--- a/mlir/test/Examples/standalone/test.toy
+++ b/mlir/test/Examples/standalone/test.toy
@@ -4,8 +4,10 @@
 # RUN: -DLLVM_ENABLE_LIBCXX=%enable_libcxx -DMLIR_DIR=%mlir_cmake_dir \
 # RUN: -DLLVM_USE_LINKER=%llvm_use_linker \
 # RUN: -DMLIR_PYTHON_PACKAGE_PREFIX=mlir_standalone \
+# RUN: -DMLIR_INCLUDE_TESTS=ON \
 # RUN: -DPython3_EXECUTABLE=%python \
 # RUN: -DPython_EXECUTABLE=%python
+# RUN: export PYTHONPATH="%mlir_obj_root/python_packages/mlir_core"
 # RUN: "%cmake_exe" --build . --target check-standalone | tee %t
 # RUN: FileCheck --input-file=%t %s
 
diff --git a/mlir/test/Examples/standalone/test.wheel.toy b/mlir/test/Examples/standalone/test.wheel.toy
index c8d188a3cacd0..f1f127fe36dab 100644
--- a/mlir/test/Examples/standalone/test.wheel.toy
+++ b/mlir/test/Examples/standalone/test.wheel.toy
@@ -14,21 +14,25 @@
 # RUN: export CMAKE_GENERATOR=%cmake_generator
 # RUN: export LLVM_USE_LINKER=%llvm_use_linker
 # RUN: export MLIR_DIR="%mlir_cmake_dir"
+# RUN: export MLIR_INCLUDE_TESTS=ON
 
 # RUN: %python -m pip wheel "%mlir_src_root/examples/standalone" -w "%mlir_obj_root/wheelhouse" -v | tee %t
 
 # RUN: rm -rf "%mlir_obj_root/standalone-python-bindings-install"
 # RUN: %python -m pip install standalone_python_bindings -f "%mlir_obj_root/wheelhouse" --target "%mlir_obj_root/standalone-python-bindings-install" -v | tee -a %t
 
-# RUN: export PYTHONPATH="%mlir_obj_root/standalone-python-bindings-install"
-# RUN: %python "%mlir_src_root/examples/standalone/test/python/smoketest.py" nanobind | tee -a %t
+# RUN: export PYTHONPATH="%mlir_obj_root/standalone-python-bindings-install:%mlir_obj_root/python_packages/mlir_core"
+# RUN: %python "%mlir_src_root/examples/standalone/test/python/smoketest.py" 2>&1 | tee -a %t
 
 # RUN: FileCheck --input-file=%t %s
 
 # CHECK: Successfully built standalone-python-bindings
 
+# CHECK-NOT: RuntimeWarning: nanobind: type '{{.*}}' was already registered!
+
 # CHECK: module {
 # CHECK:   %[[C2:.*]] = arith.constant 2 : i32
 # CHECK:   %[[V0:.*]] = standalone.foo %[[C2]] : i32
 # CHECK: }
 
+# CHECK: !standalone.custom<"foo">