[mlir][NFC] reifyResultShapes: Add extra error checking

[springerm]

Summary

This change adds a new helper function mlir::reifyResultShapes that calls the corresponding interface method and also checks the result produced by the implementation when running in debug mode. Bugs due to incorrect interface implementations can be difficult to debug.

This helper function also reduces the amount of code needed at call sites: the cast to ReifyRankedShapedTypeOpInterface is done in the helper function.

Depends On: D145702

[springerm]

update

[nicolasvasilache]

mlir/lib/Interfaces/InferTypeOpInterface.cpp

46

is this ever used in verification ?

assertion failures in verifiers lead to very hard to debug bugs ..

[springerm]

mlir/lib/Interfaces/InferTypeOpInterface.cpp

46

It's not used during verification, I checked all call sites of reifyResultShapes.

mlir/include/mlir/Interfaces/InferTypeOpInterface.h

 #include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/SmallVector.h"
 
 namespace mlir {
 
 class ShapedTypeComponents;
 using ReifiedRankedShapedTypeDims = SmallVector<SmallVector<OpFoldResult>>;
 
+/// Reify the shape of the result of an operation (typically in terms of the
+/// shape of its operands).
+LogicalResult
+reifyResultShapes(OpBuilder &b, Operation *op,
+                  ReifiedRankedShapedTypeDims &reifiedReturnShapes);
+
 /// Adaptor class to abstract the differences between whether value is from
 /// a ShapedType or ShapedTypeComponents or DenseIntElementsAttribute.
 class ShapeAdaptor {
 public:
   ShapeAdaptor(Type t) {
     if (auto st = t.dyn_cast<ShapedType>())
       val = st;
   }

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

FailureOr<Value> bufferization::allocateTensorForShapedValue(
   RankedTensorType tensorType = tensor.getType().cast<RankedTensorType>();
   SmallVector<Value> dynamicSizes;
   if (!copy) {
     // Compute the dynamic part of the shape.
     // First try to query the shape via ReifyRankedShapedTypeOpInterface.
     bool reifiedShapes = false;
     if (shapedValue.getType().isa<RankedTensorType>() &&
         shapedValue.isa<OpResult>()) {
-      if (auto rankedOp = dyn_cast_or_null<ReifyRankedShapedTypeOpInterface>(
-              shapedValue.getDefiningOp())) {
-        ReifiedRankedShapedTypeDims resultDims;
-        if (succeeded(rankedOp.reifyResultShapes(b, resultDims))) {
-          reifiedShapes = true;
-          auto &shape =
-              resultDims[shapedValue.cast<OpResult>().getResultNumber()];
-          for (const auto &dim : enumerate(tensorType.getShape()))
-            if (ShapedType::isDynamic(dim.value()))
-              dynamicSizes.push_back(shape[dim.index()].get<Value>());
-        }
-      }
-    }
+      if (succeeded(
+              reifyResultShapes(b, shapedValue.getDefiningOp(), resultDims))) {
 
     // If the shape could not be reified, create DimOps.
     if (!reifiedShapes)
       populateDynamicDimSizes(b, loc, tensor, dynamicSizes);
   }
 
   // Create AllocTensorOp.
   auto allocTensorOp = b.create<AllocTensorOp>(loc, tensorType, dynamicSizes,

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

LogicalResult matchAndRewrite(tensor::PadOp padOp,
 
     // We can only fold if the padding value is the same as the original
     // filling value.
     Value padValue = padOp.getConstantPaddingValue();
     if (!padValue || fillOp.value() != padValue)
       return failure();
 
     ReifiedRankedShapedTypeDims reifiedShape;
-    ReifyRankedShapedTypeOpInterface interface =
-        cast<ReifyRankedShapedTypeOpInterface>(padOp.getOperation());
-    if (failed(interface.reifyResultShapes(rewriter, reifiedShape)))
+    if (failed(reifyResultShapes(rewriter, padOp, reifiedShape)))
       return rewriter.notifyMatchFailure(
           padOp, "failed to reify tensor.pad op result shape");
 
     auto emptyTensor = rewriter.create<tensor::EmptyOp>(
         padOp.getLoc(), reifiedShape.front(),
         padOp.getResultType().getElementType());
     Value replacement =
         rewriter

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

 
 static SmallVector<Value> reifyOrComputeDynamicSizes(OpBuilder &b,
                                                      Value value) {
   auto tensorType = value.getType().cast<RankedTensorType>();
   if (tensorType.hasStaticShape())
     return {};
 
   // Try to reify dynamic sizes.
-  if (auto reifiableOp =
-          value.getDefiningOp<ReifyRankedShapedTypeOpInterface>()) {
-    ReifiedRankedShapedTypeDims reifiedShape;
-    if (succeeded(reifiableOp.reifyResultShapes(b, reifiedShape))) {
-      SmallVector<Value> dynSizes;
-      for (int64_t i = 0; i < tensorType.getRank(); ++i) {
-        if (tensorType.isDynamicDim(i))
-          dynSizes.push_back(
-              reifiedShape[value.cast<OpResult>().getResultNumber()][i]
-                  .get<Value>());
-      }
-      return dynSizes;
-    }
-  }
+  if (value.isa<OpResult>() &&
+      succeeded(reifyResultShapes(b, value.getDefiningOp(), reifiedShape))) {
 
   // Create tensor.dim ops.
   SmallVector<Value> dynSizes;
   for (int64_t i = 0; i < tensorType.getRank(); ++i) {
     if (tensorType.isDynamicDim(i))
       dynSizes.push_back(
           b.create<DimOp>(value.getLoc(), value,
                           b.create<arith::ConstantIndexOp>(value.getLoc(), i)));
mlir::linalg::rewriteInDestinationPassingStyle(RewriterBase &rewriter,
   // Only ops with exactly one block are supported.
   if (!padOp.getBodyRegion().hasOneBlock())
     return failure();
 
   // Create tensor.empty.
   Location loc = padOp.getLoc();
   RankedTensorType resultType = padOp.getResultType();
   ReifiedRankedShapedTypeDims reifiedShape;
-  if (failed(cast<ReifyRankedShapedTypeOpInterface>(padOp.getOperation())
-                 .reifyResultShapes(rewriter, reifiedShape)))
+  if (failed(reifyResultShapes(rewriter, padOp, reifiedShape)))
     return rewriter.notifyMatchFailure(
         padOp, "failed to reify tensor.pad op result shape");
   SmallVector<Value> dynamicSizes;
   for (int64_t i = 0; i < resultType.getRank(); ++i)
     if (resultType.isDynamicDim(i))
       dynamicSizes.push_back(reifiedShape[0][i].get<Value>());
 
   // If the `padOp` has a nofold attribute and all paddings are known to be 0,

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

if (!linalgOp) {
           padOp, "expected source to be linalg.generic op");
     }
     // All iterator types need to be parallel.
     if (linalgOp.getNumLoops() != linalgOp.getNumParallelLoops()) {
       return rewriter.notifyMatchFailure(
           padOp, "only supported for ops with all parallel iterator types");
     }
     ReifiedRankedShapedTypeDims resultShape;
-    ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
-        dyn_cast<ReifyRankedShapedTypeOpInterface>(padOp.getOperation());
-    if (failed(reifyShapedTypeInterface.reifyResultShapes(rewriter,
-                                                          resultShape)) ||
+    if (failed(reifyResultShapes(rewriter, padOp, resultShape)) ||
         resultShape.size() != 1) {
       return rewriter.notifyMatchFailure(
           padOp, "failed to get shape of pad op result");
     }
 
     Location loc = padOp.getLoc();
 
     // Create the tensor of same size as output of the pad op.

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

if (failed(paddedOperand)) {
                         << opOperand.get() << " -> FAIL\n");
       return rewriter.notifyMatchFailure(opToPad,
                                          "operand cannot be bound statically");
     }
     newOperands.push_back(*paddedOperand);
   }
 
   ReifiedRankedShapedTypeDims reifiedResultShapes;
-  if (failed(cast<ReifyRankedShapedTypeOpInterface>(opToPad.getOperation())
-                 .reifyResultShapes(rewriter, reifiedResultShapes))) {
+  if (failed(reifyResultShapes(rewriter, opToPad, reifiedResultShapes))) {
     LLVM_DEBUG(DBGS() << "--failed to reify result shapes -> FAIL\n");
     return rewriter.notifyMatchFailure(opToPad,
                                        "failed to reify result shapes");
   }
   assert(reifiedResultShapes.size() == opToPad->getNumResults() &&
          "expected same number of results");
 
   // Clone `opToPad` to operate on the statically padded shapes.

mlir/lib/Dialect/MemRef/Transforms/ResolveShapedTypeResultDims.cpp

struct DimOfReifyRankedShapedTypeOpInterface : public OpRewritePattern<OpTy> {
 
   void initialize() { OpRewritePattern<OpTy>::setHasBoundedRewriteRecursion(); }
 
   LogicalResult matchAndRewrite(OpTy dimOp,
                                 PatternRewriter &rewriter) const override {
     OpResult dimValue = dimOp.getSource().template dyn_cast<OpResult>();
     if (!dimValue)
       return failure();
-    auto rankedShapeTypeOp =
-        dyn_cast<ReifyRankedShapedTypeOpInterface>(dimValue.getOwner());
-    if (!rankedShapeTypeOp)
-      return failure();
-
     std::optional<int64_t> dimIndex = dimOp.getConstantIndex();
     if (!dimIndex)
       return failure();
 
     ReifiedRankedShapedTypeDims reifiedResultShapes;
-    if (failed(
-            rankedShapeTypeOp.reifyResultShapes(rewriter, reifiedResultShapes)))
-      return failure();
-
-    if (reifiedResultShapes.size() != rankedShapeTypeOp->getNumResults())
+    if (failed(reifyResultShapes(rewriter, dimValue.getOwner(),
+                                 reifiedResultShapes)))
       return failure();
-
     unsigned resultNumber = dimValue.getResultNumber();
-    auto sourceType = dimValue.getType().dyn_cast<RankedTensorType>();
-    if (reifiedResultShapes[resultNumber].size() !=
-        static_cast<size_t>(sourceType.getRank()))
-      return failure();
-
-    rewriter.replaceOp(dimOp,
-                       getValueOrCreateConstantIndexOp(
-                           rewriter, dimOp.getLoc(),
-                           reifiedResultShapes[resultNumber][*dimIndex]));
+    Value replacement = getValueOrCreateConstantIndexOp(
+        rewriter, dimOp.getLoc(), reifiedResultShapes[resultNumber][*dimIndex]);
+    rewriter.replaceOp(dimOp, replacement);
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Pass registration
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

FailureOr<Value> tensor::getOrCreateDestination(OpBuilder &b, Location loc,
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPoint(opResult.getDefiningOp());
 
   // Compute sizes.
   SmallVector<OpFoldResult> mixedSizes;
   if (!tensorType.hasStaticShape()) {
     // Dynamic shape: Query ReifyRankedShapedTypeOpInterface.
     ReifiedRankedShapedTypeDims reifiedShapes;
-    ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
-        dyn_cast<ReifyRankedShapedTypeOpInterface>(opResult.getDefiningOp());
-    if (!reifyShapedTypeInterface)
-      return failure();
-    if (failed(reifyShapedTypeInterface.reifyResultShapes(b, reifiedShapes)))
+    if (failed(reifyResultShapes(b, opResult.getDefiningOp(), reifiedShapes)))
       return failure();
     mixedSizes = reifiedShapes[opResult.getResultNumber()];
   } else {
     // Static shape: Take static sizes directly.
     for (int64_t sz : tensorType.getShape())
       mixedSizes.push_back(b.getIndexAttr(sz));
   }
 

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp

SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
     auto padOp = cast<PadOp>(op);
     SmallVector<utils::IteratorType> iteratorTypes(
         padOp.getResultType().getRank(), utils::IteratorType::parallel);
     return iteratorTypes;
   }
 
   SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
     ReifiedRankedShapedTypeDims reifiedShapes;
-    ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
-        dyn_cast<ReifyRankedShapedTypeOpInterface>(op);
-    (void)reifyShapedTypeInterface.reifyResultShapes(b, reifiedShapes);
-
+    (void)reifyResultShapes(b, op, reifiedShapes);
     Location loc = op->getLoc();
     Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
     Value one = b.create<arith::ConstantIndexOp>(loc, 1);
     // Initialize all the ranges to {zero, one, one}. All the `ub`s are
     // overwritten.
     SmallVector<Range> loopRanges(reifiedShapes[0].size(), {zero, one, one});
     for (const auto &ub : enumerate(reifiedShapes[0]))
       loopRanges[ub.index()].size = ub.value();
static_assert(llvm::is_one_of<OpTy, PackOp, UnPackOp>::value,
                 "applies to only pack or unpack operations");
   OpBuilder::InsertionGuard g(builder);
   Location loc = op.getLoc();
   int64_t rank = (std::is_same<OpTy, PackOp>::value) ? op.getSourceRank()
                                                      : op.getDestRank();
   Value zero = builder.create<arith::ConstantIndexOp>(loc, 0);
   Value one = builder.create<arith::ConstantIndexOp>(loc, 1);
   ReifiedRankedShapedTypeDims resultShape;
-  (void)op.reifyResultShapes(builder, resultShape);
+  (void)reifyResultShapes(builder, op, resultShape);
   SmallVector<Range> loopBounds(rank);
   for (auto dim : llvm::seq<int64_t>(0, rank)) {
     loopBounds[dim].offset = zero;
     loopBounds[dim].stride = one;
     loopBounds[dim].size = resultShape[0][dim];
   }
   return loopBounds;
 }
getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,
     auto packOp = cast<PackOp>(op);
     int64_t inputRank = packOp.getSourceRank();
     int64_t outputRank = packOp.getDestRank();
     auto zeroAttr = b.getI64IntegerAttr(0);
     resultOffsets.assign(offsets.begin(), offsets.end());
     resultOffsets.append(outputRank - inputRank, zeroAttr);
 
     ReifiedRankedShapedTypeDims outputShape;
-    (void)packOp.reifyResultShapes(b, outputShape);
+    (void)reifyResultShapes(b, packOp, outputShape);
     resultSizes.assign(sizes.begin(), sizes.end());
     for (auto dataTileDim : llvm::seq<unsigned>(inputRank, outputRank))
       resultSizes.push_back(outputShape[0][dataTileDim]);
 
     return success();
   }
 };
 

mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp

struct FoldEmptyTensorWithReshapeOp : public OpRewritePattern<ReshapeOp> {
   using OpRewritePattern<ReshapeOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ReshapeOp reshapeOp,
                                 PatternRewriter &rewriter) const override {
     if (!reshapeOp.getSrc().template getDefiningOp<EmptyOp>())
       return failure();
     Location loc = reshapeOp.getLoc();
     ReifiedRankedShapedTypeDims resultShapes;
-    ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
-        cast<ReifyRankedShapedTypeOpInterface>(reshapeOp.getOperation());
-    if (failed(reifyShapedTypeInterface.reifyResultShapes(rewriter,
-                                                          resultShapes)) ||
+    if (failed(reifyResultShapes(rewriter, reshapeOp, resultShapes)) ||
         !llvm::hasSingleElement(resultShapes))
       return failure();
     // TODO: Do not drop tensor type encoding.
     Value emptyTensor = rewriter.create<EmptyOp>(
         loc, resultShapes[0], reshapeOp.getResultType().getElementType());
     if (emptyTensor.getType() != reshapeOp.getResultType()) {
       rewriter.replaceOpWithNewOp<tensor::CastOp>(
           reshapeOp, reshapeOp.getResultType(), emptyTensor);

mlir/lib/Dialect/Tensor/Transforms/ExtractSliceFromReshapeUtils.cpp

tensor::ExtractSliceFromCollapseHelper::create(OpBuilder &b,
   // a rank-reducing extract slice.
   if (succeeded(mlir::getSimplifyCollapseShapeWithRankReducingSliceInfo(
           op.getSrcType(), op.getReassociationIndices())))
     return failure();
 
   // Materialize the output shape of the collapse_shape operation. This will
   // create IR describing the output shape in terms of the input shape.
   ReifiedRankedShapedTypeDims reifiedShapes;
-  ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
-      dyn_cast<ReifyRankedShapedTypeOpInterface>(op.getOperation());
-  if (failed(reifyShapedTypeInterface.reifyResultShapes(b, reifiedShapes)))
+  if (failed(reifyResultShapes(b, op, reifiedShapes)))
     return failure();
   SmallVector<OpFoldResult> &collapseShapeOutputShape = reifiedShapes[0];
   SmallVector<ReassociationIndices> reassociationIndices =
       op.getReassociationIndices();
 
   // Determine which of the CollapseShapeOp's result dimensions are sliced
   // and/or linearized.
   llvm::SmallBitVector linearizedDimensions =

mlir/lib/Interfaces/InferTypeOpInterface.cpp

 #include "llvm/Support/FormatVariadic.h"
 
 using namespace mlir;
 
 namespace mlir {
 #include "mlir/Interfaces/InferTypeOpInterface.cpp.inc"
 } // namespace mlir
 
+LogicalResult
+mlir::reifyResultShapes(OpBuilder &b, Operation *op,
+                        ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
+  auto reifiableOp = dyn_cast<ReifyRankedShapedTypeOpInterface>(op);
+  if (!reifiableOp)
+    return failure();
+  LogicalResult status = reifiableOp.reifyResultShapes(b, reifiedReturnShapes);
+#ifndef NDEBUG
+  if (failed(status))
+    return failure();
+  // Assert that ReifyRankedShapedTypeOpInterface::reifyResultShapes produced
+  // a correct result.
+  int64_t resultIdx = 0;
+  for (OpResult result : op->getResults()) {
+    auto shapedType = result.getType().dyn_cast<ShapedType>();
+    if (!shapedType)
+      continue;
+    if (!shapedType.hasRank()) {
+      // Nothing to check for unranked shaped values.
+      ++resultIdx;
+      continue;
+    }
+    // Assert one OpFoldResult per dimension.
+    assert(shapedType.getRank() ==
+               static_cast<int64_t>(reifiedReturnShapes[resultIdx].size()) &&
+           "incorrect implementation of ReifyRankedShapedTypeOpInterface");
+    for (int64_t dim = 0; dim < shapedType.getRank(); ++dim) {
+      // reifyResultShapes must return:
+      // * Attribute for static dimensions
+      // * Value for dynamic dimensions
+      assert(shapedType.isDynamicDim(dim) ==
+                 reifiedReturnShapes[resultIdx][dim].is<Value>() &&
+             "incorrect implementation of ReifyRankedShapedTypeOpInterface");
+    }
+    ++resultIdx;
+  }
+  // Assert that every shaped value result was reified.
+  assert(resultIdx == static_cast<int64_t>(reifiedReturnShapes.size()) &&
+         "incorrect implementation of ReifyRankedShapedTypeOpInterface");
+#endif // NDEBUG
+  return status;
+}
+
 bool ShapeAdaptor::hasRank() const {
   if (val.isNull())
     return false;
   if (auto t = val.dyn_cast<Type>())
     return t.cast<ShapedType>().hasRank();
   if (val.is<Attribute>())
     return true;
   return val.get<ShapedTypeComponents *>()->hasRank();

mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

if (succeeded(simplifiedCollapseShapeResult)) {
       // return.
       if (!newCollapseOp)
         return success();
       collapseOp = newCollapseOp;
     }
 
     // Materialize the output shape values of the slice operation.
     ReifiedRankedShapedTypeDims reifiedShapes;
-    if (failed(op.reifyResultShapes(rewriter, reifiedShapes)))
+    if (failed(reifyResultShapes(rewriter, op, reifiedShapes)))
       return rewriter.notifyMatchFailure(op, "failed to reify result shapes");
 
     // Create the destination tensor using the above values.
     Type elementType = op.getSourceType().getElementType();
     SmallVector<OpFoldResult> outputShape = reifiedShapes[0];
     Value dest = rewriter.create<tensor::EmptyOp>(op->getLoc(), outputShape,
                                                   elementType);
 

mlir/test/lib/Dialect/Test/TestDialect.cpp

LogicalResult OpWithResultShapeInterfaceOp::reifyReturnTypeShapes(
   return success();
 }
 
 LogicalResult OpWithResultShapePerDimInterfaceOp::reifyResultShapes(
     OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {
   Location loc = getLoc();
   shapes.reserve(getNumOperands());
   for (Value operand : llvm::reverse(getOperands())) {
+    auto tensorType = operand.getType().cast<RankedTensorType>();
     auto currShape = llvm::to_vector<4>(llvm::map_range(
-        llvm::seq<int64_t>(
-            0, operand.getType().cast<RankedTensorType>().getRank()),
+        llvm::seq<int64_t>(0, tensorType.getRank()),
         [&](int64_t dim) -> OpFoldResult {
-          return builder.createOrFold<tensor::DimOp>(loc, operand, dim);
+          return tensorType.isDynamicDim(dim)
+                     ? static_cast<OpFoldResult>(
+                           builder.createOrFold<tensor::DimOp>(loc, operand,
+                                                               dim))
+                     : static_cast<OpFoldResult>(
+                           builder.getIndexAttr(tensorType.getDimSize(dim)));
         }));
     shapes.emplace_back(std::move(currShape));
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Test SideEffect interfaces