└── llama-mmap.cpp


/llama-mmap.cpp:
--------------------------------------------------------------------------------
  1 | /*************************************************************************
  2 |     > File Name:     llama-mmap.cpp（基于llama.cpp b4821版本）
  3 |     > Author:        腾讯玄武实验室
  4 |     > Created Time:  2025年03月17日
  5 |     > Description:   llama.cpp部署deepseek优化；使模型加载时使用1G大页，
  6 |                      减少访存时的tlb miss，提高token生成速度。
  7 |     > Version:       1.0
  8 |  ************************************************************************/
  9 | 
 10 | #include "llama-mmap.h"
 11 | 
 12 | #include "llama-impl.h"
 13 | 
 14 | #include "ggml.h"
 15 | 
 16 | #include <cstring>
 17 | #include <climits>
 18 | #include <stdexcept>
 19 | #include <cerrno>
 20 | 
 21 | #ifdef __has_include
 22 |     #if __has_include(<unistd.h>)
 23 |         #include <unistd.h>
 24 |         #if defined(_POSIX_MAPPED_FILES)
 25 |             #include <sys/mman.h>
 26 |             #include <fcntl.h>
 27 |         #endif
 28 |         #if defined(_POSIX_MEMLOCK_RANGE)
 29 |             #include <sys/resource.h>
 30 |         #endif
 31 |     #endif
 32 | #endif
 33 | 
 34 | #if defined(_WIN32)
 35 |     #define WIN32_LEAN_AND_MEAN
 36 |     #ifndef NOMINMAX
 37 |         #define NOMINMAX
 38 |     #endif
 39 |     #include <windows.h>
 40 |     #ifndef PATH_MAX
 41 |         #define PATH_MAX MAX_PATH
 42 |     #endif
 43 |     #include <io.h>
 44 | #endif
 45 | 
 46 | // TODO: consider moving to llama-impl.h if needed in more places
 47 | #if defined(_WIN32)
 48 | static std::string llama_format_win_err(DWORD err) {
 49 |     LPSTR buf;
 50 |     size_t size = FormatMessageA(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
 51 |                                  NULL, err, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), (LPSTR)&buf, 0, NULL);
 52 |     if (!size) {
 53 |         return "FormatMessageA failed";
 54 |     }
 55 |     std::string ret(buf, size);
 56 |     LocalFree(buf);
 57 |     return ret;
 58 | }
 59 | #endif
 60 | 
 61 | // llama_file
 62 | 
 63 | struct llama_file::impl {
 64 | #if defined(_WIN32)
 65 |     HANDLE fp_win32;
 66 |     std::string GetErrorMessageWin32(DWORD error_code) const {
 67 |         std::string ret;
 68 |         LPSTR lpMsgBuf = NULL;
 69 |         DWORD bufLen = FormatMessageA(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
 70 |                                     NULL, error_code, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), (LPSTR)&lpMsgBuf, 0, NULL);
 71 |         if (!bufLen) {
 72 |             ret = format("Win32 error code: %lx", error_code);
 73 |         } else {
 74 |             ret = lpMsgBuf;
 75 |             LocalFree(lpMsgBuf);
 76 |         }
 77 | 
 78 |         return ret;
 79 |     }
 80 | 
 81 |     impl(const char * fname, const char * mode) {
 82 |         fp = ggml_fopen(fname, mode);
 83 |         if (fp == NULL) {
 84 |             throw std::runtime_error(format("failed to open %s: %s", fname, strerror(errno)));
 85 |         }
 86 |         fp_win32 = (HANDLE) _get_osfhandle(_fileno(fp));
 87 |         seek(0, SEEK_END);
 88 |         size = tell();
 89 |         seek(0, SEEK_SET);
 90 |     }
 91 | 
 92 |     size_t tell() const {
 93 |         LARGE_INTEGER li;
 94 |         li.QuadPart = 0;
 95 |         BOOL ret = SetFilePointerEx(fp_win32, li, &li, FILE_CURRENT);
 96 |         if (!ret) {
 97 |             throw std::runtime_error(format("read error: %s", GetErrorMessageWin32(GetLastError()).c_str()));
 98 |         }
 99 | 
100 |         return li.QuadPart;
101 |     }
102 | 
103 |     void seek(size_t offset, int whence) const {
104 |         static_assert(SEEK_SET == FILE_BEGIN, "SEEK_SET != FILE_BEGIN");
105 |         static_assert(SEEK_CUR == FILE_CURRENT, "SEEK_CUR != FILE_CURRENT");
106 |         static_assert(SEEK_END == FILE_END, "SEEK_END != FILE_END");
107 | 
108 |         LARGE_INTEGER li;
109 |         li.QuadPart = offset;
110 |         BOOL ret = SetFilePointerEx(fp_win32, li, NULL, whence);
111 |         if (!ret) {
112 |             throw std::runtime_error(format("read error: %s", GetErrorMessageWin32(GetLastError()).c_str()));
113 |         }
114 |     }
115 | 
116 |     void read_raw(void * ptr, size_t len) const {
117 |         size_t bytes_read = 0;
118 |         while (bytes_read < len) {
119 |             size_t chunk_size = std::min<size_t>(len - bytes_read, 64*1024*1024);
120 |             DWORD chunk_read = 0;
121 |             BOOL result = ReadFile(fp_win32, reinterpret_cast<char*>(ptr) + bytes_read, chunk_size, &chunk_read, NULL);
122 |             if (!result) {
123 |                 throw std::runtime_error(format("read error: %s", GetErrorMessageWin32(GetLastError()).c_str()));
124 |             }
125 |             if (chunk_read < chunk_size || chunk_read == 0) {
126 |                 throw std::runtime_error("unexpectedly reached end of file");
127 |             }
128 | 
129 |             bytes_read += chunk_read;
130 |         }
131 |     }
132 | 
133 |     uint32_t read_u32() const {
134 |         uint32_t val;
135 |         read_raw(&val, sizeof(val));
136 |         return val;
137 |     }
138 | 
139 |     void write_raw(const void * ptr, size_t len) const {
140 |         size_t bytes_written = 0;
141 |         while (bytes_written < len) {
142 |             size_t chunk_size = std::min<size_t>(len - bytes_written, 64*1024*1024);
143 |             DWORD chunk_written = 0;
144 |             BOOL result = WriteFile(fp_win32, reinterpret_cast<char const*>(ptr) + bytes_written, chunk_size, &chunk_written, NULL);
145 |             if (!result) {
146 |                 throw std::runtime_error(format("write error: %s", GetErrorMessageWin32(GetLastError()).c_str()));
147 |             }
148 |             if (chunk_written < chunk_size || chunk_written == 0) {
149 |                 throw std::runtime_error("unexpectedly failed to write bytes");
150 |             }
151 | 
152 |             bytes_written += chunk_written;
153 |         }
154 |     }
155 | 
156 |     void write_u32(uint32_t val) const {
157 |         write_raw(&val, sizeof(val));
158 |     }
159 | 
160 |     ~impl() {
161 |         if (fp) {
162 |             std::fclose(fp);
163 |         }
164 |     }
165 | #else
166 |     impl(const char * fname, const char * mode) {
167 |         fp = ggml_fopen(fname, mode);
168 |         if (fp == NULL) {
169 |             throw std::runtime_error(format("failed to open %s: %s", fname, strerror(errno)));
170 |         }
171 |         seek(0, SEEK_END);
172 |         size = tell();
173 |         seek(0, SEEK_SET);
174 |     }
175 | 
176 |     size_t tell() const {
177 | // TODO: this ifdef is never true?
178 | #ifdef _WIN32
179 |         __int64 ret = _ftelli64(fp);
180 | #else
181 |         long ret = std::ftell(fp);
182 | #endif
183 |         if (ret == -1) {
184 |             throw std::runtime_error(format("ftell error: %s", strerror(errno)));
185 |         }
186 | 
187 |         return (size_t) ret;
188 |     }
189 | 
190 |     void seek(size_t offset, int whence) const {
191 | // TODO: this ifdef is never true?
192 | #ifdef _WIN32
193 |         int ret = _fseeki64(fp, (__int64) offset, whence);
194 | #else
195 |         int ret = std::fseek(fp, (long) offset, whence);
196 | #endif
197 |         if (ret != 0) {
198 |             throw std::runtime_error(format("seek error: %s", strerror(errno)));
199 |         }
200 |     }
201 | 
202 |     void read_raw(void * ptr, size_t len) const {
203 |         if (len == 0) {
204 |             return;
205 |         }
206 |         errno = 0;
207 |         std::size_t ret = std::fread(ptr, len, 1, fp);
208 |         if (ferror(fp)) {
209 |             throw std::runtime_error(format("read error: %s", strerror(errno)));
210 |         }
211 |         if (ret != 1) {
212 |             throw std::runtime_error("unexpectedly reached end of file");
213 |         }
214 |     }
215 | 
216 |     uint32_t read_u32() const {
217 |         uint32_t ret;
218 |         read_raw(&ret, sizeof(ret));
219 |         return ret;
220 |     }
221 | 
222 |     void write_raw(const void * ptr, size_t len) const {
223 |         if (len == 0) {
224 |             return;
225 |         }
226 |         errno = 0;
227 |         size_t ret = std::fwrite(ptr, len, 1, fp);
228 |         if (ret != 1) {
229 |             throw std::runtime_error(format("write error: %s", strerror(errno)));
230 |         }
231 |     }
232 | 
233 |     void write_u32(uint32_t val) const {
234 |         write_raw(&val, sizeof(val));
235 |     }
236 | 
237 |     ~impl() {
238 |         if (fp) {
239 |             std::fclose(fp);
240 |         }
241 |     }
242 | #endif
243 | 
244 |     FILE * fp;
245 |     size_t size;
246 | };
247 | 
248 | llama_file::llama_file(const char * fname, const char * mode) : pimpl(std::make_unique<impl>(fname, mode)) {}
249 | llama_file::~llama_file() = default;
250 | 
251 | size_t llama_file::tell() const { return pimpl->tell(); }
252 | size_t llama_file::size() const { return pimpl->size; }
253 | 
254 | int llama_file::file_id() const {
255 | #ifdef _WIN32
256 |     return _fileno(pimpl->fp);
257 | #else
258 | #if defined(fileno)
259 |     return fileno(pimpl->fp);
260 | #else
261 |     return ::fileno(pimpl->fp);
262 | #endif
263 | #endif
264 | }
265 | 
266 | void llama_file::seek(size_t offset, int whence) const { pimpl->seek(offset, whence); }
267 | void llama_file::read_raw(void * ptr, size_t len) const { pimpl->read_raw(ptr, len); }
268 | 
269 | uint32_t llama_file::read_u32() const { return pimpl->read_u32(); }
270 | 
271 | void llama_file::write_raw(const void * ptr, size_t len) const { pimpl->write_raw(ptr, len); }
272 | void llama_file::write_u32(uint32_t val) const { pimpl->write_u32(val); }
273 | 
274 | // llama_mmap
275 | 
276 | struct llama_mmap::impl {
277 | #ifdef _POSIX_MAPPED_FILES
278 |     std::vector<std::pair<size_t, size_t>> mapped_fragments;
279 |     bool using_hugepages;
280 | 
281 |     impl(struct llama_file * file, size_t prefetch, bool numa) {
282 |         size = file->size();
283 |         int fd = file->file_id();
284 |         int flags = MAP_SHARED;
285 |         if (numa) { prefetch = 0; }
286 | #ifdef __linux__
287 |         if (posix_fadvise(fd, 0, 0, POSIX_FADV_SEQUENTIAL)) {
288 |             LLAMA_LOG_WARN("warning: posix_fadvise(.., POSIX_FADV_SEQUENTIAL) failed: %s\n",
289 |                     strerror(errno));
290 |         }
291 |         if (prefetch) { flags |= MAP_POPULATE; }
292 | #endif
293 |         //addr = mmap(NULL, file->size(), PROT_READ, flags, fd, 0);
294 |         //if (addr == MAP_FAILED) {
295 |         //    throw std::runtime_error(format("mmap failed: %s", strerror(errno)));
296 |         //}
297 | 
298 | 	// 尝试直接分配 1GB 大页匿名内存
299 | 	const size_t huge_page_size = 1 * 1024 * 1024 * 1024;
300 | 	size_t aligned_size = (file->size() + huge_page_size - 1) & ~(huge_page_size - 1);
301 | 
302 | 	void* huge_addr = mmap(
303 | 	    /*addr=*/nullptr,
304 | 	    aligned_size,
305 | 	    PROT_READ | PROT_WRITE,
306 | 	    MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB,
307 | 	    /*fd=*/-1,
308 | 	    /*offset=*/0
309 | 	);
310 | 
311 | 	if (huge_addr == MAP_FAILED) {
312 | 	    // 如果大页分配失败，回退到普通 mmap
313 | 	    LLAMA_LOG_WARN("huge page (%lx) allocation failed, using regular pages: %s\n",
314 | 			   aligned_size, strerror(errno));
315 | 	    addr = mmap(
316 | 		/*addr=*/nullptr,
317 | 		file->size(),
318 | 		PROT_READ,
319 | 		flags,
320 | 		fd,
321 | 		/*offset=*/0
322 | 	    );
323 | 	    if (addr == MAP_FAILED) {
324 | 		throw std::runtime_error(format("mmap failed: %s", strerror(errno)));
325 | 	    }
326 | 	    using_hugepages = false;
327 | 	} else {
328 | 	    using_hugepages = true;
329 | 	    // 成功分配大页后，直接将文件内容读取到大页区域，避免先普通 mmap 再 memcpy
330 | 	    size_t total_read = 0;
331 | 	    while (total_read < file->size()) {
332 | 		ssize_t ret = pread(
333 | 		    fd,
334 | 		    static_cast<char*>(huge_addr) + total_read,
335 | 		    file->size() - total_read,
336 | 		    total_read
337 | 		);
338 | 		if (ret < 0) {
339 | 		    munmap(huge_addr, aligned_size);
340 | 		    throw std::runtime_error(format("read to huge pages failed: %s", strerror(errno)));
341 | 		}
342 | 		if (ret == 0) {
343 | 		    // 意外的 EOF
344 | 		    break;
345 | 		}
346 | 		total_read += static_cast<size_t>(ret);
347 | 	    }
348 | 	    addr = huge_addr;
349 | 	}
350 | 
351 | 
352 |         if (prefetch > 0) {
353 |             if (posix_madvise(addr, std::min(file->size(), prefetch), POSIX_MADV_WILLNEED)) {
354 |                 LLAMA_LOG_WARN("warning: posix_madvise(.., POSIX_MADV_WILLNEED) failed: %s\n",
355 |                         strerror(errno));
356 |             }
357 |         }
358 |         if (numa) {
359 |             if (posix_madvise(addr, file->size(), POSIX_MADV_RANDOM)) {
360 |                 LLAMA_LOG_WARN("warning: posix_madvise(.., POSIX_MADV_RANDOM) failed: %s\n",
361 |                         strerror(errno));
362 |             }
363 |         }
364 | 
365 |         mapped_fragments.emplace_back(0,aligned_size);
366 |     }
367 | 
368 |     static void align_range(size_t * first, size_t * last, size_t page_size) {
369 |         size_t offset_in_page = *first & (page_size - 1);
370 |         size_t offset_to_page = offset_in_page == 0 ? 0 : page_size - offset_in_page;
371 |         *first += offset_to_page;
372 | 
373 |         *last = *last & ~(page_size - 1);
374 | 
375 |         if (*last <= *first) {
376 |             *last = *first;
377 |         }
378 |     }
379 | 
380 |     void unmap_fragment(size_t first, size_t last) {
381 | 	if (using_hugepages) {
382 |             // 使用大页时，禁用部分解除映射
383 |        	    return;
384 |     	}
385 |         int page_size = sysconf(_SC_PAGESIZE);
386 |         align_range(&first, &last, page_size);
387 |         size_t len = last - first;
388 | 
389 |         if (len == 0) {
390 |             return;
391 |         }
392 | 
393 |         GGML_ASSERT(first % page_size == 0);
394 |         GGML_ASSERT(last % page_size == 0);
395 |         GGML_ASSERT(last > first);
396 | 
397 |         void * next_page_start = (uint8_t *) addr + first;
398 | 
399 |         if (munmap(next_page_start, len)) {
400 |             LLAMA_LOG_WARN("warning: munmap failed: %s\n", strerror(errno));
401 |         }
402 | 
403 |         std::vector<std::pair<size_t, size_t>> new_mapped_fragments;
404 |         for (const auto & frag : mapped_fragments) {
405 |             if (frag.first < first && frag.second > last) {
406 |                 new_mapped_fragments.emplace_back(frag.first, first);
407 |                 new_mapped_fragments.emplace_back(last, frag.second);
408 |             } else if (frag.first < first && frag.second > first) {
409 |                 new_mapped_fragments.emplace_back(frag.first, first);
410 |             } else if (frag.first < last && frag.second > last) {
411 |                 new_mapped_fragments.emplace_back(last, frag.second);
412 |             } else if (frag.first >= first && frag.second <= last) {
413 |             } else {
414 |                 new_mapped_fragments.push_back(frag);
415 |             }
416 |         }
417 |         mapped_fragments = std::move(new_mapped_fragments);
418 |     }
419 | 
420 |     ~impl() {
421 |         for (const auto & frag : mapped_fragments) {
422 |             if (munmap((char *) addr + frag.first, frag.second - frag.first)) {
423 |                 LLAMA_LOG_WARN("warning: munmap failed: %s\n", strerror(errno));
424 |             }
425 |         }
426 |     }
427 | #elif defined(_WIN32)
428 |     impl(struct llama_file * file, size_t prefetch, bool numa) {
429 |         GGML_UNUSED(numa);
430 | 
431 |         size = file->size();
432 | 
433 |         HANDLE hFile = (HANDLE) _get_osfhandle(file->file_id());
434 | 
435 |         HANDLE hMapping = CreateFileMappingA(hFile, NULL, PAGE_READONLY, 0, 0, NULL);
436 | 
437 |         if (hMapping == NULL) {
438 |             DWORD error = GetLastError();
439 |             throw std::runtime_error(format("CreateFileMappingA failed: %s", llama_format_win_err(error).c_str()));
440 |         }
441 | 
442 |         addr = MapViewOfFile(hMapping, FILE_MAP_READ, 0, 0, 0);
443 |         DWORD error = GetLastError();
444 |         CloseHandle(hMapping);
445 | 
446 |         if (addr == NULL) {
447 |             throw std::runtime_error(format("MapViewOfFile failed: %s", llama_format_win_err(error).c_str()));
448 |         }
449 | 
450 |         if (prefetch > 0) {
451 | #if _WIN32_WINNT >= 0x602
452 |             BOOL (WINAPI *pPrefetchVirtualMemory) (HANDLE, ULONG_PTR, PWIN32_MEMORY_RANGE_ENTRY, ULONG);
453 |             HMODULE hKernel32 = GetModuleHandleW(L"kernel32.dll");
454 | 
455 |             pPrefetchVirtualMemory = (decltype(pPrefetchVirtualMemory))(void *) GetProcAddress(hKernel32, "PrefetchVirtualMemory");
456 | 
457 |             if (pPrefetchVirtualMemory) {
458 |                 WIN32_MEMORY_RANGE_ENTRY range;
459 |                 range.VirtualAddress = addr;
460 |                 range.NumberOfBytes = (SIZE_T) std::min(size, prefetch);
461 |                 if (!pPrefetchVirtualMemory(GetCurrentProcess(), 1, &range, 0)) {
462 |                     LLAMA_LOG_WARN("warning: PrefetchVirtualMemory failed: %s\n",
463 |                             llama_format_win_err(GetLastError()).c_str());
464 |                 }
465 |             }
466 | #else
467 |             throw std::runtime_error("PrefetchVirtualMemory unavailable");
468 | #endif
469 |         }
470 |     }
471 | 
472 |     void unmap_fragment(size_t first, size_t last) {
473 |         GGML_UNUSED(first);
474 |         GGML_UNUSED(last);
475 |     }
476 | 
477 |     ~impl() {
478 |         if (!UnmapViewOfFile(addr)) {
479 |             LLAMA_LOG_WARN("warning: UnmapViewOfFile failed: %s\n",
480 |                     llama_format_win_err(GetLastError()).c_str());
481 |         }
482 |     }
483 | #else
484 |     impl(struct llama_file * file, size_t prefetch, bool numa) {
485 |         GGML_UNUSED(file);
486 |         GGML_UNUSED(prefetch);
487 |         GGML_UNUSED(numa);
488 | 
489 |         throw std::runtime_error("mmap not supported");
490 |     }
491 | 
492 |     void unmap_fragment(size_t first, size_t last) {
493 |         GGML_UNUSED(first);
494 |         GGML_UNUSED(last);
495 | 
496 |         throw std::runtime_error("mmap not supported");
497 |     }
498 | #endif
499 | 
500 |     void * addr;
501 |     size_t size;
502 | };
503 | 
504 | llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa) : pimpl(std::make_unique<impl>(file, prefetch, numa)) {}
505 | llama_mmap::~llama_mmap() = default;
506 | 
507 | size_t llama_mmap::size() const { return pimpl->size; }
508 | void * llama_mmap::addr() const { return pimpl->addr; }
509 | 
510 | void llama_mmap::unmap_fragment(size_t first, size_t last) { pimpl->unmap_fragment(first, last); }
511 | 
512 | #if defined(_POSIX_MEMLOCK_RANGE) || defined(_WIN32)
513 | const bool llama_mmap::SUPPORTED  = true;
514 | #else
515 | const bool llama_mmap::SUPPORTED  = false;
516 | #endif
517 | 
518 | // llama_mlock
519 | 
520 | struct llama_mlock::impl {
521 | #ifdef _POSIX_MEMLOCK_RANGE
522 |     static size_t lock_granularity() {
523 |         return (size_t) sysconf(_SC_PAGESIZE);
524 |     }
525 | 
526 |     bool raw_lock(const void * addr, size_t size) const {
527 |         if (!mlock(addr, size)) {
528 |             return true;
529 |         }
530 | 
531 | #ifdef __APPLE__
532 | #define MLOCK_SUGGESTION \
533 |         "Try increasing the sysctl values 'vm.user_wire_limit' and 'vm.global_user_wire_limit' and/or " \
534 |         "decreasing 'vm.global_no_user_wire_amount'.  Also try increasing RLIMIT_MEMLOCK (ulimit -l).\n"
535 | #else
536 | #define MLOCK_SUGGESTION \
537 |         "Try increasing RLIMIT_MEMLOCK ('ulimit -l' as root).\n"
538 | #endif
539 | 
540 |         char* errmsg = std::strerror(errno);
541 |         bool suggest = (errno == ENOMEM);
542 | 
543 |         struct rlimit lock_limit;
544 |         if (suggest && getrlimit(RLIMIT_MEMLOCK, &lock_limit)) {
545 |             suggest = false;
546 |         }
547 |         if (suggest && (lock_limit.rlim_max > lock_limit.rlim_cur + size)) {
548 |             suggest = false;
549 |         }
550 | 
551 |         LLAMA_LOG_WARN("warning: failed to mlock %zu-byte buffer (after previously locking %zu bytes): %s\n%s",
552 |                 size, this->size, errmsg, suggest ? MLOCK_SUGGESTION : "");
553 |         return false;
554 |     }
555 | 
556 |     static void raw_unlock(void * addr, size_t size) {
557 |         if (munlock(addr, size)) {
558 |             LLAMA_LOG_WARN("warning: failed to munlock buffer: %s\n", std::strerror(errno));
559 |         }
560 |     }
561 | #elif defined(_WIN32)
562 |     static size_t lock_granularity() {
563 |         SYSTEM_INFO si;
564 |         GetSystemInfo(&si);
565 |         return (size_t) si.dwPageSize;
566 |     }
567 | 
568 |     bool raw_lock(void * ptr, size_t len) const {
569 |         for (int tries = 1; ; tries++) {
570 |             if (VirtualLock(ptr, len)) {
571 |                 return true;
572 |             }
573 |             if (tries == 2) {
574 |                 LLAMA_LOG_WARN("warning: failed to VirtualLock %zu-byte buffer (after previously locking %zu bytes): %s\n",
575 |                     len, size, llama_format_win_err(GetLastError()).c_str());
576 |                 return false;
577 |             }
578 | 
579 |             SIZE_T min_ws_size, max_ws_size;
580 |             if (!GetProcessWorkingSetSize(GetCurrentProcess(), &min_ws_size, &max_ws_size)) {
581 |                 LLAMA_LOG_WARN("warning: GetProcessWorkingSetSize failed: %s\n",
582 |                         llama_format_win_err(GetLastError()).c_str());
583 |                 return false;
584 |             }
585 |             size_t increment = len + 1048576;
586 |             min_ws_size += increment;
587 |             max_ws_size += increment;
588 |             if (!SetProcessWorkingSetSize(GetCurrentProcess(), min_ws_size, max_ws_size)) {
589 |                 LLAMA_LOG_WARN("warning: SetProcessWorkingSetSize failed: %s\n",
590 |                         llama_format_win_err(GetLastError()).c_str());
591 |                 return false;
592 |             }
593 |         }
594 |     }
595 | 
596 |     static void raw_unlock(void * ptr, size_t len) {
597 |         if (!VirtualUnlock(ptr, len)) {
598 |             LLAMA_LOG_WARN("warning: failed to VirtualUnlock buffer: %s\n",
599 |                     llama_format_win_err(GetLastError()).c_str());
600 |         }
601 |     }
602 | #else
603 |     static size_t lock_granularity() {
604 |         return (size_t) 65536;
605 |     }
606 | 
607 |     bool raw_lock(const void * addr, size_t len) const {
608 |         LLAMA_LOG_WARN("warning: mlock not supported on this system\n");
609 |         return false;
610 |     }
611 | 
612 |     static void raw_unlock(const void * addr, size_t len) {}
613 | #endif
614 | 
615 |     impl() : addr(NULL), size(0), failed_already(false) {}
616 | 
617 |     void init(void * ptr) {
618 |         GGML_ASSERT(addr == NULL && size == 0);
619 |         addr = ptr;
620 |     }
621 | 
622 |     void grow_to(size_t target_size) {
623 |         GGML_ASSERT(addr);
624 |         if (failed_already) {
625 |             return;
626 |         }
627 |         size_t granularity = lock_granularity();
628 |         target_size = (target_size + granularity - 1) & ~(granularity - 1);
629 |         if (target_size > size) {
630 |             if (raw_lock((uint8_t *) addr + size, target_size - size)) {
631 |                 size = target_size;
632 |             } else {
633 |                 failed_already = true;
634 |             }
635 |         }
636 |     }
637 | 
638 |     void * addr;
639 |     size_t size;
640 | 
641 |     bool failed_already;
642 | };
643 | 
644 | llama_mlock::llama_mlock() : pimpl(std::make_unique<impl>()) {}
645 | llama_mlock::~llama_mlock() = default;
646 | 
647 | void llama_mlock::init(void * ptr) { pimpl->init(ptr); }
648 | void llama_mlock::grow_to(size_t target_size) { pimpl->grow_to(target_size); }
649 | 
650 | #if defined(_POSIX_MEMLOCK_RANGE) || defined(_WIN32)
651 | const bool llama_mlock::SUPPORTED = true;
652 | #else
653 | const bool llama_mlock::SUPPORTED = false;
654 | #endif
655 | 
656 | size_t llama_path_max() {
657 |     return PATH_MAX;
658 | }
659 | 


--------------------------------------------------------------------------------