From ac02cf33bd77454ce3c3bcb5eec571b5cd6766d6 Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Sat, 24 Feb 2024 22:20:57 -0800
Subject: [PATCH] Fix some issues using MLX in C++ (#739)

* fix preamble build

* fix some issues with using MLX as a dep in C++
---
 CMakeLists.txt                    |  2 --
 mlx/backend/common/CMakeLists.txt | 10 +++++-----
 mlx/backend/metal/CMakeLists.txt  |  2 +-
 mlx/backend/metal/metal.cpp       |  6 +++++-
 mlx/backend/metal/metal.h         | 11 ++---------
 mlx/backend/no_metal/metal.cpp    |  6 +++++-
 6 files changed, 18 insertions(+), 19 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index ffccb37d4..07c642120 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -67,8 +67,6 @@ if (MLX_BUILD_METAL AND NOT METAL_LIB)
   set(MLX_BUILD_METAL OFF)
 elseif (MLX_BUILD_METAL)
   message(STATUS "Building METAL sources")
-  add_compile_definitions(_METAL_)
-
   # Throw an error if xcrun not found
   execute_process(COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version"
                   OUTPUT_VARIABLE MACOS_VERSION
diff --git a/mlx/backend/common/CMakeLists.txt b/mlx/backend/common/CMakeLists.txt
index 569d690ef..a9fcf3b71 100644
--- a/mlx/backend/common/CMakeLists.txt
+++ b/mlx/backend/common/CMakeLists.txt
@@ -9,15 +9,15 @@ add_custom_command(
               ${CMAKE_CURRENT_SOURCE_DIR}/make_compiled_preamble.sh
               ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp
               ${CMAKE_CXX_COMPILER}
-              ${CMAKE_SOURCE_DIR}
+              ${PROJECT_SOURCE_DIR}
               ${CLANG}
 
     DEPENDS make_compiled_preamble.sh
             compiled_preamble.h
-            ${CMAKE_SOURCE_DIR}/mlx/types/half_types.h
-            ${CMAKE_SOURCE_DIR}/mlx/types/fp16.h
-            ${CMAKE_SOURCE_DIR}/mlx/types/bf16.h
-            ${CMAKE_SOURCE_DIR}/mlx/types/complex.h
+            ${PROJECT_SOURCE_DIR}/mlx/types/half_types.h
+            ${PROJECT_SOURCE_DIR}/mlx/types/fp16.h
+            ${PROJECT_SOURCE_DIR}/mlx/types/bf16.h
+            ${PROJECT_SOURCE_DIR}/mlx/types/complex.h
             ops.h
 )
 
diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt
index 063c283fe..fe764c494 100644
--- a/mlx/backend/metal/CMakeLists.txt
+++ b/mlx/backend/metal/CMakeLists.txt
@@ -4,7 +4,7 @@ add_custom_command(
               ${CMAKE_CURRENT_SOURCE_DIR}/make_compiled_preamble.sh
               ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp
               ${CMAKE_C_COMPILER}
-              ${CMAKE_SOURCE_DIR}
+              ${PROJECT_SOURCE_DIR}
     DEPENDS make_compiled_preamble.sh
             kernels/compiled_preamble.h
             kernels/unary.h
diff --git a/mlx/backend/metal/metal.cpp b/mlx/backend/metal/metal.cpp
index 96cf87818..6035d2a7f 100644
--- a/mlx/backend/metal/metal.cpp
+++ b/mlx/backend/metal/metal.cpp
@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <cstdlib>
 #include <future>
@@ -10,6 +10,10 @@
 
 namespace mlx::core::metal {
 
+bool is_available() {
+  return true;
+}
+
 int max_ops_per_buffer() {
   auto get_val = []() {
     if (const char* buff_str = std::getenv("MLX_MAX_OPS_PER_BUFFER")) {
diff --git a/mlx/backend/metal/metal.h b/mlx/backend/metal/metal.h
index d249daac0..11ac510fd 100644
--- a/mlx/backend/metal/metal.h
+++ b/mlx/backend/metal/metal.h
@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #pragma once
 
@@ -11,14 +11,7 @@
 
 namespace mlx::core::metal {
 
-constexpr bool is_available() {
-#ifdef _METAL_
-  return true;
-#else
-  return false;
-#endif
-}
-
+bool is_available();
 bool cache_enabled(void);
 void set_cache_enabled(bool enabled);
 
diff --git a/mlx/backend/no_metal/metal.cpp b/mlx/backend/no_metal/metal.cpp
index 257fcbb5d..e62c05c63 100644
--- a/mlx/backend/no_metal/metal.cpp
+++ b/mlx/backend/no_metal/metal.cpp
@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <stdexcept>
 
@@ -6,6 +6,10 @@
 
 namespace mlx::core::metal {
 
+bool is_available() {
+  return false;
+}
+
 void new_stream(Stream) {}
 std::shared_ptr<void> new_scoped_memory_pool() {
   return nullptr;