修正 mxcc MACA 路径处理

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

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()