[Mlir-commits] [mlir] [mlir] Add missing pad reshape propagation patterns (PR #168888)

[Ian Wood]

================
@@ -1038,6 +1038,59 @@ class FoldWithProducerReshapeOpByExpansion
   ControlFusionFn controlFoldingReshapes;
 };
 
+/// Carries information about a padded dimension.
+struct PadDimInfo {
+  // The resulting shape after padding each dimension.
+  SmallVector<int64_t> paddedShape;
+
+  // Low and high padding amounts for each dimension.
+  SmallVector<OpFoldResult> lowPad;
+  SmallVector<OpFoldResult> highPad;
+};
+
+/// Computes the expanded padding information for the given pad operation based
+/// on the provided expanded shape and reassociation indices. Returns a list of
+/// PadDimInfo containing the low and high padding amounts and the padded
+/// size for each dimension, or failure if the expansion is not possible.
+static FailureOr<PadDimInfo>
+computeExpandedPadding(tensor::PadOp padOp, ArrayRef<int64_t> expandedShape,
+                       ArrayRef<ReassociationIndices> reassociations,
+                       PatternRewriter &rewriter) {
+  // If the padding value depends on the index values of the pad operation,
+  // then it may not be valid to expand the dimensions, since it will change
+  // the index values on which the padding value depends.
+  if (!padOp.getConstantPaddingValue())
+    return failure();
----------------
IanWood1 wrote:

It might be worth clarifying that this is unimplemented, so that it doesn't get misunderstood as a hard constraint. Here's what is done for `linalg.index` when expanding:

https://github.com/llvm/llvm-project/blob/8f1bb92bbfa45d49103953dad0d0a5dcfd388959/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp#L745-L768

https://github.com/llvm/llvm-project/pull/168888