Skip to content

RF packed keep and tighten static buffer sizes - #1867

Draft
polycrit wants to merge 17 commits into
masterfrom
sheremeta/packed-static-bounds
Draft

polycrit wants to merge 17 commits into
masterfrom
sheremeta/packed-static-bounds

Conversation

@polycrit

@polycrit polycrit commented Sep 24, 2026 •

Copy link
Copy Markdown
Contributor

Fixes packed ops whose static, bound sized buffer could end up the wrong size under tracing, either because an op quietly changed the buffer's extent or because the size it proves for a new buffer was looser than it has to be. Covers masked_select, scatter, gather and packed conv.

  • packed conv with stride 1 pads or slices its output buffer back to the input's total extent, so sibling buffers of the same layout stay the same length after the conv
  • masked_select's static buffer bound also takes a declared capacity on the result dim into account, instead of always sizing from the source's full content bound
  • outside static tracing, masked_select's result dim size is copied to the host, matching how other dim sizes live there, so combining it with another dim does not mix devices
  • a new _total_bounds cache remembers the tightest bound a packing over a dim has proven. A static scatter's buffer size now takes the tightest of that bound, the declared capacity and the summed bounds of a dim's parts when it is a sum
  • a re-laid-out gather's static buffer size prefers a proven bound on the result dim over the capacity ratio between the two dims, since that ratio only holds when the result length tracks the source length, as with a stride or a window

Tests

  • test_masked_select_static_buffer_follows_a_declared_capacity
  • test_scatter_into_a_sum_dim_bounds_both_parts_alike
  • test_masked_select_eager_result_dim_adds_to_a_host_dim
  • test_gather_relayout_static_buffer_takes_the_proven_bound_of_the_result_dim

@polycrit
polycrit force-pushed the sheremeta/packed-static-bounds branch from 1404251 to 7308ae7 Compare October 10, 2026 19:33
Base automatically changed from sheremeta/packed-scatter-gather to master October 11, 2026 22:30
robin-p-schmitt added a commit that referenced this pull request Oct 14, 2026
# Conflicts:
#	returnn/frontend/_packed_backend.py
#	returnn/frontend/array_.py
#	tests/test_rf_packed.py
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.

2 participants