[MLIR][test] introduce has_py_module_ml_dtypes

[kwk]

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:

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.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Konrad Kleine (kwk)

Changes

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:

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/123051.diff

4 Files Affected:

• (modified) mlir/test/CMakeLists.txt (+9)
• (modified) mlir/test/lit.cfg.py (+2)
• (modified) mlir/test/lit.site.cfg.py.in (+1)
• (modified) mlir/test/python/execution_engine.py (+1-1)

diff --git a/mlir/test/CMakeLists.txt b/mlir/test/CMakeLists.txt
index 58d16a657297e6..5361a450d842c0 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_dltypes" 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 *

[makslevental]

I feel like this is too broad a gate - you're disabling the entire set of tests because one or two can't run? The better thing to do is to just selectively run those tests depending on whether the import fails/succeeds:

try:
  import ml_dtypes
  HAS_ML_DTYPES = True
except ImportError:
  HAS_ML_DTYPES = False

...

if HAS_ML_DTYPES:
  ...

[kwk]

I feel like this is too broad a gate - you're disabling the entire set of tests because one or two can't run? The better thing to do is to just selectively run those tests depending on whether the import fails/succeeds:

try:
  import ml_dtypes
  HAS_ML_DTYPES = True
except ImportError:
  HAS_ML_DTYPES = False

...

if HAS_ML_DTYPES:
  ...

Agreed. Here's a PR that does that: #123061.