WIP: Avoid dense n x n allocations in upper_gather() - #1635
Open
kei51e wants to merge 4 commits into
Open
Conversation
upper_gather() dominated the cost of every pairwise function that calls it (distance, cosine similarity, KL distance, pair counting). Measured on a 4000 x 5 input, stats::dist() itself took 0.071s while the upper_gather() step took 8.3s and peaked at 1.1GB. The dist vector path built a dense n x n triangular matrix, the row() and col() index matrices for it, a logical mask, a transpose, and one more n x n copy inside mat_to_df(): six n x n allocations, about 60 * n^2 bytes, which is 6GB at n = 10000. It now builds the long form directly from the dist vector. stats::dist stores the lower half in column major order, which is the same sequence as the upper half in row major order, so the values keep their original order and no n x n object is allocated at all. The matrix path had a second copy of the same problem. With zero.rm = FALSE it evaluated is.na(tmat) | !is.na(tmat), an all TRUE matrix of the same shape, purely to hand every cell to which(). For a sparse input that forces the sparse class to store all n^2 cells, which is heavier than a plain dense logical. The selected indices are known in closed form there, so only the surviving half is generated. The output contract is unchanged: same row order, row count, column names, column types, and NA positions. bench/micro/upper_gather_equivalence.R compares this implementation against the previous one with identical() over the six argument combinations the callers use plus degenerate, non square, multibyte and duplicated name cases: 47 of 49 identical. The two remaining cases are duplicated dimension names on the vector path, where the previous implementation raised an error from tibble inside mat_to_df() and this one returns the natural result. mat_to_df() itself is deliberately left untouched; the vector path simply no longer calls it.
The five callers of upper_gather() depend on details that were not covered by any test: the full n x n grid returned when na.rm and zero.rm are FALSE, the row major order, which half of the grid carries values, the column types, and the fact that dimension names are copied into the output verbatim. Adds two test_that blocks to test_util.R covering the six argument combinations the callers use, plus the degenerate cases (n = 1, n = 2, all NA) and dimension names that are multibyte, punctuation heavy, near colliding or duplicated. Adds bench/micro/upper_gather_bench.R next to the equivalence harness. It compares both paths against the previous implementation over a range of n and reports elapsed time, gc peak and row count.
gc() inserts a "limit (Mb)" column when a memory limit is set, which shifts column 6 from "max used (Mb)" to "max used" in cells. The benchmark indexed column 6 directly, so on a machine with a limit set it reported cell counts instead of megabytes (about 9.5e6 where the value is 178 Mb). "max used (Mb)" is the last column in both shapes, so index from the end.
…void-dense-matrix
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
upper_gather()is the shared long-form builder behind every pairwise function(
do_dist.kv/do_dist.cols/do_kl_dist.kv/do_cosine_sim.kv/pair_count).It was the dominant cost of all of them, by a wide margin.
Measured on a 4000 x 5 input:
stats::dist()itself takes 0.07s, while theupper_gather()step took 8.3s and peaked at 1.1GB. That is a 117x ratio, so thedistance computation was never the bottleneck -- the reshape after it was.
The dist vector path allocated six dense n x n objects: the triangular matrix, its
row()andcol()index matrices, a logical mask, a transpose, and one more copy insidemat_to_df(). That is roughly60 * n^2bytes, i.e. ~6GB at n = 10000.The matrix path had a second copy of the same problem. With
zero.rm = FALSEit builtis.na(tmat) | !is.na(tmat)-- an all-TRUE matrix of the same shape -- purely so thatwhich()could be handed every cell. For a sparse input that forces the sparse class tostore all n^2 cells, which is heavier than a plain dense logical, and
which()thenbuilds an n^2 x 2 index matrix on top of it.
What changed
R/util.R,upper_gather()only. No caller was changed, andmat_to_df()is untouched(the vector path simply no longer calls it). The signature and the defaults are the same.
from the dist vector.
stats::diststores the lower half in column-major order, whichis the same sequence as the upper half in row-major order, so the values land in their
original order with no matrix involved.
zero.rm = FALSE: no more all-cell materialisation. The survivingindices are known in closed form, so only the half that passes the row/column comparison
is generated, in the same column-major order
Matrix::which(arr.ind = TRUE)produced.zero.rm = TRUE(thepair_countpath) is untouched --tmat != 0keeps a sparsematrix sparse, so it never had the problem.
Equivalence
bench/micro/upper_gather_equivalence.R(added) compares this implementation against theverbatim previous one -- extracted from the base revision, not retyped -- with
identical(), so row order, row count, column names, column types, NA positions and classall have to match.
49 cases: the six argument combinations the callers actually use, every caller's exact
arguments, degenerate cases (n = 1, n = 2, all-NA, all-zero, single surviving pair, 1x1),
non-square and wide matrices, sparse
dgCMatrixinputs with and without dimnames, andmultibyte / punctuation-heavy / near-colliding / duplicated dimension names.
47 of 49 are
identical(), one is an error whose message matches exactly, and two arethe intentional divergence below. The harness names those two explicitly and exits non-zero
on any other difference.
One intentional behaviour change -- please review
Duplicated dimension names on the vector path. The previous implementation raised
The `.data` argument of `add_column()` must have unique names as of tibble 3.0.0frominside
mat_to_df(). This one returns the natural result with the duplicated names kept.This also removes an inconsistency: the matrix path already accepted duplicated names
and returned a result (case D24 in the harness is
identical()between old and new, bothsucceeding). Only the vector path errored, and only because of an internal
tibbledetailleaking out, not because of anything the function guarantees. Making the two paths agree
seems right, but it is a behaviour change, so flagging it rather than burying it.
Benchmark
bench/micro/upper_gather_bench.R(added). Same arguments the callers use(
na.rm = FALSE, zero.rm = FALSE),gc()peak, row count checked equal every time.Measured locally on R 4.6.1 (Apple silicon):
Vector path (dist):
Matrix path, sparse input:
On a Linux box with more headroom the vector path reached 10.3x at n = 4000.
What this does NOT fix
Worth being explicit, so nobody reads more into the numbers than is there:
are gone, but the output itself is n^2 rows x 3 columns whenever
na.rm = FALSE, zero.rm = FALSE(which is 4 of the 5 callers), i.e. ~24 * n^2 bytes.Reducing that means changing what the function returns, which is a separate decision and
is deliberately not in this PR.
do_cmdscaleis not fixed by this. It rebuilds an n x n matrix of its own viasimple_cast()+t()+as.dist(). It does get faster input, but its own allocationis untouched.
Testing
test_util.R,test_pairwise.R,test_pair_count.R,test_stats_wrapper.Rall passlocally under
devtools::load_all()on R 4.6.1 -- zero failures. Those are the fourfiles that reach
upper_gather(30test_thatblocks).test_thatblocks added totest_util.Rpinning the output contract the callersdepend on but nothing tested: the full n x n grid, the row-major order, which half
carries values, column types, and verbatim dimension names.
Two follow-ups noticed while doing this, not addressed here:
do_kl_disthas no test coverage at all -- it is the only one of the five callerswith none.
asymmetry is pre-existing; it was preserved deliberately, but it does not look intended.
Note on R version
The new code uses
sequence(nvec =, from =), which needs R >= 4.0.0.DESCRIPTIONstill nominally says
R (>= 2.10), which was already inaccurate, and the product shipsR 4.x -- so this is not a practical constraint, but it is a real one.