[mlir][linalg] Add constant padding helper to PadTensorOp

[springerm]

Summary

• Add a helper function that returns the constant padding value (if applicable).
• Remove existing getConstantYieldValueFromBlock function, which does almost the same.
• Adapted from D103243.

[ThomasRaoux]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

1143–1148

Can you just replace that with something like this:

if(padValue.getParentRegion() == getRegion().front())
  return {};

[springerm]

address comments

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

let extraClassDeclaration = [{
 
     // Return a PadTensorOp that pads `source to `type` size with `pad` value.
     // I.e., a block will be created and the `pad` value will be yielded
     // directly. If the type passed is nullptr, it is inferred.
     static linalg::PadTensorOp createPadScalarOp(
         Type type, Value source, Value pad, ArrayRef<OpFoldResult> low,
         ArrayRef<OpFoldResult> high, Location loc, OpBuilder & builder);
 
+    // Return the pad value if it is a constant. Return null value otherwise.
+    Value getConstantPaddingValue();
+
     // Return a vector of all the static or dynamic values (low/high padding) of
     // the op.
     inline SmallVector<OpFoldResult> getMixedPadImpl(ArrayAttr staticAttrs,
                                                      ValueRange values) {
       SmallVector<OpFoldResult> res;
       unsigned numDynamic = 0;
       unsigned count = staticAttrs.size();
       for (unsigned idx = 0; idx < count; ++idx) {

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 };
 
 } // namespace
 
 void PadTensorOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
   results.add<FoldStaticZeroPadding>(context);
 }
 
+/// Return the padding value of the PadTensorOp if it constant. In this context,
+/// "constant" means an actual constant or "defined outside of the block".
+///
+/// Values are considered constant in three cases:
+///  - A ConstantLike value.
+///  - A basic block argument from a different block.
+///  - A value defined outside of the block.
+///
+/// If the padding value is not constant, an empty Value is returned.
+Value PadTensorOp::getConstantPaddingValue() {
+  auto yieldOp = dyn_cast<YieldOp>(getRegion().front().getTerminator());
+  if (!yieldOp || yieldOp.values().size() != 1)
+    return {};
+  Value padValue = yieldOp.values().front();
+  // Check if yield value is a constant.
+  if (matchPattern(padValue, m_Constant()))
+    return padValue;
+  // Check if yield value is defined inside the PadTensorOp block.
+  if (padValue.getParentBlock() == &getRegion().front())
+    return {};
+  // Else: Yield value defined outside of the PadTensorOp block.
+  return padValue;
+}
+
 //===----------------------------------------------------------------------===//
 // ReshapeOp
 //===----------------------------------------------------------------------===//
 
 Optional<SmallVector<ReassociationIndices>>
 mlir::linalg::getReassociationIndicesForReshape(ShapedType sourceType,
                                                 ShapedType targetType) {
   // Make the sourceType greater rank than the targetType. If they are same

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

mlir::linalg::vectorizeLinalgOp(OpBuilder &b, Operation *op,
   return vectorizeAsLinalgGeneric(b, linalgOp, newResults,
                                   /*broadcastToMaximalCommonShape=*/true);
 }
 
 //----------------------------------------------------------------------------//
 // Misc. vectorization patterns.
 //----------------------------------------------------------------------------//
 
-/// Given a block, return the Value that the block yields if that Value is
-/// constant. In this context, "constant" means "defined outside of the block".
-/// Should not be called on blocks that yield more than one value.
-///
-/// Values are considered constant in two cases:
-///  - A basic block argument from a different block.
-///  - A value defined outside of the block.
-///
-/// If the yielded value is not constant, an empty Value is returned.
-static Value getConstantYieldValueFromBlock(Block &block) {
-  auto yieldOp = cast<YieldOp>(block.getTerminator());
-  assert(yieldOp.getNumOperands() == 1 && "expected single operand yield");
-  Value result = yieldOp.values().front();
-  Operation *definingOp = result.getDefiningOp();
-
-  // Check if yield value is defined inside the block.
-  if (definingOp && definingOp->getBlock() == &block)
-    return Value();
-  // Check if the yield value is a BB arg of the block.
-  if (!definingOp && result.cast<BlockArgument>().getOwner() == &block)
-    return Value();
-
-  return result;
-}
-
 /// Rewrite a PadTensorOp into a sequence of InitTensorOp, TransferReadOp and
 /// TransferWriteOp. For now, this only applies when all low and high paddings
 /// are determined to be zero.
 struct GenericPadTensorOpVectorizationPattern
     : public OpRewritePattern<PadTensorOp> {
   using OpRewritePattern<PadTensorOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(PadTensorOp padOp,
                                 PatternRewriter &rewriter) const override {
     /// Given an OpFoldResult, return true if its value is guaranteed to be a
     /// zero integer.
     auto isZeroInt = [&](OpFoldResult ofr) {
       return isEqualConstantIntOrValue(ofr, rewriter.getIndexAttr(0)); };
     // Low padding must be static 0.
     if (!llvm::all_of(padOp.getMixedLowPad(), isZeroInt)) return failure();
     // High padding must be static 0.
     if (!llvm::all_of(padOp.getMixedHighPad(), isZeroInt)) return failure();
     // Pad value must be a constant.
-    auto padValue = getConstantYieldValueFromBlock(padOp.region().front());
+    auto padValue = padOp.getConstantPaddingValue();
     if (!padValue) return failure();
 
     // Bail on non-static shapes.
     auto resultShapedType = padOp.result().getType().cast<ShapedType>();
     if (!resultShapedType.hasStaticShape())
       return failure();
     VectorType vectorType = extractVectorTypeFromShapedValue(padOp.result());
     if (!vectorType)