Libertem: K2IS: efficient uint12 decoding

[sk1p]

See #63 for previous discussion, https://github.com/LiberTEM/LiberTEM/tree/master/prototypes/k2/uint12-decode for initial analysis and reference implementations.

[sk1p]

I pushed an initial implementation in d04f866187c48fee768ec3a7a10b30c03c8eb205 and the following commits. In a micro benchmark, on an old i5-2410M, it is about 3 times faster compared to the numba jit version (on a Xeon(R) W-2195, the difference is not as pronounced, factor ~1.8).

For vectorization, eliminating the output buffer was necessary, writing directly to the output works well.

[sk1p]

Using AVX512 (-march=skylake-avx512 -mtune=skylake-avx512) on the Xeon W2195 makes quite a difference (only tested on a single threaded workload for now): it is then again ~2.8 times faster than the numba jit version. This is what the compiler generates: https://godbolt.org/z/Elv9ZJ

[sk1p]

As an interesting fact, the memcpy call is optimized to 4 vmovups instructions!

[sk1p]

Together with mask application, processing a 202GB K2 dataset with the C++ implementation of uint12 decoding now takes 24s, down from 34s. Now the remaining work is to make the extension build system work as good as possible:

• should we publish the decoder as a separate Python package, and make it an optional dependency?
  
  
  
  • if we don't do this, developing on Windows may be painful: the developer would need to install a complete C++ toolchain, and make it work for Python extension building
  
  
  • on the other hand, we get some overhead: we need to handle separate releases, testing infrastructure / CI, ...
  
  
  • Maybe some overhead can be saved if we use a monorepo-approach: keep the sources in one GitHub repository, but have multiple Python packages (incl. setup.py etc.) inside.
  
  
  • Or we can keep the extension in the main LiberTEM Python package, and just make building it optional.
  
  
• how do we support different CPU instruction sets and select the right version at runtime? Is there a standard way to build a .cpp file into multiple .o files, for each ISA supported, and select the right implementation at runtime? Kind of like a fat binary for multiple optimized builds. The answer: yes, gcc FunctionMultiVersioning, more on function multi-versioning, the same in clang, clang documentation, and the FAT_RUNTIME flag of hyperscan as an example implementation (does it use the multi versioning of clang/gcc? I don't think so!).

[sk1p]

So just with something like:

 __attribute__((target_clones("avx2","arch=atom","avx512", "avx2", "default")))
void uint12_decode(...)

We can make gcc (version >=6) build multiple versions of our decoding routines that are dynamically selected based on CPU capabilities. Sadly, it doesn't work as-is with clang, but the basic building blocks should be there. Multi-versioning should also only really be considered for the distributed binary builds, and are not really useful for day-to-day development.

Travis CI does supports a current clang/gcc with something like this, so we could get away with requiring a specific compiler version for the multi versioned build.

[uellue]

@sk1p :+1: great! How possible would it be to use numba for the math? I've tried it before with my experiments, and numba went back to Python because of type casts so that it was slow. However, I haven't tried view() yet. That should do only one typecast per buffer.

What is the best way to try that out? Technically, the Python reference implementation in the Jupyter notebook should be quite possible to put in place of the current numba version?

That would save us a LOT of trouble with development and deployment!

[sk1p]

I tried to implement with numba, but didn't manage to improve on the naive version yet, that just loops over the input and processes 3 bytes at a time. A quick comparison:

First one is the numba implementation based on the jupyter notebooks of @uellue, second one is the naive version, third one is the C++ implementation.

[sk1p]

I got the "naive" version of the decoder to auto-vectorize by partially unrolling the loop (to be more specific: by helping the compiler to unroll the loop). Annotating input and output with __restrict__ is also necessary for vectorization. In my tests it was as fast as the other C++ implementation or faster. More next week.

[uellue]

Perfect! The less code, the better, the faster, the better. :-)

I'm a bit curious why the numba version is slower.

• Is it because it doesn't have the unrolling and equivalent of __restrict__ for the compiler to vectorize?
• Is it compiling for a target without SIMD?
• Is the source code or intermediate code more complicated for the compiler so that the optimizer can't do as much work as in pure C/C++, in the sense of assumptions about content and use of variables?
• Is this a result of more advanced language-specific optimizers for C++ compared to what numba can use?
• Is the resulting code doing more complicated things because of being derived from Python?

I'm just asking because being able to produce highly optimized, vectorized machine code from Python with numba would reduce the development effort @sk1p laid out in https://github.com/LiberTEM/LiberTEM/issues/108#issuecomment-419073198

In that sense it would be AWESOME if a numba version could reach the performance of a C/C++ extension.

[sk1p]

Hmm, good questions. To be clear, we are talking about the "reference" numba implementation that is currently used by the k2is reader, right?

Is it compiling for a target without SIMD?

