improve handling of padding for conv1d

Author: mdgrs

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

@@ -1153,7 +1153,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(
@@ -1168,13 +1190,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();
@@ -1193,7 +1210,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);
 
@@ -1207,6 +1224,7 @@ struct ConvertLinalgConv1DNcwFcw
     // Add the initial accumulator value.
     Value result = adaptor.getOutputs()[0];
     addBiasAndReplace(rewriter, op, finalOutput, result, layoutAttr);
+    return success();
   }
 
   LogicalResult matchAndRewrite(
@@ -1223,8 +1241,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";
@@ -1300,6 +1317,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>()));

lib/Transforms/LayoutPropagation/LayoutPropagation.cpp

@@ -49,6 +49,7 @@
 #include "mlir/include/mlir/IR/BuiltinAttributes.h"      // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinTypes.h"           // from @llvm-project
 #include "mlir/include/mlir/IR/Diagnostics.h"            // from @llvm-project
+#include "mlir/include/mlir/IR/Matchers.h"               // from @llvm-project
 #include "mlir/include/mlir/IR/Operation.h"              // from @llvm-project
 #include "mlir/include/mlir/IR/PatternMatch.h"           // from @llvm-project
 #include "mlir/include/mlir/IR/Types.h"                  // from @llvm-project
@@ -134,9 +135,95 @@ std::pair<Value, LayoutAttr> convertToLayout(
   return std::make_pair(toReplace, layoutAttr);
 }
 
+// The outcome of folding a zero `tensor.pad` on the width dim into a 1-D
+// conv's own `padding` parameter.
+struct FoldedConvPadding {
+  // What the Toeplitz matrix must be built against.
+  ConvMatrixOperand matrixOperand;
+  // The layout the padded operand must already carry for the fold to be valid.
+  IntegerRelation targetRelation;
+};
+
+// Try to fold a zero `tensor.pad` on the width dim of a conv's `data` operand
+// into the conv's own `padding` parameter.
+//
+// Returns nullopt when the pattern does not apply, in which case the caller
+// keeps the unfolded path. `dataType` must be rank 3 with N=1.
+std::optional<FoldedConvPadding> tryFoldPadIntoConvPadding(
+    Value data, RankedTensorType dataType, LayoutAttr dataLayout,
+    int64_t ciphertextSize) {
+  assert(dataType.getRank() == 3 && dataType.getDimSize(0) == 1 &&
+         "expected a rank-3 N=1 conv data operand");
+
+  // DropUnitDims rewrites a rank-3 pad into
+  // collapse_shape -> tensor.pad (rank 2) -> expand_shape, so peel any
+  // reshape/cast chain to find the pad.
+  Value cursor = data;
+  tensor::PadOp padOp;
+  while (Operation* def = cursor.getDefiningOp()) {
+    if (auto p = dyn_cast<tensor::PadOp>(def)) {
+      padOp = p;
+      break;
+    }
+    if (isa<tensor::ExpandShapeOp, tensor::CollapseShapeOp, tensor::CastOp>(
+            def)) {
+      cursor = def->getOperand(0);
+      continue;
+    }
+    break;
+  }
+  if (!padOp) return std::nullopt;
+
+  Value padValue = padOp.getConstantPaddingValue();
+  bool zeroPad = padValue && (matchPattern(padValue, m_AnyZeroFloat()) ||
+                              matchPattern(padValue, m_Zero()));
+  if (!zeroPad) return std::nullopt;
+
+  // Only a symmetric pad on the trailing (width) dim is expressible as the
+  // conv's `padding` parameter, and only when the bounds are static. The pad
+  // may be rank 2 (unit batch dim dropped) or rank 3.
+  if (!padOp.getLow().empty() || !padOp.getHigh().empty()) return std::nullopt;
+  ArrayRef<int64_t> low = padOp.getStaticLow();
+  ArrayRef<int64_t> high = padOp.getStaticHigh();
+  bool widthOnly = !low.empty() && low.size() == high.size() &&
+                   low.back() == high.back() && low.back() > 0;
+  for (size_t i = 0; widthOnly && i + 1 < low.size(); ++i) {
+    widthOnly &= low[i] == 0 && high[i] == 0;
+  }
+  if (!widthOnly) return std::nullopt;
+
+  // The reshape chain above means the pad's width dim is not guaranteed to be
+  // this conv operand's width dim, so validate before building a type from it.
+  int64_t p = low.back();
+  std::optional<ConvMatrixOperand> matrixOperand =
+      foldConvWidthPadding(dataType, p);
+  if (!matrixOperand) return std::nullopt;
+
+  // The layout we expect on the padded value: the unpadded row-major layout
+  // with the width index shifted by `p`. If the actual layout is anything else
+  // (a conversion intervened, a non-row-major producer, reshapes that did not
+  // cancel) do not fold, rather than silently mis-indexing the matrix.
+  IntegerRelation expected =
+      getRowMajorLayoutRelation(matrixOperand->dataType, ciphertextSize);
+  expected = shiftVar(
+      expected, expected.getVarKindOffset(presburger::VarKind::Domain) + 2, p);
+  if (!dataLayout.getIntegerRelation().isEqual(expected)) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "conv_1d found a pad of " << p
+               << " but the operand layout does not match the shifted "
+                  "unpadded row-major layout; not folding\n");
+    return std::nullopt;
+  }
+
+  LLVM_DEBUG(llvm::dbgs() << "conv_1d folding tensor.pad of " << p
+                          << " into the conv padding parameter\n");
+  return FoldedConvPadding{*matrixOperand, expected};
+}
+
 // Return a copy of the kernel info associated with the value and update the
 // result shape to the new result shape. If the value does not have a kernel
 // info, return an empty Attribute.
