[Static Runtime] Add pass to eliminate __getitem__/DictConstruct calls (#62429)
authorMike Iovine <mikeiovine@fb.com>
Fri, 13 Aug 2021 17:18:03 +0000 (10:18 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Fri, 13 Aug 2021 17:21:16 +0000 (10:21 -0700)
commit000e3a08817edec5681c8b10078f61ff185ce090
tree52a3e602656cc6b819cbb5859d506ca8e929f911
parentfcc1f87b6aacdfba9193691568bc405be42ac77e
[Static Runtime] Add pass to eliminate __getitem__/DictConstruct calls (#62429)

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

Introduce a new pass to eliminate calls to `prim::DictConstruct/aten::__getitem__`. Given a graph like this:
```
%2 : Dict = prim::DictConstruct(%key, %value)
%3 : Tensor = aten::__getitem__(%2, %key)
%4 : Tensor = op(%3)
```
This pass produces a graph like this (after dead code elimination):
```
%4 : Tensor = op(%value)
```

This optimization is applied in the static runtime.

Test Plan:
`buck test //caffe2/test:jit -- TestPeephole`

**local.forward performance summary**
About 3% runtime benefit. All `DictConstruct` calls optimized out, `__getitem__` calls reduced significantly (~50% of them are cut out)
P438354810

**local_request_only.forward performance summary**
About 14% runtime benefit. Again, all `DictConstruct` calls optimized out, 50% `__getitem__` calls removed.
P438359742

There is some variance with runtime measurements, so take these numbers with a grain of salt. Also note that the benefit does not exist in the shrunk model since there are no `DictConstruct` calls

Reviewed By: hlu1

Differential Revision: D29995087

fbshipit-source-id: f376376a46ff808115afd2d60446e5db8f6f752f
test/jit/test_peephole.py
tools/build_variables.bzl
torch/csrc/jit/passes/peephole.cpp
torch/csrc/jit/passes/peephole_dict_idioms.cpp [new file with mode: 0644]
torch/csrc/jit/passes/peephole_dict_idioms.h [new file with mode: 0644]
torch/csrc/jit/passes/peephole_list_idioms.h