-
Notifications
You must be signed in to change notification settings - Fork 9
/
distributor.go
74 lines (65 loc) · 1.66 KB
/
distributor.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
package driver
import (
"github.com/sarchlab/mgpusim/v3/driver/internal"
)
// A distributor can distribute a virtually consecutive memory to multiple GPUs.
type distributor interface {
Distribute(
ctx *Context,
addr, byteSize uint64,
gpuIDs []int,
) (byteAllocatedOnEachGPU []uint64)
}
type distributorImpl struct {
pageSizeAsPowerOf2 uint64
memAllocator internal.MemoryAllocator
}
func newDistributorImpl(memAllocator internal.MemoryAllocator) *distributorImpl {
return &distributorImpl{
memAllocator: memAllocator,
}
}
func (d *distributorImpl) Distribute(
ctx *Context,
addr, byteSize uint64,
gpuIDs []int,
) (byteAllocatedOnEachGPU []uint64) {
pageSize := uint64(1 << d.pageSizeAsPowerOf2)
if addr%pageSize != 0 {
panic("Address much align with pages")
}
byteAllocatedOnEachGPU = make([]uint64, len(gpuIDs))
numPages := (byteSize-1)/pageSize + 1
numGPUs := uint64(len(gpuIDs))
numPagesPerGPU := numPages / numGPUs
numGPUsToUse := uint64(0)
if numPagesPerGPU > 0 {
numGPUsToUse = numPages / numPagesPerGPU
}
if numGPUsToUse > numGPUs {
numGPUsToUse = numGPUs
}
remainingPages := numPages % numGPUs
var i uint64
var lastAllocatedGPU uint64
for i = 0; i < numGPUsToUse; i++ {
d.memAllocator.Remap(
ctx.pid,
addr+i*numPagesPerGPU*pageSize,
numPagesPerGPU*pageSize,
gpuIDs[i],
)
byteAllocatedOnEachGPU[i] += numPagesPerGPU * pageSize
lastAllocatedGPU = i
}
for i := uint64(0); i < remainingPages; i++ {
d.memAllocator.Remap(
ctx.pid,
addr+(numPagesPerGPU*numGPUsToUse+i)*pageSize,
pageSize,
gpuIDs[lastAllocatedGPU],
)
byteAllocatedOnEachGPU[lastAllocatedGPU] += pageSize
}
return byteAllocatedOnEachGPU
}