whitelist ops we can build shapes for (#49125)
Summary:
Whitelist ops we can build shapes for.
Otherwise, `buildShapeExpressions` assumes that `aten::unsqueeze` is just a regular op.
```
[DUMP tensorexpr_fuser.cpp:329] buildShapeExpressions for
[DUMP tensorexpr_fuser.cpp:329] graph(%1 : float,
[DUMP tensorexpr_fuser.cpp:329] %3 : Float(50, 28, strides=[28, 1], requires_grad=0, device=cuda:0),
[DUMP tensorexpr_fuser.cpp:329] %8 : float,
[DUMP tensorexpr_fuser.cpp:329] %10 : Float(50, strides=[1], requires_grad=0, device=cuda:0)):
[DUMP tensorexpr_fuser.cpp:329] %11 : int = prim::Constant[value=1]()
[DUMP tensorexpr_fuser.cpp:329] %12 : Float(50, 1, strides=[1, 1], requires_grad=0, device=cuda:0) = aten::unsqueeze(%10, %11)
[DUMP tensorexpr_fuser.cpp:329] %9 : Float(50, 1, strides=[1, 1], requires_grad=0, device=cuda:0) = aten::mul(%12, %8)
[DUMP tensorexpr_fuser.cpp:329] %6 : Float(50, 28, strides=[28, 1], requires_grad=0, device=cuda:0) = aten::add(%3, %9, %11)
[DUMP tensorexpr_fuser.cpp:329] %2 : Float(50, 28, strides=[28, 1], requires_grad=0, device=cuda:0) = aten::div(%6, %1)
[DUMP tensorexpr_fuser.cpp:329] return (%2, %6, %9)
[DEBUG tensorexpr_fuser.cpp:347] Adding a mapping for %3 %162 : int[] = aten::size(%27)
[DEBUG tensorexpr_fuser.cpp:347] Adding a mapping for %10 %163 : int[] = aten::size(%23)
[DEBUG tensorexpr_fuser.cpp:402] Building sizes for %12 : Float(50, 1, strides=[1, 1], requires_grad=0, device=cuda:0) = aten::unsqueeze(%10, %11)
[DEBUG tensorexpr_fuser.cpp:405] Getting aten::size for %10
[DEBUG tensorexpr_fuser.cpp:402] Building sizes for %9 : Float(50, 1, strides=[1, 1], requires_grad=0, device=cuda:0) = aten::mul(%12, %8)
[DEBUG tensorexpr_fuser.cpp:405] Getting aten::size for %12
[DEBUG tensorexpr_fuser.cpp:402] Building sizes for %6 : Float(50, 28, strides=[28, 1], requires_grad=0, device=cuda:0) = aten::add(%3, %9, %11)
[DEBUG tensorexpr_fuser.cpp:405] Getting aten::size for %3
[DEBUG tensorexpr_fuser.cpp:405] Getting aten::size for %9
[DEBUG tensorexpr_fuser.cpp:402] Building sizes for %2 : Float(50, 28, strides=[28, 1], requires_grad=0, device=cuda:0) = aten::div(%6, %1)
[DEBUG tensorexpr_fuser.cpp:405] Getting aten::size for %6
[DEBUG tensorexpr_fuser.cpp:907] Inserting a typecheck guard for a node%156 : Float(50, 28, strides=[28, 1], requires_grad=0, device=cuda:0) = prim::TensorExprGroup[Subgraph=<Graph>](%3, %27, %16, %23)
[DUMP tensorexpr_fuser.cpp:463] After guarding fusion groups:
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/49125
Reviewed By: albanD
Differential Revision: D25926997
Pulled By: Krovatkin
fbshipit-source-id: f8041bbfc12be16c329754c6d16911d12aa352ef
