diff --git a/build2cmake/src/templates/metal/compile-metal.cmake b/build2cmake/src/templates/metal/compile-metal.cmake index d7c7a365..f82874f0 100644 --- a/build2cmake/src/templates/metal/compile-metal.cmake +++ b/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) diff --git a/build2cmake/src/templates/metal/kernel.cmake b/build2cmake/src/templates/metal/kernel.cmake index 0ff280b9..ae399ba6 100644 --- a/build2cmake/src/templates/metal/kernel.cmake +++ b/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 %} \ No newline at end of file diff --git a/build2cmake/src/templates/metal/preamble.cmake b/build2cmake/src/templates/metal/preamble.cmake index c5cf4256..0f7b4f95 100644 --- a/build2cmake/src/templates/metal/preamble.cmake +++ b/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") diff --git a/build2cmake/src/templates/metal/torch-extension.cmake b/build2cmake/src/templates/metal/torch-extension.cmake index 3f81df03..9e9fa7aa 100644 --- a/build2cmake/src/templates/metal/torch-extension.cmake +++ b/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() \ No newline at end of file