Support wp.Function parameters in user functions

[shi-eric]

Description

Closes #1424.

This PR adds phase-one support for annotating user-defined @wp.func function parameters with wp.Function. A wp.Function parameter names a Warp function selected at specialization time, so helper functions can call user-defined or supported built-in Warp functions passed by kernels or other functions.

Function targets are specialization inputs, not runtime values. User functions are specialized for the concrete target, function parameters are omitted from generated runtime signatures, and targets participate in module hashing and dependency invalidation.

Changes

• Add wp.Function as the public annotation for function-valued user function parameters.
• Specialize user-defined functions for targets passed through explicit arguments, keyword arguments, defaults, local aliases, nested calls, and wp.static(...).
• Allow supported simple built-in Warp functions such as wp.sin(), wp.sqrt(), wp.add(), and wp.min() as function targets.
• Include function targets in module hashes and dependency discovery so rebinding or unloading provider modules invalidates dependent generated code.
• Update generated stubs, user guide docs, limitations, changelog, and focused regression coverage in warp/tests/test_func_parameter_targets.py.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• CHANGELOG.md is updated for any user-facing changes under the Unreleased section.

Validation summary

Function parameter coverage lives in warp/tests/test_func_parameter_targets.py and is included in the default and debug unittest suites. The suite covers user-defined targets, simple built-in targets, defaults, keyword arguments, generic helpers, local aliases, nested calls, wp.static(...) targets, wp.grad() specialization, target return validation, module hashing, and dependency invalidation across imported provider modules.

Validation performed:

• WARP_CACHE_PATH=/tmp/warp-cache-function-rename-module uv run warp/tests/test_func_parameter_targets.py
• WARP_CACHE_PATH=/tmp/warp-cache-function-rename-suite uv run --extra dev -m warp.tests -s autodetect -k TestFuncParameterTargets
• WARP_CACHE_PATH=/tmp/warp-cache-function-rename-precommit uvx pre-commit run --files CHANGELOG.md docs/user_guide/basics.rst docs/user_guide/limitations.rst warp/__init__.pyi warp/_src/codegen.py warp/_src/context.py warp/_src/types.py warp/tests/unittest_suites.py warp/tests/aux_test_function_double.py warp/tests/aux_test_function_triple.py warp/tests/test_func_parameter_targets.py
• WARP_CACHE_PATH=/tmp/warp-cache-function-marker-docs uv run --extra docs build_docs.py 2>&1 | tee /tmp/build_docs.log

Bug fix

Without this PR, Warp functions cannot be passed through user-defined @wp.func helpers like this:

import warp as wp


@wp.func
def double_it(x: float):
    return x * 2.0


@wp.func
def apply(g: wp.Function, x: float):
    return g(x)


@wp.kernel
def k(out: wp.array[float]):
    out[0] = apply(double_it, 3.0)

New feature / enhancement

wp.Function parameters can be used for generic helpers that dispatch to different Warp function targets at specialization time:

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(g: wp.Function, x: float):
    return g(x)


@wp.kernel
def k(values: wp.array[float], square_out: wp.array[float], cube_out: wp.array[float]):
    i = wp.tid()
    x = values[i]
    square_out[i] = apply(square, x)
    cube_out[i] = apply(g=cube, x=x)

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

This PR implements support for Callable-typed parameters in user-defined @wp.func functions. Callable parameters are treated as compile-time function references that trigger per-call-site specialization; callable arguments/defaults are extracted and hashed for dependency tracking, bound into adjoint symbols, and omitted from emitted C++ signatures.

Changes

Callable Parameter Support

