Merge pull request #3303 from google:mdgrs/layoutPropConvFixes

PiperOrigin-RevId: 960967691

Author: copybara-github

lib/Kernel/KernelImplementationTest.cpp

@@ -1,6 +1,7 @@
 #include <cstdint>
 #include <functional>
 #include <map>
+#include <string>
 #include <vector>
 
 #include "gtest/gtest.h"  // from @googletest
@@ -737,6 +738,74 @@ TEST_P(KernelImplementationTest, TestConv1dCwFcwStride2) {
   EXPECT_EQ(actualUnpacked, expected);
 }
 
+// End-to-end Halevi-Shoup matvec for a padded strided 1-D multichannel conv, in
+// the packing production uses when LayoutPropagation folds a zero tensor.pad on
+// the width dim into the conv's own `padding` parameter: the data ciphertext is
+// packed row-major at the *unpadded* width, and the Toeplitz matrix is built
+// with padding = p so that a window reaching into the padding contributes no
+// column. `expected` is the convolution of the zero-padded data.
+void checkPaddedConv1dCwFcw(int64_t padding, const tensor3d& expected,
+                            bool unroll, bool interchangeRows) {
+  SCOPED_TRACE("padding = " + std::to_string(padding));
+  MLIRContext context;
+  // 1x2x6 input data, 2x2x3 filter, stride 2.
+  tensor3d data = {{{0, 1, 2, 3, 4, 5}, {6, 7, 8, 9, 10, 11}}};
+  tensor3d filter = {{{3, 4, 1}, {1, 5, 2}}, {{1, 2, 3}, {2, 2, 2}}};
+  int64_t stride = 2;
+  int numSlots = 16;
+
+  RankedTensorType dataType =
+      RankedTensorType::get({1, 2, 6}, mlir::IndexType::get(&context));
+  RankedTensorType filterType =
+      RankedTensorType::get({2, 2, 3}, mlir::IndexType::get(&context));
+
+  auto dataLayout = getRowMajorLayoutRelation(dataType, numSlots);
+  std::vector<std::vector<int>> packedData =
+      evaluateLayout(dataLayout, getDataValueFn3D(data));
+
+  auto filterLayout = get1dConvCwFcwFilterDiagonalizedRelation(
+      filterType, dataType, stride, padding, numSlots, interchangeRows);
+  ASSERT_TRUE(succeeded(filterLayout));
+  std::function<int(const std::vector<int64_t>&)> getFilterValueFn =
+      [&](const std::vector<int64_t>& domainPoint) -> int {
+    return filter[domainPoint[0]][domainPoint[1]][domainPoint[2]];
+  };
+  std::vector<std::vector<int>> packedFilter =
+      evaluateLayout(filterLayout.value(), getFilterValueFn);
+  // The matrix shape must be derived the same way the filter was diagonalized,
+  // i.e. against the unpadded data type with padding = p: implementHaleviShoup
+  // sizes the squat-diagonal collapse from nextPowerOfTwo of these dims.
+  auto expandedFilterShape =
+      get1dConvCwFcwFilterExpandedType(filterType, dataType, stride, padding);
+
+  auto dag = implementHaleviShoup(
+      LiteralValue(packedData[0]), LiteralValue(packedFilter),
+      expandedFilterShape.getShape(), DagType::intTensor(32, {numSlots}),
+      /*zeroDiagonals=*/{}, unroll);
+  auto actual = std::get<std::vector<int>>(evalKernel(dag)[0].get());
+
+  int64_t outputWidth = expected[0][0].size();
+  RankedTensorType outputType = RankedTensorType::get(
+      {1, 2, outputWidth}, mlir::IndexType::get(&context));
+  auto resultLayout = get1dConvResultRelation(outputType, stride, /*padding=*/0,
+                                              numSlots, interchangeRows);
+
+  EXPECT_EQ(
+      unpackLayoutTo3DTensor<int>(resultLayout, {actual}, {1, 2, outputWidth}),
+      expected);
+}
+
+TEST_P(KernelImplementationTest, TestConv1dCwFcwStride2WithPadding) {
+  // Same shape family as TestConv1dCwFcwStride2, but with padding != 0
+  // padding 1 keeps the unpadded and padded column counts in the same
+  // power-of-two bucket (2*6=12 and 2*8=16 both round to 16); padding 2 does
+  // not (2*6=12 rounds to 16, 2*10=20 rounds to 32).
+  checkPaddedConv1dCwFcw(/*padding=*/1, {{{45, 79, 111}, {29, 62, 86}}},
+                         std::get<0>(GetParam()), std::get<1>(GetParam()));
+  checkPaddedConv1dCwFcw(/*padding=*/2, {{{12, 63, 95, 97}, {12, 50, 74, 56}}},
+                         std::get<0>(GetParam()), std::get<1>(GetParam()));
+}
+
 TEST_P(KernelImplementationTest,
        TestConv2dNchwFchwStride2InterchangedLargeSlots) {
   MLIRContext context;

lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp

@@ -1156,7 +1156,29 @@ struct ConvertLinalgConv1DNcwFcw
     return isPowerOfTwoDims && isConv1dAsMatvec;
   }
 
