Glow: [GraphOptz] Improve EliminateSliceConcat optimization

[mciprian13]

Stumbled over this when compiling some GRU network:

So the slicing is done over dimension 0 (resulting in singular dimensions) and the concatenation along dimension 1 but in this particular case all the Slice and Concat node should be eliminated and replaced with a Reshape.

[Squadrick]

I'm working on this issue, I'll open a PR soon. This is the approach I'm taking:

EliminateSliceConcat only handles the case when the Slices and Concat are along the same dimension. I'm changing it to track all consecutive Slices along some dimension. If the dimension of slicing is the same as the dimension of concat, it'll use the pre-existing logic. If the dimension of slicing is different, the multiple consecutive Slices can be replaced with Slice + Reshape.

---

This approach will handle the case when the slices don't cover the entire input. We would still want to minimize the extraneous operators. For example, consider the above example but we only slice along the first three axes and an external input:

   /-> Slice ((0, 0, 0)..(1, 5, 3)) --\
PH --> Slice ((1, 0, 0)..(2, 5, 3)) ----> Concat(Dim=1) -> Output(shape=(1, 20, 3))
   \-> Slice ((2, 0, 0)..(3, 5, 3)) --/          |
                                                 |
                                   PH2(1, 5, 3) -|    

should produce:

PH -> Slice ((0, 0, 0)..(3, 5, 3)) -> Reshape((3, 5, 3) -> (1, 15, 3)) -> Concat(Dim=1) -> Output(shape=(1, 20, 3))
                                                                           |
                                                            PH2(1, 5, 3)--/

---

In the posted example (all 4 slices), the above graph will turn into:

PH(4, 5, 3) -> Slice ((0, 0, 0)...(4, 5, 3)) -> Reshape((4, 5, 3) -> (1, 20, 3)) -> Concat(Dim=1) -> Result

Since the Slice uses the entire input, and Concat only has a single input, it should be reduced to:

PH -> Reshape -> Result

@jfix71 Does GraphOptimizer handle the last two optimizations?

1. Slice across the entire input is a no-op, and should be removed.
2. Concat with a single input should be removed.

I looked at docs/Optimizations.md, but couldn't find the pertinent info, this is all I found:

  * Optimization of concat nodes

    This optimization merges multiple consequent concat nodes into a single concat
    node.

[jfix71]

@Squadrick Thanks for the explanation! Our Optimizations.md is definitely a bit out of date.

Slice across the entire input is a no-op, and should be removed.

I'm actually not sure about this one, but it would be very simple to add if it's missing.

Concat with a single input should be removed.

Yep, this one we definitely have here, which runs as part of OptimizeConcatNodes:
https://github.com/pytorch/glow/blob/94cd0b162a68e170aab97743b9c5c83963da8242/lib/Optimizer/GraphOptimizer/GraphOptimizer.cpp#L2432-L2435

[Squadrick]

Slice across the entire input is a no-op, and should be removed.

I'm actually not sure about this one, but it would be very simple to add if it's missing.

Let me check this. If the optimization doesn't exist, I'll create an issue and fix it in a subsequent PR.

Yep, this one we definitely have here, which runs as part of OptimizeConcatNodes:
https://github.com/pytorch/glow/blob/94cd0b162a68e170aab97743b9c5c83963da8242/lib/Optimizer/GraphOptimizer/GraphOptimizer.cpp#L2432-L2435

Oops, not sure how I missed this. Thanks for pointing it out!