aten.maxpool2d shape resolution failure

[sjarus]

First some good news - we're pretty close to getting Torch->TOSA working for ResNet18 static network form. Now the bad news: the primary blocker is the strange failure to resolve the aten.maxpool2d op. This can be reproduced with the following simple e2e test:

class MaxPool2dStaticModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.mp2d = torch.nn.MaxPool2d(kernel_size=[3, 3],
                                       stride=[2, 2],
                                       padding=[1, 1],
                                       dilation=[1, 1])
 
   @export
    @annotate_args([
        None,
        ([1, 64, 112, 112], torch.float32, True),
    ])
    def forward(self, x):
        return self.mp2d(x)
@register_test_case(module_factory=lambda: MaxPool2dStaticModule())
def MaxPool2dStaticModule_basic(module, tu: TestUtils):
    module.forward(tu.rand(1, 64, 112, 112))

This emits:

%int3 = torch.constant.int 3 loc(#loc1)
%int1 = torch.constant.int 1 loc(#loc2)
%int2 = torch.constant.int 2 loc(#loc3)
%false = torch.constant.bool false loc(#loc4)
%0 = torch.prim.ListConstruct %int3, %int3 : (!torch.int, !torch.int) -> !torch.list<!torch.int> loc(#loc5)
%1 = torch.prim.ListConstruct %int2, %int2 : (!torch.int, !torch.int) -> !torch.list<!torch.int> loc(#loc5)
%2 = torch.prim.ListConstruct %int1, %int1 : (!torch.int, !torch.int) -> !torch.list<!torch.int> loc(#loc5)
%3 = torch.prim.ListConstruct %int1, %int1 : (!torch.int, !torch.int) -> !torch.list<!torch.int> loc(#loc5)
%4 = torch.aten.max_pool2d %arg0, %0, %1, %2, %3, %false : !torch.vtensor<[1,64,112,112],f32>, !torch.list<!torch.int>, !torch.list<!torch.int>, !torch.list<!torch.int>, !torch.list<!torch.int>, !torch.bool -> !torch.vtensor<[1,64,?,?],f32> loc(#loc6)

This is strange, since I expected #502 would fix it, and looking at the code, there's no reason why it would fail to do so. It appears to bail when converting op.stride() , but it's not clear why, since I'm able to independently legalize maxpool to TOSA and read stride() during that time.

Even more interestingly, the problem seems to manifest itself only in maxpool but not conv2d. I'm digging into the code further, but any insights would help.

@ramiro050 @cathyzhyi @ljfitz

[ljfitz]

I've seen this and at the time didn't understand where it could come from, but noticed that using --torch-refine-types --torch-refine-public-return on the output resolved the static shape and had not looked further at that point.

Since I've noticed that some ir optimization happens after the point RefineTypes initially runs, and suspect that if we look at the ir in this case there will still be something that inhibits the constant matching. I'll try to have a look a little later.

[ljfitz]

I just checked - there is a torch.prim.If in the ir going into RefineTypes:

module attributes {torch.debug_module_name = "MaxPool2dStaticModule"} {
  func @forward(%arg0: !torch.vtensor<[1,64,112,112],f32>) -> !torch.tensor {
    %int3 = torch.constant.int 3
    %int1 = torch.constant.int 1
    %int2 = torch.constant.int 2
    %none = torch.constant.none
    %false = torch.constant.bool false
    %0 = torch.tensor_static_info_cast %arg0 : !torch.vtensor<[1,64,112,112],f32> to !torch.vtensor
    %1 = torch.copy.to_tensor %0 : !torch.tensor
    %2 = torch.prim.ListConstruct %int3, %int3 : (!torch.int, !torch.int) -> !torch.list<!torch.int>
    %3 = torch.prim.ListConstruct %int2, %int2 : (!torch.int, !torch.int) -> !torch.list<!torch.int>
    %4 = torch.prim.ListConstruct %int1, %int1 : (!torch.int, !torch.int) -> !torch.list<!torch.int>
    %5 = torch.prim.ListConstruct %int1, %int1 : (!torch.int, !torch.int) -> !torch.list<!torch.int>
    %6 = torch.derefine %3 : !torch.list<!torch.int> to !torch.optional<!torch.list<!torch.int>>
    %7 = torch.aten.__is__ %6, %none : !torch.optional<!torch.list<!torch.int>>, !torch.none -> !torch.bool
    %8 = torch.prim.If %7 -> (!torch.list<!torch.int>) {
      %12 = torch.prim.ListConstruct  : () -> !torch.list<!torch.int>
      torch.prim.If.yield %12 : !torch.list<!torch.int>
    } else {
      %12 = torch.prim.unchecked_cast %6 : !torch.optional<!torch.list<!torch.int>> -> !torch.list<!torch.int>
      torch.prim.If.yield %12 : !torch.list<!torch.int>
    }
    %9 = torch.copy.to_vtensor %1 : !torch.vtensor
    %10 = torch.aten.max_pool2d %9, %2, %8, %4, %5, %false : !torch.vtensor, !torch.list<!torch.int>, !torch.list<!torch.int>, !torch.list<!torch.int>, !torch.list<!torch.int>, !torch.bool -> !torch.vtensor
    %11 = torch.copy.to_tensor %10 : !torch.tensor
    return %11 : !torch.tensor
  }
}

[sjarus]

Thank you for investigating, @ljfitz !

[sjarus]

This is working now. Thanks @ljfitz and @cathyzhyi .