Add JIT compiler support for Windows#3556
Open
dhiltgen wants to merge 2 commits into
Open
Conversation
Enable CPU mx.compile() on Windows by detecting and using clang-cl bundled with Visual Studio, or MSVC cl.exe, for JIT compilation. Keep GPU compile availability independent from the CPU compiler probe so CPU+GPU builds do not disable GPU mx.compile() when a host C++ compiler is unavailable. Changes: - Add clang-cl detection via vswhere and prefer a compiler matching the build toolchain - Add JitCompiler::available() to probe CPU JIT availability - Emit and load .dll JIT libraries on Windows - Support both MSVC and GCC/Clang preamble generation scripts, including optional SIMD flags - Use WIN32 shell detection and pass preamble SIMD flags through CMake - Define NOMINMAX/WIN32_LEAN_AND_MEAN on all WIN32 compilers
zcbenz
reviewed
May 17, 2026
| "{0}\\VC\\Tools\\Llvm\\{1}\\bin\\clang-cl.exe", vs_path, arch); | ||
| { | ||
| std::ifstream f(clang_cl_path); | ||
| if (f.good()) { |
Collaborator
There was a problem hiding this comment.
We usually use std::filesystem::exists.
| } | ||
| return result == 1; | ||
| #else | ||
| static int result = -1; |
Collaborator
There was a problem hiding this comment.
Suggested change
| static int result = -1; | |
| static int result = std::system("..."); | |
| return result == 0; |
| // Select the JIT compiler. The preamble was preprocessed at build time | ||
| // by whichever compiler built the library -- it contains compiler-specific | ||
| // builtins (e.g. __builtin_fpclassify for Clang, __is_same for MSVC) that | ||
| // are only valid for the same compiler family. Prefer the matching one. |
Collaborator
There was a problem hiding this comment.
You should ship the preamble headers along with your binary, which we do in python and can root out this problem completely (#3463).
| if (!info.clang_cl.empty() && info.cl_exe == info.clang_cl) { | ||
| compiler_flags += " -Wno-everything"; | ||
| } | ||
| #ifdef __AVX2__ |
Collaborator
There was a problem hiding this comment.
AVX2 should be checked at run time, you can't ensure binary built with AVX2 enabled would always run on machines with it.
| "\"{}\" /LD /EHsc /MD /Ox /nologo /std:c++17 {} \"{}\" " | ||
| "/link /out:\"{}\" 2>&1" | ||
| "cd /D \"{0}\" && " | ||
| "\"{1}\" /LD /EHsc /MD /Ox /nologo /std:c++17{5} \"{2}\" " |
Collaborator
There was a problem hiding this comment.
I prefer listing the args in increasing order otherwise I have to count the index of the args to know where they would be inserted.
Also fixed Visual Studio compiler selection conservative for the embedded prebuilt preamble fallback while preferring cl.exe when installed headers are available. Pass Visual Studio INCLUDE paths from vcvarsall into runtime JIT compilations so normal Python processes can find VC and Windows SDK headers.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Proposed changes
Split out from #3019
Enable CPU mx.compile() on Windows by detecting and using clang-cl bundled with Visual Studio, or MSVC cl.exe, for JIT compilation.
Keep GPU compile availability independent from the CPU compiler probe so CPU+GPU builds do not disable GPU mx.compile() when a host C++ compiler is unavailable.
Changes:
Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes