Repository navigation
Allow passing a key/name to the @wp.kernel #1561
Description
Description
@wp.kernel derives a kernel's key from make_full_qualified_name(func), which in turn uses the function's __qualname__. There's no way to override it.
The @wp.func decorator already exposes a name argument for the same purpose.
It would be great if we would extend this pattern to the wp.kernel, enabling overriding the kernel's name through an argument: @wp.kernel(key="custom_name")
Context
This is relevant when generating kernels using kernel closures. In this case, different kernels have the same key and only differ by their hashes:
import warp as wp
def make_scaler(factor: float) -> wp.Kernel:
@wp.kernel
def scale(x: wp.array(dtype=wp.float32)):
i = wp.tid()
x[i] = x[i] * factor
return scale
double = make_scaler(2.0)
triple = make_scaler(3.0)
print(double.key) # make_scaler__locals__scale
print(triple.key) # make_scaler__locals__scale <-- identical!In our case, we export these kernels with compile_aot_module with strip_hash=True. This results in a collision of the kernel names. We strip the hashes to enable more flexible access to the exported code.
Proposal
Expose a wp.kernel(key=...) or wp.kernel(name=...) (to be consistent with wp.func) argument, which overrides the kernel's name and skips the sanitation of the name with make_full_qualified_name.
A potential pitfall is if the user passes a name that is not a valid C name. We should make it clear in the docstring which requirements the name needs to satisfy. Alternatively, we could check the user's custom name and raise if it is not valid.
Happy to implement this, if there's interest!
Activity
thomasbbrunner commentedon Jun 19, 2026
Addressed in #1570
c0d1f1ed commentedon Jul 7, 2026
Thanks for filing this issue and providing a PR!
I'll take it for internal review and try to get it merged for the 1.16 release. I think I'll make the following refinements:
strip_hash=Truecauses distinct kernels to have the same exported native name, fail early with a clear message.wp.overloadhandling.EDIT: Let's get the PR merged. I can follow up with the refinements.
Add name argument to @wp.kernel (NVIDIAGH-1561)
Add name argument to @wp.kernel (NVIDIAGH-1561)
Add name argument to @wp.kernel (NVIDIAGH-1561)
Add name argument to @wp.kernel (NVIDIAGH-1561)
shi-eric commentedon Aug 20, 2026
@thomasbbrunner, thanks for filing this and for putting together #1570. We initially planned to merge that PR, but the kernel-naming change was ultimately folded into the related external-compilation MR instead, so #1570 was not merged separately.
I am leaving this issue open so you can verify that the merged behavior covers your closure-generated AOT use case.
Codex-generated summary
ffce876 landed experimental support for custom kernel names.
The merged implementation uses
@wp.kernel(name=...), consistent with@wp.func, and validates that custom names are valid C++ identifiers. It also handles custom-named overload lookup, including rejecting ambiguous decorator-form lookups. The existing AOT guard remains in place for generic overloads that would collide whenstrip_hash=True.Public API behavior
@wp.kernel(name="custom_name")overrides the derived kernel name.[A-Za-z_][A-Za-z0-9_]*.strip_hash=True, the custom name becomes the base of the exported native symbol.Test coverage
warp/tests/test_codegen.pyaddstest_kernel_name_preserves_warp_kernel_abi,test_generic_unique_module_hash_includes_kernel_name, andtest_kernel_name_rejects_invalid_identifier.warp/tests/test_generics.pyaddstest_custom_named_overload_stubandtest_ambiguous_custom_named_overload_stub.Behavioral verification
Two closure-generated kernels were given distinct custom names, launched successfully, and compiled together with
strip_hash=True. The resulting AOT object contained distinct exported symbols for both kernels.Follow-up to consider
strip_hash=Truewould pin the exact scenario reported here. The behavior passes in an issue-shaped probe, so this is non-blocking.Add hooks for external tools to register their functionality NVIDIAGHâĶ
thomasbbrunner commentedon Aug 27, 2026
Thanks for the update, @shi-eric. The changes seem to address this. I'll only try them out once the next release is published. We can close the issue and if this ends up not working, we can re-open it.
Add hooks for external tools to register their functionality NVIDIAGHâĶ