Layer / File(s)	Summary
Design Documentation design/callable-func-parameters.md	Design doc added: semantics, recognition rules, specialization, hashing, rejection cases, and test plan.
Callable Type Detection warp/_src/types.py	is_callable_annotation() added; get_type_code() uses it to emit "c" for typing/collections.abc callables.
Overload Instantiation warp/_src/context.py	Function.get_overload() preserved callable-template params and computes instantiated overload signature keys from callable-aware annotations (handles defaults).
Overload Matching warp/_src/codegen.py	func_match_args() updated to match callable-annotated params when bound argument is a Warp Function.
Specialization Helpers warp/_src/codegen.py	Added get_callable_arg_values(), get_default_arg_value(), specialize_callable_func(); shallow-copy/cache specialized function variants with hashed native_func suffixes.
Adjoint & Call Integration warp/_src/codegen.py	Adjoint.build() accepts callable_arg_values; adj.symbols bind callable params to concrete Function values. add_call() applies callable-aware defaults, triggers specialization, and omits callable params from runtime dispatch.
Codegen Signatures / Dispatch warp/_src/codegen.py	Forward and reverse signature generation skip callable-annotated parameters so callables are not emitted in the C++ ABI; dispatch filtering removes callable args from C++ calls.
Dependency Tracking warp/_src/context.py, warp/_src/codegen.py	Reference discovery resolves callable-typed arguments to concrete Function targets and records their modules for hashing and invalidation.
Tests warp/tests/test_func.py	Adds helpers, kernels, and unit tests covering typing/collections.abc callables, parameterized/generic/default/nested usage, module hash invalidation, return-annotation checks, dependency propagation, and rejection cases (builtins/custom grad/replay).
Docs / Changelog docs/user_guide/limitations.rst, CHANGELOG.md	Adds limitation note and changelog entry (GH-1424).

Sequence Diagram(s)

sequenceDiagram
    participant Kernel as Kernel (caller)
    participant Codegen as Codegen (add_call / func_match_args)
    participant Types as Types (is_callable_annotation)
    participant Specialize as Specialization (specialize_callable_func / get_callable_arg_values)
    participant Adjoint as Adjoint.build
    participant CppGen as C++ Generation

    Kernel->>Codegen: call apply(double_it, 3.0)
    Codegen->>Types: detect callable parameter via is_callable_annotation()
    Codegen->>Codegen: bind args, match overload (callable-aware)
    Codegen->>Specialize: extract callable args (get_callable_arg_values)
    Specialize->>Specialize: create/cache specialized clone (hash suffix)
    Codegen->>Adjoint: Adjoint.build(callable_arg_values)
    Adjoint->>Adjoint: populate adj.symbols with concrete Function values
    Codegen->>CppGen: emit call (skip callable params in forward/reverse signatures)

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 11.54% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The PR comprehensively addresses all coding requirements from issue #1424: unified Callable annotation handling, overload resolution fixes, type code generation ('c' code), callable-aware specialization, module hashing/dependency tracking, and comprehensive test coverage with documented limitations.
Out of Scope Changes check	✅ Passed	All changes are directly related to implementing Callable parameter support per GH-1424: type detection, overload matching, code generation, callable specialization, dependency tracking, tests, and documentation of limitations. No unrelated changes detected.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title captures the main change—supporting function-like parameters in user-defined functions—even though it is less specific than the Callable-based implementation.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (2)

warp/_src/codegen.py (1)

945-992: 💤 Low value

Consider documenting the specialization key stability for caching.

The specialization cache key at lines 952-954 uses a tuple of (name, callable_func) pairs, relying on Function object identity for cache hits. If the same logical function is recreated (e.g., module reload), cache misses will occur, generating duplicate specializations.

This is likely acceptable for correctness (each specialization is valid), but consider adding a brief comment explaining this design choice for future maintainers.

