diff --git a/python/tvm/contrib/mxcc.py b/python/tvm/contrib/mxcc.py index c5ad9eca5969..0031d0df5826 100644 --- a/python/tvm/contrib/mxcc.py +++ b/python/tvm/contrib/mxcc.py @@ -19,6 +19,7 @@ import re import os +import shutil import subprocess import warnings from typing import Tuple @@ -224,6 +225,12 @@ def get_maca_arch(maca_path="/opt/maca"): return gpu_arch +def _maca_path_from_mxcc(mxcc_path): + """Infer MACA root from an mxcc executable path.""" + real_path = os.path.realpath(mxcc_path) + return os.path.realpath(os.path.join(os.path.dirname(real_path), "../..")) + + def find_maca_path(): """Utility function to find MACA path @@ -234,12 +241,9 @@ def find_maca_path(): """ if "MACA_PATH" in os.environ: return os.environ["MACA_PATH"] - cmd = ["which", "mxcc"] - proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) - (out, _) = proc.communicate() - out = out.decode("utf-8").strip() - if proc.returncode == 0: - return os.path.realpath(os.path.join(out, "../../..")) + mxcc_path = shutil.which("mxcc") + if mxcc_path: + return _maca_path_from_mxcc(mxcc_path) maca_path = "/opt/maca" if os.path.exists(os.path.join(maca_path, "mxgpu_llvm/bin/mxcc")): return maca_path diff --git a/tests/python/contrib/test_mxcc_path.py b/tests/python/contrib/test_mxcc_path.py new file mode 100644 index 000000000000..54eb64fa86e8 --- /dev/null +++ b/tests/python/contrib/test_mxcc_path.py @@ -0,0 +1,28 @@ +import unittest +from pathlib import Path +from unittest.mock import patch + +from tvm.contrib import mxcc + + +class MxccPathTest(unittest.TestCase): + def test_maca_path_from_mxcc(self): + self.assertEqual( + mxcc._maca_path_from_mxcc("/opt/maca/mxgpu_llvm/bin/mxcc"), + str(Path("/opt/maca").resolve()), + ) + + def test_maca_path_from_symlink_target(self): + with patch("tvm.contrib.mxcc.os.path.realpath") as mock_realpath: + mock_realpath.side_effect = [ + "/opt/maca/mxgpu_llvm/bin/mxcc", + str(Path("/opt/maca").resolve()), + ] + self.assertEqual( + mxcc._maca_path_from_mxcc("/usr/bin/mxcc"), + str(Path("/opt/maca").resolve()), + ) + + +if __name__ == "__main__": + unittest.main()