Convnext model fails at runtime on Vulkan

[GregoryComer]

🐛 Describe the bug

The convnext_small model from torchvision fails at runtime on the Vulkan backend with "RuntimeError: Exception raised from add_reduce2d_node at /Users/gjcomer/src/executorch/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp:186: (graph.packed_dim_of(in) != reduce_dim1) 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_convnext_small_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, 144, 144),)
model = torchvision.models.convnext_small().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