Fix metal kernels support

build2cmake/src/templates/metal/compile-metal.cmake

@@ -1,15 +1,19 @@
 # Metal shader compilation function
-function(compile_metal_shaders TARGET_NAME METAL_SOURCES)
+function(compile_metal_shaders TARGET_NAME METAL_SOURCES EXTRA_INCLUDE_DIRS)
     # Find the Metal compiler
     find_program(METAL_COMPILER xcrun REQUIRED)
 
     # Set Metal compiler flags
-    set(METAL_FLAGS "-std=metal3.0" "-O2")
+    set(METAL_FLAGS "-std=metal3.2" "-O2")
 
     # Output directory for compiled metallib
     set(METALLIB_OUTPUT_DIR "${CMAKE_BINARY_DIR}/metallib")
     file(MAKE_DIRECTORY ${METALLIB_OUTPUT_DIR})
-
+
+    foreach(INC ${EXTRA_INCLUDE_DIRS})
+        list(APPEND METAL_FLAGS "-I${INC}")
+    endforeach()
+
     # Separate .metal files from .h files and compile .metal files to .air
     set(AIR_FILES)
     set(METAL_FILES)

build2cmake/src/templates/metal/kernel.cmake

@@ -39,3 +39,10 @@ list(APPEND SRC {{'"${' + kernel_name + '_CPP_SRC}"'}})
 if({{kernel_name}}_METAL_SRC)
   list(APPEND ALL_METAL_SOURCES {{'"${' + kernel_name + '_METAL_SRC}"'}})
 endif()
+
+{% if includes %}
+# Keep the includes directory for the Metal sources
+if({{kernel_name}}_METAL_SRC)
+  list(APPEND METAL_INCLUDE_DIRS {{ includes }})
+endif()
+{% endif %}

build2cmake/src/templates/metal/preamble.cmake

@@ -1,5 +1,5 @@
 cmake_minimum_required(VERSION 3.26)
-project({{name}} LANGUAGES CXX)
+project({{name}} LANGUAGES CXX C OBJC OBJCXX)
 
 set(CMAKE_OSX_DEPLOYMENT_TARGET "15.0" CACHE STRING "Minimum macOS deployment version")
 

build2cmake/src/templates/metal/torch-extension.cmake

@@ -13,5 +13,5 @@ define_gpu_extension_target(
 
 # Compile Metal shaders if any were found
 if(ALL_METAL_SOURCES)
-  compile_metal_shaders({{ ops_name }} "${ALL_METAL_SOURCES}")
+  compile_metal_shaders({{ ops_name }} "${ALL_METAL_SOURCES}" "${METAL_INCLUDE_DIRS}")
 endif()