[Mlir-commits] [mlir] [MLIR][test] introduce has_py_module_ml_dtypes (PR #123051)

[Konrad Kleine]

https://github.com/kwk updated https://github.com/llvm/llvm-project/pull/123051

>From c9893d7131739923d51af0d05445cfb5d91dfba3 Mon Sep 17 00:00:00 2001
From: Konrad Kleine <kkleine at redhat.com>
Date: Wed, 15 Jan 2025 11:09:31 +0100
Subject: [PATCH] [MLIR][test] introduce has_py_module_ml_dtypes

Some python MLIR tests require the `ml_dtypes` module which doesn't
exist on Fedora for example. In order to avoid having to exclude a test
we would like to mark certain tests that require `ml_dtypes` with:

    # REQUIRES: has_py_module_ml_dtypes

We noticed that `mlir/python/requirements.txt` lists `ml_dtypes` as a
requirement but when looking at the code in `mlir/python`, the
only `import` is guarded:

```python
try:
    import ml_dtypes
except ModuleNotFoundError:
    # The third-party ml_dtypes provides some optional low precision data-types for NumPy.
    ml_dtypes = None
```

This makes `ml_dtypes` an optional dependency. Due to that reason I was
thinking if using a `pyproject.toml` could be a better fit to mark
dependencies as optional. But that's part of a different PR.
---
 mlir/test/CMakeLists.txt             | 9 +++++++++
 mlir/test/lit.cfg.py                 | 2 ++
 mlir/test/lit.site.cfg.py.in         | 1 +
 mlir/test/python/execution_engine.py | 2 +-
 4 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/mlir/test/CMakeLists.txt b/mlir/test/CMakeLists.txt
index 58d16a657297e6..78fadf4aea62c2 100644
--- a/mlir/test/CMakeLists.txt
+++ b/mlir/test/CMakeLists.txt
@@ -85,6 +85,15 @@ llvm_canonicalize_cmake_booleans(
   MLIR_RUN_CUDA_SM90_TESTS
   )
 
+# Run python file to find out of ml_dtypes is supported or not.
+execute_process(COMMAND ${Python3_EXECUTABLE} -c "import ml_dtypes" RESULT_VARIABLE HAS_PY_MODULE_ML_DTYPES_RET)
+if(NOT HAS_PY_MODULE_ML_DTYPES_RET EQUAL "0")
+  set(HAS_PY_MODULE_ML_DTYPES 0)
+else()
+  set(HAS_PY_MODULE_ML_DTYPES 1)
+endif()
+unset(HAS_PY_MODULE_ML_DTYPES_RET)
+
 configure_lit_site_cfg(
   ${CMAKE_CURRENT_SOURCE_DIR}/lit.site.cfg.py.in
   ${CMAKE_CURRENT_BINARY_DIR}/lit.site.cfg.py
diff --git a/mlir/test/lit.cfg.py b/mlir/test/lit.cfg.py
index f162f9a00efa7c..cb42d645701c3c 100644
--- a/mlir/test/lit.cfg.py
+++ b/mlir/test/lit.cfg.py
@@ -314,6 +314,8 @@ def find_real_python_interpreter():
 else:
     config.available_features.add("noasserts")
 
+if config.has_py_module_ml_dtypes:
+    config.available_features.add("has_py_module_ml_dtypes")
 
 def have_host_jit_feature_support(feature_name):
     mlir_cpu_runner_exe = lit.util.which("mlir-cpu-runner", config.mlir_tools_dir)
diff --git a/mlir/test/lit.site.cfg.py.in b/mlir/test/lit.site.cfg.py.in
index d68d9b1a1f43a2..96be09b8036dcc 100644
--- a/mlir/test/lit.site.cfg.py.in
+++ b/mlir/test/lit.site.cfg.py.in
@@ -9,6 +9,7 @@ config.llvm_shlib_ext = "@SHLIBEXT@"
 config.llvm_shlib_dir = lit_config.substitute(path(r"@SHLIBDIR@"))
 config.python_executable = "@Python3_EXECUTABLE@"
 config.enable_assertions = @ENABLE_ASSERTIONS@
+config.has_py_module_ml_dtypes = @HAS_PY_MODULE_ML_DTYPES@
 config.native_target = "@LLVM_NATIVE_ARCH@"
 config.host_os = "@HOST_OS@"
 config.host_cc = "@HOST_CC@"
diff --git a/mlir/test/python/execution_engine.py b/mlir/test/python/execution_engine.py
index 0d12c35d96bee7..061fa94289bf17 100644
--- a/mlir/test/python/execution_engine.py
+++ b/mlir/test/python/execution_engine.py
@@ -1,5 +1,5 @@
 # RUN: env MLIR_RUNNER_UTILS=%mlir_runner_utils MLIR_C_RUNNER_UTILS=%mlir_c_runner_utils %PYTHON %s 2>&1 | FileCheck %s
-# REQUIRES: host-supports-jit
+# REQUIRES: host-supports-jit, has_py_module_ml_dtypes
 import gc, sys, os, tempfile
 from mlir.ir import *
 from mlir.passmanager import *