[Torch] Canonicalize aten.convolution with single int tuple params (#4388)

This PR is a more robust fix for the issue captured in #4380 .

Essentially, `torch.ops.aten.convolution` lowering to `tosa`, `linalg`
and `stablehlo` fail if `stride`, `padding`, `dilation` or
`output_padding` is a tuple with a singleton element, but convolution
operates on 2 or 3 spatial dimensions.

In the failing case, `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()`.

In #4380, I attempted to fix just the `torch-to-tosa` pass, but I later
realised that this is a more general bug in all the backends for
multiple params (thanks to
#4380 (comment)).

### Fix:
I followed #4250 to canonicalize aten.convolution if it operates on 2 or
3 spatial dims but params are singleton. For example, if
`aten.convolution` is 2D but `padding == [0]`, we canonicalize it to
`padding == [0, 0]`

Author: hariprasadravi

include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td

@@ -7119,6 +7119,7 @@ def Torch_AtenConvolutionOp : Torch_Op<"aten.convolution", [
       printDefaultTorchOp(printer, *this, 9, 1);
     }
   }];
+  let hasCanonicalizer = 1;
 }
 
 def Torch_Aten_ConvolutionOp : Torch_Op<"aten._convolution", [

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -6,6 +6,7 @@
 // Also available under a BSD-style license. See LICENSE.
 //
 //===----------------------------------------------------------------------===//
+#include "llvm/ADT/SmallVector.h"
 #define DEBUG_TYPE "torch-mlir-torch-dialect"
 #include "torch-mlir/Dialect/Torch/IR/TorchOps.h"
 #include "torch-mlir/Dialect/Torch/Utils/Utils.h"
@@ -5898,6 +5899,160 @@ void AtenMaxPool3dOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
   patterns.add<CanonicalizeMaxPoolWithSingleIntTuple<AtenMaxPool3dOp>>(context);
 }
 
