[AMDGPU] Address-space-at-source global memory optimization - #866
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5cba42cc23
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Tag AMDGPU device pointers as addrspace(1) where they are materialized so InferAddressSpaces propagates the tag and dependent loads/stores lower to global_load/global_store/global_atomic instead of generic flat_*. Layer A (codegen_amdgpu.cpp): source-side tagging of ExternalPtr / GlobalTemporary data pointers and address-space-preserving MatrixPtr GEPs. Layer B (codegen_llvm.cpp/.h): carry global provenance through the shared argument-pointer walk via maybe_tag_amdgpu_global_ptr() (get_args_ptr, get_struct_arg, ExternalPtrStmt) so per-iteration re-casts collapse, and preserve dest address space in atomic_op_using_cas. All tagging is gated on QD_AMDGPU_GLOBAL_AS and arch-guarded to AMDGPU, so CPU/CUDA emit byte-identical IR whether or not the patch is present. Reconstructs AMD-Ecosystem PRs Genesis-Embodied-AI#7 (6c9889f, 2d78505) + Genesis-Embodied-AI#3 (4e59c11) plus the argument-walk provenance fix, rebased onto current upstream.
get_args_ptr() is shared between the top-level kernel and @qd.real_func (Function) callees. Only the top-level kernel's arg buffer is device-staged in global memory; a Function callee receives a caller-local alloca buffer (see visit(FuncCallStmt)), so casting it to addrspace(1) under QD_AMDGPU_GLOBAL_AS made the callee read its scalar parameters through global memory pointing at private storage. Guard the args-buffer tag on the callable not being a Function. The ndarray data-pointer tags in get_struct_arg() and ExternalPtrStmt are unaffected, so top-level kernel promotion (and the measured speedups) are preserved: the franka gate-ON optimized IR is byte-identical before and after this change. Fixes the get_args_ptr provenance issue flagged in review. Co-authored-by: Cursor <cursoragent@cursor.com>
6b3b56c to
31c6a64
Compare
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 31c6a644f3
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
Checked with Opus. Here are its comments: 
Please could you address Opus concerns, in addition to addressin the codex comment I flagged earlier, and removing the env var gating. |
…params Address review on Genesis-Embodied-AI#866: - Remove the QD_AMDGPU_GLOBAL_AS env-var gate and make global address-space tagging the default AMDGPU behavior (per maintainer request to drop the gate and the old flat path). maybe_tag_amdgpu_global_ptr keeps its arch==amdgpu guard, so CPU/CUDA remain byte-identical. This also moots the cache-key and env-var-docs review comments, since there is no longer a mode toggle. - Fix get_struct_arg() to not tag Function (@qd.real_func) callee args as addrspace(1): qd.ref(...) reference parameters point at a caller-local alloca (private memory), so tagging them global was unsound (Codex P1 / test_ref). Mirrors the existing get_args_ptr() Function guard. Top-level ndarray promotion is unaffected; in-callee ndarray data pointers are still tagged at ExternalPtrStmt where the global-backed data pointer is loaded. - Remove the now-redundant AMDGPU visit(ExternalPtrStmt) override: the shared base codegen already tags ndarray data pointers for AMDGPU. - Add tests/python/test_amdgpu_global_as.py: an AMDGPU-only IR-shape regression guard asserting flat_* -> global_* promotion (global_* present, flat_* near zero) so a future refactor can't silently drop the optimization. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Pushed 593b93d addressing the review (@hughperkins, Codex, and the Opus notes). Core change
Codex P1s
Opus notes
Validation note: the code is lint-clean but I have not re-run the AMD suite on this revision yet; I'll post MI300X results for Hardware validation — done (AMD Instinct MI308X, Built from source with the AMDGPU backend ( 14 passed in 14.5s. The f16 cases ran (not skipped) — the device advertises 16-bit atomics, so Opus's f16-CAS path is actually exercised. I validated on an MI308X rather than an MI300X; both are One build fix required (pushed as |
maybe_tag_amdgpu_global_ptr() returns llvm::Value*, but `loaded` was
auto-deduced from CreateLoad() as LoadInst*, so reassigning the tagged
pointer failed to compile under clang ("incompatible pointer types
assigning to 'LoadInst *' from 'llvm::Value *'"). Declare `loaded` as
llvm::Value* so the addrspace-tagged result can be assigned back.
Caught building with the AMDGPU backend (ROCm clang 22) on gfx942.
Co-authored-by: Cursor <cursoragent@cursor.com>
|
@codex review |
|
Codex Review: Didn't find any major issues. Keep them coming! Reviewed commit: ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
Trim the address-space-at-source comments to document only the non-obvious rationale (why addrspace(1) tagging plus InferAddressSpaces, and why real_func callees are excluded), dropping restatements of the code and change history.
Remove comments deducible from the code (the tag call, the helper body, the functional-sanity label) and keep only the gotchas: why the tag exists, why MatrixPtr/atomic-CAS preserve addrspace, and why real_func callees are excluded.
|
@codex review |
|
actually, I'll probably ignore any codex cometns, since it already came back green before, and it's hard for me to test any chagnes. |
|
I think I'll just wait for amdgpu ci to pass, then merge. |
|
Codex Review: Didn't find any major issues. Another round soon, please! Reviewed commit: ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
3 participants
Summary
On the AMDGPU backend, device buffers are currently materialized as generic (flat) pointers. As a result, most memory traffic in generated kernels lowers to
flat_load/flat_storeinstructions, which go through the slower generic-addressing path even though the underlying memory is always global. This PR tags device pointers asaddrspace(1)(global) at the point where they are materialized, so that LLVM'sInferAddressSpacespass can propagate the global address space through the kernel and promoteflat_*operations toglobal_*operations.The change is fully gated behind the
QD_AMDGPU_GLOBAL_ASenvironment variable and arch-guarded to AMDGPU only. When the gate is unset, code generation is byte-for-byte identical to upstream (verified — see Test Plan). CPU and CUDA backends are completely untouched.What changes
The work is split into two layers:
Layer A — source-side tagging (
codegen/amdgpu/codegen_amdgpu.cpp)AMDGPU-specific codegen overrides that tag pointers as
addrspace(1)where they are first materialized:ExternalPtrStmt— ndarray data pointersGlobalTemporaryStmt— runtime global-temporary buffer pointersMatrixPtrStmt— preserves the origin address space through the byte-offset GEP, avoidingptrtoint/inttoptrround-trips that would otherwise strip the address space.Layer B — argument-walk provenance (
codegen/llvm/codegen_llvm.cpp,.h)Shared LLVM-layer changes that carry the global address space along the pointer-derivation chain, so the tag survives from the argument struct down to the element access and doesn't get re-stripped per iteration. A single choke-point helper,
maybe_tag_amdgpu_global_ptr(), applies the tag only when the arch is AMDGPU and the gate is set; it is a no-op otherwise. This is applied inget_struct_arg(),get_args_ptr(), andvisit(ExternalPtrStmt), andatomic_op_using_cas()is adjusted to preserve the destination's address space when forming the integer pointer for CAS loops (preventing invalid cross-addrspace bitcasts).Diff footprint: 3 files, +103 / −5.
Effect on generated IR
With the gate on, generic (flat) loads collapse dramatically as the global address space propagates through the kernel — on the franka monolith kernel, flat loads dropped from 10,768 → 503, with the remainder promoted to global-addressed operations.
Relationship to #775
#775 (
amdgpu-ieee/amdgpu-dx10-clampattribute harmonization) is complementary but not required. This PR was built and validated on upstreammainwithout #775 present, and delivers its speedups independently. The two compose cleanly and can land in either order.Test Plan
All testing was performed on an MI300X (gfx942) node under ROCm/HIP, with GPU clocks locked to 1900 MHz for measurement determinism. Correctness and performance were measured A/B (gate ON vs. gate OFF) using the genesis
test_rigidbenchmark harness.InferAddressSpacespromotes flat memory operations to global as intended (franka monolith: 10,768 → 503 flat loads).