Change Memory layout of positions from Fortran- to C-contiguous

[zemanj]

Is your feature request related to a problem? Please describe.
Currently, all functions that operate (especially in-place) on atom positions obtain a copy of atom positions via AtomGroup.positions, regardless of whether positions are modified in- or out-of-place, or they are merely input data that is not altered. In some (if not many) cases, a tremendous speed-up could be achieved if the respective functions could obtain a view instead of a copy of an AtomGroup's positions. At the moment, this is not possible for two reasons:

1. AtomGroup.positions gets its data from Universe.trajectory.ts.positions via fancy indexing, and fancy indexing means copying.
2. More importantly, Universe.trajectory.ts.positions refers to coordinates.base.Timestep._pos, which is a Fortran-contiguous array. Therefore, potentially obtained views of that array are unsuitable for use with C-level functions.

Describe the solution you'd like
Change coordinates.base.Timestep.order from 'F' to 'C' and provide a way of obtaining a (if possible, C-contiguous) view of an AtomGroup's positions.

Describe alternatives you've considered
Invent computers with two-dimensional main memory...

Additional context
Two candidates that come to my mind in terms of speed-up are *Group.center*() and *Group.wrap()
when the index array *Group.atoms.ix is a (contiguous) range.

A possible drawback of changing the positions/velocities/forces memory layout is that reading and (probably to a lesser extent) writing of DCD trajectories will be a bit slower since DCD stores positions etc. as Fortran-contiguous arrays.

If you have any additional concerns, please let me know!

[zemanj]

Trajectory Reading Performance Impact

I tested the runtimes of trajectory reading (from a RAM disk) for various trajectory formats.

tl;dr: As expected, DCD becomes significantly slower when holding positions in C-contiguous arrays. Reading any other format becomes either faster or the change is negligible.

Test code:

import os
import shutil
from collections import defaultdict
from warnings import filterwarnings
from timeit import timeit
import MDAnalysis as mda
from MDAnalysisTests.datafiles import *
filterwarnings('ignore')

###### HELPER FUNCTIONS ######

def copy_trajectories(destination, trajectories):
    for i in range(len(trajectories)):
        trjfile, name = trajectories[i]
        if isinstance(trjfile, tuple):
            trjfile = list(trjfile)
            for j in range(len(trjfile)):
                dest = os.path.join(destination, os.path.basename(trjfile[j]))
                shutil.copyfile(trjfile[j], dest)
                trjfile[j] = dest
            trjfile = tuple(trjfile)
            trajectories[i] = (trjfile, name)
        else:
            dest = os.path.join(destination, os.path.basename(trjfile))
            shutil.copyfile(trjfile, dest)
            trajectories[i] = (dest, name)

def universe(trjfile, name):
    if isinstance(trjfile, tuple):
        return mda.Universe(*trjfile, format=formats[name])
    return mda.Universe(trjfile, format=formats[name])

def readall(trjfile, name):
    trj = universe(trjfile, name).trajectory
    for _ in trj:
        pass

def get_atoms_frames(trjfile, name):
    u = universe(trjfile, name)
    return u.atoms.n_atoms, u.trajectory.n_frames

def time_trajectory_read(trjfile, name):
    n_atoms, n_frames = get_atoms_frames(trjfile, name)
    normfactor = 1.0 / (n_runs * n_frames * n_atoms)
    # Read once:
    readall(trjfile, name)
    # Get runtime per frame:
    t = timeit(lambda: readall(trjfile, name), number=n_runs) * normfactor
    # Print result:
    print("{:14s}     {:.3e}  {:>5}  {:>6}  {:>5}"
          "".format(name, t, n_runs, n_frames, n_atoms))

###### SETUP ######

n_runs = 100

formats = defaultdict(lambda: None)
formats['LAMMPSDUMP'] = 'LAMMPSDUMP'
formats['DLP'] = 'HISTORY'

ramdisk = '/media/ramdisk'

