diff --git a/examples/ingress/mlir_gen/generate-linalg-3layer-mlp.sh b/examples/ingress/mlir_gen/generate-linalg-3layer-mlp.sh index 50c117ba..e6a494a2 100755 --- a/examples/ingress/mlir_gen/generate-linalg-3layer-mlp.sh +++ b/examples/ingress/mlir_gen/generate-linalg-3layer-mlp.sh @@ -2,7 +2,7 @@ # RUN: bash %s -PROJECT_ROOT="$(dirname "$(dirname "$(dirname "$(dirname "$(readlink -fm "$0")")")")")" +PROJECT_ROOT="$(cd "$(dirname "$0")/../../.." && pwd)" CACHE_DIR=$PROJECT_ROOT/cache/ingress/mlir_gen LAYERS=1024,2048,4096,512 diff --git a/examples/xegpu/kernel_bench.py b/examples/xegpu/kernel_bench.py index c1a93425..1450bf73 100644 --- a/examples/xegpu/kernel_bench.py +++ b/examples/xegpu/kernel_bench.py @@ -1,4 +1,5 @@ # RUN: %PYTHON %s -l 2 -b 9 --dump-kernel=xegpu-wg | FileCheck %s +# REQUIRES: torch # CHECK: module attributes {gpu.container_module} { """ This script executes KernelBench benchmarks using the XEGPU lowering pipeline. diff --git a/lighthouse/execution/runner.py b/lighthouse/execution/runner.py index 2d4c132a..3512e7ff 100644 --- a/lighthouse/execution/runner.py +++ b/lighthouse/execution/runner.py @@ -21,7 +21,7 @@ from lighthouse.dialects.transform import transform_ext from lighthouse.schedule import schedule_boilerplate from lighthouse.utils.memref import to_packed_args -from lighthouse.utils.mlir import get_mlir_library_path +from lighthouse.utils.mlir import get_mlir_library_path, _SHARED_EXT from lighthouse.utils.lib_finder import find_openmp_library from .memory_manager import GPUMemoryManager, ExternalMemoryManager, MemoryManager @@ -60,7 +60,7 @@ def __init__( if shared_libs is None: shared_libs = [] # get execution engine, rtclock requires mlir_c_runner - c_runner_lib = "libmlir_c_runner_utils.so" + c_runner_lib = f"libmlir_c_runner_utils{_SHARED_EXT}" if c_runner_lib not in shared_libs: shared_libs.append(c_runner_lib) self.lib_dir = get_mlir_library_path() diff --git a/lighthouse/execution/target.py b/lighthouse/execution/target.py index ff282964..3d4d4720 100644 --- a/lighthouse/execution/target.py +++ b/lighthouse/execution/target.py @@ -81,7 +81,14 @@ def override( cls.reset_host_cache() def _get_feature_list(self) -> list[str]: - """Get features from lscpu program""" + """Get CPU features from the host system.""" + if platform.system() == "Darwin": + return self._get_feature_list_darwin() + return self._get_feature_list_linux() + + @staticmethod + def _get_feature_list_linux() -> list[str]: + """Get features from lscpu (Linux).""" flags = subprocess.run( "lscpu | grep Flags", capture_output=True, @@ -93,7 +100,40 @@ def _get_feature_list(self) -> list[str]: "Could not get CPU features from lscpu. " "Make sure lscpu is installed and available in PATH." ) - features = flags.split()[1:] # Remove the "Flags:" prefix + return flags.split()[1:] + + @staticmethod + def _get_feature_list_darwin() -> list[str]: + """Get features from sysctl (macOS).""" + result = subprocess.run( + ["sysctl", "-n", "hw.optional.cpu_features"], + capture_output=True, + text=True, + ) + if result.returncode == 0 and result.stdout.strip(): + return result.stdout.strip().split() + + # Apple Silicon: enumerate hw.optional.arm.* and hw.optional.armv8_* keys. + result = subprocess.run( + ["sysctl", "-a"], + capture_output=True, + text=True, + ) + features = [] + for line in result.stdout.splitlines(): + if not line.startswith("hw.optional."): + continue + key, _, value = line.partition(":") + if value.strip() not in ("1",): + continue + # Strip the "hw.optional." prefix. + feat = key.split(".", 2)[-1] + features.append(feat) + if not features: + raise RuntimeError( + "Could not get CPU features from sysctl. " + "Make sure sysctl is installed and available in PATH." + ) return features def has_features(self, filter: list[str]) -> list[str]: diff --git a/lighthouse/utils/mlir.py b/lighthouse/utils/mlir.py index e30e7f4e..17e626fe 100644 --- a/lighthouse/utils/mlir.py +++ b/lighthouse/utils/mlir.py @@ -5,13 +5,16 @@ from mlir import ir from mlir.dialects import func, linalg import os +import platform from pathlib import Path +_SHARED_EXT = ".dylib" if platform.system() == "Darwin" else ".so" + def get_mlir_library_path(): """Return MLIR shared library path.""" pkg_path = Path(ir.__file__).parent - run_utils_so = "libmlir_runner_utils.so" + run_utils_lib = f"libmlir_runner_utils{_SHARED_EXT}" err_msg = f"Could not find shared libs in locations relative to '{pkg_path}'" if "python_packages" in str(pkg_path): # looks like a local llvm install @@ -19,19 +22,19 @@ def get_mlir_library_path(): # LLVM_INSTALL_DIR/python_packages/mlir_core/mlir # lib location: LLVM_INSTALL_DIR/lib/ path = pkg_path.parent.parent.parent / "lib" - assert os.path.isfile(path / run_utils_so) + assert os.path.isfile(path / run_utils_lib) except AssertionError: try: # LLVM_BUILD_DIR/tools/mlir/python_packages/mlir_core/mlir # lib location: LLVM_BUILD_DIR/lib/ path = pkg_path.parent.parent.parent.parent.parent / "lib" - assert os.path.isfile(path / run_utils_so) + assert os.path.isfile(path / run_utils_lib) except AssertionError: raise ValueError(err_msg) else: # maybe installed in python path path = pkg_path / "_mlir_libs" - assert os.path.isfile(path / run_utils_so), err_msg + assert os.path.isfile(path / run_utils_lib), err_msg return path