[mlir][linalg][NFC] Cleanup padding transform

[springerm]

Summary

• Use LinalgPaddingOptions instead of passing many parameters.
• Split function into two parts.

Depends On: D153512

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 /// dimension should be padded to the specified multiple. If the derived padding
 /// sizes should not be rounded up to any multiple, use "1". Use `paddingValues`
 /// and `packPaddings` to set padding value and nofold attribute of the created
 /// tensor::PadOps, respectively. Update `paddedOp` to the cloned operation with
 /// statically shaped `paddingDimensions` and return the extracted dynamically
 /// shaped results. If padding fails, return failure.
 FailureOr<SmallVector<Value>>
 rewriteAsPaddedOp(RewriterBase &rewriter, LinalgOp opToPad,
-                  ArrayRef<int64_t> paddingDimensions,
-                  ArrayRef<int64_t> padToMultipleOf,
-                  ArrayRef<Attribute> paddingValues,
-                  ArrayRef<bool> packPaddings, LinalgOp &paddedOp);
+                  const LinalgPaddingOptions &options, LinalgOp &paddedOp);
 
 namespace detail {
 
 /// Helper struct to hold the results of building a packing loop nest.
 struct PackingResult {
   SmallVector<OpFoldResult> offsets, sizes, strides;
   SmallVector<Value> clonedLoopIvs, leadingPackedTensorIndexings;
   GenericOp maybeTransposeOp;
hoistPaddingOnTensors(tensor::PadOp opToHoist, int64_t numLoops,
                       ArrayRef<int64_t> transposeVector,
                       tensor::PadOp &hoistedOp,
                       SmallVectorImpl<GenericOp> &transposeOps);
 
 /// Apply padding and hoisting to `linalgOp` according to the configuration
 /// specified in `options`.
 FailureOr<LinalgOp> padAndHoistLinalgOp(RewriterBase &rewriter,
                                         LinalgOp linalgOp,
-                                        LinalgPaddingOptions options);
+                                        const LinalgPaddingOptions &options);
 
 /// Split the given `op` into two parts along the given iteration space
 /// `dimension` at the specified `splitPoint`, and return the two parts.
 /// If the second part is statically known to be empty, do not create it
 /// and return nullptr instead. Error state is signalled by returning
 /// a pair of nullptrs.
 ///
 /// For example, the following op:

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

transform::PadOp::applyToOne(transform::TransformRewriter &rewriter,
   }
 
   // Extract the transpose vectors.
   SmallVector<SmallVector<int64_t>> transposePaddings;
   for (Attribute transposeVector : cast<ArrayAttr>(getTransposePaddings()))
     transposePaddings.push_back(
         extractFromI64ArrayAttr(cast<ArrayAttr>(transposeVector)));
 
+  // Set up options and pad.
   LinalgOp paddedOp;
-  SmallVector<int64_t> paddingDimensions =
-      extractFromI64ArrayAttr(getPaddingDimensions());
-  SmallVector<int64_t> padToMultipleOf(paddingDimensions.size(), 1);
+  LinalgPaddingOptions options;
+  options.paddingDimensions = extractFromI64ArrayAttr(getPaddingDimensions());
+  SmallVector<int64_t> padToMultipleOf(options.paddingDimensions.size(), 1);
   if (getPadToMultipleOf().has_value())
     padToMultipleOf = extractFromI64ArrayAttr(*getPadToMultipleOf());
+  options.padToMultipleOf = padToMultipleOf;
+  options.paddingValues = paddingValues;
+  options.packPaddings = packPaddings;
   FailureOr<SmallVector<Value>> result =
-      rewriteAsPaddedOp(rewriter, target, paddingDimensions, padToMultipleOf,
-                        paddingValues, packPaddings, paddedOp);
+      rewriteAsPaddedOp(rewriter, target, options, paddedOp);
   if (succeeded(result)) {
     // We need to perform our own replacement here because this API is still
     // used in patterns that "pad and hoist", for which the replacement values
     // need to be different.
     // TODO: clean this up and stop "pad and hoist" behavior more globally now
     // that we have more composable abstractions.
     rewriter.replaceOp(target, *result);
     results.push_back(paddedOp);

mlir/lib/Dialect/Linalg/Transforms/Padding.cpp

 #define DEBUG_TYPE "linalg-padding"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")
 #define DBGSNL() (llvm::dbgs() << "\n")
 
-/// Pad the `opOperand` in the `paddingDimensions` using the padding value and
-/// the nofold flag found in `paddingValues` and `packPaddings`, respectively.
-///
-/// Exit early and return the `opOperand` value if it already has the requested
-/// shape. I.e.:
-/// - static shape
-/// - nofold is not set
-/// - dim sizes are multiples of `padToMultipleOf`
-///
-/// Otherwise, try to pad the shape dimensions that match the iterator
-/// dimensions `paddingDimensions` and return the tensor::PadOp result if
-/// padding succeeds or failure otherwise.
-static FailureOr<Value> padOperandToSmallestStaticBoundingBox(
-    RewriterBase &rewriter, linalg::LinalgOp opToPad, OpOperand *opOperand,
-    ArrayRef<int64_t> paddingDimensions, ArrayRef<int64_t> padToMultipleOf,
-    ArrayRef<Attribute> paddingValues, ArrayRef<bool> packPaddings) {
-  assert(padToMultipleOf.size() == paddingDimensions.size() &&
-         "invalid number of elements in padToMultipleOf");
-
+/// Compute the padded shape of the given operand. The operand is padded to a
+/// static bounding box according to the specified options.
+static LogicalResult computePaddedShape(linalg::LinalgOp opToPad,
+                                        OpOperand *opOperand,
+                                        const LinalgPaddingOptions &options,
+                                        SmallVector<int64_t> &paddedShape,
+                                        bool &alreadyHasRequestedShape) {
   AffineMap indexingMap = opToPad.getMatchingIndexingMap(opOperand);
   ArrayRef<int64_t> shape = opToPad.getShape(opOperand);
 
-  // Collect the shape dimensions that are a function of `paddingDimensions`,
+  // Collect the shape dimensions that are a function of "paddingDimensions",
   // along with the multiple that they should be padded to ("1" if none).
-  bool alreadyHasRequestedShape = true;
+  alreadyHasRequestedShape = true;
   DenseMap<int64_t, int64_t> shapeDimToMultiple;
-  for (const auto &dimEn : enumerate(paddingDimensions)) {
+  for (const auto &dimEn : enumerate(options.paddingDimensions)) {
     for (const auto &en : enumerate(indexingMap.getResults())) {
       if (en.value().isFunctionOfDim(dimEn.value())) {
         int64_t dimSize = shape[en.index()];
-        shapeDimToMultiple[en.index()] = padToMultipleOf[dimEn.index()];
+        if (options.padToMultipleOf.has_value()) {
+          shapeDimToMultiple[en.index()] =
+              (*options.padToMultipleOf)[dimEn.index()];
+        } else {
+          shapeDimToMultiple[en.index()] = 1;
+        }
         if (ShapedType::isDynamic(dimSize)) {
           alreadyHasRequestedShape = false;
         } else if (dimSize % shapeDimToMultiple[en.index()] != 0) {
           alreadyHasRequestedShape = false;
         }
       }
     }
   }
 
-  // Return the unpadded operand if padding to a static shape is not needed and
-  // if the nofold flag is not set.
-  bool nofold = opOperand->getOperandNumber() < packPaddings.size()
-                    ? packPaddings[opOperand->getOperandNumber()]
-                    : false;
-  if (!nofold && alreadyHasRequestedShape)
-    return opOperand->get();
-
-  // Fail if `paddingValues` specifies no padding value.
-  if (opOperand->getOperandNumber() >= paddingValues.size()) {
-    return rewriter.notifyMatchFailure(opToPad, "--no padding value specified");
-  }
-  Attribute paddingAttr = paddingValues[opOperand->getOperandNumber()];
-  Value paddingValue = rewriter.create<arith::ConstantOp>(
-      opToPad.getLoc(), cast<TypedAttr>(paddingAttr));
-
   // Helper function to round a number up to a given multiple.
   auto ceil = [](int64_t val, int64_t multiple) {
     return ((val + multiple - 1) / multiple) * multiple;
   };
 
   // Upper bound the sizes to obtain a static bounding box.
-  SmallVector<int64_t> paddedShape(shape.begin(), shape.end());
+  paddedShape.assign(shape.begin(), shape.end());
   for (int64_t i = 0, e = shape.size(); i < e; ++i) {
     LLVM_DEBUG(DBGS() << "--compute padded size for dim " << i << "\n");
     // Skip dimensions that do not require padding.
     if (!shapeDimToMultiple.contains(i)) {
       LLVM_DEBUG(DBGS() << "----dim does not require padding, SKIP\n");
       continue;
     }
     // Otherwise, try to compute a constant upper bound for the size value.
     FailureOr<int64_t> upperBound =
         ValueBoundsConstraintSet::computeConstantBound(
             presburger::BoundType::UB, opOperand->get(),
             /*dim=*/i, /*stopCondition=*/nullptr, /*closedUB=*/true);
     if (failed(upperBound)) {
       LLVM_DEBUG(DBGS() << "----count not compute a bounding box for padding");
-      return rewriter.notifyMatchFailure(
-          opToPad, "count not compute a bounding box for padding");
+      return failure();
     }
     paddedShape[i] = ceil(*upperBound, shapeDimToMultiple[i]);
     LLVM_DEBUG(DBGS() << "----new dim size: " << paddedShape[i] << "\n");
   }
 
+  return success();
+}
+
+/// Pad the `opOperand` in the "paddingDimensions" using the padding value and
+/// the nofold flag found in "paddingValues" and "packPaddings", respectively.
+///
+/// Exit early and return the `opOperand` value if it already has the requested
+/// shape. I.e.:
+/// - static shape
+/// - nofold is not set
+/// - dim sizes are multiples of "padToMultipleOf"
+///
+/// Otherwise, try to pad the shape dimensions that match the iterator
+/// dimensions "paddingDimensions" and return the tensor::PadOp result if
+/// padding succeeds or failure otherwise.
+static FailureOr<Value> padOperandToSmallestStaticBoundingBox(
+    RewriterBase &rewriter, linalg::LinalgOp opToPad, OpOperand *opOperand,
+    const LinalgPaddingOptions &options) {
+  assert(!options.padToMultipleOf.has_value() ||
+         options.padToMultipleOf->size() == options.paddingDimensions.size() &&
+             "invalid number of elements in padToMultipleOf");
+
+  // Compute padded shape.
+  SmallVector<int64_t> paddedShape;
+  bool alreadyHasRequestedShape = false;
+  if (failed(computePaddedShape(opToPad, opOperand, options, paddedShape,
+                                alreadyHasRequestedShape)))
+    return rewriter.notifyMatchFailure(opToPad,
+                                       "--failed to compute padded shape");
+
+  // Return the unpadded operand if padding to a static shape is not needed and
+  // if the nofold flag is not set.
+  bool nofold = opOperand->getOperandNumber() < options.packPaddings.size()
+                    ? options.packPaddings[opOperand->getOperandNumber()]
+                    : false;
+  if (!nofold && alreadyHasRequestedShape)
+    return opOperand->get();
+
+  // Fail if `paddingValues` specifies no padding value.
+  if (opOperand->getOperandNumber() >= options.paddingValues.size()) {
+    return rewriter.notifyMatchFailure(opToPad, "--no padding value specified");
+  }
+  Attribute paddingAttr = options.paddingValues[opOperand->getOperandNumber()];
+  Value paddingValue = rewriter.create<arith::ConstantOp>(
+      opToPad.getLoc(), cast<TypedAttr>(paddingAttr));
+
   // Pad the operand to the bounding box defined by `paddedShape`.
   auto paddedTensorType = RankedTensorType::get(
       paddedShape, getElementTypeOrSelf(opOperand->get()));
   LLVM_DEBUG(DBGS() << "--SUCCESS, makeComposedPadHighOp with type: "
                     << paddedTensorType);
   return makeComposedPadHighOp(rewriter, opToPad->getLoc(), paddedTensorType,
                                opOperand->get(), paddingValue, nofold);
 }
 
 FailureOr<SmallVector<Value>>
 linalg::rewriteAsPaddedOp(RewriterBase &rewriter, LinalgOp opToPad,
-                          ArrayRef<int64_t> paddingDimensions,
-                          ArrayRef<int64_t> padToMultipleOf,
-                          ArrayRef<Attribute> paddingValues,
-                          ArrayRef<bool> packPaddings, LinalgOp &paddedOp) {
+                          const LinalgPaddingOptions &options,
+                          LinalgOp &paddedOp) {
   LLVM_DEBUG(DBGS() << "Start rewriteAsPaddedOp : " << opToPad << "\n");
   Location loc = opToPad->getLoc();
 
   // TODO: there are cases where we may still want to pad to larger sizes.
   if (!opToPad.hasTensorSemantics())
     return rewriter.notifyMatchFailure(opToPad,
                                        "expected operation on tensors");
 
   OpBuilder::InsertionGuard g(rewriter);
   // Set IP after op because we also take the dims of the original output.
   rewriter.setInsertionPointAfter(opToPad);
 
   // Make a copy of the shaped operands and update it.
   SmallVector<Value> newOperands;
   newOperands.reserve(opToPad->getNumOperands());
   for (OpOperand &opOperand : opToPad->getOpOperands()) {
     FailureOr<Value> paddedOperand = padOperandToSmallestStaticBoundingBox(
-        rewriter, opToPad, &opOperand, paddingDimensions, padToMultipleOf,
-        paddingValues, packPaddings);
+        rewriter, opToPad, &opOperand, options);
     // Exit if `paddingDimensions` cannot be bounded statically.
     if (failed(paddedOperand)) {
       LLVM_DEBUG(DBGS() << "--operand cannot be bound statically : "
                         << opOperand.get() << " -> FAIL\n");
       return rewriter.notifyMatchFailure(opToPad,
                                          "operand cannot be bound statically");
     }
     newOperands.push_back(*paddedOperand);
paddedSubtensorResults.push_back(rewriter.create<tensor::ExtractSliceOp>(
         loc, paddedResult, offsets, reifiedResultShapes[resultNumber],
         strides));
   }
   return paddedSubtensorResults;
 }
 
 FailureOr<LinalgOp>
 mlir::linalg::padAndHoistLinalgOp(RewriterBase &rewriter, LinalgOp linalgOp,
-                                  LinalgPaddingOptions options) {
+                                  const LinalgPaddingOptions &options) {
   if (!linalgOp.hasTensorSemantics())
     return rewriter.notifyMatchFailure(
         linalgOp, "only applies to Linalg ops with tensor semantics");
 
   // Pad the operation.
   LinalgOp paddedOp;
-  SmallVector<int64_t> padToMultipleOf(options.paddingDimensions.size(), 1);
-  if (options.padToMultipleOf.has_value())
-    padToMultipleOf.assign(options.padToMultipleOf->begin(),
-                           options.padToMultipleOf->end());
-  FailureOr<SmallVector<Value>> newResults = rewriteAsPaddedOp(
-      rewriter, linalgOp, options.paddingDimensions, padToMultipleOf,
-      options.paddingValues, options.packPaddings, paddedOp);
+  FailureOr<SmallVector<Value>> newResults =
+      rewriteAsPaddedOp(rewriter, linalgOp, options, paddedOp);
   if (failed(newResults))
     return rewriter.notifyMatchFailure(linalgOp,
                                        "failed to rewrite as a padded op");
 
   // Hoist the padding.
   for (const auto &en : enumerate(options.hoistPaddings)) {
     if (static_cast<int64_t>(en.index()) >= paddedOp->getNumOperands())
       break;