From 48e4af130db5f9511aa26779c49999f1c59d8ed5 Mon Sep 17 00:00:00 2001
From: Ludovic Henry <luhenry@microsoft.com>
Date: Wed, 16 Dec 2020 08:48:22 -0800
Subject: [PATCH 01/16] [ML] Add a vectorized BLAS implementation

Whenever a native BLAS implementation isn't available on the system, Spark
automatically falls back onto a Java implementation. With the recent release
of the Vector API in the OpenJDK [1], we can use hardware acceleration for such
operations.

This patch introduces a VectorizedBLAS class which implements such
hardware-accelerated BLAS operations. This feature is hidden behind the
"vectorized" profile that you can enable by passing "-Pvectorized" to sbt or
maven.

The Vector API has been introduced in JDK 16. Following discussion on the
mailing list, this API is introduced transparently and needs to be enabled
explicitely.

[1] https://openjdk.java.net/jeps/338
---
 mllib-local/pom.xml                           |  40 ++++
 .../org/apache/spark/ml/linalg/BLAS.scala     |  46 ++--
 .../spark/ml/linalg/BLASBenchmark.scala       | 226 ++++++++++++++++++
 .../apache/spark/ml/linalg/BLASSuite.scala    |   2 +-
 .../spark/ml/linalg/VectorizedBLAS.java       | 162 +++++++++++++
 .../apache/spark/ml/recommendation/ALS.scala  |   2 +-
 project/SparkBuild.scala                      |  10 +
 7 files changed, 468 insertions(+), 20 deletions(-)
 create mode 100644 mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
 create mode 100644 mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java

diff --git a/mllib-local/pom.xml b/mllib-local/pom.xml
index 2a2c373242201..977febdc7ccb8 100644
--- a/mllib-local/pom.xml
+++ b/mllib-local/pom.xml
@@ -56,6 +56,11 @@
       <groupId>org.apache.spark</groupId>
       <artifactId>spark-tags_${scala.binary.version}</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-core_${scala.binary.version}</artifactId>
+      <version>${project.version}</version>
+    </dependency>
 
     <!--
       This spark-tags test-dep is needed even though it isn't used in this module, otherwise testing-cmds that exclude
@@ -68,6 +73,13 @@
       <scope>test</scope>
     </dependency>
 
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-core_${scala.binary.version}</artifactId>
+      <version>${project.version}</version>
+      <type>test-jar</type>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
   <profiles>
     <profile>
@@ -81,6 +93,34 @@
         </dependency>
       </dependencies>
     </profile>
+    <profile>
+      <id>vectorized</id>
+      <properties>
+        <extra.source.dir>src/vectorized/java</extra.source.dir>
+      </properties>
+      <build>
+        <plugins>
+          <plugin>
+            <groupId>org.codehaus.mojo</groupId>
+            <artifactId>build-helper-maven-plugin</artifactId>
+            <executions>
+              <execution>
+                <id>add-vectorized-sources</id>
+                <phase>generate-sources</phase>
+                <goals>
+                  <goal>add-source</goal>
+                </goals>
+                <configuration>
+                  <sources>
+                    <source>${extra.source.dir}</source>
+                  </sources>
+                </configuration>
+              </execution>
+            </executions>
+          </plugin>
+        </plugins>
+      </build>
+    </profile>
   </profiles>
   <build>
     <outputDirectory>target/scala-${scala.binary.version}/classes</outputDirectory>
diff --git a/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala b/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
index b6c1b011f004c..af3ce3b9587ac 100644
--- a/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
+++ b/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
@@ -18,28 +18,46 @@
 package org.apache.spark.ml.linalg
 
 import com.github.fommil.netlib.{BLAS => NetlibBLAS, F2jBLAS}
