@@ -30,13 +30,19 @@ struct ibv_mr {
3030#include < vector>
3131#include < list>
3232
33- #include < unistd.h>
34- #include < sys/syscall.h>
3533#include < mutex>
3634#include < thread>
3735
38- #if defined(_linux_)
39- #include < sys/mman.h>
36+ #include < cstdlib>
37+ #include < cstring>
38+ #include < cerrno>
39+
40+ #if defined(_win_)
41+ #include < malloc.h> // _aligned_malloc, _aligned_free
42+ #else
43+ #include < sys/mman.h> // madvise
44+ #include < unistd.h>
45+ #include < sys/syscall.h>
4046#endif
4147
4248static constexpr size_t HPageSz = (1 << 21 );
@@ -45,6 +51,10 @@ using ::NMonitoring::TDynamicCounters;
4551
4652namespace NInterconnect ::NRdma {
4753
54+ // Cross-platform memory management
55+ static void * allocateMemory (size_t size, size_t alignment, bool hp);
56+ static void freeMemory (void * ptr) noexcept ;
57+
4858 class TChunk : public NNonCopyable ::TMoveOnly, public TAtomicRefCount<TChunk> {
4959 public:
5060
@@ -67,7 +77,7 @@ namespace NInterconnect::NRdma {
6777#else
6878 free (MRs.front ());
6979#endif
70- std::free (addr);
80+ freeMemory (addr);
7181 MRs.clear ();
7282 }
7383
@@ -204,11 +214,29 @@ namespace NInterconnect::NRdma {
204214 );
205215 }
206216
207- void * allocateMemory (size_t size, size_t alignment, bool hp) {
217+ static void * allocateMemory (size_t size, size_t alignment, bool hp) {
208218 if (size % alignment != 0 ) {
209219 return nullptr ;
210220 }
211- void * buf = std::aligned_alloc (alignment, size);
221+
222+ void * buf = nullptr ;
223+
224+ #if defined(_win_)
225+ // Windows: use _aligned_malloc
226+ buf = _aligned_malloc (size, alignment);
227+ if (!buf) {
228+ fprintf (stderr, " Failed to allocate aligned memory on Windows\n " );
229+ return nullptr ;
230+ }
231+ #else
232+ // POSIX/C++: std::aligned_alloc (C++17)
233+ buf = std::aligned_alloc (alignment, size);
234+ if (!buf) {
235+ fprintf (stderr, " Failed to allocate aligned memory on Unix\n " );
236+ return nullptr ;
237+ }
238+ #endif
239+
212240 if (hp) {
213241#if defined(_linux_)
214242 if (madvise (buf, size, MADV_HUGEPAGE) < 0 ) {
@@ -218,12 +246,23 @@ namespace NInterconnect::NRdma {
218246#endif
219247 for (size_t i = 0 ; i < size; i += HPageSz) {
220248 // We use THP right now. We need to touch each page to promote it to HUGE.
221- (( char *) buf)[i] = 0 ;
249+ static_cast < char *>( buf)[i] = 0 ;
222250 }
223251 }
224252 return buf;
225253 }
226254
255+ static void freeMemory (void * ptr) noexcept {
256+ if (!ptr) {
257+ return ;
258+ }
259+ #if defined(_win_)
260+ _aligned_free (ptr);
261+ #else
262+ std::free (ptr);
263+ #endif
264+ }
265+
227266 std::vector<ibv_mr*> registerMemory (void * addr, size_t size, const NInterconnect::NRdma::NLinkMgr::TCtxsMap& ctxs) {
228267 std::vector<ibv_mr*> res;
229268#ifndef MEM_POOL_DISABLE_RDMA_SUPPORT
@@ -307,7 +346,7 @@ namespace NInterconnect::NRdma {
307346
308347 auto mrs = registerMemory (ptr, size, Ctxs);
309348 if (mrs.empty ()) {
310- std::free (ptr);
349+ freeMemory (ptr);
311350 return nullptr ;
312351 }
313352
0 commit comments