[MLIR][Torch] Add TorchToLinalg lowering for AtenAvgPool3dOp (#3030)

This commit also fixes the average pool op' test failing for
OnnxToLinalg lowering.

Signed-Off By: Vivek Khandelwal <[email protected]>

Author: vivekkhandelwal1

include/torch-mlir/Conversion/TorchToLinalg/Utils.h

@@ -97,6 +97,10 @@ getBackendTypeForScalarType(MLIRContext *context,
 
 bool isUnsignedTorchType(Type type);
 
+LogicalResult permuteTensor(Operation *op, PatternRewriter &rewriter,
+                            Location loc, SmallVector<int64_t> dimensions,
+                            Value input, Value &result);
+
 } // namespace torch_to_linalg
 } // namespace torch
 } // namespace mlir

lib/Conversion/TorchToLinalg/DataMovement.cpp

@@ -1800,55 +1800,15 @@ class ConvertAtenPermuteOp : public OpConversionPattern<AtenPermuteOp> {
       return rewriter.notifyMatchFailure(op, "all dimensions must be constant");
 
     Value inVector = adaptor.getSelf();
-    auto inType = cast<RankedTensorType>(inVector.getType());
-    int64_t inputRank = inType.getRank();
-    auto outType = cast<RankedTensorType>(
-        getTypeConverter()->convertType(op->getResult(0).getType()));
-    Type elementType = inType.getElementType();
-
-    // Check if the dimensions are a valid constants.
-    int64_t numDimensions = dimensions.size();
-    if (inputRank != numDimensions)
+    Value result;
+    if (failed(torch_to_linalg::permuteTensor(op, rewriter, op->getLoc(),
+                                              dimensions, inVector, result)))
       return rewriter.notifyMatchFailure(
-          op, "size of `dims` must be equal to the rank of the input");
-    for (unsigned i = 0; i < numDimensions; i++) {
-      if (dimensions[i] < 0)
-        dimensions[i] = toPositiveDim(dimensions[i], inputRank);
-      if (!isValidDim(dimensions[i], inputRank))
-        return rewriter.notifyMatchFailure(op, "dimension out of range");
-    }
-
-    Location loc = op.getLoc();
-
-    SmallVector<Value> outputDims;
-    for (unsigned i = 0; i < inputRank; i++)
-      outputDims.push_back(getDimOp(rewriter, loc, inVector, dimensions[i]));
-
-    Value outVector = rewriter.create<tensor::EmptyOp>(
-        loc, getAsOpFoldResult(outputDims), elementType);
-    SmallVector<AffineExpr> idExprs;
-    SmallVector<AffineExpr> swapExprs;
-    for (unsigned i = 0; i < inputRank; i++)
-      idExprs.push_back(getAffineDimExpr(i, rewriter.getContext()));
-    for (unsigned i = 0; i < inputRank; i++)
-      swapExprs.push_back(idExprs[dimensions[i]]);
+          op, "failed to perform permutation of tensor");
 
-    AffineMap inputMap =
-        AffineMap::get(inputRank, /*symbolCount=*/0, idExprs, op->getContext());
-    AffineMap outputMap = AffineMap::get(inputRank, /*symbolCount=*/0,
-                                         swapExprs, op->getContext());
-    SmallVector<AffineMap> indexingMaps{inputMap, outputMap};
-    SmallVector<utils::IteratorType> iteratorTypes(
-        inputRank, utils::IteratorType::parallel);
-    auto transpose = rewriter
-                         .create<linalg::GenericOp>(
-                             loc, outVector.getType(), inVector, outVector,
-                             indexingMaps, iteratorTypes,
-                             [](OpBuilder &b, Location loc, ValueRange args) {
-                               b.create<linalg::YieldOp>(loc, args[0]);
-                             })
-                         .getResult(0);
-    rewriter.replaceOpWithNewOp<tensor::CastOp>(op, outType, transpose);
+    auto outType = cast<RankedTensorType>(
+        getTypeConverter()->convertType(op->getResult(0).getType()));
+    rewriter.replaceOpWithNewOp<tensor::CastOp>(op, outType, result);
     return success();
   }
 };

