diff --git a/mlir/test/CMakeLists.txt b/mlir/test/CMakeLists.txt index 58d16a657297e..78fadf4aea62c 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 f162f9a00efa7..cb42d645701c3 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 d68d9b1a1f43a..96be09b8036dc 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 0d12c35d96bee..061fa94289bf1 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 *