From d2b28d3502fc6838856708601389bb78c03badd8 Mon Sep 17 00:00:00 2001
From: Cheng Zhao <zcbenz@gmail.com>
Date: Sun, 8 Dec 2024 08:28:03 +0000
Subject: [PATCH 1/3] Make compiled preamble work in MSVC

---
 mlx/backend/common/CMakeLists.txt             | 13 ++++++-
 mlx/backend/common/make_compiled_preamble.ps1 | 38 +++++++++++++++++++
 2 files changed, 49 insertions(+), 2 deletions(-)
 create mode 100644 mlx/backend/common/make_compiled_preamble.ps1

diff --git a/mlx/backend/common/CMakeLists.txt b/mlx/backend/common/CMakeLists.txt
index 92f6ab7da..02ccc81e1 100644
--- a/mlx/backend/common/CMakeLists.txt
+++ b/mlx/backend/common/CMakeLists.txt
@@ -5,13 +5,22 @@ else()
   set(COMPILER ${CMAKE_CXX_COMPILER})
 endif()
 
+if(WIN32)
+  set(SHELL_EXT ps1)
+  set(SHELL_CMD powershell -ExecutionPolicy Bypass -File)
+else()
+  set(SHELL_EXT sh)
+  set(SHELL_CMD /bin/bash)
+endif()
+
+message("CMAKE_CXX_COMPILER=${CMAKE_CXX_COMPILER}")
 add_custom_command(
   OUTPUT compiled_preamble.cpp
   COMMAND
-    /bin/bash ${CMAKE_CURRENT_SOURCE_DIR}/make_compiled_preamble.sh
+    ${SHELL_CMD} ${CMAKE_CURRENT_SOURCE_DIR}/make_compiled_preamble.${SHELL_EXT}
     ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp ${COMPILER}
     ${PROJECT_SOURCE_DIR} ${CLANG} ${CMAKE_SYSTEM_PROCESSOR}
-  DEPENDS make_compiled_preamble.sh
+  DEPENDS make_compiled_preamble.${SHELL_EXT}
           compiled_preamble.h
           ${PROJECT_SOURCE_DIR}/mlx/types/half_types.h
           ${PROJECT_SOURCE_DIR}/mlx/types/fp16.h
diff --git a/mlx/backend/common/make_compiled_preamble.ps1 b/mlx/backend/common/make_compiled_preamble.ps1
new file mode 100644
index 000000000..0b2248b67
--- /dev/null
+++ b/mlx/backend/common/make_compiled_preamble.ps1
@@ -0,0 +1,38 @@
+# This script generates a C++ function that provides the CPU
+# code for use with kernel generation.
+#
+# Copyright © 2024 Apple Inc.
+
+$OUTPUT_FILE = $args[0]
+$CL = $args[1]
+$SRCDIR = $args[2]
+
+# Get command result as array.
+$CONTENT = & $CL /std:c++17 /EP "/I$SRCDIR" /Tp "$SRCDIR/mlx/backend/common/compiled_preamble.h"
+# Remove empty lines.
+# Otherwise there will be too much empty lines making the result unreadable.
+$CONTENT = $CONTENT | Where-Object { $_.Trim() -ne '' }
+# Concatenate to string.
+$CONTENT = $CONTENT -join '`n'
+
+# Append extra content.
+$CONTENT = @"
+$($CONTENT)
+using namespace mlx::core;
+using namespace mlx::core::detail;
+"@
+
+# Convert each char to ASCII code.
+# Unlike the unix script that outputs string literal directly, the output from
+# MSVC is way too large to be embedded as string and compilation will fail, so
+# we store it as static array instead.
+$CHARCODES = ([System.Text.Encoding]::ASCII.GetBytes($CONTENT) -join ', ') + ', 0'
+
+$OUTPUT = @"
+const char* get_kernel_preamble() {
+  static char preamble[] = { $CHARCODES };
+  return preamble;
+}
+"@
+
+Set-Content -Path $OUTPUT_FILE -Value $OUTPUT

From 8f50f269eb26fa4c7a951ebe12e974fcafcbf01e Mon Sep 17 00:00:00 2001
From: Cheng Zhao <zcbenz@gmail.com>
Date: Sun, 8 Dec 2024 08:42:27 +0000
Subject: [PATCH 2/3] Remove logging

---
 mlx/backend/common/CMakeLists.txt | 1 -
 1 file changed, 1 deletion(-)

diff --git a/mlx/backend/common/CMakeLists.txt b/mlx/backend/common/CMakeLists.txt
index 02ccc81e1..e0512632c 100644
--- a/mlx/backend/common/CMakeLists.txt
+++ b/mlx/backend/common/CMakeLists.txt
@@ -13,7 +13,6 @@ else()
   set(SHELL_CMD /bin/bash)
 endif()
 
-message("CMAKE_CXX_COMPILER=${CMAKE_CXX_COMPILER}")
 add_custom_command(
   OUTPUT compiled_preamble.cpp
   COMMAND

From 7341a67cdf2652ab844e5cdf266434319eeb793d Mon Sep 17 00:00:00 2001
From: Cheng Zhao <zcbenz@gmail.com>
Date: Sun, 8 Dec 2024 09:54:25 +0000
Subject: [PATCH 3/3] Only use powershell for MSVC

---
 mlx/backend/common/CMakeLists.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlx/backend/common/CMakeLists.txt b/mlx/backend/common/CMakeLists.txt
index e0512632c..a412cbe7b 100644
--- a/mlx/backend/common/CMakeLists.txt
+++ b/mlx/backend/common/CMakeLists.txt
@@ -5,7 +5,7 @@ else()
   set(COMPILER ${CMAKE_CXX_COMPILER})
 endif()
 
-if(WIN32)
+if(MSVC)
   set(SHELL_EXT ps1)
   set(SHELL_CMD powershell -ExecutionPolicy Bypass -File)
 else()