-  void haleviShoupKernel(
+  // Rebuild the operand shape LayoutPropagation diagonalized the filter
+  // against, honoring any zero tensor.pad it folded into the conv's `padding`
+  // parameter.
+  FailureOr<RankedTensorType> expandedFilterShape(
+      linalg::Conv1DNcwFcwOp op) const {
+    auto filterType = cast<RankedTensorType>(op.getInputs()[1].getType());
+    auto dataType = cast<RankedTensorType>(op.getInputs()[0].getType());
+    int64_t stride =
+        llvm::to_vector(op.getStrides().getValues<int64_t>()).front();
+
+    int64_t padding = getConvFoldedPadding(op);
+    std::optional<ConvMatrixOperand> matrixOperand =
+        foldConvWidthPadding(dataType, padding);
+    if (!matrixOperand) {
+      return op.emitError()
+             << kConvFoldedPaddingAttrName << " of " << padding
+             << " does not fit this conv's data operand " << dataType;
+    }
+    return get1dConvCwFcwFilterExpandedType(filterType, matrixOperand->dataType,
+                                            stride, matrixOperand->padding);
+  }
+
+  LogicalResult haleviShoupKernel(
       linalg::Conv1DNcwFcwOp op, OpAdaptor adaptor,
       ContextAwareConversionPatternRewriter& rewriter) const {
     LLVM_DEBUG(
@@ -1171,13 +1193,8 @@ struct ConvertLinalgConv1DNcwFcw
         cast<TypedValue<RankedTensorType>>(adaptor.getInputs()[1]);
     SSAValue matrixLeaf(matrix);
 
-    // The original matrix shape is the shape of the expanded filter before
-    // diagonalization.
-    RankedTensorType expandedMatrixType = get1dConvCwFcwFilterExpandedType(
-        cast<RankedTensorType>(op.getInputs()[1].getType()),
-        cast<RankedTensorType>(op.getInputs()[0].getType()),
-        llvm::to_vector(op.getStrides().getValues<int64_t>()).front(),
-        /*padding=*/0);
+    FailureOr<RankedTensorType> expandedMatrixType = expandedFilterShape(op);
+    if (failed(expandedMatrixType)) return failure();
     // Collect any zero diagonals of the filter matrix.
     LayoutAttr filterLayout = getLayoutAttr(adaptor.getInputs()[1]);
     auto filterRelation = filterLayout.getIntegerRelation();
@@ -1196,7 +1213,7 @@ struct ConvertLinalgConv1DNcwFcw
                                              data.getType().getShape().back());
     std::shared_ptr<ArithmeticDagNode<SSAValue>> implementedKernel =
         implementHaleviShoup(vectorLeaf, matrixLeaf,
-                             expandedMatrixType.getShape(), dagType,
+                             expandedMatrixType->getShape(), dagType,
                              zeroDiagonals,
                              /*unroll=*/unrollKernels);
 
@@ -1210,6 +1227,7 @@ struct ConvertLinalgConv1DNcwFcw
     // Add the initial accumulator value.
     Value result = adaptor.getOutputs()[0];
     addBiasAndReplace(rewriter, op, finalOutput, result, layoutAttr);
+    return success();
   }
 
   LogicalResult matchAndRewrite(
@@ -1226,8 +1244,7 @@ struct ConvertLinalgConv1DNcwFcw
     }
 
     if (supportsExpandedHaleviShoup(op, adaptor)) {
-      haleviShoupKernel(op, adaptor, rewriter);
-      return success();
+      return haleviShoupKernel(op, adaptor, rewriter);
     }
 
     return op.emitError() << "unsupported layout for 1d conv";
@@ -1303,6 +1320,10 @@ struct ConvertLinalgConv2DNchwFchw
 
     // The original matrix shape is the shape of the expanded filter before
     // diagonalization.
+    // NOTE: `padding=0` is only correct because nothing folds a `tensor.pad`
+    // into a 2-D conv's padding parameter. If that changes, go through
+    // foldConvWidthPadding the way ConvertLinalgConv1DNcwFcw does; see
+    // ConvMatrixOperand.
     RankedTensorType expandedMatrixType = get2dConvChwFchwFilterExpandedType(
         cast<RankedTensorType>(op.getInputs()[1].getType()), dataType,
         /*padding=*/0, llvm::to_vector(op.getStrides().getValues<int64_t>()));