mmap file.

b4rtaz · May 27, 2024 · 1248d2b · 1248d2b
1 parent 596121e
commit 1248d2b
Show file tree

Hide file tree

Showing 3 changed files with 81 additions and 72 deletions.
diff --git a/src/transformer.cpp b/src/transformer.cpp
@@ -2,9 +2,7 @@
 #include <cmath>
 #include <cassert>
 #include <stdexcept>
-#include <fcntl.h>
 #include <string.h>
-#include <unistd.h>
 #include <ostream>
 #include <iostream>
 #include "funcs.hpp"
@@ -15,15 +13,6 @@
 #define ALLOC_WEIGHTS true
 #define IS_ROOT_SLICE(sliceIndex) (sliceIndex == 0)
 
-#ifdef _WIN32
-#include <windows.h>
-#include <io.h>
-#define fseek _fseeki64
-#define ftell _ftelli64
-#else
-#include <sys/mman.h>
-#endif
-
 RowMatmulSlice::RowMatmulSlice(FloatType type, int nSlices, int n, int d) {
     assert(d % nSlices == 0);
 
@@ -613,74 +602,19 @@ static size_t readSlicedMatmulWeights(MatmulSlice* slice, char* weights0, Socket
 }
 
 Transformer Transformer::loadRootFromFile(const char* path, TransformerSpec* spec, SocketPool* socketPool) {
-#ifdef _WIN32
-    HANDLE hFile = CreateFileA(path, GENERIC_READ, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
-    if (hFile == INVALID_HANDLE_VALUE) {
-        printf("Cannot open file %s\n", path);
-        exit(EXIT_FAILURE);
-    }
-
-    HANDLE hMapping = CreateFileMappingA(hFile, NULL, PAGE_READONLY, 0, 0, NULL);
-    if (hMapping == NULL) {
-        printf("CreateFileMappingA failed, error: %lu\n", GetLastError());
-        CloseHandle(hFile);
-        exit(EXIT_FAILURE);
-    }
-
-    char* data = (char*)MapViewOfFile(hMapping, FILE_MAP_READ, 0, 0, 0);
-    if (data == NULL) {
-        printf("MapViewOfFile failed!\n");
-        CloseHandle(hMapping);
-        CloseHandle(hFile);
-        exit(EXIT_FAILURE);
-    }
-
-    std::cout << "MapViewOfFile succeeded. data = " << static_cast<void*>(data) << std::endl;
-
-    char* weights = data + spec->headerSize;
-    std::cout << "weights = " << static_cast<void*>(weights) << std::endl;
-
-    Transformer transformer = Transformer::loadRoot(weights, spec, socketPool);
-
-#if ALLOC_WEIGHTS
-    UnmapViewOfFile(data);
-    CloseHandle(hMapping);
-    CloseHandle(hFile);
-#else
-    // TODO: handler should be released in destructor
-#endif
+    MmapFile file;
+    openMmapFile(&file, path, spec->fileSize);
 
-    return transformer;
-#else
-    int fd = open(path, O_RDONLY);
-    if (fd == -1) {
-        printf("Cannot open file %s\n", path);
-        exit(EXIT_FAILURE);
-    }
-
-    char* data = (char*)mmap(NULL, spec->fileSize, PROT_READ, MAP_PRIVATE, fd, 0);
-    if (data == MAP_FAILED) {
-        printf("Mmap failed!\n");
-        close(fd);
-        exit(EXIT_FAILURE);
-    }
-
-    std::cout << "mmap succeeded. data = " << static_cast<void*>(data) << std::endl;
-
-    char* weights = data + spec->headerSize;
-    std::cout << "weights = " << static_cast<void*>(weights) << std::endl;
-
-    Transformer transformer = Transformer::loadRoot(weights, spec, socketPool);
+    char* weights = ((char*)file.data) + spec->headerSize;
+    Transformer transformer = Transformer::loadRoot((char*)weights, spec, socketPool);
 
 #if ALLOC_WEIGHTS
-    munmap(data, spec->fileSize);
-    close(fd);
+    closeMmapFile(&file);
 #else
     // TODO: handler should be released in destructor
 #endif
 
     return transformer;
-#endif
 }
 
 Transformer Transformer::loadRoot(char* data, TransformerSpec* spec, SocketPool* socketPool) {

diff --git a/src/utils.cpp b/src/utils.cpp
@@ -8,10 +8,16 @@
 
 #define BUFFER_ALIGNMENT 16
 
-#ifdef _WIN32 
+#ifdef _WIN32
+#include <windows.h>
 #include <malloc.h>
+#include <io.h>
+#define fseek _fseeki64
+#define ftell _ftelli64
 #else
 #include <sys/mman.h>
+#include <fcntl.h>
+#include <unistd.h>
 #endif
 
 void* newBuffer(size_t size) {
@@ -61,6 +67,56 @@ float randomF32(unsigned long long *state) {
     return (randomU32(state) >> 8) / 16777216.0f;
 }
 
+void openMmapFile(MmapFile* file, const char* path, size_t size) {
+    file->size = size;
+#ifdef _WIN32
+    file->hFile = CreateFileA(path, GENERIC_READ, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
+    if (file->hFile == INVALID_HANDLE_VALUE) {
+        printf("Cannot open file %s\n", path);
+        exit(EXIT_FAILURE);
+    }
+
+    file->hMapping = CreateFileMappingA(file->hFile, NULL, PAGE_READONLY, 0, 0, NULL);
+    if (file->hMapping == NULL) {
+        printf("CreateFileMappingA failed, error: %lu\n", GetLastError());
+        CloseHandle(file->hFile);
+        exit(EXIT_FAILURE);
+    }
+
+    file->data = (char*)MapViewOfFile(file->hMapping, FILE_MAP_READ, 0, 0, 0);
+    if (file->data == NULL) {
+        printf("MapViewOfFile failed!\n");
+        CloseHandle(file->hMapping);
+        CloseHandle(file->hFile);
+        exit(EXIT_FAILURE);
+    }
+#else
+    file->fd = open(path, O_RDONLY);
+    if (file->fd == -1) {
+        printf("Cannot open file %s\n", path);
+        exit(EXIT_FAILURE);
+    }
+
+    file->data = mmap(NULL, size, PROT_READ, MAP_PRIVATE, file->fd, 0);
+    if (file->data == MAP_FAILED) {
+        printf("Mmap failed!\n");
+        close(file->fd);
+        exit(EXIT_FAILURE);
+    }
+#endif
+}
+
+void closeMmapFile(MmapFile* file) {
+#ifdef _WIN32
+    UnmapViewOfFile(file->data);
+    CloseHandle(file->hMapping);
+    CloseHandle(file->hFile);
+#else
+    munmap(file->data, file->size);
+    close(file->fd);
+#endif
+}
+
 TaskLoop::TaskLoop(unsigned int nThreads, unsigned int nTasks, unsigned int nTypes, TaskLoopTask* tasks, void* userData) {
     this->nThreads = nThreads;
     this->nTasks = nTasks;

diff --git a/src/utils.hpp b/src/utils.hpp
@@ -4,15 +4,34 @@
 #include <atomic>
 #include "common/pthread.h"
 
+#ifdef _WIN32
+#include <windows.h>
+#endif
+
 #define NEW_BUFFER(size) (char*)newBuffer(size)
 #define FREE_BUFFER(buffer) freeBuffer(buffer)
 
 void* newBuffer(size_t size);
 void freeBuffer(void* buffer);
+
 unsigned long timeMs();
 unsigned int randomU32(unsigned long long *state);
 float randomF32(unsigned long long *state);
 
+struct MmapFile {
+    void* data;
+    size_t size;
+#ifdef _WIN32
+    HANDLE hFile;
+    HANDLE hMapping;
+#else
+    int fd;
+#endif
+};
+
+void openMmapFile(MmapFile* file, const char* path, size_t size);
+void closeMmapFile(MmapFile* file);
+
 typedef void (TaskLoopHandler)(unsigned int nThreads, unsigned int threadIndex, void* userData);
 typedef struct {
     TaskLoopHandler* handler;