[AMDGPU] Feat: per-kernel LLVM function attributes via @qd.kernel(fn_attrs=...) - #774
[AMDGPU] Feat: per-kernel LLVM function attributes via @qd.kernel(fn_attrs=...)#774paveltc wants to merge 4 commits into
Conversation
…I#11) Lets users override AMDGPU codegen attributes per kernel without editing the JIT pipeline. Attributes must be pre-registered in quadrants/program/fn_attrs_registry.h; unknown backend or attribute names raise QuadrantsSyntaxError at decoration time. Currently registered: amdgpu-max-num-workgroups, amdgpu-agpr-alloc, amdgpu-waves-per-eu, amdgpu-flat-work-group-size. Examples: ```python @qd.kernel(fn_attrs={"amdgpu": {"amdgpu-max-num-workgroups": "128,1,1"}}) def k(...): ... @qd.kernel(fn_attrs={"amdgpu": {"amdgpu-waves-per-eu": "1,2"}}) def k(...): ... ``` Plumbed Python decorator -> Kernel.fn_attrs -> set_fn_attrs pybind -> codegen_llvm.cpp addFnAttr -> jit_amdgpu.cpp (defaults gated by hasFnAttribute so user values win). Included in both fastcache and frontend offline cache keys so changing fn_attrs forces a rebuild.
![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: 0f44102799
ℹ️ 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".
| F.addFnAttr("unsafe-fp-math", "true"); | ||
| F.addFnAttr("no-signed-zeros-fp-math", "true"); |
There was a problem hiding this comment.
In AMDGPU runs configured with fast_math=False, these unconditional attributes still enable unsafe FP transformations and no-signed-zero handling on every function; CompileConfig::fast_math is used just below to set strict target options, and docs/source/user_guide/init_options.md documents fast_math as controlling IEEE-relaxed rewrites. This means tests/users that disable fast math for numerical stability on AMDGPU can still get contracted or reordered arithmetic.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in commit 9cb65ab. unsafe-fp-math and no-signed-zeros-fp-math are now gated inside if (this->config_.fast_math) in jit_amdgpu.cpp, matching the TargetOptions block below and preserving IEEE semantics when qd.init(fast_math=False) is set.
| int clamped = std::max(block_dim, kAmdgpuWavefrontSize); | ||
| std::string size_str = std::to_string(clamped) + "," + std::to_string(clamped); | ||
| func->addFnAttr("amdgpu-flat-work-group-size", size_str); |
There was a problem hiding this comment.
For AMDGPU tasks whose actual block size is below 64, e.g. emit_amdgpu_gc() sets the reinit_lists task to block_dim = 1 and the launcher dispatches task.block_dim, this advertises amdgpu-flat-work-group-size as 64,64. The existing LLVM-context comment says this attribute is the min/max workgroup sizes specified at dispatch, so the IR metadata is false for those kernels and can drive invalid occupancy or optimization assumptions.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in commit 9cb65ab. The wavefront-size clamping is removed; mark_function_as_amdgpu_kernel now uses block_dim verbatim (block_dim,block_dim) so tasks with block_dim=1 correctly produce amdgpu-flat-work-group-size="1,1" instead of "64,64".
| graph: bool = False, | ||
| checkpoints: bool = False, | ||
| cuda_graph: bool = False, | ||
| fn_attrs: dict[str, dict[str, str]] | None = None, |
There was a problem hiding this comment.
This exposes fn_attrs as a new public @qd.kernel usage and allows backend-specific values, but the diff contains no docs/ changes. The repository AGENTS.md requires docs updates for public API or usage changes; without them, end users have no guidance for supported backends/attribute names or how these attributes interact with caching.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in commit 9cb65ab. Added a 'Per-kernel LLVM function attributes' section to docs/source/user_guide/optimization_passes.md covering usage, the registered AMDGPU attribute table, and cache interaction.
…s PR Four fixes from Codex review of PR Genesis-Embodied-AI#774: 1. Gate unsafe-fp-math on config_.fast_math Previously unconditional in jit_amdgpu.cpp; now matches the TargetOptions block below it. qd.init(fast_math=False) now correctly preserves IEEE semantics on AMDGPU. 2. Fix block_dim clamping in mark_function_as_amdgpu_kernel Prior impl clamped both min and max to max(block_dim, 64), so block_dim=1 produced "64,64" instead of "1,1". Now uses block_dim as-is so the IR accurately reflects the actual dispatch size. 3. Make cuda_graph a deprecated alias for graph cuda_graph=True was silently ignored because the launch path only checks use_graph. Now _kernel_impl treats cuda_graph=True as graph=True with a DeprecationWarning, and removes the now-unused use_cuda_graph field from Kernel. 4. Add fn_attrs docs to optimization_passes.md Per AGENTS.md, public API changes need docs/ updates. Adds a "Per-kernel LLVM function attributes" section documenting fn_attrs, the supported amdgpu attributes, and cache interaction. Co-authored-by: Cursor <cursoragent@cursor.com>
paveltc
changed the title
paveltc
changed the title
|
|
||
| ## Per-kernel LLVM function attributes (`fn_attrs`) | ||
|
|
||
| For fine-grained control over AMDGPU codegen on a per-kernel basis, `@qd.kernel` accepts an optional `fn_attrs` argument. This lets you pass LLVM function attributes directly to the backend JIT without editing the JIT pipeline. |
There was a problem hiding this comment.
Platform-specific configuration conflicts with the ideal that Quadrants code will run unmodified on all platforms.
The standard approach is to set appropriate defaults - or use appropriate heuristics - within each platform-specific pipeline.
There was a problem hiding this comment.
Thanks — I agree with the principle, and I think this PR is consistent with it. The sensible AMDGPU defaults already live in the backend pipeline (jit_amdgpu.cpp, each gated by hasFnAttribute); fn_attrs is just an optional override on top for when a user knows better. Portable code never needs it.
Portability is also preserved: attributes are only applied under arch == Arch::amdgpu (codegen_llvm.cpp:3244), so @qd.kernel(fn_attrs={"amdgpu": {...}}) still runs unmodified on CPU/CUDA/Metal — the block is simply ignored there. It's namespaced by backend for exactly that reason, much like CUDA __launch_bounds__ or Triton num_warps: an optional hint that only tunes the backend that understands it.
To land the right scope, which is your concern:
- the concept of any user-facing backend override, or
- specifically exposing it on
@qd.kernel(...)?
If (1), I'll narrow this PR to just the backend-default infrastructure and drop the public fn_attrs= for now. If (2), I'll move it behind a clearly-marked advanced API. Just say which unblocks the merge.
There was a problem hiding this comment.
My concern is anything platform-specific in the public API.
What I want to avoid is code that runs on one platform, and not on other; or on all platforms except one.
This includes optimizations, ideally. Code should ideally run optimally on each platform without the user needing platform-specific knowledge.
There was a problem hiding this comment.
Agreed as the default, and it holds for the common case. For peak perf I don't think it's attainable, though — no production GPU stack manages it. CUDA and HIP ship launch_bounds precisely so the user can state an occupancy/register tradeoff the compiler can't infer; Triton has num_warps/num_stages/@autotune; Halide/TVM make the schedule explicit. The vendors ship these knobs because the "sufficiently smart compiler" doesn't exist for this.
On the "AMD-only" objection: I'd argue that's expected, not a defect. Tuning is vendor-owned — AMD owns the AMD backend and its lowering and can't be expected to implement the NVIDIA/Intel/Apple equivalents; each vendor adds their own. That's how the ecosystem is already structured (torch.backends.cuda/.cudnn/.mps, LLVM's namespaced amdgpu-/nvvm. attributes).
And crucially, the current design already answers the "runs on one platform, breaks on another" worry:
-
It's namespaced by backend — fn_attrs={"amdgpu": {...}}, applied only under arch == Arch::amdgpu (codegen_llvm.cpp:3244). @qd.kernel(fn_attrs={"amdgpu": {...}}) runs unmodified on CPU/CUDA/Metal; the block is just ignored. Portable code is unaffected; only the backend that understands the hint tunes on it — same model as launch_bounds.
-
It's registry-validated, not an open escape hatch — _validate_fn_attrs rejects unknown backends and unknown attributes against get_fn_attrs_registry(), so the surface is curated and enumerable rather than "any LLVM attribute goes."
So it's an optional, explicitly-scoped, validated per-backend hint — opt-in, ignored elsewhere, and owned by the vendor whose backend it tunes. That seems to satisfy the underlying concern (no silent cross-platform breakage, no unbounded surface) while accepting that vendor-specific tuning is a legitimate and necessary thing to expose.
There was a problem hiding this comment.
Interesting. Currently, in genesis-world, we are not - I think - doing such tuning, for either CPU or CUDA. Perhaps we should. Or perhaps we are doing so implicitly.
I would much prefer to ground design and needs on real-world needs, ideally.
Do you see anywhere in genesis-world where we would be able to improve our performance, concretely, on our rigid benchmarks https://github.com/Genesis-Embodied-AI/genesis-world/blob/main/tests/benchmarks/test_rigid.py , for CUDA and AMD, by doing such tuning? And crucially, do you see anywhere where such tuning would use different numbers for each of CUDA and AMD?
This would give a concrete example to inspect and discuss, rather than discussing in the abstract.
(You could also provide a non genesis-world example, but, if it's about as much work for you to use genesis-world vs some other example, genesis-world will be a lot easier for me personally to understand and gain intuition and insight from).
There was a problem hiding this comment.
Benchmark: amdgpu-waves-per-eu via this PR's per-kernel fn_attrs
I benchmarked the knob this PR exposes on a Genesis workload — the anymal_random
benchmark in tests/benchmarks/test_rigid.py, 30,000 parallel envs, MI308X GPU.
I applied amdgpu-waves-per-eu="2,4" to the rigid-body solver kernels via
@qd.kernel(fn_attrs={"amdgpu": {"amdgpu-waves-per-eu": "2,4"}}) and compared against the
same kernels with no attribute (compiler default occupancy):
| Workload | speedup vs default | consistency |
|---|---|---|
anymal_random (30k envs, GPU) |
+48% (geomean) | faster in 5/5 rounds |
This node runs the CG constraint solver (the parametrization passes solver=None). Each
round runs default and tuned back-to-back and takes the ratio within the round, averaged
with a geometric mean (30 s warmup + 30 s record, 5 rounds).
What was modified: the feature under test is this PR's fn_attrs= API in Quadrants. To
exercise it, the only change on the Genesis side is passing fn_attrs=... on the
rigid-solver kernel decorators — the anymal_random workload and the runtime_fps metric
are the stock upstream test_rigid.py. The 5-round interleaved averaging is an external
wrapper around repeated stock runs (to cancel GPU clock drift); the longer warmup/record
windows are an optional tweak and not required to see the effect.
amdgpu-waves-per-eu directly tells the register allocator the target occupancy
(resident wavefronts per EU), trading latency-hiding against registers-per-wave. This is
a more direct handle than __launch_bounds__ (which HIP and CUDA already expose):
__launch_bounds__ constrains occupancy only indirectly, as a side effect of a block-size
cap and an optional min-blocks hint, whereas amdgpu-waves-per-eu sets the occupancy
target itself. That LLVM attribute isn't otherwise reachable from the JIT kernel path —
this PR is what makes it settable per kernel. The optimal point is kernel- and
workload-specific, so per-kernel granularity is the right level rather than a single
global flag, and the results show it has first-order performance impact on a real, public
workload — so exposing it is clearly worthwhile.
Reproduction
Requirements: Quadrants built from this PR, an AMD gfx942 GPU (MI3xx), and a Genesis build. test_rigid.py is run unmodified — the only source edit is adding the attribute to the rigid-solver kernel decorators (that's the feature under test).
Step 1. Add fn_attrs={"amdgpu": {"amdgpu-waves-per-eu": "2,4"}} to these @qd.kernel(...) decorators — kernel_step_1 and kernel_step_2 in genesis/engine/solvers/rigid/rigid_solver.py, and func_solve_init and _kernel_solve_monolith in genesis/engine/solvers/rigid/constraint/solver.py. For example, @qd.kernel(fastcache=True) becomes @qd.kernel(fastcache=True, fn_attrs={"amdgpu": {"amdgpu-waves-per-eu": "2,4"}}).
Step 2. Run the stock benchmark node once with the edit (tuned) and once without it (baseline):
python -m pytest -m benchmarks --backend gpu "tests/benchmarks/test_rigid.py::test_speed[anymal_random-None-None-30000-gpu]"
Step 3. Compare runtime_fps. A single run of each already shows the gain; averaging a few back-to-back rounds just tightens it against GPU clock noise.
hughperkins
left a comment
hughperkins
left a comment
There was a problem hiding this comment.
Platform-specific configuration conflicts with the ideal that Quadrants code will run unmodified on all platforms.
The standard approach is to set appropriate defaults - or use appropriate heuristics - within each platform-specific pipeline.
cuda_graph was never a public @qd.kernel parameter (the feature is `graph`), so introducing it only to deprecate it in favor of `graph` made no sense. It slipped in during the merge that combined the graph/checkpoints work with this PR's fn_attrs. Removes the parameter and its deprecation shim so this PR only touches fn_attrs. Co-authored-by: Cursor <cursoragent@cursor.com>
get_fn_attrs_registry() returns unordered_map<string, unordered_set<string>>, but export.h only pulls in the std::map caster. Without the unordered_map/ unordered_set casters the binding raised TypeError at runtime, so every @qd.kernel(fn_attrs=...) call failed and the fn_attrs validation tests could never pass in a real build. Adds the two caster includes to export_lang.cpp. Co-authored-by: Cursor <cursoragent@cursor.com>
Set amdgpu-ieee=false and amdgpu-dx10-clamp=false on all functions in the AMDGPU module rather than only AMDGPU_KERNEL entries. LLVM's inliner refuses to inline a callee into a caller when they carry mismatching target-specific attributes. Applying these attributes to kernels only leaves internal runtime device functions (e.g. gpu_parallel_range_for and the body functions it dispatches) with a mismatching attribute set, so they are not inlined into the kernel entry. Without that inlining, InferAddressSpaces cannot follow the pointer chain from kernel parameters to field data and cannot promote flat_load/flat_store/flat_atomic to global_*, causing flat-atomic coherency issues and a ~4% throughput regression on gfx942 (MI300X). Applying the two attributes uniformly restores inlining and lets InferAddressSpaces emit global_load/global_store/global_atomic. Each write is guarded by hasFnAttribute so it remains idempotent. Self-contained backend change: single file, no public API, no dependency on the per-kernel fn_attrs work in Genesis-Embodied-AI#774. Co-authored-by: Cursor <cursoragent@cursor.com>
4 participants
Summary
Backport of AMD-Ecosystem#11 onto the upstream
mainbranch.Lets users override AMDGPU codegen attributes per kernel without editing
the JIT pipeline. Attributes must be pre-registered in
quadrants/program/fn_attrs_registry.h; unknown backend or attribute namesraise
QuadrantsSyntaxErrorat decoration time.Currently registered attributes (
amdgpubackend):amdgpu-max-num-workgroupsamdgpu-agpr-allocamdgpu-waves-per-euamdgpu-flat-work-group-sizeamdgpu-sched-strategyExample usage:
```python
@qd.kernel(fn_attrs={"amdgpu": {"amdgpu-max-num-workgroups": "128,1,1"}})
def k(...): ...
@qd.kernel(fn_attrs={"amdgpu": {"amdgpu-waves-per-eu": "1,2"}})
def k(...): ...
```
Changes
quadrants/program/fn_attrs_registry.h(new): central registry of allowed backend/attr pairsquadrants/program/kernel.h: addsfn_attrsfield toKernelquadrants/python/export_lang.cpp: exposesset_fn_attrs+get_fn_attrs_registrypybindingsquadrants/codegen/llvm/codegen_llvm.cpp: applies per-kernel fn_attrs viaaddFnAttrat codegenquadrants/analysis/offline_cache_util.cpp: includes fn_attrs in the offline cache key (deterministic serialisation)quadrants/runtime/amdgpu/jit_amdgpu.cpp: default AMDGPU attributes (unsafe-fp-math,amdgpu-waves-per-eu,amdgpu-ieee,amdgpu-dx10-clamp, flat-work-group-size inheritance) gated byhasFnAttributeso user-supplied values always winquadrants/runtime/llvm/llvm_context.{h,cpp}:mark_function_as_amdgpu_kernelgains optionalblock_dimparameter to setamdgpu-flat-work-group-sizeat codegen timekernel.py,kernel_impl.py,src_hasher.py):fn_attrsplumbed through@qd.kernel()decorator, validated at decoration time, included in fastcache keytests/python/test_fn_attrs.py(new): validation, cache-key differentiation, and AMDGPU end-to-end IR plumbing testsMerge notes
Backported on top of commit
2e2b43b40(same base as PR #1 and PR #2). Conflicts resolved by combining HEAD'sgraph/checkpoints/use_graph/use_checkpointsadditions with this commit'sfn_attrsadditions — all parameters coexist.Test results
Full suite run in Docker image
quadrants:amd-upstream-pr3withQD_AMDGPU_FORCE_PERMLANE64_FALLBACK=1:All 21 failures are pre-existing baseline flaky tests in
test_simt.py::test_block_reduceandtest_algorithms.py::test_reduce_composition— identical failure set as the unpatchedamd-upstreambaseline. Zero new regressions.tests/python/test_fn_attrs.py(all 5 new tests) passed.Co-authored-by: kevinjosephamd kevin.joseph@amd.com
Made with Cursor