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!
Description
@wp.kernelderives a kernel's key frommake_full_qualified_name(func), which in turn uses the function's__qualname__. There's no way to override it.The
@wp.funcdecorator already exposes anameargument 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:
In our case, we export these kernels with
compile_aot_modulewithstrip_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=...)orwp.kernel(name=...)(to be consistent withwp.func) argument, which overrides the kernel's name and skips the sanitation of the name withmake_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!