Swin_v2_t fails at runtime on Vulkan

[GregoryComer]

🐛 Describe the bug

The swin_v2_t model from torchvision fails at runtime on the Vulkan backend with "RuntimeError: Exception raised from get_uniform_data at /Users/gjcomer/src/executorch/../executorch/backends/vulkan/runtime/api/containers/Tensor.h:679: (sizes_.size() <= 4) is false!".

This can be reproduced with the following test case command or standalone script.

python -m executorch.backends.test.suite.runner models --flow vulkan --filter "test_swin_v2_t_vulkan_float32$"

Standalone repro:

import torch
import torchvision
from executorch.exir import to_edge_transform_and_lower

from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner

inputs = (torch.randn(1, 3, 224, 224),)
model = torchvision.models.swin_v2_t().eval()

ep = torch.export.export(model, inputs)

model = to_edge_transform_and_lower(
    torch.export.export(model, inputs),
    partitioner=[VulkanPartitioner()],
).to_executorch()

print("Running model...")
from executorch.extension.pybindings.portable_lib import _load_for_executorch_from_buffer
loaded_model = _load_for_executorch_from_buffer(model.buffer)
loaded_model([*inputs])

Note that running the backend test case requires executorch's python bindings to be built with the Vulkan backend. An example build command is below.

CMAKE_ARGS="-DEXECUTORCH_BUILD_VULKAN=ON" ./install_executorch.sh --editable

Versions

Commit fbda3a9, M1 Mac, using MoltenVK

cc @SS-JIA @manuelcandales @cbilgin