[JIT] Freeze unrolls constant loops (#63614)
authorMike Iovine <mikeiovine@fb.com>
Tue, 7 Sep 2021 15:04:50 +0000 (08:04 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Tue, 7 Sep 2021 15:06:47 +0000 (08:06 -0700)
commit75b9e4a128496e9b2563078f62a2903f65a6d145
treecdc8aa74aad0069065ad8740b22fe474601974d9
parentadbcc819cd40deaa2755383815896d8c9dffb881
[JIT] Freeze unrolls constant loops (#63614)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/63614

There are a number of optimizations (`RemoveListMutation` in particular) that are tied to loop unrolling in `runOptimizations`. However, these were not invoked from `freeze_module` since the freezing pass should be idempotent.

This diff makes `runOptimizations` run `UnrollConstantLoops` instead of `UnrollLoops`. `freeze_module` is then able to run these optimizations.

Test Plan: Observed that `freeze_module` applies `RemoveListMutation`

Reviewed By: eellison

Differential Revision: D30437356

fbshipit-source-id: cba04bd958a48ad51b151aa3264f3d5bbb1fc2a4
test/onnx/test_pytorch_onnx_onnxruntime.py
test/onnx/test_utility_funs.py
torch/csrc/jit/passes/freeze_module.cpp
torch/csrc/jit/runtime/graph_executor.cpp
torch/csrc/jit/runtime/graph_executor_impl.h