diff --git a/CHANGELOG.md b/CHANGELOG.md index ad98f76ba7..bdbed968c0 100644 --- a/CHANGELOG.md +++ b/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)). diff --git a/warp/_src/codegen.py b/warp/_src/codegen.py index 2a525c576d..772ed223a0 100644 --- a/warp/_src/codegen.py +++ b/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 diff --git a/warp/_src/context.py b/warp/_src/context.py index 12dd12c4e6..d7ea60d022 100644 --- a/warp/_src/context.py +++ b/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 diff --git a/warp/tests/test_codegen.py b/warp/tests/test_codegen.py index 9a8ebe7ecc..f677988f65 100644 --- a/warp/tests/test_codegen.py +++ b/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`` diff --git a/warp/tests/test_context.py b/warp/tests/test_context.py index dcf932d03a..f975b01b3f 100644 --- a/warp/tests/test_context.py +++ b/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) diff --git a/warp/tests/test_generics.py b/warp/tests/test_generics.py index e1e87f2a1e..fbdf73ee0c 100644 --- a/warp/tests/test_generics.py +++ b/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 )