QNN, The performance gap between nchw and ncwh is huge. When either h or w is 1,

🐛 Describe the bug

The following is the test code. The performance gap between child_mod_type = "Conv2d_1_w" and child_mod_type = "Conv2d_h_1" is huge. Conv2d_1_w is 20 times faster than Conv2d_h_1.

import copy, os, torch
 
from torchao.quantization.pt2e.quantize_pt2e import (
    convert_pt2e,
    prepare_pt2e,
)
from executorch.backends.qualcomm.utils.utils import (
    dump_context_from_pte,
    generate_htp_compiler_spec,
    generate_qnn_executorch_compiler_spec,
    QcomChipset,
    to_edge_transform_and_lower_to_qnn,
)
from executorch.backends.qualcomm.tests.models import SimpleModel,Bmm,Add
from executorch.examples.qualcomm.utils import make_quantizer, SimpleADB
from executorch.devtools import generate_etrecord, Inspector
from executorch.devtools.inspector._inspector_utils import TimeScale

import torch
import torch.nn as nn
import torch.nn.functional as F

from pathlib import Path

from typing import Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F

from executorch.backends.qualcomm.debugger.utils import generate_optrace
from executorch.backends.qualcomm.utils.utils import get_soc_to_chipset_map


def convert_conv1d_to_conv2d_nch1(conv1d_layer):
    in_channels = conv1d_layer.in_channels
    out_channels = conv1d_layer.out_channels
    kernel_size = conv1d_layer.kernel_size[0]
    stride = conv1d_layer.stride[0]
    padding = conv1d_layer.padding[0]
    dilation = conv1d_layer.dilation[0]
    groups = conv1d_layer.groups
    bias = conv1d_layer.bias is not None
    
    conv2d = nn.Conv2d(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=(kernel_size, 1),
        stride=(stride, 1),
        padding=(padding, 0),
        dilation=(dilation, 1),
        groups=groups,
        bias=bias
    )

    with torch.no_grad():
        conv2d.weight.data = conv1d_layer.weight.data.unsqueeze(3)
        if bias:
            conv2d.bias.data = conv1d_layer.bias.data
    
    return conv2d

def convert_convtranspose1d_to_convtranspose2d_nch1(convt1d_layer):
    in_channels = convt1d_layer.in_channels
    out_channels = convt1d_layer.out_channels
    kernel_size = convt1d_layer.kernel_size[0]
    stride = convt1d_layer.stride[0]
    padding = convt1d_layer.padding[0]
    output_padding = convt1d_layer.output_padding[0]
    dilation = convt1d_layer.dilation[0]
    groups = convt1d_layer.groups
    bias = convt1d_layer.bias is not None

    convt2d = nn.ConvTranspose2d(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=(kernel_size, 1),
        stride=(stride, 1),
        padding=(padding, 0),
        output_padding=(output_padding, 0),
        dilation=(dilation, 1),
        groups=groups,
        bias=bias
    )
    
    with torch.no_grad():
        convt2d.weight.data = convt1d_layer.weight.data.unsqueeze(3)  # Add an extra dimension for the second axis
        if bias:
            convt2d.bias.data = convt1d_layer.bias.data
    
    return convt2d

def convert_model_conv1d_to_conv2d_nch1(model):
    for name, module in model.named_children():
        if isinstance(module, nn.Conv1d):
            setattr(model, name, convert_conv1d_to_conv2d_nch1(module))
        elif isinstance(module, nn.ConvTranspose1d):
            setattr(model, name, convert_convtranspose1d_to_convtranspose2d_nch1(module))
        else:
            convert_model_conv1d_to_conv2d_nch1(module)

class Conv2DVocoderWrapper_nch1(nn.Module):
    def __init__(self, original_vocoder):
        super().__init__()
        self.vocoder = original_vocoder
        convert_model_conv1d_to_conv2d_nch1(self.vocoder)
    
    def forward(self, x):
        output = self.vocoder(x)
        return output

def convert_conv1d_to_conv2d(conv1d_layer):
    in_channels = conv1d_layer.in_channels
    out_channels = conv1d_layer.out_channels
    kernel_size = conv1d_layer.kernel_size[0]
    stride = conv1d_layer.stride[0]
    padding = conv1d_layer.padding[0]
    dilation = conv1d_layer.dilation[0]
    groups = conv1d_layer.groups
    bias = conv1d_layer.bias is not None
    
    conv2d = nn.Conv2d(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=(1, kernel_size),
        stride=(1, stride),
        padding=(0, padding),
        dilation=(1, dilation),
        groups=groups,
        bias=bias
    )
    with torch.no_grad():
        conv2d.weight.data = conv1d_layer.weight.data.unsqueeze(2)
        if bias:
            conv2d.bias.data = conv1d_layer.bias.data
    return conv2d