Also, the check at lines 946-950 correctly rejects functions with custom grad/replay—good defensive validation.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@warp/_src/codegen.py` around lines 945 - 992, The specialization cache key in
specialize_callable_func currently builds specialization_key from (name,
callable_arg_values[name]) and therefore relies on Function object identity for
cache hits; add a brief explanatory comment above the specialization_key
creation stating that this is an intentional design choice (it may produce cache
misses if the same logical function object is recreated, e.g., on module
reload), why that is acceptable for correctness, and note that using
callable_func.key/native_func could be an alternative for a stable key if
desired in the future; reference specialize_callable_func, specialization_key,
and func._callable_specializations in the comment so maintainers can find the
code easily.

CHANGELOG.md (1)

32-33: ⚡ Quick win

Use imperative “Add …” phrasing for this Unreleased entry.

Line 32 currently starts with “Support …”; the changelog convention here asks for imperative present tense.

Suggested wording

-- Support passing user-defined Warp functions to `Callable` parameters in `@wp.func` functions
+- Add support for passing user-defined Warp functions to `Callable` parameters in `@wp.func` functions

As per coding guidelines: CHANGELOG.md entries in Unreleased should use imperative present tense (“Add X”).

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@CHANGELOG.md` around lines 32 - 33, Update the Unreleased changelog entry
that currently reads "Support passing user-defined Warp functions to `Callable`
parameters in `@wp.func` functions" to use imperative present-tense phrasing;
change it to something like "Add support for passing user-defined Warp functions
to `Callable` parameters in `@wp.func` functions" so the entry follows the
repository's changelog convention.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@design/callable-func-parameters.md`:
- Around line 41-43: The document's non-goal wording conflicts with the linked
GH-1424 objective about accepting built-in callables; update the text around the
callable targets section to reconcile them by either (A) removing or rephrasing
the line that excludes built-ins and explicitly stating that built-in Warp
functions like wp.sin and wp.add are supported as callable arguments, or (B) if
built-ins truly remain out of scope, change the GH-1424 reference/closure
language to reflect a narrower objective; ensure references to "wp.sin" and
"wp.add" and the GH-1424 issue ID are corrected so the scope described matches
the linked issue.
- Around line 15-18: The snippet imports Callable twice which shadows and
confuses the reader; update the examples so they don't clobber the same name by
either (a) using only one import (prefer collections.abc.Callable for modern
code) and removing the other, (b) aliasing one import (e.g., import Callable as
TypingCallable) to show the difference, or (c) present two separate, clearly
labeled snippets instead of both lines together—adjust the lines containing
"from typing import Callable" and "from collections.abc import Callable"
accordingly.

---

Nitpick comments:
In `@CHANGELOG.md`:
- Around line 32-33: Update the Unreleased changelog entry that currently reads
"Support passing user-defined Warp functions to `Callable` parameters in
`@wp.func` functions" to use imperative present-tense phrasing; change it to
something like "Add support for passing user-defined Warp functions to
`Callable` parameters in `@wp.func` functions" so the entry follows the
repository's changelog convention.

In `@warp/_src/codegen.py`:
- Around line 945-992: The specialization cache key in specialize_callable_func
currently builds specialization_key from (name, callable_arg_values[name]) and
therefore relies on Function object identity for cache hits; add a brief
explanatory comment above the specialization_key creation stating that this is
an intentional design choice (it may produce cache misses if the same logical
function object is recreated, e.g., on module reload), why that is acceptable
for correctness, and note that using callable_func.key/native_func could be an
alternative for a stable key if desired in the future; reference
specialize_callable_func, specialization_key, and func._callable_specializations
in the comment so maintainers can find the code easily.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Enterprise

Run ID: 9019ac3f-bc6a-4ec7-b04f-ab04997190aa

📥 Commits

Reviewing files that changed from the base of the PR and between bd91b99 and 230bf7f.

📒 Files selected for processing (6)

• CHANGELOG.md
• design/callable-func-parameters.md
• warp/_src/codegen.py
• warp/_src/context.py
• warp/_src/types.py
• warp/tests/test_func.py

[coderabbitai]

Actionable comments posted: 3

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

warp/_src/codegen.py (1)

4513-4539: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Track callable-parameter specializations in get_references() too.

This only finds callable targets when the AST argument itself resolves statically. For specialized higher-order functions, forwarded names like inner(f, x) and direct calls like f(x) are backed by adj.callable_arg_values, not by a global/static lookup, so the concrete callable never gets added to functions. That leaves module hashing/dependency invalidation stale when the passed function changes.

