Libertem: Mathematical formulation to allow using math libraries for applying masks to frames

[uellue]

For the mathematical operation it doesn't matter in what order or arrangement the frames are scanned. That means the "scan dimension" can be flattened to a linear array. The same way the arrangement of the pixels in a frame doesn't matter, that means masks and frames can be flattened to linear arrays as well.

With those extra dimensions flattened, applying a set of _m_ masks (_B_) to a collection of _n_ frames with _c_ pixels per frame (_A_) can be expressed as a rectangular matrix product of _A_ and _B_. _A_ is a _n x c_ matrix and _B_ a _c x m_ matrix.

That means we can use the BLAS function GEMM. For that we should be able to find a suitable fast implementation.

Notably:

• Scipy https://docs.scipy.org/doc/scipy/reference/linalg.blas.html#blas-level-3-functions
• Intel MKL https://software.intel.com/en-us/mkl-developer-reference-c-blas-level-3-routines
• CUDA http://docs.nvidia.com/cuda/cublas/index.html#cublas-lt-t-gt-gemm
• ATLAS http://math-atlas.sourceforge.net/
• OpenBLAS http://www.openblas.net/

[uellue]

In essence, we need a fast out-of-core GEMM implementation that is optimized to read from HDF5.

[uellue]

I've been playing with compiling numpy against OpenBLAS and Intel MKL and using the dot product as described above. It doesn't perform well at all (up to 3x slower than naive) when the matrices are very slim (less than 16 masks) in one direction and fat in the other. Multithreading brings a performance penalty under these circumstances and doesn't use the available cores. @sk1p will verify if this happens in a naked C version as well or if it is a numpy thing. Weird!

Edit: MKL DOES work OK for 16 masks and stacks of 64 or 128 frames.

[uellue]

Here's the result from OpenBLAS. It performs well with larger mask counts and stack heights. A single mask destroys performance. dieter-OpenBLAS.xlsx

And this is Intel MKL with similar results dieter-mkl.xlsx:

[uellue]

For comparison the manual version discussed in issue #11
result2.xlsx. Conclusion: OpenBLAS and MKL have potential for improvement for very slim rectangular matrices. The algorithms ´discussed in #11 work better for that case. Remaining question: Is this numpy-specific or does it happen in C as well?

[uellue]

Results by @sk1p show the same effect when using OpenBLAS in C. We will use BLAS for now and advertise a master thesis to integrate an improved algorithm in suitable BLAS implementations. That way it will be most useful and flexible.

[sk1p]

There was one major bug in our benchmarks that prevented them from being comparable: the C BLAS implementation used doubles, where the Python version used uint32. Now, with float64, the numpy version is almost as fast as the C BLAS version, maybe fast enough for our purposes? For a single mask, the worst case in the numpy version is still 3 times slower than the equivalent C BLAS case though. This may be fixable in numpy, or if we change the way we feed data to numpy, ...

Heatmap for numpy w/ float64:

Raw results: fastdot-numpy.xlsx

Heatmap for C BLAS:

Raw results:
fastdot-openblas.xlsx

Another problem is that our results don't currently include the largest frame size, as our buffer is too small to fit 4 frames, which is our minimum stack height. But that affects both versions so results should be comparable now.

[uellue]

In order to get a better understanding of the observed inefficiencies in hyperspy/hyperspy#1840, here's an implementation of the benchmark that uses dask for tiling and stacking. Long story short, it is exceptionally fast. The poor performance for low stackheight and maskcount comes from the known poor performance of the used BLAS library with very narrow rectangular matrices. Optimal values are reached when the input chunks fit into the L3 cache together.

dieter-dask-OpenBLAS-tiled.xlsx

[sk1p]

Here is the corresponding dask diagnostics output, for maskcount=16, framesize=512^2, tilesiez=256^2, stackheight=8:

The graph built:

Profiling output:

This is for one iter_dot call, with repeats=1. The green parts are the matrix multiplication, yellow and blue are generating the input matrixes (np.ones, transpose, ...)

[uellue]

One could use the graph as a textbook example for an optimized matrix multiplication!

[sk1p]

As we are gravitating towards using Apache Spark, I rewrote the same benchmark in Scala using DenseMatrix from spark.ml. For comparison I included an C OpenBLAS heatmap on the left. The left axis is "framesize-maskcount"

To summarize: the Scala DenseMatrix version performs as well as the C OpenBLAS implementation. This is still all local computation, it will get a little slower when distributing it via Spark.

[sk1p]

The code for the Scala/Spark benchmark is now available in the repository.

[sk1p]

The mathematical formulation from the original posting by @uellue should end up in our documentation, when that is done I think this issue can be closed.