[mlir][linalg] Prepare pad to static bounding box for scalar operands.

[gysit]

Summary

Adapt pad to static bounding box to support structured ops taking scalar operands.

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

 
 /// Try to compute a static bounding box for `operand`
 /// Return success if either:
 ///   1. The operand is already statically shaped, `result` is left unchanged.
 ///   2. The operand is (partially) dynamic, `result` is the result of a freshly
 ///      created PadTensorOp.
 /// Return failure if the operand cannot be padded to a static shape.
 static LogicalResult padOperandToSmallestStaticBoundingBox(
-    PatternRewriter &rewriter, linalg::LinalgOp opToPad, OpOperand &operand,
+    PatternRewriter &rewriter, linalg::LinalgOp opToPad, OpOperand *opOperand,
     const LinalgTilingOptions &options, Value &result) {
-  auto tensorType = operand.get().getType().cast<RankedTensorType>();
   // Already static shape, no need to pad.
-  if (tensorType.hasStaticShape())
+  if (llvm::none_of(opToPad.getShape(opOperand), ShapedType::isDynamic))
     return success();
-  auto subtensor = operand.get().getDefiningOp<SubTensorOp>();
+  auto subtensor = opOperand->get().getDefiningOp<SubTensorOp>();
   // Not a subtensor, cannot construct a static bounding box.
   if (!subtensor)
     return failure();
   SmallVector<int64_t> staticSizes;
-  staticSizes.reserve(tensorType.getRank());
+  staticSizes.reserve(opToPad.getRank(opOperand));
   auto shapedOp =
       cast<OffsetSizeAndStrideOpInterface>(subtensor.getOperation());
   for (auto size : shapedOp.getMixedSizes()) {
     auto indexAttr = size.is<Attribute>()
                          ? size.get<Attribute>().dyn_cast<IntegerAttr>()
                          : linalg::getSmallestBoundingIndex(size.get<Value>());
     // SmallestBoundingIndex must exist for all sizes.
     // For now return an error if we can't find it.
     if (!indexAttr)
       return rewriter.notifyMatchFailure(
           opToPad, "No constant bounding box can be found for padding");
     staticSizes.push_back(indexAttr.getInt());
   }
-  Value pad = options.paddingValueComputationFunction(rewriter, operand);
-  auto staticTensorType =
-      RankedTensorType::get(staticSizes, tensorType.getElementType());
+  Value pad = options.paddingValueComputationFunction(rewriter, *opOperand);
+  auto staticTensorType = RankedTensorType::get(
+      staticSizes, getElementTypeOrSelf(opOperand->get().getType()));
   result = linalg::PadTensorOp::createPadHighOp(
-      staticTensorType, operand.get(), pad, opToPad->getLoc(), rewriter);
+      staticTensorType, opOperand->get(), pad, opToPad->getLoc(), rewriter);
   return success();
 }
 
 // Try to create a static bounding box around each operand of `res.op`.
 // If successful, `res.op` is rewritten in static form with padded operands.
 // `res.op` is updated to the cloned static form of the op on success.
 static LogicalResult rewriteAsPaddedOp(PatternRewriter &rewriter,
                                        TiledLinalgOp &res,
static LogicalResult rewriteAsPaddedOp(PatternRewriter &rewriter,
   // Make a copy of the shaped operands and update it.
   SmallVector<Value> newOperands;
   newOperands.reserve(opToPad.getNumInputsAndOutputs());
   for (OpOperand *opOperand : opToPad.getInputAndOutputOperands()) {
     Value paddedOperand;
     // If padding was requested but the shape cannot be bounded statically then
     // the pattern fails to apply.
     if (failed(padOperandToSmallestStaticBoundingBox(
-            rewriter, opToPad, *opOperand, options, paddedOperand))) {
+            rewriter, opToPad, opOperand, options, paddedOperand)))
       return failure();
-    }
     newOperands.push_back(paddedOperand ? paddedOperand : opOperand->get());
   }
 
   // Clone `opToPad` to operate on the statically padded shapes.
   auto resultTensorTypes =
       ValueRange(newOperands).take_back(opToPad.getNumOutputs()).getTypes();
   ValueRange otherOperands = opToPad.getAssumedNonShapedOperands();
   newOperands.append(otherOperands.begin(), otherOperands.end());