+
 Attribute cloneKernelInfoWithResultShape(Value value,
                                          ArrayRef<int64_t> resultShape) {
   auto kernelInfo = findAttributeAssociatedWith(value, kKernelInfoAttrName);
@@ -179,6 +266,7 @@ struct LayoutPropagation : impl::LayoutPropagationBase<LayoutPropagation> {
   LogicalResult visitOperation(tensor::InsertOp op);
   LogicalResult visitOperation(tensor::InsertSliceOp op);
   LogicalResult visitOperation(tensor::ExtractSliceOp op);
+  LogicalResult visitOperation(tensor::PadOp op);
 
   // Determine if the operation arguments have compatible layouts for the
   // given op. If the check fails, the CompatibilityResult::compatible field
@@ -353,8 +441,8 @@ LogicalResult LayoutPropagation::visitOperation(Operation* op) {
       .Case<affine::AffineForOp>([&](auto op) { return visitOperation(op); })
       // tensor ops
       .Case<tensor::ExtractOp, tensor::InsertOp, tensor::InsertSliceOp,
-            tensor::ExtractSliceOp, CollapseShapeOp, ExpandShapeOp>(
-          [&](auto op) { return visitOperation(op); })
+            tensor::ExtractSliceOp, tensor::PadOp, CollapseShapeOp,
+            ExpandShapeOp>([&](auto op) { return visitOperation(op); })
       // AddI, AddF, mgmt.* all pass the layout through unchanged.
       .Default([&](Operation* op) {
         passLayoutThroughOp(op);
@@ -907,8 +995,15 @@ LogicalResult LayoutPropagation::visitOperation(Conv1DNcwFcwOp op) {
   }
 
   LayoutAttr dataLayout = getComposedLayoutAttr(data);
+
+  ConvMatrixOperand matrixOperand{dataType};
   IntegerRelation targetDataRelation =
       getRowMajorLayoutRelation(dataType, minSlotCount);
+  if (auto folded =
+          tryFoldPadIntoConvPadding(data, dataType, dataLayout, minSlotCount)) {
+    matrixOperand = folded->matrixOperand;
+    targetDataRelation = folded->targetRelation;
+  }
 
   if (!isRelationEqual(dataLayout.getIntegerRelation(), targetDataRelation)) {
     LLVM_DEBUG(llvm::dbgs() << "conv_1d data input is not row major, "
@@ -923,8 +1018,8 @@ LogicalResult LayoutPropagation::visitOperation(Conv1DNcwFcwOp op) {
   // into a larger matrix and then diagonalizing.
   LayoutAttr filterLayout = getComposedLayoutAttr(filter);
   auto convRelation = get1dConvCwFcwFilterDiagonalizedRelation(
-      filterType, dataType, stride, /*padding=*/0, minSlotCount,
-      /*interchangeRows=*/interchangeRows);
+      filterType, matrixOperand.dataType, stride, matrixOperand.padding,
+      minSlotCount, /*interchangeRows=*/interchangeRows);
   if (failed(convRelation)) {
     return failure();
   }
@@ -956,6 +1051,10 @@ LogicalResult LayoutPropagation::visitOperation(Conv1DNcwFcwOp op) {
   Attribute kernelInfoAttr =
       cloneKernelInfoWithResultShape(data, outputType.getShape());
 
+  // Record what the filter was diagonalized against, so that
+  // ConvertToCiphertextSemantics rebuilds the same expanded matrix shape.
+  setConvFoldedPadding(op, matrixOperand.padding);
+
   assignedLayouts.insert({result, resultLayoutAttr});
   setResultLayoutAttr(op, kernelInfoAttr);
   auto kernelAttr =
@@ -1395,6 +1494,48 @@ LogicalResult LayoutPropagation::visitOperation(tensor::InsertSliceOp op) {
   return success();
 }
 
+LogicalResult LayoutPropagation::visitOperation(tensor::PadOp op) {
+  // A zero-pad does not move any data: result[i + low] = source[i], so the
+  // result layout is the source layout with each domain index shifted by the
+  // low padding. Pad positions stay unmapped in the relation; unmapped points
+  // are zero-filled when a layout is materialized, which matches the
+  // zero-fill pad body. (The default layout passthrough would instead map
+  // result[i] to source[i]'s slots — off by `low`, shifting every downstream
+  // consumer's reads by one stride per pad.)
+  if (!assignedLayouts.contains(op.getSource())) {
+    return op->emitError() << "Source tensor has no assigned layout";
+  }
+  Value padValue = op.getConstantPaddingValue();
+  if (!padValue || !(matchPattern(padValue, m_AnyZeroFloat()) ||
+                     matchPattern(padValue, m_Zero()))) {
+    return op->emitError()
+           << "layout propagation only supports zero-padding tensor.pad";
+  }
+  if (!op.getLow().empty() || !op.getHigh().empty()) {
+    return op->emitError()
+           << "layout propagation requires static tensor.pad bounds";
+  }
+
+  IntegerRelation padRelation =
+      getComposedLayoutAttr(op.getSource()).getIntegerRelation();
+  auto domainVarOffset =
+      padRelation.getVarKindOffset(presburger::VarKind::Domain);
+  for (auto [dim, low] : llvm::enumerate(op.getStaticLow())) {
+    if (low != 0) {
+      padRelation = shiftVar(padRelation, domainVarOffset + dim, low);
+    }
+  }
+
+  LayoutAttr outputLayout =
+      LayoutAttr::getFromIntegerRelation(op.getContext(), padRelation);
+  Attribute kernelInfoAttr = cloneKernelInfoWithResultShape(
+      op.getSource(), op.getResultType().getShape());
+  assignedLayouts.insert({op.getResult(), outputLayout});
+  debugAssignLayout(op.getResult(), outputLayout);
+  setResultLayoutAttr(op, kernelInfoAttr);
+  return success();
+}
+
 LogicalResult LayoutPropagation::visitOperation(tensor::ExtractSliceOp op) {
   // Assign the induced layout from extracting a slice from the source tensor.
   if (!assignedLayouts.contains(op.getSource())) {