From 0a16d1a754c247199ee2a77feaea7a94a60f8dfd Mon Sep 17 00:00:00 2001
From: Maksim Levental <maksim.levental@gmail.com>
Date: Sat, 30 Aug 2025 20:52:04 -0400
Subject: [PATCH] [MLIR][Python] fix standalone example/test (#156197)

Fix some things in `standalone` in order to unblock
https://github.com/llvm/llvm-project/pull/155741.
---
 mlir/CMakeLists.txt                            |  4 ++++
 mlir/examples/standalone/CMakeLists.txt        |  4 ++++
 mlir/examples/standalone/python/CMakeLists.txt | 15 ++++++++-------
 mlir/test/Examples/standalone/test.toy         |  7 +++++--
 4 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/mlir/CMakeLists.txt b/mlir/CMakeLists.txt
index f58a4c6f506e..1a211f549576 100644
--- a/mlir/CMakeLists.txt
+++ b/mlir/CMakeLists.txt
@@ -191,6 +191,10 @@ configure_file(
 
 set(MLIR_BINDINGS_PYTHON_NB_DOMAIN "mlir"
   CACHE STRING "nanobind domain for MLIR python bindings.")
+set(MLIR_PYTHON_PACKAGE_PREFIX "mlir"
+  CACHE STRING "Specifies that all MLIR packages are co-located under the
+  `MLIR_PYTHON_PACKAGE_PREFIX` top level package (the API has been
+  embedded in a relocatable way).")
 set(MLIR_ENABLE_BINDINGS_PYTHON 0 CACHE BOOL
        "Enables building of Python bindings.")
 set(MLIR_BINDINGS_PYTHON_INSTALL_PREFIX "python_packages/mlir_core/mlir" CACHE STRING
diff --git a/mlir/examples/standalone/CMakeLists.txt b/mlir/examples/standalone/CMakeLists.txt
index 42b487fe2d40..88dfa3e5d57a 100644
--- a/mlir/examples/standalone/CMakeLists.txt
+++ b/mlir/examples/standalone/CMakeLists.txt
@@ -52,6 +52,10 @@ add_subdirectory(include)
 add_subdirectory(lib)
 if(MLIR_ENABLE_BINDINGS_PYTHON)
   message(STATUS "Enabling Python API")
+  include(MLIRDetectPythonEnv)
+  mlir_configure_python_dev_packages()
+  set(MLIR_PYTHON_PACKAGE_PREFIX "mlir_standalone" CACHE STRING "" FORCE)
+  set(MLIR_BINDINGS_PYTHON_INSTALL_PREFIX "python_packages/standalone/mlir_standalone" CACHE STRING "" FORCE)
   add_subdirectory(python)
 endif()
 add_subdirectory(test)
diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt
index 69c82fd91357..a0eca9c09577 100644
--- a/mlir/examples/standalone/python/CMakeLists.txt
+++ b/mlir/examples/standalone/python/CMakeLists.txt
@@ -2,8 +2,7 @@ include(AddMLIRPython)
 
 # Specifies that all MLIR packages are co-located under the `mlir_standalone`
 # top level package (the API has been embedded in a relocatable way).
-# TODO: Add an upstream cmake param for this vs having a global here.
-add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=mlir_standalone.")
+add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=${MLIR_PYTHON_PACKAGE_PREFIX}.")
 
 
 ################################################################################
@@ -49,8 +48,8 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension
 
 add_mlir_python_common_capi_library(StandalonePythonCAPI
   INSTALL_COMPONENT StandalonePythonModules
-  INSTALL_DESTINATION python_packages/standalone/mlir_standalone/_mlir_libs
-  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/standalone/mlir_standalone/_mlir_libs"
+  INSTALL_DESTINATION "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/_mlir_libs"
+  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/_mlir_libs"
   RELATIVE_INSTALL_ROOT "../../../.."
   DECLARED_SOURCES
     StandalonePythonSources
@@ -58,6 +57,7 @@ add_mlir_python_common_capi_library(StandalonePythonCAPI
     # available.
     MLIRPythonExtension.RegisterEverything
     MLIRPythonSources.Core
+    MLIRPythonSources.Dialects.builtin
 )
 
 ################################################################################
@@ -65,14 +65,15 @@ add_mlir_python_common_capi_library(StandalonePythonCAPI
 ################################################################################
 
 add_mlir_python_modules(StandalonePythonModules
-  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/standalone/mlir_standalone"
-  INSTALL_PREFIX "python_packages/standalone/mlir_standalone"
+  ROOT_PREFIX "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}"
+  INSTALL_PREFIX "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}"
   DECLARED_SOURCES
     StandalonePythonSources
     # TODO: Remove this in favor of showing fine grained registration once
     # available.
     MLIRPythonExtension.RegisterEverything
-    MLIRPythonSources
+    MLIRPythonSources.Core
+    MLIRPythonSources.Dialects.builtin
   COMMON_CAPI_LINK_LIBS
     StandalonePythonCAPI
   )
diff --git a/mlir/test/Examples/standalone/test.toy b/mlir/test/Examples/standalone/test.toy
index c91b3cf8a50f..e99bab5f0aff 100644
--- a/mlir/test/Examples/standalone/test.toy
+++ b/mlir/test/Examples/standalone/test.toy
@@ -2,10 +2,13 @@
 # RUN: -DCMAKE_CXX_COMPILER=%host_cxx -DCMAKE_C_COMPILER=%host_cc \
 # RUN: -DLLVM_ENABLE_LIBCXX=%enable_libcxx -DMLIR_DIR=%mlir_cmake_dir \
 # RUN: -DLLVM_USE_LINKER=%llvm_use_linker \
-# RUN: -DPython3_EXECUTABLE=%python
-# RUN: "%cmake_exe" --build . --target check-standalone | tee %t | FileCheck %s
+# RUN: -DPython3_EXECUTABLE=%python \
+# RUN: -DPython_EXECUTABLE=%python
+# RUN: "%cmake_exe" --build . --target check-standalone | tee %t
+# RUN: FileCheck --input-file=%t %s
 
 # Note: The number of checked tests is not important. The command will fail
 # if any fail.
 # CHECK: Passed
+# CHECK-NOT: Failed
 # UNSUPPORTED: target={{.*(windows|android).*}}