✨[Converter] Support symbolic shape in shape Tensor for aten::expand

[apbose]

Feature Implementation:

• Implement symbolic shape in size Tensor for aten::expand
• This addresses Issue 🐛 [Bug- FX] Matrix multiplication FX dynamo test failures #1827
• Create symint wrapper for the above test case so that it is handled in compile time and the shape dimension acts as a placeholder

[apbose]

Update-
The same occurs in aten::chunk case. aten::matmul and aten::chunk consume tensors in
Synopsis:

1. aten::chunk

 %arg0 : [#users=2] = placeholder[target=arg0]
    %sym_size : [#users=1] = call_function[target=torch.ops.aten.sym_size](args = (%arg0, 0), kwargs = {})
    %add : [#users=1] = call_function[target=operator.add](args = (%sym_size, 6), kwargs = {})
    %sub : [#users=1] = call_function[target=operator.sub](args = (%add, 1), kwargs = {})
    %floordiv : [#users=1] = call_function[target=operator.floordiv](args = (%sub, 6), kwargs = {})
    %split_tensor : [#users=6] = call_function[target=torch.ops.aten.split.Tensor](args = (%arg0, %floordiv), kwargs = {})
    %getitem : [#users=1] = call_function[target=operator.getitem](args = (%split_tensor, 0), kwargs = {})
    %getitem_1 : [#users=1] = call_function[target=operator.getitem](args = (%split_tensor, 1), kwargs = {})
    %getitem_2 : [#users=1] = call_function[target=operator.getitem](args = (%split_tensor, 2), kwargs = {})
    %getitem_3 : [#users=1] = call_function[target=operator.getitem](args = (%split_tensor, 3), kwargs = {})
    %getitem_4 : [#users=1] = call_function[target=operator.getitem](args = (%split_tensor, 4), kwargs = {})
    %getitem_5 : [#users=1] = call_function[target=operator.getitem](args = (%split_tensor, 5), kwargs = {})
    return [getitem, getitem_1, getitem_2, getitem_3, getitem_4, getitem_5]

Present status:

class Chunk(nn.Module):
            def __init__(self):
                super().__init__()
                self.chunk = 6
                self.dim = 0
            def forward(self, input):
                return torch.ops.aten.chunk(input, self.chunk, self.dim)
            
if __name__ == "__main__":
    chunk = Chunk()
    input = [torch.randn(12)

The lower_basic_pass_aten passes through lowering passes such that torch.ops.aten.split.Tensor is replaced by aten_compose_chunk. Two cases-

    #symbolic trace
    fx_module_aten_symbolic = torch.fx.symbolic_trace(chunk)

    #aten trace
    fx_module_aten, _ = aten_tracer.trace(chunk, input)
 
   #No lowering
    graph_no_lower =  convert_module(fx_module_aten.graph, input)
   #Lowering
    compose_chunk_graph = compose_chunk(fx_module_aten)
    out = compose_chunk_graph.graph_module(input[0])
    graph = convert_module(compose_chunk_graph.graph_module, input)

No lowering
Converter for torch.ops.split.Tensor is used and it encounters this case

TypeError: unsupported operand type(s) for +: 'int' and 'tensorrt.tensorrt.ITensor'

While executing %split_tensor : [#users=6] = call_function[target=torch.ops.aten.split.Tensor](args = (%arg0, %floordiv), kwargs = {_itensor_to_tensor_meta: {<tensorrt.tensorrt.ITensor object at 0x7f280d54ff30>: None, <tensorrt.tensorrt.ITensor object at 0x7f280d285c30>: None, <tensorrt.tensorrt.ITensor object at 0x7f280d287330>: None, <tensorrt.tensorrt.ITensor object at 0x7f280e6b9ef0>: None, <tensorrt.tensorrt.ITensor object at 0x7f280d17cb30>: None}})
Original traceback:
  File "chunk.py", line 19, in forward
    return torch.ops.aten.chunk(input, self.chunk, self.dim)

Lowering
Converter for torch
Chunk operator for lowering pass @tensorrt_converter(aten_compose_chunk) is used where aten_compose_chunk implemented is the chunk operator.

Code error
ValueError: tree_unflatten(values, spec): `values` has length 1 but the spec refers to a pytree that holds 6 items (TreeSpec(list, None, [*,
                      *,
                      *,
                      *,
                      *,
                      *])).

Note: out = compose_chunk_graph.graph_module(input[0]) work good and the interpretor complains

Workaround: Try to replace the getitem ops

2. aten::matmul

   %arg0 : [#users=4] = placeholder[target=arg0]
   %sym_size : [#users=4] = call_function[target=torch.ops.aten.sym_size](args = (%arg0, 0), kwargs = {})
   %_param_constant0 : [#users=1] = get_attr[target=_param_constant0]
   %expand_default : [#users=1] = call_function[target=torch.ops.aten.expand.default](args = (%_param_constant0, [%sym_size, 2, 2, 3]), kwargs = {})
   %arg0 : [#users=1] = placeholder[target=arg0]

Present status

• The matmul with inputs having rank more than 2 encounter bmm, with which lowering ops are encountered
  To do:
• All cases of bmm should be considered in Lowering
• Implementing the batched matmul operations
• Run time tensors will lead to errors with Tensor arguments

[github-actions]

This issue has not seen activity for 90 days, Remove stale label or comment or this will be closed in 10 days