Tskit: Convert tsk_id_t/tsk_size_t to 64 bit.

[jeromekelleher]

It seems that requiring metadata columns to being < 4G is an unwelcome limitation in forward simulation applications, so we should consider upgrading the metadata_offset columns (and other offset columns) to 64 bit integers. Here is what it would entail:

1. Create a tsk_offset_t typedef and go through the C tables API making sure this is used for all offset columns (metadata_offset, ancestral_state_offset, etc). Typedef this to uint64_t.
2. Increment the file-format minor version. In the file writing code, look at the last value in each offset array. If it's < UINT32_MAX, store it as a 32 bit value; if not, store as a 64 bit. In the reading code, check which type is being used to store and update accordingly. Note that this means we don't be able to use the current zero-copy behaviour where we use the memory in the kastore to back the arrays. Figure out the best approach here (note, using the memory in the kastore was motivated by using mmap for io, which we've dropped now as it's inherently dangerous and hard to do properly cross-platform).
3. Clean up _tskitmodule.c to use the correct new sizes for numpy arrays (possibly also needing to backport this over to msprime where we use the LightweightTableCollection for interchange).

From a user perspective, this will mean that old versions of tskit won't be able to read newer files. New versions of tskit will continue to read older files without issues, as we're making the code more flexible in terms of expected types.

pinging @petrelharp and @molpopgen for opinions.

[jeromekelleher]

Note: arguably we should go the whole hog here and typedef tsk_size_t to uint64_t and tsk_id_t to int64_t. It's less clear that we really do need more than 2*31 rows in a table though, and this will definitely affect memory usage and (possibly) performance. Perhaps not as much as we might think though, so as we're disrupting things it might be as well to try this out and see what the perf implication really are.

[petrelharp]

It seems like we need to allow bigger tables (see #402), but on the other hand, the implications for memory usage are pretty unfortunate. The least bad other solution I can think of is to have tables switch over from 32-bit to 64-bit once they get big, although that'd be quite annoying.

[jeromekelleher]

It seems like we need to allow bigger tables (see #402), but on the other hand, the implications for memory usage are pretty unfortunate. The least bad other solution I can think of is to have tables switch over from 32-bit to 64-bit once they get big, although that'd be quite annoying.

I think we'll probably just have to bite the bullet and make everything 64bit. Having a tsk_table_t and a tsk_table64_t would be a huge amount of extra work. Let's do a prototype and see what the actual memory costs of going to 64 bit are. It'll probably make no practical difference for the vast majority of tree sequences.

We should make the storage code more flexible though, and only store in 64 bit if we need to. This will create a little bit of forward-backward incompatability, but nothing we can't manage (the versioning infrastructure is in there just for this type of thing).

[molpopgen]

The memory impacts will be appreciable for larger sims. For example, if I modify the script from #402 to just be plain W-F:

initialize() {
        // initializeSLiMModelType("nonWF");
        initializeTreeSeq(simplificationInterval=500);
        // defineConstant("Np", 50);
        // defineConstant("K", 1000000);        // total carrying capacity
        // defineConstant("Kp", asInteger(K / Np));

        initializeMutationType("m1", 0.5, "f", 0.0);
        m1.convertToSubstitution = T;

        initializeGenomicElementType("g1", m1, 1.0);
        initializeGenomicElement(g1, 0, 600);
        initializeMutationRate(0);
        initializeRecombinationRate(0);
}
// reproduction() {
//      subpop.addCrossed(individual, subpop.sampleIndividuals(1));
// }
1 early() {
        print(time());
    sim.addSubpop("p1",1000000);
        // for (i in 1:Np)
        //      sim.addSubpop(i, Kp);
}
// early() {
        // for (subpop in sim.subpopulations)
    //     print(subpop.individualCount);
        // for (subpop in sim.subpopulations)
        //      subpop.fitnessScaling = Kp / subpop.individualCount;
// }
500 late() {
        print(time());
        sim.outputFixedMutations();
}

And run the following command:

/usr/bin/time -f "%e %M" -o bigslim_time.txt ./slim bug2.slim

The timing is:

cat bigslim_time.txt 
Command exited with non-zero status 1
6636.40 31636984

So, that's 32GB. Roughly, we can reason that the memory is mostly taken up by edges. The simulation has no mutation, so we'll assume the size of individual genomes is negligible. Changing the node integer to uint64_t bumps the total size of an edge up by 1/3, so we can guesstimate that the RAM use will go up accordingly, so about 40GB.

The same model in fwdpy11 runs through to completion, and includes the simplification and the edge table indexing. The following is run against fwdpy11/dev:

import fwdpy11
import numpy as np

N = 1000000

pop = fwdpy11.DiploidPopulation(N, 1.)

pdict = {'gvalue': fwdpy11.Multiplicative(1.),
         'nregions': [],
         'sregions': [],
         'recregions': [],
         'rates': (0, 0, 0),
         'demography': np.array([N]*500, dtype=np.int32)
         }
params = fwdpy11.ModelParams(**pdict)

rng = fwdpy11.GSLrng(42)

# the 500 is the time b/w simplifications.
# here, the parameter means once at the end.
fwdpy11.evolvets(rng, pop, params, 500)

The same type of command for timing gives:

cat bigfp11_time.txt 
687.72 67584080

So, pushing 70GB. It may be important that the back-end does not copy edge tables to sort. Rather, the sorting is in-place. I'm not convinced, though, that that difference affects peak RAM use, as the memory would be freed prior to simplification. In any case, upping that number by 1.3 is close to 90GB.

Simplifying every 100 generations is a big help:

614.79 14609144

I have not tested the data dump from the fwdpp table collection format to tskit, but I could if there is interest--maybe I'd get an over flow there? in any case, I'm not totally convinced that #402 is entirely due to overflow. Perhaps there's a subtle bug elsewhere?

[molpopgen]

Quick update. i hacked fwdpp to use 64 bit signed integers for the tree sequences and reran the example from above. The memory use wasn't as bad as I'd guessed, so my loose approximation was too loose:

cat bigfp11_time.txt 
738.32 79347220

So, that is about a 10% performance hit, but I have to admit to laziness. I merely changed a typdef. There are probably loads of narrowing conversions, etc., going on. I was just playing.

[jeromekelleher]

Thanks @molpopgen, this is good info. The more I think about it, the more I wonder whether we should be using tables for forward simulations at all though. (I'll send you an email about this)

Do we think that the simplified output of a forward simulation is every likely to need 64 bit tables?

[molpopgen]

Do we think that the _simplified_ output of a forward simulation is every likely to need 64 bit tables?

Tricky question. Personally, I am not interesting in attempting to simulate "all humans", but others may be. More generally, for tsinfer, if you think about the rate at which adding a sample adds new nodes, what does a "Continent Biobank" do to the total number?

For this particular issue, we need offsets larger than the current for metadata. For #402, we haven't totally established what is going on, or at least I'm not clear that we have.

[jeromekelleher]

Tricky question.

Indeed!

[petrelharp]

Update here; #402 was indeed due to running out of space in the metadata column; other columns (eg the node ID column) are a factor of 10 smaller, but will still be hit pretty quick in sims of size 1M. Some additional discussion there.