Generalize canonicalization for padding, stride, dilation and output_padding

Author: hariprasadravi

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -4722,122 +4722,6 @@ OpFoldResult Aten_ShapeAsTensorOp::fold(FoldAdaptor adaptor) {
   return DenseElementsAttr::get(attrty, attrs);
 }
 
-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)");
-    }
-    int64_t requiredSpatialDims = weightRank - 2;
-
-    // Validate stride, padding, output_padding, and dilation are constant
-    // lists.
-    SmallVector<int64_t> strideInts;
-    if (!matchPattern(op.getStride(), m_TorchListOfConstantInts(strideInts))) {
-      return rewriter.notifyMatchFailure(op,
-                                         "non-const int stride unsupported!");
-    }
-    SmallVector<int64_t> paddingInts;
-    if (!matchPattern(op.getPadding(),
-                      m_TorchListOfConstantInts(paddingInts))) {
-      return rewriter.notifyMatchFailure(op,
-                                         "non-const int padding unsupported!");
-    }
-    SmallVector<int64_t> outputPaddingInts;
-    if (!matchPattern(op.getOutputPadding(),
-                      m_TorchListOfConstantInts(outputPaddingInts))) {
-      return rewriter.notifyMatchFailure(
-          op, "non-const int output_padding unsupported!");
-    }
-    SmallVector<int64_t> dilationInts;
-    if (!matchPattern(op.getDilation(),
-                      m_TorchListOfConstantInts(dilationInts))) {
-      return rewriter.notifyMatchFailure(op,
-                                         "non-const int dilation unsupported!");
-    }
-
-    // Canonicalization Logic: Only rewrite if padding provided is 1 element
-    // but the convolution requires 2 or 3 elements.
-    if (strideInts.size() == static_cast<size_t>(requiredSpatialDims)) {
-      return rewriter.notifyMatchFailure(op,
-                                         "stride is already fully specified");
-    }
-    if (paddingInts.size() == static_cast<size_t>(requiredSpatialDims)) {
-      return rewriter.notifyMatchFailure(op,
-                                         "padding is already fully specified");
-    }
-    if (outputPaddingInts.size() == static_cast<size_t>(requiredSpatialDims)) {
-      return rewriter.notifyMatchFailure(
-          op, "output_padding is already fully specified");
-    }
-    if (dilationInts.size() == static_cast<size_t>(requiredSpatialDims)) {
-      return rewriter.notifyMatchFailure(op,
-                                         "dialtion is already fully specified");
-    }
-
-    // Construct the new Padding List
-    // If user provided padding=[1], and we need 2 or 3 dims, we create
-    // padding=[1, 1] or padding = [1,1,1]
-    int64_t padVal = paddingInts[0];
-    Location loc = op.getLoc();
-
-    SmallVector<Value> newPaddingValues;
-    Value paddingConst = ConstantIntOp::create(
-        rewriter, loc, rewriter.getI64IntegerAttr(padVal));
-
-    for (int i = 0; i < requiredSpatialDims; ++i) {
-      newPaddingValues.push_back(paddingConst);
-    }
-
-    // Create the list construct op
-    auto newListOp = PrimListConstructOp::create(
-        rewriter, loc, Torch::ListType::get(rewriter.getType<Torch::IntType>()),
-        newPaddingValues);
-
-    // Replace the Op
-    // We create a new convolution op, keeping all operands the same except
-    // padding
-    rewriter.replaceOpWithNewOp<AtenConvolutionOp>(
-        op, op.getType(), op.getInput(), op.getWeight(), op.getBias(),
-        op.getStride(), newListOp.getResult(), op.getDilation(),
-        op.getTransposed(), op.getOutputPadding(), op.getGroups());
-
-    return success();
-  }
-};
-} // namespace
-
-//===----------------------------------------------------------------------===//
-// AtenConvolutionOp Registration
-//===----------------------------------------------------------------------===//
-void AtenConvolutionOp::getCanonicalizationPatterns(RewritePatternSet &results,
-                                                    MLIRContext *context) {
-  results.add<CanonicalizeConvolutionWithSingleIntTuple>(context);
-}
-
 //===----------------------------------------------------------------------===//
 // AtenIntTensorOp
 //===----------------------------------------------------------------------===//
