[Misc] Harden AITER_ASM_DIR code-object loading

[fjankovi]

[Misc] Harden AITER_ASM_DIR code-object loading

Note: this change and this description were drafted with the assistance of an AI coding assistant, and reviewed by me before opening. Opened as a draft — see the unchecked hardware-testing box under Testing.

Summary

AiterAsmKernel resolves its .co from $AITER_ASM_DIR and hands the bytes to __hipRegisterFatBinary. In builds that embed their code objects at compile time, that environment variable silently took precedence over the embedded copy. This makes the embedded objects authoritative where they exist, and adds the validation the disk path was missing.

Motivation

A security scan flagged the env-to-code-object flow in aiter_hip_common.h as CWE-426 (untrusted search path). The concrete gap is the precedence order rather than the existence of the env var:

• A build that went to the trouble of embedding vetted code objects would still discard them the moment AITER_ASM_DIR was set in the environment, with no warning, and load an arbitrary file from the named directory instead.
• The consumer this affects today is PyTorch's ck_sdpa target, which compiles this header with -DAITER_EMBEDDED_HSA_HEADER and generates the embedded map from third_party/aiter/hsa at build time. Nothing in PyTorch sets AITER_ASM_DIR, so the variable is purely an external input there.

I want to be precise about the severity, because I do not think this is critical: setting an environment variable in the victim process already implies at least as much control as LD_PRELOAD or PYTHONPATH would give. No privilege boundary is crossed and aiter offers no env sandbox. This is defence in depth — the useful property is that a build which ships vetted binaries cannot be silently redirected away from them — not the closing of a live escalation path.

Separately, the disk path had a genuine crash: tellg() returns -1 on failure and was assigned straight into a size_t, so an unseekable file asked new[] for SIZE_MAX.

Changes

Precedence. In a build with an embedded map, the embedded object is used. AITER_ASM_DIR is honoured only if the builder opts in with -DAITER_ALLOW_ASM_DIR_OVERRIDE; taking the override logs a warning naming the directory and the kernel. Without the opt-in, a set AITER_ASM_DIR is ignored and logs a warning saying so, rather than being silently dropped.

No behaviour change for standalone aiter. Builds with no embedded map load from disk exactly as before, including the existing fail-closed AITER_CHECK when the variable is unset. This is the path the aiter package itself uses — aiter/jit/core.py sets AITER_ASM_DIR to the packaged hsa/ directory on import — so the normal pip install aiter flow is untouched. Gating the env var off by default would have broken it.

Disk-path validation, applied on both the standalone and opted-in override paths:

• validate_relative_hsaco_path() rejects a hsaco name that is empty, absolute, or contains a .. component, so the join onto the configured directory cannot escape it. .. as a prefix (..foo.co) is still fine; only a whole component is rejected.
• ELF magic is required, so a truncated or wrong-format file fails here with a clear message instead of somewhere inside the HIP loader.
• Size is bounded at 64 MiB (shipped objects are a few hundred KB), which is also what fixes the SIZE_MAX allocation above.

Refactor. The disk-loading body moved into load_hsaco_from_dir() so both paths share it; load_hsaco_file() is now just the source-selection policy.

Deliberate omissions

• No realpath/std::filesystem canonicalisation against an install prefix. This header is included from HIP TUs, and there is no install prefix known to the C++ layer in a standalone build — the directory comes from Python. Rejecting .. in the relative part gets the containment property that is actually available here without adding a <filesystem> dependency to a widely-included header.
• No signature verification. That would need a key-distribution story that does not exist today, and would not match how every other .so in the stack is loaded.

Testing

• Header syntax-checked in all three configurations it now has (standalone; embedded; embedded + AITER_ALLOW_ASM_DIR_OVERRIDE), -std=c++20 --offload-arch=gfx942, no warnings from the changed code.
• validate_relative_hsaco_path() and the load checks extracted into a standalone harness, 19 cases, all passing — accepts fmha_v3_bwd/foo.co, ..foo.co, a/..b/c.co; rejects "", /abs/foo.co, .., ../foo.co, a/../../etc/passwd, a/b/../c.co; accepts a well-formed ELF and rejects empty, truncated, wrong-magic and missing files.
• clang-format clean over the changed region.
• Not run on hardware — no GPU on this machine. A maintainer should confirm an ordinary standalone aiter run still loads kernels normally, since that path is the one most people exercise.

No unit tests added: the logic is in a header-only private member of a class that requires a live HIP device to instantiate, and there is no existing C++ test target here to hang it off. Happy to add one if you would like it wired up.

Documentation

• Not updated. If AITER_ALLOW_ASM_DIR_OVERRIDE looks right to you, it should probably be mentioned wherever AITER_ASM_DIR is documented for downstream embedders — op_tests/cpp/mha/README.md is the only place the variable appears in-tree today, and that flow is unaffected.

Dependencies

• No new third-party dependencies.

Breaking Changes

None for standalone aiter or for any build without an embedded map.

For a build that embeds code objects and relied on AITER_ASM_DIR to override them, the override now requires -DAITER_ALLOW_ASM_DIR_OVERRIDE at compile time. I believe PyTorch is the only such consumer today and it does not set the variable, so I expect no downstream fallout — but this is the part worth a maintainer's eye, since you will know of embedders I do not.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4739 --add-label <label>