Skip to content

Commit 73013bd

Browse files
committed
Add the section of matrix multiplication in README
1 parent 0378b97 commit 73013bd

File tree

2 files changed

+75
-10
lines changed

2 files changed

+75
-10
lines changed

README.md

Lines changed: 70 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
# Compute.scala
22

3-
**Compute.scala** is a Scala library for scientific computing with N-dimensional arrays in parallel on GPU, CPU and other devices. It will be the primary back-end of the incoming [DeepLearning.scala](http://deeplearning.thoughtworks.school/) 3.0, to address performance problems we encountered in DeepLearning.scala 2.0 with [nd4j](http://nd4j.org/).
3+
**Compute.scala** is a Scala library for scientific computing with N-dimensional arrays in parallel on GPU, CPU and other devices. It will be the primary back-end of the incoming [DeepLearning.scala](http://deeplearning.thoughtworks.school/) 3.0, to address performance problems we encountered in DeepLearning.scala 2.0 with [ND4J](http://ND4J.org/).
44

55
* Compute.scala can dynamically merge multiple operators into one kernel program, which runs significantly faster when performing complex computation.
66
* Compute.scala manages data buffers and other native resources in a determinate approach, consuming less memory.
@@ -267,16 +267,81 @@ println(merged.toString)
267267

268268
Generally, when `join`ing *n* `Tensor`s of shape *a*<sub>0</sub> × *a*<sub>1</sub> × *a*<sub>2</sub> ×  ⋯ × *a*<sub>*i*</sub> , the shape of the result `Tensor` is *a*<sub>0</sub> × *a*<sub>1</sub> × *a*<sub>2</sub> ×  ⋯ × *a*<sub>*i*</sub> × *n*
269269

270+
#### Case study: fast matrix multiplication via `split` and `join`
270271

272+
By combining `split` and `join`, you can create complex computation in the following steps:
271273

272-
#### Fast matrix multiplication from `split` and `join`
274+
1. Using `split` to create `Seq`s from some of dimensions of `Tensor`s.
275+
2. Using Scala collection functions to manipulate `Seq`s.
276+
3. Using `join` to merge transformed `Seq` back to `Tensor`.
273277

274-
TODO
278+
For example, you can implement matrix multiplication in this style.
279+
280+
``` scala
281+
def matrixMultiply1(matrix1: Tensor, matrix2: Tensor): Tensor = {
282+
val columns1 = matrix1.split(1)
283+
val columns2 = matrix2.split(1)
284+
val resultColumns = columns2.map { column2: Tensor =>
285+
(columns1 zip column2.split(0))
286+
.map {
287+
case (l: Tensor, r: Tensor) =>
288+
l * r.broadcast(l.shape)
289+
}
290+
.reduce[Tensor](_ + _)
291+
}
292+
Tensor.join(resultColumns)
293+
}
294+
```
295+
296+
You can imagine the Scala collection functions as the code generator of the kernel program, thus the loop running in Scala collection will finally create unrolled loop in the kernel program.
297+
298+
The above `matrixMultiply1` will create a kernel program that contains a unrolled loop of each row and column of `matrix2`. Thus it runs very fast when `matrix1` is big and `matrix2` is small. Our benchmark shows that the above `matrixMultiply1` runs 13 times faster than ND4J's cuBLAS back-end, on a Titan X GPU, when `matrix1` is 65536×8 and `matrix2` is 8×8.
299+
300+
---
301+
302+
You can also create another version of matrix multiplication, which only unrolls the loop of each row of `matrix2`.
303+
304+
``` scala
305+
def matrixMultiply2(matrix1: Tensor, matrix2: Tensor): Tensor = {
306+
val Array(i, j) = matrix1.shape
307+
val Array(`j`, k) = matrix2.shape
308+
val broadcastMatrix1 = matrix1.broadcast(Array(i, j, k))
309+
val broadcastMatrix2 = matrix2.reshape(Array(1, j, k)).broadcast(Array(i, j, k))
310+
val product = broadcastMatrix1 * broadcastMatrix2
311+
product.split(1).reduce[Tensor](_ + _)
312+
}
313+
```
314+
315+
`matrixMultiply2` will run faster than `matrixMultiply1` when `matrix1` is small.
316+
317+
A sophisticated matrix multiplication should dynamically switch the two version of implementation according to matrix size.
318+
319+
``` scala
320+
val UnrollThreshold = 4000
321+
322+
def matrixMultiply(matrix1: Tensor, matrix2: Tensor): Tensor = {
323+
if (matrix1.shape.head >= UnrollThreshold) {
324+
matrixMultiply1(matrix1, matrix2)
325+
} else {
326+
matrixMultiply2(matrix1, matrix2)
327+
}
328+
}
329+
```
330+
331+
The final version of `matrixMultiply` will have good performance for both small and big matrixes.
275332

276333
## Benchmark
277334

278-
* [Compute.scala vs Nd4j on NVIDIA GPU](http://jmh.morethan.io/?source=https://thoughtworksinc.github.io/Compute.scala/benchmarks/nvidia-gpu.json)
279-
* [Compute.scala on AMD GPU](http://jmh.morethan.io/?source=https://thoughtworksinc.github.io/Compute.scala/benchmarks/amd-gpu.json)
335+
* [Compute.scala vs ND4J on a NVIDIA Titan X GPU](http://jmh.morethan.io/?source=https://thoughtworksinc.github.io/Compute.scala/benchmarks/nvidia-gpu.json)
336+
* [Compute.scala on a AMD RX480 GPU](http://jmh.morethan.io/?source=https://thoughtworksinc.github.io/Compute.scala/benchmarks/amd-gpu.json)
337+
338+
Some information can be found in the benchmark result:
339+
340+
* Apparently, Compute.scala supports both NVIDIA GPU and AMD GPU, while ND4J does not support AMD GPU.
341+
* Compute.scala is faster than ND4J on large arrays or complex expressions.
342+
* ND4J is faster than Compute.scala when performing one simple primary operation on very small arrays.
343+
* ND4J's reduced sum is faster than Compute.scala.
344+
* ND4J's `permute` and `broadcast` are extremely slow, causing very low score in the convolution benchmark (unlike this benchmark, Deeplearning4j's convolution operation internally uses some undocumented variant of `permute` and `broadcast` in ND4J, which are not extremely slow).
280345

281346
## Future work
282347

Tensors/src/test/scala/com/thoughtworks/compute/TensorsSpec.scala

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -337,17 +337,17 @@ class TensorsSpec extends AsyncFreeSpec with Matchers {
337337
import tensors._
338338

339339
def matrixMultiply(matrix1: Tensor, matrix2: Tensor): Tensor = {
340-
341340
val columns1 = matrix1.split(1)
342-
343-
Tensor.join(matrix2.split(1).map { column2: Tensor =>
344-
(columns1 zip column2.split(0))
341+
val columns2 = matrix2.split(1)
342+
val resultColumns = columns2.map { column2: Tensor =>
343+
(columns1.view zip column2.split(0))
345344
.map {
346345
case (l: Tensor, r: Tensor) =>
347346
l * r.broadcast(l.shape)
348347
}
349348
.reduce[Tensor](_ + _)
350-
})
349+
}
350+
Tensor.join(resultColumns)
351351
}
352352

353353
matrixMultiply(

0 commit comments

Comments
 (0)