Vulkan partially delegated replication_pad2d fails to run

[GregoryComer]

🐛 Describe the bug

Graphs that contain replication_pad2d and are delegated to Vulkan fail to run with an error in portable index select. This error does not occur when not delegating to Vulkan. It appears that the Vulkan partition is giving invalid outputs.

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.replication_pad = torch.nn.ReplicationPad2d(
            padding=padding,
        )
        
    def forward(self, x):
        return self.replication_pad(x)
        
model = Model()
inputs = (
    torch.randn(1, 3, 8, 8),
)
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:

[advanced_index_util.cpp:463] Check failed (index_val >= 0 && index_val < in.size(i)): Index 4294967296 is out of bounds for input dimension 2 with size 8.
[op_index.cpp:104] Check failed (success): 
[method.cpp:1322] KernelCall failed at instruction 0:1 in operator aten::index.Tensor_out: 0x12
[method.cpp:1328] arg 0 with type id 1
[method.cpp:1328] arg 1 with type id 11
[method.cpp:1328] arg 2 with type id 1
[method.cpp:1328] arg 3 with type id 1

Versions

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

cc @SS-JIA @manuelcandales @cbilgin