支持环境变量注入 mxcc 额外参数

python/tvm/contrib/mxcc.py

@@ -19,6 +19,7 @@
 
 import re
 import os
+import shlex
 import subprocess
 import warnings
 from typing import Tuple
@@ -31,6 +32,13 @@
 from . import utils
 
 
+def _get_env_mxcc_options():
+    value = os.environ.get("TVM_MXCC_OPTIONS")
+    if not value:
+        return []
+    return shlex.split(value)
+
+
 def compile_maca(
     code, target_format="mcbin", arch=None, options=None, path_target=None
 ):  # pylint: disable=unused-argument
@@ -99,6 +107,7 @@ def compile_maca(
             cmd += options
         else:
             raise ValueError("options must be str of list of str")
+    cmd += _get_env_mxcc_options()
 
     cmd += ["-D__FAST_HALF_CVT__", "-o", file_target]
     cmd += [temp_code]

tests/python/contrib/test_mxcc_env_options.py

@@ -0,0 +1,18 @@
+import os
+from unittest.mock import patch
+from tvm.contrib.mxcc import _get_env_mxcc_options
+
+
+def test_env_mxcc_options_are_shell_split():
+    with patch.dict(os.environ, {"TVM_MXCC_OPTIONS": "--flag 'two words'"}, clear=True):
+        assert _get_env_mxcc_options() == ["--flag", "two words"]
+
+
+def test_env_mxcc_options_default_empty():
+    with patch.dict(os.environ, {}, clear=True):
+        assert _get_env_mxcc_options() == []
+
+
+if __name__ == "__main__":
+    test_env_mxcc_options_are_shell_split()
+    test_env_mxcc_options_default_empty()