Fix block-sparse sub-block assignment when RHS threshold is looser than dest#572
Open
evaleev wants to merge 1 commit into
Open
Fix block-sparse sub-block assignment when RHS threshold is looser than dest#572evaleev wants to merge 1 commit into
evaleev wants to merge 1 commit into
Conversation
…an dest Expr::eval_to(BlkTsrExpr) forms the assigned sub-block's shape via SparseShape::update_block, which screens the updated region with the DESTINATION array's threshold, but the tile-write loop wrote every tile the RHS (dist_eval) shape kept. When the RHS shape carries a looser screening threshold than the destination, a tile the RHS keeps yet the merged result screens to zero was written into a shape-zero slot, tripping ArrayImpl::set's !is_zero assertion. Make the destination shape authoritative: still consume (get) every RHS non-zero tile -- skipping the get strands the producer and deadlocks dist_eval.wait() -- but only set tiles the merged result shape keeps. Add a regression test in expressions_sparse using a real DistArray whose sub-block source keeps a small-norm tile the stricter destination threshold screens out.
There was a problem hiding this comment.
Pull request overview
This PR fixes a correctness issue in sparse sub-block assignment (dest.block(...) = rhs) where the block-write loop could attempt to write tiles that the destination’s (stricter) sparse shape screens to zero, triggering ArrayImpl::set assertions (or hanging under throwing assert policy in tasks). The fix makes the destination block shape authoritative when deciding which RHS tiles to actually set, while still unconditionally consuming RHS tiles from dist_eval to avoid deadlocking dist_eval.wait().
Changes:
- Guard sub-block tile writes so only tiles kept by the merged destination
resultshape are written. - Always
dist_eval.get()consumed tiles even when they’ll be dropped, preventing evaluator deadlock. - Add a regression test covering threshold-mismatch behavior for sparse block assignment.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
src/TiledArray/expressions/expr.h |
Drops RHS tiles that the destination-updated block shape screens to zero, preventing set into shape-zero slots while still consuming RHS futures. |
tests/expressions_sparse.cpp |
Adds a regression test verifying block assignment behavior when RHS and destination screening thresholds differ. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+48
to
+49
| const float saved_threshold = Shape::threshold(); | ||
|
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
a("i,j").block(lo,hi) = rhs("i,j")aborts withTA_ASSERT(!is_zero(i))inArrayImpl<..., SparsePolicy>::set(or hangs, under the throwing assert policy inside a task) when the RHS expression's shape carries a looser screening threshold than the destination array.Expr::eval_to(BlkTsrExpr)builds the assigned sub-block's shape viaSparseShape::update_block, which screens the updated region with the destination array'smy_threshold_. The tile-write loop, however, wrote every tile the RHS (dist_eval) shape kept. A tile the RHS keeps but the mergedresultshape screens to zero was then written into a shape-zero slot, tripping the!is_zeroassertion.This surfaces downstream (e.g. a sub-block scatter that reconstructs a sparse result from independently-screened partials, where the partials were screened at a lower threshold than the pre-sized destination).
Fix
Make the destination shape authoritative in the block-write loop: only
settiles the mergedresultshape keeps. The RHS tile is stillget()-consumed unconditionally — skipping theget()would strand the scheduled producer and deadlockdist_eval.wait().DensePolicyis unaffected:DenseShape::is_zerois always false, so the added guard is inert and behavior is byte-identical.Test
Adds
expressions_sparse_block_assign_suite/block_assign_threshold_mismatch: a realTSpArrayDsub-block assignment whose source keeps a small-norm tile that the stricter destination threshold screens out. It asserts the assignment does not throw and that the destination threshold decides which tiles survive. The case fails (assertion) without the fix.Verified locally: new case +
expressions_sparse(52), denseexpressions(52), and ToT expression suites all pass.