[Metal] Enable optimize_for_mobile on Linux (#46384)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/46384
Currently, the optimize_for_mobile binary only works on macOS, which is not very convenient to use. This diff introduces a new buck target that separates out the objective-c code. The goal here is to be able to export models for metal on linux machines.
ghstack-source-id: 114499418
Test Plan:
- set `enable_mpscnn` to 1 in pt_defs.bzl
- buck build //xplat/caffe2:optimize_for_mobile --show-full-output
- ./optimize_for_mobile --model="./model.pt" --backend="metal"
- CI
```
[taox@devvm2780.vll0 ~/mobilenetv2] ./optimize_for_mobile --model="./model.pt" --backend="metal"
pt_operator_library(
name = "old_op_library",
ops = [
"aten::Int.Tensor",
"aten::_convolution.deprecated",
"aten::adaptive_avg_pool2d",
"aten::add.Tensor",
"aten::addmm",
"aten::batch_norm",
"aten::dropout",
"aten::hardtanh_",
"aten::reshape",
"aten::size.int",
"aten::t",
"prim::NumToTensor.Scalar",
],
)
find output:
%411 : Tensor = aten::addmm(%self.classifier.1.bias, %input0.1, %694, %26, %26) # /Users/taox/anaconda/lib/python3.7/site-packages/torch/nn/functional.py:1674:0
insert: %817 : Tensor = metal::copy_to_host(%411)
pt_operator_library(
name = "new_op_library",
ops = [
"aten::adaptive_avg_pool2d",
"aten::add.Tensor",
"aten::addmm",
"aten::reshape",
"aten::size.int",
"metal::copy_to_host",
"metal_prepack::conv2d_run",
],
)
```
Reviewed By: linbinyu
Differential Revision: D24322017
fbshipit-source-id: 2b8062b069659bfec78ca4f9f9a5bb9dfbd693d2