Generalize MaxUnpool lowering including 2d case

llvm · Oct 9, 2024 · 6e8caa0 · 6e8caa0
1 parent 057f51d
commit 6e8caa0
Show file tree

Hide file tree

Showing 6 changed files with 82 additions and 23 deletions.
diff --git a/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td b/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
@@ -7159,6 +7159,33 @@ def Torch_AtenMaxPool2dOp : Torch_Op<"aten.max_pool2d", [
   }];
 }
 
+def Torch_AtenMaxUnpool2dOp : Torch_Op<"aten.max_unpool2d", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::max_unpool2d : (Tensor, Tensor, int[], int[], int[]) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    AnyTorchTensorType:$indices,
+    AnyTorchListOfTorchIntType:$output_size,
+    AnyTorchListOfTorchIntType:$stride,
+    AnyTorchListOfTorchIntType:$padding
+  );
+  let results = (outs
+    AnyTorchOptionalTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenMaxUnpool2dOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 5, 1);
+    }
+    void AtenMaxUnpool2dOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 5, 1);
+    }
+  }];
+}
+
 def Torch_AtenMaxPool2dWithIndicesOp : Torch_Op<"aten.max_pool2d_with_indices", [
     AllowsTypeRefinement,
     HasValueSemantics,

diff --git a/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp b/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp
@@ -3464,6 +3464,12 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
         Value paddingList = createConstantIntList(binder, rewriter, padding);
         Value stridesList = createConstantIntList(binder, rewriter, strides);
 
+        if (rank == 4) {
+          rewriter.replaceOpWithNewOp<Torch::AtenMaxUnpool2dOp>(
+              binder.op, resultType, data, indices, resultShapeList,
+              stridesList, paddingList);
+          return success();
+        }
         rewriter.replaceOpWithNewOp<Torch::AtenMaxUnpool3dOp>(
             binder.op, resultType, data, indices, resultShapeList, stridesList,
             paddingList);

diff --git a/lib/Conversion/TorchToLinalg/Pooling.cpp b/lib/Conversion/TorchToLinalg/Pooling.cpp
@@ -596,37 +596,51 @@ namespace {
 // input_size=2, output_size=5 and stride=2, kernel_size can be either 2 or 3).
 // What worse, without knowing kernel size we cannot even reliably detect such
 // cases and this conversion will just return invalid values.
-class ConvertAtenMaxUnpool3dOp final
-    : public OpConversionPattern<AtenMaxUnpool3dOp> {
-public:
-  using OpConversionPattern::OpConversionPattern;
-  LogicalResult
-  matchAndRewrite(AtenMaxUnpool3dOp op, OpAdaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const override {
+
+template <> struct DimensionTraits<AtenMaxUnpool2dOp> {
+  static constexpr int64_t Dim = 2;
+  // unused const variable warning suppression:
+  static_assert(Dim == Dim);
+};
+
+template <> struct DimensionTraits<AtenMaxUnpool3dOp> {
+  static constexpr int64_t Dim = 3;
+  // unused const variable warning suppression:
+  static_assert(Dim == Dim);
+};
+
+template <typename OpTy>
+class ConvertAtenMaxUnpoolOp : public OpConversionPattern<OpTy> {
+  using OpConversionPattern<OpTy>::OpConversionPattern;
+
+private:
+  static const int64_t Dim = DimensionTraits<OpTy>::Dim;
+
+  LogicalResult createUnpoolOp(OpTy &op, typename OpTy::Adaptor adaptor,
+                               ConversionPatternRewriter &rewriter) const {
     if (failed(verifyLinalgCompatibleTypes(op, rewriter)))
       return failure();
 
     Location loc = op->getLoc();
-    const TypeConverter *typeConverter = getTypeConverter();
+    const TypeConverter *typeConverter = this->getTypeConverter();
     Value self = adaptor.getSelf();
     auto selfType = cast<RankedTensorType>(self.getType());
 
-    size_t spatial = selfType.getRank() - 2;
-    ArrayRef<int64_t> inputSize = selfType.getShape().take_back(spatial);
+    ArrayRef<int64_t> inputSize = selfType.getShape().take_back(Dim);
     if (ShapedType::isDynamicShape(inputSize))
       return rewriter.notifyMatchFailure(op,
                                          "input type must be of static shape");
 
     Value indices = adaptor.getIndices();
     auto indicesType = cast<RankedTensorType>(indices.getType());
-    if (inputSize != indicesType.getShape().take_back(spatial))
+    if (inputSize != indicesType.getShape().take_back(Dim))
       return rewriter.notifyMatchFailure(op, "input/indices shape mismatch");
 
     auto resType = typeConverter->convertType<RankedTensorType>(op.getType());
     if (!resType)
       return rewriter.notifyMatchFailure(op, "invalid result type");
 
-    ArrayRef<int64_t> inferredOutSize = resType.getShape().take_back(spatial);
+    ArrayRef<int64_t> inferredOutSize = resType.getShape().take_back(Dim);
     if (ShapedType::isDynamicShape(inferredOutSize))
       return rewriter.notifyMatchFailure(op,
                                          "output type must be of static shape");
@@ -637,7 +651,7 @@ class ConvertAtenMaxUnpool3dOp final
         return rewriter.notifyMatchFailure(op,
                                            "only support constant int output");
 
-      if (inferredOutSize != ArrayRef(output).take_back(spatial))
+      if (inferredOutSize != ArrayRef(output).take_back(Dim))
         return rewriter.notifyMatchFailure(op, "Invalid output size");
     }
     SmallVector<int64_t> stride;
@@ -653,12 +667,12 @@ class ConvertAtenMaxUnpool3dOp final
 
     // TODO: add support for asymmetric padding coming from "onnx.MaxUnpool"
     // (padding.size() == 6).
-    if (stride.size() != spatial || padding.size() != spatial)
+    if (stride.size() != Dim || padding.size() != Dim)
       return rewriter.notifyMatchFailure(
-          op, "stride and padding must be of size 3");
+          op, "stride and padding must be of size Dim");
 
     int64_t outRank = resType.getRank();
-    int64_t NC = outRank - spatial;
+    int64_t NC = outRank - Dim;
 
     for (auto &&[inDim, outDim, str, pad] :
          llvm::zip_equal(inputSize, inferredOutSize, stride, padding)) {
@@ -695,7 +709,7 @@ class ConvertAtenMaxUnpool3dOp final
     // (e.g. pooling_input_size=5, kernel_size=2, stride=2, output_size=2)
     // pad self and indices tensors to avoid out of bounds access.
     SmallVector<int64_t> expectedInputShape =
-        llvm::to_vector(resType.getShape().drop_back(spatial));
+        llvm::to_vector(resType.getShape().drop_back(Dim));
     for (auto &&[str, pad, resSize] :
          llvm::zip_equal(stride, padding, inferredOutSize))
       expectedInputShape.emplace_back(ceilDiv(resSize, str) + pad * 2);
@@ -708,7 +722,7 @@ class ConvertAtenMaxUnpool3dOp final
       SmallVector<int64_t> low(outRank, 0);
       SmallVector<int64_t> high(NC, 0);
       for (auto &&[inpSize, outSize] : llvm::zip_equal(
-               inputSize, ArrayRef(expectedInputShape).take_back(spatial))) {
+               inputSize, ArrayRef(expectedInputShape).take_back(Dim))) {
         high.emplace_back(outSize - inpSize);
       }
 
@@ -827,6 +841,13 @@ class ConvertAtenMaxUnpool3dOp final
     rewriter.replaceOp(op, result);
     return success();
   }
+
+public:
+  LogicalResult
+  matchAndRewrite(OpTy op, typename OpTy::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    return createUnpoolOp(op, adaptor, rewriter);
+  }
 };
 } // namespace
 
@@ -1527,8 +1548,12 @@ void mlir::torch::torch_to_linalg::populatePoolingPatternsAndLegality(
   patterns.add<ConvertAtenMaxPoolOp<AtenMaxPool3dWithIndicesOp>>(typeConverter,
                                                                  context);
 
+  target.addIllegalOp<AtenMaxUnpool2dOp>();
   target.addIllegalOp<AtenMaxUnpool3dOp>();
-  patterns.add<ConvertAtenMaxUnpool3dOp>(typeConverter, context);
+  patterns.add<ConvertAtenMaxUnpoolOp<AtenMaxUnpool2dOp>>(typeConverter,
+                                                          context);
+  patterns.add<ConvertAtenMaxUnpoolOp<AtenMaxUnpool3dOp>>(typeConverter,
+                                                          context);
 
   target.addIllegalOp<AtenAvgPool1dOp, AtenAvgPool2dOp, AtenAvgPool3dOp>();
   patterns

diff --git a/projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py b/projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py
@@ -622,6 +622,7 @@ def emit_with_mutating_variants(key, **kwargs):
         "aten::max_pool1d_with_indices : (Tensor, int[], int[], int[], int[], bool) -> (Tensor, Tensor)"
     )
     emit("aten::max_pool2d : (Tensor, int[], int[], int[], int[], bool) -> (Tensor)")
+    emit("aten::max_unpool2d : (Tensor, Tensor, int[]) -> (Tensor)")
     emit(
         "aten::max_pool2d_with_indices : (Tensor, int[], int[], int[], int[], bool) -> (Tensor, Tensor)",
         has_canonicalizer=True,

diff --git a/test/Conversion/TorchOnnxToTorch/simple_ops_g_to_p.mlir b/test/Conversion/TorchOnnxToTorch/simple_ops_g_to_p.mlir
@@ -1680,7 +1680,7 @@ func.func @test_maxunpool_2d_export_without_output_shape(%arg0: !torch.vtensor<[
   // CHECK: %[[INT2:.*]] = torch.constant.int 2
   // CHECK: %[[INT2_1:.*]] = torch.constant.int 2
   // CHECK: %[[STRIDE:.*]] = torch.prim.ListConstruct %[[INT2]], %[[INT2_1]] : (!torch.int, !torch.int) -> !torch.list<int>
-  // CHECK: %[[RESULT:.*]] = torch.aten.max_unpool3d %arg0, %arg1, %[[OUTPUT_SHAPE]], %[[STRIDE]], %[[PADDING]] : !torch.vtensor<[1,1,2,2],f32>, !torch.vtensor<[1,1,2,2],si64>, !torch.list<int>, !torch.list<int>, !torch.list<int> -> !torch.vtensor<[1,1,4,4],f32>
+  // CHECK: %[[RESULT:.*]] = torch.aten.max_unpool2d %arg0, %arg1, %[[OUTPUT_SHAPE]], %[[STRIDE]], %[[PADDING]] : !torch.vtensor<[1,1,2,2],f32>, !torch.vtensor<[1,1,2,2],si64>, !torch.list<int>, !torch.list<int>, !torch.list<int> -> !torch.vtensor<[1,1,4,4],f32>
   // return %[[RESULT]] : !torch.vtensor<[1,1,4,4],f32>
   %0 = torch.operator "onnx.MaxUnpool"(%arg0, %arg1) {torch.onnx.kernel_shape = [2 : si64, 2 : si64], torch.onnx.strides = [2 : si64, 2 : si64]} : (!torch.vtensor<[1,1,2,2],f32>, !torch.vtensor<[1,1,2,2],si64>) -> !torch.vtensor<[1,1,4,4],f32>
   return %0 : !torch.vtensor<[1,1,4,4],f32>

diff --git a/test/Conversion/TorchToLinalg/pooling.mlir b/test/Conversion/TorchToLinalg/pooling.mlir
@@ -100,8 +100,8 @@ func.func @forward_max_pool3d(%arg0: !torch.vtensor<[?,?,?,?,?],f32>) -> !torch.
 
 // CHECK: #map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 floordiv 2, d3 floordiv 2)>
 // CHECK: #map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
-// CHECK-LABEL: func @forward_max_unpool
-func.func @forward_max_unpool(%arg0: !torch.vtensor<[1,1,2,2],f32>, %arg1: !torch.vtensor<[1,1,2,2],si64>) -> !torch.vtensor<[1,1,4,4],f32> attributes {torch.onnx_meta.ir_version = 6 : si64, torch.onnx_meta.opset_version = 11 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
+// CHECK-LABEL: func @forward_max_unpool2d
+func.func @forward_max_unpool2d(%arg0: !torch.vtensor<[1,1,2,2],f32>, %arg1: !torch.vtensor<[1,1,2,2],si64>) -> !torch.vtensor<[1,1,4,4],f32> attributes {torch.onnx_meta.ir_version = 6 : si64, torch.onnx_meta.opset_version = 11 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
   %int1 = torch.constant.int 1
   %int1_0 = torch.constant.int 1
   %int4 = torch.constant.int 4
@@ -113,7 +113,7 @@ func.func @forward_max_unpool(%arg0: !torch.vtensor<[1,1,2,2],f32>, %arg1: !torc
   %int2 = torch.constant.int 2
   %int2_3 = torch.constant.int 2
   %2 = torch.prim.ListConstruct %int2, %int2_3 : (!torch.int, !torch.int) -> !torch.list<int>
-  %3 = torch.aten.max_unpool3d %arg0, %arg1, %0, %2, %1 : !torch.vtensor<[1,1,2,2],f32>, !torch.vtensor<[1,1,2,2],si64>, !torch.list<int>, !torch.list<int>, !torch.list<int> -> !torch.vtensor<[1,1,4,4],f32>
+  %3 = torch.aten.max_unpool2d %arg0, %arg1, %0, %2, %1 : !torch.vtensor<[1,1,2,2],f32>, !torch.vtensor<[1,1,2,2],si64>, !torch.list<int>, !torch.list<int>, !torch.list<int> -> !torch.vtensor<[1,1,4,4],f32>
 
   // CHECK: %[[INDICES:.*]] = torch_c.to_builtin_tensor %arg1 : !torch.vtensor<[1,1,2,2],si64> -> tensor<1x1x2x2xi64>
   // CHECK: %[[INPUT:.*]] = torch_c.to_builtin_tensor %arg0 : !torch.vtensor<[1,1,2,2],f32> -> tensor<1x1x2x2xf32>