Link python extension with mlx statically on Windows (#1716)

* Link python extension with mlx statically on Windows

* More readable code

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})
 

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.