Added HSA header flag passthrough to compile.py#2245
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
There was a problem hiding this comment.
Pull request overview
Adds support for optionally compiling the MHA C++ test modules with an “embedded HSA” header, controlled by an environment variable, so builds can bundle/load HSACO data without relying on AITER_ASM_DIR.
Changes:
- Introduces
get_embedded_hsa_build_args()to derive additional compiler flags fromAITER_EMBEDDED_HSA_HEADER_PATH. - Plumbs the derived flags into both forward (
libmha_fwd) and backward (libmha_bwd) compile configurations.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| header_path_var = os.getenv("AITER_EMBEDDED_HSA_HEADER_PATH", "").strip() | ||
| flags_extra_cc = [] | ||
|
|
||
| # If header is provided, then make sure that include_dir is also provided, | ||
| # otherwise the embedded HSA header won't be found during compilation. | ||
| if header_path_var: | ||
| header_path = Path(header_path_var) | ||
| # aiter_hip_common.h uses: #include AITER_EMBEDDED_HSA_HEADER | ||
| # so the macro value must be a quoted header token. | ||
| flags_extra_cc.append(f'-DAITER_EMBEDDED_HSA_HEADER=\\"{header_path.name}\\"') | ||
|
|
||
| # Keep optCompilerConfig's default include list intact by passing include dir | ||
| # as a compile flag instead of overriding `extra_include` in custom args. | ||
| flags_extra_cc.append(f"-I{header_path.parent}") | ||
| return flags_extra_cc |
There was a problem hiding this comment.
AITER_EMBEDDED_HSA_HEADER_PATH may be provided as a relative path, but the actual compilation happens from the JIT build directory (ninja build dir), so -I{header_path.parent} can end up pointing at the wrong location and the header won’t be found. Consider normalizing to an absolute path (e.g., expanduser/resolve) before deriving parent, and validate that the path exists/is a file to fail fast with a clear error if misconfigured.
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Motivation
TE relies on
aiter/op_tests/cpp/mha/compile.pyto generatelibmha_*wd.so. These libraries are used in multi-threaded contexts, whereas their reliance on theAITER_ASM_DIRenvironment variable is thread-unsafe. To mitigate this, we would like to use the embedded HSA header option, but that requires some minor modifications tocompile.py.Technical Details
AITER_EMBEDDED_HSA_HEADER_PATHenvironment variable to generate necessary flags for utilizing AITER's HSA embedded header strategy.Test Plan
Test with TE build.
Test Result
Unable to test on exact commit yet, since TE is still pending updates to newer compatible versions of AITER, however these changes work on older commits of AITER on TE branches.
Submission Checklist