[AMDGPU] Feat: per-kernel LLVM function attributes via @qd.kernel(fn_attrs=...)

docs/source/user_guide/optimization_passes.md

@@ -81,6 +81,38 @@ All of these are fields of `CompileConfig`, so you set them at `qd.init(...)` (o
 
 For everyday use, leave them at their defaults. The most common deliberate change is `cfg_optimization=False` when iterating on a kernel whose compile time is in your way. Note that, in general, changing these options is relatively fragile since the Quadrants tests run assuming the default values.
 
+## 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.
+
+```python
+@qd.kernel(fn_attrs={"amdgpu": {"amdgpu-waves-per-eu": "1,4"}})
+def my_kernel():
+    ...
+
+@qd.kernel(fn_attrs={"amdgpu": {"amdgpu-max-num-workgroups": "128,1,1"}})
+def another_kernel():
+    ...
+```
+
+**How it works:** attributes are applied inside `codegen_llvm.cpp` before the JIT sees the module. The JIT's own defaults (e.g. `amdgpu-waves-per-eu`, `amdgpu-ieee`) are set only when the kernel does *not* already carry the attribute, so `fn_attrs` values always win.
+
+**Supported backends and attribute names** are registered in `quadrants/program/fn_attrs_registry.h`. Passing an unknown backend name or an unregistered attribute name raises `QuadrantsSyntaxError` at decoration time (not at compile time), so mistakes surface early.
+
+Currently registered attributes for the `amdgpu` backend:
+
+| Attribute | Effect |
+|-----------|--------|
+| `amdgpu-waves-per-eu` | Target waves-per-execution-unit range, e.g. `"1,2"` |
+| `amdgpu-flat-work-group-size` | Flat work-group size range, e.g. `"64,256"` |
+| `amdgpu-max-num-workgroups` | Maximum number of workgroups per dispatch, e.g. `"128,1,1"` |
+| `amdgpu-agpr-alloc` | AGPR allocation hint |
+| `amdgpu-sched-strategy` | Instruction scheduling strategy |
+
+**Cache interaction:** changing `fn_attrs` produces a different fastcache key, so kernels with different attribute sets are compiled and cached independently.
+
+**Other backends:** only `amdgpu` is currently registered. Adding a new backend means adding it to `fn_attrs_registry.h` and ensuring the consuming codegen path picks it up.
+
 ## Inspecting what the compiler did
 
 These environment variables dump the IR so you can see the effect of each pass. Files are written to `debug_dump_path` (default `/tmp/ir/`):

python/quadrants/lang/_fast_caching/src_hasher.py

@@ -27,13 +27,15 @@ def create_cache_key(
     kernel_source_info: FunctionSourceInfo,
     args: Sequence[Any],
     arg_metas: Sequence[ArgMetadata],
+    fn_attrs: dict[str, dict[str, str]] | None = None,
 ) -> str | None:
     """
     cache key takes into account:
     - arg types
     - cache value arg values
     - kernel function (but not sub functions)
     - compilation config (which includes arch, and debug)
+    - per-kernel fn_attrs (set via @qd.kernel(fn_attrs=...))
     """
     args_hash = args_hasher.hash_args(raise_on_templated_floats, args, arg_metas)
     if isinstance(args_hash, FastcacheSkip):
@@ -47,6 +49,7 @@ def create_cache_key(
         return None
     kernel_hash = function_hasher.hash_kernel(kernel_source_info)
     config_hash = config_hasher.hash_compile_config()
+    fn_attrs_hash = json.dumps(fn_attrs or {}, sort_keys=True)
     cache_key = hash_iterable_strings(
         (
             quadrants.__version_str__,
@@ -60,6 +63,7 @@ def create_cache_key(
             # Fast-cache value schema version. Bump when CacheValue's stored fields change so stale entries are not
             # mis-read. v2: graph_do_while single-arg -> nested level table.
             _CACHE_VALUE_SCHEMA_VERSION,
+            fn_attrs_hash,
         )
     )
     return cache_key

python/quadrants/lang/kernel.py

@@ -336,6 +336,7 @@ def __init__(self, _func: Callable, autodiff_mode: AutodiffMode, _is_classkernel
         self.use_checkpoints: bool = False
         # Legacy single-loop arg name, kept for reporting/back-compat; equals the outermost level's condition arg for
         # nested kernels. The authoritative data is `graph_do_while_levels`.
+        self.fn_attrs: dict[str, dict[str, str]] | None = None
         self.graph_do_while_arg: str | None = None
         # Nested graph_do_while level table, indexed by level id (outer before inner). Rebuilt each compilation pass by
         # the AST transformer; serialized to the launch context at launch.
@@ -394,7 +395,7 @@ def _try_load_fastcache(self, args: tuple[Any, ...], key: "CompiledKernelKeyType
         if self.runtime.src_ll_cache and self.quadrants_callable and self.quadrants_callable.is_pure:
             kernel_source_info, _src = get_source_info_and_src(self.func)
             self.fast_checksum = src_hasher.create_cache_key(
-                self.raise_on_templated_floats, kernel_source_info, args, self.arg_metas
+                self.raise_on_templated_floats, kernel_source_info, args, self.arg_metas, self.fn_attrs
             )
             used_py_dataclass_parameters = None
             cached_graph_do_while_levels: list[tuple[str, int]] | None = None
@@ -508,6 +509,8 @@ def materialize(self, key: "CompiledKernelKeyType | None", py_args: tuple[Any, .
                 quadrants_kernel = impl.get_runtime().prog.create_kernel(
                     quadrants_ast_generator, kernel_name, self.autodiff_mode
                 )
+                if self.fn_attrs:
+                    quadrants_kernel.set_fn_attrs(self.fn_attrs)
                 if _pass == 1:
                     assert key not in self.materialized_kernels
                     self.materialized_kernels[key] = quadrants_kernel

python/quadrants/lang/kernel_impl.py

@@ -5,6 +5,7 @@
 from functools import update_wrapper, wraps
 from typing import Any, Callable, TypeVar, cast, overload
 
+from quadrants._lib import core as _qd_core
 from quadrants.lang import impl
 from quadrants.lang.exception import (
     QuadrantsCompilationError,
@@ -153,13 +154,34 @@ def _inside_class(level_of_class_stackframe: int) -> bool:
     return False
 
 
+def _validate_fn_attrs(fn_attrs: dict[str, dict[str, str]], func_name: str) -> None:
+    registry = _qd_core.get_fn_attrs_registry()
+    for backend, attrs in fn_attrs.items():
+        if backend not in registry:
+            raise QuadrantsSyntaxError(
+                f"@qd.kernel(fn_attrs=...) on {func_name!r}: unknown backend "
+                f"{backend!r}. Registered backends: {sorted(registry)}."
+            )
+        allowed = registry[backend]
+        for key in attrs:
+            if key not in allowed:
+                raise QuadrantsSyntaxError(
+                    f"@qd.kernel(fn_attrs=...) on {func_name!r}: unknown "
+                    f"{backend} attribute {key!r}. Registered attributes: "
+                    f"{sorted(allowed)}."
+                )
+
+
 def _kernel_impl(
     _func: Callable,
     level_of_class_stackframe: int,
     verbose: bool = False,
     graph: bool = False,
     checkpoints: bool = False,
+    fn_attrs: dict[str, dict[str, str]] | None = None,
 ) -> QuadrantsCallable:
+    if fn_attrs:
+        _validate_fn_attrs(fn_attrs, _func.__name__)
     # Can decorators determine if a function is being defined inside a class?
     # https://stackoverflow.com/a/8793684/12003165
     is_classkernel = _inside_class(level_of_class_stackframe + 1)
@@ -171,6 +193,8 @@ def _kernel_impl(
     primal.use_graph = graph
     primal.use_checkpoints = checkpoints
     adjoint.use_checkpoints = checkpoints
+    primal.fn_attrs = fn_attrs
+    adjoint.fn_attrs = fn_attrs
     # Having |primal| contains |grad| makes the tape work.
     primal.grad = adjoint
 
@@ -213,7 +237,12 @@ def wrapped_classkernel(*args, **kwargs):
 # TODO: This callable should be Callable[[F], F].
 # See comments below.
 def kernel(
-    _fn: None = None, *, pure: bool = False, graph: bool = False, checkpoints: bool = False
+    _fn: None = None,
+    *,
+    pure: bool = False,
+    graph: bool = False,
+    checkpoints: bool = False,
+    fn_attrs: dict[str, dict[str, str]] | None = None,
 ) -> Callable[[Any], Any]: ...
 
 
@@ -224,7 +253,14 @@ def kernel(
 # However, by making it return Any, we can make the pure parameter
 # change now, without breaking pyright.
 @overload
-def kernel(_fn: Any, *, pure: bool = False, graph: bool = False, checkpoints: bool = False) -> Any: ...
+def kernel(
+    _fn: Any,
+    *,
+    pure: bool = False,
+    graph: bool = False,
+    checkpoints: bool = False,
+    fn_attrs: dict[str, dict[str, str]] | None = None,
+) -> Any: ...
 
 
 def kernel(
@@ -234,6 +270,7 @@ def kernel(
     fastcache: bool = False,
     graph: bool = False,
     checkpoints: bool = False,
+    fn_attrs: dict[str, dict[str, str]] | None = None,
 ):
     """
     Marks a function as a Quadrants kernel.
@@ -278,7 +315,13 @@ def decorator(fn: F, has_kernel_params: bool = True) -> F:
                 f"@qd.kernel({fn.__name__!r}, checkpoints=True) requires graph=True; "
                 "the checkpoint resume model is only meaningful for graph kernels."
             )
-        wrapped = _kernel_impl(fn, level_of_class_stackframe=level, graph=graph, checkpoints=checkpoints)
+        wrapped = _kernel_impl(
+            fn,
+            level_of_class_stackframe=level,
+            graph=graph,
+            checkpoints=checkpoints,
+            fn_attrs=fn_attrs,
+        )
         wrapped.is_pure = pure is not None and pure or fastcache
         if pure is not None:
             warnings_helper.warn_once(

quadrants/analysis/offline_cache_util.cpp

@@ -10,6 +10,7 @@
 
 #include "picosha2.h"
 
+#include <algorithm>
 #include <vector>
 
 namespace quadrants::lang {
@@ -184,14 +185,44 @@ std::string get_hashed_offline_cache_key(const CompileConfig &config,
 
   auto compile_config_key = get_offline_cache_key_of_compile_config(config);
   auto device_caps_key = get_offline_cache_key_of_device_caps(caps);
-  std::string autodiff_mode = std::to_string(static_cast<std::size_t>(kernel->autodiff_mode));
+  std::string autodiff_mode =
+      std::to_string(static_cast<std::size_t>(kernel->autodiff_mode));
+  // fn_attrs (set via @qd.kernel(fn_attrs=...)) affect codegen and must
+  // participate in the cache key, otherwise two kernels with identical IR
+  // but different attribute values collide. Iterate in sorted order so the
+  // hash is deterministic across unordered_map rehashes.
+  std::string fn_attrs_string;
+  {
+    std::vector<std::string> backends;
+    backends.reserve(kernel->fn_attrs.size());
+    for (const auto &kv : kernel->fn_attrs)
+      backends.push_back(kv.first);
+    std::sort(backends.begin(), backends.end());
+    for (const auto &backend : backends) {
+      fn_attrs_string += backend;
+      fn_attrs_string += '\0';
+      const auto &attrs = kernel->fn_attrs.at(backend);
+      std::vector<std::string> keys;
+      keys.reserve(attrs.size());
+      for (const auto &kv : attrs)
+        keys.push_back(kv.first);
+      std::sort(keys.begin(), keys.end());
+      for (const auto &k : keys) {
+        fn_attrs_string += k;
+        fn_attrs_string += '\0';
+        fn_attrs_string += attrs.at(k);
+        fn_attrs_string += '\0';
+      }
+    }
+  }
   picosha2::hash256_one_by_one hasher;
   hasher.process(compile_config_key.begin(), compile_config_key.end());
   hasher.process(device_caps_key.begin(), device_caps_key.end());
   hasher.process(kernel_params_string.begin(), kernel_params_string.end());
   hasher.process(kernel_rets_string.begin(), kernel_rets_string.end());
   hasher.process(kernel_body_string.begin(), kernel_body_string.end());
   hasher.process(autodiff_mode.begin(), autodiff_mode.end());
+  hasher.process(fn_attrs_string.begin(), fn_attrs_string.end());
   hasher.finish();
 
   auto res = picosha2::get_hash_hex_string(hasher);

quadrants/codegen/llvm/codegen_llvm.cpp

@@ -3242,10 +3242,16 @@ LLVMCompiledTask TaskCodeGenLLVM::run_compilation() {
       tlctx->mark_function_as_cuda_kernel(func, task.block_dim);
     }
   } else if (compile_config.arch == Arch::amdgpu) {
+    const auto fn_attrs_it = kernel->fn_attrs.find("amdgpu");
     for (const auto &task : offloaded_tasks) {
       llvm::Function *func = module->getFunction(task.name);
       QD_ASSERT(func);
-      tlctx->mark_function_as_amdgpu_kernel(func);
+      tlctx->mark_function_as_amdgpu_kernel(func, task.block_dim);
+      if (fn_attrs_it != kernel->fn_attrs.end()) {
+        for (const auto &kv : fn_attrs_it->second) {
+          func->addFnAttr(kv.first, kv.second);
+        }
+      }
     }
 #if defined(QD_WITH_AMDGPU)
     llvm::legacy::FunctionPassManager fpm(module.get());

quadrants/program/fn_attrs_registry.h

@@ -0,0 +1,33 @@
+#pragma once
+
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+
+namespace quadrants::lang {
+
+// Central registry of LLVM function attributes that users may set per-kernel
+// via `@qd.kernel(fn_attrs={...})`. Keyed by backend short name (matching
+// the namespacing used in the Python dict, e.g. "amdgpu"). Adding a new
+// attribute means: (1) register it here, (2) ensure the consuming backend
+// honors it (for AMDGPU: codegen_llvm.cpp applies it via addFnAttr; if the
+// backend has a hardcoded default for the same key, gate that default with
+// `!F.hasFnAttribute(key)` in the corresponding jit_*.cpp).
+inline const std::unordered_map<std::string, std::unordered_set<std::string>> &
+get_fn_attrs_registry() {
+  static const std::unordered_map<std::string,
+                                  std::unordered_set<std::string>>
+      kRegistry = {
+          {"amdgpu",
+           {
+               "amdgpu-max-num-workgroups",
+               "amdgpu-agpr-alloc",
+               "amdgpu-waves-per-eu",
+               "amdgpu-flat-work-group-size",
+               "amdgpu-sched-strategy",
+           }},
+      };
+  return kRegistry;
+}
+
+}  // namespace quadrants::lang

quadrants/program/kernel.h

@@ -1,5 +1,8 @@
 #pragma once
 
+#include <string>
+#include <unordered_map>
+
 #include "quadrants/util/lang_util.h"
 #include "quadrants/ir/snode.h"
 #include "quadrants/ir/ir.h"
@@ -16,6 +19,15 @@ class QD_DLL_EXPORT Kernel : public Callable {
  public:
   std::vector<SNode *> no_activate;
 
+  // User-supplied per-backend LLVM function attributes, set via
+  // `@qd.kernel(fn_attrs=...)`. Outer key is the backend short name
+  // ("amdgpu", ...); inner map is attr_name -> attr_value. Allowed keys
+  // are validated against fn_attrs_registry.h. Each backend's codegen is
+  // responsible for picking up its own entry; entries for inactive
+  // backends are stored but unused.
+  std::unordered_map<std::string, std::unordered_map<std::string, std::string>>
+      fn_attrs;
+
   bool is_accessor{false};
 
   Kernel(Program &program,

quadrants/python/export_lang.cpp

@@ -11,12 +11,18 @@
 
 #include <nanobind/eigen/dense.h>
 #include <nanobind/eigen/sparse.h>
+// get_fn_attrs_registry() returns std::unordered_map<std::string,
+// std::unordered_set<std::string>>; export.h only pulls in the std::map caster,
+// so these are required for the return value to convert to a Python dict/set.
+#include <nanobind/stl/unordered_map.h>
+#include <nanobind/stl/unordered_set.h>
 
 #include "quadrants/ir/expression_ops.h"
 #include "quadrants/ir/frontend_ir.h"
 #include "quadrants/ir/statements.h"
 #include "quadrants/program/adstack_size_expr_eval.h"
 #include "quadrants/program/extension.h"
+#include "quadrants/program/fn_attrs_registry.h"
 #include "quadrants/program/ndarray.h"
 #include "quadrants/rhi/device_capability.h"
 #include "quadrants/program/matrix.h"
@@ -80,6 +86,11 @@ void export_lang(nb::module_ &m) {
 
   m.def("arch_name", arch_name);
   m.def("arch_from_name", arch_from_name);
+  m.def("get_fn_attrs_registry", []() {
+    // Returns {backend_name: {allowed_attr_names...}} for validating
+    // @qd.kernel(fn_attrs=...) at decoration time.
+    return get_fn_attrs_registry();
+  });
 
   nb::enum_<SNodeType>(m, "SNodeType", nb::is_arithmetic())
 #define PER_SNODE(x) .value(#x, SNodeType::x)
@@ -553,6 +564,31 @@ void export_lang(nb::module_ &m) {
              self->no_activate.push_back(snode);
            })
       .def("to_string", &Kernel::to_string)
+      .def("set_fn_attrs",
+           [](Kernel *self,
+              const std::unordered_map<
+                  std::string,
+                  std::unordered_map<std::string, std::string>> &fn_attrs) {
+             const auto &registry = get_fn_attrs_registry();
+             for (const auto &backend_kv : fn_attrs) {
+               const auto reg_it = registry.find(backend_kv.first);
+               if (reg_it == registry.end()) {
+                 throw std::invalid_argument(
+                     "Unknown fn_attrs backend: '" + backend_kv.first +
+                     "'. See quadrants/program/fn_attrs_registry.h.");
+               }
+               const auto &allowed = reg_it->second;
+               for (const auto &attr_kv : backend_kv.second) {
+                 if (!allowed.count(attr_kv.first)) {
+                   throw std::invalid_argument(
+                       "Unknown " + backend_kv.first + " fn_attr: '" +
+                       attr_kv.first +
+                       "'. See quadrants/program/fn_attrs_registry.h.");
+                 }
+               }
+             }
+             self->fn_attrs = fn_attrs;
+           })
       .def("insert_scalar_param", &Kernel::insert_scalar_param)
       .def("insert_arr_param", &Kernel::insert_arr_param)
       .def("insert_ndarray_param", &Kernel::insert_ndarray_param)