[AMDGPU] Feat: per-kernel LLVM function attributes via @qd.kernel(fn_attrs=...) by paveltc · Pull Request #774 · Genesis-Embodied-AI/quadrants

paveltc · 2026-07-10T05:14:31Z

Summary

Backport of ROCm#11 onto the upstream main branch.

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 attributes (amdgpu backend):

amdgpu-max-num-workgroups
amdgpu-agpr-alloc
amdgpu-waves-per-eu
amdgpu-flat-work-group-size
amdgpu-sched-strategy

Example 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 pairs
quadrants/program/kernel.h: adds fn_attrs field to Kernel
quadrants/python/export_lang.cpp: exposes set_fn_attrs + get_fn_attrs_registry pybindings
quadrants/codegen/llvm/codegen_llvm.cpp: applies per-kernel fn_attrs via addFnAttr at codegen
quadrants/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 by hasFnAttribute so user-supplied values always win
quadrants/runtime/llvm/llvm_context.{h,cpp}: mark_function_as_amdgpu_kernel gains optional block_dim parameter to set amdgpu-flat-work-group-size at codegen time
Python layer (kernel.py, kernel_impl.py, src_hasher.py): fn_attrs plumbed through @qd.kernel() decorator, validated at decoration time, included in fastcache key
tests/python/test_fn_attrs.py (new): validation, cache-key differentiation, and AMDGPU end-to-end IR plumbing tests

Merge notes

Backported on top of commit 2e2b43b40 (same base as PR #1 and PR #2). Conflicts resolved by combining HEAD's graph/checkpoints/use_graph/use_checkpoints additions with this commit's fn_attrs additions — all parameters coexist.

Test results

Full suite run in Docker image quadrants:amd-upstream-pr3 with QD_AMDGPU_FORCE_PERMLANE64_FALLBACK=1:

21 failed, 4992 passed, 142 skipped, 28 xfailed  (1:46:01)

All 21 failures are pre-existing baseline flaky tests in test_simt.py::test_block_reduce and test_algorithms.py::test_reduce_composition — identical failure set as the unpatched amd-upstream baseline. 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

…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

💡 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".

chatgpt-codex-connector · 2026-07-10T05:18:44Z

+    F.addFnAttr("unsafe-fp-math", "true");
+    F.addFnAttr("no-signed-zeros-fp-math", "true");


Honor fast_math=False before forcing unsafe AMDGPU math

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 👍 / 👎.

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.

chatgpt-codex-connector · 2026-07-10T05:18:44Z

+    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);


Keep flat-work-group-size within the launched block size

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 👍 / 👎.

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".

chatgpt-codex-connector · 2026-07-10T05:18:44Z

    graph: bool = False,
    checkpoints: bool = False,
+    cuda_graph: bool = False,
+    fn_attrs: dict[str, dict[str, str]] | None = None,


Add docs for the new fn_attrs kernel API

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 👍 / 👎.

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>

hughperkins · 2026-07-13T13:51:24Z


+## 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.


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.

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.

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.

hughperkins

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>

chatgpt-codex-connector Bot reviewed Jul 10, 2026

View reviewed changes

paveltc changed the title ~~feat(amdgpu): per-kernel LLVM function attributes via @qd.kernel(fn_attrs=...)~~ [AMDGPU] feat: per-kernel LLVM function attributes via @qd.kernel(fn_attrs=...) Jul 10, 2026

paveltc changed the title ~~[AMDGPU] feat: per-kernel LLVM function attributes via @qd.kernel(fn_attrs=...)~~ [AMDGPU] Feat: per-kernel LLVM function attributes via @qd.kernel(fn_attrs=...) Jul 10, 2026

paveltc mentioned this pull request Jul 10, 2026

[AMDGPU] Fix: propagate amdgpu-ieee and amdgpu-dx10-clamp to all functions #775

Closed

hughperkins reviewed Jul 13, 2026

View reviewed changes

Comment thread python/quadrants/lang/kernel_impl.py Outdated

hughperkins reviewed Jul 13, 2026

View reviewed changes

hughperkins requested changes Jul 13, 2026

View reviewed changes

hughperkins added the awaiting-contributor-action awaiting-contributor-action label Jul 13, 2026

ptcherni and others added 2 commits July 13, 2026 16:14

paveltc requested a review from hughperkins July 13, 2026 23:21

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[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=...)#774
paveltc wants to merge 4 commits into
Genesis-Embodied-AI:mainfrom
paveltc:feat/amdgpu-per-kernel-fn-attrs

paveltc commented Jul 10, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Jul 10, 2026

Uh oh!

paveltc Jul 10, 2026

Uh oh!

chatgpt-codex-connector Bot Jul 10, 2026

Uh oh!

paveltc Jul 10, 2026

Uh oh!

Uh oh!

chatgpt-codex-connector Bot Jul 10, 2026

Uh oh!

paveltc Jul 10, 2026

Uh oh!

Uh oh!

hughperkins Jul 13, 2026

Uh oh!

paveltc Jul 13, 2026

Uh oh!

hughperkins Jul 14, 2026

Uh oh!

hughperkins left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		F.addFnAttr("unsafe-fp-math", "true");
		F.addFnAttr("no-signed-zeros-fp-math", "true");


		## 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.

Uh oh!

Conversation

paveltc commented Jul 10, 2026

Summary

Changes

Merge notes

Test results

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

paveltc Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

paveltc Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

chatgpt-codex-connector Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

paveltc Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

hughperkins Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

paveltc Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

hughperkins Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

hughperkins left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants