Reject non-None kernel return annotations

[snico432]

Description

Raises an error when a kernel has a non-None return annotation. Closes #1471.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

Test plan

Tested on CPU with:

uv run warp/tests/test_codegen.py TestCodeGen.test_error_kernel_return_value_cpu

Bug fix

Previously, the float return annotation was silently ignored by Warp:

import warp as wp

@wp.kernel
def annotated_kernel(x: float) -> float:
    return

wp.launch(annotated_kernel, dim=1, inputs=[1.0], device="cpu")

This PR makes that code raise:

WarpCodegenTypeError: 'annotated_kernel': Error, kernels can't have return values; write results to output arguments and either omit the return annotation or use '-> None'

Summary by CodeRabbit

• Bug Fixes

  • Kernels with non-None return type annotations now raise a clear error, with guidance to write results to output arguments or use a None return.

• Documentation

  • Changelog updated with an Unreleased entry describing the behavior change.

• Tests

  • Tests updated to cover and assert the new error behavior.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR enforces that Warp kernels cannot declare non-None return annotations. A validation check in ModuleBuilder.build_kernel() detects return type annotations and raises an error with guidance to use output arguments instead. The change includes a test case and changelog documentation.

Changes

Kernel return annotation validation

Layer / File(s)	Summary
Return annotation validation and error handling warp/_src/context.py, warp/tests/test_codegen.py, CHANGELOG.md	ModuleBuilder.build_kernel() detects return type annotations via kernel.adj.arg_types.get("return") and raises WarpCodegenTypeError instructing users to write results to output arguments and either omit the return annotation or use -> None. New test kernel f4 verifies the error is raised for a -> float annotation with no explicit return. Changelog documents the new constraint.

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% 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
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Linked Issues check	✅ Passed	All changes implement the core requirements from #1471: reject non-None return annotations, continue rejecting value-returning kernels, and provide clear diagnostics guiding users to omit return annotations or use -> None.
Out of Scope Changes check	✅ Passed	All changes are scoped to implementing the kernel return annotation validation. CHANGELOG.md documents the change, context.py adds the validation logic, and test_codegen.py adds the corresponding test case.
Title check	✅ Passed	The title 'Reject non-None kernel return annotations' directly and clearly summarizes the main change: preventing kernels from having non-None return type annotations.

_{✏️ 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 and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 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 `@CHANGELOG.md`:
- Around line 5-6: The CHANGELOG entry "Raise an error when a kernel is defined
with a non-`None` return annotation." is missing the GitHub issue reference;
update that line to include " (fixes `#1471`)" or similar so the entry closes the
referenced issue (PR closes `#1471`) and follows the repository's changelog
guideline.

🪄 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: bf3be72d-c55d-406c-b060-71c6eeae7902

📥 Commits

Reviewing files that changed from the base of the PR and between 0187490 and b93d722.

📒 Files selected for processing (3)

• CHANGELOG.md
• warp/_src/context.py
• warp/tests/test_codegen.py

[greptile-apps]

Greptile Summary

This PR fixes a silent bug where a kernel annotated with a non-None return type (e.g., -> float) but with a bare return statement was accepted without error. A new guard in build_kernel now raises WarpCodegenTypeError with a helpful message directing users to output arguments.

• warp/_src/context.py: Adds a check on kernel.adj.arg_types.get(\"return\") after the existing return_var check; the fix is safe because arg_types is built during annotation parsing and already strips -> None entries, so -> None and unannotated kernels are unaffected.
• warp/tests/test_codegen.py: Activates the previously commented-out f4 test case with module=\"unique\" isolation and asserts the new, more descriptive error message.

Confidence Score: 5/5

Safe to merge — the change adds a new error guard that is narrowly scoped to the kernel build path and does not touch any existing code paths.

The fix is small and surgical: one new conditional in build_kernel that fires only when a kernel carries a non-None return annotation but no actual return value. The underlying arg_types dict already strips -> None at parse time, so the guard correctly excludes every valid kernel form. The previously commented-out test case is now activated with proper module isolation, giving direct coverage of the new branch.

No files require special attention.

Important Files Changed

Filename	Overview
warp/_src/context.py	Adds a second guard in build_kernel that catches non-None return annotations on kernels that don't actually return a value; the fix is correct and relies on the existing arg_types dict that already strips -> None entries at parse time.
warp/tests/test_codegen.py	Activates the previously commented-out f4 test case (annotation with no return value) and adds the expected error message assertion; module=unique is correctly used for isolation.
CHANGELOG.md	Adds a changelog entry for the new error; placement within the Fixed section looks correct.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[build_kernel called] --> B[kernel.adj.build]
    B --> C{return_var is not None?}
    C -- Yes --> D[raise WarpCodegenTypeError\n'kernels can't have return values']
    C -- No --> E{arg_types.get 'return' is not None?}
    E -- Yes\n e.g. '-> float' annotation --> F[raise WarpCodegenTypeError\n'kernels can't have return values;\nwrite results to output arguments...']
    E -- No\n '-> None' or no annotation --> G[Build proceeds normally]

Loading

_{Reviews (2): Last reviewed commit: "Reject non-None kernel return annotation..." | Re-trigger Greptile}