pytorch

3d6f85c9 - [Vulkan] Enable Codegen ShaderInfo Registry from GLSLT + Params YAML files (conv2d_pw) (#91916)

[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