lib/Conversion/TorchToLinalg/Pooling.cpp

@@ -168,11 +168,42 @@ static LogicalResult createPoolingOp(
   Value windowTensor = rewriter.create<tensor::EmptyOp>(
       loc, getAsOpFoldResult(shape), elementType);
 
-  result = rewriter
-               .create<OpTy>(loc, outTensorInitialized.getType(),
-                             ValueRange{paddedInput, windowTensor},
-                             outTensorInitialized, stridesAttr, dilationAttr)
-               .getResult(0);
+  Value permutedInput = paddedInput, permutedOutput = outTensorInitialized;
+  if (dimensionality == 3) {
+    // Permute input and output tensor as follows:
+    // (n,c,d,h,w) -> (n,d,h,w,c)
+    SmallVector<int64_t> dimensions = {0, 2, 3, 4, 1};
+    if (failed(torch_to_linalg::permuteTensor(op, rewriter, op->getLoc(),
+                                              dimensions, paddedInput,
+                                              permutedInput)))
+      return rewriter.notifyMatchFailure(
+          op, "failed to perform permutation of tensor");
+
+    if (failed(torch_to_linalg::permuteTensor(op, rewriter, op->getLoc(),
+                                              dimensions, outTensorInitialized,
+                                              permutedOutput)))
+      return rewriter.notifyMatchFailure(
+          op, "failed to perform permutation of tensor");
+  }
+
+  Value poolingResult =
+      rewriter
+          .create<OpTy>(loc, permutedOutput.getType(),
+                        ValueRange{permutedInput, windowTensor}, permutedOutput,
+                        stridesAttr, dilationAttr)
+          .getResult(0);
+
+  result = poolingResult;
+  if (dimensionality == 3) {
+    // Permute output tensor as follows:
+    // (n,d,h,w,c) -> (n,c,d,h,w)
+    SmallVector<int64_t> dimensions = {0, 4, 1, 2, 3};
+    if (failed(torch_to_linalg::permuteTensor(
+            op, rewriter, op->getLoc(), dimensions, poolingResult, result)))
+      return rewriter.notifyMatchFailure(
+          op, "failed to perform permutation of tensor");
+  }
+
   return success();
 }
 
@@ -604,15 +635,17 @@ class ConvertAtenAvgPoolOp : public OpConversionPattern<OpTy> {
             paddingInts, dilationInts, rewriter.getZeroAttr(inputElementType),
             outTensorShape, paddedInput, sumPool)))
       return rewriter.notifyMatchFailure(op, "unable to compute sumpool");
