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