[AMDGPU] Address-space-at-source global memory optimization - #866
Open
paveltc wants to merge 2 commits into
Open
[AMDGPU] Address-space-at-source global memory optimization#866paveltc wants to merge 2 commits into
paveltc wants to merge 2 commits into
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>
paveltc
force-pushed
the
feat/amdgpu-global-as
branch
from
August 14, 2026 21:13
6b3b56c to
31c6a64
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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).