[TOSA] Fix conv2d.padding decomposition to emit 2D padding for aten.convolution

[hariprasadravi]

torch.nn.Conv2d with padding = 'valid' generates a torch.ops.aten.conv2d.padding op in ExportedProgram, which is later decomposed to torch.ops.aten.convolution.default with a single padding value of [0] after running ep.run_decompositions().

This causes a compiler failure in torch-to-tosa because the code expects padding to have two elements.

Here is the print of the ExportedProgram

>>> print(prog)
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, p_0_weight: "f32[1, 1, 1, 1]", p_0_bias: "f32[1]", input: "f32[1, 1, 5, 5]"):
            input_1 = input             
            conv2d: "f32[1, 1, 5, 5]" = torch.ops.aten.conv2d.padding(input_1, p_0_weight, p_0_bias);  input_1 = p_0_weight = p_0_bias = None
            return (conv2d,)

>>> prog_decomposed = prog.run_decompositions()
>>> print(prog_decomposed)
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, p_0_weight: "f32[1, 1, 1, 1]", p_0_bias: "f32[1]", input: "f32[1, 1, 5, 5]"):
            input_1 = input
            convolution: "f32[1, 1, 5, 5]" = torch.ops.aten.convolution.default(input_1, p_0_weight, p_0_bias, [1, 1], [0], [1, 1], False, [0], 1);  input_1 = p_0_weight = p_0_bias = None
            return (convolution,)

Note: A previous, correct decomposition for a similar issue exists and is locked down by e2e tests (refer to relevant context like #3883). This fix aligns the conv2d.padding decomposition with that correct behavior.

Since locking this specific decomposition down with end-to-end (e2e) tests was problematic, a dedicated LIT test has been added to the test suite to ensure this specific padding regression cannot reoccur.

[hariprasadravi]

@sahas3 @sjarus Kindly review this PR.

[sjarus]

Is there an e2e test that exercises this functionality ?

[hariprasadravi]

Is there an e2e test that exercises this functionality ?

I did not write an e2e test that directly calls torch.nn.Conv2d with padding='valid' because such a test will not trigger this issue. This issue can only be reproduced if run_decompositions is run on the ExportedProgram before lowering to Torch IR.

In fact, an e2e test for valid padding already exists and was added as part of #3883 here but it does not catch this issue as ExportedProgram decompositions are not run as part of the test.

Is there another way to lock this down with an e2e test?

[zjgarvey]

I believe we run decompositions in the fx_importer tests. Let me double check.

[zjgarvey]

Oh, we do not seem to run decompositions...
https://github.com/llvm/torch-mlir/blob/main/projects%2Fpt1%2Fpython%2Ftorch_mlir_e2e_test%2Fconfigs%2Ffx_importer_backend.py#L140

This doesn't seem right. I think we almost certainly want to run certain decompositions. There should be a default list somewhere in the codebase.

Do you know in which versions of pytorch the issue referenced in this PR gets reproduced?

[sahas3]

Does linalg path need a similar fix? If so, it may be better to add a canonicalization pattern for torch.aten.convolution so that padding is extended to be 2 dimension for all TorchToXXX paths in one place.

[hariprasadravi]

torch-to-linalg asserts due the same assumption here:

torch-mlir/lib/Conversion/TorchToLinalg/Utils.cpp

Line 77 in b834f94

assert(static_cast<int64_t>(unpaddedDims + padding.size()) ==

torch-to-stablehlo also asserts due to the same assumption here:

torch-mlir/lib/Conversion/TorchToStablehlo/Linear.cpp

Line 788 in b834f94

assert(padding.size() == dilation.size() &&

So it looks like this is similar to #3885. Shall I proceed with a cononicalization pattern similar to the proposed fix in #4250 ?

Thoughts @sjarus and @zjgarvey ?

[hariprasadravi]

Do you know in which versions of pytorch the issue referenced in this PR gets reproduced?

I used 2.10.0.dev20251016+cpu to repro this issue.

[sahas3]

Oh, we do not seem to run decompositions... https://github.com/llvm/torch-mlir/blob/main/projects%2Fpt1%2Fpython%2Ftorch_mlir_e2e_test%2Fconfigs%2Ffx_importer_backend.py#L140

I think we do run some decompositions by default

torch-mlir/python/torch_mlir/extras/fx_decomp_util.py

Line 5 in b834f94

DEFAULT_DECOMPOSITIONS = [

and

torch-mlir/python/torch_mlir/fx.py

Line 103 in b834f94

if decomposition_table is None:

but the conv.padding variant is not in that list. It doesn't seem trivial to specify a non-default decomp list from the e2e list.