trajectories = [(CRD, 'CRD'),
                (DCD, 'DCD'),
                (DLP_HISTORY, 'DLP'),
                (DMS, 'DMS'),
                (GMS_ASYMSURF, 'GMS'),
                (GRO, 'GRO'),
                (GSD, 'GSD'),
                (LAMMPSdata, 'LAMMPS'),
                (LAMMPSDUMP, 'LAMMPSDUMP'),
                (MMTF, 'MMTF'),
                (MMTF_gz, 'MMTF_gz'),
                (mol2_molecules, 'MOL2'),
                (NCDF, 'NCDF'),
                (PDB_multiframe, 'PDB_multiframe'),
                (PDBQT_input, 'PDBQT'),
                (PDB_xvf, 'PDB_xvf'),
                (PQR, 'PQR'),
                (TRR, 'TRR'),
                ((TRZ_psf, TRZ), 'TRZ'),
                (TXYZ, 'TXYZ'),
                (XTC, 'XTC'),
                (XYZ, 'XYZ')]

copy_trajectories(ramdisk, trajectories)

###### RUNTIME MEASUREMENT ######

print("Format          s/frame/atom   runs  frames  atoms")
print("--------------------------------------------------")
for trjfile, name in trajectories:
    time_trajectory_read(trjfile, name)

Comparison of Test Results: new (C-contiguous) vs. old (F-contiguous)

             new, C-cont.  old, F-cont.  (old-new)/old
Format       s/frame/atom  s/frame/atom  relative gain
------------------------------------------------------
CRD             1.060E-05     1.084E-05          2.21%
DCD             2.554E-08     2.171E-08        -17.64%
DLP             1.284E-05     1.311E-05          2.06%
DMS             1.725E-05     1.821E-05          5.27%
GMS             6.183E-05     6.282E-05          1.58%
GRO             7.737E-06     7.962E-06          2.83%
GSD             5.728E-07     5.732E-07          0.07%
LAMMPS          1.462E-05     1.477E-05          1.02%
LAMMPSDUMP      3.420E-05     3.415E-05         -0.15%
MMTF            7.418E-06     7.407E-06         -0.15%
MMTF_gz         4.304E-06     4.347E-06          0.99%
MOL2            3.149E-06     3.175E-06          0.82%
NCDF            2.225E-08     2.262E-08          1.64%
PDB_multiframe  2.276E-06     2.315E-06          1.68%
PDBQT           5.230E-06     5.253E-06          0.44%
PDB_xvf         1.015E-05     1.024E-05          0.88%
PQR             4.238E-06     4.245E-06          0.16%
TRR             1.537E-07     1.555E-07          1.16%
TRZ             7.711E-07     8.087E-07          4.65%
TXYZ            1.495E-04     1.485E-04         -0.67%
XTC             4.778E-08     4.868E-08          1.85%
XYZ             1.512E-06     1.518E-06          0.40%
------------------------------------------------------
Average         1.584E-05     1.593E-05          0.58%

[orbeckst]

Thorough impact analysis!!!

• Is the impact on DCD still as big when you read from a spinning platter disk?
• What is the speed benefit on the other operations that you mentioned?

Note that historically the speed of DCD reading went down after the transition to cython, so DCD has already taken a hit compared to pre 0.16 (?).

[zemanj]

@orbeckst I tried running the tests on a HDD, SSD and also from an NFS but the test files are so small that the OS caches the files, therefore yielding identical results. I don't have a large DCD file available (i.e., larger than my 16 GB main memory). Or is there a way to enforce reads from disk?

[zemanj]

I'm pretty sure that the impact will be much lower if the file cannot be cached. The real bottleneck should then be disk I/O and the in-memory array transposition penalty will probably be hidden.

[zemanj]

@orbeckst Concerning the speedup of other operations, I noticed that wrapping positions of an entire medium-sized system becomes seven times faster. Of course, that speedup won't be available if the atomgroup indices are strided or shuffled. In such cases, one cannot get a C-contiguous view of the coordinates (i.e., a copy has to be made) and the performance is likely the same as before.
I will provide some more robust numbers in PR #2211 soon.

[zemanj]

@orbeckst

Note that historically the speed of DCD reading went down after the transition to cython, so DCD has already taken a hit compared to pre 0.16 (?).

Do you recall how it used to be implemented before?

I'm not a big fan of cython tbh. I'd prefer MDAnalysis being mainly a C++-based framework with a Python interface connected via pybind11. But that would basically require re-implementing major parts of the code base, and that's unlikely to happen anytime soon.

[kain88-de]

What is the effect of this on distance calculations and complex selections. The FORTRAN order is chosen because it fits to the compute intensive applications we have. Before a change is made this should be checked as well.

[zemanj]

