pytorch/cmake/VulkanCodegen.cmake
Sam Spilsbury 111b2bf9da [cmake] Use list(APPEND instead of string(APPEND for vulkan codegen args
This was causing the shaders to be incorrectly templated because
both the precision argument and the format argument were being treated
as a single argument by argparse and therefore pasted into shaders
incorrectly. In turn this meant that shaders couldn't be compiled
when the precision or format options were turned on.

Fixes #76195

Pull Request resolved: https://github.com/pytorch/pytorch/pull/76196
Approved by: https://github.com/dagitses
2022-04-25 17:36:36 +00:00

80 lines
2.3 KiB
CMake

# Shaders processing
if(NOT USE_VULKAN)
return()
endif()
set(VULKAN_GEN_OUTPUT_PATH "${CMAKE_BINARY_DIR}/vulkan/ATen/native/vulkan")
set(VULKAN_GEN_ARG_ENV "")
if(USE_VULKAN_RELAXED_PRECISION)
list(APPEND VULKAN_GEN_ARG_ENV "precision=mediump")
endif()
if(USE_VULKAN_FP16_INFERENCE)
list(APPEND VULKAN_GEN_ARG_ENV "format=rgba16f")
endif()
if(USE_VULKAN_SHADERC_RUNTIME)
set(PYTHONPATH "$ENV{PYTHONPATH}")
set(ENV{PYTHONPATH} "$ENV{PYTHONPATH}:${CMAKE_CURRENT_LIST_DIR}/..")
execute_process(
COMMAND
"${PYTHON_EXECUTABLE}"
${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/gen_vulkan_glsl.py
--glsl-path ${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/vulkan/glsl
--output-path ${VULKAN_GEN_OUTPUT_PATH}
--tmp-dir-path=${CMAKE_BINARY_DIR}/vulkan/glsl
--env ${VULKAN_GEN_ARG_ENV}
RESULT_VARIABLE error_code)
set(ENV{PYTHONPATH} "$PYTHONPATH")
if(error_code)
message(FATAL_ERROR "Failed to gen glsl.h and glsl.cpp with shaders sources for Vulkan backend")
endif()
set(vulkan_generated_cpp ${VULKAN_GEN_OUTPUT_PATH}/glsl.cpp)
return()
endif()
if(NOT USE_VULKAN_SHADERC_RUNTIME)
# Precompiling shaders
if(ANDROID)
if(NOT ANDROID_NDK)
message(FATAL_ERROR "ANDROID_NDK not set")
endif()
set(GLSLC_PATH "${ANDROID_NDK}/shader-tools/${ANDROID_NDK_HOST_SYSTEM_NAME}/glslc")
else()
find_program(
GLSLC_PATH glslc
PATHS
ENV VULKAN_SDK
PATHS "$ENV{VULKAN_SDK}/${CMAKE_HOST_SYSTEM_PROCESSOR}/bin"
PATHS "$ENV{VULKAN_SDK}/bin"
)
if(NOT GLSLC_PATH)
message(FATAL_ERROR "USE_VULKAN glslc not found")
endif(NOT GLSLC_PATH)
endif()
set(PYTHONPATH "$ENV{PYTHONPATH}")
set(ENV{PYTHONPATH} "$ENV{PYTHONPATH}:${CMAKE_CURRENT_LIST_DIR}/..")
execute_process(
COMMAND
"${PYTHON_EXECUTABLE}"
${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/gen_vulkan_spv.py
--glsl-path ${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/vulkan/glsl
--output-path ${VULKAN_GEN_OUTPUT_PATH}
--glslc-path=${GLSLC_PATH}
--tmp-dir-path=${CMAKE_BINARY_DIR}/vulkan/spv
--env ${VULKAN_GEN_ARG_ENV}
RESULT_VARIABLE error_code)
set(ENV{PYTHONPATH} "$PYTHONPATH")
if(error_code)
message(FATAL_ERROR "Failed to gen spv.h and spv.cpp with precompiled shaders for Vulkan backend")
endif()
set(vulkan_generated_cpp ${VULKAN_GEN_OUTPUT_PATH}/spv.cpp)
endif()