pytorch

1
# Shaders processing
2
if(NOT USE_VULKAN)
3
  return()
4
endif()
5

6
set(VULKAN_GEN_OUTPUT_PATH "${CMAKE_BINARY_DIR}/vulkan/ATen/native/vulkan")
7
set(VULKAN_GEN_ARG_ENV "")
8

9
if(USE_VULKAN_RELAXED_PRECISION)
10
  list(APPEND VULKAN_GEN_ARG_ENV "PRECISION=mediump")
11
endif()
12
if(USE_VULKAN_FP16_INFERENCE)
13
  list(APPEND VULKAN_GEN_ARG_ENV "FLOAT_IMAGE_FORMAT=rgba16f")
14
else()
15
  list(APPEND VULKAN_GEN_ARG_ENV "FLOAT_IMAGE_FORMAT=rgba32f")
16
endif()
17

18
# Precompiling shaders
19
if(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")
25
else()
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)
37
endif()
38

39
set(PYTHONPATH "$ENV{PYTHONPATH}")
40
set(NEW_PYTHONPATH ${PYTHONPATH})
41
list(APPEND NEW_PYTHONPATH "${CMAKE_CURRENT_LIST_DIR}/..")
42
set(ENV{PYTHONPATH} ${NEW_PYTHONPATH})
43
execute_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)
53
set(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

59
set(vulkan_generated_cpp ${VULKAN_GEN_OUTPUT_PATH}/spv.cpp)
60