Repository: pytorch
Issue: Adding a linear layer leads to failure of `optimize_for_mobile`
Labels: triage review,oncall: jit,module: mkldnn

### 🐛 Describe the bug

Hi! I'm trying to optimize my model with `optimize_for_mobile`, and find that simply adding an additional linear layer to the model will cause failure, though a model with only a linear layer works fine (which means the linear layer itself will not cause this issue). See the codes below.

```python
import torch
from torch.utils.mobile_optimizer import optimize_for_mobile
import torch.nn as nn

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(
            2, 1, kernel_size=(1,1),
        )
        self.linear = nn.Linear(2, 1)
        
    def forward(self, x):
        x = torch.nn.functional.pad(
            x, pad=[0, 1, 0, 1, 1, 0], mode="constant", value=0.5,
        )
        x = self.conv(x)
        x = self.linear(x) # <-- works fine if we remove this line
        return x

mod = MyModule()

x = torch.zeros((2,1,1,1), dtype=torch.float32)
out = mod(x)
print(f'eager: out = {out}') # <-- works fine

exported = torch.jit.trace(mod, [x])
exported = torch.jit.optimize_for_inference(exported) # <-- works fine
exported = optimize_for_mobile(exported) # <-- Exception here!

eout = exported(x)
print(f'JIT: eout = {eout}')

assert torch.allclose(out, eout)
```

<details><summary>Click to expand a long log!</summary>

```python
Traceback (most recent call last):
  File "/home/colin/code/test/nnsmith/backends/factory.py", line 92, in checked_exec
    return executable(input)
  File "/home/colin/code/test/nnsmith/backends/torchjit.py", line 50, in closure
    output = exported(*input_ts.values())
  File "/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
NotImplementedError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/torch/nn/modules/linear.py(114): forward
/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/torch/nn/modules/module.py(1118): _slow_forward
/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/torch/nn/modules/module.py(1130): _call_impl
/home/colin/code/test/nnsmith/materialize/torch/symbolnet.py(348): forward
/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/torch/nn/modules/module.py(1118): _slow_forward
/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/torch/nn/modules/module.py(1130): _call_impl
/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/torch/jit/_trace.py(967): trace_module
/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/torch/jit/_trace.py(750): trace
/home/colin/code/test/nnsmith/backends/torchjit.py(40): make_backend
/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/multipledispatch/dispatcher.py(435): __call__
/home/colin/code/test/nnsmith/backends/factory.py(65): checked_make_backend
/home/colin/code/test/nnsmith/backends/factory.py(69): checked_compile
/home/colin/code/test/nnsmith/backends/factory.py(112): checked_compile_and_exec
/home/colin/code/test/nnsmith/cli/model_exec.py(70): verify_testcase
/home/colin/code/test/nnsmith/cli/fuzz.py(159): validate_and_report
/home/colin/code/test/nnsmith/cli/fuzz.py(194): run
/home/colin/code/test/nnsmith/cli/fuzz.py(201): main
/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/hydra/core/utils.py(186): run_job
/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/hydra/_internal/hydra.py(119): run
/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/hydra/_internal/utils.py(453): <lambda>
/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/hydra/_internal/utils.py(213): run_and_report
/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/hydra/_internal/utils.py(452): _run_app
/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/hydra/_internal/utils.py(389): _run_hydra
/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/hydra/main.py(90): decorated_main
/home/colin/code/test/nnsmith/cli/fuzz.py(205): <module>
RuntimeError: Could not run 'prepacked::linear_clamp_prepack' with arguments from the 'MkldnnCPU' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'prepacked::linear_clamp_prepack' is only available for these backends: [Dense, UNKNOWN_TENSOR_TYPE_ID, QuantizedXPU, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, SparseCPU, SparseCUDA, SparseHIP, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, SparseVE, UNKNOWN_TENSOR_TYPE_ID, NestedTensorCUDA, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID].

CPU: registered at ../aten/src/ATen/native/xnnpack/RegisterOpContextClass.cpp:95 [kernel]
BackendSelect: fallthrough registered at ../aten/src/ATen/core/BackendSelectFallbackKernel.cpp:3 [backend fallback]
Python: registered at ../aten/src/ATen/core/PythonFallbackKernel.cpp:133 [backend fallback]
Named: registered at ../aten/src/ATen/core/NamedRegistrations.cpp:7 [backend fallback]
Conjugate: registered at ../aten/src/ATen/ConjugateFallback.cpp:18 [backend fallback]
Negative: registered at ../aten/src/ATen/native/NegateFallback.cpp:18 [backend fallback]
ZeroTensor: registered at ../aten/src/ATen/ZeroTensorFallback.cpp:86 [backend fallback]
ADInplaceOrView: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:64 [backend fallback]
AutogradOther: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:35 [backend fallback]
AutogradCPU: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:39 [backend fallback]
AutogradCUDA: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:47 [backend fallback]
AutogradXLA: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:51 [backend fallback]
AutogradMPS: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:59 [backend fallback]
AutogradXPU: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:43 [backend fallback]
AutogradHPU: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:68 [backend fallback]
AutogradLazy: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:55 [backend fallback]
Tracer: registered at ../torch/csrc/autograd/TraceTypeManual.cpp:295 [backend fallback]
AutocastCPU: fallthrough registered at ../aten/src/ATen/autocast_mode.cpp:481 [backend fallback]
Autocast: fallthrough registered at ../aten/src/ATen/autocast_mode.cpp:324 [backend fallback]
Batched: registered at ../aten/src/ATen/BatchingRegistrations.cpp:1064 [backend fallback]
VmapMode: fallthrough registered at ../aten/src/ATen/VmapModeRegistrations.cpp:33 [backend fallback]
Functionalize: registered at ../aten/src/ATen/FunctionalizeFallbackKernel.cpp:89 [backend fallback]
PythonTLSSnapshot: registered at ../aten/src/ATen/core/PythonFallbackKernel.cpp:137 [backend fallback]
```

