Support Callable parameters in user-defined functions

CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ### Added
 
+- Add support for `Callable`-typed parameters in user-defined `@wp.func` functions, allowing user-defined Warp functions
+  and simple built-in Warp functions such as `wp.sin()` and `wp.min()` to be used as callable targets from kernels or
+  other functions, including through defaults ([GH-1424](https://github.com/NVIDIA/warp/issues/1424)).
 - 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)).

docs/user_guide/basics.rst

@@ -335,6 +335,69 @@ User functions may also be overloaded by defining multiple function signatures w
     def custom(x: wp.vec3):
         return x + wp.vec3(1.0, 0.0, 0.0)
 
+.. _callable-parameters:
+
+Callable Parameters
+^^^^^^^^^^^^^^^^^^^
+
+User functions can accept another user-defined Warp function or simple built-in
+Warp function by annotating the parameter as ``Callable`` from
+``collections.abc`` or ``typing``.
+Callable targets are specialization inputs, not runtime values: the target must
+be known during code generation, and Warp generates a separate specialization
+for each target function. The callable target is chosen where the user function
+is called and can be invoked directly inside the user function body:
+
+.. code-block:: python
+
+    from collections.abc import Callable
+
+    import warp as wp
+
+    @wp.func
+    def square(x: float):
+        return x * x
+
+
+    @wp.func
+    def cube(x: float):
+        return x * x * x
+
+
+    @wp.func
+    def apply(f: Callable, x: float):
+        return f(x)
+
+
+    @wp.kernel
+    def apply_kernel(
+        values: wp.array[float],
+        square_out: wp.array[float],
+        cube_out: wp.array[float],
+    ):
+        i = wp.tid()
+        square_out[i] = apply(square, values[i])
+        cube_out[i] = apply(cube, values[i])
+
+Parameterized callable annotations such as ``Callable[[float], float]`` are
+accepted, but Warp currently treats them the same as bare ``Callable``. The
+argument and return types in the annotation are not validated against the target
+function signature. The callable target is checked only through the actual calls
+made in the function body during code generation.
+
+Callable parameters may also use defaults and keyword arguments:
+
+.. code-block:: python
+
+    @wp.func
+    def apply_default(f: Callable = square, x: float = 0.0):
+        return f(x)
+
+Pass only user-defined :func:`@wp.func <warp.func>` functions or simple built-in
+functions such as ``wp.sin``, ``wp.cos``, ``wp.sqrt``, ``wp.add``, and ``wp.min``
+as callable targets. See :doc:`limitations` for unsupported callable targets and
+other restrictions.
+
 Tiles may also be passed to user functions. The function signature tile argument should include
 dtype and shape parameters to match the tile type intended to be used in the function. For example:
 

docs/user_guide/limitations.rst

@@ -35,6 +35,15 @@ Kernels and User Functions
   (e.g., ``wp.float64(wp.PI)`` or ``wp.int64(large_value)``).
 * Python ``IntFlag`` values behave like raw integers in Warp kernels: bitwise negation (``~``)
   produces the integer negation, not a masked combination of flags as in standard Python ``IntFlag`` behavior.
+* :ref:`Callable parameters <callable-parameters>` in user functions only support direct inline calls with
+  specialization-time targets: user-defined :func:`@wp.func <warp.func>` functions and simple built-in Warp functions
+  such as ``wp.sin``, ``wp.cos``, ``wp.sqrt``, ``wp.add``, and ``wp.min``.
+  Arbitrary Python callables and built-in Warp functions that require special code-generation behavior, dispatch logic,
+  or side effects, such as ``wp.printf``, are not supported.
+  Rebinding a function-valued local to a different function or to a non-function value is not supported.
+  User functions with ``Callable`` parameters also cannot define custom gradient or replay functions.
+  Callable argument and return types in annotations such as ``Callable[[float], float]`` are accepted but are not
+  validated against the target function signature.
 
 A limitation of Warp is that each dimension of the grid used to launch a kernel must be representable as a 32-bit
 signed integer. Therefore, no single dimension of a grid should exceed :math:`2^{31}-1`.