[Mlir-commits] [mlir] [mlir][python] fix flatnamespace (PR #174057)

[Maksim Levental]

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

>From afa0c29ee8a7be1dd01a92be29931d0624e7c65a 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/examples/standalone/CMakeLists.txt           |  4 +++-
 mlir/examples/standalone/test/lit.cfg.py          | 12 +++++-------
 mlir/examples/standalone/test/python/smoketest.py |  7 ++++++-
 mlir/test/Examples/standalone/test.toy            |  2 ++
 mlir/test/Examples/standalone/test.wheel.toy      |  8 ++++++--
 5 files changed, 22 insertions(+), 11 deletions(-)

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..3ba2e2e4698de 100644
--- a/mlir/examples/standalone/test/python/smoketest.py
+++ b/mlir/examples/standalone/test/python/smoketest.py
@@ -1,8 +1,10 @@
-# RUN: %python %s nanobind | FileCheck %s
+# RUN: %python %s 2>&1 | FileCheck %s
 
 from mlir_standalone.ir import *
 from mlir_standalone.dialects import standalone_nanobind as standalone_d
 
+# CHECK-NOT: RuntimeWarning: nanobind: type '{{.*}}' was already registered!
+
 with Context():
     standalone_d.register_dialects()
     module = Module.parse(
@@ -14,3 +16,6 @@
     # CHECK: %[[C:.*]] = arith.constant 2 : i32
     # CHECK: standalone.foo %[[C]] : i32
     print(str(module))
+
+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">