</details>

Something interesting is that, code snippets below works fine.

1. Remove the linear layer in forward pass.

```python
    def forward(self, x):
        x = torch.nn.functional.pad(
            x, pad=[0, 1, 0, 1, 1, 0], mode="constant", value=0.5,
        )
        x = self.conv(x)
        # x = self.linear(x) # <-- works fine if we remove this line
        return x
```

2. Use a single linear layer.

```python
class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 1)
        
    def forward(self, x):
        x = self.linear(x)
        return x

mod = MyModule()

x = torch.zeros((2,1,1,2), dtype=torch.float32)
out = mod(x)
print(f'eager: out = {out}') # <-- works fine

exported = torch.jit.trace(mod, [x])
exported = torch.jit.optimize_for_inference(exported)
exported = optimize_for_mobile(exported) # <-- works fine

eout = exported(x)
print(f'JIT: eout = {eout}')

assert torch.allclose(out, eout)
```

3. While making no change to the original model, remove `optimize_for_inference`.

```python
exported = torch.jit.trace(mod, [x])
# exported = torch.jit.optimize_for_inference(exported) # <-- works fine if we remove this line
exported = optimize_for_mobile(exported) # <-- works fine if we remove the above line
```


### Versions

```python
PyTorch version: 1.12.1+cu116
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.1 LTS (x86_64)
GCC version: (Ubuntu 11.2.0-19ubuntu1) 11.2.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.9.13 (main, Aug 25 2022, 23:26:10)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-48-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 11.6.124
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090
GPU 2: NVIDIA GeForce RTX 3090

Nvidia driver version: 515.65.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.3
[pip3] torch==1.12.1+cu116
[pip3] torchaudio==0.12.1+cu116
[pip3] torchvision==0.13.1+cu116
[conda] numpy                     1.23.3                   pypi_0    pypi
[conda] torch                     1.12.1+cu116             pypi_0    pypi
[conda] torchaudio                0.12.1+cu116             pypi_0    pypi
[conda] torchvision               0.13.1+cu116             pypi_0    pypi
```

cc @gujinghui @PenghuiCheng @XiaobingSuper @jianyuh @VitalyFedyunin