[TorchToLinalg] Fix avg_pool divisor to account for dilation. (#4671)

PoolSizeCalculator::getPoolSize() ignored dilation when computing the
average pooling divisor. With dilation > 1, edge positions produced
incorrect averages because the divisor did not reflect the actual number
of valid dilated taps within bounds.

Note adding an e2e test is not possible as PyTorch doesn't allow
creating a pool op with dilation. The bug was found as part of stress
testing onnx importer path by creating individual onnx ops.

Author: sahas3

lib/Conversion/TorchToLinalg/Pooling.cpp

@@ -960,7 +960,8 @@ template <int NumOfDims> class PoolSizeCalculator {
   // height and width labels in variables.
   Value getPoolSize(OpBuilder &b, SmallVectorImpl<Value> &kernelSizeIntValues,
                     SmallVectorImpl<int64_t> &strideInts,
-                    SmallVectorImpl<int64_t> &paddingInts);
+                    SmallVectorImpl<int64_t> &paddingInts,
+                    SmallVectorImpl<int64_t> &dilationInts);
 
 private:
   int64_t SumPoolTypeDimIndex[NumOfDims];
@@ -994,12 +995,14 @@ PoolSizeCalculator<NumOfDims>::PoolSizeCalculator(
 template <int NumOfDims>
 Value PoolSizeCalculator<NumOfDims>::getPoolSize(
     OpBuilder &b, SmallVectorImpl<Value> &kernelDimSizes,
-    SmallVectorImpl<int64_t> &strideInts,
-    SmallVectorImpl<int64_t> &paddingInts) {
+    SmallVectorImpl<int64_t> &strideInts, SmallVectorImpl<int64_t> &paddingInts,
+    SmallVectorImpl<int64_t> &dilationInts) {
   Value poolSize;
 
   Value cstZero =
       b.createOrFold<arith::ConstantOp>(location, b.getI64IntegerAttr(0));
+  Value cstOne =
+      b.createOrFold<arith::ConstantOp>(location, b.getI64IntegerAttr(1));
 
   for (int i = 0; i < NumOfDims; ++i) {
     // See the link below for the PyTorch implementation where this is
@@ -1016,25 +1019,60 @@ Value PoolSizeCalculator<NumOfDims>::getPoolSize(
         location, b.getI64IntegerAttr(strideInts[i]));
     Value PadDim = b.createOrFold<arith::ConstantOp>(
         location, b.getI64IntegerAttr(paddingInts[i]));
+    Value DilDim = b.createOrFold<arith::ConstantOp>(
+        location, b.getI64IntegerAttr(dilationInts[i]));
     Value ODimDDim = b.createOrFold<arith::MulIOp>(location, ODim, DDim);
     Value IDim0 = b.createOrFold<arith::SubIOp>(location, ODimDDim, PadDim);
     Value IDim = castIndexToInt64(b, location, InputSpatialDimSizes[i]);
+
+    // Effective window end: IDim0 + (kernel - 1) * dilation + 1
+    Value KernelM1 =
+        b.createOrFold<arith::SubIOp>(location, kernelDimSizes[i], cstOne);
+    Value KernelM1Dil =
+        b.createOrFold<arith::MulIOp>(location, KernelM1, DilDim);
+    Value EffectiveKernel =
+        b.createOrFold<arith::AddIOp>(location, KernelM1Dil, cstOne);
     Value IDim0KDim =
-        b.createOrFold<arith::AddIOp>(location, IDim0, kernelDimSizes[i]);
+        b.createOrFold<arith::AddIOp>(location, IDim0, EffectiveKernel);
     Value IDimPadDim = b.createOrFold<arith::AddIOp>(location, IDim, PadDim);
     Value IDim1 =
         b.createOrFold<arith::MinSIOp>(location, IDim0KDim, IDimPadDim);
 
-    Value IDim0Clamped =
-        b.createOrFold<arith::MaxSIOp>(location, IDim0, cstZero);
     Value IDim1Clamped = b.createOrFold<arith::MinSIOp>(location, IDim1, IDim);
-    Value IDim1_IDim0_Clamped =
-        b.createOrFold<arith::SubIOp>(location, IDim1Clamped, IDim0Clamped);
 
-    Value poolSizeDim =
-        !isCountIncludePad
-            ? IDim1_IDim0_Clamped
-            : b.createOrFold<arith::SubIOp>(location, IDim1, IDim0);
+    // Count valid taps using k_min/k_max approach:
+    // Tap k is valid when IDim0 + k*dilation is in [0, IDim1Clamped), where
+    // IDim1Clamped = min(window end, IDim) is the input extent clamped to the
+    // window.
+    // k_min = ceil(max(-IDim0, 0) / dilation)  -- first valid k
+    // k_max_excl = ceil((IDim1Clamped - IDim0) / dilation)  -- first k past the
+    // end ValidTaps = max(k_max_excl - k_min, 0)
+    Value NegIDim0 = b.createOrFold<arith::SubIOp>(location, cstZero, IDim0);
+    Value KMinNumer =
+        b.createOrFold<arith::MaxSIOp>(location, NegIDim0, cstZero);
+    Value KMin =
+        b.createOrFold<arith::CeilDivSIOp>(location, KMinNumer, DilDim);
+    Value IDim1ClampedMinusIDim0 =
+        b.createOrFold<arith::SubIOp>(location, IDim1Clamped, IDim0);
+    Value KMaxExcl = b.createOrFold<arith::CeilDivSIOp>(
+        location, IDim1ClampedMinusIDim0, DilDim);
+    Value KDiff = b.createOrFold<arith::SubIOp>(location, KMaxExcl, KMin);
+    Value ValidTaps = b.createOrFold<arith::MaxSIOp>(location, KDiff, cstZero);
+
+    // For count_include_pad: count every dilated tap position in the full
+    // effective window [IDim0, IDim1), including padded positions. Same as
+    // ValidTaps but WITHOUT the low-side clamp (padded taps are counted), so
+    // it reduces to k_min == 0:
+    //   FullTaps = max(ceil((IDim1 - IDim0) / dilation), 0)
+    // The max(.., 0) keeps this in agreement with ValidTaps on a degenerate
+    // empty window (range <= 0 => 0 taps).
+    Value FullRange = b.createOrFold<arith::SubIOp>(location, IDim1, IDim0);
+    Value FullTapsRaw =
+        b.createOrFold<arith::CeilDivSIOp>(location, FullRange, DilDim);
+    Value FullTaps =
+        b.createOrFold<arith::MaxSIOp>(location, FullTapsRaw, cstZero);
+
+    Value poolSizeDim = !isCountIncludePad ? ValidTaps : FullTaps;
     if (i == 0) {
       poolSize = poolSizeDim;
     } else {
@@ -1060,7 +1098,8 @@ class ConvertAtenAvgPoolOp : public OpConversionPattern<OpTy> {
   static bool
   doesAvgPoolDivisorNeedsClamping(bool ceilMode, bool countIncludePad,
                                   SmallVectorImpl<int64_t> &strideInts,
-                                  SmallVectorImpl<int64_t> &paddingInts);
+                                  SmallVectorImpl<int64_t> &paddingInts,
+                                  SmallVectorImpl<int64_t> &dilationInts);
 
   // Creates the average pooling operation value with a clamped
   // divisor. The clamped divisor is the product of kernel
@@ -1073,6 +1112,7 @@ class ConvertAtenAvgPoolOp : public OpConversionPattern<OpTy> {
       SmallVectorImpl<Value> &kernelDimSizes,
       SmallVectorImpl<int64_t> &strideInts,
       SmallVectorImpl<int64_t> &paddingInts,
+      SmallVectorImpl<int64_t> &dilationInts,
       SmallVector<AffineMap> &indexingMapsAvg,
       SmallVector<utils::IteratorType> &iteratorTypesAvg);
 
@@ -1147,11 +1187,11 @@ LogicalResult ConvertAtenAvgPoolOp<OpTy, PoolingOpTy, Dim>::matchAndRewrite(
       Dim + 2, utils::IteratorType::parallel);
 
   if (doesAvgPoolDivisorNeedsClamping(ceilMode, countIncludePad, strideInts,
-                                      paddingInts)) {
+                                      paddingInts, dilationInts)) {
     return createAveragePoolValueWithClampedDivisor(
         ceilMode, countIncludePad, op, adaptor, rewriter, self, sumPool,
         outputTensor, resultType, kernelSizeIntValues, strideInts, paddingInts,
-        indexingMapsAvg, iteratorTypesAvg);
+        dilationInts, indexingMapsAvg, iteratorTypesAvg);
   }
 
   return createAveragePoolValueWithRegularDivisor(
@@ -1163,7 +1203,8 @@ template <typename OpTy, typename PoolingOpTy, int Dim>
 bool ConvertAtenAvgPoolOp<OpTy, PoolingOpTy, Dim>::
     doesAvgPoolDivisorNeedsClamping(bool ceilMode, bool countIncludePad,
                                     SmallVectorImpl<int64_t> &strideInts,
-                                    SmallVectorImpl<int64_t> &paddingInts) {
+                                    SmallVectorImpl<int64_t> &paddingInts,
+                                    SmallVectorImpl<int64_t> &dilationInts) {
   // Determines whether the average pooling divisor needs to be clamped
   // (i.e., adjusted to exclude padded or out-of-bounds elements).
   //
@@ -1193,8 +1234,11 @@ bool ConvertAtenAvgPoolOp<OpTy, PoolingOpTy, Dim>::
       !llvm::all_of(paddingInts, [](int64_t p) { return p == 0; });
   bool allStridesUnitary =
       llvm::all_of(strideInts, [](int64_t s) { return s == 1; });
+  bool allDilationsUnitary =
+      llvm::all_of(dilationInts, [](int64_t d) { return d == 1; });
 
-  return (!countIncludePad && hasPadding) || (ceilMode && !allStridesUnitary);
+  return (!countIncludePad && hasPadding) || (ceilMode && !allStridesUnitary) ||
+         !allDilationsUnitary;
 }
 
 template <typename OpTy, typename PoolingOpTy, int Dim>
@@ -1206,6 +1250,7 @@ LogicalResult ConvertAtenAvgPoolOp<OpTy, PoolingOpTy, Dim>::
         SmallVectorImpl<Value> &kernelDimSizes,
         SmallVectorImpl<int64_t> &strideInts,
         SmallVectorImpl<int64_t> &paddingInts,
+        SmallVectorImpl<int64_t> &dilationInts,
         SmallVector<AffineMap> &indexingMapsAvg,
         SmallVector<utils::IteratorType> &iteratorTypesAvg) {
   Location loc = op->getLoc();
@@ -1241,7 +1286,7 @@ LogicalResult ConvertAtenAvgPoolOp<OpTy, PoolingOpTy, Dim>::
           [&](OpBuilder &b, Location loc, ValueRange args) {
             if (!poolSize) {
               poolSize = poolSizeCalculator.getPoolSize(
-                  b, kernelDimSizes, strideInts, paddingInts);
+                  b, kernelDimSizes, strideInts, paddingInts, dilationInts);
             }
             Value divisor =
                 convertScalarToDtype(b, loc, poolSize, resultElementType);

test/Conversion/TorchToLinalg/pooling.mlir

@@ -291,3 +291,105 @@ func.func @forward_avgpool_2d_ceil(%arg0: !torch.vtensor<[1,1,4,4],f32>) -> !tor
   %3 = torch.aten.avg_pool2d %arg0, %0, %2, %1, %true, %false, %none : !torch.vtensor<[1,1,4,4],f32>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[1,1,2,2],f32>
   return %3 : !torch.vtensor<[1,1,2,2],f32>
 }
+
+// -----
+
+// CHECK: #[[$MAP:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+// CHECK-LABEL: func @forward_avgpool_2d_ceil_dilated
+func.func @forward_avgpool_2d_ceil_dilated(%arg0: !torch.vtensor<[1,1,7,7],f32>) -> !torch.vtensor<[1,1,2,2],f32> {
+  // Sum pool uses dilation=2:
+  // CHECK: linalg.pooling_nchw_sum {dilations = dense<2> : vector<2xi64>, strides = dense<3> : vector<2xi64>}
+
+  // The divisor generic uses dilation to count valid taps per dim:
+  // effectiveKernel = (3-1)*2+1 = 5, then ceildivsi to count valid taps.
+  // CHECK: linalg.generic {indexing_maps = [#[[$MAP]], #[[$MAP]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+  // CHECK:   %[[C5:.*]] = arith.constant 5 : i64
+  // CHECK:   arith.addi %{{.*}}, %[[C5]] : i64
+  // CHECK:   arith.ceildivsi {{.*}} : i64
+  // CHECK:   arith.divf
+  // CHECK:   linalg.yield
+
+  %int3 = torch.constant.int 3
+  %int3_0 = torch.constant.int 3
+  %int0 = torch.constant.int 0
+  %int0_1 = torch.constant.int 0
+  %int3_s = torch.constant.int 3
+  %int3_s2 = torch.constant.int 3
+  %int2_d = torch.constant.int 2
+  %int2_d2 = torch.constant.int 2
+  // stride list = [stride_h, stride_w, dilation_h, dilation_w]
+  %0 = torch.prim.ListConstruct %int3, %int3_0 : (!torch.int, !torch.int) -> !torch.list<int>
+  %1 = torch.prim.ListConstruct %int0, %int0_1 : (!torch.int, !torch.int) -> !torch.list<int>
+  %2 = torch.prim.ListConstruct %int3_s, %int3_s2, %int2_d, %int2_d2 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
+  %true = torch.constant.bool true
+  %false = torch.constant.bool false
+  %none = torch.constant.none
+  %3 = torch.aten.avg_pool2d %arg0, %0, %2, %1, %true, %false, %none : !torch.vtensor<[1,1,7,7],f32>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[1,1,2,2],f32>
+  return %3 : !torch.vtensor<[1,1,2,2],f32>
+}
+
+// -----
+
+// CHECK: #[[$MAP:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+// CHECK-LABEL: func @forward_avgpool_2d_count_include_pad_dilated
+func.func @forward_avgpool_2d_count_include_pad_dilated(%arg0: !torch.vtensor<[1,1,7,7],f32>) -> !torch.vtensor<[1,1,3,3],f32> {
+  // CHECK: linalg.pooling_nchw_sum {dilations = dense<2> : vector<2xi64>, strides = dense<2> : vector<2xi64>}
+
+  // count_include_pad=true with dilation: ValidTaps uses ceildivsi; FullTaps uses divsi.
+  // CHECK: linalg.generic {indexing_maps = [#[[$MAP]], #[[$MAP]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+  // CHECK:   %[[C5:.*]] = arith.constant 5 : i64
+  // CHECK:   arith.addi %{{.*}}, %[[C5]] : i64
+  // CHECK:   arith.ceildivsi {{.*}} : i64
+  // CHECK:   arith.divf
+  // CHECK:   linalg.yield
+
+  %int3 = torch.constant.int 3
+  %int3_0 = torch.constant.int 3
+  %int1 = torch.constant.int 1
+  %int1_1 = torch.constant.int 1
+  %int2_s = torch.constant.int 2
+  %int2_s2 = torch.constant.int 2
+  %int2_d = torch.constant.int 2
+  %int2_d2 = torch.constant.int 2
+  // stride list = [stride_h, stride_w, dilation_h, dilation_w]
+  %0 = torch.prim.ListConstruct %int3, %int3_0 : (!torch.int, !torch.int) -> !torch.list<int>
+  %1 = torch.prim.ListConstruct %int1, %int1_1 : (!torch.int, !torch.int) -> !torch.list<int>
+  %2 = torch.prim.ListConstruct %int2_s, %int2_s2, %int2_d, %int2_d2 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
+  %false = torch.constant.bool false
+  %true = torch.constant.bool true
+  %none = torch.constant.none
+  %3 = torch.aten.avg_pool2d %arg0, %0, %2, %1, %false, %true, %none : !torch.vtensor<[1,1,7,7],f32>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[1,1,3,3],f32>
+  return %3 : !torch.vtensor<[1,1,3,3],f32>
+}
+
+// -----
+
+// CHECK-LABEL: func @forward_avgpool_2d_exclude_pad_dilated_edge
+func.func @forward_avgpool_2d_exclude_pad_dilated_edge(%arg0: !torch.vtensor<[1,1,11,11],f32>) -> !torch.vtensor<[1,1,1,1],f32> {
+  // kernel=4, dilation=4, padding=1, stride=1, count_include_pad=false.
+  // effective_kernel = (4-1)*4+1 = 13; output = (11+2-13)/1+1 = 1.
+  // At output position 0: start_idx = 0*1 - 1 = -1. Taps land at -1, 3, 7, 11.
+  // Valid taps (in [0,11)): 3 and 7 -> ValidTaps = 2.
+  // Correct divisor is 2.
+  // CHECK: linalg.pooling_nchw_sum {dilations = dense<4> : vector<2xi64>, strides = dense<1> : vector<2xi64>}
+  // CHECK: linalg.generic
+  // CHECK:   arith.ceildivsi
+  // CHECK:   arith.divf
+  // CHECK:   linalg.yield
+  %int4   = torch.constant.int 4
+  %int4_0 = torch.constant.int 4
+  %int1   = torch.constant.int 1
+  %int1_1 = torch.constant.int 1
+  %int1_s = torch.constant.int 1
+  %int1_s2 = torch.constant.int 1
+  %int4_d = torch.constant.int 4
+  %int4_d2 = torch.constant.int 4
+  %0 = torch.prim.ListConstruct %int4, %int4_0 : (!torch.int, !torch.int) -> !torch.list<int>
+  %1 = torch.prim.ListConstruct %int1, %int1_1 : (!torch.int, !torch.int) -> !torch.list<int>
+  %2 = torch.prim.ListConstruct %int1_s, %int1_s2, %int4_d, %int4_d2 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
+  %false = torch.constant.bool false
+  %false_1 = torch.constant.bool false
+  %none = torch.constant.none
+  %3 = torch.aten.avg_pool2d %arg0, %0, %2, %1, %false, %false_1, %none : !torch.vtensor<[1,1,11,11],f32>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[1,1,1,1],f32>
+  return %3 : !torch.vtensor<[1,1,1,1],f32>
+}