fix: rewrite avg pool as 2-d nchw fchw convolution

This fixes the bug where before the avg pooling was split into per-channel 2d convolutions, but that was incorrect since that doesn't properly compute the average since linalg.conv_2d doesnt support striding.

This temporarily disables lenet since that includes the pooling layer

Fixes #2702

PiperOrigin-RevId: 888710468

Author: asraa

lib/Transforms/LinalgCanonicalizations/LinalgCanonicalizations.cpp

@@ -653,8 +653,9 @@ struct RewriteAvgPoolAsConv2D
     auto filterTy = cast<RankedTensorType>(poolOp.getInputs()[1].getType());
     auto outputTy = cast<RankedTensorType>(poolOp.getResultTypes()[0]);
 
-    auto kernelShape =
-        SmallVector<int64_t>{filterTy.getDimSize(0), filterTy.getDimSize(1)};
+    auto c = inputTy.getDimSize(1);
+    auto kernelShape = SmallVector<int64_t>{c, c, filterTy.getDimSize(0),
+                                            filterTy.getDimSize(1)};
     auto kernelTy =
         RankedTensorType::get(kernelShape, filterTy.getElementType());
     TypedAttr kernelVals = rewriter.getOneAttr(kernelTy);
@@ -681,55 +682,17 @@ struct RewriteAvgPoolAsConv2D
     }
     auto kernel =
         arith::ConstantOp::create(rewriter, poolOp.getLoc(), kernelVals);
-
-    // Rewrite the 2D avg pool output shape is N x C x H' x W'. Apply the kernel
-    // to each channel separately, and insert into the output.
-    auto outputVal = poolOp.getOutputs()[0];
-    //  The filter must ensure each output channel i is only the sum of values
-    //  from input channel i. So the filter uses an identity matrix when f ==
-    //  c and 0 otherwise.
-    RankedTensorType twoDOutputType =
-        RankedTensorType::get({outputTy.getDimSize(2), outputTy.getDimSize(3)},
-                              outputTy.getElementType());
-    RankedTensorType twoDInputType =
-        RankedTensorType::get({inputTy.getDimSize(2), inputTy.getDimSize(3)},
-                              inputTy.getElementType());
-    Value convOutput = tensor::EmptyOp::create(rewriter, poolOp.getLoc(),
-                                               twoDOutputType.getShape(),
-                                               twoDOutputType.getElementType());
-    for (int n = 0; n < inputTy.getDimSize(0); ++n) {
-      for (int c = 0; c < inputTy.getDimSize(1); ++c) {
-        // Compute the 2-D constant convolution.
-        SmallVector<OpFoldResult> offsets = {
-            rewriter.getIndexAttr(n), rewriter.getIndexAttr(c),
-            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-        SmallVector<OpFoldResult> inputSizes = {
-            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
-            rewriter.getIndexAttr(inputTy.getDimSize(2)),
-            rewriter.getIndexAttr(inputTy.getDimSize(3))};
-        SmallVector<OpFoldResult> strides(4, rewriter.getIndexAttr(1));
-        auto extractInputOp = tensor::ExtractSliceOp::create(
-            rewriter, poolOp.getLoc(), twoDInputType, poolOp.getInputs()[0],
-            offsets, inputSizes, strides);
-
-        auto convOp = linalg::Conv2DOp::create(
-            rewriter, poolOp.getLoc(), twoDOutputType,
-            ValueRange{extractInputOp, kernel}, ValueRange{convOutput});
-        // Insert into the outputVal.
-        SmallVector<OpFoldResult> outputSizes = {
-            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
-            rewriter.getIndexAttr(outputTy.getDimSize(2)),
-            rewriter.getIndexAttr(outputTy.getDimSize(3))};
-        outputVal = tensor::InsertSliceOp::create(
-            rewriter, poolOp.getLoc(), convOp.getResult(0), outputVal, offsets,
-            outputSizes, strides);
-      }
-    }
+    Value conv = linalg::Conv2DNchwFchwOp::create(
+                     rewriter, poolOp.getLoc(), outputTy,
+                     ValueRange{poolOp.getInputs()[0], kernel},
+                     ValueRange{poolOp.getOutputs()[0]}, poolOp.getStrides(),
+                     poolOp.getDilations())
+                     .getResult(0);
 
     if (avgPoolOutput) {
-      rewriter.replaceAllUsesWith(avgPoolOutput, outputVal);
+      rewriter.replaceAllUsesWith(avgPoolOutput, conv);
     } else {
-      rewriter.replaceOp(poolOp, outputVal);
+      rewriter.replaceOp(poolOp, conv);
     }
     return success();
   }