+namespace {
+class CanonicalizeConvolutionWithSingleIntTuple
+    : public OpRewritePattern<AtenConvolutionOp> {
+public:
+  using OpRewritePattern<AtenConvolutionOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AtenConvolutionOp op,
+                                PatternRewriter &rewriter) const override {
+
+    auto weight = op.getWeight();
+    auto weightType = dyn_cast<ValueTensorType>(weight.getType());
+
+    if (!weightType) {
+      return rewriter.notifyMatchFailure(op, "weight is not a vtensor");
+    }
+    auto optionalSizes = weightType.getOptionalSizes();
+    if (!optionalSizes.has_value()) {
+      return rewriter.notifyMatchFailure(op,
+                                         "unranked weight tensor unsupported!");
+    }
+
+    // The rank is the size of the dimensions array
+    int64_t weightRank = optionalSizes.value().size();
+
+    // We canonicalize Rank 4 (2D Conv) or Rank 5 (3D Conv).
+    if (weightRank < 4 || weightRank > 5) {
+      return rewriter.notifyMatchFailure(
+          op, "unsupported weight rank (must be 4 or 5)");
+    }
+    int requiredSpatialDims = weightRank - 2;
+
+    // Validate stride, padding, output_padding, and dilation are constant
+    // lists.
+    SmallVector<int64_t, 3> strideInts;
+    if (!matchPattern(op.getStride(), m_TorchListOfConstantInts(strideInts))) {
+      return rewriter.notifyMatchFailure(op,
+                                         "non-const int stride unsupported!");
+    }
+    SmallVector<int64_t, 3> paddingInts;
+    if (!matchPattern(op.getPadding(),
+                      m_TorchListOfConstantInts(paddingInts))) {
+      return rewriter.notifyMatchFailure(op,
+                                         "non-const int padding unsupported!");
+    }
+
+    SmallVector<int64_t, 3> dilationInts;
+    if (!matchPattern(op.getDilation(),
+                      m_TorchListOfConstantInts(dilationInts))) {
+      return rewriter.notifyMatchFailure(op,
+                                         "non-const int dilation unsupported!");
+    }
+
+    bool transposed;
+    if (!matchPattern(op.getTransposed(), m_TorchConstantBool(&transposed))) {
+      return rewriter.notifyMatchFailure(
+          op, "non-const int tranposed unsupported!");
+    }
+
+    SmallVector<int64_t, 3> outputPaddingInts;
+    if (!matchPattern(op.getOutputPadding(),
+                      m_TorchListOfConstantInts(outputPaddingInts))) {
+      return rewriter.notifyMatchFailure(
+          op, "non-const int output_padding unsupported!");
+    }
+
+    // Canonicalization Logic: Only rewrite if convolution attribute provided is
+    // 1 element but the convolution requires 2 or 3 elements.
+    auto isCanonical = [requiredSpatialDims](ArrayRef<int64_t> param) {
+      return param.size() == static_cast<size_t>(requiredSpatialDims);
+    };
+
+    if (isCanonical(strideInts) && isCanonical(paddingInts) &&
+        isCanonical(dilationInts)) {
+      return rewriter.notifyMatchFailure(
+          op, "stride, padding, dialtion and outputPadding is already fully "
+              "specified");
+    }
+
+    if (transposed && isCanonical(outputPaddingInts)) {
+      return rewriter.notifyMatchFailure(
+          op, "output_padding is already fully specified");
+    }
+
+    expand(strideInts, requiredSpatialDims);
+    expand(paddingInts, requiredSpatialDims);
+    expand(dilationInts, requiredSpatialDims);
+
+    if (transposed)
+      expand(outputPaddingInts, requiredSpatialDims);
+
+    // Construct the new List
+    // For example: If user provided padding=[1], and we need 2 or 3 dims, we
+    // create padding=[1, 1] or padding = [1,1,1]
+    Location loc = op.getLoc();
+    SmallVector<Value> cstPadding, cstStrides, cstDilation, cstOutputPadding;
+
+    for (auto dim : llvm::seq<int>(0, requiredSpatialDims)) {
+
+      cstStrides.push_back(Torch::ConstantIntOp::create(
+          rewriter, loc, rewriter.getI64IntegerAttr(strideInts[dim])));
+
+      cstPadding.push_back(Torch::ConstantIntOp::create(
+          rewriter, loc, rewriter.getI64IntegerAttr(paddingInts[dim])));
+
+      cstDilation.push_back(Torch::ConstantIntOp::create(
+          rewriter, loc, rewriter.getI64IntegerAttr(dilationInts[dim])));
+
+      if (transposed)
+        cstOutputPadding.push_back(Torch::ConstantIntOp::create(
+            rewriter, loc, rewriter.getI64IntegerAttr(outputPaddingInts[dim])));
+    }
+
+    auto targetListType =
+        Torch::ListType::get(Torch::IntType::get(op->getContext()));
+
+    // Create the list construct op
+    auto stridesList = Torch::PrimListConstructOp::create(
+        rewriter, loc, targetListType, cstStrides);
+    auto paddingList = Torch::PrimListConstructOp::create(
+        rewriter, loc, targetListType, cstPadding);
+    auto dilationsList = Torch::PrimListConstructOp::create(
+        rewriter, loc, targetListType, cstDilation);
+
+    Value outputPaddingList;
+    if (transposed) {
+      outputPaddingList = Torch::PrimListConstructOp::create(
+          rewriter, loc, targetListType, cstOutputPadding);
+    } else {
+      outputPaddingList = op.getOutputPadding();
+    }
+
+    // Replace the Op
+    // We create a new convolution op, keeping all operands the same except
+    // stride, padding,dilation, and output_padding which are now fully
+    // specified
+    rewriter.replaceOpWithNewOp<AtenConvolutionOp>(
+        op, op.getType(), op.getInput(), op.getWeight(), op.getBias(),
+        stridesList.getResult(), paddingList.getResult(),
+        dilationsList.getResult(), op.getTransposed(), outputPaddingList,
+        op.getGroups());
+
+    return success();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// AtenConvolutionOp Registration
+//===----------------------------------------------------------------------===//
+void AtenConvolutionOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                                    MLIRContext *context) {
+  results.add<CanonicalizeConvolutionWithSingleIntTuple>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // AtenLinalgCrossOp
 //===----------------------------------------------------------------------===//

projects/pt1/e2e_testing/xfail_sets.py

@@ -1130,8 +1130,10 @@
     "Conv2dWithPaddingDilationStrideStaticModule_grouped",
     "Conv2dWithPaddingDilationStrideStaticModule_grouped_multiplier",
     "Convolution2DStaticModule_basic",
+    "Convolution2DSingleIntTupleModule_basic",
     "ConvolutionBackwardModule2DStatic_basic",
     "ConvolutionModule2DTransposeStridedStatic_basic",
+    "ConvolutionModule2DTransposeScalarTupleParams_basic",
     "Conv_Transpose1dStaticModule_basic",
     "ConstantPad2dStaticModule_basic",
     "ConstantPadNdModule_basic",
@@ -2163,6 +2165,7 @@
     "Conv2dWithValidPaddingModule_basic",
     "Conv2dWithSamePaddingModule_basic",
     "Convolution2DStaticModule_basic",
+    "Convolution2DSingleIntTupleModule_basic",
     "Conv3dModule_basic",
     "Conv3dWithSamePaddingModule_basic",
     "Conv3dWithValidPaddingModule_basic",
@@ -2912,6 +2915,13 @@
     "Conv2dWithPaddingModule_basic",
     "Conv2dWithSamePaddingModule_basic",
     "Conv2dWithValidPaddingModule_basic",
+    "Conv3dModule_basic",
+    "Conv3dModuleScalarTupleParams_basic",
+    "Conv3dWithSamePaddingModule_basic",
+    "Conv3dWithValidPaddingModule_basic",
+    "ConvolutionModule3DGroups_basic",
+    "ConvolutionModule3DGroupsStrided_basic",
+    "ConvolutionModule3DGroupsDilated_basic",
     "ConvTbcModule_basic",
     "ConvTranspose2DQInt8_basic",
     "Conv_Transpose2dModule_basic",
@@ -2922,7 +2932,9 @@
     "ConvolutionBackwardModule2DStrided_basic",
     "ConvolutionBackwardModule2D_basic",
     "ConvolutionModule2DGroups_basic",
+    "Convolution2DSingleIntTupleModule_basic",
     "ConvolutionModule2DTransposeNonUnitOutputPadding_basic",
+    "ConvolutionModule2DTransposeScalarTupleParams_basic",
     "ConvolutionModule2DTransposeStrided_basic",
     "ConvolutionModule2DTranspose_basic",
     # Error: onnx lowering,
@@ -4310,18 +4322,24 @@
     "Conv2dWithPaddingModule_basic",
     "Conv2dWithSamePaddingModule_basic",
     "Conv2dWithValidPaddingModule_basic",
+    "Conv3dModule_basic",
+    "Conv3dModuleScalarTupleParams_basic",
+    "Conv3dWithSamePaddingModule_basic",
+    "Conv3dWithValidPaddingModule_basic",
     "ConvTbcModule_basic",
     "ConvTranspose2DQInt8_basic",
     "Conv_Transpose2dModule_basic",
     "Convolution2DModule_basic",
     "Convolution2DStridedModule_basic",
+    "Convolution2DSingleIntTupleModule_basic",
     "ConvolutionBackwardModule2DPadded_basic",
     "ConvolutionBackwardModule2DStatic_basic",
     "ConvolutionBackwardModule2DStrided_basic",
     "ConvolutionBackwardModule2D_basic",
     "ConvolutionModule2DGroups_basic",
     "ConvolutionModule2DTransposeNonUnitOutputPadding_basic",
     "ConvolutionModule2DTransposeStridedStatic_basic",
+    "ConvolutionModule2DTransposeScalarTupleParams_basic",
     "ConvolutionModule2DTransposeStrided_basic",
     "ConvolutionModule2DTranspose_basic",
     "ConvolutionModule2DGroupedTranspose_basic",

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py

@@ -612,7 +612,8 @@ def emit_with_mutating_variants(key, **kwargs):
         "aten::conv_tbc_backward : (Tensor, Tensor, Tensor, Tensor, int) -> (Tensor, Tensor, Tensor)"
     )
     emit(
-        "aten::convolution : (Tensor, Tensor, Tensor?, int[], int[], int[], bool, int[], int) -> (Tensor)"
+        "aten::convolution : (Tensor, Tensor, Tensor?, int[], int[], int[], bool, int[], int) -> (Tensor)",
+        has_canonicalizer=True,
     )
     emit(
         "aten::_convolution : (Tensor, Tensor, Tensor?, int[], int[], int[], bool, int[], int, bool, bool, bool, bool) -> (Tensor)"

projects/pt1/python/torch_mlir_e2e_test/test_suite/conv.py

@@ -304,6 +304,37 @@ def Convolution2DStaticModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(3, 3, 10, 10), tu.rand(3, 3, 2, 2))
 
 
+class Convolution2DSingleIntTupleModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([3, 3, 10, 10], torch.float32, True),
+            ([3, 3, 2, 2], torch.float32, True),
+        ]
+    )
+    def forward(self, inputVec, weight):
+        return torch.ops.aten.convolution(
+            inputVec,
+            weight,
+            bias=None,
+            stride=(1,),
+            padding=(0,),
+            dilation=(1,),
+            transposed=False,
+            output_padding=[0, 0],
+            groups=1,
+        )
+
+
+@register_test_case(module_factory=lambda: Convolution2DSingleIntTupleModule())
+def Convolution2DSingleIntTupleModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(3, 3, 10, 10), tu.rand(3, 3, 2, 2))
+
+
 class Convolution2DStridedModule(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -901,6 +932,39 @@ def ConvolutionModule2DTransposeNonUnitOutputPadding_basic(module, tu: TestUtils
     module.forward(tu.rand(1, 2, 4, 4), tu.rand(2, 2, 3, 3))
 
 
+class ConvolutionModule2DTransposeScalarTupleParams(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([5, 2, 5, 6], torch.float32, True),
+            ([2, 5, 2, 2], torch.float32, True),
+        ]
+    )
+    def forward(self, inputVec, weight):
+        return torch.ops.aten.convolution(
+            inputVec,
+            weight,
+            bias=None,
+            stride=(1,),
+            padding=(1,),
+            dilation=(1,),
+            transposed=True,
+            output_padding=(0,),
+            groups=1,
+        )
+
+
+@register_test_case(
+    module_factory=lambda: ConvolutionModule2DTransposeScalarTupleParams()
+)
+def ConvolutionModule2DTransposeScalarTupleParams_basic(module, tu: TestUtils):
+    module.forward(tu.rand(5, 2, 5, 6), tu.rand(2, 5, 2, 2))
+
+
 class Conv_Transpose1dModule(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -1569,6 +1633,39 @@ def Conv3dWithValidPaddingModule_basic(module, tu: TestUtils):
     module.forward(inputVec, weight, bias)
 
 
+class Conv3dModuleScalarTupleParams(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([-1, -1, -1, -1, -1], torch.float32, True),
+            ([-1, -1, -1, -1, -1], torch.float32, True),
+            ([-1], torch.float32, True),
+        ]
+    )
+    def forward(self, inputVec, weight, bias):
+        return torch.ops.aten.conv3d(
+            inputVec,
+            weight,
+            bias=bias,
+            stride=(1,),
+            padding=(0,),
+            dilation=(1,),
+            groups=1,
+        )
+
+
+@register_test_case(module_factory=lambda: Conv3dModuleScalarTupleParams())
+def Conv3dModuleScalarTupleParams_basic(module, tu: TestUtils):
+    inputVec = tu.rand(2, 2, 6, 6, 6)
+    weight = torch.randn(8, 2, 3, 3, 3)
+    bias = torch.randn(8)
+    module.forward(inputVec, weight, bias)
+
+
 class ConvTbcModule(torch.nn.Module):
     def __init__(self):
         super().__init__()