Fix bsr_get_diag zero fill for missing blocks

[Cucchi01]

Description

To my understanding, bsr_get_diag should produce the same result whether or not an out array is provided. Previously, when out was passed, it was not zero-initialized, causing missing diagonal blocks to retain whatever values were already in out instead of being set to 0.0. This PR fixes that and adds a regression test.

Before your PR is "Ready for review"

• All commits are signed-off to indicate that your contribution adheres to the Developer Certificate of Origin requirements
• Necessary tests have been added
• Documentation is up-to-date
• Auto-generated files modified by compiling Warp and building the documentation have been updated (e.g. __init__.pyi, docs/api_reference/, docs/language_reference/)
• Code passes formatting and linting checks with pre-commit run -a

Summary by CodeRabbit

• Bug Fixes

  • Fixed bsr_get_diag() to properly initialize output buffers when provided by users, ensuring correct results when writing diagonal blocks to an external buffer.

• Tests

  • Added test coverage validating output buffer behavior for the diagonal extraction function.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[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

The bsr_get_diag function in the sparse module now clears user-provided output arrays by calling out.zero_() after device and shape validation, before launching the diagonal extraction kernel. A corresponding test validates that results written to an external buffer match internal allocation behavior.

Changes

Cohort / File(s)	Summary
Core sparse function update warp/_src/sparse.py	Added out.zero_() call in bsr_get_diag to reset the output array after validation and before kernel execution
Test coverage warp/tests/test_sparse.py	Added test case exercising the out parameter: allocates sentinel buffer, invokes bsr_get_diag with external output, and validates parity with internal allocation

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

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

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the main change: fixing bsr_get_diag to zero the output buffer and adding a corresponding test.

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

✨ Finishing touches

• 📝 Generate docstrings

---

📜 Recent review details

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 7961180 and 3d8dc3b.

📒 Files selected for processing (2)

• warp/_src/sparse.py
• warp/tests/test_sparse.py

🧰 Additional context used

🧬 Code graph analysis (1)

warp/tests/test_sparse.py (2)

warp/_src/sparse.py (8)

• shape (126-129)
• shape (705-706)
• nrow (673-674)
• dtype (132-134)
• dtype (709-710)
• device (137-139)
• device (717-718)
• bsr_get_diag (1243-1273)

warp/tests/unittest_utils.py (1)

• assert_np_equal (241-247)

🪛 Ruff (0.14.10)

warp/tests/test_sparse.py

247-247: assert_np_equal may be undefined, or defined from star imports

(F405)

🔇 Additional comments (2)

warp/_src/sparse.py (1)

1264-1264: LGTM! Correct fix for ensuring consistent behavior.

Zeroing the user-provided out buffer ensures positions without diagonal blocks are set to zero, matching the behavior when out is allocated internally via wp.zeros().

warp/tests/test_sparse.py (1)

244-248: Good regression test for the out parameter fix.

The test correctly validates the fix by:

1. Using a non-zero sentinel value (7.0) to detect whether positions without diagonal blocks get properly zeroed
2. Asserting parity between the externally provided buffer and internal allocation results

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[greptile-apps]

Greptile Overview

Greptile Summary

This PR fixes a bug in bsr_get_diag where providing an out parameter would not zero-initialize it, causing missing diagonal blocks to retain pre-existing values instead of being set to 0.0.

The Fix:

• Adds out.zero_() call in bsr_get_diag when an out array is provided
• This ensures consistent behavior whether or not out is passed

The Test:

• Creates a sentinel array filled with 7.0 values
• Calls bsr_get_diag with this array as output
• Verifies the result matches the internally-allocated case
• This regression test ensures missing diagonal blocks are properly zeroed

The kernel _bsr_get_diag_kernel only writes to out[row] when bsr_block_index(row, row, ...) returns a valid block (not -1). Without zeroing, missing diagonal blocks would retain garbage values.

Confidence Score: 5/5

• Safe to merge - simple, well-tested bug fix with no side effects
• The fix is minimal (adds one line), addresses a clear bug, includes proper regression testing, and follows the pattern used elsewhere in the codebase. The kernel only writes to positions with diagonal blocks, so zeroing is necessary for correctness.
• No files require special attention

Important Files Changed

File Analysis

Filename	Score	Overview
warp/_src/sparse.py	5/5	Adds out.zero_() to properly initialize the output buffer before extracting diagonal blocks
warp/tests/test_sparse.py	5/5	Adds regression test verifying out buffer produces same result as internal allocation

Sequence Diagram

sequenceDiagram
    participant User
    participant bsr_get_diag
    participant out array
    participant _bsr_get_diag_kernel
    
    User->>bsr_get_diag: Call with BSR matrix A and out buffer
    
    alt out is None
        bsr_get_diag->>out array: Create zeros(dim)
    else out is provided
        bsr_get_diag->>bsr_get_diag: Validate dtype, device, shape
        bsr_get_diag->>out array: out.zero_() [NEW FIX]
    end
    
    bsr_get_diag->>_bsr_get_diag_kernel: Launch kernel(scale, offsets, columns, values, out)
    
    loop For each row
        _bsr_get_diag_kernel->>_bsr_get_diag_kernel: diag = bsr_block_index(row, row, ...)
        
        alt diag != -1 (diagonal block exists)
            _bsr_get_diag_kernel->>out array: out[row] = scale * A_values[diag]
        else diag == -1 (no diagonal block)
            Note over _bsr_get_diag_kernel,out array: out[row] remains zero (from zero_() call)
        end
    end
    
    bsr_get_diag-->>User: Return out array

Loading

[greptile-apps]

_{No files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[gdaviet]

Good catch, thanks for the PR!

[shi-eric]

Thanks @Cucchi01, this is now merged and will be a part of the next Warp release.