Suggested direction

 def get_references(adj) -> tuple[dict[str, Any], dict[Any, Any], dict[warp._src.context.Function, Any]]:
     """Traverses ``adj.tree`` and returns referenced constants, types, and user-defined functions."""

     local_variables = set()  # Track local variables appearing on the LHS so we know when variables are shadowed

     constants: dict[str, Any] = {}
     types: dict[Struct | type, Any] = {}
     functions: dict[warp._src.context.Function, Any] = {}
+    callable_arg_values = getattr(adj, "callable_arg_values", None) or {}

     for node in ast.walk(adj.tree):
         ...
         elif isinstance(node, ast.Call):
             func, _ = adj.resolve_static_expression(node.func, eval_types=False)
+            if func is None and isinstance(node.func, ast.Name):
+                func = callable_arg_values.get(node.func.id)
+
+            if isinstance(func, warp._src.context.Function) and not func.is_builtin():
+                functions[func] = None
+
                 try:
                     bound_args = func.signature.bind(*node.args, **{kw.arg: kw.value for kw in node.keywords})
                 except TypeError:
                     bound_arg_nodes = {}
                 else:
                     ...
                 for arg_name, arg_node in bound_arg_nodes.items():
                     if not warp._src.types.is_callable_annotation(func.input_types.get(arg_name)):
                         continue

                     if isinstance(arg_node, warp._src.context.Function):
                         callable_func = arg_node
+                    elif isinstance(arg_node, ast.Name):
+                        callable_func = callable_arg_values.get(arg_node.id)
                     else:
                         callable_func, _ = adj.resolve_static_expression(arg_node, eval_types=False)

                     if isinstance(callable_func, warp._src.context.Function) and not callable_func.is_builtin():
                         functions[callable_func] = None

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@warp/_src/codegen.py` around lines 4513 - 4539, get_references() currently
only discovers callable arguments when the AST arg resolves statically via
adj.resolve_static_expression, so specialized higher-order calls backed by
adj.callable_arg_values (e.g., when func is forwarded into inner/f and invoked)
are missed; update the Call-handling branch in get_references() (the block that
binds arguments with func.signature.bind, apply_defaults, and iterates
bound_arg_nodes) to also check adj.callable_arg_values for the bound argument
node (and for raw arg names where resolution failed) and add any concrete
warp._src.context.Function entries there into the functions dict (same guard:
isinstance(..., Function) and not is_builtin()). Ensure you reference
adj.callable_arg_values lookup when bound_arg_nodes contains nodes that didn't
resolve statically or that are names/attributes representing forwarded callables
so the concrete callable specializations are included.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@design/callable-func-parameters.md`:
- Around line 43-46: The PR currently conflicts with the design non-goal that
excludes built-in Warp callables (e.g., wp.sin, wp.add) from being accepted as
callable parameters for user-defined `@wp.func` per callable-func-parameters.md;
either update the PR metadata so it does not close GH-1424 (change "Closes
`#1424`" to "Addresses/Related to/Partial implementation of `#1424`" in the PR
description/commit message) or implement full support for built-in callables
(update the `@wp.func` callable-argument handling, specialization hashing, and
dependency tracking code paths that validate callables so built-ins are
recognized and included) and then remove/update the non-goal text and document
the addition.

In `@warp/_src/codegen.py`:
- Around line 972-975: The specialized callable currently clears its value_func
(specialized_func.value_func = None) which breaks downstream use where
add_call() invokes func.value_func(...); restore the original value_func on the
specialized copy instead of nulling it (e.g., assign specialized_func.value_func
= func.value_func or simply remove the line that sets it to None) so the
specialized_func keeps the callable-return resolver used later by add_call().

