diff --git a/python/tvm/contrib/mxcc.py b/python/tvm/contrib/mxcc.py index c5ad9eca5969..9012b3632504 100644 --- a/python/tvm/contrib/mxcc.py +++ b/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] diff --git a/tests/python/contrib/test_mxcc_env_options.py b/tests/python/contrib/test_mxcc_env_options.py new file mode 100644 index 000000000000..25778270ddba --- /dev/null +++ b/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()