Introduce graph runtime shader library that enables dynamic shapes (#2366)
Summary:
Pull Request resolved: https://github.com/pytorch/executorch/pull/2366
## Context
https://github.com/pytorch/pytorch/pull/121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that https://github.com/pytorch/pytorch/pull/115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
ghstack-source-id: 218429132
exported-using-ghexport
bypass-github-export-checks
bypass-github-pytorch-ci-checks
bypass-github-executorch-ci-checks
Reviewed By: jorgep31415
Differential Revision: D54754545
fbshipit-source-id: 7e2074699b61f8358358775a8b790d34dcb99ee6 S
Stephen Jia committed
835279e11f2cfd8787419d365d0a5e96f11d077d
Parent: 341d2d9
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
on 3/13/2024, 6:47:33 AM