@kain88-de It has the potential to improve the performance of all operations on positions/velocities/forces that obtain their data from an AtomGroup whose indices can be represented as a slice. Since many operations are performed on residues or molecules, whose atoms are often stored consecutively in a trajectory, one can avoid unnecessary copies of the respective position data.

What is the effect of this on distance calculations and complex selections.

For disttance calculations, which usually don't alter positions, one can potentially feed the algorithms with C-contiguous views instead of copies. Note that, however, it is not just the copying that slows down the operations, but also the fancy indexing itself that can be avoided if it is possible to represent an AtomGroup's indices (ix) by a slice.
Complex selections that are distance-based usually lead to AtomGroups with non-contiguous or non-uniformly strided indices. For such groups, the memory layout is irrelevant.

The FORTRAN order is chosen because it fits to the compute intensive applications we have.

Could you please elaborate? Do you have an example of an operation that requires a Fortran-contiguous memory layout?

I know that having positions as xxx...xxxyyy...yyyzzz...zzz in memory can yield huge performance boosts when optimizing for SSE/AVX vector ops. However, this is usually only true as long as data fits in the CPU cache. For larger amounts of data, it is better to adopt an L1-optimized and cache-aligned hybrid scheme where positions are stored in small Fortran-contiguous blocks.

AtomGroup.positions, by the way, always returns a C-contiguous copy of the current Timestep's positions. That is (up to now) independent of the memory layout of coordinates.base.Timestep._pos.

Before a change is made this should be checked as well.

I absolutely agree!

[tylerjereddy]

I'd prefer MDAnalysis being mainly a C++-based framework with a Python interface connected via pybind11.

-1 on that; in my view the language choices are also to enable the broader Python community to make contributions & not to hit us too hard with reviewer burden.

[zemanj]

@tylerjereddy Point taken. Doesn't change my opinion on cython, though. 😉

[richardjgowers]

I think Tyler is right wrt this package, it’s a python package so really we should use Cython. That said, I’d be curious what speed up we can get if we rewrite our distance code. I know you’ve tinkered with this before, maybe we could do a C++/pybind package that provides basic distance calculations (like BLAS/lapack) in a non MDA specific way.

…

On Thu, Feb 28, 2019 at 4:05 PM, Johannes Zeman ***@***.***> wrote: @tylerjereddy <https://github.com/tylerjereddy> Point taken. Doesn't change my opinion on cython, though. 😉 — You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHub <#2210 (comment)>, or mute the thread <https://github.com/notifications/unsubscribe-auth/AI0jB8L6wOgAgQFRvqGjcdOyxeZihPCmks5vSESsgaJpZM4bV6Zt> .

[zemanj]

Maybe I should clarify (and backpedal) a bit: I don't generally think that using Cython is a bad choice.
I just had quite bad experiences when allowing it to automatically mangle C++ templates into its auto-generated code, or when trying to use LTO between different compilation units. I also saw reproducible changes in performance of parts of the code when changing something (seemingly) unrelated a few lines away.

I therefore usually prefer having low-level C/C++ implementations in separate compilation units and use Cython just as glue code. That does not apply to functions which should be inlined (as I mentioned, tackling inlining with LTO can be problematic).

On the other hand, having the possibility to call numpy methods (most of which are hard to beat in terms of performance) from Cython is something invaluable, so there's usually a trade-off between pure Cython and hybrid C/C++/Cython implementations.

[tylerjereddy]

Slightly off-topic, but another thing I'm curious about are possible MD applications for __array_function__ and its low-level equivalent. Perhaps some of you have thought about that already wrt custom handling of coordinates or residues or units, etc.

[kain88-de]

Could you please elaborate? Do you have an example of an operation that requires a Fortran-contiguous memory layout?

Have a look at this loop. Here we go over the z-coordinate last. So we read a whole bunch of z values before y or x are updated. This allows us to make optimal use of a cache line on the z values, i.e. we can to several iterations of the loop without reading another cache line. This means the CPU can prefetch memory and keep the cores busy. Using a C-layout would require us to read from memory much more often as x and y values we do not use in the loop need to be read. Fetching memory is slow. This pattern keeps repeating for almost all compute-intensive calculations that MDAnalysis does. Therefore an F-layout makes some sense for us.

I do understand that it can have a performance improvement for slicing. And it's great you made the benchmarks to show how much!

Because compute and slicing needs orthogonal optimizations I think it's beneficial to benchmark both when we optimize one or the other.