In `@warp/_src/context.py`:
- Around line 491-506: The overload_annotations map is being populated with the
template annotation (template_type) for callable parameters, causing all
Callable args to collapse to the same specialization; instead, when
warp._src.types.is_callable_annotation(template_type) and the runtime type
(arg_type or default_type) is a Function, store the concrete Function object
into overload_annotations (use arg_type for positional args and default_type for
defaults) so specializations keyed later (lines ~514-516) use the actual
Function instance rather than the generic annotation; update the branches that
currently assign template_type to assign arg_type or default_type accordingly
and keep using get_arg_type/strip_reference and is_callable_annotation checks as
present.

---

Outside diff comments:
In `@warp/_src/codegen.py`:
- Around line 4513-4539: get_references() currently only discovers callable
arguments when the AST arg resolves statically via
adj.resolve_static_expression, so specialized higher-order calls backed by
adj.callable_arg_values (e.g., when func is forwarded into inner/f and invoked)
are missed; update the Call-handling branch in get_references() (the block that
binds arguments with func.signature.bind, apply_defaults, and iterates
bound_arg_nodes) to also check adj.callable_arg_values for the bound argument
node (and for raw arg names where resolution failed) and add any concrete
warp._src.context.Function entries there into the functions dict (same guard:
isinstance(..., Function) and not is_builtin()). Ensure you reference
adj.callable_arg_values lookup when bound_arg_nodes contains nodes that didn't
resolve statically or that are names/attributes representing forwarded callables
so the concrete callable specializations are included.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Enterprise

Run ID: 075e71fa-9c64-4710-9b9b-afd2eee2f986

📥 Commits

Reviewing files that changed from the base of the PR and between af680cc and 99a8d88.

📒 Files selected for processing (7)

• CHANGELOG.md
• design/callable-func-parameters.md
• docs/user_guide/limitations.rst
• warp/_src/codegen.py
• warp/_src/context.py
• warp/_src/types.py
• warp/tests/test_func.py

✅ Files skipped from review due to trivial changes (2)

• docs/user_guide/limitations.rst
• CHANGELOG.md

🚧 Files skipped from review as they are similar to previous changes (2)

• warp/_src/types.py
• warp/tests/test_func.py

[coderabbitai]

♻️ Duplicate comments (1)

warp/_src/context.py (1)

493-506: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Specialize callable overload keys by the concrete Function.

Line 496 and Line 506 still write the annotation back into overload_annotations. Since Lines 516-518 key the instantiated overload from that map, every Callable argument collapses to the same specialization and a later call can reuse the wrong overload/body/hash.

Proposed fix

                 overload_annotations = {}
                 for name, arg_type, template_type in zip(arg_names, arg_types, template_types, strict=False):
                     if warp._src.types.is_callable_annotation(template_type) and isinstance(arg_type, Function):
-                        overload_annotations[name] = template_type
+                        overload_annotations[name] = arg_type
                     else:
                         overload_annotations[name] = arg_type
@@
                         template_type = f.input_types[k]
                         default_type = warp._src.codegen.strip_reference(warp._src.codegen.get_arg_type(d))
                         if warp._src.types.is_callable_annotation(template_type) and isinstance(default_type, Function):
-                            overload_annotations[k] = template_type
+                            overload_annotations[k] = default_type
                         else:
                             overload_annotations[k] = default_type

