Maxvit_t model lowering segfaults on QNN

[GregoryComer]

🐛 Describe the bug

The maxvit_t model in torchvision segfaults on QNN when lowering. It doesn't seem to be a 100% repro, but it's pretty close on my local machine.

Output excerpt (note that the segfault happens several pages of logs in):

...
[INFO] [Qnn ExecuTorch]: Initialize Qnn backend parameters for Qnn executorch backend type 2
[INFO] [Qnn ExecuTorch]: Caching: Caching is in RESTORE MODE.
[INFO] [Qnn ExecuTorch]: QnnContextCustomProtocol expected magic number: 0x5678abcd but get: 0x2000000
[INFO] [Qnn ExecuTorch]: Running level=3 optimization.
Segmentation fault

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

python -m executorch.backends.test.suite.runner models --flow qnn --filter "test_maxvit_t_qnn_float32$"

Standalone repro:

from typing import Tuple

import executorch
import torch
import torchvision

from executorch.backends.qualcomm.utils.utils import (
    generate_qnn_executorch_compiler_spec,
    generate_htp_compiler_spec,
    QcomChipset,
    to_edge_transform_and_lower_to_qnn,
)

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

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

backend_options = generate_htp_compiler_spec(
    use_fp16=True,
)

compile_spec = generate_qnn_executorch_compiler_spec(
    soc_model=QcomChipset.SM8650,
    backend_options=backend_options,
)

model = to_edge_transform_and_lower_to_qnn(
    model,
    inputs,
    compile_spec
).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 QNN backend. An example build command is below, Note that it will still need the library paths to be set up properly as described in the ET QNN docs.

CMAKE_ARGS="-DEXECUTORCH_BUILD_QNN=ON -DQNN_SDK_ROOT=$QNN_SDK_ROOT" ./install_executorch.sh --editable

Versions

Commit fbda3a9, x86-64 simulator, WSL

cc @cccclai @winskuo-quic @shewu-quic @cbilgin

[cccclai]

@haowhsu-quic can you help with triaging this issue? Note that it's reported from the backend tester, and we can cherry pick changes after branch cut. One thing to note is that, we don't need to have the full model lowered. As long as the repro command work, that's good.

[haowhsu-quic]

Will look into it, this might take some time to resolve.