[mlir][Interfaces] ReifyRankedShapedTypeOpInterface returns OpFoldResults

[springerm]

Summary

reifyResultShapes now returns OpFoldResults instead of Values. This is often more efficient because many transformations immediately attempt to extract a constant from the reified values.

[springerm]

update

[mravishankar]

Awesome! Thanks!

[jpienaar]

Should this be renamed now? E.g., reifyResultShapes now only sometimes reify. (I think folks are also reaching here for something closer to what onnx-mlir and others are doing)

[springerm]

In D145250#4176525, @jpienaar wrote:

Should this be renamed now? E.g., reifyResultShapes now only sometimes reify. (I think folks are also reaching here for something closer to what onnx-mlir and others are doing)

Maybe, depending on what "reification" means. If you interpret it as "building some handle that you can pass around in C++" (Value or Attribute), the name is still correct.

Btw, I am preparing additional changes to this interface that would allow users to reify specific dimensions and populate a FlatAffineValueConstraints (so that the "reification results" of multiple chained ops can be combined and analyzed).

[jpienaar]

In D145250#4177310, @springerm wrote:

In D145250#4176525, @jpienaar wrote:

Should this be renamed now? E.g., reifyResultShapes now only sometimes reify. (I think folks are also reaching here for something closer to what onnx-mlir and others are doing)

Maybe, depending on what "reification" means. If you interpret it as "building some handle that you can pass around in C++" (Value or Attribute), the name is still correct.

Btw, I am preparing additional changes to this interface that would allow users to reify specific dimensions and populate a FlatAffineValueConstraints (so that the "reification results" of multiple chained ops can be combined and analyzed).

