Libertem: Efficient Matrix-Matrix-Multiplication in Apache Spark

[sk1p]

To implement our basic mask-application-and-summation operation from #14 in Apache Spark, we require a solution that has the following properties:

1) As little communication as possible. There should be no global shuffle step that moves the whole set of frames or something like that
1) The solution needs to exploit the inherent locality of the problem (see also #11) - it needs to be efficient even for frame/mask sizes that don't fit into the L3 cache
1) If possible, prevent any memory copies. For example, read a whole stack of frames into a contiguous block of memory, operate on it and then throw it away, maybe keeping the block of memory allocated for the next stack of frames.

Also, we need to efficiently handle the masks. Ideally, we need to prevent shipping the masks to all nodes; for a larger number of masks it is very likely more efficient to generate them on all nodes from their parameters.

Possible solutions from spark.mllib.linalg.distributed:

• BlockMatrix.multiply(...) - was quite slow in my tests, but I didn't test all parameters yet; does not allow multiplication against a local matrix, the local matrix needs to be converted to a BlockMatrix first which then can be multiplied against. Construction of a BlockMatrix is not free, so this is additional overhead.
• IndexedRowMatrix.multiply(...) - uses RowMatrix.multiply(...) under the hood
• RowMatrix.multiply(...) - splits the matrix into row vectors and calculates each entry of the result matrix by doing a dot product. Not as slow as BlockMatrix, but still about 10× slower than local BLAS gemm operations. As it is using the dot product, it is not cache efficient at all. I still need to find out how much of the inefficiency comes from the dot product vs. from other overheads. Overheads include copying the mask data for broadcasting (DenseMatrix.toDenseVector, which is called by RowMatrix.multiply creates a copy of the underlying data).

Other possibilities:

• implement or use flyby matrix multiplication. It may be too complicated for our case, because we know that the masks matrix is always quite small (i.e. fits into memory) and does not need to be partitioned in any way
• implement our own specialized matrix multiplication that is optimized for our case.

[sk1p]

Here is the heatmap for a RowMatrix multiplication benchmark. This was run on a local Spark instance, without any network communication overhead involved, using 4 cores. There is a more extreme dependence on maskcount compared to the local BLAS version.

For comparison, here is the local result again using a native local BLAS implementation with Scala:

[sk1p]

Looking at a flamegraph of the RowMatrix benchmark acquired using async-profiler:

We can see four 'columns' that represent different parts of the program execution (from left to right):

1) The main program, creating the input data matrices, broadcasting the masks matrix etc. - includes a bit of serialization for broadcasting.
1) Numeric computation using native BLAS functions
1) The largest part: serialization/deserialization! At the top are the ObjectInputStream.bytesToDouble / ObjectOutputStream.doublesToBytes (cropped from the top on this screenshot) functions
1) JVM GC. The impact of the garbage collector can possibly be reduced by implementing the configuration changes outlined in the Spark Tuning Guide. This is especially important when operating on large buffer sizes

I suspect that the serialization/deserialization happens because we generate the input frames on the Driver, so they get transferred as tasks to the executors. Because this is a local test, we don't pay the network overhead for the transfer, only serialization.

In a more realistic setting, the frames are read from HFDS and this should remove the serialization overhead. It remains to be seen how efficient reading from parquet is, as others have concluded that Reading from Parquet is CPU-intensive. More about that in #17

[sk1p]

I also tested another way to optimize the basic matrix multiplication operation: instead of using the RowMatrix and calculating it row by row with the dot product, a whole partition can be multiplied at once:

  def iter_dot(data: RDD[Array[Double]], masks: DenseMatrix, params: Params): Array[DenseMatrix] = {
    var result : Array[DenseMatrix] = null
    for(i <- 1 to params.repeats) {
      result = data.glom().map(partition => {
        val dataMatrix : DenseMatrix = new DenseMatrix(
          numCols = params.framesize,
          numRows = partition.length,
          values = partition.flatten
        )
        dataMatrix.multiply(masks)
      }).collect()
    }
    result
  }

partition is of type Array[Array[Double]] here, so for constructing the matrix it needs to be flattened to contiguous memory, that is Array[Double]. This operation is very expensive but can be circumvented by having a whole stack of frames as contiguous memory in the first place, for example.

[uellue]

@sk1p That means we'd store something like a HDF5 chunk as a blob. That should work well if we store the frames in their natural order.

If we want to chunk data into square tiles on the scan level and perhaps on the frame level, too, in order to have nice L3 "bites" for our processing that contain the local environment, it gets complicated in two ways:

1. We need to keep track which portion of the data belongs where. The contents of the file will be pretty mangled. Like HDF5 we'd need an abstraction layer that keeps an overview for the naive programmer that just wants frames at position X
2. If we acquire directly into a Parquet file we'll have to keep a lot of buffers open since we can only write entire finished blocks and can't change them after, if I am correct. We can only append blocks. AFAIK HDF5 does that more efficiently.

In general I have a feeling that Parquet + Spark in its current form can only be an intermediate solution for numerical data like ours. In the future we need one of these two solutions:

1. Proper HDF5 reading and writing support for Hadoop FS and Spark. Probably that'll rather be "HDF6" for writing.
2. A proper intermediate layer to turn Parquet data into matrices efficiently, with nice selectors to specify chunking and such. I.e. "Turn my data set into a block matrix based on this SQL query, where this vector goes here, this one there, and so forth". Doing that without a shuffle could be "fun".

[sk1p]

I agree that things get messy when we need a local environment on the scan level.

Just to see what kind of performance we _can_ get with natural frame order, this is a flamegraph of reading and processing raw binary files using SparkContext.binaryFiles:

This was done using our EMPAD dataset, cropped to framesize=128x128, converted to doubles and chopped into 32 partitions, each 256MB. The partitions are not read each as a whole but in stacks of 32 frames, which fit comfortably into the L3 cache.

This benchmark reads and processes 8GB of frames in about 3 seconds, so about 2.666GB/s. A large part (~40%) is still deserialization, but it's not as bad as using some of the other file formats. 36% are spent doing 'useful' work, that is, reading the file and doing the matrix multiplication. There is close to no GC overhead, and only some compilation overhead, which should be amortized over the processing time of a large file.

Here is a flamegraph of reading and processing parquet files:

This is again our EMPAD dataset converted to doubles and saved as parquet files, here 16 partitions, 512MB each. Multiplication is done using the less-than-optimal RowMatrix method, which maybe can be improved upon by .glom()ing into a DenseMatrix.

Reading parquet is about 10 times slower than raw binary files and takes about 30 seconds to process the 8GB dataset. Only about 5% are spent reading from the file system and doing matrix multiplication. Most of the time, about 55%, is spent somewhere in parquet converting data / decoding etc. The Magenta parts of the graph are GC overhead, and take about 26% of samples.

Another test, this time with Avro:

This takes about 20 seconds to process the 8GB dataset. The main difference to parquet appears to be significantly better memory handling, i.e. buffer re-use: there is close to no GC overhead. There is also a some prior work for speeding up Avro, see 1, 2.

Bear in mind that all of these don't include the overheads that will be introduced with using HDFS.

[sk1p]

On the other hand, the encoding done by parquet might make sense for different datasets, in particular with integer as pixel type, and for more sparse datasets. I think having a dense, noisy, float/double dataset could be the worst case for parquet.

[sk1p]

As performing multiple copies as part of deserialization is a fundamental problem of JVM based solutions, we have decided not to further pursue using apache spark for our project. Apache spark and the JVM are obviously not as optimized as we hoped in the sense of raw I/O throughput.