1# -*- Python -*- 2 3import os 4import platform 5import re 6import shutil 7import subprocess 8import tempfile 9 10import lit.formats 11import lit.util 12 13from lit.llvm import llvm_config 14from lit.llvm.subst import ToolSubst 15from lit.llvm.subst import FindTool 16 17# Configuration file for the 'lit' test runner. 18 19# name: The name of this test suite. 20config.name = "MLIR" 21 22# TODO: Consolidate the logic for turning on the internal shell by default for all LLVM test suites. 23# See https://github.com/llvm/llvm-project/issues/106636 for more details. 24# 25# We prefer the lit internal shell which provides a better user experience on failures 26# unless the user explicitly disables it with LIT_USE_INTERNAL_SHELL=0 env var. 27use_lit_shell = True 28lit_shell_env = os.environ.get("LIT_USE_INTERNAL_SHELL") 29if lit_shell_env: 30 use_lit_shell = lit.util.pythonize_bool(lit_shell_env) 31 32config.test_format = lit.formats.ShTest(execute_external=not use_lit_shell) 33 34# suffixes: A list of file extensions to treat as test files. 35config.suffixes = [ 36 ".td", 37 ".mlir", 38 ".toy", 39 ".ll", 40 ".tc", 41 ".py", 42 ".yaml", 43 ".test", 44 ".pdll", 45 ".c", 46] 47 48# test_source_root: The root path where tests are located. 49config.test_source_root = os.path.dirname(__file__) 50 51# test_exec_root: The root path where tests should be run. 52config.test_exec_root = os.path.join(config.mlir_obj_root, "test") 53 54config.substitutions.append(("%PATH%", config.environment["PATH"])) 55config.substitutions.append(("%shlibext", config.llvm_shlib_ext)) 56config.substitutions.append(("%llvm_src_root", config.llvm_src_root)) 57config.substitutions.append(("%mlir_src_root", config.mlir_src_root)) 58config.substitutions.append(("%host_cxx", config.host_cxx)) 59config.substitutions.append(("%host_cc", config.host_cc)) 60 61 62# Searches for a runtime library with the given name and returns the found path. 63# Correctly handles the platforms shared library directory and naming conventions. 64def find_runtime(name): 65 path = "" 66 for prefix in ["", "lib"]: 67 path = os.path.join( 68 config.llvm_shlib_dir, f"{prefix}{name}{config.llvm_shlib_ext}" 69 ) 70 if os.path.isfile(path): 71 break 72 return path 73 74 75# Searches for a runtime library with the given name and returns a tool 76# substitution of the same name and the found path. 77def add_runtime(name): 78 return ToolSubst(f"%{name}", find_runtime(name)) 79 80 81# Provide the path to asan runtime lib 'libclang_rt.asan_osx_dynamic.dylib' if 82# available. This is darwin specific since it's currently only needed on darwin. 83# Stolen from llvm/test/lit.cfg.py with a few modifications 84def get_asan_rtlib(): 85 if not "asan" in config.available_features: 86 return "" 87 88 if "Darwin" in config.host_os: 89 # Find the asan rt lib 90 resource_dir = ( 91 subprocess.check_output([config.host_cc.strip(), "-print-resource-dir"]) 92 .decode("utf-8") 93 .strip() 94 ) 95 return os.path.join( 96 resource_dir, "lib", "darwin", "libclang_rt.asan_osx_dynamic.dylib" 97 ) 98 if "Linux" in config.host_os: 99 return ( 100 subprocess.check_output( 101 [ 102 config.host_cxx.strip(), 103 f"-print-file-name=libclang_rt.asan-{config.host_arch}.so", 104 ] 105 ) 106 .decode("utf-8") 107 .strip() 108 ) 109 110 return "" 111 112 113# On macOS, we can't do the DYLD_INSERT_LIBRARIES trick with a shim python 114# binary as the ASan interceptors get loaded too late. Also, when SIP is 115# enabled, we can't inject libraries into system binaries at all, so we need a 116# copy of the "real" python to work with. 117# Stolen from lldb/test/API/lit.cfg.py with a few modifications 118def find_real_python_interpreter(): 119 # If we're running in a virtual environment, we have to copy Python into 120 # the virtual environment for it to work. 121 if sys.prefix != sys.base_prefix: 122 copied_python = os.path.join(sys.prefix, "bin", "copied-python") 123 else: 124 copied_python = os.path.join(config.mlir_obj_root, "copied-python") 125 126 # Avoid doing any work if we already copied the binary. 127 if os.path.isfile(copied_python): 128 return copied_python 129 130 # Find the "real" python binary. 131 real_python = ( 132 subprocess.check_output( 133 [ 134 config.python_executable, 135 os.path.join( 136 os.path.dirname(os.path.realpath(__file__)), 137 "get_darwin_real_python.py", 138 ), 139 ] 140 ) 141 .decode("utf-8") 142 .strip() 143 ) 144 145 shutil.copy(real_python, copied_python) 146 147 # Now make sure the copied Python works. The Python in Xcode has a relative 148 # RPATH and cannot be copied. 149 try: 150 # We don't care about the output, just make sure it runs. 151 subprocess.check_call([copied_python, "-V"]) 152 except subprocess.CalledProcessError: 153 # The copied Python didn't work. Assume we're dealing with the Python 154 # interpreter in Xcode. Given that this is not a system binary SIP 155 # won't prevent us form injecting the interceptors, but when running in 156 # a virtual environment, we can't use it directly. Create a symlink 157 # instead. 158 os.remove(copied_python) 159 os.symlink(real_python, copied_python) 160 161 # The copied Python works. 162 return copied_python 163 164 165llvm_config.with_system_environment(["HOME", "INCLUDE", "LIB", "TMP", "TEMP"]) 166 167llvm_config.use_default_substitutions() 168 169# excludes: A list of directories to exclude from the testsuite. The 'Inputs' 170# subdirectories contain auxiliary inputs for various tests in their parent 171# directories. 172config.excludes = [ 173 "Inputs", 174 "CMakeLists.txt", 175 "README.txt", 176 "LICENSE.txt", 177 "lit.cfg.py", 178 "lit.site.cfg.py", 179 "get_darwin_real_python.py", 180] 181 182# Tweak the PATH to include the tools dir. 183llvm_config.with_environment("PATH", config.mlir_tools_dir, append_path=True) 184llvm_config.with_environment("PATH", config.llvm_tools_dir, append_path=True) 185 186tool_dirs = [config.mlir_tools_dir, config.llvm_tools_dir] 187tools = [ 188 "mlir-tblgen", 189 "mlir-translate", 190 "mlir-lsp-server", 191 "mlir-capi-execution-engine-test", 192 "mlir-capi-ir-test", 193 "mlir-capi-irdl-test", 194 "mlir-capi-llvm-test", 195 "mlir-capi-pass-test", 196 "mlir-capi-pdl-test", 197 "mlir-capi-quant-test", 198 "mlir-capi-rewrite-test", 199 "mlir-capi-sparse-tensor-test", 200 "mlir-capi-transform-test", 201 "mlir-capi-transform-interpreter-test", 202 "mlir-capi-translation-test", 203 "mlir-runner", 204 add_runtime("mlir_runner_utils"), 205 add_runtime("mlir_c_runner_utils"), 206 add_runtime("mlir_async_runtime"), 207 add_runtime("mlir_float16_utils"), 208 "mlir-linalg-ods-yaml-gen", 209 "mlir-reduce", 210 "mlir-pdll", 211 "not", 212] 213 214if config.enable_vulkan_runner: 215 tools.extend([add_runtime("mlir_vulkan_runtime")]) 216 217if config.enable_rocm_runner: 218 tools.extend([add_runtime("mlir_rocm_runtime")]) 219 220if config.enable_cuda_runner: 221 tools.extend([add_runtime("mlir_cuda_runtime")]) 222 223if config.enable_sycl_runner: 224 tools.extend([add_runtime("mlir_sycl_runtime")]) 225 226if config.enable_spirv_cpu_runner: 227 tools.extend([add_runtime("mlir_spirv_cpu_runtime")]) 228 229if config.mlir_run_arm_sve_tests or config.mlir_run_arm_sme_tests: 230 tools.extend([add_runtime("mlir_arm_runner_utils")]) 231 232if config.mlir_run_arm_sme_tests: 233 config.substitutions.append( 234 ( 235 "%arm_sme_abi_shlib", 236 # Use passed Arm SME ABI routines, if not present default to stubs. 237 config.arm_sme_abi_routines_shlib or find_runtime("mlir_arm_sme_abi_stubs"), 238 ) 239 ) 240 241# The following tools are optional 242tools.extend( 243 [ 244 ToolSubst("toyc-ch1", unresolved="ignore"), 245 ToolSubst("toyc-ch2", unresolved="ignore"), 246 ToolSubst("toyc-ch3", unresolved="ignore"), 247 ToolSubst("toyc-ch4", unresolved="ignore"), 248 ToolSubst("toyc-ch5", unresolved="ignore"), 249 ToolSubst("toyc-ch6", unresolved="ignore"), 250 ToolSubst("toyc-ch7", unresolved="ignore"), 251 ToolSubst("transform-opt-ch2", unresolved="ignore"), 252 ToolSubst("transform-opt-ch3", unresolved="ignore"), 253 ToolSubst("transform-opt-ch4", unresolved="ignore"), 254 ToolSubst("mlir-transform-opt", unresolved="ignore"), 255 ToolSubst("%mlir_lib_dir", config.mlir_lib_dir, unresolved="ignore"), 256 ToolSubst("%mlir_src_dir", config.mlir_src_root, unresolved="ignore"), 257 ] 258) 259 260python_executable = config.python_executable 261# Python configuration with sanitizer requires some magic preloading. This will only work on clang/linux/darwin. 262# TODO: detect Windows situation (or mark these tests as unsupported on these platforms). 263if "asan" in config.available_features: 264 if "Linux" in config.host_os: 265 python_executable = ( 266 f"env LD_PRELOAD={get_asan_rtlib()} {config.python_executable}" 267 ) 268 if "Darwin" in config.host_os: 269 # Ensure we use a non-shim Python executable, for the `DYLD_INSERT_LIBRARIES` 270 # env variable to take effect 271 real_python_executable = find_real_python_interpreter() 272 if real_python_executable: 273 python_executable = real_python_executable 274 lit_config.note( 275 "Using {} instead of {}".format( 276 python_executable, config.python_executable 277 ) 278 ) 279 280 asan_rtlib = get_asan_rtlib() 281 lit_config.note("Using ASan rtlib {}".format(asan_rtlib)) 282 config.environment["MallocNanoZone"] = "0" 283 config.environment["ASAN_OPTIONS"] = "detect_stack_use_after_return=1" 284 config.environment["DYLD_INSERT_LIBRARIES"] = asan_rtlib 285 286 287# On Windows the path to python could contains spaces in which case it needs to be provided in quotes. 288# This is the equivalent of how %python is setup in llvm/utils/lit/lit/llvm/config.py. 289elif "Windows" in config.host_os: 290 python_executable = '"%s"' % (python_executable) 291tools.extend( 292 [ 293 ToolSubst("%PYTHON", python_executable, unresolved="ignore"), 294 ] 295) 296 297if "MLIR_OPT_CHECK_IR_ROUNDTRIP" in os.environ: 298 tools.extend( 299 [ 300 ToolSubst("mlir-opt", "mlir-opt --verify-roundtrip", unresolved="fatal"), 301 ] 302 ) 303else: 304 tools.extend(["mlir-opt"]) 305 306llvm_config.add_tool_substitutions(tools, tool_dirs) 307 308 309# FileCheck -enable-var-scope is enabled by default in MLIR test 310# This option avoids to accidentally reuse variable across -LABEL match, 311# it can be explicitly opted-in by prefixing the variable name with $ 312config.environment["FILECHECK_OPTS"] = "-enable-var-scope --allow-unused-prefixes=false" 313 314# Add the python path for both the source and binary tree. 315# Note that presently, the python sources come from the source tree and the 316# binaries come from the build tree. This should be unified to the build tree 317# by copying/linking sources to build. 318if config.enable_bindings_python: 319 config.environment["PYTHONPATH"] = os.getenv("MLIR_LIT_PYTHONPATH", "") 320 llvm_config.with_environment( 321 "PYTHONPATH", 322 [ 323 os.path.join(config.mlir_obj_root, "python_packages", "mlir_core"), 324 os.path.join(config.mlir_obj_root, "python_packages", "mlir_test"), 325 ], 326 append_path=True, 327 ) 328 329if config.enable_assertions: 330 config.available_features.add("asserts") 331else: 332 config.available_features.add("noasserts") 333 334 335def have_host_jit_feature_support(feature_name): 336 mlir_runner_exe = lit.util.which("mlir-runner", config.mlir_tools_dir) 337 338 if not mlir_runner_exe: 339 return False 340 341 try: 342 mlir_runner_cmd = subprocess.Popen( 343 [mlir_runner_exe, "--host-supports-" + feature_name], 344 stdout=subprocess.PIPE, 345 ) 346 except OSError: 347 print("could not exec mlir-runner") 348 return False 349 350 mlir_runner_out = mlir_runner_cmd.stdout.read().decode("ascii") 351 mlir_runner_cmd.wait() 352 353 return "true" in mlir_runner_out 354 355 356if have_host_jit_feature_support("jit"): 357 config.available_features.add("host-supports-jit") 358 359if config.run_nvptx_tests: 360 config.available_features.add("host-supports-nvptx") 361 362if config.run_rocm_tests: 363 config.available_features.add("host-supports-amdgpu") 364 365if config.arm_emulator_executable: 366 config.available_features.add("arm-emulator") 367