[Torch] Canonicalize pool ops with single int tuple params. #4250
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -5617,6 +5617,184 @@ void Aten_AdaptiveAvgPool2dOp::getCanonicalizationPatterns( | |
| }); | ||
| } | ||
|
|
||
| namespace { | ||
|
|
||
| void expand(SmallVectorImpl<int64_t> ¶ms, int numSpatialDims) { | ||
| if (params.size() == 1) { | ||
| for ([[maybe_unused]] int dim : llvm::seq<int>(0, numSpatialDims - 1)) { | ||
| params.push_back(params[0]); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| template <typename AtenPoolOpT> | ||
| LogicalResult expandPoolParams(AtenPoolOpT op, int numSpatialDims, | ||
| mlir::PatternRewriter &rewriter, | ||
| Value &kernelSizeList, Value &stridesList, | ||
| Value &paddingList, Value &dilationsList) { | ||
|
|
||
| SmallVector<int64_t, 3> kernelSizeInts, strideInts, paddingInts, dilationInts; | ||
| if (!matchPattern(op.getKernelSize(), | ||
| m_TorchListOfConstantInts(kernelSizeInts))) | ||
| return rewriter.notifyMatchFailure( | ||
| op, "Non-const kernel_size for pooling op unsupported"); | ||
|
|
||
| if (!matchPattern(op.getPadding(), m_TorchListOfConstantInts(paddingInts))) | ||
| return rewriter.notifyMatchFailure( | ||
| op, "Non-const padding factor for pooling op unsupported"); | ||
|
|
||
| if (!matchPattern(op.getStride(), m_TorchListOfConstantInts(strideInts))) | ||
| return rewriter.notifyMatchFailure( | ||
| op, "Non-const stride for pooling op unsupported"); | ||
|
|
||
| if constexpr (std::is_same<AtenPoolOpT, AtenMaxPool2dOp>() || | ||
| std::is_same<AtenPoolOpT, AtenMaxPool3dOp>()) { | ||
| if (!matchPattern(op.getDilation(), | ||
| m_TorchListOfConstantInts(dilationInts))) | ||
| return rewriter.notifyMatchFailure( | ||
| op, "Non-const dilation for pooling op unsupported"); | ||
|
|
||
| if (kernelSizeInts.size() != 1 && paddingInts.size() != 1 && | ||
| strideInts.size() != 1 && dilationInts.size() != 1) { | ||
| return rewriter.notifyMatchFailure( | ||
| op, | ||
| "Expected one of kernel/stride/padding/dilation to be singleton."); | ||
| } | ||
|
|
||
| expand(dilationInts, numSpatialDims); | ||
|
|
||
| } else if (kernelSizeInts.size() != 1 && paddingInts.size() != 1 && | ||
| strideInts.size() != 1) { | ||
| return rewriter.notifyMatchFailure( | ||
| op, "Expected one of kernel/stride/padding to be singleton."); | ||
| } | ||
|
Comment on lines
+5650
to
+5670
There was a problem hiding this comment. Instead of this, you can also do: There was a problem hiding this comment. I think the singleton check for dilation cannot be decoupled from the singleton check for the other values for So the code will look like: Is that your suggestion? There was a problem hiding this comment. Hi @vivekkhandelwal1 can you please clarify if I understood your suggestion as captured in my previous comment? If yes, is it to make the code easier to read? There was a problem hiding this comment. Hi @vivekkhandelwal1 any further thoughts on this? |
||
|
|
||
| // expand singleton elements | ||
| expand(kernelSizeInts, numSpatialDims); | ||
| expand(paddingInts, numSpatialDims); | ||
| expand(strideInts, numSpatialDims); | ||
|
|
||
| Location loc = op.getLoc(); | ||
|
|
||
| SmallVector<Value> cstKernel, cstPadding, cstStrides, cstDilations; | ||
| for (auto dim : llvm::seq<int>(0, kernelSizeInts.size())) { | ||
| cstKernel.push_back(rewriter.create<Torch::ConstantIntOp>( | ||
| loc, rewriter.getI64IntegerAttr(kernelSizeInts[dim]))); | ||
| cstPadding.push_back(rewriter.create<Torch::ConstantIntOp>( | ||
| loc, rewriter.getI64IntegerAttr(paddingInts[dim]))); | ||
| cstStrides.push_back(rewriter.create<Torch::ConstantIntOp>( | ||
| loc, rewriter.getI64IntegerAttr(strideInts[dim]))); | ||
| } | ||
|
|
||
| // set dilations separately as for AvgPool op it won't be set | ||
| for (auto dim : llvm::seq<int>(0, dilationInts.size())) { | ||
| cstDilations.push_back(rewriter.create<Torch::ConstantIntOp>( | ||
| loc, rewriter.getI64IntegerAttr(dilationInts[dim]))); | ||
| } | ||
|
|
||
| auto targetListType = | ||
| Torch::ListType::get(Torch::IntType::get(op->getContext())); | ||
| kernelSizeList = rewriter.create<Torch::PrimListConstructOp>( | ||
| loc, targetListType, cstKernel); | ||
| paddingList = rewriter.create<Torch::PrimListConstructOp>(loc, targetListType, | ||
| cstPadding); | ||
| stridesList = rewriter.create<Torch::PrimListConstructOp>(loc, targetListType, | ||
| cstStrides); | ||
| dilationsList = rewriter.create<Torch::PrimListConstructOp>( | ||
| loc, targetListType, cstDilations); | ||
|
|
||
| return success(); | ||
| } | ||
|
|
||
| template <typename AvgPoolOpT> | ||
| struct CanonicalizeAvgPoolWithSingleIntTuple | ||
| : public mlir::OpRewritePattern<AvgPoolOpT> { | ||
| CanonicalizeAvgPoolWithSingleIntTuple(mlir::MLIRContext *context) | ||
| : OpRewritePattern<AvgPoolOpT>(context, /*benefit=*/1) {} | ||
|
|
||
| LogicalResult | ||
| matchAndRewrite(AvgPoolOpT op, | ||
| mlir::PatternRewriter &rewriter) const override { | ||
| Value kernel, stride, pad, dilations; | ||
|
|
||
| auto numSpatialDims = 2; | ||
| if constexpr (std::is_same<AvgPoolOpT, AtenAvgPool3dOp>()) | ||
| numSpatialDims = 3; | ||
|
|
||
| // Attempt to expand params if necessary. | ||
| if (failed(expandPoolParams(op, numSpatialDims, rewriter, kernel, stride, | ||
| pad, dilations))) | ||
| return rewriter.notifyMatchFailure( | ||
| op, "Failed to expand params for AvgPooling"); | ||
|
|
||
| rewriter.replaceOpWithNewOp<AvgPoolOpT>( | ||
| op, op.getResult().getType(), op.getSelf(), kernel, stride, pad, | ||
| op.getCeilMode(), op.getCountIncludePad(), op.getDivisorOverride()); | ||
| return success(); | ||
| } | ||
| }; | ||
|
|
||
| template <typename MaxPoolOpT> | ||
| struct CanonicalizeMaxPoolWithSingleIntTuple | ||
| : public mlir::OpRewritePattern<MaxPoolOpT> { | ||
| CanonicalizeMaxPoolWithSingleIntTuple(mlir::MLIRContext *context) | ||
| : OpRewritePattern<MaxPoolOpT>(context, /*benefit=*/1) {} | ||
|
|
||
| LogicalResult | ||
| matchAndRewrite(MaxPoolOpT op, | ||
| mlir::PatternRewriter &rewriter) const override { | ||
| Value kernel, stride, pad, dilations; | ||
|
|
||
| auto numSpatialDims = 2; | ||
| if constexpr (std::is_same<MaxPoolOpT, AtenMaxPool3dOp>()) | ||
| numSpatialDims = 3; | ||
|
|
||
| // Attempt to expand params if necessary. | ||
| if (failed(expandPoolParams(op, numSpatialDims, rewriter, kernel, stride, | ||
| pad, dilations))) | ||
| return rewriter.notifyMatchFailure( | ||
| op, "Failed to expand params for MaxPooling"); | ||
|
|
||
| rewriter.replaceOpWithNewOp<MaxPoolOpT>(op, op.getResult().getType(), | ||
| op.getSelf(), kernel, stride, pad, | ||
| dilations, op.getCeilMode()); | ||
| return success(); | ||
| } | ||
| }; | ||
| } // namespace | ||
|
|
||
| //===----------------------------------------------------------------------===// | ||
| // AtenAvgPool2dOp | ||
| //===----------------------------------------------------------------------===// | ||
| void AtenAvgPool2dOp::getCanonicalizationPatterns(RewritePatternSet &patterns, | ||
| MLIRContext *context) { | ||
| patterns.add<CanonicalizeAvgPoolWithSingleIntTuple<AtenAvgPool2dOp>>(context); | ||
| } | ||
|
|
||
| //===----------------------------------------------------------------------===// | ||
| // AtenAvgPool3dOp | ||
| //===----------------------------------------------------------------------===// | ||
| void AtenAvgPool3dOp::getCanonicalizationPatterns(RewritePatternSet &patterns, | ||
| MLIRContext *context) { | ||
| patterns.add<CanonicalizeAvgPoolWithSingleIntTuple<AtenAvgPool3dOp>>(context); | ||
| } | ||
|
|
||
| //===----------------------------------------------------------------------===// | ||
| // AtenMaxPool2dOp | ||
| //===----------------------------------------------------------------------===// | ||
| void AtenMaxPool2dOp::getCanonicalizationPatterns(RewritePatternSet &patterns, | ||
| MLIRContext *context) { | ||
| patterns.add<CanonicalizeMaxPoolWithSingleIntTuple<AtenMaxPool2dOp>>(context); | ||
| } | ||
|
|
||
| //===----------------------------------------------------------------------===// | ||
| // AtenMaxPool3dOp | ||
| //===----------------------------------------------------------------------===// | ||
| void AtenMaxPool3dOp::getCanonicalizationPatterns(RewritePatternSet &patterns, | ||
| MLIRContext *context) { | ||
| patterns.add<CanonicalizeMaxPoolWithSingleIntTuple<AtenMaxPool3dOp>>(context); | ||
| } | ||
|
|
||
| //===----------------------------------------------------------------------===// | ||
| // AtenLinalgCrossOp | ||
| //===----------------------------------------------------------------------===// | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.