def convert_convtranspose1d_to_convtranspose2d(convt1d_layer):
    in_channels = convt1d_layer.in_channels
    out_channels = convt1d_layer.out_channels
    kernel_size = convt1d_layer.kernel_size[0]
    stride = convt1d_layer.stride[0]
    padding = convt1d_layer.padding[0]
    output_padding = convt1d_layer.output_padding[0]
    dilation = convt1d_layer.dilation[0]
    groups = convt1d_layer.groups
    bias = convt1d_layer.bias is not None

    convt2d = nn.ConvTranspose2d(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=(1, kernel_size),
        stride=(1, stride),
        padding=(0, padding),
        output_padding=(0, output_padding),
        dilation=(1, dilation),
        groups=groups,
        bias=bias
    )
    
    with torch.no_grad():
        convt2d.weight.data = convt1d_layer.weight.data.unsqueeze(2)
        if bias:
            convt2d.bias.data = convt1d_layer.bias.data
    
    return convt2d

def convert_model_conv1d_to_conv2d(model):
    for name, module in model.named_children():
        if isinstance(module, nn.Conv1d):
            setattr(model, name, convert_conv1d_to_conv2d(module))
        elif isinstance(module, nn.ConvTranspose1d):
            setattr(model, name, convert_convtranspose1d_to_convtranspose2d(module))
        else:
            convert_model_conv1d_to_conv2d(module)

class Conv2DVocoderWrapper_nc1w(nn.Module):
    def __init__(self, original_vocoder):
        super().__init__()
        self.vocoder = original_vocoder
        convert_model_conv1d_to_conv2d(self.vocoder)
    
    def forward(self, x):
        output = self.vocoder(x)
        return output


