gdeep/allreduce.go at master · kuroko1t/gdeep · GitHub

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
// +build mpi

package gdeep

import (
	//"github.com/kuroko1t/gmpi"
	"C"
	"github.com/kuroko1t/gmat"
)

func InitAllreduce() {
	mpi.Init()
}

func Allreduce(layer []LayerInterface) {
	for _, v := range layer {
		v.allreduce()
	}
	return
}

func Comm_rank() int {
	rank := mpi.Comm_rank(mpi.COMM_WORLD)
	return rank
}

func Comm_size() int {
	size := mpi.Comm_size(mpi.COMM_WORLD)
	return size
}

func performReduce(x gmat.Tensor) gmat.Tensor {
	n, c := gmat.Shape2D(x)
	x1D := gmat.Reshape2D1D(x)
	allreduceSum := make([]float64, n*c)
	mpi.Barrier(mpi.COMM_WORLD)
	mpi.Allreduce(&x1D[0], &allreduceSum[0], n*c,
		mpi.Float64, mpi.SUM, mpi.COMM_WORLD)
	commsize := mpi.Comm_size(mpi.COMM_WORLD)
	allreduceValue := make([]float64, n*c)
	for i := range allreduceSum {
		allreduceValue[i] = allreduceSum[i] / float64(commsize)
	}
	return gmat.Reshape1D2D(allreduceValue, n, c)
}

func (dense *Dense) allreduce() {
	dense.Dw = performReduce(dense.Dw)
	dense.Db = performReduce(dense.Db)
	return
}

func (relu *Relu) allreduce() {
	return
}

func (drop *Dropout) allreduce() {
	return
}

func (softmaxWithLoss *SoftmaxWithLoss) allreduce() {
	return
}