-
Notifications
You must be signed in to change notification settings - Fork 6
/
matrix.go
83 lines (73 loc) · 1.71 KB
/
matrix.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
package functions
import (
"golang.org/x/exp/constraints"
"runtime"
"sync"
)
var numCPU int = runtime.NumCPU()
// matMulParallel runs matrix multiply in parallel by dividing the input rows
func matMulParallel[T constraints.Float](
dst, x, y []T, m, n, p int,
vecMul func(dst, x, y []T, m, n int),
matMul func(dst, x, y []T, m, n, p int),
) {
if m < 4 || m*p*n < 100_000 {
if p == 1 {
vecMul(dst, x, y, m, n)
} else {
matMul(dst, x, y, m, n, p)
}
return
}
rowsPerCPU, rem := m/numCPU, m%numCPU
i := 0
var wg sync.WaitGroup
for c := 0; c < numCPU && i < m; c++ {
numRows := rowsPerCPU
if c < rem {
numRows += 1
}
dstStart := i * p
dstEnd := (i + numRows) * p
xStart := i * n
xEnd := (i + numRows) * n
wg.Add(1)
go func() {
if p == 1 {
vecMul(dst[dstStart:dstEnd], x[xStart:xEnd], y, numRows, n)
} else {
matMul(dst[dstStart:dstEnd], x[xStart:xEnd], y, numRows, n, p)
}
wg.Done()
}()
i += numRows
}
wg.Wait()
}
func MatMul_Go[T constraints.Float](dst, x, y []T, m, n, p int) {
for i := 0; i < m; i++ {
for k := 0; k < n; k++ {
for j := 0; j < p; j++ { // dst not set to zero
dst[i*p+j] += x[i*n+k] * y[k*p+j]
}
}
}
}
func MatMulVec_Go[T constraints.Float](dst, x, y []T, m, n int) {
for i := 0; i < m; i++ {
for k := 0; k < n; k++ { // note: dst is not set to zero
dst[i] += x[i*n+k] * y[k]
}
}
}
func Mat4Mul_Go[T constraints.Float](dst, x, y []T) {
for i := 0; i < 4; i++ {
for j := 0; j < 4; j++ {
dst[i*4+j] = x[i*4]*y[j] + x[i*4+1]*y[1*4+j] +
x[i*4+2]*y[2*4+j] + x[i*4+3]*y[3*4+j]
}
}
}
func MatMul_Parallel_Go[T constraints.Float](dst, x, y []T, m, n, p int) {
matMulParallel(dst, x, y, m, n, p, MatMulVec_Go[T], MatMul_Go[T])
}