An optimized K-Means based on DenseMatrix and GEMM

The main computations in K-Means are calculating distances between individual points and center points. Currently K-Means implementation is vector-based which can't take advantage of optimized native BLAS libraries.

When the original points are represented as dense vectors, our approach is to modify the original input data structures to a DenseMatrix-based one by grouping several points together. The original distance calculations can be translated into a Matrix multiplication then optimized native GEMM routines (Intel MKL, OpenBLAS etc.) can be used. This approach can also work with sparse vectors despite having larger memory consumption when translating sparse vectors to dense matrix.

Our preliminary benchmark shows this DenseMatrix+GEMM approach can boost the training performance by 3.5x with Intel MKL, looks very promising!

To minimize end user impact, proposed changes are to use config parameters to control if turn on this implementation without modifying public interfaces. Parameter rowsPerMatrix is used to control how many points are grouped together to build a DenseMatrix. An example:

$ spark-submit --master $SPARK_MASTER \

    --conf "spark.ml.kmeans.matrixImplementation.enabled=true" \

    --conf "spark.ml.kmeans.matrixImplementation.rowsPerMatrix=5000" \

    --class org.apache.spark.examples.ml.KMeansExample 

Several code changes are made in "spark.ml" namespace as we think "spark.mllib" is in maintenance mode, some are duplications from spark.mllib for using private definitions in the same package: 

• Modified: KMeans.scala, DatasetUtils.scala

• Added: KMeansMatrixImpl.scala

• Duplications: DistanceMeasure.scala, LocalKMeans.scala

If this general idea is accepted by community, we are willing to contribute our code to upstream and polish the implementation according to feedbacks and produce benchmarks.