Support wp.Function parameters (GH-1424)

wp.Function parameters let user functions accept another Warp
function as a specialization argument. Before this, helpers had to
hard-code the callee, which made higher-order patterns unavailable and
left specialization-sensitive behavior untested.

Use wp.Function as the public annotation so the API describes the
supported target set directly. Keep Callable handling internal for
built-in signatures while updating codegen, module hashing, dependency
tracking, generated stubs, docs, and tests around function target
specialization.

Signed-off-by: Eric Shi <ershi@nvidia.com>

Author: shi-eric

CHANGELOG.md

@@ -16,6 +16,9 @@
   to an addressable expression for use with native snippets, and allow `@wp.func_native` ref-parameter functions to
   provide manual adjoints with `adj_snippet`
   ([GH-1277](https://github.com/NVIDIA/warp/issues/1277)).
+- Add support for `wp.Function`-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 function 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

@@ -360,6 +360,67 @@ 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:
+
+Function Parameters
+^^^^^^^^^^^^^^^^^^^
+
+User functions can accept another user-defined Warp function or simple built-in
+Warp function by annotating the parameter as :class:`warp.Function`.
+Function targets are chosen when Warp generates code, not while the kernel is
+running. Warp compiles a separate version of the user function for each distinct
+set of target functions used at a call site. This is similar to generic
+specialization: many target combinations mean more specialized versions to
+compile, including for ``wp.grad()`` calls. The function target is chosen where
+the user function is called and can be invoked directly inside the user function
+body:
+
+.. code-block:: python
+
+    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: wp.Function, 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])
+
+The :class:`warp.Function` annotation is type-erased. It does not encode or
+validate the target function signature. The target is checked only through the
+actual calls made in the function body during code generation.
+
+Function parameters may also use defaults and keyword arguments:
+
+.. code-block:: python
+
+    @wp.func
+    def apply_default(f: wp.Function = 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 function targets. See :doc:`limitations` for unsupported function 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

@@ -37,6 +37,13 @@ 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:`Function parameters <callable-parameters>` in user functions only support direct inline calls with
+  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 are not supported. Some built-in Warp functions, such as ``wp.printf``, cannot be used
+  as ``wp.Function`` arguments because they need special handling during kernel compilation.
+  Rebinding a function-valued local to a different function or to a non-function value is not supported.
+  User functions with ``wp.Function`` parameters also cannot define custom gradient or replay functions.
 
 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`.