Fix DecomposeMultiDimSqueeze callback failure

[meenakshiramanathan1]

Fixes #642

[dgolubovicTT]

Can you give some context in PR why the change? What is it solving...

[meenakshiramanathan1]

@dgolubovicTT , this change fixes #642, the current status of this model is mentioned in the description here.

[dgolubovicTT]

From what I understand, if any dimension is tvm.tir.expr.Any dimension size cannot be determined at compile-time. In that case we don't decompose multidimsqueeze? Am I right?
If that is true, shouldn't we give some warning? Also, I am curious what this multi-dim squeeze lowers later in the compiler? What is the last compile stage of this densenet model with this change included?

[meenakshiramanathan1]

Yes, in that case we skip removal of squeeze containing dynamic shapes, warning can be added. Squeeze op is supported in mlir as mentioned here.
Densenet model is failing with unsupported ops - argwhere and scatter_elements after this change

[dgolubovicTT]

Ok, can you add a single squeeze test that will test this case, just to be sure everything after tvm will work.

[meenakshiramanathan1]

These test cases for squeeze op were added as a part of this PR when DecomposeMultiDimSqueeze callback was added.

[dgolubovicTT]

But is any of them actually testing this code path where this condition is not satisfied:

 if any([isinstance(dim, tvm.tir.expr.Any) for dim in pre.checked_type.shape]):

[meenakshiramanathan1]

I will try this and let you know.

[nvukobratTT]

@meenakshiramanathan1 can you provide a bit more context to this PR. Here are few questions I'm interested in:

1. The one @dgolubovicTT already asked, are we testing this path? If no, let's try to reproduce this with simple op test.
2. What are the results with and without your fix?
   • Before the change, you end up with failure in DecomposeMultiDimSqueeze
   • What about when you add this fix? Howe this multi-dim squeeze is further lowered in FFE, and latter on into MLIR? Can you share some reportify graphs for few compile stages (both TVM and FFE), as well a TTIR?
   • If you're unable to create simple unit test to cover this one, let's see screenshots from the original model.

[meenakshiramanathan1]

@meenakshiramanathan1 can you provide a bit more context to this PR. Here are few questions I'm interested in:

1. The one @dgolubovicTT already asked, are we testing this path? If no, let's try to reproduce this with simple op test.

2. What are the results with and without your fix?

   • Before the change, you end up with failure in DecomposeMultiDimSqueeze
   • What about when you add this fix? Howe this multi-dim squeeze is further lowered in FFE, and latter on into MLIR? Can you share some reportify graphs for few compile stages (both TVM and FFE), as well a TTIR?
   • If you're unable to create simple unit test to cover this one, let's see screenshots from the original model.

1. Tried testing this path with dynamic inputs but in the all cases of trying to create dynamic inputs , it is ending up with formation of argwhere which is not supported in forge yet
   For example,

def test_dynamic_multi_squeeze(test_device):
    class Multi_Squeeze(torch.nn.Module):
        def __init__(self):
            super().__init__()

        def forward(self, input):
            mask = input > 5
            indices = torch.nonzero(mask)
            squeezed_tensor = indices.squeeze(0)
            return squeezed_tensor

Here non_zero is mapped to argwhere which was one of the cases which was tried, so if we want to test this out further , argwhere op tvm decomposition should be brought in.

Before fix , graps are dumped till relay passes.

While after fix, it is progressing till forge_passes

2. With my fix , model is ending up with argwhere unsupported and without my fix, it would end up with this error.

3. Screenshots from original model which is causing dynamic shapes:

[nvukobratTT]

@meenakshiramanathan1

Tried testing this path with dynamic inputs but in the all cases of trying to create dynamic inputs , it is ending up with formation of argwhere which is not supported in forge yet

Thanks for the details! We're pushing the effort on outlining what and how our indexing should work.

Here non_zero is mapped to argwhere which was one of the cases which was tried, so if we want to test this out further , argwhere op tvm decomposition should be brought in.

Before fix , graps are dumped till relay passes.

Can you send actually screenshot of the graph? I'm interested in graph structure changes, rather than compile depth.

[meenakshiramanathan1]

@meenakshiramanathan1

Tried testing this path with dynamic inputs but in the all cases of trying to create dynamic inputs , it is ending up with formation of argwhere which is not supported in forge yet

Thanks for the details! We're pushing the effort on outlining what and how our indexing should work.

Here non_zero is mapped to argwhere which was one of the cases which was tried, so if we want to test this out further , argwhere op tvm decomposition should be brought in.
Before fix , graps are dumped till relay passes.

Can you send actually screenshot of the graph? I'm interested in graph structure changes, rather than compile depth.

I will send this information.