Large scale perf - #194
Conversation
The genetic value computation is slow for a trait with a large number of weakly causal sites, so add a benchmark that measures how it scales. The script simulates 100k sample nodes over 5Mb, caching the tree sequence, and times genetic_value for 1, 100, 1000, 10,000 and 100,000 causal sites at each of the individual, node and edge levels. The mutation rate is raised to ten times the human rate so that there are enough sites in a genome short enough to simulate quickly; this leaves the allele frequency spectrum, and so the weakly causal regime, intact. sim_trait is timed separately because it has a per-site Python loop of its own, the numba kernels are compiled by an untimed warm up call, and --max-seconds gives up on the larger points of the grid once a single call goes over budget. On a tree sequence with 217k nodes the cost per causal site is flat from 1000 sites upwards, at 444us for individuals, 384us for nodes and 1807us for edges, so 100,000 causal sites take 44s, 38s and 181s respectively. Holding the number of causal sites fixed and varying the sample size, the cost per site normalised by num_nodes is constant across a sixteen fold range, so the computation is O(num_causal * num_nodes): every causal site pays several passes over every node in the tree sequence. The median causal site is carried by 0.4% of the nodes.
Replace the tree by tree loop in genetic_value with one descent over the ARG that accumulates every causal site at once, using the tskit child index to find the child edges of a node and carrying a range of causal site indexes down so that a subtree spanning no causal site is pruned. The causal allele state change of a mutation telescopes down a path, so values are additive from root to node and nested and back mutations need no special handling; the has_mutation pruning is gone, and so is the separate accumulation pass for edges, which now fall out of the descent by crediting the edge that a value arrives through. The value carried down is a sum over the sites in the current range, so when an edge narrows that range it is not the value the child should inherit: the mutations above at the sites that dropped out do not apply below. The path back to the root is therefore kept, and the value is recomputed over the narrowed range when that happens. It is rare, a hundredth of a percent of descents on a large tree sequence. This is slower than what it replaces and is committed as a checkpoint only. On 100k samples at level="node", 1000 causal sites go from 0.436s to 14.4s and 10,000 from 3.761s to 51.7s. Cost is independent of allele frequency as intended, rare and uniform causal sites taking the same time, but the per visit constant is around 40 times worse than a prototype that only counted visits, once the value lookups, the per edge mutation searches and the output writes are in the loop. Add naive_genetic_value, the tree by tree implementation kept as a reference oracle, and TestGeneticValueReference comparing against it over the tests/data.py tree sequences, all_trees_ts(2..5) with recurrent and back mutations, and simulations with and without recombination, at each of the individual, node and edge levels and with multiple traits, along with isolated samples, multiple roots and mutations above a root. Give the benchmark a causal site selection axis, since drawing sites uniformly is dominated by the common variants in the tail of the frequency spectrum and the rare ones behave quite differently.
Replace the range carrying descent with a single sweep of the nodes from the past to the present. Each node conceptually holds a set of (causal site, effect) tuples, seeded at the node of each causal mutation. A parent is always older than its child, so time order is a topological order and every node is reached after all of its ancestors: for each tuple the outbound edges spanning its causal site are found, and the tuple is credited to the edge and to the node at the far end before being passed on to it. Nodes are credited when a tuple arrives rather than when they are swept, so a node that is never a parent never holds a tuple, which is about half of the work on a large tree sequence. Each tuple carries a single causal site rather than a range, so there is nothing to narrow and none of the path recomputation the descent needed. Tuples live in an arena threaded as a per node linked list, with a free list, since every tuple of a node is dead once that node has been swept and only those in flight need to be held. The causal allele being the ancestral state needs no separate treatment either: there is no mutation to seed from, so the roots are seeded instead and the effect reaches exactly the nodes of the tree. Mutations above a root need nothing at all, since seeding at the node and pushing down is already right. Cost is the number of nodes that carry a causal allele, rather than the size of the tree sequence, so the gain depends entirely on allele frequency. On 100k samples at level="node", against the tree by tree implementation, causal sites drawn from those below a frequency of 0.001 go from 0.436s to 0.060s at 1000 sites, 3.761s to 0.099s at 10,000, and 38.374s to 0.665s at 100,000. Causal sites drawn uniformly over all sites are around 2.8 times slower throughout, because the common variants in the tail of the frequency spectrum touch 7.4% of the nodes on average where the median site touches 0.39%, and a dense pass over every node is sequential where following carriers is not. Retarget the jit tests at the new kernel, dropping the binary search helpers that only existed for the range searches.
Replace the hand rolled array of linked lists with a typed list per node,
holding indexes into the seed arrays. The causal site and effect size of
a seed are fixed for the whole sweep, so an item is one integer and is
looked up rather than copied. A node's list is made when something first
reaches it, since most nodes are never reached when the causal alleles
are rare, and dropping the reference once the node has been swept hands
the storage straight back. This removes the arena, the free list and the
doubling block that was written out twice.
All figures below are level="node" on a tree sequence of 100,000 samples,
217,091 nodes, 285,979 edges, 22,523 trees and 235,458 sites, taking
causal sites either uniformly over all sites or from those with an allele
frequency below 0.001.
Against the array of linked lists it replaces, per causal site count:
1,000 10,000 100,000
uniform typed 0.656s 5.615s 54.8s
manual 1.121s 10.427s 109.7s
1.71x 1.86x 2.00x
rare typed 0.076s 0.163s 0.558s
manual 0.060s 0.099s 0.665s
0.79x 0.61x 1.19x
It is faster wherever the structure is under any pressure, and slower
only where the lists are a few items long and the sweep is short enough
that making them is a noticeable part of it, a difference of tens of
milliseconds. Peak memory at 100,000 uniformly drawn causal sites falls
from 4.57GB to 1.22GB, because each node's storage is released as the
sweep passes it rather than being held to the high water mark with up to
twice as much again in slack.
Against the tree by tree implementation this branch started from, which
took 0.436s, 3.761s and 38.374s for 1000, 10,000 and 100,000 causal
sites, and whose cost comes from passes over every node rather than from
the number of carriers:
1,000 10,000 100,000
rare 5.7x 23.1x 68.8x
uniform 0.66x 0.67x 0.70x
The remaining loss on uniformly drawn causal sites is the common variants
in the tail of the frequency spectrum, which touch 7.4% of the nodes on
average where the median site touches 0.39%. Following carriers is random
access at around 50ns each, while a dense pass over every node is
sequential at under a nanosecond, so the two cross over at a carrier
fraction of one or two percent.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #194 +/- ##
=========================================
Coverage 100.00% 100.00%
=========================================
Files 9 9
Lines 417 510 +93
Branches 50 63 +13
=========================================
+ Hits 417 510 +93
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
The tree sequence the ARG sweep was written against takes nine minutes for a full grid and 48 seconds for its longest single call, which is too slow to work against, and the benchmark measured wall time and nothing else. The peak memory in the last commit message and the carrier fractions in the one before it were both produced out of tree. Add a --preset flag and default it to a simulation of 30,000 samples over 1Mb, which has 63,287 nodes against 217,091 and runs the grid in thirty seconds. What makes it a fair substitute is not the timings but the distribution underneath them, since the sweep costs the number of nodes that carry a causal allele: the mean fraction of the nodes that a causal site reaches is 7.19% against 5.99%, and the median 0.406% against 0.184%, which is as close as 150 sampled sites from a distribution with that tail can show. Both presets then have a flat per call floor, a uniform cost per causal site that falls to a plateau from 1000 sites upwards, a rare cost still falling at the top of the grid, a uniform to rare ratio growing to two orders of magnitude, and parity between the three levels. The plateau is 160us against 483us per causal site, a factor of 3.0 on a node count ratio of 3.4. Sixteen times smaller was tried and is measurably less faithful, and the grid stops at 10,000 causal sites because raising the mutation rate to lift the ceiling would take sites per node from 0.68 to 2.7 against the large preset's 1.08 and inflate the setup floor. Measure four things besides the wall clock. --phases times the check, the setup, the kernel and the dataframe separately, without which a flat setup cost and a kernel that grows with the causal sites cannot be told apart; setup goes from 8ms to 17ms across the whole small grid and the dataframe is 1ms, so at 10,000 causal sites the kernel is 98% of the call. --counters runs a counting only copy of the sweep, since perf cannot see inside the kernel. --structure reports the carrier fractions above. --memory reports peak RSS, resetting VmHWM through clear_refs because it never falls on its own. The counters say three things that were assumed otherwise. The scan of a node's out edges is not where the waste is, because 94% of the trips find an edge spanning the seed's causal site. Uniform selection costs 178 times as many scans as rare at 10,000 causal sites but only 46 times the time, because rare costs 63ns a scan against 28ns, which is the random access against sequential difference measured directly rather than inferred. And the kernel has a floor in the number of nodes: it builds a pending entry for every node and sweeps every node whether or not anything reached it, so one rare causal site, ten scans, still takes 2ms. Add profile_genetic_value.py, which profiles one cell of the grid with cProfile for the setup and sets up a cell to run just the kernel under perf, printing the commands. Two things had to be established for that to be worth anything. Source lines inside the kernel are not available, because this llvmlite has no LLVM PerfJITEventListener and so nothing writes a perf map or a jitdump; what perf does give is the split between the kernel, the numba runtime with its typed list functions named, the interpreter and compilation. And NUMBA_ENABLE_PROFILING=1, the documented way to profile numba, is the wrong thing to use here: it would only help through the listener that is missing, and it defaults NUMBA_DEBUGINFO to 1, which measured 2.62s a call against 1.61s. perf finds the JIT mappings by itself. Record that _GeneticValue takes the _root_runs branch, a Python loop over every tree costing 7ms on the small preset and 38ms on the large one. It is not the edge case it looks like: drawing 10,000 sites uniformly makes it near certain that one of them has the ancestral state as its causal allele, and it fires at 10,000 and 100,000 causal sites on both presets, so on the large one it is over a third of the setup. Check in the small preset grid to diff against. The counts in it are identical between runs and are the part worth treating as a regression test; the timings are specific to the machine they were taken on.
The three levels differ only in which slot a contribution is credited to, and the kernel already discards a negative slot, so individuals need no more than a node to individual mapping: nodes_individual is exactly that array, and tskit.NULL for a node belonging to no individual is already negative. That removes the separate O(num_nodes) accumulation pass the individual level ran after the descent, one per trait, and with it _accumulate_individual_values. Accumulate into the caller's output array rather than allocating one per trait, so that a second implementation can write into the same table and combining the two is nothing at all. Split the causal mutation expansion in two. A descent of the trees needs every mutation at a causal site, because a mutation blocks the inheritance of the allele above it whatever it changes the state to, where a sum of state changes needs only those that change the state. _row_mutations returns all of them and _causal_mutations filters, which is what the push down seeds and edge_effect keep using. Keep the row each seed came from for the same reason: a subset of the rows can then be selected without recomputing the seeds. No change in what is computed. Verified against the tree by tree implementation at 31fec9a on 500 causal sites over a 30,000 sample tree sequence, at each level and for one and three traits, agreeing to 3e-14.
A descent of a tree needs the tree, so add the incremental construction that a second implementation of genetic values will run over: apply the edges leaving and then entering at each tree of tskit's TreeIndex, maintaining the quintuply linked encoding along with the edge above each node. Attaching and detaching a child are both constant time, so a full pass costs one insertion and one removal per edge: 2.25ms over the 75,586 edges of a 30,000 sample tree sequence and 12.35ms over the 285,979 edges of a 100,000 sample one, which is less than the setup that the push down already pays on the same data. There is no virtual root. The descent starts at the mutations of a causal site and needs the roots only when the ancestral state is the causal allele, which is 0.04% of the sites of the smaller tree sequence and 0.07% of the larger. Maintaining the virtual root's children costs something on every edge, and knowing which nodes belong in that list means tracking the samples below every node; walking up from each sample and taking the top of the path finds the same roots on demand, and marking the nodes already walked through keeps the whole thing to one visit per node however many samples there are. That leaves one difference from tskit: it threads the roots together as children of the virtual root, so they are siblings of each other there and have none here. Nothing else differs, and the descent never walks the siblings of a root, so the test asserts the sibling arrays match at every node that has a parent and are null at the roots, along with the parent, left child and edge arrays and the root set itself, tree by tree over every fixture in tests.data, all_trees_ts(2..5), comb trees, multiple roots, isolated samples and a tree sequence with no nodes at all.
The push down of the ARG costs the nodes that carry a causal allele, at
around 28ns each because following carriers is random access. A descent
of the trees costs the same nodes at around 20ns, since a child is the
next link of a list rather than a search of the out edges of a node, and
it needs no typed list per node to hold what is in flight. Against that
it has to build the trees, which the push down never does.
So neither wins everywhere, and which one a causal site should take
follows from its allele frequency: the tree pass is a fixed cost that
common sites amortise and rare ones do not. Add the descent as a second
implementation and route each row to one or the other on allele_freq,
which sim_trait already returns and _check_trait_df now keeps when it is
there. A trait dataframe assembled by hand has no such column and takes
the push down for everything, which is what it did before.
Measured on 30,000 samples over 1Mb at level="node", end to end:
1,000 10,000
uniform before 195ms 1624ms
after 129ms 1080ms
1.51x 1.50x
rare before 16.3ms 31.0ms
after 16.6ms 31.1ms
Rare sites are below the threshold and take the same path as before, so
the gain is on the common ones, which is where the time was: sites above
a frequency of 0.03 are under a third of the sites and 94 to 97% of the
time.
The descent walks the rows in step with the trees, which works because
_check_trait_df already requires the rows to be sorted by site. A row's
mutations are marked with the row's own index rather than into an array
that has to be cleared, so a descent costs the nodes it reaches and
nothing per node of the tree sequence. All of the traits share the pass,
where the push down runs once per trait.
Two things the descent has to do that the push down does not. It takes
every mutation at a causal site, not only the state changing ones, since
a mutation blocks the inheritance of the allele above it whatever it
changes the state to. And where the causal allele is the ancestral state
it seeds the roots, skipping any root that carries a mutation: the push
down seeds every root and lets the root's own mutation cancel it, and the
descent has no such cancellation, so seeding an ancestral root that
carries the allele would count it twice. Removing that skip fails
test_mutation_on_root and two of the reference comparisons.
Run every comparison against the tree by tree reference through both
implementations and a mix of the two, and add tests for the cases where
they are most likely to part company: a mutation on a root under both
causal alleles, a causal site carrying no mutations at all, a site
exactly on a breakpoint, and that a middling threshold really does divide
the rows rather than quietly sending them all the same way.
Compare the two implementations over the whole benchmark grid, on both
presets, at all three levels, for one and three traits, and for causal
sites drawn both uniformly and from the rare ones. A threshold of zero,
sending everything to the descent of the trees, was fastest or tied in
every cell:
uniform rare
1,000 10,000 1,000 10,000
small descent 122ms 1058ms 15ms 23ms
push down 196ms 1619ms 16ms 35ms
large descent 427ms 3709ms 71ms 90ms
push down 592ms 4917ms 68ms 108ms
The push down led only where the causal sites are few and rare, so the
tree pass is not amortised, and there it led by a millisecond or two on a
call taking a few. With three traits the gap widens the other way, to
2.27x on rare causal sites, since the push down sweeps once per trait
where the descent takes all of them in one pass.
So set the threshold to zero. Every allele frequency is at or above zero,
so there is nothing to compare against and nothing to look up, which is
what lets the default apply to a trait dataframe assembled by hand as
well as to one from sim_trait. The push down is kept, and --threshold inf
still runs it, only so that the two can go on being compared.
Against the tree by tree implementation this branch started from, at
level="node" on 100,000 samples, the whole of it is now ahead:
1,000 10,000 100,000
tree by tree 0.436s 3.761s 38.374s
rare 0.072s 0.091s 0.216s
6.1x 41x 178x
uniform 0.443s 3.848s 32.589s
0.98x 0.98x 1.18x
which closes the regression on uniformly drawn causal sites that the push
down left behind at 0.70x.
Report the work the descent does rather than the work the push down does,
since that is the code that runs. The descent kernel returns the number
of nodes it visited, so there is no second copy of the loop to keep in
step with the first, and the count is the thing its run time is
proportional to: seconds over visits is around 20ns once there are enough
causal sites to amortise the tree pass. Visits per row as a fraction of
the nodes comes to 8.4% for uniformly drawn causal sites and 0.02% for
rare ones, against the 7.2% and 0.4% that --structure measures from the
other end.
The _root_runs branch is only taken by rows that go to the push down, so
at the default threshold it never runs; say so rather than reporting that
any ancestral causal allele triggers it.
|
The seed is also reflected well, and the two simulations produced identical phenotype and trait dataframe. |
The push down holds the seeds still in flight, which grows with the number of causal sites; the descent holds a fixed handful of arrays the length of the nodes whatever the trait looks like. At 100,000 uniformly drawn causal sites on 100,000 samples, that is 0.01GB over the baseline against 0.89GB, and a peak of 0.42GB against 1.30GB.
|
Excellent, thanks @daikitag. I'm working on further improvements, which should bring it down more. |
Two implementations of the genetic value computation have been carried
side by side so that they could be compared. The comparison is done, so
delete the one that lost, along with the threshold that chose between
them and everything built only to feed it.
Measured over the whole benchmark grid, on both tree sequences, at each
level, for one and three traits, and for causal sites drawn both
uniformly and from the rare ones, the push down was slower everywhere the
cost mattered, by 1.33 to 1.70 times on uniformly drawn causal sites and
up to 2.27 times on rare ones with three traits, where it swept once per
trait against a single pass for all of them:
uniform rare
1,000 10,000 1,000 10,000
small descent 122ms 1058ms 15ms 23ms
push down 196ms 1619ms 16ms 35ms
large descent 427ms 3709ms 71ms 90ms
push down 592ms 4917ms 68ms 108ms
It was ahead in ten of the thirty six cells, every one of them rare
causal sites at a thousand or fewer, where the tree building the descent
starts with has too little to amortise it against. In none of them was it
ahead by more than a few milliseconds on a call taking a few. The note
left in jit.py says as much, since a reader can check it against the same
benchmark.
Deleting it takes the setup with it, which was built on every call
whether the push down ran or not: an argsort over the nodes, two
searchsorted over the edge table, the tskit child index, and the seed
arrays. _GeneticValue setup on 100,000 samples goes from 98ms to 42ms,
almost all of what is left being tskit's jitwrap, and every cell of the
small preset grid gets between 1.01 and 1.41 times faster with no change
to the kernel at all.
_root_runs goes too. It was a Python loop over every tree, finding the
roots for causal sites whose ancestral state is the causal allele, and
only the push down needed it: the descent is already standing in the tree
and asks _tree_roots. The threshold argument, the optional allele_freq
column that it read, and the row selection the descent kernel took all go
with it, so _check_trait_df is back to the four required columns and
sim_trait's allele_freq is once again only an output.
The thirty five topology tests carry over unchanged by repointing the one
fixture they share; the jit and nojit parametrisation survives, though
nojit now interprets only the descent's own loop, since the tree building
kernels it calls are compiled. Those are covered instead by TestTreeState
against tskit.Tree. The tests that existed to compare the two
implementations go, and the reference comparison against the tree by tree
oracle stops looping over thresholds.
Verified against that same tree by tree implementation, recovered from
31fec9a, on both tree sequences at all three levels for one and three
traits.
|
I've changed the algorithm, which should be a bit better. Can you try it again please @daikitag ? |
Coverage of jit.py fell to 62% when the push down went. It had called no other kernel, so running it through py_func interpreted everything it did; the descent calls five kernels to build the trees, and a numba function calls whatever the name is bound to in the module, so those five stayed compiled and the coverage of their bodies went with them. Swap the whole set for their py_func together when the nojit parameter asks for it, so that a kernel is interpreted the whole way down rather than only in its own loop, and give the tree state test the same parameter as the rest. That is what the module docstring has always claimed the convention is; correct the caveat that said otherwise. jit.py is back to 100%, and so is every other module, on branches as well as lines. Running only the nojit half covers jit.py completely and only the jit half covers 16% of it, so the measurement is coming from the interpreted path rather than from somewhere incidental.
|
I think we can merge this, assuming that it does in fact lead to performance improvement over the previous version. |
|
CPU Utilized: 02:03:34 The new codes took a slightly longer time to simulate traits but with a lower memory. |
|
Huh, 124m vs 45 minutes is quite a big difference. That's completely against what my local benchmarks were saying. Always good to try on the real data! I'm going to try parallelising the current algorithm though. If that scales over multiple threads, it's a much simpler algorithm with predictable memory usage so could still win. Thanks @daikitag - I'll ping you again when there's something to try out. |
|
Thank you for modifying the codes, Jerome! There is still a substantial improvement in the current codes compared with the previous codes, so it will definitely improve the usability of the tstrait package on large tree sequences. |
The rows of the trait dataframe are independent of each other, so add a
num_threads argument that gives each of that many worker threads a
contiguous range of them. It defaults to 0, which does the work on the
calling thread with no pool created, matching what tskit's
divergence_matrix and its neighbours mean by the same argument.
Releasing the GIL is what makes it worth anything: a numba kernel holds
it for its whole execution by default, so the threads would have taken
turns. With nogil=True on the descent, four of them report a busy to wall
ratio of 3.8 out of 4.
A range needs nothing from any other. The kernel takes row_start and
row_stop rather than a slice, so nothing is copied or rebased: pair_offset
is indexed by absolute row, and the arrays a thread only reads are built
once and shared. Each range accumulates into its own table, since the
kernel adds into it and that is not atomic, and the tables are summed at
the end, which is where the answer stops being bit for bit what one
thread would have produced. The rows are marked with their own index, and
two ranges cannot collide.
Minimum of two replicates at level="node" on four cores:
preset selection num_causal 1 thread 2 3 4
small uniform 10,000 980ms 1.94x 2.86x 3.41x
small uniform 300 41ms 1.36x 1.66x 1.69x
small rare 10,000 17ms 0.94x 0.91x 0.89x
large uniform 10,000 3338ms 1.47x 1.64x 1.64x
large rare 10,000 76ms 1.01x 1.00x 0.99x
Two things in there are worth saying out loud, since neither is what one
would assume. A bigger trait parallelises and a bigger tree sequence does
not: the kernel holds around eleven arrays the length of the nodes per
thread, 49 bytes a node, which is 3.1MB on 30,000 samples so four threads
sit inside a 16MB L3, and 10.6MB on 100,000 samples so four want 42MB and
are held up by memory bandwidth instead. Going from 10,000 causal sites
to 100,000 on the larger one moves four threads from 1.50x only to 1.61x,
so this is cache capacity rather than the division of work. And threads
cost more than they save on a rare or a small trait, because a thread
walks the whole tree sequence whatever range it takes: there is no
seeking to the first tree a range wants, and that pass is nothing against
a second of descending and everything against seventeen milliseconds.
Equal row counts turned out to need no cost model behind them. They were
the thing most likely to want one, since rows differ enormously in how
many nodes they reach, but four ranges came out within 1.11x of each other
on time and 1.18x on nodes visited.
The sequential path is unchanged, within 0.96x to 1.09x of the committed
baseline over five replicates.
|
I've just added a It seems to be cache-sensitive, so it would be good to try num_threads=4 and num_threads=8 separately to see how well it scales with such a large ARG. |
|
There was a slight improvement in computational time when we increase the number of threads. |
|
Hmm, well we halved the time with 4 threads, but it's still slower than the arg algorithm. Seems like it's better to revert to the arg algorithm after all. Can you tszip that tree sequence file and send it to me on slack please? My benchmarks aren't capturing your use case. |
|
I tried to zip the tree sequence file, but it was too big to download on my computer. |
|
Why can't you run tszip on the server? |
|
Just to clarify here @daikitag - are you using tstrait end-to-end to simulate the traits and compute the genetic values, or are you computing genetic values based on traits simulated by SLiM? On slack you said you loaded the effect size dataframe and used This will matter as the performance here depends a lot on the allele frequency, which could differ significantly. |
|
@jeromekelleher and I did not use SLiM information to simulate the traits. |
|
Can you profile |
|
I used the following codes to measure the times.
|
|
Ok, great. We're getting a reasonable return from threads here, so I think we'll commit this and tag a release. Thanks for your help @daikitag ! |
|
Thank you very much for modifying the algorithm @jeromekelleher ! |


Currently performance is very poor on large tree sequences with large numbers of causal sites. This PR changes the core algorithm from a tree-by-tree traversal approach to a whole-ARG traversal.
@daikitag would you mind trying this out to see if it helps in your case? There may be some memory issues, so do report back on what you see.