[mlir][Linalg] Use ReifyRankedShapedTypeOpInterface for pad transforms.

[hanchung]

Summary

The information is not tied to tensor.empty op and tensor.extract_slice
op. We can infer smallest static bounding box for pad transform if
they implement ReifyRankedShapedTypeOpInterface. The revision extends
the usability for downstream projects. No tests are added because the
existing tests cover the change, and most of MLIR
ReifyRankedShapedTypeOpInterface ops are covered in the tests, except
tensor.generate and bufferization.alloc_tensor ops.

[mravishankar]

I think this looks right to me. Not sure why the droppedDims were there to begin with. Whats done here should make it unnecessary.

[hanchung]

In D150227#4333760, @mravishankar wrote:

I think this looks right to me. Not sure why the droppedDims were there to begin with. Whats done here should make it unnecessary.

The droppedDims was needed because of tensor.extract_slice op. If it is a rank-reduced op, the number of elements in sizes does not match the rank of result type. The previous implementations uses getMixedSizes() and droppedDims to model it. Today the extract_slice op implements ReifyRankedShapedTypeOpInterface, so we can get the result sizes from the reifyResultShapes method. The logic is considered in ExtractSliceOp::reifyResultShapes, so we no longer need it.

https://github.com/llvm/llvm-project/blob/f85a8456f114d071fb5985ef2c4cd922ba8f29cf/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp#L1848-L1860

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

static FailureOr<Value> padOperandToSmallestStaticBoundingBox(
 
   // Follow the use-def chain if `currOpOperand` is defined by a LinalgOp.
   OpOperand *currOpOperand = opOperand;
   while (auto linalgOp = currOpOperand->get().getDefiningOp<LinalgOp>()) {
     OpResult result = currOpOperand->get().cast<OpResult>();
     currOpOperand = linalgOp.getDpsInitOperand(result.getResultNumber());
   }
 
-  // Fail if `currOpOperand` is not defined by an ExtractSliceOp or EmptyOp.
-  auto sliceOp = currOpOperand->get().getDefiningOp<tensor::ExtractSliceOp>();
-  auto emptyOp = currOpOperand->get().getDefiningOp<tensor::EmptyOp>();
-
-  llvm::SmallBitVector droppedDims;
   SmallVector<OpFoldResult> mixedSizes;
-  if (sliceOp) {
-    // Compute the dropped dimensions if `sliceOp` is rank-reducing.
-    droppedDims = sliceOp.getDroppedDims();
-    mixedSizes = sliceOp.getMixedSizes();
-  } else if (emptyOp) {
-    mixedSizes = emptyOp.getMixedSizes();
-    droppedDims.resize(mixedSizes.size());
+  if (auto reifiableOp =
+          llvm::dyn_cast_or_null<ReifyRankedShapedTypeOpInterface>(
+              currOpOperand->get().getDefiningOp())) {
+    ReifiedRankedShapedTypeDims reifiedReturnShapes;
+    LogicalResult status =
+        reifiableOp.reifyResultShapes(rewriter, reifiedReturnShapes);
+    mixedSizes = reifiedReturnShapes[0];
+    if (failed(status)) {
+      LLVM_DEBUG(DBGS() << "--failed to reify result shapes\n");
+      return rewriter.notifyMatchFailure(opToPad,
+                                         "failed to reify result shapes");
+    }
   } else if (hasStaticShape) {
     mixedSizes = getAsIndexOpFoldResult(rewriter.getContext(), shape);
-    droppedDims.resize(mixedSizes.size());
   } else {
     // TODO: may want to add support for going through loop iter args.
     // This is not strictly necessary as we can pad before hoisting but it would
     // make the system more resilient to minor transformation reordering.
-    LLVM_DEBUG(DBGS() << "--not defined by an extractSlice or emptyOp\n");
+    LLVM_DEBUG(DBGS() << "--not a ReifyRankedShapedTypeOpInterface op\n");
     return rewriter.notifyMatchFailure(
-        opToPad, "not defined by an extractSlice or emptyOp");
+        opToPad, "not a ReifyRankedShapedTypeOpInterface op");
   }
   LLVM_DEBUG(llvm::interleaveComma(mixedSizes, DBGS() << "--mixedSizes:  ");
              llvm::dbgs() << "\n");
 
   // Upper bound the sizes to obtain a static bounding box.
   SmallVector<int64_t> paddedShape(shape.begin(), shape.end());
   int64_t shapeIdx = 0;
   for (const auto &en : enumerate(mixedSizes)) {
     LLVM_DEBUG(DBGS() << "----mixedSizes:  " << en.value() << "\n");
-    // Skip dropped dimensions.
-    if (droppedDims.test(en.index())) {
-      LLVM_DEBUG(DBGS() << "------dim is dropped, SKIP\n");
-      continue;
-    }
     // Skip dimensions that do not require padding.
     if (!shapeDimsToPad.contains(shapeIdx)) {
       shapeIdx++;
       LLVM_DEBUG(DBGS() << "------dim does not require padding, SKIP\n");
       continue;
     }
     // If the size is an attribute add it directly to `paddedShape`.
     if (en.value().is<Attribute>()) {