Repository navigation
Refactor - #8
Merged
Merged
Refactor#8
Conversation
Add late-stage tsinfer workloads (1000 Genomes chr20p ancestor groups 1483-1485 and sample matching) to the compression experiments, and measure the proposed delta and affine encodings on them. On these workloads a quarter to a third of the stored Viterbi entries were redundant: when edges change, nodes get explicit copies of their inherited values, and the engine stored the traceback before coalescing them away. That churn limited a delta encoding of the traceback to 1.1-1.7x. Viterbi now coalesces before storing the traceback, and no longer after. A node is removed only if its value and its recombination_required flag both equal the ones it inherits, so the traceback, in which unlisted nodes take the flag of their nearest listed ancestor, is unchanged. Nodes kept for their flag are removed at a later site. Forward and backward still coalesce after storing, on values only. What was tried: 1. Flag-aware coalescing before storing plus the old pass after: minimal lists, but two passes per site made tsls 1.36-1.48x tsinfer-hmm. 2. A single flag-aware pass before storing, with an uncached walk to the nearest listed ancestor: +2-4% instructions over the old engine. 3. As 2, with the inherited flag cached along tree paths alongside the inherited likelihood (recombination_required_cache): landed. It uses no more instructions than the old engine (callgrind: -0.2% on 20 group 1485 ancestors, -0.3% on one sample), and head-to-head wall times agree within noise. Results: the stored lists now equal the change points, Viterbi churn falls 5-8x, engine traceback entries fall 31-47%, and a delta encoding on top saves a further 5-8x. tsls still finds tsinfer's exact path on all benchmark queries. The compression experiments gain: - a tsinfer case in capture.py, with --cases and merged manifests; - churn, value change, flag change and affine residual measures in analysis.py; - an Encodings section and work plots in report.py. The README records the findings and the sequence of attempts.
tsls stores a smaller traceback than tsinfer-hmm but was still slower, because storing the traceback is under 1% of the time in both engines. Profiling both on the same 1000 Genomes chr20p queries (callgrind, limited to ancestor_matcher_find_path) showed tsls doing 1.26x (samples) to 1.49x (late ancestors) the instructions, mostly in the allele lookup and in building the first tree of each query. What was changed, and how it landed: 1. Allele lookup. tsinfer-hmm sets root 0 to the ancestral state, so its walk up the tree checks only for a set allele; tsls allows several roots and checked for the root at every step. tsls now sets every root's allele (roots are all in the likelihood list) before each site's update, and the Viterbi, forward and backward loops use get_marked_allelic_state. A first version kept an assert in the walk; because the Python module is built with asserts on, it cost as much as the check it replaced (12.71 G against 12.50 G instructions on one sample), so the walk has no assert. The traceback keeps get_allelic_state. 2. First-tree build. For Viterbi, init_sweep now builds the first tree with plain edge insertions and removals (ancestor_matcher_next_tree_topology), as tsinfer-hmm does, and no longer initialises num_below. Forward and backward, which track num_below, are unchanged. Instructions, tsinfer-hmm / tsls before / tsls after: 20 group 1485 ancestors 6.65 / 9.89 / 7.44 G; one sample 9.92 / 12.50 / 10.96 G. Paired wall times (best of 3, alternating engines per query), tsls over tsinfer-hmm: group 1485 ancestors 1.35 -> 1.10, samples 1.06 -> 1.06. Change 2 is the win; change 1 saves instructions but no time, because the walk is limited by memory access. run.py ratios proved unreliable for comparing engines, since it runs the engines in blocks and machine speed drifts; the README documents this and the paired method. Paths are unchanged: tsls still finds tsinfer's path on every query checked, and the tsls, C and benchmarking tests pass.
benchmarking/callgrind.py runs the first queries of each tsinfer chr20p workload with one engine under callgrind, collecting only inside the engine's find_path call, and prints total Ir per workload and engine. Baseline before the tsls_hmm_t refactor (tsls at bafe726), 20 ancestors per group and 1 sample, paths identical between engines: | workload | queries | tsinfer | tsls | ratio | |---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.608 G | 2.055 | | ancestors_g1000 | 20 | 0.621 G | 0.975 G | 1.571 | | ancestors_g1483 | 20 | 2.388 G | 2.845 G | 1.192 | | ancestors_g1484 | 20 | 9.911 G | 11.137 G | 1.124 | | ancestors_g1485 | 20 | 6.653 G | 7.445 G | 1.119 | | samples | 1 | 9.922 G | 10.955 G | 1.104 | On g84, 57% of tsls's instructions are memset: resetting per-site and per-node arrays on every query, which dominates for short ancestors.
The engine (ancestor_matcher_t in matcher.c) becomes tsls_hmm_t in hmm.c, and works directly on a borrowed tsk_treeseq_t instead of the copied edge, site and mutation data in matcher_indexes_t, which is removed. - Trees are visited with tskit's tsk_tree_position_t. The first tree of a sweep is found from a null position with tsk_tree_position_seek_forward or _seek_backward (whichever end is closer, as tsk_tree_seek_from_null does), inserting only the edges that cover the first site, instead of applying every edge from the start (or end) of the sequence. The traceback does the same from the last site of the window, instead of walking in from the end of the sequence. - Alleles are read from the site and mutation tables and encoded by the HMM: by default numbered at each site as tskit does (ancestral 0, then derived states in order of first appearance), with at least two alleles per site so that a tsinfer ancestor's focal sites (which have no mutations yet) can still be matched; or as indexes into a given list (e.g. ACGT for sc2ts). encode_alleles and the ancestral/derived state arrays are gone, along with the error codes tskit now guarantees. - Coordinates are doubles; integer coordinates are no longer required, and path_left/path_right are doubles. - Kept as simple as possible: the root-marked allele lookup and the topology-only first tree from bafe726 are dropped (the seek replaces the latter). The core algorithm (likelihood lists, coalescing with flags, caches, traceback reuse, pruning, forward/backward) is unchanged. Python: MatcherIndexes is replaced by tsls.TreeSequence (a minimal wrapper of tsk_treeseq_t built from a LightweightTableCollection, shareable between HMMs), and AncestorMatcher by tsls.Hmm(ts, recombination, mismatch, *, alleles=None, ...). The extension now compiles all of tskit's C library. Callers in benchmarking, verification and experiments are updated. Tests: new C and Python tests for allele encoding, windows starting and ending throughout the sequence against the dense optimum (exercising seeks from both ends), non-integer coordinates, and sites without mutations. tsls finds tsinfer's path for all 1,914 queries in the six chr20p workloads; traceback sizes and memory are unchanged. Instruction counts (callgrind, find_path only) on the reduced profiling set (see the next commit): | workload | queries | tsinfer | tsls | ratio | |---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.486 G | 1.644 | | ancestors_g1000 | 20 | 0.621 G | 0.686 G | 1.105 | | ancestors_g1483 | 10 | 1.355 G | 1.381 G | 1.019 | | ancestors_g1484 | 3 | 0.646 G | 0.736 G | 1.139 | | ancestors_g1485 | 5 | 1.131 G | 1.117 G | 0.988 | | samples | 1 (100000 sites) | 2.064 G | 2.578 G | 1.249 | On the full set the old engine's ratios were 2.06, 1.57, 1.19, 1.12, 1.12 and 1.10, so the seek more than pays for the dropped special cases on ancestors, while the sample (where the first tree is cheap anyway) is worse, mostly from dropping the root-marked allele lookup.
The first harness took over 20 minutes per engine, mostly in setup: under callgrind, loading a large reference tree sequence and building the engine took 4-5 minutes per workload before any matching. Now: - Setup runs uninstrumented (--instr-atstart=no), and the driver switches instrumentation on with callgrind_control just before the queries. Setup for g1485 fell from 235 s to 35 s. Collection is still limited to find_path, so the counts are the same. - Fewer queries, sized at about 2G tsinfer-hmm instructions per workload (matching runs at about 36M instructions/s under callgrind): 20 g84, 20 g1000, 10 g1483, 3 g1484 and 5 g1485 ancestors, and 1 sample restricted to the middle 100,000 of its sites (a whole sample is 10G). - tsinfer-hmm's counts never change, so they are measured once and reused from callgrind/tsinfer/; each run measures only tsls. - Paths are saved next to the callgrind files and compared between the engines. A run measuring both engines takes 2.6 minutes with 4 jobs; tsls alone is about half that.
find_path_matrix reset the traceback and max_likelihood_node arrays over every site of the sequence after each call. Each holds an entry per site (581k for chr20p, 16 MB together), so for a short tsinfer ancestor the resets cost far more than the matching: on g84 (1,390 sites per ancestor) memset was 71% of all instructions. Only sites in [start, end) can have been set, so only these are reset now, including when the run fails. Instruction counts (find_path only), before -> after: | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.486 G | 0.162 G | 0.547 | | ancestors_g1000 | 20 | 0.621 G | 0.686 G | 0.362 G | 0.583 | | ancestors_g1483 | 10 | 1.355 G | 1.381 G | 1.221 G | 0.901 | | ancestors_g1484 | 3 | 0.646 G | 0.736 G | 0.688 G | 1.065 | | ancestors_g1485 | 5 | 1.131 G | 1.117 G | 1.036 G | 0.916 | | samples | 1 (100000 sites) | 2.064 G | 2.578 G | 2.565 G | 1.242 | Callgrind counts each iteration of memset's inner loop, so the wall time saving is likely smaller than these counts suggest, but the resets did touch 16 MB per query. Paths unchanged.
Phase 1 dropped this special case from bafe726 to keep the refactor simple. The profile showed the root check in get_allelic_state's walk up the tree costing 0.57G of the sample's 2.57G, more than the whole gap to tsinfer-hmm, so it earns its place: every root's allele is set to the ancestral state from the likelihood list (roots always have likelihoods) before the site's update, and the three likelihood loops walk with get_marked_allelic_state, which has no root check and no assert. | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.162 G | 0.159 G | 0.537 | | ancestors_g1000 | 20 | 0.621 G | 0.362 G | 0.340 G | 0.547 | | ancestors_g1483 | 10 | 1.355 G | 1.221 G | 1.171 G | 0.864 | | ancestors_g1484 | 3 | 0.646 G | 0.688 G | 0.608 G | 0.941 | | ancestors_g1485 | 5 | 1.131 G | 1.036 G | 0.963 G | 0.852 | | samples | 1 (100000 sites) | 2.064 G | 2.565 G | 2.232 G | 1.081 | Paths unchanged.
The walks up the tree in get_inherited_likelihood, its _and_flag variant and insert_mutation_likelihoods had an assert at every step. The Python module builds with asserts on, and tsinfer-hmm's equivalent walks have none. The invariant (every root has a likelihood) is still checked by tsls_hmm_check_state under extended checks, and the once-per-call value asserts remain. The NULL_NODE test in reset_likelihood_cache's walk looks similar but is needed: within one tree update a node with a cached value can become a root (when a later edge in the batch is removed), so the walk can reach the top of the tree. | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.159 G | 0.158 G | 0.534 | | ancestors_g1000 | 20 | 0.621 G | 0.340 G | 0.337 G | 0.543 | | ancestors_g1483 | 10 | 1.355 G | 1.171 G | 1.153 G | 0.851 | | ancestors_g1484 | 3 | 0.646 G | 0.608 G | 0.602 G | 0.932 | | ancestors_g1485 | 5 | 1.131 G | 0.963 G | 0.955 G | 0.845 | | samples | 1 (100000 sites) | 2.064 G | 2.232 G | 2.200 G | 1.066 | Paths unchanged.
The profiling sample was restricted to the middle 100,000 of its sites, which made both engines build the first tree half way along the sequence. A whole sample (as tsinfer matches it) starts at the first site, so this measured seek costs that the real workload doesn't have: about 100M of tsls's 2.20G. The window is now the first 100,000 sites. tsinfer-hmm's counts are remeasured (its cache files are keyed by the window). | window | tsinfer | tsls | ratio | |---|---|---|---| | middle 100k sites | 2.064 G | 2.200 G | 1.066 | | first 100k sites | 1.642 G | 1.736 G | 1.057 | Most of the gap remains, so it is not mainly the seek. Parity will also be checked on a whole sample once at the end.
insert_edge was called out of line 778k times on the sample (36M instructions), and get_inherited_likelihood 770k times, where tsinfer-hmm does the same work inline in its tree loop. Marking these and remove_edge/reset_likelihood_cache inline saves a little. | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.158 G | 0.158 G | 0.533 | | ancestors_g1000 | 20 | 0.621 G | 0.337 G | 0.336 G | 0.541 | | ancestors_g1483 | 10 | 1.355 G | 1.153 G | 1.147 G | 0.846 | | ancestors_g1484 | 3 | 0.646 G | 0.602 G | 0.600 G | 0.928 | | ancestors_g1485 | 5 | 1.131 G | 0.955 G | 0.951 G | 0.841 | | samples | 1 (first 100k sites) | 1.642 G | 1.736 G | 1.716 G | 1.045 | Investigated and not the cause: the cache-reset walk's NULL_NODE test. Both engines spend the same on cache resets (tsls 107M; tsinfer-hmm 80M + 24M), so making that walk cheaper would not close the gap. Paths unchanged.
Each site marked every root's allele before its update and unmarked them after, with two passes over the likelihood list (38M of the sample's 1.72G; tsinfer-hmm has a single root and no such passes). Now every node without a parent is permanently marked with ROOT_ALLELE during the sweeps: init_sweep sets the marks with a memset, insert_edge and remove_edge update the child's, and unsetting a mutation restores the mark if the node is a root. get_marked_allelic_state turns ROOT_ALLELE into the site's ancestral state with one compare per walk. The traceback, which doesn't maintain marks, starts from a cleared array. | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.158 G | 0.161 G | 0.544 | | ancestors_g1000 | 20 | 0.621 G | 0.336 G | 0.339 G | 0.547 | | ancestors_g1483 | 10 | 1.355 G | 1.147 G | 1.157 G | 0.854 | | ancestors_g1484 | 3 | 0.646 G | 0.600 G | 0.599 G | 0.928 | | ancestors_g1485 | 5 | 1.131 G | 0.951 G | 0.963 G | 0.852 | | samples | 1 (first 100k sites) | 1.642 G | 1.716 G | 1.691 G | 1.030 | This trades about 1% on ancestors (marking the first tree's edges and clearing for the traceback) for 1.5% on the sample, the one workload still above tsinfer-hmm. A first version set the marks in init_sweep's per-node loop, which cost 2-4% on ancestors (0.973G on g1485); a memset fixed most of that. Dead end, reverted: marking the traceback tree's roots from the site's traceback list (which contains every root), so the traceback's allele lookup could also skip the root check. The two extra passes over the list cost more than the check: the sample went from 1.716G to 1.734G, and every ancestor workload got worse too. Paths unchanged.
The walks in coalesce_likelihoods cost the same as tsinfer-hmm's (133M + 106M against 134M + 100M on the sample), but tsls spent about 40M more around them: get_inherited_likelihood_and_flag was not inlined despite 'inline' (17M of per-call setup), each call asserted value >= 0 (5.8M), and the loop tested the invariant check_recombination_required for every node (5.8M). The helper and coalesce_likelihoods are now always_inline, so each call site gets a copy specialised for its constant flag, and the per-call value asserts are gone (as for the walks in the last change). The flag cache stores (10M) remain; they are the cost of coalescing with recombination flags, which halves the traceback. | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.161 G | 0.160 G | 0.542 | | ancestors_g1000 | 20 | 0.621 G | 0.339 G | 0.338 G | 0.545 | | ancestors_g1483 | 10 | 1.355 G | 1.157 G | 1.153 G | 0.851 | | ancestors_g1484 | 3 | 0.646 G | 0.599 G | 0.597 G | 0.924 | | ancestors_g1485 | 5 | 1.131 G | 0.963 G | 0.961 G | 0.849 | | samples | 1 (first 100k sites) | 1.642 G | 1.691 G | 1.680 G | 1.023 | Paths unchanged.
The traceback looked up the matched allele with get_allelic_state, which checks for a root at every step of the walk (34M on the sample, against 23M for tsinfer-hmm's walk to its single root). Now every root of the traceback tree has the virtual root N as its parent (the tree arrays already have N + 1 entries), and the virtual root's allelic state is set to the site's ancestral state, so the walk needs no root check. Setting up the parent array is one fill per traceback instead of a memset. get_allelic_state is no longer used and is removed. This replaces the dead end in 4c56547 (marking the traceback roots from the site's traceback list), which cost two passes over the list per site. | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.160 G | 0.156 G | 0.529 | | ancestors_g1000 | 20 | 0.621 G | 0.338 G | 0.328 G | 0.528 | | ancestors_g1483 | 10 | 1.355 G | 1.153 G | 1.133 G | 0.836 | | ancestors_g1484 | 3 | 0.646 G | 0.597 G | 0.591 G | 0.915 | | ancestors_g1485 | 5 | 1.131 G | 0.961 G | 0.950 G | 0.840 | | samples | 1 (first 100k sites) | 1.642 G | 1.680 G | 1.670 G | 1.017 | Paths unchanged.
reset_likelihood_cache tested u != NULL_NODE at every step, because within one tree update a node with a cached value could become a root: when an edge above a cached node is removed later in the same batch. Now remove_edge clears the cut child's cached value. That value is redundant, since the child is given the same value as its likelihood. So no root ever has a cached value. Values are only cached on nodes without likelihoods, and every root has one, so every cached chain ends below a root and the walk can stop at the first unset value. This is the same loop as tsinfer-hmm's (which relies on its single root never being cached). Checked under valgrind memcheck with the dense comparisons, window, Viterbi matrix and pruning tests (138 tests with extended checks, which also verify that every cached value is unset after each update): no invalid reads, so the walks never leave the tree. | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.156 G | 0.155 G | 0.523 | | ancestors_g1000 | 20 | 0.621 G | 0.328 G | 0.323 G | 0.520 | | ancestors_g1483 | 10 | 1.355 G | 1.133 G | 1.110 G | 0.819 | | ancestors_g1484 | 3 | 0.646 G | 0.591 G | 0.583 G | 0.902 | | ancestors_g1485 | 5 | 1.131 G | 0.950 G | 0.939 G | 0.830 | | samples | 1 (first 100k sites) | 1.642 G | 1.670 G | 1.638 G | 0.998 | tsls now uses no more instructions than tsinfer-hmm on every workload of the profiling set. Paths unchanged.
The previous change matched the first 100k sites of the sample against the whole reference, and showed tsls at parity (0.998). A one-off run on the original full set (20 ancestors per group, 1 whole sample) disagreed for the sample: | workload | tsinfer | tsls | ratio | |---|---|---|---| | ancestors_g84 (20) | 0.296 G | 0.155 G | 0.523 | | ancestors_g1000 (20) | 0.621 G | 0.323 G | 0.520 | | ancestors_g1483 (20) | 2.388 G | 1.877 G | 0.786 | | ancestors_g1484 (20) | 9.911 G | 9.575 G | 0.966 | | ancestors_g1485 (20) | 6.653 G | 5.940 G | 0.893 | | samples (1 whole) | 9.922 G | 10.287 G | 1.037 | The window was biased: tsinfer-hmm's traceback walks in from the end of the sequence, which costs about 60M for a window at the start but nothing extra for a whole sample, while tsls seeks to the window's last site. So the first-100k window overstated tsinfer-hmm's cost by about 4%. (The ancestor windows don't have this problem: real ancestors are matched that way.) Now the sample's reference is cut at site 100,000 (keep_intervals and rtrim, node IDs unchanged) and the sample is matched over the whole of it, so both engines start at the first site and finish at the last, as for a whole sample. This gives 1.561G against 1.618G (1.036), in line with the whole sample's 1.037, so parity has not yet been reached on samples.
When nodes left the tree at a tree transition, each was removed from the likelihood list with its own pass over the list (delete_likelihood, as in tsinfer-hmm): 96M on the sample, with about 5 nodes leaving per transition. They are now marked NOT_IN_TREE and removed together with a single pass, only when some node has left. | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.155 G | 0.154 G | 0.522 | | ancestors_g1000 | 20 | 0.621 G | 0.323 G | 0.320 G | 0.516 | | ancestors_g1483 | 10 | 1.355 G | 1.110 G | 1.067 G | 0.787 | | ancestors_g1484 | 3 | 0.646 G | 0.583 G | 0.567 G | 0.878 | | ancestors_g1485 | 5 | 1.131 G | 0.939 G | 0.919 G | 0.812 | | samples | 1 (cut at 100k sites) | 1.561 G | 1.618 G | 1.525 G | 0.976 | (The ancestor 'before' numbers are from 0cd8ab3; the sample's is on the cut reference from ea952ad.) Considered and not done: tsls's remaining extra work on the sample is mostly the recombination flag bookkeeping (about 35M: the flag cache stores, the flag compare and the traceback reuse check), which buys the halved traceback, and tskit's tree position scans (about 20M more than tsinfer-hmm's inline edge loops), which come with using tsk_tree_position_t. Packing the flag into the cached value's sign bit would save perhaps 10-15M but needs a new CACHE_UNSET sentinel and care with -0.0; not worth it now. Paths unchanged.
benchmarking/paired.py times each query three times with each engine, alternating, and keeps the best (as in the README's earlier paired measurements), so drift in machine speed doesn't bias one engine. The README now records the instruction counts on the original set after the refactor and optimisations (tsls 0.52-0.97 of tsinfer-hmm, including a whole sample), and the paired wall times, which go the other way: tsls takes 1.14-1.43 times as long. Diagnosis (in the README): building a tsls.Hmm re-encodes every site's alleles (4-8 ms per query against 1.5 ms), and find_path misses the cache more, despite fewer instructions: tsls has 25% more L1 read misses, 2.2x the L1 write misses and 35% more last-level read misses on the profiling sample (callgrind --cache-sim). The last-level misses are mostly tskit's tree position scans reading edges through the index arrays, and the tsk_mutation_t reads; the write misses come from the per-edge marks added by 4c56547 and 0cd8ab3. Also marks the earlier 'where the time goes' section as describing the engine before the refactor.
Paired timing showed that building a tsls.Hmm took 4-8 ms per query against tsinfer-hmm's 1.5 ms, because every HMM encoded the alleles of every site from the tables, which dominated for short ancestors. The read-only data the HMM derives from a tree sequence is now a tsls_index_t, built once (tsls_index_init) and borrowed by any number of HMMs (tsls_hmm_alloc takes the index). It holds the allele encoding and compact per-site mutation arrays (node and encoded derived state, with site offsets), which replace reading the 80-byte tsk_mutation_t structs. In Python, tsls.TreeSequence(ts, alleles=None) owns the tsk_treeseq_t and its index, and Hmm(ts, ..., alleles=...) builds one only when given a tskit tree sequence (alleles with a tsls.TreeSequence is an error). The benchmark engine builds its tsls.TreeSequence once per workload. Cache simulation (callgrind --cache-sim, find_path, profiling sample), before -> after: last-level read misses 1.53M -> 1.31M (tsinfer-hmm 1.14M); L1 read misses 78.9M -> 78.8M (63.0M); L1 write misses unchanged at 16.5-16.9M (7.7M). Instructions 1.525G -> 1.526G. Paired wall time (paired.py --quick, tsls/tsinfer-hmm): ancestors g84 0.735 (1.43 before, on 50 queries), g1485 1.143, samples 1.35 on 2 samples. Paired ratios on samples vary by about +-3% between runs of the same build (1.23 and 1.17 on 6 samples), so cache-miss counts are used to judge the next steps. Instruction counts are unchanged: | workload | queries | tsinfer | tsls | ratio | |---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.155 G | 0.523 | | ancestors_g1000 | 20 | 0.621 G | 0.321 G | 0.517 | | ancestors_g1483 | 10 | 1.355 G | 1.068 G | 0.788 | | ancestors_g1484 | 3 | 0.646 G | 0.568 G | 0.878 | | ancestors_g1485 | 5 | 1.131 G | 0.919 G | 0.813 | | samples | 1 (cut at 100k sites) | 1.561 G | 1.526 G | 0.977 |
The index now holds the edges in edge_insertion_order and edge_removal_order as 16-byte records (parent, child, and one coordinate), so that the ranges of a tsk_tree_position_t index them directly: tree updates, the traceback and the first-tree seek read them sequentially instead of reading four edge table columns through the order arrays. The seek only needs one coordinate: the edges in its range all have left <= x in insertion order (forward seek) and right > x in removal order (backward seek). Memory: 32 bytes per edge, the same as tsinfer-hmm's copied edge indexes. Paired wall time (paired.py, full set), tsls / tsinfer-hmm: | workload | queries | before (db2415e) | after | |---|---|---|---| | ancestors_g84 | 50 | 1.43 | 0.74 | | ancestors_g1000 | 50 | 1.19 | 0.78 | | ancestors_g1483 | 100 | 1.19 | 0.91 | | ancestors_g1484 | 50 | 1.23 | 1.02 | | ancestors_g1485 | 100 | 1.25 | 0.79 | | samples | 6 | 1.14 | 1.16 | ('before' is before the shared index too, which made most of the difference for g84.) Late ancestors gain most because their seek scans most of the edges, now sequentially. Callgrind's cache simulation is misleading here: it shows last-level read misses on the sample rising from 1.31M to 1.83M, because tskit's tree position scans still read the edge table columns, and the copies add a second stream (about 190k misses per loop). But the simulator doesn't model hardware prefetching, which hides most of the cost of a sequential stream, while the old reads through the order arrays were random. Wall time is the judge from here; the simulator is still useful for locating misses. Dropped: asserts in the seek loops checking the one-coordinate test against the edge table. The Python module builds with asserts on, and they added 25% to the instructions of late ancestors (the range covers most edges). The window tests check the seek against the dense model. Instruction counts: | workload | queries | tsinfer | tsls | ratio | |---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.153 G | 0.518 | | ancestors_g1000 | 20 | 0.621 G | 0.319 G | 0.514 | | ancestors_g1483 | 10 | 1.355 G | 1.060 G | 0.783 | | ancestors_g1484 | 3 | 0.646 G | 0.564 G | 0.873 | | ancestors_g1485 | 5 | 1.131 G | 0.910 G | 0.804 | | samples | 1 (cut at 100k sites) | 1.561 G | 1.522 G | 0.975 | Paths unchanged.
Documents the paired wall times after the shared index (ancestors now 0.74-1.02 of tsinfer-hmm, samples 1.07-1.16), two reverted experiments (removing the per-edge cache clear, and marking roots at each site again: no measurable change on samples), and a perf profile locating the samples' remaining gap in the likelihood cache walks, which execute about the same instructions as tsinfer-hmm's but miss the cache more.
The likelihood walks jump around per-node arrays of several MB each
(433k nodes for the chr20p ancestors), so with 4 KB pages most steps miss
the TLB. Large arrays (>= 1 MiB: the per-node arrays and the index's edge
copies) are now mapped directly, aligned to 2 MiB, with
madvise(MADV_HUGEPAGE) before first use (Linux only; elsewhere and for
small arrays, malloc). They are mapped directly because glibc reuses
memory already faulted in with small pages for large requests in a
long-running process (its mmap threshold grows after large frees), and
madvise doesn't change pages that are already there: with
posix_memalign + madvise, no huge pages were used at all.
Paired wall time, tsls / tsinfer-hmm (paired.py), two rounds run
alternately with the previous build (U) to share machine conditions:
| workload | U round 1 | U round 2 | this, round 1 | this, round 2 |
|---|---|---|---|---|
| ancestors_g84 | 0.75 | 0.75 | 0.72 | 0.75 |
| ancestors_g1000 | 0.76 | 0.75 | 0.76 | 0.74 |
| ancestors_g1483 | 0.92 | 0.91 | 0.91 | 0.80 |
| ancestors_g1484 | 0.99 | 1.00 | 0.96 | 0.79 |
| ancestors_g1485 | 1.06 | 1.04 | 1.00 | 1.00 |
| samples | 0.99 | 1.11 | 0.87 | 0.81 |
tsls's own times also vary less between runs (samples 16.99 and 19.24 s
before; 14.67 and 14.22 s now), while tsinfer-hmm's were steady
throughout. Run-to-run variation of tsls of up to 20% had made earlier
single comparisons unreliable; it fits page placement varying between
processes.
Dead ends (not committed), each run twice on the full paired set,
alternating with the then-current build:
- One 24-byte record per node {likelihood, likelihood_cache, parent,
allelic_state, recombination_required, its cache}, so that each step of
a walk reads one record: samples 1.52 and 1.52, ancestors 0.80-1.64,
all worse. Indexing 24-byte records also costs an address computation
per step: instructions on the profiling sample rose from 1.52G to 1.82G
(the allele walk alone 538M -> 804M).
- Split records, an 8-byte {parent, allelic_state, flags} and a 16-byte
{likelihood, likelihood_cache}: no consistent difference from separate
arrays (samples 1.16 and 0.94 against 1.17 and 1.12; ancestors within
noise).
Paths unchanged. Instruction counts are unaffected.
Replaces the interim wall-time notes with: how paired.py measures, the current result (tsls faster than tsinfer-hmm on samples by 13-19% and on most ancestor groups, level on group 1485), how each change moved it (after the refactor, + shared index, + huge pages), why single runs and callgrind's cache simulation mislead on this machine, and the node layouts that didn't help.
The per-node arrays the walks read are now two record arrays:
tsls_node_t, 8 bytes {parent, allelic_state, recombination_required, its
cache}, and tsls_likelihood_t, 16 bytes {value, cache}. The allele walk
reads one record per step instead of two arrays, and the likelihood cache
walks read two instead of four. Both arrays are backed by huge pages as
before; the tree's child and sibling links and the forward/backward
counts stay as separate arrays.
Without huge pages this layout made no consistent difference (see
9806f55). With them it is faster where most of the matching time goes.
Paired wall time, tsls / tsinfer-hmm, two full runs of each build in
alternation (this run's H rounds varied more than in 9806f55):
| workload | separate arrays (9806f55) | two records |
|---|---|---|
| ancestors_g84 | 0.76, 0.76 | 0.80, 0.81 |
| ancestors_g1000 | 0.75, 0.76 | 0.82, 0.82 |
| ancestors_g1483 | 0.80, 0.96 | 0.78, 0.78 |
| ancestors_g1484 | 0.81, 0.92 | 0.84, 0.82 |
| ancestors_g1485 | 0.99, 0.95 | 0.90, 0.87 |
| samples | 1.11, 0.95 | 0.79, 0.82 |
The early groups (4-8 ms per query) are slightly slower; the late groups
and the samples (0.1-3 s per query), which dominate tsinfer's matching
time, are faster and steadier.
The cost is instructions: the profiling sample goes from 1.52G to 1.76G
(tsinfer-hmm 1.56G), so tsls is no longer at instruction-count parity on
samples. GCC forms a pointer to the record at each step of the allele
walk (an extra lea), which is 538M -> 669M instructions there, but the
step still costs one memory access.
| workload | queries | tsinfer | tsls | ratio |
|---|---|---|---|---|
| ancestors_g84 | 20 | 0.296 G | 0.157 G | 0.532 |
| ancestors_g1000 | 20 | 0.621 G | 0.346 G | 0.557 |
| ancestors_g1483 | 10 | 1.355 G | 1.155 G | 0.852 |
| ancestors_g1484 | 3 | 0.646 G | 0.629 G | 0.974 |
| ancestors_g1485 | 5 | 1.131 G | 0.968 G | 0.856 |
| samples | 1 (cut at 100k sites) | 1.561 G | 1.764 G | 1.130 |
Paths unchanged.
The current result is now 0.78-0.90 of tsinfer-hmm's wall time on every workload, with a column for the node records in the progression table, and the layouts section notes that the two-record layout only helped once huge pages were used.
Later runs of the huge pages build varied as much as before (samples 14.2-18.0 s over four runs); only the two-record build's two runs were close (14.1 and 14.5 s).
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.
Refactor the engine to use as much of tskit as possible, then sequentially optimise to match the tsinfer HMM engine's performance.