diff --git a/python/triton/runtime/build.py b/python/triton/runtime/build.py index 3c8accef3c4f..a3a08a7460e4 100644 --- a/python/triton/runtime/build.py +++ b/python/triton/runtime/build.py @@ -22,6 +22,7 @@ def quiet(): def _build(name, src, srcdir, library_dirs, include_dirs, libraries): suffix = sysconfig.get_config_var('EXT_SUFFIX') system = platform.system() + machine = platform.machine() so = os.path.join(srcdir, '{name}{suffix}'.format(name=name, suffix=suffix)) # try to avoid setuptools if possible cc = os.environ.get("CC") @@ -55,6 +56,9 @@ def _build(name, src, srcdir, library_dirs, include_dirs, libraries): cc_cmd += ["-std=c++17", "-fopenmp"] if src.endswith(".s"): cc_cmd += ["-gdwarf-5"] + if system == "Linux" and machine in ("aarch64", "arm64"): + # On Arm backend, some CPU (neoverse-v2) needs to be specified through -mcpu + cc_cmd += ["-mcpu=native"] ret = subprocess.check_call(cc_cmd) if ret == 0: return so