From ed4ec81bca3962a7f2bee0200fbd60bf69574c30 Mon Sep 17 00:00:00 2001
From: Cheng <git@zcbenz.com>
Date: Thu, 19 Dec 2024 12:26:04 +0900
Subject: [PATCH] Link python extension with mlx statically on Windows (#1716)

* Link python extension with mlx statically on Windows

* More readable code
---
 CMakeLists.txt |  5 +----
 setup.py       | 11 +++++++++--
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 60cf657e2..78bae582b 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -266,10 +266,7 @@ install(
   EXPORT MLXTargets
   LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
   ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
-  RUNTIME
-    DESTINATION
-      # On Windows, DLLs must be put in the same dir with the python bindings.
-      $<IF:$<BOOL:${WIN32}>,${CMAKE_INSTALL_PREFIX},${CMAKE_INSTALL_BINDIR}>
+  RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}
   INCLUDES
   DESTINATION ${CMAKE_INSTALL_INCLUDEDIR})
 
diff --git a/setup.py b/setup.py
index f6fa171ba..ddff007a3 100644
--- a/setup.py
+++ b/setup.py
@@ -2,6 +2,7 @@
 
 import datetime
 import os
+import platform
 import re
 import subprocess
 import sys
@@ -60,7 +61,6 @@ def build_extension(self, ext: CMakeExtension) -> None:
         cmake_args = [
             f"-DCMAKE_INSTALL_PREFIX={extdir}{os.sep}",
             f"-DCMAKE_BUILD_TYPE={cfg}",
-            "-DBUILD_SHARED_LIBS=ON",
             "-DMLX_BUILD_PYTHON_BINDINGS=ON",
             "-DMLX_BUILD_TESTS=OFF",
             "-DMLX_BUILD_BENCHMARKS=OFF",
@@ -77,11 +77,18 @@ def build_extension(self, ext: CMakeExtension) -> None:
         # Pass version to C++
         cmake_args += [f"-DMLX_VERSION={self.distribution.get_version()}"]  # type: ignore[attr-defined]
 
-        if sys.platform.startswith("darwin"):
+        if platform.system() == "Darwin":
             # Cross-compile support for macOS - respect ARCHFLAGS if set
             archs = re.findall(r"-arch (\S+)", os.environ.get("ARCHFLAGS", ""))
             if archs:
                 cmake_args += ["-DCMAKE_OSX_ARCHITECTURES={}".format(";".join(archs))]
+        if platform.system() == "Windows":
+            # On Windows DLLs must be put in the same dir with the extension
+            # while cmake puts mlx.dll into the "bin" sub-dir. Link with mlx
+            # statically to work around it.
+            cmake_args += ["-DBUILD_SHARED_LIBS=OFF"]
+        else:
+            cmake_args += ["-DBUILD_SHARED_LIBS=ON"]
 
         # Set CMAKE_BUILD_PARALLEL_LEVEL to control the parallel build level
         # across all generators.