mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
545d2126f6
Summary:
This change makes two major improvements to PyTorch Vulkan's shader authoring workflow.
## Review Guide
There are a lot of changed files because every GLSL shader had to be touched. The majority of changes is changing
```
#define PRECISION $precision
#define FORMAT $format
```
to
```
#define PRECISION ${PRECISION}
#define FORMAT ${FORMAT}
```
due to changes in how shader templates are processed.
For reviewers, the primary functional changes to review are:
* `gen_vulkan_spv.py`
* Majority of functional changes are in this file, which controls how shader templates are processed.
* `shader_params.yaml`
* controls how shader variants are generated
## Python Codeblocks in Shader Templates
From now on, every compute shader (i.e. `.glsl`) is treated as a shader template. To this effect, the `templates/` folder has been removed and there is now a global `shader_params.yaml` file to describe the shader variants that should be generated for all shader templates.
**Taking inspiration from XNNPACK's [`xngen` tool](https://github.com/google/XNNPACK/blob/master/tools/xngen.py), shader templates can now use Python codeblocks**. One example is:
```
$if not INPLACE:
layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict writeonly image3D uOutput;
layout(set = 0, binding = 1) uniform PRECISION sampler3D uInput;
layout(set = 0, binding = 2) uniform PRECISION sampler3D uOther;
layout(set = 0, binding = 3) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 input_sizes;
ivec4 other_sizes;
float alpha;
}
uArgs;
$else:
layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict image3D uOutput;
layout(set = 0, binding = 1) uniform PRECISION sampler3D uOther;
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
uArgs;
```
Another is:
```
// PYTHON CODEBLOCK
$if not IS_DIV:
const int c_index = (pos.z % ((uArgs.output_sizes.z + 3) / 4)) * 4;
if (uArgs.other_sizes.z != 1 && c_index + 3 >= uArgs.output_sizes.z) {
ivec4 c_ind = ivec4(c_index) + ivec4(0, 1, 2, 3);
vec4 mask = vec4(lessThan(c_ind, ivec4(uArgs.output_sizes.z)));
other_texel = other_texel * mask + vec4(1, 1, 1, 1) - mask;
}
// PYTHON CODEBLOCK
$if not INPLACE:
ivec3 input_pos =
map_output_pos_to_input_pos(pos, uArgs.output_sizes, uArgs.input_sizes);
const vec4 in_texel =
load_texel(input_pos, uArgs.output_sizes, uArgs.input_sizes, uInput);
imageStore(uOutput, pos, OP(in_texel, other_texel, uArgs.alpha));
$else:
const vec4 in_texel = imageLoad(uOutput, pos);
imageStore(uOutput, pos, OP(in_texel, other_texel, uArgs.alpha));
```
In addition to making it easier and clearer to write shader templates, this enables shaders that were previously unable to be consolidated into a single template to now be represented using a single template, such as non inplace and inplace variants of the same shader.
## `generate_variant_forall` in shader variant YAML configuration
YAML files that describe how shader variants should be generated can now use a `generate_variant_forall` field to iterate over various settings for a specific parameter for each variant defined. Example:
```
unary_op:
parameter_names_with_default_values:
OPERATOR: exp(X)
INPLACE: 0
generate_variant_forall:
INPLACE:
- VALUE: 0
SUFFIX: ""
- VALUE: 1
SUFFIX: "inplace"
shader_variants:
- NAME: exp
OPERATOR: exp(X)
- NAME: sqrt
OPERATOR: sqrt(X)
- NAME: log
OPERATOR: log(X)
```
Previously, the `inplace` variants would need to have separate `shader_variants` entries. If there are multiple variables that need to be iterated across, then all possible combinations will be generated. Would be good to take a look to see how the new YAML configuration works.
Test Plan:
There is no functional change to this diff; we only need to make sure that the generated shaders are still correct. Therefore, we only need to run `vulkan_api_test`.
```
# On Mac Laptop
buck run --target-platforms ovr_config//platform/macos:arm64-fbsource //xplat/caffe2:pt_vulkan_api_test_binAppleMac\#macosx-arm64 -c pt.vulkan_full_precision=1 -- --gtest_filter="*"
```
Reviewed By: digantdesai
Differential Revision: D52087084
Pull Request resolved: https://github.com/pytorch/pytorch/pull/115948
Approved by: https://github.com/manuelcandales
60 lines
1.6 KiB
CMake
60 lines
1.6 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 "FLOAT_IMAGE_FORMAT=rgba16f")
|
|
else()
|
|
list(APPEND VULKAN_GEN_ARG_ENV "FLOAT_IMAGE_FORMAT=rgba32f")
|
|
endif()
|
|
|
|
# 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(NEW_PYTHONPATH ${PYTHONPATH})
|
|
list(APPEND NEW_PYTHONPATH "${CMAKE_CURRENT_LIST_DIR}/..")
|
|
set(ENV{PYTHONPATH} ${NEW_PYTHONPATH})
|
|
execute_process(
|
|
COMMAND
|
|
"${PYTHON_EXECUTABLE}"
|
|
${CMAKE_CURRENT_LIST_DIR}/../tools/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)
|