That sounds similar to what onnx-mlir does, have you looked at what they are doing? ( https://github.com/onnx/onnx-mlir/blob/main/src/Interface/ShapeHelperOpInterface.hpp may be good entry).

[springerm]

In D145250#4184524, @jpienaar wrote:

That sounds similar to what onnx-mlir does, have you looked at what they are doing? ( https://github.com/onnx/onnx-mlir/blob/main/src/Interface/ShapeHelperOpInterface.hpp may be good entry).

Thx for the pointer, it definitely looks similar. My approach is a bit different because we want to support lower/upper bounds in addition to EQ bounds. I posted an RFC on Discourse as this may be interesting for various downstream projects.

mlir/include/mlir/Interfaces/InferTypeOpInterface.h

 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/SmallVector.h"
 
 namespace mlir {
 
 class ShapedTypeComponents;
-using ReifiedRankedShapedTypeDims = SmallVector<SmallVector<Value>>;
+using ReifiedRankedShapedTypeDims = SmallVector<SmallVector<OpFoldResult>>;
 
 /// 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/include/mlir/Interfaces/InferTypeOpInterface.td

let description = [{
     the result is a ranked shape type, i.e. `RankedTensorType` or
     `MemRefType`.
   }];
   let cppNamespace = "::mlir";
 
   let methods = [
     InterfaceMethod<
       /*desc=*/[{
-        Reify the shape of the result of an operation (typically in
-        terms of shape of its operands)
+        Reify the shape of the result of an operation (typically in terms of the
+        shape of its operands).
 
-        Insert operations using the given `OpBuilder` that computes
-        the result shape. The `reifiedReturnShapes` is expected to be
-        populated with as many vectors as the number of results of the
-        op. Each of these vectors is expected to be of size equal to
-        rank of the corresponding result. If the shape of a particular
-        result cannot be computed it must be empty.
+        `reifiedReturnShapes` is populated with one vector per op result. Each
+        of those vectors contains an OpFoldResult for each dimension of the
+        shaped type. In case a dimension in the type is static, the
+        corresponding entry is an IntegerAttr. Otherwise, it is a Value. The
+        given builder may be used to insert ops that compute result shapes.
+
+        If the shape of a particular result cannot be computed it must be empty.
       }],
       /*retTy=*/"::mlir::LogicalResult",
       /*methodName=*/"reifyResultShapes",
       /*args=*/(ins "::mlir::OpBuilder &":$builder,
         "::mlir::ReifiedRankedShapedTypeDims &":$reifiedReturnShapes)
     >
   ];
 }
 
 // Op has the same operand and result type.
 // TODO: Change from hard coded to utilizing type inference trait.
 def SameOperandsAndResultType : NativeOpTrait<"SameOperandsAndResultType">;
 
 #endif // MLIR_INFERTYPEOPINTERFACE

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

if (shapedValue.getType().isa<RankedTensorType>() &&
               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()]);
+              dynamicSizes.push_back(shape[dim.index()].get<Value>());
         }
       }
     }
 
     // If the shape could not be reified, create DimOps.
     if (!reifiedShapes)
       populateDynamicDimSizes(b, loc, tensor, dynamicSizes);
   }

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

 
 void AllocTensorOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *ctx) {
   results.add<FoldDimOfAllocTensorOp, ReplaceStaticShapeDims>(ctx);
 }
 
 LogicalResult AllocTensorOp::reifyResultShapes(
     OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
-  auto shapes = llvm::to_vector<4>(llvm::map_range(
-      llvm::seq<int64_t>(0, getType().getRank()), [&](int64_t dim) -> Value {
+  auto shapes = llvm::to_vector<4>(
+      llvm::map_range(llvm::seq<int64_t>(0, getType().getRank()),
+                      [&](int64_t dim) -> OpFoldResult {
                         if (isDynamicDim(dim))
                           return getDynamicSize(builder, dim);
-        return builder.create<arith::ConstantIndexOp>(getLoc(),
-                                                      getStaticSize(dim));
+                        return builder.getIndexAttr(getStaticSize(dim));
                       }));
   reifiedReturnShapes.emplace_back(std::move(shapes));
   return success();
 }
 
 ParseResult AllocTensorOp::parse(OpAsmParser &parser, OperationState &result) {
   SmallVector<OpAsmParser::UnresolvedOperand> dynamicSizesOperands;
   if (parser.parseLParen() || parser.parseOperandList(dynamicSizesOperands) ||
       parser.parseRParen())

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

LinalgOp::reifyResultShapes(OpBuilder &b,
   IRRewriter rewriter(b);
   SmallVector<OpFoldResult> allResultDimValues =
       makeComposedFoldedMultiResultAffineApply(
           rewriter, loc, resultShapesFromInputShapesMap,
           createFlatListOfOperandDims(b, loc));
   int64_t pos = 0;
   ArrayRef<AffineExpr> shapeExprs = resultShapesFromInputShapesMap.getResults();
   for (OpOperand *opOperand : getDpsInitOperands()) {
-    SmallVector<Value> shapes;
+    SmallVector<OpFoldResult> shapes;
     for (int64_t dim : llvm::seq<int64_t>(0, getRank(opOperand))) {
       if (checkDimExpr.visit(shapeExprs[pos]))
         shapes.push_back(createOrFoldDimOp(b, loc, opOperand->get(), dim));
       else
-        shapes.push_back(
-            getValueOrCreateConstantIndexOp(b, loc, allResultDimValues[pos]));
+        shapes.push_back(allResultDimValues[pos]);
       pos++;
     }
     reifiedReturnShapes.emplace_back(std::move(shapes));
   }
   return success();
 }
 
 /// Return the index in the indexingMaps vector that corresponds to this

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

LogicalResult matchAndRewrite(tensor::PadOp padOp,
 
     ReifiedRankedShapedTypeDims reifiedShape;
     ReifyRankedShapedTypeOpInterface interface =
         cast<ReifyRankedShapedTypeOpInterface>(padOp.getOperation());
     if (failed(interface.reifyResultShapes(rewriter, reifiedShape)))
       return rewriter.notifyMatchFailure(
           padOp, "failed to reify tensor.pad op result shape");
 
-    SmallVector<OpFoldResult> newShape =
-        getAsOpFoldResult(reifiedShape.front());
     auto emptyTensor = rewriter.create<tensor::EmptyOp>(
-        padOp.getLoc(), newShape, padOp.getResultType().getElementType());
+        padOp.getLoc(), reifiedShape.front(),
+        padOp.getResultType().getElementType());
     Value replacement =
         rewriter
             .create<FillOp>(fillOp.getLoc(), ValueRange{padValue},
                             ValueRange{emptyTensor})
             .getResult(0);
     if (replacement.getType() != padOp.getResultType()) {
       replacement = rewriter.create<tensor::CastOp>(
           fillOp.getLoc(), padOp.getResultType(), replacement);

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

 //===- ConvertToDestinationStyle.cpp - Convert non-DPS to DPS ops ---------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains patterns to convert non-DPS ops to DPS ops. New
 // tensor.empty ops are inserted as a destination. Such tensor.empty can be
 // eliminated with "empty tensor elimination", allowing them to bufferize
 // without an allocation (assuming there are no further conflicts).
 //
 //===----------------------------------------------------------------------===//
 //
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
static SmallVector<Value> reifyOrComputeDynamicSizes(OpBuilder &b,
   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]);
+              reifiedShape[value.cast<OpResult>().getResultNumber()][i]
+                  .get<Value>());
       }
       return dynSizes;
     }
   }
 
   // Create tensor.dim ops.
   SmallVector<Value> dynSizes;
   for (int64_t i = 0; i < tensorType.getRank(); ++i) {
mlir::linalg::rewriteInDestinationPassingStyle(RewriterBase &rewriter,
   ReifiedRankedShapedTypeDims reifiedShape;
   if (failed(cast<ReifyRankedShapedTypeOpInterface>(padOp.getOperation())
                  .reifyResultShapes(rewriter, 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]);
+      dynamicSizes.push_back(reifiedShape[0][i].get<Value>());
 
   // If the `padOp` has a nofold attribute and all paddings are known to be 0,
   // explicitly insert a `linalg.copy`.
   if (padOp.getNofoldAttr() &&
       llvm::all_of(padOp.getMixedLowPad(), isZeroIndex) &&
       llvm::all_of(padOp.getMixedHighPad(), isZeroIndex)) {
     using bufferization::AllocTensorOp;
     Value allocated =

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

if (failed(reifyShapedTypeInterface.reifyResultShapes(rewriter,
       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.
     RankedTensorType padResultType = padOp.getResultType();
-    auto resultSizes = getAsOpFoldResult(resultShape[0]);
+    auto resultSizes = resultShape[0];
     auto emptyTensor = rewriter.create<tensor::EmptyOp>(
         loc, resultSizes, padResultType.getElementType());
 
     // Fill the tensor with the pad value.
     // TODO: There is an option to fill only the boundaries. For now just
     // filling the whole tensor.
     auto fillTensor =
         rewriter.create<linalg::FillOp>(loc, padValue, emptyTensor.getResult());

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

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);
   }
 
-  SmallVector<SmallVector<Value>> reifiedResultShapes;
+  ReifiedRankedShapedTypeDims reifiedResultShapes;
   if (failed(cast<ReifyRankedShapedTypeOpInterface>(opToPad.getOperation())
                  .reifyResultShapes(rewriter, 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");
linalg::rewriteAsPaddedOp(RewriterBase &rewriter, LinalgOp opToPad,
   SmallVector<Value> paddedSubtensorResults;
   paddedSubtensorResults.reserve(opToPad->getNumResults());
   for (const auto &en : llvm::enumerate(paddedOp->getResults())) {
     Value paddedResult = en.value();
     int64_t resultNumber = en.index();
     int64_t rank = paddedResult.getType().cast<RankedTensorType>().getRank();
     SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
     SmallVector<OpFoldResult> sizes;
-    for (Value v : reifiedResultShapes[resultNumber])
-      sizes.push_back(getAsOpFoldResult(v));
     SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
     paddedSubtensorResults.push_back(rewriter.create<tensor::ExtractSliceOp>(
-        loc, paddedResult, offsets, sizes, strides));
+        loc, paddedResult, offsets, reifiedResultShapes[resultNumber],
+        strides));
   }
   return paddedSubtensorResults;
 }
 
 //===----------------------------------------------------------------------===//
 // peelLoop transformation.
 //===----------------------------------------------------------------------===//
 

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

 // shapes of their operands using the `InferShapedTypeOpInterface`.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 namespace mlir {
 namespace memref {
 #define GEN_PASS_DEF_RESOLVERANKEDSHAPETYPERESULTDIMS
auto rankedShapeTypeOp =
         dyn_cast<ReifyRankedShapedTypeOpInterface>(dimValue.getOwner());
     if (!rankedShapeTypeOp)
       return failure();
 
     std::optional<int64_t> dimIndex = dimOp.getConstantIndex();
     if (!dimIndex)
       return failure();
 
-    SmallVector<SmallVector<Value>> reifiedResultShapes;
+    ReifiedRankedShapedTypeDims reifiedResultShapes;
     if (failed(
             rankedShapeTypeOp.reifyResultShapes(rewriter, reifiedResultShapes)))
       return failure();
 
     if (reifiedResultShapes.size() != rankedShapeTypeOp->getNumResults())
       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, reifiedResultShapes[resultNumber][*dimIndex]);
+    rewriter.replaceOp(dimOp,
+                       getValueOrCreateConstantIndexOp(
+                           rewriter, dimOp.getLoc(),
+                           reifiedResultShapes[resultNumber][*dimIndex]));
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Pass registration
 //===----------------------------------------------------------------------===//

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

getExpandedDimToCollapsedDimMap(ArrayRef<AffineMap> reassociation) {
   }
   return expandedDimToCollapsedDim;
 }
 
 /// For reshape op compute the shape at dimension `dimIndex` of the output in
 /// terms of shape of the `src`, when the reshape op is a collapsing
 /// operation. It is the product of the shape of the collapsed dimensions of the
 /// `src`.
-static OpFoldResult
-getCollapsedOutputDimFromInputShape(OpBuilder &builder, Location loc,
-                                    int64_t dimIndex, Value src,
-                                    ArrayRef<AffineMap> reassociationMap) {
+static OpFoldResult getCollapsedOutputDimFromInputShape(
+    OpBuilder &builder, Location loc, int64_t dimIndex, Value src,
+    ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociationMap) {
+  if (!ShapedType::isDynamic(dstStaticShape[dimIndex])) {
+    return builder.getIndexAttr(dstStaticShape[dimIndex]);
+  }
   AffineMap map = reassociationMap[dimIndex];
   unsigned startPos =
       map.getResults().front().cast<AffineDimExpr>().getPosition();
   unsigned endPos = map.getResults().back().cast<AffineDimExpr>().getPosition();
   AffineExpr expr;
   SmallVector<Value, 2> dynamicDims;
   for (auto dim : llvm::seq_inclusive(startPos, endPos)) {
     dynamicDims.push_back(builder.createOrFold<tensor::DimOp>(loc, src, dim));
     AffineExpr currExpr = builder.getAffineSymbolExpr(dim - startPos);
     expr = (expr ? expr * currExpr : currExpr);
   }
   return applyMapToValues(builder, loc,
                           AffineMap::get(0, endPos - startPos + 1, expr),
                           dynamicDims)[0];
 }
 
 /// Given the `src` of a collapsing reshape op and its reassociation maps,
 /// compute the shape of the result of the reshape.
 static SmallVector<OpFoldResult, 4> getCollapsedOutputShapeFromInputShape(
     OpBuilder &builder, Location loc, Value src,
     ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation) {
   return llvm::to_vector<4>(llvm::map_range(
       llvm::seq<int64_t>(0, dstStaticShape.size()), [&](int64_t dim) {
-        return getCollapsedOutputDimFromInputShape(builder, loc, dim, src,
-                                                   reassociation);
+        return getCollapsedOutputDimFromInputShape(
+            builder, loc, dim, src, dstStaticShape, reassociation);
       }));
 }
 
 /// For an expanding reshape op, compute the value for a dimension of the output
 /// from the shape of the input.
 static OpFoldResult getExpandedOutputDimFromInputShape(
     OpBuilder &builder, Location loc, int64_t dimIndex, Value src,
     ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation,
     llvm::DenseMap<int64_t, int64_t> &expandedDimToCollapsedDim) {
   if (!ShapedType::isDynamic(dstStaticShape[dimIndex])) {
-    return builder.getI64IntegerAttr(dstStaticShape[dimIndex]);
+    return builder.getIndexAttr(dstStaticShape[dimIndex]);
   }
   unsigned sourceDimPos = expandedDimToCollapsedDim[dimIndex];
   unsigned startPos = reassociation[sourceDimPos]
                           .getResults()
                           .front()
                           .cast<AffineDimExpr>()
                           .getPosition();
   unsigned endPos = reassociation[sourceDimPos]

 struct ReifyExpandOrCollapseShapeOp
     : public ReifyRankedShapedTypeOpInterface::ExternalModel<
           ReifyExpandOrCollapseShapeOp<OpTy>, OpTy> {
   LogicalResult
   reifyResultShapes(Operation *op, OpBuilder &b,
                     ReifiedRankedShapedTypeDims &reifiedReturnShapes) const {
     auto loc = op->getLoc();
     auto reshapeOp = cast<OpTy>(op);
-    auto resultShape = getReshapeOutputShapeFromInputShape(
+    reifiedReturnShapes.push_back(getReshapeOutputShapeFromInputShape(
         b, loc, reshapeOp.getSrc(), reshapeOp.getResultType().getShape(),
-        reshapeOp.getReassociationMaps());
-    reifiedReturnShapes.push_back(
-        getValueOrCreateConstantIndexOp(b, loc, resultShape));
+        reshapeOp.getReassociationMaps()));
     return success();
   }
 };
 
 namespace {
 
 struct ReifyPadOp
     : public ReifyRankedShapedTypeOpInterface::ExternalModel<ReifyPadOp,
                                                              PadOp> {
   LogicalResult
   reifyResultShapes(Operation *op, OpBuilder &b,
                     ReifiedRankedShapedTypeDims &reifiedReturnShapes) const {
     auto padOp = cast<PadOp>(op);
     Location loc = padOp.getLoc();
     auto lowPad = padOp.getMixedLowPad();
     auto highPad = padOp.getMixedHighPad();
-    SmallVector<Value> shapes;
+    SmallVector<OpFoldResult> shapes;
     for (auto dim : llvm::seq<int64_t>(0, padOp.getSourceType().getRank())) {
+      if (!padOp.getResultType().isDynamicDim(dim)) {
+        shapes.push_back(b.getIndexAttr(padOp.getResultType().getDimSize(dim)));
+        continue;
+      }
+
       // Shape along each dimension is source dim + low pad + high pad.
       SmallVector<Value> mapOperands;
       mapOperands.push_back(
           b.createOrFold<tensor::DimOp>(loc, padOp.getSource(), dim));
       AffineExpr expr = b.getAffineDimExpr(0);
       unsigned numSymbols = 0;
       auto addOpFoldResult = [&](OpFoldResult valueOrAttr) {
         if (Value v = valueOrAttr.dyn_cast<Value>()) {

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

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)))
       return failure();
-    mixedSizes = getAsOpFoldResult(reifiedShapes[opResult.getResultNumber()]);
+    mixedSizes = reifiedShapes[opResult.getResultNumber()];
   } else {
     // Static shape: Take static sizes directly.
     for (int64_t sz : tensorType.getShape())
       mixedSizes.push_back(b.getIndexAttr(sz));
   }
 
   // Create empty tensor.
   Value emptyTensor =
return emitOpError("incorrect number of dynamic sizes, has ")
            << getDynamicSizes().size() << ", expected "
            << getType().getNumDynamicDims();
   return success();
 }
 
 LogicalResult
 EmptyOp::reifyResultShapes(OpBuilder &builder,
                            ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
-  reifiedReturnShapes.resize(1, SmallVector<Value>(getType().getRank()));
+  reifiedReturnShapes.resize(1, SmallVector<OpFoldResult>(getType().getRank()));
   unsigned ctr = 0;
   for (int64_t i = 0; i < getType().getRank(); ++i) {
     if (getType().isDynamicDim(i)) {
       reifiedReturnShapes[0][i] = getDynamicSizes()[ctr++];
     } else {
-      reifiedReturnShapes[0][i] =
-          builder.create<arith::ConstantIndexOp>(getLoc(), i);
+      reifiedReturnShapes[0][i] = builder.getIndexAttr(getType().getDimSize(i));
     }
   }
   return success();
 }
 
 Value EmptyOp::getDynamicSize(unsigned idx) {
   assert(getType().isDynamicDim(idx) && "expected dynamic dim");
   unsigned ctr = 0;

 
 void GenerateOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   setNameFn(getResult(), "generated");
 }
 
 LogicalResult GenerateOp::reifyResultShapes(
     OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
-  reifiedReturnShapes.resize(1, SmallVector<Value>(getType().getRank()));
+  reifiedReturnShapes.resize(1, SmallVector<OpFoldResult>(getType().getRank()));
   int idx = 0;
   for (auto dim : llvm::seq<int64_t>(0, getType().getRank())) {
     if (getType().isDynamicDim(dim)) {
       reifiedReturnShapes[0][dim] = getOperand(idx++);
     } else {
-      reifiedReturnShapes[0][dim] = builder.create<arith::ConstantIndexOp>(
-          getLoc(), getType().getDimSize(dim));
+      reifiedReturnShapes[0][dim] =
+          builder.getIndexAttr(getType().getDimSize(dim));
     }
   }
   return success();
 }
 
 LogicalResult GenerateOp::verify() {
   // Ensure that the tensor type has as many dynamic dimensions as are
   // specified by the operands.

 }
 
 LogicalResult ExtractSliceOp::reifyResultShapes(
     OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
   reifiedReturnShapes.resize(1);
   reifiedReturnShapes[0].reserve(getType().getRank());
   SmallVector<OpFoldResult> mixedSizes = getMixedSizes();
   llvm::SmallBitVector droppedDims = getDroppedDims();
-  Location loc = getLoc();
   for (const auto &size : enumerate(mixedSizes)) {
     if (droppedDims.test(size.index()))
       continue;
-    if (auto attr = size.value().dyn_cast<Attribute>()) {
-      reifiedReturnShapes[0].push_back(builder.create<arith::ConstantIndexOp>(
-          loc, attr.cast<IntegerAttr>().getInt()));
-      continue;
-    }
-    reifiedReturnShapes[0].push_back(size.value().get<Value>());
+    reifiedReturnShapes[0].push_back(size.value());
   }
   return success();
 }
 
 namespace {
 /// Pattern to rewrite an extract_slice op with tensor::Cast arguments.
 /// This essentially pushes memref_cast past its consuming slice when
 /// `canFoldIntoConsumerOp` is true.
if (succeeded(foldInsertAfterInsertSlice(*this)))
     return getResult();
   if (auto result = foldInsertAfterExtractSlice(*this))
     return result;
   return OpFoldResult();
 }
 
 LogicalResult InsertSliceOp::reifyResultShapes(
     OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
-  reifiedReturnShapes.resize(1, SmallVector<Value>(getType().getRank()));
+  reifiedReturnShapes.resize(1, SmallVector<OpFoldResult>(getType().getRank()));
   for (auto dim : llvm::seq<int64_t>(0, getType().getRank())) {
     reifiedReturnShapes[0][dim] =
         builder.createOrFold<tensor::DimOp>(getLoc(), getDest(), dim);
   }
   return success();
 }
 
 namespace {

 
 template <typename OpTy>
 static LogicalResult
 reifyResultShapesImpl(OpTy op, OpBuilder &builder,
                       ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
   static_assert(llvm::is_one_of<OpTy, PackOp, UnPackOp>::value,
                 "applies to only pack or unpack operations");
   int64_t destRank = op.getDestRank();
-  reifiedReturnShapes.resize(1, SmallVector<Value>(destRank));
+  reifiedReturnShapes.resize(1, SmallVector<OpFoldResult>(destRank));
   for (auto dim : llvm::seq<int64_t>(0, destRank)) {
     reifiedReturnShapes[0][dim] =
         builder.createOrFold<tensor::DimOp>(op.getLoc(), op.getDest(), dim);
   }
   return success();
 }
 
 template <typename OpTy>

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

getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,
     auto zeroAttr = b.getI64IntegerAttr(0);
     resultOffsets.assign(offsets.begin(), offsets.end());
     resultOffsets.append(outputRank - inputRank, zeroAttr);
 
     ReifiedRankedShapedTypeDims outputShape;
     (void)packOp.reifyResultShapes(b, outputShape);
     resultSizes.assign(sizes.begin(), sizes.end());
     for (auto dataTileDim : llvm::seq<unsigned>(inputRank, outputRank))
-      resultSizes.push_back(getAsOpFoldResult(outputShape[0][dataTileDim]));
+      resultSizes.push_back(outputShape[0][dataTileDim]);
 
     return success();
   }
 };
 
 struct UnpackTileDimInfo {
   bool isAlignedToInnerTileSize;
   OpFoldResult sourceOffset;

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

LogicalResult matchAndRewrite(ReshapeOp reshapeOp,
     ReifiedRankedShapedTypeDims resultShapes;
     ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
         cast<ReifyRankedShapedTypeOpInterface>(reshapeOp.getOperation());
     if (failed(reifyShapedTypeInterface.reifyResultShapes(rewriter,
                                                           resultShapes)) ||
         !llvm::hasSingleElement(resultShapes))
       return failure();
     // TODO: Do not drop tensor type encoding.
-    Value emptyTensor =
-        rewriter.create<EmptyOp>(loc, getAsOpFoldResult(resultShapes[0]),
-                                 reshapeOp.getResultType().getElementType());
+    Value emptyTensor = rewriter.create<EmptyOp>(
+        loc, resultShapes[0], reshapeOp.getResultType().getElementType());
     if (emptyTensor.getType() != reshapeOp.getResultType()) {
       rewriter.replaceOpWithNewOp<tensor::CastOp>(
           reshapeOp, reshapeOp.getResultType(), emptyTensor);
     } else {
       rewriter.replaceOp(reshapeOp, emptyTensor);
     }
     return success();
   }

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

tensor::ExtractSliceFromCollapseHelper::create(OpBuilder &b,
 
   // 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)))
     return failure();
-  SmallVector<OpFoldResult> collapseShapeOutputShape =
-      getAsOpFoldResult(reifiedShapes[0]);
+  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 =
       getLinearizedDimensions(reassociationIndices);
   llvm::SmallBitVector slicedDimensions =

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

LogicalResult matchAndRewrite(tensor::ExtractSliceOp op,
 
     // Materialize the output shape values of the slice operation.
     ReifiedRankedShapedTypeDims reifiedShapes;
     if (failed(op.reifyResultShapes(rewriter, 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 = getAsOpFoldResult(reifiedShapes[0]);
+    SmallVector<OpFoldResult> outputShape = reifiedShapes[0];
     Value dest = rewriter.create<tensor::EmptyOp>(op->getLoc(), outputShape,
                                                   elementType);
 
     // Calculate the parameters for the tile loop nest.
     FailureOr<tensor::ExtractSliceFromCollapseHelper> params =
         tensor::ExtractSliceFromCollapseHelper::create(rewriter, collapseOp,
                                                        op);
     if (failed(params))

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

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