xref: /aosp_15_r20/external/pytorch/cmake/VulkanCodegen.cmake (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1# Shaders processing
2if(NOT USE_VULKAN)
3  return()
4endif()
5
6set(VULKAN_GEN_OUTPUT_PATH "${CMAKE_BINARY_DIR}/vulkan/ATen/native/vulkan")
7set(VULKAN_GEN_ARG_ENV "")
8
9if(USE_VULKAN_RELAXED_PRECISION)
10  list(APPEND VULKAN_GEN_ARG_ENV "PRECISION=mediump")
11endif()
12if(USE_VULKAN_FP16_INFERENCE)
13  list(APPEND VULKAN_GEN_ARG_ENV "FLOAT_IMAGE_FORMAT=rgba16f")
14else()
15  list(APPEND VULKAN_GEN_ARG_ENV "FLOAT_IMAGE_FORMAT=rgba32f")
16endif()
17
18# Precompiling shaders
19if(ANDROID)
20  if(NOT ANDROID_NDK)
21    message(FATAL_ERROR "ANDROID_NDK not set")
22  endif()
23
24  set(GLSLC_PATH "${ANDROID_NDK}/shader-tools/${ANDROID_NDK_HOST_SYSTEM_NAME}/glslc")
25else()
26  find_program(
27    GLSLC_PATH glslc
28    PATHS
29    ENV VULKAN_SDK
30    PATHS "$ENV{VULKAN_SDK}/${CMAKE_HOST_SYSTEM_PROCESSOR}/bin"
31    PATHS "$ENV{VULKAN_SDK}/bin"
32  )
33
34  if(NOT GLSLC_PATH)
35    message(FATAL_ERROR "USE_VULKAN glslc not found")
36  endif(NOT GLSLC_PATH)
37endif()
38
39set(PYTHONPATH "$ENV{PYTHONPATH}")
40set(NEW_PYTHONPATH ${PYTHONPATH})
41list(APPEND NEW_PYTHONPATH "${CMAKE_CURRENT_LIST_DIR}/..")
42set(ENV{PYTHONPATH} ${NEW_PYTHONPATH})
43execute_process(
44  COMMAND
45  "${Python_EXECUTABLE}"
46  ${CMAKE_CURRENT_LIST_DIR}/../tools/gen_vulkan_spv.py
47  --glsl-path ${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/vulkan/glsl
48  --output-path ${VULKAN_GEN_OUTPUT_PATH}
49  --glslc-path=${GLSLC_PATH}
50  --tmp-dir-path=${CMAKE_BINARY_DIR}/vulkan/spv
51  --env ${VULKAN_GEN_ARG_ENV}
52  RESULT_VARIABLE error_code)
53set(ENV{PYTHONPATH} ${PYTHONPATH})
54
55  if(error_code)
56    message(FATAL_ERROR "Failed to gen spv.h and spv.cpp with precompiled shaders for Vulkan backend")
57  endif()
58
59set(vulkan_generated_cpp ${VULKAN_GEN_OUTPUT_PATH}/spv.cpp)
60