Fix wp.static() capturing global variables instead of loop variables#1167
Fix wp.static() capturing global variables instead of loop variables#1167shi-eric merged 1 commit intoNVIDIA:mainfrom
Conversation
📝 WalkthroughWalkthroughTracks loop-local variable names during codegen so wp.static() inside for-loops no longer captures shadowed globals; adds tests exercising simple, expression, nested, and shadowed-variable loop scenarios and a changelog entry. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
📜 Recent review detailsConfiguration used: Path: .coderabbit.yml Review profile: CHILL Plan: Pro 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (1)
🧰 Additional context used🧠 Learnings (1)📚 Learning: 2026-01-13T06:46:05.869ZApplied to files:
🪛 Ruff (0.14.11)warp/tests/test_static.py710-710: Unused function argument: (ARG001) 776-776: (F405) ⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)
🔇 Additional comments (6)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
warp/_src/codegen.py (1)
3744-3753: Loop variable tracking logic is correct.The implementation properly adds and removes loop variables when entering/exiting
Fornode scope. The use ofdiscard()instead ofremove()is a good defensive choice.Optional: Consider handling tuple unpacking in loop targets.
While the comment notes that tuple unpacking is rare, the current implementation silently ignores patterns like
for x, y in .... Consider either:
- Extracting all names from tuple targets and tracking them
- Raising a warning if tuple unpacking is used with
wp.static()This could prevent subtle bugs if users do use tuple unpacking.
warp/tests/test_static.py (1)
691-716: Test logic is correct, but parameter may be unused.The test function correctly validates that
wp.static()captures loop variables instead of shadowed globals. The expected outputs are accurate:
- Simple loop:
[0, 1, 2]from iteration values- Expression:
[1, 3, 5]fromi*2+1- Nested:
[0, 1, 10, 11]fromi*10+jConsider addressing the unused
testparameter.Static analysis indicates the
testparameter is unused. If it's required by the test framework interface, consider prefixing it with an underscore (_test) to indicate it's intentionally unused. Otherwise, you may be able to use it for assertions instead of relying on NumPy'sassert_array_equal.♻️ Optional: Use test parameter for assertions
If the test framework provides assertion methods, you could use them instead:
- np.testing.assert_array_equal(results.numpy(), np.array([0, 1, 2]), err_msg="Simple loop variable test failed") + test.assertTrue(np.array_equal(results.numpy(), np.array([0, 1, 2])), "Simple loop variable test failed")Or if not needed, rename to
_testto indicate it's intentionally unused.
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
CHANGELOG.mdwarp/_src/codegen.pywarp/tests/test_static.py
🧰 Additional context used
🪛 Ruff (0.14.10)
warp/tests/test_static.py
691-691: Unused function argument: test
(ARG001)
748-748: add_function_test may be undefined, or defined from star imports
(F405)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)
- GitHub Check: ci / test-warp (ubuntu-24.04, ubuntu-x86_64)
- GitHub Check: ci / test-warp (macos-latest, macos)
- GitHub Check: ci / test-warp (windows-2025, windows)
- GitHub Check: ci / test-warp (ubuntu-24.04-arm, ubuntu-aarch64)
- GitHub Check: ci / build-docs
- GitHub Check: cppcheck
🔇 Additional comments (6)
CHANGELOG.md (1)
18-19: Changelog entry looks good.The fix is documented clearly and accurately. The entry follows the established format, includes the GitHub issue reference, and provides sufficient context for users encountering this issue.
warp/_src/codegen.py (2)
3739-3742: LGTM!The initialization of
loop_varsas a set is appropriate for tracking unique loop variable names. The comment clearly explains the purpose of preventingwp.static()from capturing shadowed global variables.
3758-3765: Correct implementation of loop variable detection.The code properly:
- Walks the AST of the static expression to find all referenced names
- Checks intersection with tracked loop variables
- Defers evaluation by setting the
has_unresolved_static_expressionsflag- Continues traversal to allow later codegen-time evaluation
This ensures that loop constants are available when the static expression is eventually evaluated.
warp/tests/test_static.py (3)
662-665: LGTM!The global test variables are well-named and their purpose is clearly documented. The values (999, 888) are deliberately different from loop iteration values to make it obvious if
wp.static()incorrectly captures the globals instead of loop variables.
667-689: Comprehensive test coverage for loop variable handling.The three test kernels effectively cover:
- Simple loop variable capture
- Loop variables in complex expressions
- Nested loops with multiple loop variables
Each kernel has a clear docstring explaining its purpose, and the test cases progressively increase in complexity. This provides good validation that
wp.static()correctly handles loop variables in various scenarios.
748-753: LGTM!The test registration follows the same pattern as other tests in the file and correctly:
- Associates the test with the TestStatic class
- Uses a descriptive test name
- Passes the test function and devices parameter
The static analysis warning about
add_function_testbeing undefined is a false positive—it's imported via the star import fromwarp.tests.unittest_utilsat line 24.
Greptile SummaryThis PR fixes a bug where Key Changes:
Technical Correctness: Confidence Score: 5/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant User as User Code
participant AST as AST Transformer
participant Adj as Adjoint
participant Eval as Static Evaluator
participant Codegen as Code Generator
User->>AST: Define kernel with wp.static(loop_var)
AST->>AST: visit_For(): Track loop variable in loop_vars
AST->>AST: visit_Call(): Detect wp.static() referencing loop_var
AST->>Adj: Set has_unresolved_static_expressions = True
AST->>AST: Defer evaluation (don't replace with constant)
Note over AST,Adj: wp.static() left as call node in AST
Adj->>Codegen: emit_For(): Begin loop unrolling
loop For each iteration i
Codegen->>Adj: Add loop constant to symbols[loop_var]
Codegen->>Eval: Evaluate wp.static(loop_var) with loop constant
Eval->>Eval: Build evaluation context from adj.symbols
Eval->>Eval: Find loop_var constant in symbols
Eval-->>Codegen: Return correct loop iteration value
Codegen->>Codegen: Generate code for loop body with constant
end
Note over Codegen: Loop unrolled with correct constants
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
shi-eric
force-pushed
the
ershi/fix-static-loop-var
branch
from
b7e5a81 to
3db7004
Compare
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Fix all issues with AI agents
In @warp/tests/test_static.py:
- Around line 776-781: Ruff flags F405 for the bare usage of add_function_test
due to star-imports; fix by adding a local noqa to that call site — append a
trailing comment "# noqa: F405" to the add_function_test invocation for
test_static_loop_variable_not_shadowed_by_global so the single-line suppression
silences the warning without changing imports.
- Around line 674-680: The triple-quoted docstring inside the @wp.kernel
function static_loop_var_in_expr_kernel can end up as a runtime string literal
inside the kernel body and interfere with static-analysis; remove or move that
docstring out of the kernel (e.g., convert it to a regular comment above the
function or place the descriptive string at module level) so the kernel body
contains no top-level string literal, keeping the loop and the use of
wp.static(_global_test_idx * 2 + 1) unchanged.
- Around line 682-689: The kernel static_nested_loop_kernel contains a Python
docstring inside the kernel body which can be misinterpreted by the compiler;
remove the triple-quoted docstring from inside the kernel (or convert it to a
regular comment) and place any description outside the kernel function so the
loop variables (_global_test_idx, _global_test_j), idx calculation, and
wp.static usage remain unchanged and the test still expects [0, 1, 10, 11].
- Around line 710-743: The test function
test_static_loop_variable_not_shadowed_by_global has an unused parameter `test`
causing Ruff ARG001—rename it to `_test` or `_` (or add `# noqa: ARG001`) to
silence the linter; additionally, replace uses of
`np.testing.assert_array_equal(...)` with the file's existing
`assert_np_equal(...)` helper for consistency (update the four assertions: the
ones checking results, results2, results3, results4) while preserving the same
expected arrays and messages.
- Around line 667-672: The kernel function static_loop_var_kernel currently
contains a leading docstring inside the @wp.kernel function which can become an
ast.Expr and break Warp codegen; remove the triple-quoted string literal from
inside the function and instead place the explanatory text as a regular comment
immediately above the def line (or convert it to inline comments within the
body), ensuring the function body only contains executable statements that use
results and wp.static(_global_test_idx) with the loop variable _global_test_idx.
🧹 Nitpick comments (1)
warp/tests/test_static.py (1)
662-665: Make the “shadowed global” intent clearer and safer from collisions.These globals are fine for reproducing the bug, but consider making them even more unique (e.g., include
static_loop_shadow_) to reduce any chance of name collision with other tests/modules.
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
CHANGELOG.mdwarp/_src/codegen.pywarp/tests/test_static.py
🚧 Files skipped from review as they are similar to previous changes (2)
- CHANGELOG.md
- warp/_src/codegen.py
🧰 Additional context used
🪛 Ruff (0.14.10)
warp/tests/test_static.py
710-710: Unused function argument: test
(ARG001)
776-776: add_function_test may be undefined, or defined from star imports
(F405)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: ci / build-warp (windows-2025, windows, ["cudart", "nvcc", "nvrtc_dev", "nvjitlink"])
🔇 Additional comments (1)
warp/tests/test_static.py (1)
691-708: Test correctly asserts the intended behavior; docstring usage is valid.The test properly validates that
wp.static()captures the loop variable's final value (2) after the inner loop completes, not the global value (999). The assertionresults == [2, 2]is correct: the outer loop executes twice, and each iterationwp.static(_global_test_idx)captures 2, matching Python's variable shadowing semantics.Regarding docstrings in kernels: docstrings are explicitly allowed in Warp kernels per the framework documentation ("Only pass, ellipsis (...), comments, or docstrings are allowed"). This kernel follows the same pattern as other kernels in the same test file (
static_loop_var_kernel,static_nested_loop_kernel, etc.), all with docstrings, making the approach consistent and correct.Likely an incorrect or invalid review comment.
| @wp.kernel | ||
| def static_loop_var_kernel(results: wp.array(dtype=int)): | ||
| """Kernel where wp.static() should capture the loop variable, not the global.""" | ||
| for _global_test_idx in range(3): | ||
| results[_global_test_idx] = wp.static(_global_test_idx) | ||
|
|
There was a problem hiding this comment.
Potential Warp compile hazard: docstring inside @wp.kernel.
If Warp’s AST/codegen doesn’t explicitly ignore function docstrings, the leading string literal becomes an ast.Expr statement and can break compilation. Safer to replace kernel docstrings with comments above the function.
Proposed fix
@wp.kernel
def static_loop_var_kernel(results: wp.array(dtype=int)):
- """Kernel where wp.static() should capture the loop variable, not the global."""
+ # wp.static() should capture the loop variable, not the global.
for _global_test_idx in range(3):
results[_global_test_idx] = wp.static(_global_test_idx)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| @wp.kernel | |
| def static_loop_var_kernel(results: wp.array(dtype=int)): | |
| """Kernel where wp.static() should capture the loop variable, not the global.""" | |
| for _global_test_idx in range(3): | |
| results[_global_test_idx] = wp.static(_global_test_idx) | |
| @wp.kernel | |
| def static_loop_var_kernel(results: wp.array(dtype=int)): | |
| # wp.static() should capture the loop variable, not the global. | |
| for _global_test_idx in range(3): | |
| results[_global_test_idx] = wp.static(_global_test_idx) |
🤖 Prompt for AI Agents
In @warp/tests/test_static.py around lines 667 - 672, The kernel function
static_loop_var_kernel currently contains a leading docstring inside the
@wp.kernel function which can become an ast.Expr and break Warp codegen; remove
the triple-quoted string literal from inside the function and instead place the
explanatory text as a regular comment immediately above the def line (or convert
it to inline comments within the body), ensuring the function body only contains
executable statements that use results and wp.static(_global_test_idx) with the
loop variable _global_test_idx.
There was a problem hiding this comment.
This is explicitly handled in codegen.py:2019:
# Skip variable creation for standalone constants, including docstrings
if isinstance(f, ast.Expr) and isinstance(f.value, ast.Constant):
continueWe skip the docstring, I've also confirmed this in practice.
There was a problem hiding this comment.
@eric-heiden, thank you for the clarification! You're absolutely right—I apologize for the incorrect flag. Since Warp's codegen explicitly skips standalone ast.Constant expressions (including docstrings) at that line in codegen.py, there's no compilation hazard here. I appreciate you confirming this works in practice as well.
/)_/)
( •.•)
/>📚 Noted!
✏️ Learnings added
Learnt from: eric-heiden
Repo: NVIDIA/warp PR: 1167
File: warp/tests/test_static.py:667-672
Timestamp: 2026-01-22T19:39:17.536Z
Learning: In the NVIDIA/warp repository, docstrings inside `wp.kernel` and `wp.func` decorated functions are safe and explicitly handled by the codegen. The code at `codegen.py` skips standalone `ast.Constant` expressions (including docstrings) with the check `if isinstance(f, ast.Expr) and isinstance(f.value, ast.Constant): continue`, so they do not cause compilation issues.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
🧠 Learnings used
Learnt from: shi-eric
Repo: NVIDIA/warp PR: 1167
File: warp/tests/test_static.py:776-781
Timestamp: 2026-01-13T06:46:08.692Z
Learning: In the NVIDIA/warp repository, star imports like `from warp.tests.unittest_utils import *` in test files are acceptable per the project's Ruff configuration, and F405 warnings should not be flagged if pre-commit checks pass.
When a global Python variable had the same name as a kernel for-loop variable, wp.static() incorrectly used the global value. The fix tracks loop variables during AST traversal and defers wp.static() evaluation when the expression references a loop variable. Fixes NVIDIAGH-1139 Signed-off-by: Eric Shi <ershi@nvidia.com>
shi-eric
force-pushed
the
ershi/fix-static-loop-var
branch
from
3db7004 to
11764fd
Compare
eric-heiden
left a comment
eric-heiden
left a comment
There was a problem hiding this comment.
LGTM, it fixes the issue.
2 participants
When a global Python variable had the same name as a kernel for-loop variable, wp.static() incorrectly used the global value. The fix tracks loop variables during AST traversal and defers wp.static() evaluation when the expression references a loop variable.
Fixes GH-1139
Description
Before your PR is "Ready for review"
__init__.pyi,docs/api_reference/,docs/language_reference/)pre-commit run -aSummary by CodeRabbit
Bug Fixes
Tests
✏️ Tip: You can customize this high-level summary in your review settings.