diff --git a/docs/source/user_guide/optimization_passes.md b/docs/source/user_guide/optimization_passes.md index c0c90ba318..d7d973ec00 100644 --- a/docs/source/user_guide/optimization_passes.md +++ b/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/`): diff --git a/python/quadrants/lang/_fast_caching/src_hasher.py b/python/quadrants/lang/_fast_caching/src_hasher.py index 76a6a0160a..f19261beff 100644 --- a/python/quadrants/lang/_fast_caching/src_hasher.py +++ b/python/quadrants/lang/_fast_caching/src_hasher.py @@ -27,6 +27,7 @@ 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: @@ -34,6 +35,7 @@ def create_cache_key( - 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 diff --git a/python/quadrants/lang/kernel.py b/python/quadrants/lang/kernel.py index e663c6ed9e..b194beb26e 100644 --- a/python/quadrants/lang/kernel.py +++ b/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 diff --git a/python/quadrants/lang/kernel_impl.py b/python/quadrants/lang/kernel_impl.py index d52222a6e9..c6d5b571a3 100644 --- a/python/quadrants/lang/kernel_impl.py +++ b/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,43 @@ 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, + cuda_graph: bool = False, + fn_attrs: dict[str, dict[str, str]] | None = None, ) -> QuadrantsCallable: + if cuda_graph and not graph: + import warnings + warnings.warn( + "@qd.kernel(cuda_graph=True) is deprecated; use graph=True instead.", + DeprecationWarning, + stacklevel=3, + ) + graph = True + 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 +202,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 +246,13 @@ 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, + cuda_graph: bool = False, + fn_attrs: dict[str, dict[str, str]] | None = None, ) -> Callable[[Any], Any]: ... @@ -224,7 +263,15 @@ 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, + cuda_graph: bool = False, + fn_attrs: dict[str, dict[str, str]] | None = None, +) -> Any: ... def kernel( @@ -234,6 +281,8 @@ def kernel( fastcache: bool = False, graph: bool = False, checkpoints: bool = False, + cuda_graph: bool = False, + fn_attrs: dict[str, dict[str, str]] | None = None, ): """ Marks a function as a Quadrants kernel. @@ -278,7 +327,14 @@ 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, + cuda_graph=cuda_graph, + fn_attrs=fn_attrs, + ) wrapped.is_pure = pure is not None and pure or fastcache if pure is not None: warnings_helper.warn_once( diff --git a/quadrants/analysis/offline_cache_util.cpp b/quadrants/analysis/offline_cache_util.cpp index e5bf0ddc39..791456c001 100644 --- a/quadrants/analysis/offline_cache_util.cpp +++ b/quadrants/analysis/offline_cache_util.cpp @@ -10,6 +10,7 @@ #include "picosha2.h" +#include #include namespace quadrants::lang { @@ -184,7 +185,36 @@ 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(kernel->autodiff_mode)); + std::string autodiff_mode = + std::to_string(static_cast(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 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 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()); @@ -192,6 +222,7 @@ std::string get_hashed_offline_cache_key(const CompileConfig &config, 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); diff --git a/quadrants/codegen/llvm/codegen_llvm.cpp b/quadrants/codegen/llvm/codegen_llvm.cpp index d085c5c6b3..6693051bdc 100644 --- a/quadrants/codegen/llvm/codegen_llvm.cpp +++ b/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()); diff --git a/quadrants/program/fn_attrs_registry.h b/quadrants/program/fn_attrs_registry.h new file mode 100644 index 0000000000..b494afdbaa --- /dev/null +++ b/quadrants/program/fn_attrs_registry.h @@ -0,0 +1,33 @@ +#pragma once + +#include +#include +#include + +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> & +get_fn_attrs_registry() { + static const std::unordered_map> + 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 diff --git a/quadrants/program/kernel.h b/quadrants/program/kernel.h index 73c5926029..9b2e3c43e4 100644 --- a/quadrants/program/kernel.h +++ b/quadrants/program/kernel.h @@ -1,5 +1,8 @@ #pragma once +#include +#include + #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 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> + fn_attrs; + bool is_accessor{false}; Kernel(Program &program, diff --git a/quadrants/python/export_lang.cpp b/quadrants/python/export_lang.cpp index e589ea907f..6452fc802b 100644 --- a/quadrants/python/export_lang.cpp +++ b/quadrants/python/export_lang.cpp @@ -17,6 +17,7 @@ #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 +81,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_(m, "SNodeType", nb::is_arithmetic()) #define PER_SNODE(x) .value(#x, SNodeType::x) @@ -553,6 +559,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> &fn_attrs) { + const auto ®istry = 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) diff --git a/quadrants/runtime/amdgpu/jit_amdgpu.cpp b/quadrants/runtime/amdgpu/jit_amdgpu.cpp index 7a2dccf783..673a1a9ad6 100644 --- a/quadrants/runtime/amdgpu/jit_amdgpu.cpp +++ b/quadrants/runtime/amdgpu/jit_amdgpu.cpp @@ -31,6 +31,102 @@ std::string JITSessionAMDGPU::compile_module_to_hsaco(std::unique_ptrconfig_.fast_math) { + F.addFnAttr("unsafe-fp-math", "true"); + F.addFnAttr("no-signed-zeros-fp-math", "true"); + } + + // Apply amdgpu-ieee and amdgpu-dx10-clamp to ALL functions, not just + // AMDGPU_KERNEL entries. LLVM's inliner refuses to inline a callee into a + // caller when they carry mismatching target-specific attributes, so keeping + // these on kernels only blocks runtime device functions (e.g. + // gpu_parallel_range_for) from being inlined. Without that inlining, + // InferAddressSpaces cannot see the full pointer chain from kernel params + // to field data and cannot promote flat_load/store/atomic to global_*, + // causing a ~4% throughput regression and flat-atomic coherency issues on + // MI300X. The hasFnAttribute guard ensures user-supplied fn_attrs still win. + if (!F.hasFnAttribute("amdgpu-ieee")) { + F.addFnAttr("amdgpu-ieee", "false"); + } + if (!F.hasFnAttribute("amdgpu-dx10-clamp")) { + F.addFnAttr("amdgpu-dx10-clamp", "false"); + } + + if (F.getCallingConv() == llvm::CallingConv::AMDGPU_KERNEL) { + const std::string kernel_name = F.getName().str(); + const bool is_lightweight_cg_subkernel = + kernel_name.find("_kernel_cg_only_save_prev_grad") != + std::string::npos || + kernel_name.find("_kernel_update_constraint_forces") != + std::string::npos || + kernel_name.find("_kernel_update_constraint_qfrc") != + std::string::npos || + kernel_name.find("_kernel_update_constraint_cost") != + std::string::npos || + kernel_name.find("_kernel_update_search_direction") != + std::string::npos; + + // Each default below is skipped if the kernel already carries that + // attribute (set upstream in codegen_llvm.cpp from user-supplied + // @qd.kernel(fn_attrs={...})). User values win. + if (!is_lightweight_cg_subkernel && + !F.hasFnAttribute("amdgpu-waves-per-eu")) { + F.addFnAttr("amdgpu-waves-per-eu", "1,2"); + } + if (!F.hasFnAttribute("uniform-work-group-size")) { + F.addFnAttr("uniform-work-group-size", "true"); + } + } + } + + for (auto &F : *llvm_module) { + if (F.isDeclaration() || F.empty()) + continue; + if (F.getCallingConv() == llvm::CallingConv::AMDGPU_KERNEL) + continue; + if (F.hasFnAttribute("amdgpu-flat-work-group-size")) + continue; // already set (e.g., on runtime kernels via + // mark_function_as_amdgpu_kernel-equivalent paths) + llvm::StringRef inherited; + for (auto *U : F.users()) { + auto *CB = llvm::dyn_cast(U); + if (!CB) + continue; + // Direct call only — function-pointer args (e.g., body fn passed + // as `RangeForTaskFunc *func` to gpu_parallel_range_for) are + // skipped because the use is the function pointer itself, not a + // call to it. `alwaysinline` on gpu_parallel_range_for + // collapses the function-pointer indirection so the body fn ends + // up with direct callers in the kernel entry. + if (CB->getCalledOperand() != &F) + continue; + auto *Caller = CB->getFunction(); + if (Caller && Caller->getCallingConv() == llvm::CallingConv::AMDGPU_KERNEL && + Caller->hasFnAttribute("amdgpu-flat-work-group-size")) { + inherited = + Caller->getFnAttribute("amdgpu-flat-work-group-size").getValueAsString(); + break; + } + } + if (inherited.empty()) + inherited = "1,128"; // conservative fallback + F.addFnAttr("amdgpu-flat-work-group-size", inherited); + } + + auto *daz_type = llvm::Type::getInt8Ty(llvm_module->getContext()); + auto *daz_init = llvm::ConstantInt::get(daz_type, 1); + auto *daz_var = new llvm::GlobalVariable( + *llvm_module, daz_type, true, llvm::GlobalValue::LinkOnceODRLinkage, + daz_init, "__oclc_daz_opt"); + daz_var->setVisibility(llvm::GlobalValue::HiddenVisibility); + + if (llvm::verifyModule(*llvm_module, &llvm::errs())) { llvm_module->print(llvm::errs(), nullptr); QD_WARN("Module broken"); diff --git a/quadrants/runtime/llvm/llvm_context.cpp b/quadrants/runtime/llvm/llvm_context.cpp index f67b84fbf4..8fdd5a26cf 100644 --- a/quadrants/runtime/llvm/llvm_context.cpp +++ b/quadrants/runtime/llvm/llvm_context.cpp @@ -1017,8 +1017,16 @@ void QuadrantsLLVMContext::mark_function_as_cuda_kernel(llvm::Function *func, in } } -void QuadrantsLLVMContext::mark_function_as_amdgpu_kernel(llvm::Function *func) { +void QuadrantsLLVMContext::mark_function_as_amdgpu_kernel(llvm::Function *func, int block_dim) { func->setCallingConv(llvm::CallingConv::AMDGPU_KERNEL); + if (block_dim > 0) { + // Record the actual dispatch size as the flat work-group-size range. + // Use block_dim as-is: advertising a larger range than what the launcher + // dispatches gives the backend false headroom to assume wider warps and + // can produce incorrect occupancy or wrong codegen for small block sizes. + std::string size_str = std::to_string(block_dim) + "," + std::to_string(block_dim); + func->addFnAttr("amdgpu-flat-work-group-size", size_str); + } } void QuadrantsLLVMContext::eliminate_unused_functions(llvm::Module *module, diff --git a/quadrants/runtime/llvm/llvm_context.h b/quadrants/runtime/llvm/llvm_context.h index 407e69775b..da7a8fb279 100644 --- a/quadrants/runtime/llvm/llvm_context.h +++ b/quadrants/runtime/llvm/llvm_context.h @@ -105,7 +105,7 @@ class QuadrantsLLVMContext { void mark_function_as_cuda_kernel(llvm::Function *func, int block_dim = 0); - void mark_function_as_amdgpu_kernel(llvm::Function *func); + void mark_function_as_amdgpu_kernel(llvm::Function *func, int block_dim = 0); void fetch_this_thread_struct_module(); llvm::Module *get_this_thread_runtime_module(); diff --git a/tests/python/test_fn_attrs.py b/tests/python/test_fn_attrs.py new file mode 100644 index 0000000000..24c1f66858 --- /dev/null +++ b/tests/python/test_fn_attrs.py @@ -0,0 +1,179 @@ +"""Tests for the @qd.kernel(fn_attrs=...) plumbing. + +Three layers of coverage: + +1. Validation: unknown backend / unknown attr name raise at decoration time. + Cheap, runs on any backend. +2. Cache differentiation: identical kernels with different fn_attrs values + produce different fastcache keys. Verifies the hash plumbing in + `src_hasher` (Python) and `offline_cache_util` (C++). +3. End-to-end IR plumbing: registered attribute appears verbatim in the + LLVM IR that jit_amdgpu sees. Gated on AMDGPU backend; runs the kernel + in a subprocess with print_kernel_llvm_ir=True and greps the dump. +""" + +import os +import pathlib +import subprocess +import sys + +import pytest + +import quadrants as qd +from quadrants.lang._fast_caching import src_hasher +from quadrants.lang._wrap_inspect import get_source_info_and_src +from quadrants.lang.exception import QuadrantsSyntaxError + +from tests import test_utils + +RET_SUCCESS = 42 + + +# --------------------------------------------------------------------------- +# 1. Validation +# --------------------------------------------------------------------------- + + +@test_utils.test(arch=[qd.cpu]) +def test_fn_attrs_unknown_backend_raises(): + with pytest.raises(QuadrantsSyntaxError, match="unknown backend 'nvidia'"): + + @qd.kernel(fn_attrs={"nvidia": {"some-attr": "x"}}) + def k(): + pass + + +@test_utils.test(arch=[qd.cpu]) +def test_fn_attrs_unknown_amdgpu_attr_raises(): + with pytest.raises(QuadrantsSyntaxError, match="unknown amdgpu attribute"): + + @qd.kernel(fn_attrs={"amdgpu": {"definitely-not-a-real-attr": "0,0"}}) + def k(): + pass + + +@test_utils.test(arch=[qd.cpu]) +def test_fn_attrs_registered_attr_accepts(): + # Should not raise; uses an attribute we know is in the registry. + @qd.kernel(fn_attrs={"amdgpu": {"amdgpu-max-num-workgroups": "1,1,1"}}) + def k(): + pass + + assert k._primal.fn_attrs == {"amdgpu": {"amdgpu-max-num-workgroups": "1,1,1"}} + + +# --------------------------------------------------------------------------- +# 2. Cache key differentiation (Python fastcache layer) +# --------------------------------------------------------------------------- + + +@test_utils.test(arch=[qd.cpu]) +def test_fn_attrs_changes_fastcache_key(): + # Same source, same args, same config — only fn_attrs differs. + # Cache keys must differ, otherwise a fn_attrs change would silently + # serve a stale CompiledKernelData. + def _kernel_source(): + @qd.kernel + def k(): + pass + + return k + + k = _kernel_source() + info, _ = get_source_info_and_src(k._primal.func) + + key_none = src_hasher.create_cache_key(False, info, (), (), None) + key_a = src_hasher.create_cache_key( + False, info, (), (), {"amdgpu": {"amdgpu-max-num-workgroups": "1,1,1"}} + ) + key_b = src_hasher.create_cache_key( + False, info, (), (), {"amdgpu": {"amdgpu-max-num-workgroups": "256,1,1"}} + ) + + assert key_none is not None and key_a is not None and key_b is not None + assert key_none != key_a, "fn_attrs=None must hash differently from fn_attrs={...}" + assert key_a != key_b, "different fn_attr values must produce different cache keys" + + +# --------------------------------------------------------------------------- +# 3. End-to-end IR plumbing — AMDGPU only +# --------------------------------------------------------------------------- + + +# Marker value — unique enough that grep finds it only on the kernel we set +# it on, not on any incidental codegen output. +_FN_ATTRS_MARKER = "1,7,42" + + +def _e2e_child(args: list[str]) -> None: + dump_dir = args[0] + os.chdir(dump_dir) # print_kernel_llvm_ir writes to CWD + qd.init(arch=qd.amdgpu, print_kernel_llvm_ir=True, offline_cache=False) + + # Use a non-trivial body so the kernel has a real offloaded task. An + # empty kernel may produce zero offloaded tasks, in which case the + # per-task addFnAttr loop in codegen_llvm.cpp's Arch::amdgpu branch + # has nothing to attach the attribute to and the test would falsely + # fail. + buf = qd.ndarray(qd.f32, shape=(64,)) + + @qd.kernel(fn_attrs={"amdgpu": {"amdgpu-max-num-workgroups": _FN_ATTRS_MARKER}}) + def with_attr(x: qd.types.ndarray(qd.f32, 1)): + for i in x: + x[i] = 1.0 + + with_attr(buf) + sys.exit(RET_SUCCESS) + + +@test_utils.test(arch=[qd.amdgpu]) +def test_fn_attrs_reaches_amdgpu_jit_ir(tmp_path: pathlib.Path): + # Run the kernel in a fresh subprocess so the LLVM IR dump lives in a + # clean directory we can grep without interference from other tests. + cmd = [sys.executable, __file__, _e2e_child.__name__, str(tmp_path)] + env = dict(os.environ) + env["PYTHONPATH"] = env.get("PYTHONPATH", "") + os.pathsep + "." + proc = subprocess.run(cmd, capture_output=True, text=True, env=env) + if proc.returncode != RET_SUCCESS: + print(proc.stdout) + print("-" * 80) + print(proc.stderr) + assert proc.returncode == RET_SUCCESS + + ll_files = list(tmp_path.glob("quadrants_kernel_amdgpu_llvm_ir_*.ll")) + assert ll_files, f"no LLVM IR dumps produced in {tmp_path}" + + # The marker should appear in an attribute group at the bottom of the IR + # as `"amdgpu-max-num-workgroups"="1,7,42"`. Search loosely (just the key + # and the value) so we tolerate any escaping/whitespace LLVM may apply. + matches = [] + for p in ll_files: + text = p.read_text() + if "amdgpu-max-num-workgroups" in text and _FN_ATTRS_MARKER in text: + matches.append(p) + + if not matches: + # Dump every `attributes #N` group from each .ll so the failure is + # actionable instead of opaque. + debug = [] + for p in ll_files: + attr_lines = [ + ln for ln in p.read_text().splitlines() if ln.startswith("attributes #") + ] + debug.append(f"--- {p.name} ---\n" + "\n".join(attr_lines[:20])) + pytest.fail( + f"expected 'amdgpu-max-num-workgroups' with value {_FN_ATTRS_MARKER!r} " + f"in at least one of {[p.name for p in ll_files]}, but absent.\n" + f"Attribute groups found:\n" + "\n\n".join(debug) + ) + + +# --------------------------------------------------------------------------- +# Subprocess dispatch (matches the pattern used in test_config.py). +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + name, *rest = sys.argv[1:] + { + _e2e_child.__name__: _e2e_child, + }[name](rest)