Skip to content

[Torch] Canonicalize pool ops with single int tuple params. #4250

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Nov 24, 2025
Merged

[Torch] Canonicalize pool ops with single int tuple params. #4250

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
f6f5617
[Torch] Canonicalize pool ops with single int tuple params.
sahas3 Jul 1, 2025
7403809
Add [[maybe_unused]] attribute
sahas3 Jul 1, 2025
1f45f0c
mend
sahas3 Jul 1, 2025
026ac79
Add lit test.
sahas3 Jul 2, 2025
b667200
Add canonicalizer in torch_ods_gen.py + update error message
sahas3 Jul 9, 2025
3e003bb
Update test names
sahas3 Aug 28, 2025
4138144
[Task] : Fix create calls.
sahas3 Nov 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7744,6 +7744,7 @@ def Torch_AtenMaxPool2dOp : Torch_Op<"aten.max_pool2d", [
printDefaultTorchOp(printer, *this, 6, 1);
}
}];
let hasCanonicalizer = 1;
}

def Torch_AtenMaxUnpool2dOp : Torch_Op<"aten.max_unpool2d", [
Expand Down Expand Up @@ -7857,6 +7858,7 @@ def Torch_AtenMaxPool3dOp : Torch_Op<"aten.max_pool3d", [
printDefaultTorchOp(printer, *this, 6, 1);
}
}];
let hasCanonicalizer = 1;
}

def Torch_AtenMaxUnpool3dOp : Torch_Op<"aten.max_unpool3d", [
Expand Down Expand Up @@ -8001,6 +8003,7 @@ def Torch_AtenAvgPool2dOp : Torch_Op<"aten.avg_pool2d", [
printDefaultTorchOp(printer, *this, 7, 1);
}
}];
let hasCanonicalizer = 1;
}

def Torch_AtenAvgPool2dBackwardOp : Torch_Op<"aten.avg_pool2d_backward", [
Expand Down Expand Up @@ -8060,6 +8063,7 @@ def Torch_AtenAvgPool3dOp : Torch_Op<"aten.avg_pool3d", [
printDefaultTorchOp(printer, *this, 7, 1);
}
}];
let hasCanonicalizer = 1;
}

def Torch_AtenAvgPool3dBackwardOp : Torch_Op<"aten.avg_pool3d_backward", [
Expand Down
6 changes: 0 additions & 6 deletions lib/Conversion/TorchToTosa/TorchToTosa.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6275,12 +6275,6 @@ void expandPoolParams(AtenOpT op, SmallVectorImpl<int64_t> &params,
if constexpr (std::is_same<AtenOpT, AtenMaxPool1dOp>() ||
std::is_same<AtenOpT, AtenAvgPool1dOp>())
params.push_back(val);

if constexpr (std::is_same<AtenOpT, AtenMaxPool2dOp>() ||
std::is_same<AtenOpT, AtenAvgPool2dOp>()) {
if (params.size() == 1)
params.push_back(params[0]);
}
}

// Checks the validity of pooling parameters and stores them in the respective
Expand Down
178 changes: 178 additions & 0 deletions lib/Dialect/Torch/IR/TorchOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5720,6 +5720,184 @@ void Aten_AdaptiveAvgPool2dOp::getCanonicalizationPatterns(
});
}

