支持配置 MACA kernel 输出目录#18
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the MACA kernel output directory retrieval in python/tvm/contrib/mxcc.py by introducing a helper function _get_kernels_output_dir that falls back to the TVM_MACA_KERNELS_OUTPUT_DIR environment variable, and adds corresponding unit tests. The review feedback suggests adding a safety check to handle cases where pass_context is None to avoid an AttributeError, and simplifying the unit tests by directly importing the helper function instead of using complex AST parsing and dynamic execution.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| def _get_kernels_output_dir(pass_context): | ||
| if "maca.kernels_output_dir" in pass_context.config: | ||
| return pass_context.config["maca.kernels_output_dir"] | ||
| return os.environ.get("TVM_MACA_KERNELS_OUTPUT_DIR") |
There was a problem hiding this comment.
如析 pass_context 为 None（当没有活跃的 pass context 时可能会发生），访问 pass_context.config 将会抛出 AttributeError。我们应该在访问其属性之前安全地检查 pass_context 是否不为 None，以确保环境变量的备用逻辑能够正常工作。
| def _get_kernels_output_dir(pass_context): | |
| if "maca.kernels_output_dir" in pass_context.config: | |
| return pass_context.config["maca.kernels_output_dir"] | |
| return os.environ.get("TVM_MACA_KERNELS_OUTPUT_DIR") | |
| def _get_kernels_output_dir(pass_context): | |
| if pass_context is not None and "maca.kernels_output_dir" in pass_context.config: | |
| return pass_context.config["maca.kernels_output_dir"] | |
| return os.environ.get("TVM_MACA_KERNELS_OUTPUT_DIR") |
| import ast | ||
| import os | ||
| from pathlib import Path | ||
| from types import SimpleNamespace | ||
| from unittest.mock import patch | ||
|
|
||
|
|
||
| def _load_output_dir_helper(): | ||
| source_path = ( | ||
| Path(__file__).resolve().parents[3] / "python" / "tvm" / "contrib" / "mxcc.py" | ||
| ) | ||
| tree = ast.parse(source_path.read_text(encoding="utf-8")) | ||
| nodes = [ | ||
| node | ||
| for node in tree.body | ||
| if ( | ||
| isinstance(node, ast.Import) | ||
| and all(alias.name == "os" for alias in node.names) | ||
| ) | ||
| or (isinstance(node, ast.FunctionDef) and node.name == "_get_kernels_output_dir") | ||
| ] | ||
| module = ast.Module(body=nodes, type_ignores=[]) | ||
| ast.fix_missing_locations(module) | ||
| namespace = {} | ||
| exec(compile(module, str(source_path), "exec"), namespace) | ||
| return namespace["_get_kernels_output_dir"] | ||
|
|
||
|
|
||
| def test_kernel_output_dir_uses_pass_context_first(): | ||
| get_output_dir = _load_output_dir_helper() | ||
| pass_context = SimpleNamespace(config={"maca.kernels_output_dir": "/ctx"}) | ||
| with patch.dict(os.environ, {"TVM_MACA_KERNELS_OUTPUT_DIR": "/env"}, clear=True): | ||
| assert get_output_dir(pass_context) == "/ctx" | ||
|
|
||
|
|
||
| def test_kernel_output_dir_falls_back_to_env(): | ||
| get_output_dir = _load_output_dir_helper() | ||
| pass_context = SimpleNamespace(config={}) | ||
| with patch.dict(os.environ, {"TVM_MACA_KERNELS_OUTPUT_DIR": "/env"}, clear=True): | ||
| assert get_output_dir(pass_context) == "/env" | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| test_kernel_output_dir_uses_pass_context_first() | ||
| test_kernel_output_dir_falls_back_to_env() |
There was a problem hiding this comment.
使用 AST 解析和动态执行（exec）来从 mxcc.py 加载函数非常脆弱、复杂且难以维护。由于这是一个标准的 TVM 测试套件，我们可以直接从 tvm.contrib.mxcc 导入 _get_kernels_output_dir。
import os
from unittest.mock import patch
from types import SimpleNamespace
from tvm.contrib.mxcc import _get_kernels_output_dir
def test_kernel_output_dir_uses_pass_context_first():
pass_context = SimpleNamespace(config={"maca.kernels_output_dir": "/ctx"})
with patch.dict(os.environ, {"TVM_MACA_KERNELS_OUTPUT_DIR": "/env"}, clear=True):
assert _get_kernels_output_dir(pass_context) == "/ctx"
def test_kernel_output_dir_falls_back_to_env():
pass_context = SimpleNamespace(config={})
with patch.dict(os.environ, {"TVM_MACA_KERNELS_OUTPUT_DIR": "/env"}, clear=True):
assert _get_kernels_output_dir(pass_context) == "/env"
if __name__ == "__main__":
test_kernel_output_dir_uses_pass_context_first()
test_kernel_output_dir_falls_back_to_env()1 participant
该 PR 增加通过环境变量配置 MACA kernel 输出目录的能力，方便在容器和 CI 中隔离编译产物并归档调试文件。
这个修改面向沐曦 GPU 适配场景中比较容易影响开发、构建或验证稳定性的环节，把原来需要人工排查的问题前移到工具链、运行前检查或基准脚本中处理。实现上保持对现有默认行为的兼容，只在检测到明确配置、输入或环境异常时给出更直接的诊断，避免引入额外运行依赖，也方便维护者独立审阅该分支。
已在沐曦算力环境中完成对应分支验证，验证记录包含真实运行日志、命令输出和失败路径检查，本地归档目录为：E:/Documents/muxi/测试报告/mcTVM_new_toolchain_validation_20260608。提交分支：
mengz/env-maca-kernel-output-dir，目标仓库：MetaX-MACA/mcTVM。