It is definitely compiling for an AVX2 target, it is using AVX instructions for some data movement. I'll try to help the compiler unroll the loop on the numba version, too.

Is it because it doesn't have the unrolling and equivalent of __restrict__ for the compiler to vectorize?

I tried to find if __restrict__ is necessary with numba, or if there is some syntax for that but didn't find anything yet.

[sk1p]

The numba guys pointed me in the direction of https://github.com/numba/numba/issues/3205 - "Numba is possibly generating code that handles the wrap-around negative indexing that is unnecessary because the index expressions are all non-negative. We've observed this in a few situations, and are trying to come up with a fix."

[sk1p]

See also: https://stackoverflow.com/questions/44735756/python-reading-12-bit-binary-files/45070947#45070947

[sk1p]

The version by @max9111 (I hope I'm tagging the right user :grin:) is vectorized by numba and almost as fast as the C++ version (15 µs vs 18 µs)! I think the crucial difference to our numba impl is that the loop is incremented by 1 instead of 3.

The encoding is a bit different - the left part of the middle byte belongs to the left uint12 vs the right uint12. I'll try to reformulate our decoder using increments of 1 in the loop and see if it gets vectorized or if the encoding plays a role here.

[sk1p]

In d074b4a3b8d531e394a718bc39bec8058c416d26 I landed a new optimized numba implementation of uint12 decoding.

In August I measured the following performance:

Performance Comparison (note: these numbers were obtained on a 4 core Xeon E3-1505M):

• K2IS reader: 15s for 8GB
• converted to raw, dtype=uint16: ~3.7s (delegating the tiling to BLAS)
• converted to raw, dtype=uint16: ~11s (with tileshape=(1, 1, 930, 16))
• converted to raw, dtype=uint16: ~1.7s (with tileshape=(1, 16, 8, 2048))
• converted to raw, dtype=uint16: ~4.5s (with tileshape=(1, 1, 1860, 2048))
• converted to raw, dtype=float32: ~1.5s (delegating the tiling to BLAS)
• converted to raw, dtype=float32: ~1.7s (tileshape=(1, 16, 8, 2048))

Dataset converted to uint16 takes up 8.5GB, converted to float32 ~17GB.

And a bit later:

With the changes in 291b890 we now process the 8GB dataset in ~3.4 seconds, which is about 2.4GB/s.

Now, reading and processing the 8GB dataset takes 2s, which is nearing the performance of using the dtype=uint16 converted dataset.

This translates into processing the 202GB dataset in 21.8s on the 18 core machine, once loaded into fs cache (previously took 34s).

I think we should measure how much time is still spent in decoding and if it makes sense to optimize it further.

[max9111]

Did you also try a parallelized version [nb.njit(parallel=True) and changing the loop to nb.prange]?

This breaks SIMD-vectorization but is about 20% faster on a 4C/8T. On a 18C/32T machine this could be much faster than the single threaded SIMD-vectorized version.

A multithreaded and SIMD-vectorized version may be possible by using Pythons multithreading module (Numba releases the GIL) or some compiler hints which force vectorization.

[sk1p]

Thanks for the pointers!

Indeed, at first I tested the version as-is on your stackoverflow post, with nb.prange and parallel=True, and it was ~10x slower than the vectorized single-threaded version we are using now.

In the context of our program, we are parallelizing on a much coarser level - we split our data into partitions, and each partition is submitted as a future into a dask.distributed cluster (currently a local "cluster" using multiple processes, in the near future it will run on a "real" cluster!).

Btw: the numbers above are not the naked uint12 decoding, but include the processing. In theory, the single-threaded decoding should process about 4.8GB/s determined by a micro-benchmark, but I like to measure the real impact :grin:

[sk1p]

Thinking about it some more: I measured the parallel version using %timeit in a jupyter notebook, and on a quite small buffer (as I'm doing the same in the real workload), so maybe there was some overhead starting up new threads for each call, which could be amortized by working on larger buffers. Also, I'm still using numba 0.39, and I heard the new version has quite some improvements in that area. Nevertheless, I'm quite happy with the results now, as the single-threaded version also fits well into our overall architecture.

[uellue]

That is awesome, I'm impressed. I think we'll rely more on Numba in the future. :-)

[uellue]

The current Numba version of the decoder has all the properties we want:

• High throughput, enabled through vectorization by the compiler
• Simple, understandable, compact code
• Avoids a C build chain for development or installation

It reaches the same order of magnitude in throughput as the fastest dedicated bit packing and unpacking implementations like https://github.com/lemire/simdcomp and https://github.com/powturbo/TurboPFor that go through a lot more trouble and use hand-optimized architecture-specific intrinsics.

I'm closing this issue for now since @sk1p and I are satisfied with the current version and right now can't think of a better way to do this. :-) If you do have a better idea, let us know!