windows-support

bench: 4571956
connormcmonigle · Oct 22, 2022 · 02fd267 · 02fd267
1 parent 0365f01
commit 02fd267
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/include/nnue_util.h b/include/nnue_util.h
@@ -21,7 +21,6 @@
 
 #include <algorithm>
 #include <cmath>
-#include <cstdlib>
 #include <cstring>
 #include <iostream>
 #include <type_traits>
@@ -297,7 +296,7 @@ struct big_affine {
   }
 
   big_affine(const big_affine<T, dim0, dim1>& other) {
-    W = static_cast<T*>(std::aligned_alloc(simd::alignment, sizeof(T) * W_numel));
+    W = static_cast<T*>(simd::aligned_alloc(simd::alignment, sizeof(T) * W_numel));
 #pragma omp simd
     for (size_t i = 0; i < W_numel; ++i) { W[i] = other.W[i]; }
     for (size_t i = 0; i < b_numel; ++i) { b[i] = other.b[i]; }
@@ -308,9 +307,9 @@ struct big_affine {
     std::swap(b, other.b);
   }
 
-  big_affine() { W = static_cast<T*>(std::aligned_alloc(simd::alignment, sizeof(T) * W_numel)); }
+  big_affine() { W = static_cast<T*>(simd::aligned_alloc(simd::alignment, sizeof(T) * W_numel)); }
   ~big_affine() {
-    if (W != nullptr) { std::free(W); }
+    if (W != nullptr) { simd::aligned_free(W); }
   }
 };
 

diff --git a/include/simd.h b/include/simd.h
@@ -20,11 +20,28 @@
 #include <x86intrin.h>
 
 #include <cstdint>
+#include <cstdlib>
 #include <type_traits>
 #include <utility>
 
 namespace simd {
 
+void* aligned_alloc(size_t alignment, size_t size) {
+#if defined(_WIN32)
+  return _mm_malloc(size, alignment);
+#else
+  return std::aligned_alloc(alignment, size);
+#endif
+}
+
+void aligned_free(void* ptr) {
+#if defined(_WIN32)
+  _mm_free(ptr);
+#else
+  free(ptr);
+#endif
+}
+
 #if defined(__AVX2__)
 constexpr size_t alignment = 32;
 #elif defined(__SSSE3__)