From a78e361b178eca4a9b89a25dc5067e3e0732dc33 Mon Sep 17 00:00:00 2001 From: Ivan Kylchik Date: Sun, 13 Feb 2022 16:01:05 +0300 Subject: [PATCH] Implement much faster dot product algorithm for tensors --- .../kscience/kmath/benchmarks/DotBenchmark.kt | 10 ++++++++++ .../kmath/tensors/core/DoubleTensorAlgebra.kt | 2 +- .../kmath/tensors/core/internal/linUtils.kt | 18 +++++++++++++----- .../tensors/core/TestDoubleTensorAlgebra.kt | 9 +++++++++ 4 files changed, 33 insertions(+), 6 deletions(-) diff --git a/benchmarks/src/jvmMain/kotlin/space/kscience/kmath/benchmarks/DotBenchmark.kt b/benchmarks/src/jvmMain/kotlin/space/kscience/kmath/benchmarks/DotBenchmark.kt index 63165baaa..53d7ee9b6 100644 --- a/benchmarks/src/jvmMain/kotlin/space/kscience/kmath/benchmarks/DotBenchmark.kt +++ b/benchmarks/src/jvmMain/kotlin/space/kscience/kmath/benchmarks/DotBenchmark.kt @@ -15,7 +15,9 @@ import space.kscience.kmath.linear.invoke import space.kscience.kmath.linear.linearSpace import space.kscience.kmath.multik.multikAlgebra import space.kscience.kmath.operations.DoubleField +import space.kscience.kmath.operations.invoke import space.kscience.kmath.structures.Buffer +import space.kscience.kmath.tensors.core.DoubleTensorAlgebra import kotlin.random.Random @State(Scope.Benchmark) @@ -32,6 +34,9 @@ internal class DotBenchmark { random.nextDouble() } + val tensor1 = DoubleTensorAlgebra.randomNormal(shape = intArrayOf(dim, dim), 12224) + val tensor2 = DoubleTensorAlgebra.randomNormal(shape = intArrayOf(dim, dim), 12225) + val cmMatrix1 = CMLinearSpace { matrix1.toCM() } val cmMatrix2 = CMLinearSpace { matrix2.toCM() } @@ -78,4 +83,9 @@ internal class DotBenchmark { fun doubleDot(blackhole: Blackhole) = with(DoubleField.linearSpace) { blackhole.consume(matrix1 dot matrix2) } + + @Benchmark + fun doubleTensorDot(blackhole: Blackhole) = DoubleTensorAlgebra.invoke { + blackhole.consume(tensor1 dot tensor2) + } } diff --git a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/DoubleTensorAlgebra.kt b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/DoubleTensorAlgebra.kt index 50252ad31..a75e5a8e3 100644 --- a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/DoubleTensorAlgebra.kt +++ b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/DoubleTensorAlgebra.kt @@ -421,7 +421,7 @@ public open class DoubleTensorAlgebra : for ((res, ab) in resTensor.matrixSequence().zip(newThis.matrixSequence().zip(newOther.matrixSequence()))) { val (a, b) = ab - dotTo(a.as2D(), b.as2D(), res.as2D(), l, m1, n) + dotTo(a, b, res, l, m1, n) } return if (penultimateDim) { diff --git a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/internal/linUtils.kt b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/internal/linUtils.kt index 2fb5b949f..aba6167ce 100644 --- a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/internal/linUtils.kt +++ b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/internal/linUtils.kt @@ -54,18 +54,26 @@ internal val BufferedTensor.matrices: VirtualBuffer> internal fun BufferedTensor.matrixSequence(): Sequence> = matrices.asSequence() internal fun dotTo( - a: MutableStructure2D, - b: MutableStructure2D, - res: MutableStructure2D, + a: BufferedTensor, + b: BufferedTensor, + res: BufferedTensor, l: Int, m: Int, n: Int, ) { + val aStart = a.bufferStart + val bStart = b.bufferStart + val resStart = res.bufferStart + + val aBuffer = a.mutableBuffer + val bBuffer = b.mutableBuffer + val resBuffer = res.mutableBuffer + for (i in 0 until l) { for (j in 0 until n) { var curr = 0.0 for (k in 0 until m) { - curr += a[i, k] * b[k, j] + curr += aBuffer[aStart + i * m + k] * bBuffer[bStart + k * n + j] } - res[i, j] = curr + resBuffer[resStart + i * n + j] = curr } } } diff --git a/kmath-tensors/src/commonTest/kotlin/space/kscience/kmath/tensors/core/TestDoubleTensorAlgebra.kt b/kmath-tensors/src/commonTest/kotlin/space/kscience/kmath/tensors/core/TestDoubleTensorAlgebra.kt index 03357f1e1..205ae2fee 100644 --- a/kmath-tensors/src/commonTest/kotlin/space/kscience/kmath/tensors/core/TestDoubleTensorAlgebra.kt +++ b/kmath-tensors/src/commonTest/kotlin/space/kscience/kmath/tensors/core/TestDoubleTensorAlgebra.kt @@ -107,6 +107,8 @@ internal class TestDoubleTensorAlgebra { val tensor11 = fromArray(intArrayOf(3, 2), doubleArrayOf(1.0, 2.0, 3.0, 4.0, 5.0, 6.0)) val tensor2 = fromArray(intArrayOf(3), doubleArrayOf(10.0, 20.0, 30.0)) val tensor3 = fromArray(intArrayOf(1, 1, 3), doubleArrayOf(-1.0, -2.0, -3.0)) + val tensor4 = fromArray(intArrayOf(2, 3, 3), (1..18).map { it.toDouble() }.toDoubleArray()) + val tensor5 = fromArray(intArrayOf(2, 3, 3), (1..18).map { 1 + it.toDouble() }.toDoubleArray()) val res12 = tensor1.dot(tensor2) assertTrue(res12.mutableBuffer.array() contentEquals doubleArrayOf(140.0, 320.0)) @@ -123,6 +125,13 @@ internal class TestDoubleTensorAlgebra { val res11 = tensor1.dot(tensor11) assertTrue(res11.mutableBuffer.array() contentEquals doubleArrayOf(22.0, 28.0, 49.0, 64.0)) assertTrue(res11.shape contentEquals intArrayOf(2, 2)) + + val res45 = tensor4.dot(tensor5) + assertTrue(res45.mutableBuffer.array() contentEquals doubleArrayOf( + 36.0, 42.0, 48.0, 81.0, 96.0, 111.0, 126.0, 150.0, 174.0, + 468.0, 501.0, 534.0, 594.0, 636.0, 678.0, 720.0, 771.0, 822.0 + )) + assertTrue(res45.shape contentEquals intArrayOf(2, 3, 3)) } @Test