@@ -746,6 +709,14 @@ struct LowerConv2DNchwFchw
 
   LogicalResult matchAndRewrite(mlir::linalg::Conv2DNchwFchwOp convOp,
                                 PatternRewriter& rewriter) const override {
+    // Fails is strides > 1.
+    if (!llvm::all_of(convOp.getStrides(), [](const APInt& element) {
+          return element.getSExtValue() == 1;
+        })) {
+      return rewriter.notifyMatchFailure(convOp,
+                                         "expected all ones for strides");
+    }
+
     Location loc = convOp.getLoc();
     Value image = convOp.getInputs()[0];
     Value filter = convOp.getInputs()[1];

tests/Examples/openfhe/ckks/lenet/BUILD

@@ -1,5 +1,6 @@
-load("@heir//tools:heir-opt.bzl", "heir_opt")
 load("@pybind11_bazel//:build_defs.bzl", "pybind_extension")
+# load("@heir//tools:heir-opt.bzl", "heir_opt")
+
 load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
 load("@rules_cc//cc:cc_library.bzl", "cc_library")
 load("@rules_python//python:py_test.bzl", "py_test")
@@ -39,20 +40,21 @@ cc_binary(
     ],
 )
 
-heir_opt(
-    name = "lenet_mlir_opt",
-    src = "lenet.mlir",
-    generated_filename = "lenet.openfhe.mlir",
-    pass_flags = [
-        "--annotate-module=backend=openfhe scheme=ckks",
-        "--torch-linalg-to-ckks=ciphertext-degree=1024",
-        "--scheme-to-openfhe",
-    ],
-    tags = [
-        "nofastbuild",
-        "requires-mem:28g",
-    ],
-)
+# TODO(#2702): Re-enable once pooling is fully supported.
+# heir_opt(
+#     name = "lenet_mlir_opt",
+#     src = "lenet.mlir",
+#     generated_filename = "lenet.openfhe.mlir",
+#     pass_flags = [
+#         "--annotate-module=backend=openfhe scheme=ckks",
+#         "--torch-linalg-to-ckks=ciphertext-degree=1024",
+#         "--scheme-to-openfhe",
+#     ],
+#     tags = [
+#         "nofastbuild",
+#         "requires-mem:28g",
+#     ],
+# )
 
 cc_library(
     name = "interpreter_shim",

tests/Transforms/linalg_canonicalizations/average_pooling.mlir

@@ -6,11 +6,11 @@
 module {
   // CHECK: func.func @main
   // CHECK-SAME: (%[[arg0:.*]]: tensor<1x6x28x28xf32>)
+  // CHECK-DAG: %[[out:.*]] = arith.constant dense<2.0
   // CHECK-DAG: %[[divided_cst:.*]] = arith.constant dense<2.500000e-01>
-  // CHECK-DAG: %[[out:.*]] = tensor.empty() : tensor<14x14xf32>
-  // CHECK-DAG: %[[extracted:.*]] = tensor.extract_slice %[[arg0]]
-  // CHECK: linalg.conv_2d ins(%[[extracted]], %[[divided_cst]] : tensor<28x28xf32>, tensor<2x2xf32>) outs(%[[out]] : tensor<14x14xf32>)
-  // CHECK-COUNT-5: linalg.conv_2d
+  // CHECK: linalg.conv_2d_nchw_fchw
+  // CHECK-SAME: strides = dense<2> : vector<2xi64>
+  // CHECK-SAME: ins(%[[arg0]], %[[divided_cst]] : tensor<1x6x28x28xf32>, tensor<6x6x2x2xf32>) outs(%[[out]] : tensor<1x6x14x14xf32>)
   // CHECK: return
   func.func @main(%arg0: tensor<1x6x28x28xf32>) -> tensor<1x6x14x14xf32> {
     %cst_0 = arith.constant 2.000000e+00 : f32