Vulkan replication_pad3d pte fails to load with missing shader error

[GregoryComer]

🐛 Describe the bug

Models with replication_pad3d decomposed and partially delegated to vulkan error out with a missing shader error (see below).

import torch
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
from executorch.exir import to_edge_transform_and_lower, EdgeCompileConfig, to_edge
from executorch.extension.pybindings.portable_lib import _load_for_executorch_from_buffer
from typing import Callable, List, Optional, Tuple, Union

class Model(torch.nn.Module):
    def __init__(
        self,
        padding=2,
    ):
        super().__init__()
        self.reflection_pad = torch.nn.ReflectionPad3d(
            padding=padding,
        )
        
    def forward(self, x):
        return self.reflection_pad(x)
        
model = Model()
inputs = (
    torch.randn(2, 3, 4, 4, 4),
)
eager_outputs = model(*inputs)

ep = torch.export.export(model.eval(), inputs)
print(ep)
lowered = to_edge_transform_and_lower(
    ep,
    partitioner=[VulkanPartitioner()],
    compile_config=EdgeCompileConfig(_check_ir_validity=False)
).to_executorch()
print(lowered.exported_program())

et_model = _load_for_executorch_from_buffer(lowered.buffer)
et_outputs = et_model([*inputs])[0]

print(f"Inputs: {inputs}")
print(f"Eager: {eager_outputs}")
print(f"ET:    {et_outputs}")

Outputs:

RuntimeError                              Traceback (most recent call last)
Cell In[18], line 35
     28 lowered = to_edge_transform_and_lower(
     29     ep,
     30     partitioner=[VulkanPartitioner()],
     31     compile_config=EdgeCompileConfig(_check_ir_validity=False)
     32 ).to_executorch()
     33 print(lowered.exported_program())
---> 35 et_model = _load_for_executorch_from_buffer(lowered.buffer)
     36 et_outputs = et_model([*inputs])[0]
     38 print(f"Inputs: {inputs}")
RuntimeError: Exception raised from get_shader_info at fbcode/executorch/backends/vulkan/runtime/api/ShaderRegistry.cpp:54: (it != listings_.end()) is false! Could not find ShaderInfo with name abs_int32_texture3d

Versions

Run on Meta internal master, Jul 3, fbcode/SwiftShader

cc @SS-JIA @manuelcandales @cbilgin