[JIT] Add variadic stack op (#63578)
authorMike Iovine <mikeiovine@fb.com>
Tue, 24 Aug 2021 15:19:38 +0000 (08:19 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Tue, 24 Aug 2021 15:20:54 +0000 (08:20 -0700)
commit1385f9fb12e6607c98d2d9d5edaaaab2bc07386f
tree5ecfeb80b3178839ae13f21303fb87898b9a1c45
parentf4aff3a346a0525e37d6071f318f7a4c54d5e1fb
[JIT] Add variadic stack op (#63578)

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

Added a new op `prim::VarStack` and a pass that transforms instances of `aten::stack(list, dim)` into `prim::VarStack(list[0], ..., list[n], dim)`. Also provided a JIT interpreter implementation.

Most of the implementation/tests are the same as `prim::VarConcat`.

Test Plan: `buck test caffe2/test/cpp/jit:jit -- TestStackOpt`

Reviewed By: navahgar

Differential Revision: D30426232

fbshipit-source-id: 9829a7db6e0a5038c9b7528c43c25b0c221aa2ce
aten/src/ATen/core/interned_strings.h
test/cpp/jit/CMakeLists.txt
test/cpp/jit/test_stack_opt.cpp [new file with mode: 0644]
torch/csrc/jit/passes/variadic_ops.cpp
torch/csrc/jit/passes/variadic_ops.h
torch/csrc/jit/runtime/register_prim_ops.cpp
torch/csrc/jit/runtime/static/ops.cpp