Windows compatibility fix (#368)

* Windows compatibility fix * Missing header * Define NOMINMAX to avoid polluting std::max
dmlc · Feb 17, 2022 · 9b416be · 9b416be
1 parent 60b732a
commit 9b416be
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/include/treelite/predictor.h b/include/treelite/predictor.h
@@ -16,6 +16,11 @@
 #include <mutex>
 #include <cstdint>
 
+#ifdef _WIN32
+#define NOMINMAX
+#include <windows.h>
+#endif  // _WIN32
+
 namespace treelite {
 namespace predictor {
 
@@ -71,8 +76,13 @@ union Entry {
 
 class SharedLibrary {
  public:
+#ifdef _WIN32
+  using LibraryHandle = HMODULE;
+  using FunctionHandle = FARPROC;
+#else  // _WIN32
   using LibraryHandle = void*;
   using FunctionHandle = void*;
+#endif  // _WIN32
   SharedLibrary();
   ~SharedLibrary();
   void Load(const char* libpath);

diff --git a/src/predictor/thread_pool/thread_pool.h b/src/predictor/thread_pool/thread_pool.h
@@ -81,12 +81,7 @@ class ThreadPool {
 
   inline void SetAffinity() {
 #ifdef _WIN32
-    /* Windows */
-    SetThreadAffinityMask(GetCurrentThread(), 0x1);
-    for (int i = 0; i < num_worker_; ++i) {
-      const int core_id = i + 1;
-      SetThreadAffinityMask(thread_[i].native_handle(), (1ULL << core_id));
-    }
+    /* Do not set thread affinity on Windows */
 #elif defined(__APPLE__) && defined(__MACH__)
 #include <TargetConditionals.h>
 #if TARGET_OS_MAC == 1