Support convolution with valid padding.

[sahas3]

Convolution created with valid padding produces the aten.convolution op in the following fashion:

module {
  func.func @main(%arg0: !torch.vtensor<[1,64,57],f32>) -> !torch.vtensor<[1,64,57],f32> attributes {torch.assume_strict_symbolic_shapes} {
    %false = torch.constant.bool false
    %int1 = torch.constant.int 1
    %0 = torch.vtensor.literal(dense<0.536443591> : tensor<1xf32>) : !torch.vtensor<[1],f32>
    %1 = torch.vtensor.literal(dense<-7.486820e-03> : tensor<1x1x1x1xf32>) : !torch.vtensor<[1,1,1,1],f32>
    %int0 = torch.constant.int 0
    %2 = torch.aten.unsqueeze %arg0, %int0 : !torch.vtensor<[1,64,57],f32>, !torch.int -> !torch.vtensor<[1,1,64,57],f32>
    %3 = torch.prim.ListConstruct %int1, %int1 : (!torch.int, !torch.int) -> !torch.list<int>
    %4 = torch.prim.ListConstruct %int0 : (!torch.int) -> !torch.list<int>
    %5 = torch.prim.ListConstruct %int1, %int1 : (!torch.int, !torch.int) -> !torch.list<int>
    %6 = torch.prim.ListConstruct %int0 : (!torch.int) -> !torch.list<int>
    %7 = torch.aten.convolution %2, %1, %0, %3, %4, %5, %false, %6, %int1 : !torch.vtensor<[1,1,64,57],f32>, !torch.vtensor<[1,1,1,1],f32>, !torch.vtensor<[1],f32>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool, !torch.list<int>, !torch.int -> !torch.vtensor<[1,1,64,57],f32>
    %8 = torch.aten.squeeze.dim %7, %int0 : !torch.vtensor<[1,1,64,57],f32>, !torch.int -> !torch.vtensor<[1,64,57],f32>
    return %8 : !torch.vtensor<[1,64,57],f32>
  }
}

Note that the padding input to aten.convolution is 1-element whereas the lowerings expect them to be same as number of spatial dims in the input. This results in hitting the assertion in https://github.com/sahas3/torch-mlir/blob/dc7a1ff7d9134758128a637dca976f72c2366e59/lib/Conversion/TorchToLinalg/Utils.cpp#L78 for the TorchToLinalg pass. The failure modes for lowering to tosa and stablehlo are different but stems from the same root cause.

[sahas3]

Hi @vivekkhandelwal1 and @qingyunqu tagging you as potential reviewers since you've contributed to the convolution lowering. Can you please take a look at this PR or add other reviewers? Thanks!

[sahas3]

Hello @vivekkhandelwal1 any thoughts on this PR? Also, I am a bit puzzled about the CI failure as I cannot reproduce it locally. Any suggestions on debugging it? Thanks!

[vivekkhandelwal1]

Hi @sahas3, sorry for the late reply. Why did you close this PR? Don't you need it now?

[sahas3]

Ah sorry @vivekkhandelwal1, I thought I commented when I closed the PR.

Looks like some pytorch behavior has changed which is now leaving the convolution op with same or valid padding as torch.aten.conv2d.padding. During fx.import_and_export this op is being left as torch.operator "torch.aten.conv2d.padding" resulting in "Lowering TorchFX IR -> Torch Backend IR" pipeline to fail. I think this can be handled by first adding "torch.aten.conv2d.padding" to torch_ods_gen.py so that lowering to Torch Backend IR doesn't fail. And then "torch.aten.conv2d.padding" can be expressed as ""torch.aten.conv2d". Is that the correct way forward?

[vivekkhandelwal1]

Ah sorry @vivekkhandelwal1, I thought I commented when I closed the PR.

Looks like some pytorch behavior has changed which is now leaving the convolution op with same or valid padding as torch.aten.conv2d.padding. During fx.import_and_export this op is being left as torch.operator "torch.aten.conv2d.padding" resulting in "Lowering TorchFX IR -> Torch Backend IR" pipeline to fail. I think this can be handled by first adding "torch.aten.conv2d.padding" to torch_ods_gen.py so that lowering to Torch Backend IR doesn't fail. And then "torch.aten.conv2d.padding" can be expressed as ""torch.aten.conv2d". Is that the correct way forward?

Yeah, the error description that you gave means that the op torch.aten.conv2d.padding is missing torch-mlir lowering. You've add to add the op lowering or decomposition to support this. You can follow the steps here: https://github.com/llvm/torch-mlir/blob/main/docs/add_ops.md or let me know if you face any issue in doing this.