namespace {

void expand(SmallVectorImpl<int64_t> &params, 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.");
}

// 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(Torch::ConstantIntOp::create(
rewriter, loc, rewriter.getI64IntegerAttr(kernelSizeInts[dim])));
cstPadding.push_back(Torch::ConstantIntOp::create(
rewriter, loc, rewriter.getI64IntegerAttr(paddingInts[dim])));
cstStrides.push_back(Torch::ConstantIntOp::create(
rewriter, 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(Torch::ConstantIntOp::create(
rewriter, loc, rewriter.getI64IntegerAttr(dilationInts[dim])));
}

auto targetListType =
Torch::ListType::get(Torch::IntType::get(op->getContext()));
kernelSizeList = Torch::PrimListConstructOp::create(
rewriter, loc, targetListType, cstKernel);
paddingList = Torch::PrimListConstructOp::create(rewriter, loc,
targetListType, cstPadding);
stridesList = Torch::PrimListConstructOp::create(rewriter, loc,
targetListType, cstStrides);
dilationsList = Torch::PrimListConstructOp::create(
rewriter, 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
//===----------------------------------------------------------------------===//
Expand Down
11 changes: 7 additions & 4 deletions projects/pt1/e2e_testing/xfail_sets.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -533,8 +533,6 @@
"Aten_TrilinearModuleVaryingRanksUnorderedExpands_basic",
"Aten_TrilinearModuleSumAllDims_basic",
"Aten_TrilinearModuleSumdims_basic",
"AvgPool2dSingleIntTupleParamsIncludePadModule_basic",
"AvgPool2dSingleIntTupleParamsModule_basic",
"SliceOutOfLowerBoundEndIndexModule_basic",
"RollModule_basic",
"AdaptiveAvgPool2dDynamicNoBatch_basic",
Expand Down Expand Up @@ -990,8 +988,6 @@
}

FX_IMPORTER_STABLEHLO_CRASHING_SET = {
"AvgPool2dSingleIntTupleParamsIncludePadModule_basic",
"AvgPool2dSingleIntTupleParamsModule_basic",
"BatchNorm1DModule_basic",
"BatchNorm2DModule_basic",
"BatchNorm3DModule_basic",
Expand Down Expand Up @@ -2863,6 +2859,7 @@
"AvgPool1dPadCeilPadNotIncluded_basic",
"AvgPool2dDiffKernelsStridesPadCeilPadNotIncluded_basic",
"AvgPool3dDiffKernelsStridesPadCeilPadNotIncluded_basic",
"AvgPool3dSingleIntTupleStrideModule_basic",
"BatchMlpLayerModule_basic",
"BincountMinlengthModule_basic",
"BincountModule_basic",
Expand Down Expand Up @@ -3059,6 +3056,7 @@
"MaxPool2dWithIndicesNonDefaultDilationModule_basic",
"MaxPool2dWithIndicesNonDefaultParamsModule_basic",
"MaxPool2dWithIndicesNonDefaultStrideModule_basic",
"MaxPool2dSingleIntTupleKernelModule_basic",
"MaxPool3dCeilModeTrueModule_basic",
"MaxPool3dLargeDatadModule_basic",
"MaxPool3dModuleRandomSimple_basic",
Expand All @@ -3070,6 +3068,7 @@
"MaxPool3dWithIndicesNonDefaultDilationModule_basic",
"MaxPool3dWithIndicesNonDefaultParamsModule_basic",
"MaxPool3dWithIndicesNonDefaultStrideModule_basic",
"MaxPool3dSingleIntTupleDilationModule_basic",
"MaxUnpool3dModule_basic",
"MaxUnpool3dModulePad0_basic",
"MaxUnpool2dModule_basic",
Expand Down Expand Up @@ -3579,6 +3578,7 @@
"AvgPool3dStaticModule_basic",
"AvgPool3dCountIncludePadFalse_basic",
"AvgPool3dCountIncludePadFalseWithoutPadding_basic",
"AvgPool3dSingleIntTupleStrideModule_basic",
"Conv_Transpose1dModule_basic",
"Conv_Transpose1dStaticModule_basic",
"Conv_Transpose3dModule_basic",
Expand Down Expand Up @@ -3821,6 +3821,7 @@
"MaxPool3dWithIndicesNonDefaultParamsModule_basic",
"MaxPool3dWithIndicesNonDefaultStrideModule_basic",
"MaxPool3dWithIndicesStaticModule_basic",
"MaxPool3dSingleIntTupleDilationModule_basic",
"MeanDimEmptyDimModule_basic",
"MlGroupNormManualModule_basic",
"MlGroupNormModule_basic",
Expand Down Expand Up @@ -4270,6 +4271,7 @@
"AvgPool2dIntModule_basic",
"AvgPool2dStaticModule_basic",
"AvgPool2dWithoutPadModule_basic",
"AvgPool3dSingleIntTupleStrideModule_basic",
"BatchMlpLayerModule_basic",
"BernoulliFloatModule_basic",
"BernoulliModule_basic",
Expand Down Expand Up @@ -4684,6 +4686,7 @@
"MaxPool3dWithIndicesNonDefaultParamsModule_basic",
"MaxPool3dWithIndicesNonDefaultStrideModule_basic",
"MaxPool3dWithIndicesStaticModule_basic",
"MaxPool3dSingleIntTupleDilationModule_basic",
"MeanDimAllReduceKeepdimModule_basic",
"MeanDimAllReduceModule_basic",
"MeanDimDtypeModule_basic",
Expand Down
16 changes: 12 additions & 4 deletions projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -657,7 +657,10 @@ def emit_with_mutating_variants(key, **kwargs):
emit(
"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_pool2d : (Tensor, int[], int[], int[], int[], bool) -> (Tensor)",
has_canonicalizer=True,
)
emit("aten::max_unpool2d : (Tensor, Tensor, int[]) -> (Tensor)")
emit(
"aten::max_pool2d_with_indices : (Tensor, int[], int[], int[], int[], bool) -> (Tensor, Tensor)",
Expand All @@ -666,7 +669,10 @@ def emit_with_mutating_variants(key, **kwargs):
emit(
"aten::max_pool2d_with_indices_backward : (Tensor, Tensor, int[], int[], int[], int[], bool, Tensor) -> (Tensor)"
)
emit("aten::max_pool3d : (Tensor, int[], int[], int[], int[], bool) -> (Tensor)")
emit(
"aten::max_pool3d : (Tensor, int[], int[], int[], int[], bool) -> (Tensor)",
has_canonicalizer=True,
)
emit("aten::max_unpool3d : (Tensor, Tensor, int[], int[], int[]) -> (Tensor)")
emit(
"aten::max_pool3d_with_indices : (Tensor, int[], int[], int[], int[], bool) -> (Tensor, Tensor)",
Expand All @@ -677,13 +683,15 @@ def emit_with_mutating_variants(key, **kwargs):
)
emit("aten::avg_pool1d : (Tensor, int[], int[], int[], bool, bool) -> (Tensor)")
emit(
"aten::avg_pool2d : (Tensor, int[], int[], int[], bool, bool, int?) -> (Tensor)"
"aten::avg_pool2d : (Tensor, int[], int[], int[], bool, bool, int?) -> (Tensor)",
has_canonicalizer=True,
)
emit(
"aten::avg_pool2d_backward : (Tensor, Tensor, int[], int[], int[], bool, bool, int?) -> (Tensor)"
)
emit(
"aten::avg_pool3d : (Tensor, int[], int[], int[], bool, bool, int?) -> (Tensor)"
"aten::avg_pool3d : (Tensor, int[], int[], int[], bool, bool, int?) -> (Tensor)",
has_canonicalizer=True,
)
emit(
"aten::avg_pool3d_backward : (Tensor, Tensor, int[], int[], int[], bool, bool, int?) -> (Tensor)"
Expand Down
Loading
Loading