@@ -6015,6 +5899,138 @@ 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> outputPaddingInts;
+    if (!matchPattern(op.getOutputPadding(),
+                      m_TorchListOfConstantInts(outputPaddingInts))) {
+      return rewriter.notifyMatchFailure(
+          op, "non-const int output_padding unsupported!");
+    }
+    SmallVector<int64_t, 3> dilationInts;
+    if (!matchPattern(op.getDilation(),
+                      m_TorchListOfConstantInts(dilationInts))) {
+      return rewriter.notifyMatchFailure(op,
+                                         "non-const int dilation unsupported!");
+    }
+
+    // Canonicalization Logic: Only rewrite if padding 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) && isCanonical(outputPaddingInts)) {
+      return rewriter.notifyMatchFailure(
+          op, "stride, padding, dialtion and outputPadding is already fully "
+              "specified");
+    }
+
+    expand(strideInts, requiredSpatialDims);
+    expand(paddingInts, requiredSpatialDims);
+    expand(dilationInts, requiredSpatialDims);
+    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])));
+
+      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);
+    auto outputPaddingList = Torch::PrimListConstructOp::create(
+        rewriter, loc, targetListType, cstOutputPadding);
+
+    // 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.getResult(), 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

@@ -1131,9 +1131,10 @@
     "Conv2dWithPaddingDilationStrideStaticModule_grouped",
     "Conv2dWithPaddingDilationStrideStaticModule_grouped_multiplier",
     "Convolution2DStaticModule_basic",
-    "Convolution2DSingleIntTuplePaddingModule_basic",
+    "Convolution2DSingleIntTupleModule_basic",
     "ConvolutionBackwardModule2DStatic_basic",
     "ConvolutionModule2DTransposeStridedStatic_basic",
+    "ConvolutionModule2DTransposeScalarTupleParams_basic",
     "Conv_Transpose1dStaticModule_basic",
     "Conv_Transpose2dStaticModule_basic",
     "Conv_Transpose3dStaticModule_basic",
@@ -2167,7 +2168,7 @@
     "Conv2dWithValidPaddingModule_basic",
     "Conv2dWithSamePaddingModule_basic",
     "Convolution2DStaticModule_basic",
-    "Convolution2DSingleIntTuplePaddingModule_basic",
+    "Convolution2DSingleIntTupleModule_basic",
     "CosineSimilarityStaticModule_basic",
     "DetachModule_basic",
     "DropoutEvalFloatModule_basic",
@@ -2906,6 +2907,7 @@
     "Conv2dWithSamePaddingModule_basic",
     "Conv2dWithValidPaddingModule_basic",
     "Conv3dModule_basic",
+    "Conv3dModuleScalarTupleParams_basic",
     "Conv3dWithSamePaddingModule_basic",
     "Conv3dWithValidPaddingModule_basic",
     "ConvolutionModule3DGroups_basic",
@@ -2921,7 +2923,9 @@
     "ConvolutionBackwardModule2DStrided_basic",
     "ConvolutionBackwardModule2D_basic",
     "ConvolutionModule2DGroups_basic",
+    "Convolution2DSingleIntTupleModule_basic",
     "ConvolutionModule2DTransposeNonUnitOutputPadding_basic",
+    "ConvolutionModule2DTransposeScalarTupleParams_basic",
     "ConvolutionModule2DTransposeStrided_basic",
     "ConvolutionModule2DTranspose_basic",
     # Error: onnx lowering,
@@ -3691,6 +3695,7 @@
     "Conv2dWithPaddingDilationStrideStaticModule_grouped",
     "Conv2dWithPaddingDilationStrideStaticModule_grouped_multiplier",
     "Conv3dModule_basic",
+    "Conv3dModuleScalarTupleParams_basic",
     "Conv3dWithSamePaddingModule_basic",
     "Conv3dWithValidPaddingModule_basic",
     "ConvTbcModule_basic",
@@ -4332,20 +4337,23 @@
     "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",