Also applies to: 516-518

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@warp/_src/context.py` around lines 493 - 506, The overload_annotations map
currently stores the generic callable annotation (template_type) when a
parameter is declared as Callable, causing all callable params to collapse to
the same specialization; instead, when
warp._src.types.is_callable_annotation(template_type) and the runtime value is a
concrete Function, assign the concrete Function type (arg_type for parameters,
default_type for defaults) into overload_annotations[name] so each callable
parameter is specialized by its actual Function; update both the zip loop that
handles arg_types (overload_annotations[name] = arg_type) and the defaults loop
that handles f.defaults (overload_annotations[k] = default_type) and ensure the
later instantiation that keys overloads uses those concrete Function entries.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Duplicate comments:
In `@warp/_src/context.py`:
- Around line 493-506: The overload_annotations map currently stores the generic
callable annotation (template_type) when a parameter is declared as Callable,
causing all callable params to collapse to the same specialization; instead,
when warp._src.types.is_callable_annotation(template_type) and the runtime value
is a concrete Function, assign the concrete Function type (arg_type for
parameters, default_type for defaults) into overload_annotations[name] so each
callable parameter is specialized by its actual Function; update both the zip
loop that handles arg_types (overload_annotations[name] = arg_type) and the
defaults loop that handles f.defaults (overload_annotations[k] = default_type)
and ensure the later instantiation that keys overloads uses those concrete
Function entries.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Enterprise

Run ID: 57311a82-690c-4a18-9a42-229586f4bef2

📥 Commits

Reviewing files that changed from the base of the PR and between 99a8d88 and 5414529.

📒 Files selected for processing (7)

• CHANGELOG.md
• design/callable-func-parameters.md
• docs/user_guide/limitations.rst
• warp/_src/codegen.py
• warp/_src/context.py
• warp/_src/types.py
• warp/tests/test_func.py

✅ Files skipped from review due to trivial changes (3)

• docs/user_guide/limitations.rst
• CHANGELOG.md
• design/callable-func-parameters.md

🚧 Files skipped from review as they are similar to previous changes (3)

• warp/_src/types.py
• warp/_src/codegen.py
• warp/tests/test_func.py

[greptile-apps]

Greptile Summary

This PR adds phase-one support for wp.Function-annotated parameters in user-defined @wp.func functions. Warp functions can now be passed at specialization time, enabling generic helper functions that dispatch to different function targets without runtime indirection.

• Adds a specialize_callable_func path that shallow-copies the callee, attaches concrete targets to the adjoint's callable_arg_values, and mints a unique native name via a SHA-256 suffix so each target combination compiles to its own C++ function.
• Extends get_references() and _find_references() with iter_call_callable_arg_targets so kernel-level scans propagate callable targets into both module hashes and the module dependency graph; built-in targets are hashed by identity while user-defined targets participate in the full dependency chain.
• Covers user-defined targets, simple built-in targets (wp.sin, wp.add, etc.), keyword arguments, defaults, local aliases, nested calls, wp.static(...), wp.grad() specialization, ref-parameter interaction, deterministic-counter compatibility, and cross-module dependency invalidation across 832 lines of focused regression tests.

Confidence Score: 5/5

The change correctly implements a new specialization layer on top of the existing codegen infrastructure; all dependency and hash paths are exercised by the new test suite, and no existing behaviour is modified for non-callable-parameter callers.

Specialization creation, native-name uniqueness, wp.Function arg exclusion from C++ signatures, hash propagation through iter_call_callable_arg_targets, and cross-module dependency invalidation all appear correct. The previously flagged ast.Assign gap is benign: the kernel-level scan covers hash soundness and module-dep edges for every concrete callable target. The comprehensive test suite raises confidence further.

No files require special attention. The three core files (codegen.py, context.py, types.py) each have targeted changes that are well-supported by the regression suite.

Important Files Changed

Filename	Overview
warp/_src/codegen.py	Core of the feature: adds specialize_callable_func, resolve_call refactor, iter_call_callable_arg_targets, and get_references updates. Specialization key is object-identity-safe, native-name uniqueness via SHA-256 prefix is sound, and wp.Function args are correctly skipped in C++ forward/reverse signatures.
warp/_src/context.py	Updates Function.get_overload generic instantiation to preserve wp.Function annotations rather than substituting concrete types, adds hash_builtin_function for identity-based builtin hashing, and wires iter_call_callable_arg_targets into _find_references for correct cross-module dependency edges.
warp/_src/types.py	Adds is_warp_function_annotation (identity check on warp.Function / warp._src.context.Function), is_builtin_callable_annotation (catches Callable-annotated builtin slots), and a dedicated 'c' type-code entry early in get_type_code — correctly replacing the over-broad isinstance(arg_type, Callable) catch-all.
warp/tests/test_func_parameter_targets.py	832-line regression suite covering runtime correctness, hash sensitivity, module dependency invalidation (including cross-module unload cascades), rejection paths, and wp.grad() interaction. All major call patterns exercised.
warp/tests/deterministic/test_deterministic_counter.py	Adds a counter kernel that delegates a bitwise atomic through a wp.Function parameter, verifying that specialization preserves the unconsumed-return semantics required by deterministic lowering.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant K as Kernel (emit_Call)
    participant AC as Adjoint.resolve_call
    participant CAV as get_callable_arg_values
    participant SCF as specialize_callable_func
    participant SF as SpecializedFunc.adj.build
    participant CG as codegen_func

    K->>AC: "func=apply, args=[double_it, 3.0]"
    AC->>AC: resolve_func → exact overload
    AC->>CAV: "bound_args={g: double_it, x: 3.0}"
    CAV-->>AC: "{g: double_it}"
    AC->>SCF: "specialize apply for {g: double_it}"
    SCF-->>AC: specialized clone (unique native_func)
    AC-->>K: (specialized_func, bound_args)
    K->>SF: "build(callable_arg_values={g: double_it})"
    SF->>SF: "symbols[g] = double_it (Function object)"
    SF->>SF: emit g(x) → direct call to double_it
    SF->>CG: codegen_func (skip wp.Function args in C++ sig)
    CG-->>SF: C++ function without g parameter

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant K as Kernel (emit_Call)
    participant AC as Adjoint.resolve_call
    participant CAV as get_callable_arg_values
    participant SCF as specialize_callable_func
    participant SF as SpecializedFunc.adj.build
    participant CG as codegen_func

    K->>AC: "func=apply, args=[double_it, 3.0]"
    AC->>AC: resolve_func → exact overload
    AC->>CAV: "bound_args={g: double_it, x: 3.0}"
    CAV-->>AC: "{g: double_it}"
    AC->>SCF: "specialize apply for {g: double_it}"
    SCF-->>AC: specialized clone (unique native_func)
    AC-->>K: (specialized_func, bound_args)
    K->>SF: "build(callable_arg_values={g: double_it})"
    SF->>SF: "symbols[g] = double_it (Function object)"
    SF->>SF: emit g(x) → direct call to double_it
    SF->>CG: codegen_func (skip wp.Function args in C++ sig)
    CG-->>SF: C++ function without g parameter

Loading

_{Reviews (17): Last reviewed commit: "Support wp.Function parameters (GH-1424)" | Re-trigger Greptile}

[thomasbbrunner]

We'd be super interested in having something like this!

[shi-eric]

We'd be super interested in having something like this!

Hey @thomasbbrunner, do you have any extra requirements for what the feature needs to be able to do?

[thomasbbrunner]

@shi-eric I don't think there's anything that wouldn't be covered by the current examples in the PR's description.

Maybe a small difference of our use-case would be that the functions themselves (e.g., double_it and triple_it in your example) would be located in different modules. Unsure if this matters. To illustrate:

# in ops/double.py

@wp.func
def double_it(x: float):
    return x * 2.0

# in ops/triple.py

@wp.func
def triple_it(x: float):
    return x * 3.0

# in main.py
from ops import double, triple

@wp.func
def apply(g: Callable, x: float):
    return g(x)

@wp.kernel
def k(cond: wp.array[bool], out: wp.array[float]):
    if cond[0]:
        out[0] = apply(double.double_it, 3.0)
    else:
        out[0] = apply(triple.triple_it, 3.0)

On another note, I'm curious about the motivation for using the generic Callable type instead of wp.Function. As in, not all callables could be passed as an argument, only Warp functions, no? To illustrate, this is what I'd have expected:

@wp.func
def apply(g: wp.Function[[float], float], x: float):
    return g(x)