Skip to content

Add triangular solve function for sparse CSR tensors in XPU#3261

Open
tszulist-hbn wants to merge 1 commit intointel:mainfrom
tszulist-hbn:dev/tszulist/issue3167
Open

Add triangular solve function for sparse CSR tensors in XPU#3261
tszulist-hbn wants to merge 1 commit intointel:mainfrom
tszulist-hbn:dev/tszulist/issue3167

Conversation

@tszulist-hbn
Copy link
Copy Markdown
Contributor

Resolves: #3167

This PR adds the SparseCsrXPU dispatch for triangular_solve by converting the sparse matrix to dense and delegating to the existing dense XPU triangular_solve implementation. This follows the same pattern used by other XPU sparse ops (addmm, baddbmm, etc.).

Changes:

yaml/native/native_functions.yaml — Register SparseCsrXPU: triangular_solve_out_sparse_csr_xpu dispatch for both triangular_solve.X (structured out variant) and triangular_solve (structured delegate).
src/ATen/native/sparse/xpu/SparseCsrTensorMath.cpp — Implement triangular_solve_out_sparse_csr_xpu() which handles the zero-nnz edge case (fills X with NaN, matching CUDA behavior) and otherwise converts sparse A to dense before calling at::triangular_solve.

Tests verified:

test_block_triangular_solve — 64/64 variants pass (block_size 2/3, int32/int64, contiguous/noncontiguous, float32/float64/complex64/complex128)
test_sparse_triangular_solve_xpu — 4/4 variants pass (float32/float64/complex64/complex128)

Copilot AI review requested due to automatic review settings April 2, 2026 15:59
@tszulist-hbn tszulist-hbn force-pushed the dev/tszulist/issue3167 branch from a47240d to fb6f866 Compare April 2, 2026 16:00
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds SparseCsrXPU support for triangular_solve by wiring a SparseCsrXPU dispatch and implementing an XPU sparse CSR fallback that densifies A and calls the existing dense XPU solve.

Changes:

  • Register SparseCsrXPU dispatch for triangular_solve.X (and delegate triangular_solve to it).
  • Implement triangular_solve_out_sparse_csr_xpu in the XPU sparse CSR math file using A.to_dense() + at::triangular_solve.
  • Add an _nnz()==0 fast-path intended to match CUDA behavior (fills X with NaNs).

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File Description
yaml/native/native_functions.yaml Adds SparseCsrXPU dispatch registration for triangular_solve.X and structured delegate entry for triangular_solve.
src/ATen/native/sparse/xpu/SparseCsrTensorMath.cpp Implements SparseCsrXPU out-kernel that densifies A and delegates to dense at::triangular_solve, with an _nnz()==0 special-case.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

@tszulist-hbn tszulist-hbn marked this pull request as draft April 2, 2026 16:14
@tszulist-hbn tszulist-hbn marked this pull request as ready for review April 2, 2026 16:16
@tszulist-hbn tszulist-hbn force-pushed the dev/tszulist/issue3167 branch from fb6f866 to 2a8ed8f Compare April 3, 2026 08:12
Copilot AI review requested due to automatic review settings April 3, 2026 15:37
@tszulist-hbn tszulist-hbn force-pushed the dev/tszulist/issue3167 branch from 2a8ed8f to 1abc210 Compare April 3, 2026 15:37
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.


💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

@tszulist-hbn tszulist-hbn force-pushed the dev/tszulist/issue3167 branch 2 times, most recently from 6b2c07b to f242fda Compare April 7, 2026 13:40
Copilot AI review requested due to automatic review settings April 7, 2026 13:40
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.


💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

@tszulist-hbn tszulist-hbn force-pushed the dev/tszulist/issue3167 branch from f242fda to 2ab054a Compare April 9, 2026 08:44
Copy link
Copy Markdown
Contributor

@guangyey guangyey left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall LGTM. Could you please add a test case? Maybe test/regressions is a good place.
Defer to @CuiYifeng

@guangyey guangyey requested a review from CuiYifeng April 9, 2026 16:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

NotImplementedError: Could not run 'aten::triangular_solve.X' with arguments from the 'SparseCsrXPU' backend

3 participants