-import com.github.fommil.netlib.BLAS.{getInstance => NativeBLAS}
+import scala.util.Try
+
+import org.apache.spark.util.Utils
 
 /**
  * BLAS routines for MLlib's vectors and matrices.
  */
 private[spark] object BLAS extends Serializable {
 
-  @transient private var _f2jBLAS: NetlibBLAS = _
+  @transient private var _javaBLAS: NetlibBLAS = _
   @transient private var _nativeBLAS: NetlibBLAS = _
   private val nativeL1Threshold: Int = 256
 
-  // For level-1 function dspmv, use f2jBLAS for better performance.
-  private[ml] def f2jBLAS: NetlibBLAS = {
-    if (_f2jBLAS == null) {
-      _f2jBLAS = new F2jBLAS
+  // For level-1 function dspmv, use javaBLAS for better performance.
+  private[ml] def javaBLAS: NetlibBLAS = {
+    if (_javaBLAS == null) {
+      _javaBLAS = Try(
+        Utils.classForName("org.apache.spark.ml.linalg.VectorizedBLAS")
+             .newInstance()
+             .asInstanceOf[NetlibBLAS]).getOrElse(new F2jBLAS)
+    }
+    _javaBLAS
+  }
+
+  // For level-3 routines, we use the native BLAS.
+  private[ml] def nativeBLAS: NetlibBLAS = {
+    if (_nativeBLAS == null) {
+      _nativeBLAS =
+        if (NetlibBLAS.getInstance.isInstanceOf[F2jBLAS]) {
+          javaBLAS
+        } else {
+          NetlibBLAS.getInstance
+        }
     }
-    _f2jBLAS
+    _nativeBLAS
   }
 
   private[ml] def getBLAS(vectorSize: Int): NetlibBLAS = {
     if (vectorSize < nativeL1Threshold) {
-      f2jBLAS
+      javaBLAS
     } else {
       nativeBLAS
     }
@@ -235,14 +253,6 @@ private[spark] object BLAS extends Serializable {
     }
   }
 
-  // For level-3 routines, we use the native BLAS.
-  private[ml] def nativeBLAS: NetlibBLAS = {
-    if (_nativeBLAS == null) {
-      _nativeBLAS = NativeBLAS
-    }
-    _nativeBLAS
-  }
-
   /**
    * Adds alpha * x * x.t to a matrix in-place. This is the same as BLAS's ?SPR.
    *
@@ -267,7 +277,7 @@ private[spark] object BLAS extends Serializable {
       x: DenseVector,
       beta: Double,
       y: DenseVector): Unit = {
-    f2jBLAS.dspmv("U", n, alpha, A.values, x.values, 1, beta, y.values, 1)
+    javaBLAS.dspmv("U", n, alpha, A.values, x.values, 1, beta, y.values, 1)
   }
 
   /**
@@ -279,7 +289,7 @@ private[spark] object BLAS extends Serializable {
     val n = v.size
     v match {
       case DenseVector(values) =>
-        NativeBLAS.dspr("U", n, alpha, values, 1, U)
+        nativeBLAS.dspr("U", n, alpha, values, 1, U)
       case SparseVector(size, indices, values) =>
         val nnz = indices.length
         var colStartIdx = 0
diff --git a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
new file mode 100644
index 0000000000000..48ad5a71aae31
--- /dev/null
+++ b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
@@ -0,0 +1,226 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.ml.linalg
+
+import com.github.fommil.netlib.{BLAS => NetlibBLAS, F2jBLAS}
+import scala.util.Try
+
+import org.apache.spark.benchmark.{Benchmark, BenchmarkBase}
+import org.apache.spark.util.Utils
+
+/**
+ * Serialization benchmark for BLAS.
+ * To run this benchmark:
+ * {{{
+ * 1. without sbt: bin/spark-submit --class <this class> <spark mllib test jar>
+ * 2. build/sbt "mllib/test:runMain <this class>"
+ * 3. generate result: SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt "mllib/test:runMain <this class>"
+ *    Results will be written to "benchmarks/UDTSerializationBenchmark-results.txt".
+ * }}}
+ */
+object BLASBenchmark extends BenchmarkBase {
+
+  override def runBenchmarkSuite(mainArgs: Array[String]): Unit = {
+
+    val iters = 1e2.toInt
+    val rnd = new scala.util.Random(0)
+
+    val f2jBLAS = new F2jBLAS
+    val nativeBLAS = NetlibBLAS.getInstance
+    val vectorBLAS = Try(
+        Utils.classForName("org.apache.spark.ml.linalg.VectorizedBLAS")
+             .newInstance()
+             .asInstanceOf[NetlibBLAS]).getOrElse(new F2jBLAS)
+
+    // scalastyle:off println
+    println("nativeBLAS = " + nativeBLAS.getClass.getName)
+    println("f2jBLAS    = " + f2jBLAS.getClass.getName)
+    println("vectorBLAS = " + vectorBLAS.getClass.getName)
+    // scalastyle:on println
+
+    runBenchmark("daxpy") {
+      val n = 1e7.toInt
+      val a = rnd.nextDouble
+      val x = Array.fill(n) { rnd.nextDouble }
+      val y = Array.fill(n) { rnd.nextDouble }
+
+      val benchmark = new Benchmark("daxpy", n, iters, output = output)
+
+      benchmark.addCase("f2j") { _ =>
+        f2jBLAS.daxpy(n, a, x, 1, y, 1)
+      }
+
+      if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("native") { _ =>
+          nativeBLAS.daxpy(n, a, x, 1, y, 1)
+        }
+      }
+
+      if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("vector") { _ =>
+          vectorBLAS.daxpy(n, a, x, 1, y, 1)
+        }
+      }
+
+      benchmark.run()
+    }
+
+    runBenchmark("sdot") {
+      val n = 1e7.toInt
+      val x = Array.fill(n) { rnd.nextFloat }
+      val y = Array.fill(n) { rnd.nextFloat }
+
+      val benchmark = new Benchmark("sdot", n, iters, output = output)
+
+      benchmark.addCase("f2j") { _ =>
+        f2jBLAS.sdot(n, x, 1, y, 1)
+      }
+
+      if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("native") { _ =>
+          nativeBLAS.sdot(n, x, 1, y, 1)
+        }
+      }
+
+      if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("vector") { _ =>
+          vectorBLAS.sdot(n, x, 1, y, 1)
+        }
+      }
+
+      benchmark.run()
+    }
+
+    runBenchmark("ddot") {
+      val n = 1e7.toInt
+      val x = Array.fill(n) { rnd.nextDouble }
+      val y = Array.fill(n) { rnd.nextDouble }
+
+      val benchmark = new Benchmark("ddot", n, iters, output = output)
+
+      benchmark.addCase("f2j") { _ =>
+        f2jBLAS.ddot(n, x, 1, y, 1)
+      }
+
+      if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("native") { _ =>
+          nativeBLAS.ddot(n, x, 1, y, 1)
+        }
+      }
+
+      if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("vector") { _ =>
+          vectorBLAS.ddot(n, x, 1, y, 1)
+        }
+      }
+
+      benchmark.run()
+    }
+
+    runBenchmark("dscal") {
+      val n = 1e7.toInt
+      val a = rnd.nextDouble
+      val x = Array.fill(n) { rnd.nextDouble }
+
+      val benchmark = new Benchmark("dscal", n, iters, output = output)
+
+      benchmark.addCase("f2j") { _ =>
+        f2jBLAS.dscal(n, a, x, 1)
+      }
+
+      if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("native") { _ =>
+          nativeBLAS.dscal(n, a, x, 1)
+        }
+      }
+
+      if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("vector") { _ =>
+          vectorBLAS.dscal(n, a, x, 1)
+        }
+      }
+
+      benchmark.run()
+    }
+
+    runBenchmark("dgemv[T]") {
+      val m = 1e4.toInt
+      val n = 1e3.toInt
+      val alpha = rnd.nextDouble
+      val a = Array.fill(n * m) { rnd.nextDouble }
+      val lda = m
+      val x = Array.fill(m) { rnd.nextDouble }
+      val beta = rnd.nextDouble
+      val y = Array.fill(n) { rnd.nextDouble }
+
+      val benchmark = new Benchmark("dgemv[T]", n, iters, output = output)
+
+      benchmark.addCase("f2j") { _ =>
+        f2jBLAS.dgemv("T", m, n, alpha, a, lda, x, 1, beta, y, 1)
+      }
+
+      if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("native") { _ =>
+          nativeBLAS.dgemv("T", m, n, alpha, a, lda, x, 1, beta, y, 1)
+        }
+      }
+
+      if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("vector") { _ =>
+          vectorBLAS.dgemv("T", m, n, alpha, a, lda, x, 1, beta, y, 1)
+        }
+      }
+
+      benchmark.run()
+    }
+
+    runBenchmark("dgemm[T,N]") {
+      val m = 1e3.toInt
+      val n = 1e2.toInt
+      val k = 1e3.toInt
+      val alpha = rnd.nextDouble
+      val a = Array.fill(k * m) { rnd.nextDouble }
+      val lda = k
+      val b = Array.fill(n * k) { rnd.nextDouble }
+      val ldb = k
+      val beta = rnd.nextDouble
+      val c = Array.fill(m * n) { rnd.nextDouble }
+      var ldc = m
+
+      val benchmark = new Benchmark("dgemm[T,N]", m*n, iters, output = output)
+
+      benchmark.addCase("f2j") { _ =>
+        f2jBLAS.dgemm("T", "N", m, n, k, alpha, a, lda, b, ldb, beta, c, ldc)
+      }
+
+      if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("native") { _ =>
+          nativeBLAS.dgemm("T", "N", m, n, k, alpha, a, lda, b, ldb, beta, c, ldc)
+        }
+      }
+
+      if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("vector") { _ =>
+          vectorBLAS.dgemm("T", "N", m, n, k, alpha, a, lda, b, ldb, beta, c, ldc)
+        }
+      }
+
+      benchmark.run()
+    }
+  }
+}
diff --git a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASSuite.scala b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASSuite.scala
index 781f3da313d82..ce177666f34e6 100644
--- a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASSuite.scala
+++ b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASSuite.scala
@@ -24,7 +24,7 @@ import org.apache.spark.ml.util.TestingUtils._
 class BLASSuite extends SparkMLFunSuite {
 
   test("nativeL1Threshold") {
-    assert(getBLAS(128) == BLAS.f2jBLAS)
+    assert(getBLAS(128) == BLAS.javaBLAS)
     assert(getBLAS(256) == BLAS.nativeBLAS)
     assert(getBLAS(512) == BLAS.nativeBLAS)
   }
diff --git a/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java b/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java
new file mode 100644
index 0000000000000..9e77c647f52ec
--- /dev/null
+++ b/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java
@@ -0,0 +1,162 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.ml.linalg;
+
+import com.github.fommil.netlib.F2jBLAS;
+import jdk.incubator.vector.DoubleVector;
+import jdk.incubator.vector.FloatVector;
+import jdk.incubator.vector.VectorOperators;
+import jdk.incubator.vector.VectorSpecies;
+
+public class VectorizedBLAS extends F2jBLAS {
+
+  private final static VectorSpecies<Float>  FMAX = FloatVector.SPECIES_MAX;
+  private final static VectorSpecies<Double> DMAX = DoubleVector.SPECIES_MAX;
+
+  // y += alpha * x
+  @Override
+  public void daxpy(int n, double alpha, double[] x, int incx, double[] y, int incy) {
+    if (incx == 1 && incy == 1 && n <= x.length && n <= y.length) {
+      if (alpha != 0.) {
+        int i = 0;
+        for (; i < DMAX.loopBound(n); i += DMAX.length()) {
+          DoubleVector vx = DoubleVector.fromArray(DMAX, x, i);
+          DoubleVector vy = DoubleVector.fromArray(DMAX, y, i);
+          vx.lanewise(VectorOperators.MUL, alpha).add(vy)
+            .intoArray(y, i);
+        }
+        for (; i < n; i += 1) {
+          y[i] += alpha * x[i];
+        }
+      }
+    } else {
+      super.daxpy(n, alpha, x, incx, y, incy);
+    }
+  }
+
+  // sum(x * y)
+  @Override
+  public float sdot(int n, float[] x, int incx, float[] y, int incy) {
+    if (incx == 1 && incy == 1) {
+      float sum = 0.0f;
+      int i = 0;
+      for (; i < FMAX.loopBound(n); i += FMAX.length()) {
+        FloatVector vx = FloatVector.fromArray(FMAX, x, i);
+        FloatVector vy = FloatVector.fromArray(FMAX, y, i);
+        sum += vx.mul(vy).reduceLanes(VectorOperators.ADD);
+      }
+      for (; i < n; i += 1) {
+        sum += x[i] * y[i];
+      }
+      return sum;
+    } else {
+      return super.sdot(n, x, incx, y, incy);
+    }
+  }
+
+  // sum(x * y)
+  @Override
+  public double ddot(int n, double[] x, int incx, double[] y, int incy) {
+    if (incx == 1 && incy == 1) {
+      double sum = 0.;
+      int i = 0;
+      for (; i < DMAX.loopBound(n); i += DMAX.length()) {
+        DoubleVector vx = DoubleVector.fromArray(DMAX, x, i);
+        DoubleVector vy = DoubleVector.fromArray(DMAX, y, i);
+        sum += vx.mul(vy).reduceLanes(VectorOperators.ADD);
+      }
+      for (; i < n; i += 1) {
+        sum += x[i] * y[i];
+      }
+      return sum;
+    } else {
+      return super.ddot(n, x, incx, y, incy);
+    }
+  }
+
+  // x = alpha * x
+  @Override
+  public void dscal(int n, double alpha, double[] x, int incx) {
+    if (incx == 1) {
+      if (alpha != 1.) {
+        int i = 0;
+        for (; i < DMAX.loopBound(n); i += DMAX.length()) {
+          DoubleVector vx = DoubleVector.fromArray(DMAX, x, i);
+          vx.lanewise(VectorOperators.MUL, alpha)
+            .intoArray(x, i);
+        }
+        for (; i < n; i += 1) {
+          x[i] *= alpha;
+        }
+      }
+    } else {
+      super.dscal(n, alpha, x, incx);
+    }
+  }
+
+  // y = alpha * A * x + beta * y
+  @Override
+  public void dgemv(String trans, int m, int n, double alpha, double[] a, int lda, double[] x, int incx, double beta, double[] y, int incy) {
+    if (trans == "T" && incx == 1 && incy == 1 && lda == m) {
+      // y = beta * y
+      dscal(n, beta, y, 1);
+      // y += alpha * A * x
+      if (alpha != 0.) {
+        for (int col = 0; col < n; col += 1) {
+          int row = 0;
+          for (; row < DMAX.loopBound(m); row += DMAX.length()) {
+            DoubleVector vx = DoubleVector.fromArray(DMAX, x, row);
+            DoubleVector va = DoubleVector.fromArray(DMAX, a, row + col * m);
+            y[col] += alpha * vx.mul(va).reduceLanes(VectorOperators.ADD);
+          }
+          for (; row < m; row += 1) {
+            y[col] += alpha * x[row] * a[row + col * m];
+          }
+        }
+      }
+    } else {
+      super.dgemv(trans, m, n, alpha, a, lda, x, incx, beta, y, incy);
+    }
+  }
+
+  @Override
+  public void dgemm(String transa, String transb, int m, int n, int k, double alpha, double[] a, int lda, double[] b, int ldb, double beta, double[] c, int ldc) {
+    if (transa == "T" && transb == "N" && lda == k && ldb == k && ldc == m) {
+      // C = beta * C
+      dscal(m * n, beta, c, 1);
+      // C += alpha * A * B
+      if (alpha != 0.) {
+        for (int col = 0; col < n; col += 1) {
+          for (int row = 0; row < m; row += 1) {
+            int i = 0;
+            for (; i < DMAX.loopBound(k); i += DMAX.length()) {
+              DoubleVector va = DoubleVector.fromArray(DMAX, a, i + row * k);
+              DoubleVector vb = DoubleVector.fromArray(DMAX, b, i + col * k);
+              c[row + col * m] += alpha * va.mul(vb).reduceLanes(VectorOperators.ADD);
+            }
+            for (; i < k; i += 1) {
+              c[row + col * m] += alpha * a[i + row * k] * b[i + col * k];
+            }
+          }
+        }
+      }
+    } else {
+      super.dgemm(transa, transb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc);
+    }
+  }
+}
diff --git a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
index 088f6a682be82..217bbea0e82a1 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
@@ -470,7 +470,7 @@ class ALSModel private[ml] (
         srcIter.foreach { case (srcId, srcFactor) =>
           dstIter.foreach { case (dstId, dstFactor) =>
             // We use F2jBLAS which is faster than a call to native BLAS for vector dot product
-            val score = BLAS.f2jBLAS.sdot(rank, srcFactor, 1, dstFactor, 1)
+            val score = BLAS.javaBLAS.sdot(rank, srcFactor, 1, dstFactor, 1)
             pq += dstId -> score
           }
           pq.foreach { case (dstId, score) =>
diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala
index 23fb73d228e01..52e1426ecdd83 100644
--- a/project/SparkBuild.scala
+++ b/project/SparkBuild.scala
@@ -415,6 +415,11 @@ object SparkBuild extends PomBuild {
     outputStrategy in run := Some (StdoutOutput),
 
     javaOptions += "-Xmx2g",
+    javaOptions += {
+      val versionParts = System.getProperty("java.version").split("[+.\\-]+", 3)
+      var major = versionParts(0).toInt
+      if (major >= 16) "--add-modules=jdk.incubator.vector" else ""
+    },
 
     sparkShell := {
       (runMain in Compile).toTask(" org.apache.spark.repl.Main -usejavacp").value
@@ -1027,6 +1032,11 @@ object TestSettings {
     javaOptions in Test ++= System.getProperties.asScala.filter(_._1.startsWith("spark"))
       .map { case (k,v) => s"-D$k=$v" }.toSeq,
     javaOptions in Test += "-ea",
+    javaOptions in Test += {
+      val versionParts = System.getProperty("java.version").split("[+.\\-]+", 3)
+      var major = versionParts(0).toInt
+      if (major >= 16) "--add-modules=jdk.incubator.vector" else ""
+    },
     // SPARK-29282 This is for consistency between JDK8 and JDK11.
     javaOptions in Test ++= "-Xmx4g -Xss4m -XX:+UseParallelGC -XX:-UseDynamicNumberOfGCThreads"
       .split(" ").toSeq,

From c949d444eb2c6f8980c94baa8ebf9f5dfed1a90e Mon Sep 17 00:00:00 2001
From: Ludovic Henry <luhenry@microsoft.com>
Date: Fri, 18 Dec 2020 01:57:15 -0800
Subject: [PATCH 02/16] Improve performance

All implmentations beat f2j on microbenchmarks on x86 (w/ AVX-2).

See https://github.com/luhenry/vectorizedblas/releases/tag/v0.1.5 for more details
---
 .../spark/ml/linalg/VectorizedBLAS.java       | 100 ++++++++++++++++--
 1 file changed, 90 insertions(+), 10 deletions(-)

diff --git a/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java b/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java
index 9e77c647f52ec..2af168d21d921 100644
--- a/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java
+++ b/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java
@@ -25,8 +25,8 @@
 
 public class VectorizedBLAS extends F2jBLAS {
 
-  private final static VectorSpecies<Float>  FMAX = FloatVector.SPECIES_MAX;
-  private final static VectorSpecies<Double> DMAX = DoubleVector.SPECIES_MAX;
+  private static final VectorSpecies<Float>  FMAX = FloatVector.SPECIES_MAX;
+  private static final VectorSpecies<Double> DMAX = DoubleVector.SPECIES_MAX;
 
   // y += alpha * x
   @Override
@@ -34,11 +34,11 @@ public void daxpy(int n, double alpha, double[] x, int incx, double[] y, int inc
     if (incx == 1 && incy == 1 && n <= x.length && n <= y.length) {
       if (alpha != 0.) {
         int i = 0;
+        DoubleVector valpha = DoubleVector.broadcast(DMAX, alpha);
         for (; i < DMAX.loopBound(n); i += DMAX.length()) {
           DoubleVector vx = DoubleVector.fromArray(DMAX, x, i);
           DoubleVector vy = DoubleVector.fromArray(DMAX, y, i);
-          vx.lanewise(VectorOperators.MUL, alpha).add(vy)
-            .intoArray(y, i);
+          vx.fma(valpha, vy).intoArray(y, i);
         }
         for (; i < n; i += 1) {
           y[i] += alpha * x[i];
@@ -55,11 +55,13 @@ public float sdot(int n, float[] x, int incx, float[] y, int incy) {
     if (incx == 1 && incy == 1) {
       float sum = 0.0f;
       int i = 0;
+      FloatVector vsum = FloatVector.zero(FMAX);
       for (; i < FMAX.loopBound(n); i += FMAX.length()) {
         FloatVector vx = FloatVector.fromArray(FMAX, x, i);
         FloatVector vy = FloatVector.fromArray(FMAX, y, i);
-        sum += vx.mul(vy).reduceLanes(VectorOperators.ADD);
+        vsum = vx.fma(vy, vsum);
       }
+      sum += vsum.reduceLanes(VectorOperators.ADD);
       for (; i < n; i += 1) {
         sum += x[i] * y[i];
       }
@@ -75,11 +77,13 @@ public double ddot(int n, double[] x, int incx, double[] y, int incy) {
     if (incx == 1 && incy == 1) {
       double sum = 0.;
       int i = 0;
+      DoubleVector vsum = DoubleVector.zero(DMAX);
       for (; i < DMAX.loopBound(n); i += DMAX.length()) {
         DoubleVector vx = DoubleVector.fromArray(DMAX, x, i);
         DoubleVector vy = DoubleVector.fromArray(DMAX, y, i);
-        sum += vx.mul(vy).reduceLanes(VectorOperators.ADD);
+        vsum = vx.fma(vy, vsum);
       }
+      sum += vsum.reduceLanes(VectorOperators.ADD);
       for (; i < n; i += 1) {
         sum += x[i] * y[i];
       }
@@ -95,10 +99,10 @@ public void dscal(int n, double alpha, double[] x, int incx) {
     if (incx == 1) {
       if (alpha != 1.) {
         int i = 0;
+        DoubleVector valpha = DoubleVector.broadcast(DMAX, alpha);
         for (; i < DMAX.loopBound(n); i += DMAX.length()) {
           DoubleVector vx = DoubleVector.fromArray(DMAX, x, i);
-          vx.lanewise(VectorOperators.MUL, alpha)
-            .intoArray(x, i);
+          vx.mul(valpha).intoArray(x, i);
         }
         for (; i < n; i += 1) {
           x[i] *= alpha;
@@ -109,10 +113,86 @@ public void dscal(int n, double alpha, double[] x, int incx) {
     }
   }
 
+  // y := alpha * a * x + beta * y
+  @Override
+  public void dspmv(String uplo, int n, double alpha, double[] a, double[] x, int incx, double beta, double[] y, int incy) {
+    if (uplo.equals("U") && incx == 1 && incy == 1) {
+      // y = beta * y
+      dscal(n, beta, y, 1);
+      // y += alpha * A * x
+      if (alpha != 0.) {
+        for (int row = 0; row < n; row += 1) {
+          int col = 0;
+          DoubleVector valphaxrow = DoubleVector.broadcast(DMAX, alpha * x[row]);
+          for (; col < DMAX.loopBound(row); col += DMAX.length()) {
+            DoubleVector vx = DoubleVector.fromArray(DMAX, x, col);
+            DoubleVector vy = DoubleVector.fromArray(DMAX, y, col);
+            DoubleVector va = DoubleVector.fromArray(DMAX, a, col + row * (row + 1) / 2);
+            y[row] += alpha * vx.mul(va).reduceLanes(VectorOperators.ADD);
+            valphaxrow.fma(va, vy).intoArray(y, col);
+          }
+          for (; col < row; col += 1) {
+            y[row] += alpha * x[col] * a[col + row * (row + 1) / 2];
+            y[col] += alpha * x[row] * a[col + row * (row + 1) / 2];
+          }
+          y[row] += alpha * x[col] * a[col + row * (row + 1) / 2];
+        }
+      }
+    } else {
+      super.dspmv(uplo, n, alpha, a, x, incx, beta, y, incy);
+    }
+  }
+
+  // a += alpha * x * x.t
+  @Override
+  public void dspr(String uplo, int n, double alpha, double[] x, int incx, double[] a) {
+    if (uplo.equals("U") && incx == 1) {
+      if (alpha != 0.) {
+        for (int row = 0; row < n; row += 1) {
+          int col = 0;
+          DoubleVector valphax = DoubleVector.broadcast(DMAX, alpha * x[row]);
+          for (; col < DMAX.loopBound(row + 1); col += DMAX.length()) {
+            DoubleVector vx = DoubleVector.fromArray(DMAX, x, col);
+            DoubleVector va = DoubleVector.fromArray(DMAX, a, col + row * (row + 1) / 2);
+            vx.fma(valphax, va).intoArray(a, col + row * (row + 1) / 2);
+          }
+          for (; col < row + 1; col += 1) {
+            a[col + row * (row + 1) / 2] += alpha * x[row] * x[col];
+          }
+        }
+      }
+    } else {
+      super.dspr(uplo, n, alpha, x, incx, a);
+    }
+  }
+
+  // a += alpha * x * x.t
+  @Override
+  public void dsyr(String uplo, int n, double alpha, double[] x, int incx, double[] a, int lda) {
+    if (uplo.equals("U") && incx == 1) {
+      if (alpha != 0.) {
+        for (int row = 0; row < n; row += 1) {
+          int col = 0;
+          DoubleVector valphax = DoubleVector.broadcast(DMAX, alpha * x[row]);
+          for (; col < DMAX.loopBound(row + 1); col += DMAX.length()) {
+            DoubleVector vx = DoubleVector.fromArray(DMAX, x, col);
+            DoubleVector va = DoubleVector.fromArray(DMAX, a, col + row * n);
+            vx.fma(valphax, va).intoArray(a, col + row * n);
+          }
+          for (; col < row + 1; col += 1) {
+            a[col + row * n] += alpha * x[row] * x[col];
+          }
+        }
+      }
+    } else {
+      super.dsyr(uplo, n, alpha, x, incx, a, lda);
+    }
+  }
+
   // y = alpha * A * x + beta * y
   @Override
   public void dgemv(String trans, int m, int n, double alpha, double[] a, int lda, double[] x, int incx, double beta, double[] y, int incy) {
-    if (trans == "T" && incx == 1 && incy == 1 && lda == m) {
+    if (trans.equals("T") && incx == 1 && incy == 1 && lda == m) {
       // y = beta * y
       dscal(n, beta, y, 1);
       // y += alpha * A * x
@@ -136,7 +216,7 @@ public void dgemv(String trans, int m, int n, double alpha, double[] a, int lda,
 
   @Override
   public void dgemm(String transa, String transb, int m, int n, int k, double alpha, double[] a, int lda, double[] b, int ldb, double beta, double[] c, int ldc) {
-    if (transa == "T" && transb == "N" && lda == k && ldb == k && ldc == m) {
+    if (transa.equals("T") && transb.equals("N") && lda == k && ldb == k && ldc == m) {
       // C = beta * C
       dscal(m * n, beta, c, 1);
       // C += alpha * A * B

From bc9cfb066430155f3d3a8911b4e47c4a5b6c717d Mon Sep 17 00:00:00 2001
From: Ludovic Henry <luhenry@microsoft.com>
Date: Sun, 20 Dec 2020 13:52:28 -0800
Subject: [PATCH 03/16] Use dev.ludovic.vectorizedblas package directly

It simplifies the build process, and will allow for faster iterations on my end
---
 mllib-local/pom.xml                           |  32 +--
 .../org/apache/spark/ml/linalg/BLAS.scala     |   2 +-
 .../spark/ml/linalg/VectorizedBLAS.java       | 242 ------------------
 3 files changed, 8 insertions(+), 268 deletions(-)
 delete mode 100644 mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java

diff --git a/mllib-local/pom.xml b/mllib-local/pom.xml
index 977febdc7ccb8..481d985f744da 100644
--- a/mllib-local/pom.xml
+++ b/mllib-local/pom.xml
@@ -95,31 +95,13 @@
     </profile>
     <profile>
       <id>vectorized</id>
-      <properties>
-        <extra.source.dir>src/vectorized/java</extra.source.dir>
-      </properties>
-      <build>
-        <plugins>
-          <plugin>
-            <groupId>org.codehaus.mojo</groupId>
-            <artifactId>build-helper-maven-plugin</artifactId>
-            <executions>
-              <execution>
-                <id>add-vectorized-sources</id>
-                <phase>generate-sources</phase>
-                <goals>
-                  <goal>add-source</goal>
-                </goals>
-                <configuration>
-                  <sources>
-                    <source>${extra.source.dir}</source>
-                  </sources>
-                </configuration>
-              </execution>
-            </executions>
-          </plugin>
-        </plugins>
-      </build>
+      <dependencies>
+        <dependency>
+          <groupId>dev.ludovic</groupId>
+          <artifactId>vectorizedblas-blas</artifactId>
+          <version>0.1.7</version>
+        </dependency>
+      </dependencies>
     </profile>
   </profiles>
   <build>
diff --git a/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala b/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
index af3ce3b9587ac..8c2bbe2f6f757 100644
--- a/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
+++ b/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
@@ -35,7 +35,7 @@ private[spark] object BLAS extends Serializable {
   private[ml] def javaBLAS: NetlibBLAS = {
     if (_javaBLAS == null) {
       _javaBLAS = Try(
-        Utils.classForName("org.apache.spark.ml.linalg.VectorizedBLAS")
+        Utils.classForName("dev.ludovic.blas.VectorizedBLAS")
              .newInstance()
              .asInstanceOf[NetlibBLAS]).getOrElse(new F2jBLAS)
     }
diff --git a/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java b/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java
deleted file mode 100644
index 2af168d21d921..0000000000000
--- a/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java
+++ /dev/null
@@ -1,242 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.ml.linalg;
-
-import com.github.fommil.netlib.F2jBLAS;
-import jdk.incubator.vector.DoubleVector;
-import jdk.incubator.vector.FloatVector;
-import jdk.incubator.vector.VectorOperators;
-import jdk.incubator.vector.VectorSpecies;
-
-public class VectorizedBLAS extends F2jBLAS {
-
-  private static final VectorSpecies<Float>  FMAX = FloatVector.SPECIES_MAX;
-  private static final VectorSpecies<Double> DMAX = DoubleVector.SPECIES_MAX;
-
-  // y += alpha * x
-  @Override
-  public void daxpy(int n, double alpha, double[] x, int incx, double[] y, int incy) {
-    if (incx == 1 && incy == 1 && n <= x.length && n <= y.length) {
-      if (alpha != 0.) {
-        int i = 0;
-        DoubleVector valpha = DoubleVector.broadcast(DMAX, alpha);
-        for (; i < DMAX.loopBound(n); i += DMAX.length()) {
-          DoubleVector vx = DoubleVector.fromArray(DMAX, x, i);
-          DoubleVector vy = DoubleVector.fromArray(DMAX, y, i);
-          vx.fma(valpha, vy).intoArray(y, i);
-        }
-        for (; i < n; i += 1) {
-          y[i] += alpha * x[i];
-        }
-      }
-    } else {
-      super.daxpy(n, alpha, x, incx, y, incy);
-    }
-  }
-
-  // sum(x * y)
-  @Override
-  public float sdot(int n, float[] x, int incx, float[] y, int incy) {
-    if (incx == 1 && incy == 1) {
-      float sum = 0.0f;
-      int i = 0;
-      FloatVector vsum = FloatVector.zero(FMAX);
-      for (; i < FMAX.loopBound(n); i += FMAX.length()) {
-        FloatVector vx = FloatVector.fromArray(FMAX, x, i);
-        FloatVector vy = FloatVector.fromArray(FMAX, y, i);
-        vsum = vx.fma(vy, vsum);
-      }
-      sum += vsum.reduceLanes(VectorOperators.ADD);
-      for (; i < n; i += 1) {
-        sum += x[i] * y[i];
-      }
-      return sum;
-    } else {
-      return super.sdot(n, x, incx, y, incy);
-    }
-  }
-
-  // sum(x * y)
-  @Override
-  public double ddot(int n, double[] x, int incx, double[] y, int incy) {
-    if (incx == 1 && incy == 1) {
-      double sum = 0.;
-      int i = 0;
-      DoubleVector vsum = DoubleVector.zero(DMAX);
-      for (; i < DMAX.loopBound(n); i += DMAX.length()) {
-        DoubleVector vx = DoubleVector.fromArray(DMAX, x, i);
-        DoubleVector vy = DoubleVector.fromArray(DMAX, y, i);
-        vsum = vx.fma(vy, vsum);
-      }
-      sum += vsum.reduceLanes(VectorOperators.ADD);
-      for (; i < n; i += 1) {
-        sum += x[i] * y[i];
-      }
-      return sum;
-    } else {
-      return super.ddot(n, x, incx, y, incy);
-    }
-  }
-
-  // x = alpha * x
-  @Override
-  public void dscal(int n, double alpha, double[] x, int incx) {
-    if (incx == 1) {
-      if (alpha != 1.) {
-        int i = 0;
-        DoubleVector valpha = DoubleVector.broadcast(DMAX, alpha);
-        for (; i < DMAX.loopBound(n); i += DMAX.length()) {
-          DoubleVector vx = DoubleVector.fromArray(DMAX, x, i);
-          vx.mul(valpha).intoArray(x, i);
-        }
-        for (; i < n; i += 1) {
-          x[i] *= alpha;
-        }
-      }
-    } else {
-      super.dscal(n, alpha, x, incx);
-    }
-  }
-
-  // y := alpha * a * x + beta * y
-  @Override
-  public void dspmv(String uplo, int n, double alpha, double[] a, double[] x, int incx, double beta, double[] y, int incy) {
-    if (uplo.equals("U") && incx == 1 && incy == 1) {
-      // y = beta * y
-      dscal(n, beta, y, 1);
-      // y += alpha * A * x
-      if (alpha != 0.) {
-        for (int row = 0; row < n; row += 1) {
-          int col = 0;
-          DoubleVector valphaxrow = DoubleVector.broadcast(DMAX, alpha * x[row]);
-          for (; col < DMAX.loopBound(row); col += DMAX.length()) {
-            DoubleVector vx = DoubleVector.fromArray(DMAX, x, col);
-            DoubleVector vy = DoubleVector.fromArray(DMAX, y, col);
-            DoubleVector va = DoubleVector.fromArray(DMAX, a, col + row * (row + 1) / 2);
-            y[row] += alpha * vx.mul(va).reduceLanes(VectorOperators.ADD);
-            valphaxrow.fma(va, vy).intoArray(y, col);
-          }
-          for (; col < row; col += 1) {
-            y[row] += alpha * x[col] * a[col + row * (row + 1) / 2];
-            y[col] += alpha * x[row] * a[col + row * (row + 1) / 2];
-          }
-          y[row] += alpha * x[col] * a[col + row * (row + 1) / 2];
-        }
-      }
-    } else {
-      super.dspmv(uplo, n, alpha, a, x, incx, beta, y, incy);
-    }
-  }
-
-  // a += alpha * x * x.t
-  @Override
-  public void dspr(String uplo, int n, double alpha, double[] x, int incx, double[] a) {
-    if (uplo.equals("U") && incx == 1) {
-      if (alpha != 0.) {
-        for (int row = 0; row < n; row += 1) {
-          int col = 0;
-          DoubleVector valphax = DoubleVector.broadcast(DMAX, alpha * x[row]);
-          for (; col < DMAX.loopBound(row + 1); col += DMAX.length()) {
-            DoubleVector vx = DoubleVector.fromArray(DMAX, x, col);
-            DoubleVector va = DoubleVector.fromArray(DMAX, a, col + row * (row + 1) / 2);
-            vx.fma(valphax, va).intoArray(a, col + row * (row + 1) / 2);
-          }
-          for (; col < row + 1; col += 1) {
-            a[col + row * (row + 1) / 2] += alpha * x[row] * x[col];
-          }
-        }
-      }
-    } else {
-      super.dspr(uplo, n, alpha, x, incx, a);
-    }
-  }
-
-  // a += alpha * x * x.t
-  @Override
-  public void dsyr(String uplo, int n, double alpha, double[] x, int incx, double[] a, int lda) {
-    if (uplo.equals("U") && incx == 1) {
-      if (alpha != 0.) {
-        for (int row = 0; row < n; row += 1) {
-          int col = 0;
-          DoubleVector valphax = DoubleVector.broadcast(DMAX, alpha * x[row]);
-          for (; col < DMAX.loopBound(row + 1); col += DMAX.length()) {
-            DoubleVector vx = DoubleVector.fromArray(DMAX, x, col);
-            DoubleVector va = DoubleVector.fromArray(DMAX, a, col + row * n);
-            vx.fma(valphax, va).intoArray(a, col + row * n);
-          }
-          for (; col < row + 1; col += 1) {
-            a[col + row * n] += alpha * x[row] * x[col];
-          }
-        }
-      }
-    } else {
-      super.dsyr(uplo, n, alpha, x, incx, a, lda);
-    }
-  }
-
-  // y = alpha * A * x + beta * y
-  @Override
-  public void dgemv(String trans, int m, int n, double alpha, double[] a, int lda, double[] x, int incx, double beta, double[] y, int incy) {
-    if (trans.equals("T") && incx == 1 && incy == 1 && lda == m) {
-      // y = beta * y
-      dscal(n, beta, y, 1);
-      // y += alpha * A * x
-      if (alpha != 0.) {
-        for (int col = 0; col < n; col += 1) {
-          int row = 0;
-          for (; row < DMAX.loopBound(m); row += DMAX.length()) {
-            DoubleVector vx = DoubleVector.fromArray(DMAX, x, row);
-            DoubleVector va = DoubleVector.fromArray(DMAX, a, row + col * m);
-            y[col] += alpha * vx.mul(va).reduceLanes(VectorOperators.ADD);
-          }
-          for (; row < m; row += 1) {
-            y[col] += alpha * x[row] * a[row + col * m];
-          }
-        }
-      }
-    } else {
-      super.dgemv(trans, m, n, alpha, a, lda, x, incx, beta, y, incy);
-    }
-  }
-
-  @Override
-  public void dgemm(String transa, String transb, int m, int n, int k, double alpha, double[] a, int lda, double[] b, int ldb, double beta, double[] c, int ldc) {
-    if (transa.equals("T") && transb.equals("N") && lda == k && ldb == k && ldc == m) {
-      // C = beta * C
-      dscal(m * n, beta, c, 1);
-      // C += alpha * A * B
-      if (alpha != 0.) {
-        for (int col = 0; col < n; col += 1) {
-          for (int row = 0; row < m; row += 1) {
-            int i = 0;
-            for (; i < DMAX.loopBound(k); i += DMAX.length()) {
-              DoubleVector va = DoubleVector.fromArray(DMAX, a, i + row * k);
-              DoubleVector vb = DoubleVector.fromArray(DMAX, b, i + col * k);
-              c[row + col * m] += alpha * va.mul(vb).reduceLanes(VectorOperators.ADD);
-            }
-            for (; i < k; i += 1) {
-              c[row + col * m] += alpha * a[i + row * k] * b[i + col * k];
-            }
-          }
-        }
-      }
-    } else {
-      super.dgemm(transa, transb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc);
-    }
-  }
-}

From 494254b0322689cd8f6e968d01eb6f94e2b31bf4 Mon Sep 17 00:00:00 2001
From: Ludovic Henry <luhenry@microsoft.com>
Date: Sun, 20 Dec 2020 14:05:50 -0800
Subject: [PATCH 04/16] Make sure we don't depend on spark-core in mllib-local

We still use it for the benchmarks but it should go away at some point
---
 mllib-local/pom.xml                                   |  5 -----
 .../main/scala/org/apache/spark/ml/linalg/BLAS.scala  | 11 +++++++----
 .../org/apache/spark/ml/linalg/BLASBenchmark.scala    |  7 ++++---
 3 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/mllib-local/pom.xml b/mllib-local/pom.xml
index 481d985f744da..fbdc0d8578863 100644
--- a/mllib-local/pom.xml
+++ b/mllib-local/pom.xml
@@ -56,11 +56,6 @@
       <groupId>org.apache.spark</groupId>
       <artifactId>spark-tags_${scala.binary.version}</artifactId>
     </dependency>
-    <dependency>
-      <groupId>org.apache.spark</groupId>
-      <artifactId>spark-core_${scala.binary.version}</artifactId>
-      <version>${project.version}</version>
-    </dependency>
 
     <!--
       This spark-tags test-dep is needed even though it isn't used in this module, otherwise testing-cmds that exclude
diff --git a/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala b/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
index 8c2bbe2f6f757..b8a4ec0a1e93c 100644
--- a/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
+++ b/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
@@ -20,8 +20,6 @@ package org.apache.spark.ml.linalg
 import com.github.fommil.netlib.{BLAS => NetlibBLAS, F2jBLAS}
 import scala.util.Try
 
-import org.apache.spark.util.Utils
-
 /**
  * BLAS routines for MLlib's vectors and matrices.
  */
@@ -34,10 +32,15 @@ private[spark] object BLAS extends Serializable {
   // For level-1 function dspmv, use javaBLAS for better performance.
   private[ml] def javaBLAS: NetlibBLAS = {
     if (_javaBLAS == null) {
+      // scalastyle:off classforname
       _javaBLAS = Try(
-        Utils.classForName("dev.ludovic.blas.VectorizedBLAS")
+        Class.forName("dev.ludovic.blas.VectorizedBLAS", true,
+                        Option(Thread.currentThread().getContextClassLoader)
+                          .getOrElse(getClass.getClassLoader))
              .newInstance()
-             .asInstanceOf[NetlibBLAS]).getOrElse(new F2jBLAS)
+             .asInstanceOf[NetlibBLAS])
+          .getOrElse(new F2jBLAS)
+      // scalastyle:on classforname
     }
     _javaBLAS
   }
diff --git a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
index 48ad5a71aae31..b13f98c552c8f 100644
--- a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
+++ b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
@@ -43,9 +43,10 @@ object BLASBenchmark extends BenchmarkBase {
     val f2jBLAS = new F2jBLAS
     val nativeBLAS = NetlibBLAS.getInstance
     val vectorBLAS = Try(
-        Utils.classForName("org.apache.spark.ml.linalg.VectorizedBLAS")
-             .newInstance()
-             .asInstanceOf[NetlibBLAS]).getOrElse(new F2jBLAS)
+      Utils.classForName("dev.ludovic.blas.VectorizedBLAS")
+            .newInstance()
+            .asInstanceOf[NetlibBLAS])
+        .getOrElse(new F2jBLAS)
 
     // scalastyle:off println
     println("nativeBLAS = " + nativeBLAS.getClass.getName)

From 628af4fa70f66df5e61db79036ea10a8a8b1bb6a Mon Sep 17 00:00:00 2001
From: Ludovic Henry <luhenry@microsoft.com>
Date: Mon, 21 Dec 2020 07:35:32 -0800
Subject: [PATCH 05/16] Remove dependency on scala.utils.Try

---
 .../org/apache/spark/ml/linalg/BLAS.scala      | 18 ++++++++++--------
 .../apache/spark/ml/linalg/BLASBenchmark.scala | 14 +++++++++-----
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala b/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
index b8a4ec0a1e93c..85cc1e26262e6 100644
--- a/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
+++ b/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
@@ -18,7 +18,6 @@
 package org.apache.spark.ml.linalg
 
 import com.github.fommil.netlib.{BLAS => NetlibBLAS, F2jBLAS}
-import scala.util.Try
 
 /**
  * BLAS routines for MLlib's vectors and matrices.
@@ -33,13 +32,16 @@ private[spark] object BLAS extends Serializable {
   private[ml] def javaBLAS: NetlibBLAS = {
     if (_javaBLAS == null) {
       // scalastyle:off classforname
-      _javaBLAS = Try(
-        Class.forName("dev.ludovic.blas.VectorizedBLAS", true,
-                        Option(Thread.currentThread().getContextClassLoader)
-                          .getOrElse(getClass.getClassLoader))
-             .newInstance()
-             .asInstanceOf[NetlibBLAS])
-          .getOrElse(new F2jBLAS)
+      _javaBLAS =
+        try {
+          Class.forName("dev.ludovic.blas.VectorizedBLAS", true,
+                          Option(Thread.currentThread().getContextClassLoader)
+                            .getOrElse(getClass.getClassLoader))
+              .newInstance()
+              .asInstanceOf[NetlibBLAS]
+        } catch {
+          case _: Throwable => new F2jBLAS
+        }
       // scalastyle:on classforname
     }
     _javaBLAS
diff --git a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
index b13f98c552c8f..c4d025046129f 100644
--- a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
+++ b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
@@ -21,7 +21,6 @@ import com.github.fommil.netlib.{BLAS => NetlibBLAS, F2jBLAS}
 import scala.util.Try
 
 import org.apache.spark.benchmark.{Benchmark, BenchmarkBase}
-import org.apache.spark.util.Utils
 
 /**
  * Serialization benchmark for BLAS.
@@ -42,11 +41,16 @@ object BLASBenchmark extends BenchmarkBase {
 
     val f2jBLAS = new F2jBLAS
     val nativeBLAS = NetlibBLAS.getInstance
-    val vectorBLAS = Try(
-      Utils.classForName("dev.ludovic.blas.VectorizedBLAS")
+    val vectorBLAS =
+      try {
+        Class.forName("dev.ludovic.blas.VectorizedBLAS", true,
+                        Option(Thread.currentThread().getContextClassLoader)
+                          .getOrElse(getClass.getClassLoader))
             .newInstance()
-            .asInstanceOf[NetlibBLAS])
-        .getOrElse(new F2jBLAS)
+            .asInstanceOf[NetlibBLAS]
+      } catch {
+        case _: Throwable => new F2jBLAS
+      }
 
     // scalastyle:off println
     println("nativeBLAS = " + nativeBLAS.getClass.getName)

From 53bf41b63bd5f8c5c5e24dff9187e855671071ac Mon Sep 17 00:00:00 2001
From: Ludovic Henry <luhenry@microsoft.com>
Date: Mon, 21 Dec 2020 16:20:44 -0800
Subject: [PATCH 06/16] Fix indentation

---
 .../src/main/scala/org/apache/spark/ml/linalg/BLAS.scala  | 8 ++++----
 .../scala/org/apache/spark/ml/linalg/BLASBenchmark.scala  | 7 ++++---
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala b/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
index 85cc1e26262e6..4482e118dce0f 100644
--- a/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
+++ b/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
@@ -31,18 +31,18 @@ private[spark] object BLAS extends Serializable {
   // For level-1 function dspmv, use javaBLAS for better performance.
   private[ml] def javaBLAS: NetlibBLAS = {
     if (_javaBLAS == null) {
-      // scalastyle:off classforname
       _javaBLAS =
         try {
+          // scalastyle:off classforname
           Class.forName("dev.ludovic.blas.VectorizedBLAS", true,
                           Option(Thread.currentThread().getContextClassLoader)
                             .getOrElse(getClass.getClassLoader))
-              .newInstance()
-              .asInstanceOf[NetlibBLAS]
+               .newInstance()
+               .asInstanceOf[NetlibBLAS]
+          // scalastyle:on classforname
         } catch {
           case _: Throwable => new F2jBLAS
         }
-      // scalastyle:on classforname
     }
     _javaBLAS
   }
diff --git a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
index c4d025046129f..f40ac8a93223c 100644
--- a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
+++ b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
@@ -18,7 +18,6 @@
 package org.apache.spark.ml.linalg
 
 import com.github.fommil.netlib.{BLAS => NetlibBLAS, F2jBLAS}
-import scala.util.Try
 
 import org.apache.spark.benchmark.{Benchmark, BenchmarkBase}
 
@@ -43,11 +42,13 @@ object BLASBenchmark extends BenchmarkBase {
     val nativeBLAS = NetlibBLAS.getInstance
     val vectorBLAS =
       try {
+        // scalastyle:off classforname
         Class.forName("dev.ludovic.blas.VectorizedBLAS", true,
                         Option(Thread.currentThread().getContextClassLoader)
                           .getOrElse(getClass.getClassLoader))
-            .newInstance()
-            .asInstanceOf[NetlibBLAS]
+             .newInstance()
+             .asInstanceOf[NetlibBLAS]
+        // scalastyle:on classforname
       } catch {
         case _: Throwable => new F2jBLAS
       }

From 001810902cf04d545a2b28d6352378aa2f5db668 Mon Sep 17 00:00:00 2001
From: Ludovic Henry <luhenry@microsoft.com>
Date: Mon, 21 Dec 2020 16:21:54 -0800
Subject: [PATCH 07/16] Extend benchmark coverage

---
 .../spark/ml/linalg/BLASBenchmark.scala       | 217 ++++++++++++++++--
 1 file changed, 202 insertions(+), 15 deletions(-)

diff --git a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
index f40ac8a93223c..eb4ab518f6221 100644
--- a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
+++ b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
@@ -86,6 +86,32 @@ object BLASBenchmark extends BenchmarkBase {
       benchmark.run()
     }
 
+    runBenchmark("ddot") {
+      val n = 1e7.toInt
+      val x = Array.fill(n) { rnd.nextDouble }
+      val y = Array.fill(n) { rnd.nextDouble }
+
+      val benchmark = new Benchmark("ddot", n, iters, output = output)
+
+      benchmark.addCase("f2j") { _ =>
+        f2jBLAS.ddot(n, x, 1, y, 1)
+      }
+
+      if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("native") { _ =>
+          nativeBLAS.ddot(n, x, 1, y, 1)
+        }
+      }
+
+      if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("vector") { _ =>
+          vectorBLAS.ddot(n, x, 1, y, 1)
+        }
+      }
+
+      benchmark.run()
+    }
+
     runBenchmark("sdot") {
       val n = 1e7.toInt
       val x = Array.fill(n) { rnd.nextFloat }
@@ -112,52 +138,83 @@ object BLASBenchmark extends BenchmarkBase {
       benchmark.run()
     }
 
-    runBenchmark("ddot") {
+    runBenchmark("dscal") {
       val n = 1e7.toInt
+      val a = rnd.nextDouble
       val x = Array.fill(n) { rnd.nextDouble }
-      val y = Array.fill(n) { rnd.nextDouble }
 
-      val benchmark = new Benchmark("ddot", n, iters, output = output)
+      val benchmark = new Benchmark("dscal", n, iters, output = output)
 
       benchmark.addCase("f2j") { _ =>
-        f2jBLAS.ddot(n, x, 1, y, 1)
+        f2jBLAS.dscal(n, a, x, 1)
       }
 
       if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
         benchmark.addCase("native") { _ =>
-          nativeBLAS.ddot(n, x, 1, y, 1)
+          nativeBLAS.dscal(n, a, x, 1)
         }
       }
 
       if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
         benchmark.addCase("vector") { _ =>
-          vectorBLAS.ddot(n, x, 1, y, 1)
+          vectorBLAS.dscal(n, a, x, 1)
         }
       }
 
       benchmark.run()
     }
 
-    runBenchmark("dscal") {
+    runBenchmark("sscal") {
       val n = 1e7.toInt
-      val a = rnd.nextDouble
+      val a = rnd.nextFloat
+      val x = Array.fill(n) { rnd.nextFloat }
+
+      val benchmark = new Benchmark("sscal", n, iters, output = output)
+
+      benchmark.addCase("f2j") { _ =>
+        f2jBLAS.sscal(n, a, x, 1)
+      }
+
+      if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("native") { _ =>
+          nativeBLAS.sscal(n, a, x, 1)
+        }
+      }
+
+      if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("vector") { _ =>
+          vectorBLAS.sscal(n, a, x, 1)
+        }
+      }
+
+      benchmark.run()
+    }
+
+    runBenchmark("dgemv[N]") {
+      val m = 1e4.toInt
+      val n = 1e3.toInt
+      val alpha = rnd.nextDouble
+      val a = Array.fill(m * n) { rnd.nextDouble }
+      val lda = m
       val x = Array.fill(n) { rnd.nextDouble }
+      val beta = rnd.nextDouble
+      val y = Array.fill(m) { rnd.nextDouble }
 
-      val benchmark = new Benchmark("dscal", n, iters, output = output)
+      val benchmark = new Benchmark("dgemv[N]", n, iters, output = output)
 
       benchmark.addCase("f2j") { _ =>
-        f2jBLAS.dscal(n, a, x, 1)
+        f2jBLAS.dgemv("N", m, n, alpha, a, lda, x, 1, beta, y, 1)
       }
 
       if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
         benchmark.addCase("native") { _ =>
-          nativeBLAS.dscal(n, a, x, 1)
+          nativeBLAS.dgemv("N", m, n, alpha, a, lda, x, 1, beta, y, 1)
         }
       }
 
       if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
         benchmark.addCase("vector") { _ =>
-          vectorBLAS.dscal(n, a, x, 1)
+          vectorBLAS.dgemv("N", m, n, alpha, a, lda, x, 1, beta, y, 1)
         }
       }
 
@@ -168,7 +225,7 @@ object BLASBenchmark extends BenchmarkBase {
       val m = 1e4.toInt
       val n = 1e3.toInt
       val alpha = rnd.nextDouble
-      val a = Array.fill(n * m) { rnd.nextDouble }
+      val a = Array.fill(m * n) { rnd.nextDouble }
       val lda = m
       val x = Array.fill(m) { rnd.nextDouble }
       val beta = rnd.nextDouble
@@ -195,14 +252,144 @@ object BLASBenchmark extends BenchmarkBase {
       benchmark.run()
     }
 
+    runBenchmark("sgemv[N]") {
+      val m = 1e4.toInt
+      val n = 1e3.toInt
+      val alpha = rnd.nextFloat
+      val a = Array.fill(m * n) { rnd.nextFloat }
+      val lda = m
+      val x = Array.fill(n) { rnd.nextFloat }
+      val beta = rnd.nextFloat
+      val y = Array.fill(m) { rnd.nextFloat }
+
+      val benchmark = new Benchmark("sgemv[N]", n, iters, output = output)
+
+      benchmark.addCase("f2j") { _ =>
+        f2jBLAS.sgemv("N", m, n, alpha, a, lda, x, 1, beta, y, 1)
+      }
+
+      if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("native") { _ =>
+          nativeBLAS.sgemv("N", m, n, alpha, a, lda, x, 1, beta, y, 1)
+        }
+      }
+
+      if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("vector") { _ =>
+          vectorBLAS.sgemv("N", m, n, alpha, a, lda, x, 1, beta, y, 1)
+        }
+      }
+
+      benchmark.run()
+    }
+
+    runBenchmark("sgemv[T]") {
+      val m = 1e4.toInt
+      val n = 1e3.toInt
+      val alpha = rnd.nextFloat
+      val a = Array.fill(m * n) { rnd.nextFloat }
+      val lda = m
+      val x = Array.fill(m) { rnd.nextFloat }
+      val beta = rnd.nextFloat
+      val y = Array.fill(n) { rnd.nextFloat }
+
+      val benchmark = new Benchmark("sgemv[T]", n, iters, output = output)
+
+      benchmark.addCase("f2j") { _ =>
+        f2jBLAS.sgemv("T", m, n, alpha, a, lda, x, 1, beta, y, 1)
+      }
+
+      if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("native") { _ =>
+          nativeBLAS.sgemv("T", m, n, alpha, a, lda, x, 1, beta, y, 1)
+        }
+      }
+
+      if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("vector") { _ =>
+          vectorBLAS.sgemv("T", m, n, alpha, a, lda, x, 1, beta, y, 1)
+        }
+      }
+
+      benchmark.run()
+    }
+
+    runBenchmark("dgemm[N,N]") {
+      val m = 1e3.toInt
+      val n = 1e2.toInt
+      val k = 1e3.toInt
+      val alpha = rnd.nextDouble
+      val a = Array.fill(m * k) { rnd.nextDouble }
+      val lda = m
+      val b = Array.fill(k * n) { rnd.nextDouble }
+      val ldb = k
+      val beta = rnd.nextDouble
+      val c = Array.fill(m * n) { rnd.nextDouble }
+      var ldc = m
+
+      val benchmark = new Benchmark("dgemm[N,N]", m*n, iters, output = output)
+
+      benchmark.addCase("f2j") { _ =>
+        f2jBLAS.dgemm("N", "N", m, n, k, alpha, a, lda, b, ldb, beta, c, ldc)
+      }
+
+      if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("native") { _ =>
+          nativeBLAS.dgemm("N", "N", m, n, k, alpha, a, lda, b, ldb, beta, c, ldc)
+        }
+      }
+
+      if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("vector") { _ =>
+          vectorBLAS.dgemm("N", "N", m, n, k, alpha, a, lda, b, ldb, beta, c, ldc)
+        }
+      }
+
+      benchmark.run()
+    }
+
+    runBenchmark("dgemm[N,T]") {
+      val m = 1e3.toInt
+      val n = 1e2.toInt
+      val k = 1e3.toInt
+      val alpha = rnd.nextDouble
+      val a = Array.fill(m * k) { rnd.nextDouble }
+      val lda = m
+      val b = Array.fill(k * n) { rnd.nextDouble }
+      val ldb = n
+      val beta = rnd.nextDouble
+      val c = Array.fill(m * n) { rnd.nextDouble }
+      var ldc = m
+
+      val benchmark = new Benchmark("dgemm[N,T]", m*n, iters, output = output)
+
+      benchmark.addCase("f2j") { _ =>
+        f2jBLAS.dgemm("N", "T", m, n, k, alpha, a, lda, b, ldb, beta, c, ldc)
+      }
+
+      if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("native") { _ =>
+          nativeBLAS.dgemm("N", "T", m, n, k, alpha, a, lda, b, ldb, beta, c, ldc)
+        }
+      }
+
+      if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("vector") { _ =>
+          vectorBLAS.dgemm("N", "T", m, n, k, alpha, a, lda, b, ldb, beta, c, ldc)
+        }
+      }
+
+      benchmark.run()
+    }
+
     runBenchmark("dgemm[T,N]") {
       val m = 1e3.toInt
       val n = 1e2.toInt
       val k = 1e3.toInt
       val alpha = rnd.nextDouble
-      val a = Array.fill(k * m) { rnd.nextDouble }
+      val a = Array.fill(m * k) { rnd.nextDouble }
       val lda = k
-      val b = Array.fill(n * k) { rnd.nextDouble }
+      val b = Array.fill(k * n) { rnd.nextDouble }
       val ldb = k
       val beta = rnd.nextDouble
       val c = Array.fill(m * n) { rnd.nextDouble }

From d4b737cdca52a113ac39121117c4fe3655205850 Mon Sep 17 00:00:00 2001
From: Ludovic Henry <luhenry@microsoft.com>
Date: Mon, 21 Dec 2020 16:22:35 -0800
Subject: [PATCH 08/16] Fix how --add-modules=jdk.incubator.vector is passed

---
 project/SparkBuild.scala | 16 ++++++----------
 1 file changed, 6 insertions(+), 10 deletions(-)

diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala
index ac6b68cbb121b..04eef554049b2 100644
--- a/project/SparkBuild.scala
+++ b/project/SparkBuild.scala
@@ -288,6 +288,12 @@ object SparkBuild extends PomBuild {
     publishLocal in SbtCompile := publishTask(publishLocalConfiguration in SbtCompile).value,
     publishLocal := Seq(publishLocal in MavenCompile, publishLocal in SbtCompile).dependOn.value,
 
+    javaOptions ++= {
+      val versionParts = System.getProperty("java.version").split("[+.\\-]+", 3)
+      var major = versionParts(0).toInt
+      if (major >= 16) Seq("--add-modules=jdk.incubator.vector") else Seq.empty
+    },
+
     javacOptions in (Compile, doc) ++= {
       val versionParts = System.getProperty("java.version").split("[+.\\-]+", 3)
       var major = versionParts(0).toInt
@@ -425,11 +431,6 @@ object SparkBuild extends PomBuild {
     outputStrategy in run := Some (StdoutOutput),
 
     javaOptions += "-Xmx2g",
-    javaOptions += {
-      val versionParts = System.getProperty("java.version").split("[+.\\-]+", 3)
-      var major = versionParts(0).toInt
-      if (major >= 16) "--add-modules=jdk.incubator.vector" else ""
-    },
 
     sparkShell := {
       (runMain in Compile).toTask(" org.apache.spark.repl.Main -usejavacp").value
@@ -1065,11 +1066,6 @@ object TestSettings {
     javaOptions in Test ++= System.getProperties.asScala.filter(_._1.startsWith("spark"))
       .map { case (k,v) => s"-D$k=$v" }.toSeq,
     javaOptions in Test += "-ea",
-    javaOptions in Test += {
-      val versionParts = System.getProperty("java.version").split("[+.\\-]+", 3)
-      var major = versionParts(0).toInt
-      if (major >= 16) "--add-modules=jdk.incubator.vector" else ""
-    },
     // SPARK-29282 This is for consistency between JDK8 and JDK11.
     javaOptions in Test ++= "-Xmx4g -Xss4m -XX:+UseParallelGC -XX:-UseDynamicNumberOfGCThreads"
       .split(" ").toSeq,

From 3ac7e0e2b7be5cc169511e94845d60470f615f47 Mon Sep 17 00:00:00 2001
From: Ludovic Henry <luhenry@microsoft.com>
Date: Mon, 21 Dec 2020 16:36:58 -0800
Subject: [PATCH 09/16] Add java 17 build to CI

---
 .github/workflows/build_and_test.yml | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml
index 0048bc7ffba0d..ad5d18c9274cf 100644
--- a/.github/workflows/build_and_test.yml
+++ b/.github/workflows/build_and_test.yml
@@ -394,6 +394,31 @@ jobs:
         ./build/mvn $MAVEN_CLI_OPTS -DskipTests -Pyarn -Pmesos -Pkubernetes -Phive -Phive-thriftserver -Phadoop-cloud -Djava.version=11 install
         rm -rf ~/.m2/repository/org/apache/spark
 
+  java-17:
+    name: Java 17 build with Maven
+    runs-on: ubuntu-20.04
+    steps:
+    - name: Checkout Spark repository
+      uses: actions/checkout@v2
+    - name: Cache Maven local repository
+      uses: actions/cache@v2
+      with:
+        path: ~/.m2/repository
+        key: java17-maven-${{ hashFiles('**/pom.xml') }}
+        restore-keys: |
+          java17-maven-
+    - name: Install Java 17
+      uses: actions/setup-java@v1
+      with:
+        java-version: 17-ea
+    - name: Build with Maven
+      run: |
+        export MAVEN_OPTS="-Xmx2g -XX:ReservedCodeCacheSize=1g -Dorg.slf4j.simpleLogger.defaultLogLevel=WARN"
+        export MAVEN_CLI_OPTS="--no-transfer-progress"
+        # It uses Maven's 'install' intentionally, see https://github.com/apache/spark/pull/26414.
+        ./build/mvn $MAVEN_CLI_OPTS -DskipTests -Pyarn -Pmesos -Pkubernetes -Phive -Phive-thriftserver -Phadoop-cloud -Djava.version=11 install
+        rm -rf ~/.m2/repository/org/apache/spark
+
   scala-213:
     name: Scala 2.13 build with SBT
     runs-on: ubuntu-20.04

From 0c52da59e3bf1011dadffd5b237fc1df4081f607 Mon Sep 17 00:00:00 2001
From: Ludovic Henry <luhenry@microsoft.com>
Date: Tue, 22 Dec 2020 01:33:42 -0800
Subject: [PATCH 10/16] Bump dev.ludovic:vectorizedblas-blas to 0.1.9

Brings acceleration for sscal, dgemm[N,N], dgemm[N,T], dgemv[N], sgemv[N] and sgemv[T]
---
 mllib-local/pom.xml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mllib-local/pom.xml b/mllib-local/pom.xml
index fbdc0d8578863..d5822e20523c8 100644
--- a/mllib-local/pom.xml
+++ b/mllib-local/pom.xml
@@ -94,7 +94,7 @@
         <dependency>
           <groupId>dev.ludovic</groupId>
           <artifactId>vectorizedblas-blas</artifactId>
-          <version>0.1.7</version>
+          <version>0.1.9</version>
         </dependency>
       </dependencies>
     </profile>

From e23a0fd6d40bd3d9cac54903ba5a04b5abe5d140 Mon Sep 17 00:00:00 2001
From: Ludovic Henry <luhenry@microsoft.com>
Date: Tue, 22 Dec 2020 05:13:27 -0800
Subject: [PATCH 11/16] Bring back VectorizedBLAS in the repo

This reverts commit bc9cfb066430155f3d3a8911b4e47c4a5b6c717d.
---
 mllib-local/pom.xml                           |  32 +-
 .../org/apache/spark/ml/linalg/BLAS.scala     |   2 +-
 .../spark/ml/linalg/BLASBenchmark.scala       |   2 +-
 .../spark/ml/linalg/VectorizedBLAS.java       | 485 ++++++++++++++++++
 4 files changed, 512 insertions(+), 9 deletions(-)
 create mode 100644 mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java

diff --git a/mllib-local/pom.xml b/mllib-local/pom.xml
index d5822e20523c8..3665678157ee2 100644
--- a/mllib-local/pom.xml
+++ b/mllib-local/pom.xml
@@ -90,13 +90,31 @@
     </profile>
     <profile>
       <id>vectorized</id>
-      <dependencies>
-        <dependency>
-          <groupId>dev.ludovic</groupId>
-          <artifactId>vectorizedblas-blas</artifactId>
-          <version>0.1.9</version>
-        </dependency>
-      </dependencies>
+      <properties>
+        <extra.source.dir>src/vectorized/java</extra.source.dir>
+      </properties>
+      <build>
+        <plugins>
+          <plugin>
+            <groupId>org.codehaus.mojo</groupId>
+            <artifactId>build-helper-maven-plugin</artifactId>
+            <executions>
+              <execution>
+                <id>add-vectorized-sources</id>
+                <phase>generate-sources</phase>
+                <goals>
+                  <goal>add-source</goal>
+                </goals>
+                <configuration>
+                  <sources>
+                    <source>${extra.source.dir}</source>
+                  </sources>
+                </configuration>
+              </execution>
+            </executions>
+          </plugin>
+        </plugins>
+      </build>
     </profile>
   </profiles>
   <build>
diff --git a/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala b/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
index 4482e118dce0f..518c71129a970 100644
--- a/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
+++ b/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
@@ -34,7 +34,7 @@ private[spark] object BLAS extends Serializable {
       _javaBLAS =
         try {
           // scalastyle:off classforname
-          Class.forName("dev.ludovic.blas.VectorizedBLAS", true,
+          Class.forName("org.apache.spark.ml.linalg.VectorizedBLAS", true,
                           Option(Thread.currentThread().getContextClassLoader)
                             .getOrElse(getClass.getClassLoader))
                .newInstance()
diff --git a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
index eb4ab518f6221..29ad7cbc65f0e 100644
--- a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
+++ b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
@@ -43,7 +43,7 @@ object BLASBenchmark extends BenchmarkBase {
     val vectorBLAS =
       try {
         // scalastyle:off classforname
-        Class.forName("dev.ludovic.blas.VectorizedBLAS", true,
+        Class.forName("org.apache.spark.ml.linalg.VectorizedBLAS", true,
                         Option(Thread.currentThread().getContextClassLoader)
                           .getOrElse(getClass.getClassLoader))
              .newInstance()
diff --git a/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java b/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java
new file mode 100644
index 0000000000000..a1c1e1228783a
--- /dev/null
+++ b/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java
@@ -0,0 +1,485 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.ml.linalg;
+
+import com.github.fommil.netlib.F2jBLAS;
+import jdk.incubator.vector.DoubleVector;
+import jdk.incubator.vector.FloatVector;
+import jdk.incubator.vector.VectorOperators;
+import jdk.incubator.vector.VectorSpecies;
+
+public class VectorizedBLAS extends F2jBLAS {
+
+  private static final VectorSpecies<Float>  FMAX = FloatVector.SPECIES_MAX;
+  private static final VectorSpecies<Double> DMAX = DoubleVector.SPECIES_MAX;
+
+  // y += alpha * x
+  @Override
+  public void daxpy(int n, double alpha, double[] x, int incx, double[] y, int incy) {
+    if (n >= 0
+        && x != null && x.length >= n && incx == 1
+        && y != null && y.length >= n && incy == 1) {
+      if (alpha != 0.) {
+        DoubleVector valpha = DoubleVector.broadcast(DMAX, alpha);
+        int i = 0;
+        for (; i < DMAX.loopBound(n); i += DMAX.length()) {
+          DoubleVector vx = DoubleVector.fromArray(DMAX, x, i);
+          DoubleVector vy = DoubleVector.fromArray(DMAX, y, i);
+          vx.fma(valpha, vy).intoArray(y, i);
+        }
+        for (; i < n; i += 1) {
+          y[i] += alpha * x[i];
+        }
+      }
+    } else {
+      super.daxpy(n, alpha, x, incx, y, incy);
+    }
+  }
+
+  // sum(x * y)
+  @Override
+  public float sdot(int n, float[] x, int incx, float[] y, int incy) {
+    if (n >= 0
+        && x != null && x.length >= n && incx == 1
+        && y != null && y.length >= n && incy == 1) {
+      float sum = 0.0f;
+      int i = 0;
+      FloatVector vsum = FloatVector.zero(FMAX);
+      for (; i < FMAX.loopBound(n); i += FMAX.length()) {
+        FloatVector vx = FloatVector.fromArray(FMAX, x, i);
+        FloatVector vy = FloatVector.fromArray(FMAX, y, i);
+        vsum = vx.fma(vy, vsum);
+      }
+      sum += vsum.reduceLanes(VectorOperators.ADD);
+      for (; i < n; i += 1) {
+        sum += x[i] * y[i];
+      }
+      return sum;
+    } else {
+      return super.sdot(n, x, incx, y, incy);
+    }
+  }
+
+  // sum(x * y)
+  @Override
+  public double ddot(int n, double[] x, int incx, double[] y, int incy) {
+    if (n >= 0
+        && x != null && x.length >= n && incx == 1
+        && y != null && y.length >= n && incy == 1) {
+      double sum = 0.;
+      int i = 0;
+      DoubleVector vsum = DoubleVector.zero(DMAX);
+      for (; i < DMAX.loopBound(n); i += DMAX.length()) {
+        DoubleVector vx = DoubleVector.fromArray(DMAX, x, i);
+        DoubleVector vy = DoubleVector.fromArray(DMAX, y, i);
+        vsum = vx.fma(vy, vsum);
+      }
+      sum += vsum.reduceLanes(VectorOperators.ADD);
+      for (; i < n; i += 1) {
+        sum += x[i] * y[i];
+      }
+      return sum;
+    } else {
+      return super.ddot(n, x, incx, y, incy);
+    }
+  }
+
+  @Override
+  public void dscal(int n, double alpha, double[] x, int incx) {
+    dscal(n, alpha, x, 0, incx);
+  }
+
+  // x = alpha * x
+  @Override
+  public void dscal(int n, double alpha, double[] x, int offsetx, int incx) {
+    if (n >= 0 && x != null && x.length >= offsetx + n && incx == 1) {
+      if (alpha != 1.) {
+        DoubleVector valpha = DoubleVector.broadcast(DMAX, alpha);
+        int i = 0;
+        for (; i < DMAX.loopBound(n); i += DMAX.length()) {
+          DoubleVector vx = DoubleVector.fromArray(DMAX, x, offsetx + i);
+          vx.mul(valpha).intoArray(x, offsetx + i);
+        }
+        for (; i < n; i += 1) {
+          x[offsetx + i] *= alpha;
+        }
+      }
+    } else {
+      super.dscal(n, alpha, x, offsetx, incx);
+    }
+  }
+
+  @Override
+  public void sscal(int n, float alpha, float[] x, int incx) {
+    sscal(n, alpha, x, 0, incx);
+  }
+
+  // x = alpha * x
+  @Override
+  public void sscal(int n, float alpha, float[] x, int offsetx, int incx) {
+    if (n >= 0 && x != null && x.length >= offsetx + n && incx == 1) {
+      if (alpha != 1.) {
+        FloatVector valpha = FloatVector.broadcast(FMAX, alpha);
+        int i = 0;
+        for (; i < FMAX.loopBound(n); i += FMAX.length()) {
+          FloatVector vx = FloatVector.fromArray(FMAX, x, offsetx + i);
+          vx.mul(valpha).intoArray(x, offsetx + i);
+        }
+        for (; i < n; i += 1) {
+          x[offsetx + i] *= alpha;
+        }
+      }
+    } else {
+      super.sscal(n, alpha, x, offsetx, incx);
+    }
+  }
+
+  // y = alpha * a * x + beta * y
+  @Override
+  public void dspmv(String uplo, int n, double alpha, double[] a,
+      double[] x, int incx, double beta, double[] y, int incy) {
+    if ("U".equals(uplo)
+        && n >= 0
+        && a != null && a.length >= n * (n + 1) / 2
+        && x != null && x.length >= n && incx == 1
+        && y != null && y.length >= n && incy == 1) {
+      // y = beta * y
+      dscal(n, beta, y, 1);
+      // y += alpha * A * x
+      if (alpha != 0.) {
+        DoubleVector valpha = DoubleVector.broadcast(DMAX, alpha);
+        for (int row = 0; row < n; row += 1) {
+          int col = 0;
+          DoubleVector vyrowsum = DoubleVector.zero(DMAX);
+          DoubleVector valphaxrow = DoubleVector.broadcast(DMAX, alpha * x[row]);
+          for (; col < DMAX.loopBound(row); col += DMAX.length()) {
+            DoubleVector vx = DoubleVector.fromArray(DMAX, x, col);
+            DoubleVector vy = DoubleVector.fromArray(DMAX, y, col);
+            DoubleVector va = DoubleVector.fromArray(DMAX, a, col + row * (row + 1) / 2);
+            vyrowsum = valpha.mul(vx).fma(va, vyrowsum);
+            valphaxrow.fma(va, vy).intoArray(y, col);
+          }
+          y[row] += vyrowsum.reduceLanes(VectorOperators.ADD);
+          for (; col < row; col += 1) {
+            y[row] += alpha * x[col] * a[col + row * (row + 1) / 2];
+            y[col] += alpha * x[row] * a[col + row * (row + 1) / 2];
+          }
+          y[row] += alpha * x[col] * a[col + row * (row + 1) / 2];
+        }
+      }
+    } else {
+      super.dspmv(uplo, n, alpha, a, x, incx, beta, y, incy);
+    }
+  }
+
+  // a += alpha * x * x.t
+  @Override
+  public void dspr(String uplo, int n, double alpha, double[] x, int incx, double[] a) {
+    if ("U".equals(uplo)
+        && n >= 0
+        && x != null && x.length >= n && incx == 1
+        && a != null && a.length >= n * (n + 1) / 2) {
+      if (alpha != 0.) {
+        for (int row = 0; row < n; row += 1) {
+          int col = 0;
+          DoubleVector valphaxrow = DoubleVector.broadcast(DMAX, alpha * x[row]);
+          for (; col < DMAX.loopBound(row + 1); col += DMAX.length()) {
+            DoubleVector vx = DoubleVector.fromArray(DMAX, x, col);
+            DoubleVector va = DoubleVector.fromArray(DMAX, a, col + row * (row + 1) / 2);
+            vx.fma(valphaxrow, va).intoArray(a, col + row * (row + 1) / 2);
+          }
+          for (; col < row + 1; col += 1) {
+            a[col + row * (row + 1) / 2] += alpha * x[row] * x[col];
+          }
+        }
+      }
+    } else {
+      super.dspr(uplo, n, alpha, x, incx, a);
+    }
+  }
+
+  // a += alpha * x * x.t
+  @Override
+  public void dsyr(String uplo, int n, double alpha, double[] x, int incx, double[] a, int lda) {
+    if ("U".equals(uplo)
+        && n >= 0
+        && x != null && x.length >= n && incx == 1
+        && a != null && a.length >= n * n && lda == n) {
+      if (alpha != 0.) {
+        for (int row = 0; row < n; row += 1) {
+          int col = 0;
+          DoubleVector valphaxrow = DoubleVector.broadcast(DMAX, alpha * x[row]);
+          for (; col < DMAX.loopBound(row + 1); col += DMAX.length()) {
+            DoubleVector vx = DoubleVector.fromArray(DMAX, x, col);
+            DoubleVector va = DoubleVector.fromArray(DMAX, a, col + row * n);
+            vx.fma(valphaxrow, va).intoArray(a, col + row * n);
+          }
+          for (; col < row + 1; col += 1) {
+            a[col + row * n] += alpha * x[row] * x[col];
+          }
+        }
+      }
+    } else {
+      super.dsyr(uplo, n, alpha, x, incx, a, lda);
+    }
+  }
+
+  @Override
+  public void dgemv(String trans, int m, int n,
+      double alpha, double[] a, int lda, double[] x, int incx,
+      double beta, double[] y, int incy) {
+    dgemv(trans, m, n, alpha, a, 0, lda, x, 0, incx, beta, y, 0, incy);
+  }
+
+  // y = alpha * A * x + beta * y
+  @Override
+  public void dgemv(String trans, int m, int n,
+      double alpha, double[] a, int offseta, int lda, double[] x, int offsetx, int incx,
+      double beta, double[] y, int offsety, int incy) {
+    if ("N".equals(trans)
+        && m >= 0 && n >= 0
+        && a != null && a.length >= offseta + m * n && lda == m
+        && x != null && x.length >= offsetx + n && incx == 1
+        && y != null && y.length >= offsety + m && incy == 1) {
+      // y = beta * y
+      dscal(m, beta, y, offsety, 1);
+      // y += alpha * A * x
+      if (alpha != 0.) {
+        DoubleVector valpha = DoubleVector.broadcast(DMAX, alpha);
+        for (int col = 0; col < n; col += 1) {
+          int row = 0;
+          for (; row < DMAX.loopBound(m); row += DMAX.length()) {
+            DoubleVector va = DoubleVector.fromArray(DMAX, a, offseta + row + col * m);
+            DoubleVector vy = DoubleVector.fromArray(DMAX, y, offsety + row);
+            valpha.mul(x[offsetx + col]).fma(va, vy)
+                  .intoArray(y, offsety + row);
+          }
+          for (; row < m; row += 1) {
+            y[offsety + row] += alpha * x[offsetx + col] * a[offseta + row + col * m];
+          }
+        }
+      }
+    } else if ("T".equals(trans)
+        && m >= 0 && n >= 0
+        && a != null && a.length >= offseta + m * n && lda == m
+        && x != null && x.length >= offsetx + m && incx == 1
+        && y != null && y.length >= offsety + n && incy == 1) {
+      if (alpha != 0. || beta != 1.) {
+        for (int col = 0; col < n; col += 1) {
+          double sum = 0.;
+          int row = 0;
+          DoubleVector vsum = DoubleVector.zero(DMAX);
+          for (; row < DMAX.loopBound(m); row += DMAX.length()) {
+            DoubleVector va = DoubleVector.fromArray(DMAX, a, offseta + row + col * m);
+            DoubleVector vx = DoubleVector.fromArray(DMAX, x, offsetx + row);
+            vsum = va.fma(vx, vsum);
+          }
+          sum += vsum.reduceLanes(VectorOperators.ADD);
+          for (; row < m; row += 1) {
+            sum += x[offsetx + row] * a[offseta + row + col * m];
+          }
+          y[offsety + col] = alpha * sum + beta * y[offsety + col];
+        }
+      }
+    } else {
+      super.dgemv(trans, m, n, alpha, a, offseta, lda, x, offsetx, incx, beta, y, offsety, incy);
+    }
+  }
+
+  @Override
+  public void sgemv(String trans, int m, int n,
+      float alpha, float[] a, int lda, float[] x, int incx,
+      float beta, float[] y, int incy) {
+    sgemv(trans, m, n, alpha, a, 0, lda, x, 0, incx, beta, y, 0, incy);
+  }
+
+  // y = alpha * A * x + beta * y
+  @Override
+  public void sgemv(String trans, int m, int n,
+      float alpha, float[] a, int offseta, int lda, float[] x, int offsetx, int incx,
+      float beta, float[] y, int offsety, int incy) {
+    if ("N".equals(trans)
+        && m >= 0 && n >= 0
+        && a != null && a.length >= offseta + m * n && lda == m
+        && x != null && x.length >= offsetx + n && incx == 1
+        && y != null && y.length >= offsety + m && incy == 1) {
+      // y = beta * y
+      sscal(m, beta, y, offsety, 1);
+      // y += alpha * A * x
+      if (alpha != 0.f) {
+        FloatVector valpha = FloatVector.broadcast(FMAX, alpha);
+        for (int col = 0; col < n; col += 1) {
+          int row = 0;
+          for (; row < FMAX.loopBound(m); row += FMAX.length()) {
+            FloatVector va = FloatVector.fromArray(FMAX, a, offseta + row + col * m);
+            FloatVector vy = FloatVector.fromArray(FMAX, y, offsety + row);
+            valpha.mul(x[offsetx + col]).fma(va, vy)
+                  .intoArray(y, offsety + row);
+          }
+          for (; row < m; row += 1) {
+            y[offsety + row] += alpha * x[offsetx + col] * a[offseta + row + col * m];
+          }
+        }
+      }
+    } else if ("T".equals(trans)
+        && m >= 0 && n >= 0
+        && a != null && a.length >= offseta + m * n && lda == m
+        && x != null && x.length >= offsetx + m && incx == 1
+        && y != null && y.length >= offsety + n && incy == 1) {
+      if (alpha != 0. || beta != 1.) {
+        for (int col = 0; col < n; col += 1) {
+          float sum = 0.f;
+          int row = 0;
+          FloatVector vsum = FloatVector.zero(FMAX);
+          for (; row < FMAX.loopBound(m); row += FMAX.length()) {
+            FloatVector va = FloatVector.fromArray(FMAX, a, offseta + row + col * m);
+            FloatVector vx = FloatVector.fromArray(FMAX, x, offsetx + row);
+            vsum = va.fma(vx, vsum);
+          }
+          sum += vsum.reduceLanes(VectorOperators.ADD);
+          for (; row < m; row += 1) {
+            sum += x[offsetx + row] * a[offseta + row + col * m];
+          }
+          y[offsety + col] = alpha * sum + beta * y[offsety + col];
+        }
+      }
+    } else {
+      super.sgemv(trans, m, n, alpha, a, offseta, lda, x, offsetx, incx, beta, y, offsety, incy);
+    }
+  }
+
+  @Override
+  public void dgemm(String transa, String transb, int m, int n, int k,
+      double alpha, double[] a, int lda, double[] b, int ldb,
+      double beta, double[] c, int ldc) {
+    dgemm(transa, transb, m, n, k, alpha, a, 0, lda, b, 0, ldb, beta, c, ldc, 0);
+  }
+
+  // c = alpha * a * b + beta * c
+  @Override
+  public void dgemm(String transa, String transb, int m, int n, int k,
+      double alpha, double[] a, int offseta, int lda, double[] b, int offsetb, int ldb,
+      double beta, double[] c, int ldc, int offsetc) {
+    // System.out.println(String.format("dgemm(transa=%s, transb=%s, m=%s, n=%s, k=%s, alpha=%s, a=%s, lda=%s, b=%s, ldb=%s, beta=%s, c=%s, ldc=%s)",
+    //     transa, transb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc));
+    if ("N".equals(transa) && "N".equals(transb)
+        && m >= 0 && n >= 0 && k >= 0
+        && a != null && a.length >= offseta + m * k && lda == m
+        && b != null && b.length >= offsetb + k * n && ldb == k
+        && c != null && c.length >= offsetc + m * n && ldc == m) {
+      // C = beta * C
+      dscal(m * n, beta, c, offsetc, 1);
+      // C += alpha * A * B
+      if (alpha != 0.) {
+        DoubleVector valpha = DoubleVector.broadcast(DMAX, alpha);
+        for (int col = 0; col < n; col += 1) {
+          for (int i = 0; i < k; i += 1) {
+            int row = 0;
+            for (; row < DMAX.loopBound(m); row += DMAX.length()) {
+              DoubleVector va = DoubleVector.fromArray(DMAX, a, offseta + i * m + row);
+              DoubleVector vc = DoubleVector.fromArray(DMAX, c, offsetc + col * m + row);
+              valpha.mul(b[offsetb + col * k + i]).fma(va, vc)
+                    .intoArray(c, offsetc + col * m + row);
+            }
+            for (; row < m; row += 1) {
+              c[offsetc + col * m + row] += alpha * a[offseta + i * m + row] * b[offsetb + col * k + i];
+            }
+          }
+        }
+      }
+    } else if ("N".equals(transa) && "T".equals(transb)
+        && m >= 0 && n >= 0 && k >= 0
+        && a != null && a.length >= offseta + m * k && lda == m
+        && b != null && b.length >= offsetb + k * n && ldb == n
+        && c != null && c.length >= offsetc + m * n && ldc == m) {
+      // C = beta * C
+      dscal(m * n, beta, c, offsetc, 1);
+      // C += alpha * A * B
+      if (alpha != 0.) {
+        DoubleVector valpha = DoubleVector.broadcast(DMAX, alpha);
+        for (int i = 0; i < k; i += 1) {
+          for (int col = 0; col < n; col += 1) {
+            int row = 0;
+            for (; row < DMAX.loopBound(m); row += DMAX.length()) {
+              DoubleVector va = DoubleVector.fromArray(DMAX, a, offseta + i * m + row);
+              DoubleVector vc = DoubleVector.fromArray(DMAX, c, offsetc + col * m + row);
+              valpha.mul(b[offsetb + col + i * n]).fma(va, vc)
+                    .intoArray(c, offsetc + col * m + row);
+            }
+            for (; row < m; row += 1) {
+              c[offsetc + col * m + row] += alpha * a[offseta + i * m + row] * b[offsetb + col + i * n];
+            }
+          }
+        }
+      }
+    } else if ("T".equals(transa) && "N".equals(transb)
+        && m >= 0 && n >= 0 && k >= 0
+        && a != null && a.length >= offseta + m * k && lda == k
+        && b != null && b.length >= offsetb + k * n && ldb == k
+        && c != null && c.length >= offsetc + m * n && ldc == m) {
+      if (alpha != 0. || beta != 1.) {
+        for (int col = 0; col < n; col += 1) {
+          for (int row = 0; row < m; row += 1) {
+            double sum = 0.;
+            int i = 0;
+            DoubleVector vsum = DoubleVector.zero(DMAX);
+            for (; i < DMAX.loopBound(k); i += DMAX.length()) {
+              DoubleVector va = DoubleVector.fromArray(DMAX, a, offseta + i + row * k);
+              DoubleVector vb = DoubleVector.fromArray(DMAX, b, offsetb + col * k + i);
+              vsum = va.fma(vb, vsum);
+            }
+            sum += vsum.reduceLanes(VectorOperators.ADD);
+            for (; i < k; i += 1) {
+              sum += a[offseta + i + row * k] * b[offsetb + col * k + i];
+            }
+            if (beta != 0.) {
+              c[offsetc + col * m + row] = alpha * sum + beta * c[offsetc + col * m + row];
+            } else {
+              c[offsetc + col * m + row] = alpha * sum;
+            }
+          }
+        }
+      }
+    } else if ("T".equals(transa) && "T".equals(transb)
+        && m >= 0 && n >= 0 && k >= 0
+        && a != null && a.length >= offseta + m * k && lda == k
+        && b != null && b.length >= offsetb + k * n && ldb == n
+        && c != null && c.length >= offsetc + m * n && ldc == m) {
+      if (alpha != 0. || beta != 1.) {
+        // FIXME: do block by block
+        for (int col = 0; col < n; col += 1) {
+          for (int row = 0; row < m; row += 1) {
+            double sum = 0.;
+            for (int i = 0; i < k; i += 1) {
+              sum += a[offseta + i + row * k] * b[offsetb + col + i * n];
+            }
+            if (beta != 0.) {
+              c[offsetc + col * m + row] = alpha * sum + beta * c[offsetc + col * m + row];
+            } else {
+              c[offsetc + col * m + row] = alpha * sum;
+            }
+          }
+        }
+      }
+    } else {
+      super.dgemm(transa, transb, m, n, k,
+                  alpha, a, offseta, lda, b, offsetb, ldb,
+                  beta, c, offsetc, ldc);
+    }
+  }
+}

From 18e62f92fc9479095537a92b2a35f9d91cb64533 Mon Sep 17 00:00:00 2001
From: Ludovic Henry <git@ludovic.dev>
Date: Tue, 22 Dec 2020 20:36:43 +0000
Subject: [PATCH 12/16] Update according to review

---
 .github/workflows/build_and_test.yml          | 25 -------------------
 .../spark/ml/linalg/BLASBenchmark.scala       |  2 +-
 .../spark/ml/linalg/VectorizedBLAS.java       |  6 ++---
 3 files changed, 3 insertions(+), 30 deletions(-)

diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml
index ad5d18c9274cf..0048bc7ffba0d 100644
--- a/.github/workflows/build_and_test.yml
+++ b/.github/workflows/build_and_test.yml
@@ -394,31 +394,6 @@ jobs:
         ./build/mvn $MAVEN_CLI_OPTS -DskipTests -Pyarn -Pmesos -Pkubernetes -Phive -Phive-thriftserver -Phadoop-cloud -Djava.version=11 install
         rm -rf ~/.m2/repository/org/apache/spark
 
-  java-17:
-    name: Java 17 build with Maven
-    runs-on: ubuntu-20.04
-    steps:
-    - name: Checkout Spark repository
-      uses: actions/checkout@v2
-    - name: Cache Maven local repository
-      uses: actions/cache@v2
-      with:
-        path: ~/.m2/repository
-        key: java17-maven-${{ hashFiles('**/pom.xml') }}
-        restore-keys: |
-          java17-maven-
-    - name: Install Java 17
-      uses: actions/setup-java@v1
-      with:
-        java-version: 17-ea
-    - name: Build with Maven
-      run: |
-        export MAVEN_OPTS="-Xmx2g -XX:ReservedCodeCacheSize=1g -Dorg.slf4j.simpleLogger.defaultLogLevel=WARN"
-        export MAVEN_CLI_OPTS="--no-transfer-progress"
-        # It uses Maven's 'install' intentionally, see https://github.com/apache/spark/pull/26414.
-        ./build/mvn $MAVEN_CLI_OPTS -DskipTests -Pyarn -Pmesos -Pkubernetes -Phive -Phive-thriftserver -Phadoop-cloud -Djava.version=11 install
-        rm -rf ~/.m2/repository/org/apache/spark
-
   scala-213:
     name: Scala 2.13 build with SBT
     runs-on: ubuntu-20.04
diff --git a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
index 29ad7cbc65f0e..f70e1f1067db2 100644
--- a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
+++ b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
@@ -28,7 +28,7 @@ import org.apache.spark.benchmark.{Benchmark, BenchmarkBase}
  * 1. without sbt: bin/spark-submit --class <this class> <spark mllib test jar>
  * 2. build/sbt "mllib/test:runMain <this class>"
  * 3. generate result: SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt "mllib/test:runMain <this class>"
- *    Results will be written to "benchmarks/UDTSerializationBenchmark-results.txt".
+ *    Results will be written to "benchmarks/BLASBenchmark-results.txt".
  * }}}
  */
 object BLASBenchmark extends BenchmarkBase {
diff --git a/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java b/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java
index a1c1e1228783a..7db1bb9111a00 100644
--- a/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java
+++ b/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java
@@ -367,16 +367,14 @@ public void sgemv(String trans, int m, int n,
   public void dgemm(String transa, String transb, int m, int n, int k,
       double alpha, double[] a, int lda, double[] b, int ldb,
       double beta, double[] c, int ldc) {
-    dgemm(transa, transb, m, n, k, alpha, a, 0, lda, b, 0, ldb, beta, c, ldc, 0);
+    dgemm(transa, transb, m, n, k, alpha, a, 0, lda, b, 0, ldb, beta, c, 0, ldc);
   }
 
   // c = alpha * a * b + beta * c
   @Override
   public void dgemm(String transa, String transb, int m, int n, int k,
       double alpha, double[] a, int offseta, int lda, double[] b, int offsetb, int ldb,
-      double beta, double[] c, int ldc, int offsetc) {
-    // System.out.println(String.format("dgemm(transa=%s, transb=%s, m=%s, n=%s, k=%s, alpha=%s, a=%s, lda=%s, b=%s, ldb=%s, beta=%s, c=%s, ldc=%s)",
-    //     transa, transb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc));
+      double beta, double[] c, int offsetc, int ldc) {
     if ("N".equals(transa) && "N".equals(transb)
         && m >= 0 && n >= 0 && k >= 0
         && a != null && a.length >= offseta + m * k && lda == m

From 1aea88aa917b424a6c843276962d8ee0ba3bb17c Mon Sep 17 00:00:00 2001
From: Ludovic Henry <git@ludovic.dev>
Date: Tue, 22 Dec 2020 20:37:30 +0000
Subject: [PATCH 13/16] Access com.github.fommil.netlib.BLAS  through
 org.apache.spark.ml.linalg.BLAS instead

---
 .../scala/org/apache/spark/ml/linalg/Matrices.scala |  3 +--
 .../scala/org/apache/spark/ml/ann/BreezeUtil.scala  |  7 ++++---
 .../spark/ml/classification/GBTClassifier.scala     |  5 ++---
 .../org/apache/spark/ml/recommendation/ALS.scala    | 13 ++++++-------
 .../apache/spark/ml/regression/GBTRegressor.scala   |  5 ++---
 .../ml/classification/GBTClassifierSuite.scala      |  8 ++------
 .../apache/spark/ml/recommendation/ALSSuite.scala   |  7 +++----
 7 files changed, 20 insertions(+), 28 deletions(-)

diff --git a/mllib-local/src/main/scala/org/apache/spark/ml/linalg/Matrices.scala b/mllib-local/src/main/scala/org/apache/spark/ml/linalg/Matrices.scala
index 1254ed747aeb4..3272a50508a3d 100644
--- a/mllib-local/src/main/scala/org/apache/spark/ml/linalg/Matrices.scala
+++ b/mllib-local/src/main/scala/org/apache/spark/ml/linalg/Matrices.scala
@@ -22,7 +22,6 @@ import java.util.{Arrays, Random}
 import scala.collection.mutable.{ArrayBuffer, ArrayBuilder => MArrayBuilder, HashSet => MHashSet}
 
 import breeze.linalg.{CSCMatrix => BSM, DenseMatrix => BDM, Matrix => BM}
-import com.github.fommil.netlib.BLAS.{getInstance => blas}
 
 import org.apache.spark.annotation.Since
 
@@ -457,7 +456,7 @@ class DenseMatrix @Since("2.0.0") (
     if (isTransposed) {
       Iterator.tabulate(numCols) { j =>
         val col = new Array[Double](numRows)
-        blas.dcopy(numRows, values, j, numCols, col, 0, 1)
+        BLAS.nativeBLAS.dcopy(numRows, values, j, numCols, col, 0, 1)
         new DenseVector(col)
       }
     } else {
diff --git a/mllib/src/main/scala/org/apache/spark/ml/ann/BreezeUtil.scala b/mllib/src/main/scala/org/apache/spark/ml/ann/BreezeUtil.scala
index 6bbe7e1cb2134..ce9093a6238ee 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/ann/BreezeUtil.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/ann/BreezeUtil.scala
@@ -18,7 +18,8 @@
 package org.apache.spark.ml.ann
 
 import breeze.linalg.{DenseMatrix => BDM, DenseVector => BDV}
-import com.github.fommil.netlib.BLAS.{getInstance => NativeBLAS}
+
+import org.apache.spark.ml.linalg.BLAS
 
 /**
  * In-place DGEMM and DGEMV for Breeze
@@ -41,7 +42,7 @@ private[ann] object BreezeUtil {
     require(A.cols == B.rows, "A & B Dimension mismatch!")
     require(A.rows == C.rows, "A & C Dimension mismatch!")
     require(B.cols == C.cols, "A & C Dimension mismatch!")
-    NativeBLAS.dgemm(transposeString(A), transposeString(B), C.rows, C.cols, A.cols,
+    BLAS.nativeBLAS.dgemm(transposeString(A), transposeString(B), C.rows, C.cols, A.cols,
       alpha, A.data, A.offset, A.majorStride, B.data, B.offset, B.majorStride,
       beta, C.data, C.offset, C.rows)
   }
@@ -57,7 +58,7 @@ private[ann] object BreezeUtil {
   def dgemv(alpha: Double, A: BDM[Double], x: BDV[Double], beta: Double, y: BDV[Double]): Unit = {
     require(A.cols == x.length, "A & x Dimension mismatch!")
     require(A.rows == y.length, "A & y Dimension mismatch!")
-    NativeBLAS.dgemv(transposeString(A), A.rows, A.cols,
+    BLAS.nativeBLAS.dgemv(transposeString(A), A.rows, A.cols,
       alpha, A.data, A.offset, A.majorStride, x.data, x.offset, x.stride,
       beta, y.data, y.offset, y.stride)
   }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala
index ca43ef863003a..453e28609a0c2 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala
@@ -17,14 +17,13 @@
 
 package org.apache.spark.ml.classification
 
-import com.github.fommil.netlib.BLAS.{getInstance => blas}
 import org.json4s.{DefaultFormats, JObject}
 import org.json4s.JsonDSL._
 
 import org.apache.spark.annotation.Since
 import org.apache.spark.internal.Logging
 import org.apache.spark.ml.feature.Instance
-import org.apache.spark.ml.linalg.{DenseVector, SparseVector, Vector, Vectors}
+import org.apache.spark.ml.linalg.{BLAS, DenseVector, SparseVector, Vector, Vectors}
 import org.apache.spark.ml.param.ParamMap
 import org.apache.spark.ml.regression.DecisionTreeRegressionModel
 import org.apache.spark.ml.tree._
@@ -371,7 +370,7 @@ class GBTClassificationModel private[ml](
   /** Raw prediction for the positive class. */
   private def margin(features: Vector): Double = {
     val treePredictions = _trees.map(_.rootNode.predictImpl(features).prediction)
-    blas.ddot(getNumTrees, treePredictions, 1, _treeWeights, 1)
+    BLAS.nativeBLAS.ddot(getNumTrees, treePredictions, 1, _treeWeights, 1)
   }
 
   /** (private[ml]) Convert to a model in the old API */
diff --git a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
index adfea09909fc6..9a8dd803ca457 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
@@ -26,7 +26,6 @@ import scala.reflect.ClassTag
 import scala.util.{Sorting, Try}
 import scala.util.hashing.byteswap64
 
-import com.github.fommil.netlib.BLAS.{getInstance => blas}
 import com.google.common.collect.{Ordering => GuavaOrdering}
 import org.apache.hadoop.fs.Path
 import org.json4s.DefaultFormats
@@ -894,9 +893,9 @@ object ALS extends DefaultParamsReadable[ALS] with Logging {
       require(c >= 0.0)
       require(a.length == k)
       copyToDouble(a)
-      blas.dspr(upper, k, c, da, 1, ata)
+      BLAS.nativeBLAS.dspr(upper, k, c, da, 1, ata)
       if (b != 0.0) {
-        blas.daxpy(k, b, da, 1, atb, 1)
+        BLAS.nativeBLAS.daxpy(k, b, da, 1, atb, 1)
       }
       this
     }
@@ -904,8 +903,8 @@ object ALS extends DefaultParamsReadable[ALS] with Logging {
     /** Merges another normal equation object. */
     def merge(other: NormalEquation): NormalEquation = {
       require(other.k == k)
-      blas.daxpy(ata.length, 1.0, other.ata, 1, ata, 1)
-      blas.daxpy(atb.length, 1.0, other.atb, 1, atb, 1)
+      BLAS.nativeBLAS.daxpy(ata.length, 1.0, other.ata, 1, ata, 1)
+      BLAS.nativeBLAS.daxpy(atb.length, 1.0, other.atb, 1, atb, 1)
       this
     }
 
@@ -1279,8 +1278,8 @@ object ALS extends DefaultParamsReadable[ALS] with Logging {
           val random = new XORShiftRandom(byteswap64(seed ^ srcBlockId))
           val factors = Array.fill(inBlock.srcIds.length) {
             val factor = Array.fill(rank)(random.nextGaussian().toFloat)
-            val nrm = blas.snrm2(rank, factor, 1)
-            blas.sscal(rank, 1.0f / nrm, factor, 1)
+            val nrm = BLAS.nativeBLAS.snrm2(rank, factor, 1)
+            BLAS.nativeBLAS.sscal(rank, 1.0f / nrm, factor, 1)
             factor
           }
           (srcBlockId, factors)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala
index 78d5ddaa2758b..fd8af71d43568 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala
@@ -17,13 +17,12 @@
 
 package org.apache.spark.ml.regression
 
-import com.github.fommil.netlib.BLAS.{getInstance => blas}
 import org.json4s.{DefaultFormats, JObject}
 import org.json4s.JsonDSL._
 
 import org.apache.spark.annotation.Since
 import org.apache.spark.internal.Logging
-import org.apache.spark.ml.linalg.Vector
+import org.apache.spark.ml.linalg.{BLAS, Vector}
 import org.apache.spark.ml.param.ParamMap
 import org.apache.spark.ml.tree._
 import org.apache.spark.ml.tree.impl.GradientBoostedTrees
@@ -299,7 +298,7 @@ class GBTRegressionModel private[ml](
     // TODO: When we add a generic Boosting class, handle transform there?  SPARK-7129
     // Classifies by thresholding sum of weighted tree predictions
     val treePredictions = _trees.map(_.rootNode.predictImpl(features).prediction)
-    blas.ddot(getNumTrees, treePredictions, 1, _treeWeights, 1)
+    BLAS.nativeBLAS.ddot(getNumTrees, treePredictions, 1, _treeWeights, 1)
   }
 
   @Since("1.4.0")
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/GBTClassifierSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/GBTClassifierSuite.scala
index a2208edcb8394..9ab59990ff9dd 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/GBTClassifierSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/GBTClassifierSuite.scala
@@ -17,12 +17,10 @@
 
 package org.apache.spark.ml.classification
 
-import com.github.fommil.netlib.BLAS
-
 import org.apache.spark.{SparkException, SparkFunSuite}
 import org.apache.spark.ml.classification.LinearSVCSuite.generateSVMInput
 import org.apache.spark.ml.feature.{Instance, LabeledPoint}
-import org.apache.spark.ml.linalg.{Vector, Vectors}
+import org.apache.spark.ml.linalg.{BLAS, Vector, Vectors}
 import org.apache.spark.ml.param.ParamsSuite
 import org.apache.spark.ml.regression.DecisionTreeRegressionModel
 import org.apache.spark.ml.tree._
@@ -170,8 +168,6 @@ class GBTClassifierSuite extends MLTest with DefaultReadWriteTest {
     val numFeatures = trainingDataset.select(featuresCol).first().getAs[Vector](0).size
     assert(gbtModel.numFeatures === numFeatures)
 
-    val blas = BLAS.getInstance()
-
     val validationDataset = validationData.toDF(labelCol, featuresCol)
     testTransformer[(Double, Vector)](validationDataset, gbtModel,
       "rawPrediction", "features", "probability", "prediction") {
@@ -179,7 +175,7 @@ class GBTClassifierSuite extends MLTest with DefaultReadWriteTest {
         assert(raw.size === 2)
         // check that raw prediction is tree predictions dot tree weights
         val treePredictions = gbtModel.trees.map(_.rootNode.predictImpl(features).prediction)
-        val prediction = blas.ddot(gbtModel.getNumTrees, treePredictions, 1,
+        val prediction = BLAS.nativeBLAS.ddot(gbtModel.getNumTrees, treePredictions, 1,
           gbtModel.treeWeights, 1)
         assert(raw ~== Vectors.dense(-prediction, prediction) relTol eps)
 
diff --git a/mllib/src/test/scala/org/apache/spark/ml/recommendation/ALSSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/recommendation/ALSSuite.scala
index 28275eb06cf0d..cebd8cac057f1 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/recommendation/ALSSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/recommendation/ALSSuite.scala
@@ -24,14 +24,13 @@ import scala.collection.JavaConverters._
 import scala.collection.mutable
 import scala.collection.mutable.{ArrayBuffer, WrappedArray}
 
-import com.github.fommil.netlib.BLAS.{getInstance => blas}
 import org.apache.commons.io.FileUtils
 import org.apache.commons.io.filefilter.TrueFileFilter
 import org.scalatest.BeforeAndAfterEach
 
 import org.apache.spark._
 import org.apache.spark.internal.Logging
-import org.apache.spark.ml.linalg.Vectors
+import org.apache.spark.ml.linalg.{BLAS, Vectors}
 import org.apache.spark.ml.recommendation.ALS._
 import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTest, MLTestingUtils}
 import org.apache.spark.ml.util.TestingUtils._
@@ -296,7 +295,7 @@ class ALSSuite extends MLTest with DefaultReadWriteTest with Logging {
     for ((userId, userFactor) <- userFactors; (itemId, itemFactor) <- itemFactors) {
       val x = random.nextDouble()
       if (x < totalFraction) {
-        val rating = blas.sdot(rank, userFactor, 1, itemFactor, 1)
+        val rating = BLAS.nativeBLAS.sdot(rank, userFactor, 1, itemFactor, 1)
         if (x < trainingFraction) {
           val noise = noiseStd * random.nextGaussian()
           training += Rating(userId, itemId, rating + noise.toFloat)
@@ -1194,7 +1193,7 @@ object ALSSuite extends Logging {
     val training = ArrayBuffer.empty[Rating[Int]]
     val test = ArrayBuffer.empty[Rating[Int]]
     for ((userId, userFactor) <- userFactors; (itemId, itemFactor) <- itemFactors) {
-      val rating = blas.sdot(rank, userFactor, 1, itemFactor, 1)
+      val rating = BLAS.nativeBLAS.sdot(rank, userFactor, 1, itemFactor, 1)
       val threshold = if (rating > 0) positiveFraction else negativeFraction
       val observed = random.nextDouble() < threshold
       if (observed) {

From 43a82b17a75fb4aeab31771552ae88a9e46525d9 Mon Sep 17 00:00:00 2001
From: Ludovic Henry <git@ludovic.dev>
Date: Tue, 22 Dec 2020 21:28:01 +0000
Subject: [PATCH 14/16] Add missing benchmarks for dspmv, dspr, and dsyr

---
 .../spark/ml/linalg/BLASBenchmark.scala       | 107 ++++++++++++++++--
 1 file changed, 95 insertions(+), 12 deletions(-)

diff --git a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
index f70e1f1067db2..5a58f33365ecd 100644
--- a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
+++ b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
@@ -61,25 +61,25 @@ object BLASBenchmark extends BenchmarkBase {
 
     runBenchmark("daxpy") {
       val n = 1e7.toInt
-      val a = rnd.nextDouble
+      val alpha = rnd.nextDouble
       val x = Array.fill(n) { rnd.nextDouble }
       val y = Array.fill(n) { rnd.nextDouble }
 
       val benchmark = new Benchmark("daxpy", n, iters, output = output)
 
       benchmark.addCase("f2j") { _ =>
-        f2jBLAS.daxpy(n, a, x, 1, y, 1)
+        f2jBLAS.daxpy(n, alpha, x, 1, y, 1)
       }
 
       if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
         benchmark.addCase("native") { _ =>
-          nativeBLAS.daxpy(n, a, x, 1, y, 1)
+          nativeBLAS.daxpy(n, alpha, x, 1, y, 1)
         }
       }
 
       if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
         benchmark.addCase("vector") { _ =>
-          vectorBLAS.daxpy(n, a, x, 1, y, 1)
+          vectorBLAS.daxpy(n, alpha, x, 1, y, 1)
         }
       }
 
@@ -140,24 +140,24 @@ object BLASBenchmark extends BenchmarkBase {
 
     runBenchmark("dscal") {
       val n = 1e7.toInt
-      val a = rnd.nextDouble
+      val alpha = rnd.nextDouble
       val x = Array.fill(n) { rnd.nextDouble }
 
       val benchmark = new Benchmark("dscal", n, iters, output = output)
 
       benchmark.addCase("f2j") { _ =>
-        f2jBLAS.dscal(n, a, x, 1)
+        f2jBLAS.dscal(n, alpha, x, 1)
       }
 
       if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
         benchmark.addCase("native") { _ =>
-          nativeBLAS.dscal(n, a, x, 1)
+          nativeBLAS.dscal(n, alpha, x, 1)
         }
       }
 
       if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
         benchmark.addCase("vector") { _ =>
-          vectorBLAS.dscal(n, a, x, 1)
+          vectorBLAS.dscal(n, alpha, x, 1)
         }
       }
 
@@ -166,24 +166,107 @@ object BLASBenchmark extends BenchmarkBase {
 
     runBenchmark("sscal") {
       val n = 1e7.toInt
-      val a = rnd.nextFloat
+      val alpha = rnd.nextFloat
       val x = Array.fill(n) { rnd.nextFloat }
 
       val benchmark = new Benchmark("sscal", n, iters, output = output)
 
       benchmark.addCase("f2j") { _ =>
-        f2jBLAS.sscal(n, a, x, 1)
+        f2jBLAS.sscal(n, alpha, x, 1)
+      }
+
+      if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("native") { _ =>
+          nativeBLAS.sscal(n, alpha, x, 1)
+        }
+      }
+
+      if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("vector") { _ =>
+          vectorBLAS.sscal(n, alpha, x, 1)
+        }
+      }
+
+      benchmark.run()
+    }
+
+    runBenchmark("dspmv[U]") {
+      val n = 1e4.toInt
+      val alpha = rnd.nextDouble
+      val a = Array.fill(n * (n + 1) / 2) { rnd.nextDouble }
+      val x = Array.fill(n) { rnd.nextDouble }
+      val beta = rnd.nextDouble
+      val y = Array.fill(n) { rnd.nextDouble }
+
+      val benchmark = new Benchmark("dspmv[U]", n, iters, output = output)
+
+      benchmark.addCase("f2j") { _ =>
+        f2jBLAS.dspmv("U", n, alpha, a, x, 1, beta, y, 1)
+      }
+
+      if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("native") { _ =>
+          nativeBLAS.dspmv("U", n, alpha, a, x, 1, beta, y, 1)
+        }
+      }
+
+      if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("vector") { _ =>
+          vectorBLAS.dspmv("U", n, alpha, a, x, 1, beta, y, 1)
+        }
+      }
+
+      benchmark.run()
+    }
+
+    runBenchmark("dspr[U]") {
+      val n = 1e4.toInt
+      val alpha = rnd.nextDouble
+      val x = Array.fill(n) { rnd.nextDouble }
+      val a = Array.fill(n * (n + 1) / 2) { rnd.nextDouble }
+
+      val benchmark = new Benchmark("dspr[U]", n, iters, output = output)
+
+      benchmark.addCase("f2j") { _ =>
+        f2jBLAS.dspr("U", n, alpha, x, 1, a)
+      }
+
+      if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("native") { _ =>
+          nativeBLAS.dspr("U", n, alpha, x, 1, a)
+        }
+      }
+
+      if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
+        benchmark.addCase("vector") { _ =>
+          vectorBLAS.dspr("U", n, alpha, x, 1, a)
+        }
+      }
+
+      benchmark.run()
+    }
+
+    runBenchmark("dsyr[U]") {
+      val n = 1e4.toInt
+      val alpha = rnd.nextDouble
+      val x = Array.fill(n) { rnd.nextDouble }
+      val a = Array.fill(n * n) { rnd.nextDouble }
+
+      val benchmark = new Benchmark("dsyr[U]", n, iters, output = output)
+
+      benchmark.addCase("f2j") { _ =>
+        f2jBLAS.dsyr("U", n, alpha, x, 1, a, n)
       }
 
       if (!nativeBLAS.getClass.equals(classOf[F2jBLAS])) {
         benchmark.addCase("native") { _ =>
-          nativeBLAS.sscal(n, a, x, 1)
+          nativeBLAS.dsyr("U", n, alpha, x, 1, a, n)
         }
       }
 
       if (!vectorBLAS.getClass.equals(classOf[F2jBLAS])) {
         benchmark.addCase("vector") { _ =>
-          vectorBLAS.sscal(n, a, x, 1)
+          vectorBLAS.dsyr("U", n, alpha, x, 1, a, n)
         }
       }
 

From 14a0e60b5c459ae85e187ae900c2441f40b960cd Mon Sep 17 00:00:00 2001
From: Ludovic Henry <luhenry@microsoft.com>
Date: Mon, 4 Jan 2021 06:20:18 -0800
Subject: [PATCH 15/16] Correct typo in command to run benchmark

---
 .../test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
index 5a58f33365ecd..1dcfcf9ebb034 100644
--- a/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
+++ b/mllib-local/src/test/scala/org/apache/spark/ml/linalg/BLASBenchmark.scala
@@ -26,7 +26,7 @@ import org.apache.spark.benchmark.{Benchmark, BenchmarkBase}
  * To run this benchmark:
  * {{{
  * 1. without sbt: bin/spark-submit --class <this class> <spark mllib test jar>
- * 2. build/sbt "mllib/test:runMain <this class>"
+ * 2. build/sbt "mllib-local/test:runMain <this class>"
  * 3. generate result: SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt "mllib/test:runMain <this class>"
  *    Results will be written to "benchmarks/BLASBenchmark-results.txt".
  * }}}

From d93b27410b8223605f4891eae3f675914c8d0346 Mon Sep 17 00:00:00 2001
From: Ludovic Henry <git@ludovic.dev>
Date: Thu, 8 Apr 2021 15:00:18 +0200
Subject: [PATCH 16/16] Rename profile to jvm-vectorized

---
 mllib-local/pom.xml                                           | 4 ++--
 .../java/org/apache/spark/ml/linalg/VectorizedBLAS.java       | 0
 2 files changed, 2 insertions(+), 2 deletions(-)
 rename mllib-local/src/{vectorized => jvm-vectorized}/java/org/apache/spark/ml/linalg/VectorizedBLAS.java (100%)

diff --git a/mllib-local/pom.xml b/mllib-local/pom.xml
index 3665678157ee2..16cd55c8a45f0 100644
--- a/mllib-local/pom.xml
+++ b/mllib-local/pom.xml
@@ -89,9 +89,9 @@
       </dependencies>
     </profile>
     <profile>
-      <id>vectorized</id>
+      <id>jvm-vectorized</id>
       <properties>
-        <extra.source.dir>src/vectorized/java</extra.source.dir>
+        <extra.source.dir>src/jvm-vectorized/java</extra.source.dir>
       </properties>
       <build>
         <plugins>
diff --git a/mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java b/mllib-local/src/jvm-vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java
similarity index 100%
rename from mllib-local/src/vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java
rename to mllib-local/src/jvm-vectorized/java/org/apache/spark/ml/linalg/VectorizedBLAS.java