[Vulkan] Enable Codegen ShaderInfo Registry from GLSLT + Params YAML files (conv2d_pw) (#91916)
@bypass-github-export-checks
This diff allows for adding entries to the shader registry by specifying which op names and registry keys should map to a template codegen Shader in the codegen Shader's glslt and params yaml files.
This can be done by
- adding a REGISTER_FOR entry which maps to either a tuple of (op name, list of registry keys) or null to the YAML file, and
- adding a ```REGISTER_FOR = $REGISTER_FOR``` line to the ShaderInfo comment in the glslt file
Ex.
YAML File:
```
conv2d_pw:
parameter_names_with_default_values:
...
REGISTER_FOR:
- !!python/tuple ["conv2d_pw", ["catchall"]]
parameter_values:
- ...
REGISTER_FOR: null
```
GLSLT File:
```
...
* REGISTER_FOR = $REGISTER_FOR
...
```
This diff also registers the conv2d_pw_2x2 Shader under ```'conv2d_pw → 'catchall'``` in the registry and uses ```VK_REGISTRY_KERNEL``` to retrieve the shader by look up in the registry
The shader registry generated in spv.cpp now looks like
```
ShaderRegistry shader_registry = {
{"conv2d", {{"catchall", "conv2d"}}},
{"conv2d_pw", {{"catchall", "conv2d_pw_2x2"}}}};
```
and the generated conv2d_p2_KxK.glsl files look like:
K=1
```
...
/*
* TILE_SIZE = (1, 1, 1)
* WEIGHT_STORAGE = TEXTURE_2D
* WEIGHT_STORAGE_LAYOUT = OC4,IC4,4ic,4oc
* BIAS_STORAGE = TEXTURE_2D
* REGISTER_FOR = None
*/
...
```
K=2
```
...
/*
* TILE_SIZE = (2, 2, 1)
* WEIGHT_STORAGE = TEXTURE_2D
* WEIGHT_STORAGE_LAYOUT = OC4,IC4,4ic,4oc
* BIAS_STORAGE = TEXTURE_2D
* REGISTER_FOR = ('conv2d_pw', ['catchall'])
*/
...
```
Differential Revision: [D42198560](https://our.internmc.facebook.com/intern/diff/D42198560/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/91916
Approved by: https://github.com/mcr229
