[Mlir-commits] [mlir] [MLIR][test] Check for ml_dtypes before running tests (PR #123061)

[Konrad Kleine]

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

>From 8db4f46694df096ce0e1e967ebcf8477403fbddc Mon Sep 17 00:00:00 2001
From: Konrad Kleine <kkleine at redhat.com>
Date: Wed, 15 Jan 2025 15:12:44 +0100
Subject: [PATCH] [MLIR][test] Check for ml_dtypes before running tests

---
 mlir/test/python/execution_engine.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/mlir/test/python/execution_engine.py b/mlir/test/python/execution_engine.py
index 0d12c35d96bee7..e3f41815800d58 100644
--- a/mlir/test/python/execution_engine.py
+++ b/mlir/test/python/execution_engine.py
@@ -5,7 +5,14 @@
 from mlir.passmanager import *
 from mlir.execution_engine import *
 from mlir.runtime import *
-from ml_dtypes import bfloat16, float8_e5m2
+
+try:
+    from ml_dtypes import bfloat16, float8_e5m2
+
+    HAS_ML_DTYPES = True
+except ModuleNotFoundError:
+    HAS_ML_DTYPES = False
+
 
 MLIR_RUNNER_UTILS = os.getenv(
     "MLIR_RUNNER_UTILS", "../../../../lib/libmlir_runner_utils.so"
@@ -564,7 +571,8 @@ def testBF16Memref():
         log(npout)
 
 
-run(testBF16Memref)
+if HAS_ML_DTYPES:
+    run(testBF16Memref)
 
 
 # Test f8E5M2 memrefs
@@ -603,7 +611,8 @@ def testF8E5M2Memref():
         log(npout)
 
 
-run(testF8E5M2Memref)
+if HAS_ML_DTYPES:
+    run(testF8E5M2Memref)
 
 
 #  Test addition of two 2d_memref