pnnx: support PyTorch ExportedProgram .pt2 models

.github/workflows/pnnx.yml

@@ -56,9 +56,18 @@ jobs:
       run: |
         cd tools/pnnx
         mkdir build && cd build
-        cmake -DCMAKE_BUILD_TYPE=Release ..
+        cmake -DCMAKE_BUILD_TYPE=Release "-DCMAKE_RUNTIME_OUTPUT_DIRECTORY_RELEASE=${{ github.workspace }}/tools/pnnx/build/src" ..
         cmake --build . --config Release -j 4
 
+    - name: setup-libtorch-runtime
+      shell: bash
+      run: echo "$(python3 -c 'import os, torch; print(os.path.join(os.path.dirname(torch.__file__), "lib"))')" >> "$GITHUB_PATH"
+
+    - name: quick-pt2-test
+      run: |
+        cd tools/pnnx/build
+        ctest -C Release --output-on-failure -R "^test_exported_program(_ops|_pytree_roundtrip|_invalid)?$"
+
     - name: quick-test
       if: matrix.os != 'windows-latest'
       run: |
@@ -365,7 +374,19 @@ jobs:
         export MKL_NUM_THREADS=1
         export MKL_ENABLE_INSTRUCTIONS=SSE4_2
         cd tools/pnnx/build
-        ctest --output-on-failure -j 8
+        ctest --output-on-failure -j 8 -E '_pt2$'
+
+    - name: pt2-test
+      if: matrix.torch == '2.12.1'
+      run: |
+        export PATH=${{ env.PYTHONUSERBASE }}/bin:$PATH
+        chmod +x tools/pnnx/build/src/pnnx
+        export OMP_THREAD_LIMIT=1
+        export OMP_NUM_THREADS=1
+        export MKL_NUM_THREADS=1
+        export MKL_ENABLE_INSTRUCTIONS=SSE4_2
+        cd tools/pnnx/build
+        ctest --output-on-failure -j 8 -R '_pt2$'
 
     - name: python-pnnx
       run: |

tools/pnnx/README.md

@@ -31,7 +31,7 @@ PNNX tries to define a set of operators and a simple and easy-to-use format that
 9. [Model optimization](#pnnx-model-optimization)
 10. [Custom operator support](#pnnx-custom-operator)
 
-# Build TorchScript to PNNX converter
+# Build PyTorch model to PNNX converter
 
 1. Install PyTorch and TorchVision c++ library
 2. Build PNNX with cmake
@@ -56,12 +56,39 @@ mod = torch.jit.trace(net, x)
 mod.save("resnet18.pt")
 ```
 
-2. Convert TorchScript to PNNX
+Alternatively, with PyTorch 2.2 or later, export an `ExportedProgram`:
+
+```python
+import torch
+import torchvision.models as models
+
+net = models.resnet18(weights="DEFAULT").eval()
+x = torch.rand(1, 3, 224, 224)
+
+exported_program = torch.export.export(net, (x,))
+torch.export.save(exported_program, "resnet18.pt2")
+```
+
+2. Convert TorchScript or ExportedProgram to PNNX
 
 ```shell
 pnnx resnet18.pt inputshape=[1,3,224,224]
 ```
 
+or
+
+```shell
+pnnx resnet18.pt2
+```
+
+The ExportedProgram archive already contains input tensor metadata, so
+`inputshape` is normally not required for `.pt2` input.
+
+The generated `*_pnnx.py` also provides `export_exported_program(example_inputs=None)`.
+It returns a new `torch.export.ExportedProgram` and saves it as `*_pnnx.pt2`.
+Pass a tuple of example inputs to override the generated inputs; positional
+tuple/list input and output structures are preserved.
+
 Normally, you will get seven files
 
 ```resnet18.pnnx.param``` PNNX graph definition
@@ -85,7 +112,7 @@ Open https://netron.app/ in browser, and drag resnet18.pnnx.param into it.
 4. PNNX command line options
 
 ```
-Usage: pnnx [model.pt] [(key=value)...]
+Usage: pnnx [model.pt | model.pt2] [(key=value)...]
   pnnxparam=model.pnnx.param
   pnnxbin=model.pnnx.bin
   pnnxpy=model_pnnx.py
@@ -101,6 +128,7 @@ Usage: pnnx [model.pt] [(key=value)...]
   customop=/home/nihui/.cache/torch_extensions/fused/fused.so,...
   moduleop=models.common.Focus,models.yolo.Detect,...
 Sample usage: pnnx mobilenet_v2.pt inputshape=[1,3,224,224]
+              pnnx exported_model.pt2
               pnnx yolov5s.pt inputshape=[1,3,640,640] inputshape2=[1,3,320,320] device=gpu moduleop=models.common.Focus,models.yolo.Detect
 ```
 

tools/pnnx/src/CMakeLists.txt

@@ -114,6 +114,7 @@ set(pnnx_pass_level2_SRCS
     pass_level2/eliminate_size_numtotensor_int.cpp
     pass_level2/functionize.cpp
     pass_level2/fuse_constantlist.cpp
+    pass_level2/fuse_exported_rnn.cpp
 
     pass_level2/F_adaptive_avg_pool1d.cpp
     pass_level2/F_adaptive_avg_pool2d.cpp
@@ -196,6 +197,7 @@ set(pnnx_pass_level2_SRCS
     pass_level2/Tensor_index_put.cpp
     pass_level2/Tensor_masked_fill.cpp
     pass_level2/Tensor_new_empty.cpp
+    pass_level2/Tensor_new_full.cpp
     pass_level2/Tensor_new_ones.cpp
     pass_level2/Tensor_new_zeros.cpp
     pass_level2/Tensor_permute.cpp
@@ -276,6 +278,7 @@ set(pnnx_pass_level2_SRCS
     pass_level2/torch_repeat_interleave.cpp
     pass_level2/torch_roll.cpp
     pass_level2/torch_scatter_add.cpp
+    pass_level2/torch_select_scatter.cpp
     pass_level2/torch_slice_scatter.cpp
     pass_level2/torch_split.cpp
     pass_level2/torch_squeeze.cpp
@@ -286,12 +289,16 @@ set(pnnx_pass_level2_SRCS
     pass_level2/torch_tensor_split.cpp
     pass_level2/torch_tile.cpp
     pass_level2/torch_topk.cpp
+    pass_level2/torch_tril.cpp
     pass_level2/torch_transpose.cpp
     pass_level2/torch_unbind.cpp
     pass_level2/torch_unsqueeze.cpp
     pass_level2/torch_var.cpp
     pass_level2/torch_view_as_complex.cpp
     pass_level2/torch_view_as_real.cpp
+    pass_level2/torch_window.cpp
+    pass_level2/torchvision_DeformConv2d.cpp
+    pass_level2/torchvision_RoIAlign.cpp
     pass_level2/torch_where.cpp
     pass_level2/torch_zeros.cpp
     pass_level2/torch_zeros_like.cpp
@@ -322,6 +329,7 @@ set(pnnx_pass_level2_SRCS
     pass_level2/torchaudio_F_spectrogram.cpp
 
     pass_level2/nn_GRU.cpp
+    pass_level2/nn_ChannelShuffle.cpp
     pass_level2/nn_LSTM.cpp
     pass_level2/nn_RNN.cpp
 )
@@ -670,6 +678,7 @@ set(torch2pnnx_SRCS
     ${pnnx_pass_level0_SRCS}
     ${pnnx_pass_level1_SRCS}
 
+    load_exported_program.cpp
     load_torchscript.cpp
 )