/
memmap.go
161 lines (140 loc) · 4.44 KB
/
memmap.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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
package axlarod
/*
#include <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <errno.h>
int createAndMapTmpFile(char* fileName, size_t fileSize, void** mappedAddr, int* fd) {
*fd = mkstemp(fileName);
if (*fd < 0) {
return errno;
}
if (ftruncate(*fd, (off_t)fileSize) < 0) {
close(*fd);
return errno;
}
if (unlink(fileName)) {
return errno;
}
*mappedAddr = mmap(NULL, fileSize, PROT_READ | PROT_WRITE, MAP_SHARED, *fd, 0);
if (*mappedAddr == MAP_FAILED) {
close(*fd);
return errno;
}
return 0; // Success
}
*/
import "C"
import (
"fmt"
"math/rand"
"os"
"strings"
"unsafe"
)
// MemMapFile represents a memory mapped file.
type MemMapFile struct {
MemoryAddress unsafe.Pointer
File *os.File
Size uint
FilePattern string
UsePitch0Size bool
}
// MemMapConfiguration represents the configuration for memory mapped files.
type MemMapConfiguration struct {
InputTmpMapFiles map[int]*MemMapFile
OutputTmpMapFiles map[int]*MemMapFile
}
// configureMemMapFile configures the memory mapped file for the tensor.
func (model *LarodModel) configureMemMapFile(file_map map[int]*MemMapFile, tensors []*LarodTensor, pitches *LarodTensorPitches) error {
var err error
for i, f := range file_map {
// Create file other wise reuse fd
if f.File == nil {
f.FilePattern = generateRandomMapFilePattern(fmt.Sprintf("%s-in", model.Name), i)
if f.UsePitch0Size {
f.Size = pitches.Pitches[0]
}
f.MemoryAddress, f.File, err = CreateAndMapTmpFile(f.FilePattern, f.Size)
if err != nil {
return err
}
}
tensors[i].SetTensorFd(f.File.Fd())
tensors[i].MemMapFile = f
}
return nil
}
// MapModelTmpFiles maps the temporary files for the model.
func (model *LarodModel) MapModelTmpFiles(m *MemMapConfiguration) error {
if err := model.configureMemMapFile(m.InputTmpMapFiles, model.Inputs, model.InputPitches); err != nil {
return err
}
if err := model.configureMemMapFile(m.OutputTmpMapFiles, model.Outputs, model.OutputPitches); err != nil {
return err
}
return nil
}
// UnmapMemory unmaps the memory mapped file
func (t *MemMapFile) UnmapMemory() error {
if ret := C.munmap(t.MemoryAddress, C.size_t(t.Size)); ret != 0 {
return fmt.Errorf("failed to unmap memory: return code %d", int(ret))
}
return nil
}
// Seek the memory mapped file to the beginning
func (t *MemMapFile) Rewind() error {
_, err := t.File.Seek(0, 0)
return err
}
// CreateAndMapTmpFile creates a temporary file and maps it to memory
func CreateAndMapTmpFile(fileNamePattern string, fileSize uint) (unsafe.Pointer, *os.File, error) {
cFileName := C.CString(fileNamePattern)
defer C.free(unsafe.Pointer(cFileName))
var mappedAddr unsafe.Pointer
var fd C.int
errCode := C.createAndMapTmpFile(cFileName, C.size_t(fileSize), &mappedAddr, &fd)
if errCode != 0 {
return nil, nil, fmt.Errorf("error creating and mapping file: %s", C.GoString(C.strerror(errCode)))
}
file := os.NewFile(uintptr(fd), C.GoString(cFileName))
return mappedAddr, file, nil
}
// CopyDataToMappedMemory copies data to the mapped memory
func CopyDataToMappedMemory(mappedAddr unsafe.Pointer, data []byte) error {
dataLen := len(data)
if dataLen == 0 {
return fmt.Errorf("data slice is empty")
}
C.memcpy(mappedAddr, unsafe.Pointer(&data[0]), C.size_t(dataLen))
return nil
}
// CopyDataFromMappedMemory copies data from the mapped memory
func CopyDataFromMappedMemory(mappedAddr unsafe.Pointer, size int) ([]byte, error) {
if size <= 0 {
return nil, fmt.Errorf("size must be positive")
}
dataSlice := unsafe.Slice((*byte)(mappedAddr), size)
copiedData := make([]byte, size)
copy(copiedData, dataSlice)
return copiedData, nil
}
// randomString generates a random string of a specified length using the provided character set.
func randomString(length int, charset string) string {
var output strings.Builder
charsetLength := len(charset)
for i := 0; i < length; i++ {
randomIndex := rand.Intn(charsetLength)
output.WriteByte(charset[randomIndex])
}
return output.String()
}
// generatePath creates a path string with a random 10 character string.
func generateRandomMapFilePattern(prefix string, tensor_io_index int) string {
const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
randomPart := randomString(10, charset)
return fmt.Sprintf("/tmp/larod-%s-%d.%s-XXXXXX", prefix, tensor_io_index, randomPart)
}