Add name argument to @wp.kernel (GH-1561)

CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Added
 
+- Add a `name` argument to `@wp.kernel` to override the kernel's generated key.
+  ([GH-1561](https://github.com/NVIDIA/warp/issues/1561)).
 - Add mipmap (texture level-of-detail) support to `wp.Texture1D`, `wp.Texture2D`, and `wp.Texture3D` via the new
   `num_mip_levels` and `mip_filter_mode` constructor parameters, and allow `wp.texture_sample()` to accept an optional
   trailing `lod` argument for controlling sampled detail level ([GH-1409](https://github.com/NVIDIA/warp/issues/1409)).

warp/_src/codegen.py

@@ -5131,6 +5131,14 @@ def make_full_qualified_name(func: str | Callable) -> str:
     return re.sub("[^0-9a-zA-Z_]+", "", func.replace(".", "__"))
 
 
+# matches a valid C identifier: a letter or underscore followed by letters, digits, or underscores
+_C_IDENTIFIER_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+
+
+def is_valid_c_identifier(name: str) -> bool:
+    return _C_IDENTIFIER_RE.match(name) is not None
+
+
 def codegen_struct(struct, device="cpu", indent_size=4):
     name = struct.native_name
 

warp/_src/context.py

@@ -805,8 +805,13 @@ def __init__(self, func, key=None, module=None, options=None, code_transformers=
         else:
             self.module = module
 
+        # Keep track of the qualified name of the Python function used to create
+        # this kernel, so it can still be identified even when the kernel is given
+        # a custom name (key).
+        self.qualname = warp._src.codegen.make_full_qualified_name(func)
+
         if key is None:
-            self.key = warp._src.codegen.make_full_qualified_name(func)
+            self.key = self.qualname
         else:
             self.key = key
 
@@ -1371,6 +1376,7 @@ def kernel2(
 def kernel(
     f: Callable | None = None,
     *,
+    name: str | None = None,
     enable_backward: bool | None = None,
     launch_bounds: tuple[int, ...] | int | None = None,
     module: Module | Literal["unique"] | str | None = None,
@@ -1419,6 +1425,9 @@ def my_kernel_fast(a: wp.array[float], b: wp.array[float]):
 
     Args:
         f: The function to be registered as a kernel.
+        name: Overrides the kernel's key. If ``None``, the function's
+            qualified name will be used as the key. The name must be a valid
+            C identifier; raises a ``ValueError`` otherwise.
         enable_backward: If False, the backward pass will not be
             generated.
         launch_bounds: CUDA ``__launch_bounds__`` attribute for the
@@ -1448,6 +1457,13 @@ def my_kernel_fast(a: wp.array[float], b: wp.array[float]):
     def wrapper(f, *args, **kwargs):
         kernel_options = {}
 
+        if name is None:
+            key = warp._src.codegen.make_full_qualified_name(f)
+        else:
+            if not warp._src.codegen.is_valid_c_identifier(name):
+                raise ValueError(f"@wp.kernel for '{f.__name__}': name '{name}' is not a valid C identifier.")
+            key = name
+
         if enable_backward is not None:
             kernel_options["enable_backward"] = enable_backward
 
@@ -1493,7 +1509,7 @@ def wrapper(f, *args, **kwargs):
         # Create the kernel object and register it with the module
         k = Kernel(
             func=f,
-            key=warp._src.codegen.make_full_qualified_name(f),
+            key=key,
             module=m,
             options=kernel_options,
         )
@@ -2997,7 +3013,27 @@ def _detach_references(self):
     # find kernel corresponding to a Python function
     def _find_kernel(self, func):
         qualname = warp._src.codegen.make_full_qualified_name(func)
-        return self.kernels.get(qualname)
+
+        # A kernel's key matches the qualified name unless the user specifies a
+        # custom name for the kernel.
+        # First, check the fast path: kernels registered under their ``qualname``.
+        kernel = self.kernels.get(qualname)
+        if kernel is not None:
+            return kernel
+
+        # Otherwise, the kernel was renamed. Fall back to matching via the
+        # ``qualname`` attribute.
+        # In this case, several kernels may share the same qualified name (e.g.
+        # kernels generated in a closure). We take the most recently registered
+        # kernel as the intended target. This assumes, that the ``@wp.overload``
+        # is defined just after the overloaded kernel, ``self.kernels`` preserves
+        # insertion order and that the ``self.kernels`` keys aren't overwritten.
+        match = None
+        for kernel in self.kernels.values():
+            if kernel.qualname == qualname:
+                match = kernel
+
+        return match
 
     # collect all referenced functions / structs
     # given the AST of a function or kernel

warp/tests/test_codegen.py

@@ -1645,6 +1645,14 @@ def test_line_directive_preserves_normalized_path(self):
 
         self.assertEqual(directive, '#line 1 "C:/warp/kernels/example.py"')
 
+    def test_is_valid_c_identifier(self):
+        for name in ("foo", "_foo", "Foo123", "_", "a1_b2"):
+            with self.subTest(name=name):
+                self.assertTrue(codegen.is_valid_c_identifier(name))
+        for name in ("", "1leading_digit", "has space", "has-dash", "a.b", "a+b", "naïve"):
+            with self.subTest(name=name):
+                self.assertFalse(codegen.is_valid_c_identifier(name))
+
     def test_extract_function_source_slow_path_when_fast_returns_none(self):
         """When ``_try_extract_function_source`` returns ``None`` (e.g. for an
         ``exec``-defined function with no linecache entry), ``extract_function_source``

warp/tests/test_context.py

@@ -16,6 +16,37 @@ def test_context_type_str(self):
         self.assertEqual(wp._src.context.type_str(tuple[int, float]), "tuple[int, float]")
         self.assertEqual(wp._src.context.type_str(tuple[int, ...]), "tuple[int, ...]")
 
+    def test_kernel_name_override(self):
+        """@wp.kernel(name=...) should override the kernel's key."""
+
+        @wp.kernel(name="custom_kernel_name")
+        def some_kernel(x: wp.array(dtype=float)):
+            i = wp.tid()
+            x[i] = x[i]
+
+        self.assertEqual(some_kernel.key, "custom_kernel_name")
+
+    def test_kernel_name_default(self):
+        """Omitting name should fall back to the qualified name."""
+
+        @wp.kernel
+        def default_named_kernel(x: wp.array(dtype=float)):
+            i = wp.tid()
+            x[i] = x[i]
+
+        self.assertTrue(default_named_kernel.key.endswith("default_named_kernel"))
+
+    def test_kernel_name_invalid(self):
+        """A name that is not a valid C identifier should raise ValueError."""
+        with self.assertRaises(ValueError) as cm:
+
+            @wp.kernel(name="has space")
+            def bad_kernel(x: wp.array(dtype=float)):
+                i = wp.tid()
+                x[i] = x[i]
+
+        self.assertIn("not a valid C identifier", str(cm.exception))
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

warp/tests/test_generics.py

@@ -330,6 +330,47 @@ def nested_overload_stub_kernel(x: float): ...
     test.assertEqual(len(nested_overload_stub_kernel.overloads), 1)
 
 
+def test_overload_with_renamed_kernel(test, device):
+    @wp.kernel(name="renamed_overload_kernel")
+    def scale(x: Any):
+        return
+
+    @wp.overload
+    def scale(x: float): ...
+
+    @wp.overload
+    def scale(x: int): ...
+
+    test.assertEqual(scale.key, "renamed_overload_kernel")
+    test.assertEqual(len(scale.overloads), 2)
+
+
+def test_overload_renamed_kernel_closure(test, device):
+    # A closure that yields one renamed kernel per call produces several kernels
+    # sharing the same qualified name. Decorator-form @wp.overload is defined
+    # inline next to each kernel, so it must bind to that call's kernel.
+    def make(suffix):
+        @wp.kernel(name=f"closure_overload_kernel_{suffix}")
+        def scale(x: Any):
+            return
+
+        @wp.overload
+        def scale(x: float): ...
+
+        @wp.overload
+        def scale(x: int): ...
+
+        return scale
+
+    first = make("a")
+    second = make("b")
+
+    test.assertEqual(first.key, "closure_overload_kernel_a")
+    test.assertEqual(second.key, "closure_overload_kernel_b")
+    test.assertEqual(len(first.overloads), 2)
+    test.assertEqual(len(second.overloads), 2)
+
+
 @wp.kernel
 def overload_stub_return_annotation_kernel(x: Any):
     return
@@ -622,6 +663,10 @@ class TestGenerics(unittest.TestCase):
 add_function_test(TestGenerics, "test_generic_conditional_setter", test_generic_conditional_setter, devices=devices)
 add_function_test(TestGenerics, "test_generic_transform_kernel", test_generic_transform_kernel, devices=devices)
 add_function_test(TestGenerics, "test_nested_overload_stub", test_nested_overload_stub, devices=devices)
+add_function_test(TestGenerics, "test_overload_with_renamed_kernel", test_overload_with_renamed_kernel, devices=devices)
+add_function_test(
+    TestGenerics, "test_overload_renamed_kernel_closure", test_overload_renamed_kernel_closure, devices=devices
+)
 add_function_test(
     TestGenerics, "test_overload_stub_return_annotation", test_overload_stub_return_annotation, devices=devices
 )