-    Value divisor;
-    if constexpr (std::is_same<OpTy, AtenAvgPool2dOp>()) {
-      Value kHtimeskW = rewriter.create<arith::MulIOp>(
-          loc, kernelSizeIntValues[0], kernelSizeIntValues[1]);
+    // }
+
+    Value divisor = kernelSizeIntValues[0];
+    for (uint32_t i = 1; i < kernelSizeIntValues.size(); i++) {
+      divisor =
+          rewriter.create<arith::MulIOp>(loc, divisor, kernelSizeIntValues[i]);
+    }
+    if constexpr (!std::is_same<OpTy, AtenAvgPool1dOp>()) {
       divisor = isa<Torch::NoneType>(op.getDivisorOverride().getType())
-                    ? kHtimeskW
+                    ? divisor
                     : adaptor.getDivisorOverride();
-    } else {
-      divisor = kernelSizeIntValues[0];
     }
     divisor = convertScalarToDtype(rewriter, loc, divisor, resultElementType);
 
@@ -1115,13 +1148,16 @@ void mlir::torch::torch_to_linalg::populatePoolingPatternsAndLegality(
 
   target.addIllegalOp<AtenMaxPool2dWithIndicesOp>();
   patterns.add<ConvertAtenMaxPool2dWithIndicesOp>(typeConverter, context);
-  target.addIllegalOp<AtenAvgPool1dOp, AtenAvgPool2dOp>();
+  target.addIllegalOp<AtenAvgPool1dOp, AtenAvgPool2dOp, AtenAvgPool3dOp>();
   patterns
       .add<ConvertAtenAvgPoolOp<AtenAvgPool1dOp, linalg::PoolingNcwSumOp, 1>>(
           typeConverter, context);
   patterns
       .add<ConvertAtenAvgPoolOp<AtenAvgPool2dOp, linalg::PoolingNchwSumOp, 2>>(
           typeConverter, context);
+  patterns
+      .add<ConvertAtenAvgPoolOp<AtenAvgPool3dOp, linalg::PoolingNdhwcSumOp, 3>>(
+          typeConverter, context);
   target.addIllegalOp<AtenAdaptiveAvgPool1dOp, AtenAdaptiveAvgPool2dOp,
                       AtenAdaptiveAvgPool3dOp, Aten_AdaptiveAvgPool3dOp>();
   patterns.add<ConvertAtenAdaptivePoolOp<AtenAdaptiveAvgPool1dOp>>(

lib/Conversion/TorchToLinalg/Utils.cpp

@@ -572,3 +572,55 @@ bool torch_to_linalg::isUnsignedTorchType(Type type) {
   llvm_unreachable("Unknown type checked for signedness");
   return false;
 }
+
+LogicalResult torch_to_linalg::permuteTensor(Operation *op,
+                                             PatternRewriter &rewriter,
+                                             Location loc,
+                                             SmallVector<int64_t> dimensions,
+                                             Value input, Value &result) {
+  auto inType = cast<RankedTensorType>(input.getType());
+  int64_t inputRank = inType.getRank();
+  Type elementType = inType.getElementType();
+
+  // Check if the dimensions are a valid constants.
+  int64_t numDimensions = dimensions.size();
+  if (inputRank != numDimensions)
+    return rewriter.notifyMatchFailure(
+        op, "size of `dims` must be equal to the rank of the input");
+  for (uint32_t i = 0; i < numDimensions; i++) {
+    if (dimensions[i] < 0)
+      dimensions[i] = toPositiveDim(dimensions[i], inputRank);
+    if (!isValidDim(dimensions[i], inputRank))
+      return rewriter.notifyMatchFailure(op, "dimension out of range");
+  }
+
+  SmallVector<Value> outputDims;
+  for (uint32_t i = 0; i < inputRank; i++)
+    outputDims.push_back(getDimOp(rewriter, loc, input, dimensions[i]));
+
+  Value outVector = rewriter.create<tensor::EmptyOp>(
+      loc, getAsOpFoldResult(outputDims), elementType);
+  SmallVector<AffineExpr> idExprs;
+  SmallVector<AffineExpr> swapExprs;
+  for (uint32_t i = 0; i < inputRank; i++)
+    idExprs.push_back(getAffineDimExpr(i, rewriter.getContext()));
+  for (uint32_t i = 0; i < inputRank; i++)
+    swapExprs.push_back(idExprs[dimensions[i]]);
+
+  AffineMap inputMap =
+      AffineMap::get(inputRank, /*symbolCount=*/0, idExprs, op->getContext());
+  AffineMap outputMap =
+      AffineMap::get(inputRank, /*symbolCount=*/0, swapExprs, op->getContext());
+  SmallVector<AffineMap> indexingMaps{inputMap, outputMap};
+  SmallVector<utils::IteratorType> iteratorTypes(inputRank,
+                                                 utils::IteratorType::parallel);
+  result = rewriter
+               .create<linalg::GenericOp>(
+                   loc, outVector.getType(), input, outVector, indexingMaps,
+                   iteratorTypes,
+                   [](OpBuilder &b, Location loc, ValueRange args) {
+                     b.create<linalg::YieldOp>(loc, args[0]);
+                   })
+               .getResult(0);
+  return success();
+}