[Vulkan] Add Basic Shader Registry (#91914)
@bypass-github-export-checks
We want to be able to look-up which shader to use in a registry given a particular op/algorithm name, which is what this diff enables. This is done with the newly added ```shader_registry``` map and ```look_up_shader_info``` function.
After this change, Shaders can be retrieved either with the ```VK_KERNEL``` macro, which gets the Shader with a specified name directly, or with the ```VK_REGISTRY_KERNEL``` macro, which looks up what Shader should be used for a specified algorithm name in the registry.
For now, the registry is empty and unused. In the next diffs in this stack, I will be adding support for registering a shader in the registry in GLSL and GLSLT + Params Yaml files.
I also
- Adjusted the formatting of spv.h and spv.cpp so that they are closer to what clang wants, which makes them easier to read. (proper indentation, proper order of includes, etc.)
- Moved the codegen spv/registry code from at::native::vulkan to at::native::vulkan::api (since registry.cpp / .h are in ```ATen/native/vulkan/api```)
Now spv.h looks like
```
#pragma once
#include <ATen/native/vulkan/api/Types.h>
#include <ATen/native/vulkan/api/vk_api.h>
#include <c10/util/flat_hash_map.h>
#include <string>
namespace at {
namespace native {
namespace vulkan {
namespace api {
struct ShaderInfo;
} // namespace api
typedef ska::flat_hash_map<std::string, api::ShaderInfo> ShaderListing;
typedef ska::flat_hash_map<std::string, std::string> RegistryKeyMap;
typedef ska::flat_hash_map<std::string, RegistryKeyMap> ShaderRegistry;
extern const ShaderListing shader_infos;
extern ShaderRegistry shader_registry;
inline const ShaderListing& get_shader_infos() {
return shader_infos;
}
inline ShaderRegistry& get_shader_registry() {
return shader_registry;
}
} // namespace vulkan
} // namespace native
} // namespace at
```
and spv.cpp looks like
```
#include <ATen/native/vulkan/api/Shader.h>
#include <ATen/native/vulkan/spv.h>
#include <stdint.h>
#include <vector>
namespace at {
namespace native {
namespace vulkan {
namespace {
const uint32_t adaptive_avg_pool2d_bin[] = {
119734787,
...
};
...
const uint32_t conv2d_pw_2x2_bin[] = {
119734787,
...
};
} // namespace
const ShaderListing shader_infos = {
{"adaptive_avg_pool2d",
api::ShaderInfo(
"vulkan.adaptive_avg_pool2d",
adaptive_avg_pool2d_bin,
3204,
{VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER},
std::vector<uint32_t>(),
api::StorageType::UNKNOWN,
api::StorageType::UNKNOWN)},
...
{"conv2d_pw_2x2",
api::ShaderInfo(
"vulkan.conv2d_pw_2x2",
conv2d_pw_2x2_bin,
7736,
{VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER},
{2, 2, 1},
api::StorageType::TEXTURE_2D,
api::StorageType::TEXTURE_2D)}};
ShaderRegistry shader_registry = {
};
} // namespace vulkan
} // namespace native
} // namespace at
```
(Full File: P594112814)
Differential Revision: [D41594453](https://our.internmc.facebook.com/intern/diff/D41594453/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/91914
Approved by: https://github.com/mcr229