class ResidualBlock(nn.Module):
    def __init__(self, kernel_size):
        super().__init__()
        self.block = nn.Sequential(
            nn.LeakyReLU(0.1),
            nn.Conv1d(128, 128, kernel_size=kernel_size, padding=kernel_size // 2),
            nn.LeakyReLU(0.1),
            nn.Conv1d(128, 128, kernel_size=kernel_size, padding=kernel_size // 2),
        )

    def forward(self, x):
        return self.block(x) + x

class CustomBranch(nn.Module):
    def __init__(self, kernel_size):
        super().__init__()
        self.resblock1 = ResidualBlock(kernel_size)
        self.resblock2 = ResidualBlock(kernel_size)
        self.resblock3 = ResidualBlock(kernel_size)

    def forward(self, x):
        x = self.resblock1(x)
        x = self.resblock2(x)
        x = self.resblock3(x)
        x = x / 3.0
        return x

class FullModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv_in = nn.Conv1d(80, 1024, kernel_size=7, padding=3)
        self.leaky_relu_in = nn.LeakyReLU(0.1)
        self.deconv = nn.ConvTranspose1d(256, 128, kernel_size=7, stride=3, padding=2)

        self.branch1 = CustomBranch(kernel_size=3)
        self.branch2 = CustomBranch(kernel_size=7)
        self.branch3 = CustomBranch(kernel_size=11)

        self.final_relu = nn.LeakyReLU(0.1)
        self.tanh = nn.Tanh()
    def forward(self, x):
        x = self.deconv(x)

        out1 = self.branch1(x)
        out2 = self.branch2(x)
        out3 = self.branch3(x)

        out = out1 + out2 + out3
        out = self.final_relu(out)

        out = self.tanh(out)

        return out

def get_model_and_input(mod_type: str) -> Tuple[nn.Module, Tuple[torch.Tensor]]:
    if mod_type == "Conv2d_h_1":
        input = torch.randn(1, 256, 12800, 1)
        return Conv2DVocoderWrapper_nch1(FullModel().eval()), (input,)
    elif mod_type == "Conv2d_1_w":
        input = torch.randn(1, 256, 1, 12800)
        return Conv2DVocoderWrapper_nc1w(FullModel().eval()), (input,)
    else:
        raise ValueError(f"Unknown mod_type: {mod_type}")

torch.manual_seed(42)
if __name__ == '__main__':

    child_mod_type = "Conv2d_1_w"
    # child_mod_type = "Conv2d_h_1"
    model_class, sample_input = get_model_and_input(child_mod_type)
    
    input_list = ""
    for i in range(10):
        current_input = ""
        for j in range(1):
            file_name = f"input_{i}_{j}.pt"
            torch.save(torch.randn(1, 256, 1, 12800), file_name)
            current_input += f"{file_name} "
        input_list += f"{current_input.strip()}\n"

    with open(f"input_list", 'w') as f:
        f.write(input_list)
    
    exported = torch.export.export(model_class, sample_input)
    
    onnx_path = f"{child_mod_type}.onnx"
    torch.onnx.export(model_class, 
        sample_input,
        onnx_path,
        export_params=True,
        opset_version=17,
        do_constant_folding=True,
        input_names=['x'],
    )
    
    pt2_path = f"{child_mod_type}.pt2"
    torch.export.save(exported, pt2_path)
    
    model = exported.module()
    
    # # generate QnnQuantizer
    # quantizer = make_quantizer()
    # prepared = prepare_pt2e(model, quantizer)
    # # perform calibration
    # prepared(*sample_input)
    # converted = convert_pt2e(prepared)
     
    # setup compile spec for HTP backend
    backend_options = generate_htp_compiler_spec(use_fp16=True)
    compiler_specs = generate_qnn_executorch_compiler_spec(
        soc_model=QcomChipset.SM8650,
        backend_options=backend_options,
        # profile=True,
        # online_prepare=True,
        # optrace=True,
    )
    # lower to QNN ExecuTorch Backend
    edge_prog_mgr = to_edge_transform_and_lower_to_qnn(
        module=model_class,
        inputs=sample_input,
        compiler_specs=compiler_specs,
    )

    # for inspector API
    edge_copy = copy.deepcopy(edge_prog_mgr)
    # store pte file
    exec_prog = edge_prog_mgr.to_executorch()
    pte_name = f"{child_mod_type}.pte"
    with open(pte_name, "wb") as f:
        exec_prog.write_to_file(f)
 
    # setup ADB for on-device execution
    adb = SimpleADB(
        qnn_sdk=os.getenv("QNN_SDK_ROOT"),
        build_path="/workspace/executorch/build-android",
        pte_path=pte_name,
        workspace="/data/local/tmp/simple_example",
        device_id="30.21.26.11:5555",
        soc_model="SM8650",
    )
    # binaries_trace = generate_optrace(
    #     "/workspace/visualizer_Conv2d_h_1",
    #     get_soc_to_chipset_map()["SM8650"],
    #     adb,
    #     pte_name,
    #     sample_input,
    # )

    user_inputs, input_list = [], ""
    with open("input_list", "r") as f:
        for line in f.read().split("\n")[:-1]:
            inputs, input_names = [], ""
            for data in line.split(" "):
                input_names += f"{Path(data).stem}.raw "
                inputs.append(torch.load(data, weights_only=True))
            user_inputs.append(inputs)
            input_list += input_names.strip() + "\n"
    adb.push(inputs=user_inputs, input_list=input_list)

    adb.execute()
 
    # pull etdump back and display the statistics
    adb.pull_etdump(".")
    generate_etrecord("etrecord.bin", edge_copy, exec_prog)
    inspector = Inspector(
        etdump_path="etdump.etdp",
        etrecord="etrecord.bin",
        source_time_scale=TimeScale.CYCLES,
        target_time_scale=TimeScale.CYCLES,
    )
    df = inspector.to_dataframe()
    print(df)
    inspector.print_data_tabular()
    inspector.save_data_to_tsv("Conv2d_1_w.tsv")

Versions

Collecting environment information...
PyTorch version: 2.9.0.dev20250725+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.5 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 14.0.0-1ubuntu1.1
CMake version: version 3.31.6
Libc version: glibc-2.35

Python version: 3.10.0 (default, Mar  3 2022, 09:58:08) [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.241-1-tlinux4-0017.5-x86_64-with-glibc2.35
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Address sizes:                   46 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          10
On-line CPU(s) list:             0-9
Vendor ID:                       GenuineIntel
BIOS Vendor ID:                  Smdbmds
Model name:                      Intel(R) Xeon(R) Platinum 8255C CPU @ 2.50GHz
BIOS Model name:                 3.0
CPU family:                      6
Model:                           85
Thread(s) per core:              2
Core(s) per socket:              5
Socket(s):                       1
Stepping:                        5
BogoMIPS:                        4988.28
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single pti fsgsbase bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 arat avx512_vnni
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       320 KiB (10 instances)
L1i cache:                       320 KiB (10 instances)
L2 cache:                        20 MiB (5 instances)
L3 cache:                        35.8 MiB (1 instance)
NUMA node(s):                    1
NUMA node0 CPU(s):               0-9
Vulnerability Itlb multihit:     KVM: Vulnerable
Vulnerability L1tf:              Mitigation; PTE Inversion
Vulnerability Mds:               Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Meltdown:          Mitigation; PTI
Vulnerability Mmio stale data:   Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Retbleed:          Vulnerable
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown

Versions of relevant libraries:
[pip3] executorch==0.8.0a0+2c84f70
[pip3] numpy==2.2.6
[pip3] nvidia-cublas-cu12==12.6.4.1
[pip3] nvidia-cuda-cupti-cu12==12.6.80
[pip3] nvidia-cuda-nvrtc-cu12==12.6.77
[pip3] nvidia-cuda-runtime-cu12==12.6.77
[pip3] nvidia-cudnn-cu12==9.5.1.17
[pip3] nvidia-cufft-cu12==11.3.0.4
[pip3] nvidia-curand-cu12==10.3.7.77
[pip3] nvidia-cusolver-cu12==11.7.1.2
[pip3] nvidia-cusparse-cu12==12.5.4.2
[pip3] nvidia-cusparselt-cu12==0.6.3
[pip3] nvidia-nccl-cu12==2.26.2
[pip3] nvidia-nvjitlink-cu12==12.6.85
[pip3] nvidia-nvtx-cu12==12.6.77
[pip3] onnx==1.18.0
[pip3] onnxconverter-common==1.15.0
[pip3] onnxruntime==1.22.0
[pip3] onnxsim==0.4.36
[pip3] pytorch_tokenizers==0.1.0
[pip3] torch==2.9.0.dev20250725+cpu
[pip3] torchao==0.13.0+git2eb4f9762
[pip3] torchaudio==2.8.0.dev20250725+cpu
[pip3] torchdata==0.11.0
[pip3] torchsr==1.0.4
[pip3] torchsummary==1.5.1
[pip3] torchtune==0.6.1
[pip3] torchvision==0.24.0.dev20250725+cpu
[pip3] triton==3.3.1
[conda] executorch                0.8.0a0+2c84f70          pypi_0    pypi
[conda] numpy                     2.2.6                    pypi_0    pypi
[conda] nvidia-cublas-cu12        12.6.4.1                 pypi_0    pypi
[conda] nvidia-cuda-cupti-cu12    12.6.80                  pypi_0    pypi
[conda] nvidia-cuda-nvrtc-cu12    12.6.77                  pypi_0    pypi
[conda] nvidia-cuda-runtime-cu12  12.6.77                  pypi_0    pypi
[conda] nvidia-cudnn-cu12         9.5.1.17                 pypi_0    pypi
[conda] nvidia-cufft-cu12         11.3.0.4                 pypi_0    pypi
[conda] nvidia-curand-cu12        10.3.7.77                pypi_0    pypi
[conda] nvidia-cusolver-cu12      11.7.1.2                 pypi_0    pypi
[conda] nvidia-cusparse-cu12      12.5.4.2                 pypi_0    pypi
[conda] nvidia-cusparselt-cu12    0.6.3                    pypi_0    pypi
[conda] nvidia-nccl-cu12          2.26.2                   pypi_0    pypi
[conda] nvidia-nvjitlink-cu12     12.6.85                  pypi_0    pypi
[conda] nvidia-nvtx-cu12          12.6.77                  pypi_0    pypi
[conda] pytorch-tokenizers        0.1.0                    pypi_0    pypi
[conda] torch                     2.9.0.dev20250725+cpu          pypi_0    pypi
[conda] torchao                   0.13.0+git2eb4f9762          pypi_0    pypi
[conda] torchaudio                2.8.0.dev20250725+cpu          pypi_0    pypi
[conda] torchdata                 0.11.0                   pypi_0    pypi
[conda] torchsr                   1.0.4                    pypi_0    pypi
[conda] torchsummary              1.5.1                    pypi_0    pypi
[conda] torchtune                 0.6.1                    pypi_0    pypi
[conda] torchvision               0.24.0.dev20250725+cpu          pypi_0    pypi
[conda] triton                    3.3.1                    pypi_0    pypi

cc @cccclai @cbilgin