[mlir][tensor] Clean up tensor::DimOp usage

[springerm]

Summary

• Remove duplicate functions. tensor::getMixedSize and tensor::getMixedSizes should be used.
• Use tensor::getMixedSize instead of createOrFold<tensor::DimOp>. This is more efficient. createOrFold will create an op an immediately try to fold it. In case of a static dimension size, an attribute can be used directly.

Depends On: D153324

mlir/include/mlir/Dialect/Tensor/IR/Tensor.h

 /// Not all ops might be canonicalizable this way, but for those that can be,
 /// this method provides a check that it is worth doing the canonicalization.
 bool canFoldIntoProducerOp(CastOp castOp);
 
 /// Performs folding of any operand of `op` if it comes from a tensor::CastOp
 /// that can be folded.
 LogicalResult foldTensorCast(Operation *op);
 
+/// Return the dimension of the given tensor value.
+OpFoldResult getMixedSize(OpBuilder &builder, Location loc, Value value,
+                          int64_t dim);
+
 /// Return the dimensions of the given tensor value.
 SmallVector<OpFoldResult> getMixedSizes(OpBuilder &builder, Location loc,
                                         Value value);
 
 /// Create a rank-reducing ExtractSliceOp @[0 .. 0] with strides [1 .. 1] and
 /// appropriate sizes (i.e. `tensor.getSizes()`) to reduce the rank of `tensor`
 /// to that of `targetType`.
 Value createCanonicalRankReducingExtractSliceOp(OpBuilder &b, Location loc,

mlir/include/mlir/Dialect/Tensor/Utils/Utils.h

 PadOp createPadHighOp(RankedTensorType type, Value source, Value pad,
                       bool nofold, Location loc, OpBuilder &builder);
 
 // Creates dim ops for each dynamic dimension of the ranked tensor argument and
 // returns these as values.
 SmallVector<Value> createDynamicDimValues(OpBuilder &b, Location loc,
                                           Value rankedTensor);
 
-// Returns the tensor extent along dimension `dim` if `rankedTensor` is of
-// `RankedTensorType`. Returns `failure()` otherwise.
-FailureOr<OpFoldResult> createDimValue(OpBuilder &b, Location loc,
-                                       Value rankedTensor, int64_t dim);
-
-// Creates dim ops or constant ops for each dimension of the ranked tensor
-// argument and returns these as values.
-SmallVector<OpFoldResult> createDimValues(OpBuilder &b, Location loc,
-                                          Value rankedTensor);
-
 /// Returns the transposed `rankedTensorType` if `transposeVector` is non-empty.
 /// Fail if `transposeVector` is not a permutation matching the tensor rank.
 FailureOr<RankedTensorType>
 computeTransposedType(RankedTensorType rankedTensorType,
                       ArrayRef<int64_t> transposeVector);
 
 /// A tensor.insert_slice is a cast-like operation if it merely rank-extends the
 /// source tensor or inserts the source tensor into a destination tensor with

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

public:
 
   static llvm::SmallVector<Value> inferDynamicDimsForGather(OpBuilder &builder,
                                                             Location loc,
                                                             Value values,
                                                             Value indices) {
     llvm::SmallVector<Value> results;
 
     auto addDynamicDimension = [&](Value source, int64_t dim) {
-      auto dynamicDim = tensor::createDimValue(builder, loc, source, dim);
-      if (auto dimValue = llvm::dyn_cast_if_present<Value>(dynamicDim.value()))
+      auto sz = tensor::getMixedSize(builder, loc, source, dim);
+      if (auto dimValue = llvm::dyn_cast_if_present<Value>(sz))
         results.push_back(dimValue);
     };
 
     addDynamicDimension(values, 0);
     addDynamicDimension(indices, 1);
     addDynamicDimension(values, 2);
     return results;
   }

mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp

matchAndRewrite(tosa::ConcatOp op, OpAdaptor adaptor,
     Location loc = op.getLoc();
     int axis = op.getAxis();
     Value axisValue = rewriter.createOrFold<arith::ConstantOp>(
         loc, rewriter.getIndexAttr(axis));
     int64_t rank = resultType.getRank();
 
     SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
     SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
-    SmallVector<OpFoldResult> sizes = tensor::createDimValues(
-        rewriter, op.getLoc(), adaptor.getOperands()[0]);
+    SmallVector<OpFoldResult> sizes =
+        tensor::getMixedSizes(rewriter, op.getLoc(), adaptor.getOperands()[0]);
 
     // Pre-compute the offsets along the axis dimension.
     // The axisOffsets will be of size rank + 1, where the last value
     // will hold the total size of the tensor along the 'axis' dimension.
     SmallVector<OpFoldResult> axisOffsets;
     axisOffsets.push_back(rewriter.getIndexAttr(0));
     axisOffsets.push_back(sizes[axis]);
 
for (int64_t i = 0; i < rank; ++i) {
             getValueOrCreateConstantIndexOp(rewriter, loc, sizes[i]));
       }
     }
 
     Value result = rewriter.create<tensor::EmptyOp>(
         loc, resultType.getShape(), resultType.getElementType(), dynDims);
 
     for (auto [arg, offset] : llvm::zip(adaptor.getOperands(), axisOffsets)) {
-      auto sizes = tensor::createDimValues(rewriter, op.getLoc(), arg);
+      auto sizes = tensor::getMixedSizes(rewriter, op.getLoc(), arg);
       offsets[axis] = offset;
       result = rewriter.createOrFold<tensor::InsertSliceOp>(
           loc, arg, result, offsets, sizes, strides);
     }
     rewriter.replaceOp(op, result);
     return success();
   }
 };

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

LogicalResult matchAndRewrite(GenericOp genericOp,
 
     Location loc = genericOp.getLoc();
     SmallVector<Value> newOutputOperands = outputOperands;
     for (OpOperand *op : candidates) {
       OpBuilder::InsertionGuard guard(rewriter);
       rewriter.setInsertionPointAfterValue(op->get());
       auto elemType = cast<ShapedType>(op->get().getType()).getElementType();
       auto empty = rewriter.create<tensor::EmptyOp>(
-          loc, tensor::createDimValues(rewriter, loc, op->get()), elemType);
+          loc, tensor::getMixedSizes(rewriter, loc, op->get()), elemType);
 
       auto [start, end] = genericOp.getDpsInitsPositionRange();
       newOutputOperands[op->getOperandNumber() - start] = empty.getResult();
     }
 
     auto newOp = rewriter.create<GenericOp>(
         loc, genericOp.getResultTypes(), newInputOperands, newOutputOperands,
         newIndexingMaps, genericOp.getIteratorTypesArray(),

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

for (OpOperand *opOperand : op.getDpsInitOperands()) {
         if (sparse_tensor::getSparseTensorEncoding(operandVal.getType()))
           continue;
 
         // If outs is already an `empty` operation, nothing to do.
         auto definingOp = operandVal.getDefiningOp<tensor::EmptyOp>();
         if (definingOp)
           continue;
         modifiedOutput = true;
-        SmallVector<Value> dynamicDims;
-        for (const auto &dim : llvm::enumerate(operandType.getShape())) {
-          if (dim.value() != ShapedType::kDynamic)
-            continue;
-          dynamicDims.push_back(rewriter.createOrFold<tensor::DimOp>(
-              loc, operandVal, dim.index()));
-        }
+        SmallVector<OpFoldResult> mixedSizes =
+            tensor::getMixedSizes(rewriter, loc, operandVal);
         Value emptyTensor = rewriter.create<tensor::EmptyOp>(
-            loc, operandType.getShape(), operandType.getElementType(),
-            dynamicDims);
+            loc, mixedSizes, operandType.getElementType());
         op->setOperand(opOperand->getOperandNumber(), emptyTensor);
       }
     }
     if (!modifiedOutput) {
       rewriter.cancelRootUpdate(op);
       return failure();
     }
     rewriter.finalizeRootUpdate(op);

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

for (int64_t idx : llvm::seq<int64_t>(0, oldShape.size() + 1)) {
       if (idx == insertSplitDimension) {
         dispatchIndexOpFoldResults(sizes[idx], dynamicDims, newOutputShape);
         continue;
       }
       int64_t oldIdx = idx < insertSplitDimension ? idx : idx - 1;
       int64_t dim = oldShape[oldIdx];
       newOutputShape.push_back(dim);
       if (ShapedType::isDynamic(dim))
-        dynamicDims.push_back(b.createOrFold<tensor::DimOp>(
+        dynamicDims.push_back(b.create<tensor::DimOp>(
             loc, linalgOp.getDpsInitOperand(0)->get(), oldIdx));
     }
     Value emptyTensor = b.create<tensor::EmptyOp>(
         loc, newOutputShape, linalgOp.getRegionOutputArgs()[0].getType(),
         dynamicDims);
     Value constantOp = b.create<arith::ConstantOp>(loc, *identity);
     auto identityTensor =
         b.create<linalg::FillOp>(loc, constantOp, emptyTensor);

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

FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
   SmallVector<OpFoldResult> lows(packOp.getSourceRank(),
                                  rewriter.getIndexAttr(0));
   SmallVector<OpFoldResult> highs(packOp.getSourceRank(),
                                   rewriter.getIndexAttr(0));
   for (auto [pos, innerSize] :
        llvm::zip_equal(packOp.getInnerDimsPos(), packOp.getMixedTiles())) {
     int outerPos =
         packedToStripMinedShapePerm[packingMetadata.outerPositions[pos]];
-    OpFoldResult origSize = rewriter.createOrFold<tensor::DimOp>(
-        loc, packOp.getSource(),
-        rewriter.create<arith::ConstantIndexOp>(loc, pos));
-    OpFoldResult outerSize = rewriter.createOrFold<tensor::DimOp>(
-        loc, packOp.getDest(),
-        rewriter.create<arith::ConstantIndexOp>(loc, outerPos));
+    OpFoldResult origSize =
+        tensor::getMixedSize(rewriter, loc, packOp.getSource(), pos);
+    OpFoldResult outerSize =
+        tensor::getMixedSize(rewriter, loc, packOp.getDest(), outerPos);
     AffineExpr s0, d0, d1;
     bindDims(rewriter.getContext(), d0, d1);
     bindSymbols(rewriter.getContext(), s0);
     auto map = AffineMap::get(/*dimCount=*/2, /*symbolCount=*/1, d0 * s0 - d1);
     highs[pos] = affine::makeComposedFoldedAffineApply(
         rewriter, loc, map, {outerSize, origSize, innerSize});
   }
   RankedTensorType collapsed = tensor::CollapseShapeOp::inferCollapsedType(
GeneralizePadOpPattern::matchAndRewrite(tensor::PadOp padOp,
   };
 
   auto resultType = padOp.getResultType();
   // Compute size of EmptyOp. Any combination of static/dynamic is supported.
   SmallVector<Value> dynSizes;
   SmallVector<int64_t> staticSizes;
   for (unsigned dim = 0; dim < resultType.getRank(); ++dim) {
     if (resultType.isDynamicDim(dim)) {
-      auto srcSize = rewriter.createOrFold<tensor::DimOp>(
-          padOp.getLoc(), padOp.getSource(), dim);
+      auto srcSize = getIdxValue(tensor::getMixedSize(rewriter, padOp.getLoc(),
+                                                      padOp.getSource(), dim));
       // Add low and high padding value.
       auto plusLow = rewriter.createOrFold<arith::AddIOp>(
           padOp.getLoc(), srcSize, getIdxValue(padOp.getMixedLowPad()[dim]));
       auto plusHigh = rewriter.createOrFold<arith::AddIOp>(
           padOp.getLoc(), plusLow, getIdxValue(padOp.getMixedHighPad()[dim]));
       dynSizes.push_back(plusHigh);
     }
     staticSizes.push_back(resultType.getDimSize(dim));
   }
 
   // Init tensor and fill it with padding.
   Value emptyTensor = rewriter.create<tensor::EmptyOp>(
       padOp.getLoc(), staticSizes, resultType.getElementType(), dynSizes);
   Value fill = createFillOrGenerateOp(rewriter, padOp, emptyTensor, dynSizes);
 
   // Try optimize the copy of source.
   if (optimizeCopyFn && optimizeCopyFn(rewriter, padOp, fill).succeeded())
     return success();
 
   // tensor::PadOps cannot be optimized. Generate a InsertSliceOp instead
   // for copying the PadOp source.
   auto sourceType = padOp.getSourceType();
   // Compute size of source of tensor::PadOp.
-  SmallVector<OpFoldResult> srcSizes;
-  for (unsigned dim = 0; dim < sourceType.getRank(); ++dim) {
-    if (sourceType.isDynamicDim(dim)) {
-      srcSizes.push_back(rewriter.createOrFold<tensor::DimOp>(
-          padOp.getLoc(), padOp.getSource(), dim));
-    } else {
-      srcSizes.push_back(rewriter.getIndexAttr(sourceType.getDimSize(dim)));
-    }
-  }
+  SmallVector<OpFoldResult> srcSizes =
+      tensor::getMixedSizes(rewriter, padOp.getLoc(), padOp.getSource());
   // Strides of InsertSliceOp are all 1.
   SmallVector<OpFoldResult> strides(sourceType.getRank(),
                                     rewriter.getIndexAttr(1));
   rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
       padOp, padOp.getSource(), fill, padOp.getMixedLowPad(), srcSizes,
       strides);
 
   return success();
LogicalResult GeneralizeOuterUnitDimsUnPackOpPattern::matchAndRewrite(
   // transposed tile.
   int numLoops = transpShape.size();
   SmallVector<OpFoldResult> tileStrides(numLoops, oneIdxAttr);
   SmallVector<OpFoldResult> tileOffsets(numLoops, zeroIdxAttr);
   SmallVector<OpFoldResult> tileSizes;
   ArrayRef<int64_t> destShape = unpackOp.getDestType().getShape();
   for (auto i : llvm::seq<unsigned>(0, destRank)) {
     if (dimAndTileMapping.count(i) || destShape[i] != 1)
-      tileSizes.push_back(getAsOpFoldResult(
-          rewriter.createOrFold<tensor::DimOp>(loc, unpackOp.getDest(), i)));
+      tileSizes.push_back(
+          tensor::getMixedSize(rewriter, loc, unpackOp.getDest(), i));
   }
 
   auto partialTile = rewriter.create<tensor::ExtractSliceOp>(
       loc, transposedOp.getResult()[0], tileOffsets, tileSizes, tileStrides);
 
   // 4. Insert the result to the destination tensor.
   SmallVector<OpFoldResult> writeSizes;
   SmallVector<OpFoldResult> writeStrides(destRank, oneIdxAttr);

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

mlir::scf::tileReductionUsingScf(RewriterBase &b,
   // 3. Generate the tiled implementation within the inner most loop.
   b.setInsertionPoint(loops.back().getBody()->getTerminator());
   Operation *parallelOp = op.tileToPartialReduction(
       b, loc, (*identityTensor)->getResults(), offsets, sizes, reductionDim);
 
   SmallVector<OpFoldResult> resultSizesList;
   for (size_t i = 0; i < offsets.size(); i++)
     resultSizesList.push_back(
-        b.createOrFold<tensor::DimOp>(loc, parallelOp->getResult(0), i));
+        tensor::getMixedSize(b, loc, parallelOp->getResult(0), i));
   SmallVector<OpFoldResult> outOffsets(offsets.size(), b.getIndexAttr(0));
   SmallVector<Value> replacements = yieldTiledValues(
       b, (*identityTensor)->getResults(), parallelOp->getResults(), outOffsets,
       resultSizesList, loops);
 
   auto dstOp = cast<DestinationStyleOpInterface>(parallelOp);
   auto innerMostLoop = loops.back();
   SmallVector<Value> destinationTensors = dstOp.getDpsInitOperands();

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

Operation *TensorDialect::materializeConstant(OpBuilder &builder,
   if (auto op = arith::ConstantOp::materialize(builder, value, type, loc))
     return op;
   if (complex::ConstantOp::isBuildableWith(value, type))
     return builder.create<complex::ConstantOp>(loc, type,
                                                llvm::cast<ArrayAttr>(value));
   return nullptr;
 }
 
+OpFoldResult tensor::getMixedSize(OpBuilder &builder, Location loc, Value value,
+                                  int64_t dim) {
+  auto tensorType = llvm::cast<RankedTensorType>(value.getType());
+  SmallVector<OpFoldResult> result;
+  if (tensorType.isDynamicDim(dim))
+    return builder.createOrFold<tensor::DimOp>(loc, value, dim);
+
+  return builder.getIndexAttr(tensorType.getDimSize(dim));
+}
+
 SmallVector<OpFoldResult> tensor::getMixedSizes(OpBuilder &builder,
                                                 Location loc, Value value) {
   auto tensorType = llvm::cast<RankedTensorType>(value.getType());
   SmallVector<OpFoldResult> result;
-  for (int64_t i = 0; i < tensorType.getRank(); ++i) {
-    if (tensorType.isDynamicDim(i)) {
-      Value size = builder.create<tensor::DimOp>(loc, value, i);
-      result.push_back(size);
-    } else {
-      result.push_back(builder.getIndexAttr(tensorType.getDimSize(i)));
-    }
-  }
+  for (int64_t i = 0; i < tensorType.getRank(); ++i)
+    result.push_back(getMixedSize(builder, loc, value, i));
   return result;
 }
 
 FailureOr<Value> tensor::getOrCreateDestination(OpBuilder &b, Location loc,
                                                 OpResult opResult) {
   auto tensorType = llvm::dyn_cast<TensorType>(opResult.getType());
   assert(tensorType && "expected tensor type");
 
OpFoldResult InsertSliceOp::fold(FoldAdaptor) {
   if (auto result = foldInsertAfterExtractSlice(*this))
     return result;
   return OpFoldResult();
 }
 
 LogicalResult InsertSliceOp::reifyResultShapes(
     OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
   reifiedReturnShapes.resize(1, SmallVector<OpFoldResult>(getType().getRank()));
-  for (auto dim : llvm::seq<int64_t>(0, getType().getRank())) {
-    if (getType().isDynamicDim(dim)) {
-      reifiedReturnShapes[0][dim] =
-          builder.createOrFold<tensor::DimOp>(getLoc(), getDest(), dim);
-    } else {
-      reifiedReturnShapes[0][dim] =
-          builder.getIndexAttr(getType().getDimSize(dim));
-    }
-  }
+  reifiedReturnShapes[0] = tensor::getMixedSizes(builder, getLoc(), getDest());
   return success();
 }
 
 namespace {
 /// Pattern to rewrite a insert_slice op with constant arguments.
 ///
 /// This pattern works with both InsertSliceOp and ParallelInsertSliceOp.
 template <typename InsertOpTy>

 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<OpFoldResult>(destRank));
-  ShapedType resultType = llvm::cast<ShapedType>(op.getResult().getType());
-  for (auto dim : llvm::seq<int64_t>(0, destRank)) {
-    if (resultType.isDynamicDim(dim)) {
-      reifiedReturnShapes[0][dim] =
-          builder.createOrFold<tensor::DimOp>(op.getLoc(), op.getDest(), dim);
-    } else {
-      reifiedReturnShapes[0][dim] =
-          builder.getIndexAttr(resultType.getDimSize(dim));
-    }
-  }
+  reifiedReturnShapes[0] =
+      tensor::getMixedSizes(builder, op.getLoc(), op.getDest());
   return success();
 }
 
 template <typename OpTy>
 static DenseMap<int64_t, OpFoldResult> getDimAndTileMappingImpl(OpTy op) {
   static_assert(llvm::is_one_of<OpTy, PackOp, UnPackOp>::value,
                 "applies to only pack or unpack operations");
   DenseMap<int64_t, OpFoldResult> dimAndTileMapping;

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

getTiledImplementation(Operation *op, OpBuilder &b,
     SmallVector<OpFoldResult> origOffsets(offsets.begin(), offsets.end());
     SmallVector<OpFoldResult> origSizes(sizes.begin(), sizes.end());
     applyPermToRange(origOffsets, origSizes,
                      invertPermutationVector(packOp.getOuterDimsPerm()));
 
     DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
         packOp.getDimAndTileMapping();
     SmallVector<OpFoldResult> srcDimValues =
-        tensor::createDimValues(b, loc, packOp.getSource());
+        tensor::getMixedSizes(b, loc, packOp.getSource());
     SmallVector<OpFoldResult> inputIndices, inputSizes;
     for (auto dim : llvm::seq<int64_t>(0, inputRank)) {
       using AV = affine::AffineValueExpr;
       affine::AffineBuilder ab(b, loc);
       AffineExpr dim0, dim1, sym;
       bindDims(b.getContext(), dim0, dim1);
       bindSymbols(b.getContext(), sym);
       if (dimAndTileMapping.count(dim)) {
FailureOr<TilingResult> tensor::bubbleUpPadSlice(OpBuilder &b,
   int64_t rank = padOp.getSourceType().getRank();
   for (unsigned dim = 0; dim < rank; ++dim) {
     auto low = padOp.getMixedLowPad()[dim];
     bool hasLowPad = !isConstantIntValue(low, 0);
     auto high = padOp.getMixedHighPad()[dim];
     bool hasHighPad = !isConstantIntValue(high, 0);
     auto offset = offsets[dim];
     auto length = sizes[dim];
-    auto srcSize =
-        tensor::createDimValue(b, loc, padOp.getSource(), dim).value();
+    auto srcSize = tensor::getMixedSize(b, loc, padOp.getSource(), dim);
 
     // The new amount of low padding is `low - offset`. Except for the case
     // where none of the low padding is read. In that case, the new amount of
     // low padding is zero.
     //
     // Optimization: If low = 0, then newLow = 0.
     OpFoldResult newLow = hasLowPad ? max(zero, sub(low, offset)) : zero;
     newLows.push_back(newLow);

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

 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
 #include "llvm/ADT/STLExtras.h"
 
 using namespace mlir;
 using namespace mlir::affine;
 using namespace mlir::tensor;
 
-/// Get the dimension size of a value of RankedTensor type at the
-static OpFoldResult getShapeDimSize(OpBuilder &b, Location loc,
-                                    Value rankedTensor, int64_t dimIdx) {
-  RankedTensorType tensorType = cast<RankedTensorType>(rankedTensor.getType());
-  if (!tensorType.isDynamicDim(dimIdx)) {
-    return b.getIndexAttr(tensorType.getDimSize(dimIdx));
-  }
-  Value idxValue = b.create<arith::ConstantIndexOp>(loc, dimIdx);
-  return b.createOrFold<tensor::DimOp>(loc, rankedTensor, idxValue);
-}
-
-/// Get all the dimension sizes of a value of RankedTensor type.
-static SmallVector<OpFoldResult> getShapeDimSizes(OpBuilder &b, Location loc,
-                                                  Value rankedTensor) {
-  SmallVector<OpFoldResult> dimSizes;
-  RankedTensorType tensorType = cast<RankedTensorType>(rankedTensor.getType());
-  for (unsigned i = 0; i < tensorType.getRank(); i++)
-    dimSizes.push_back(getShapeDimSize(b, loc, rankedTensor, i));
-  return dimSizes;
-}
-
 /// A tuple that represents (dimension number, dimension value).
 using DimAndIndex = std::tuple<unsigned, Value>;
 
 /// Transform `dimAndIndex` from the output index space of a (non-rank-reducing)
 /// slice described by `sliceParams` into the input index space.
 static DimAndIndex invertSliceIndexing(OpBuilder &b, Location loc,
                                        ArrayRef<Range> sliceParams,
                                        const DimAndIndex &dimAndIndex) {
tensor::ExtractSliceFromCollapseHelper::create(OpBuilder &b,
 
   // Determine which of the CollapseShapeOp's result dimensions are sliced
   // and/or linearized.
   llvm::SmallBitVector linearizedDimensions =
       getLinearizedDimensions(reassociationIndices);
   llvm::SmallBitVector slicedDimensions =
       getSlicedDimensions(collapseShapeOutputShape, sliceParams);
 
-  auto collapseShapeInputShape = getShapeDimSizes(b, op.getLoc(), op.getSrc());
+  auto collapseShapeInputShape =
+      tensor::getMixedSizes(b, op.getLoc(), op.getSrc());
 
   SmallVector<Value> tileSizes;
   for (unsigned i = 0; i < sliceParams.size(); i++) {
     if (slicedDimensions[i] && linearizedDimensions[i])
       tileSizes.push_back(
           getValueOrCreateConstantIndexOp(b, op.getLoc(), sliceParams[i].size));
   }
 
tensor::simplifyCollapseShapeWithRankReducingExtractSlice(
   if (failed(info))
     return failure();
 
   // Create the rank-reducing extract slice op.
   auto zero = rewriter.getIndexAttr(0);
   auto one = rewriter.getIndexAttr(1);
   SmallVector<OpFoldResult> offsets(sourceType.getRank(), zero);
   SmallVector<OpFoldResult> sizes =
-      getShapeDimSizes(rewriter, op.getLoc(), op.getSrc());
+      tensor::getMixedSizes(rewriter, op.getLoc(), op.getSrc());
   SmallVector<OpFoldResult> strides(sourceType.getRank(), one);
   auto sliceOp = rewriter.create<tensor::ExtractSliceOp>(
       op.getLoc(), info->sliceResultType, op.getSrc(), offsets, sizes, strides);
 
   if (!info->newReassociationIndices.has_value()) {
     rewriter.replaceOp(op, sliceOp.getResult());
     return sliceOp.getOperation();
   }
 
   return rewriter
       .replaceOpWithNewOp<tensor::CollapseShapeOp>(
           op, sliceOp.getResult(), *info->newReassociationIndices)
       .getOperation();
 }

mlir/lib/Dialect/Tensor/Utils/Utils.cpp

PadOp mlir::tensor::createPadHighOp(RankedTensorType type, Value source,
   SmallVector<OpFoldResult> high(type.getRank(), zero);
   for (const auto &en : enumerate(type.getShape())) {
     // Pad only the static dimensions of the result tensor type.
     if (ShapedType::isDynamic(en.value()))
       continue;
     // Compute the padding width.
     AffineExpr d0;
     bindDims(b.getContext(), d0);
-    auto dimOp = b.createOrFold<tensor::DimOp>(loc, source, en.index());
+    OpFoldResult sz = tensor::getMixedSize(b, loc, source, en.index());
     high[en.index()] =
-        affine::makeComposedAffineApply(b, loc, en.value() - d0, {dimOp})
+        affine::makeComposedAffineApply(b, loc, en.value() - d0, {sz})
             .getResult();
   }
   return b.create<PadOp>(loc, type, source, low, high, pad, nofold);
 }
 
 SmallVector<Value> mlir::tensor::createDynamicDimValues(OpBuilder &b,
                                                         Location loc,
                                                         Value rankedTensor) {
   auto tensorTy = cast<RankedTensorType>(rankedTensor.getType());
   SmallVector<Value> dynamicDims;
   for (const auto &en : llvm::enumerate(tensorTy.getShape())) {
     if (en.value() == ShapedType::kDynamic)
       dynamicDims.push_back(
           b.create<tensor::DimOp>(loc, rankedTensor, en.index()));
   }
   return dynamicDims;
 }
 
-FailureOr<OpFoldResult> mlir::tensor::createDimValue(OpBuilder &b, Location loc,
-                                                     Value rankedTensor,
-                                                     int64_t dim) {
-  auto tensorTy = dyn_cast<RankedTensorType>(rankedTensor.getType());
-  if (!tensorTy)
-    return failure();
-  auto shape = tensorTy.getShape();
-  if (dim >= static_cast<int64_t>(shape.size()))
-    return failure();
-  if (ShapedType::isDynamic(shape[dim]))
-    return OpFoldResult(b.createOrFold<tensor::DimOp>(loc, rankedTensor, dim));
-  return OpFoldResult(b.getIndexAttr(shape[dim]));
-}
-
-SmallVector<OpFoldResult>
-mlir::tensor::createDimValues(OpBuilder &b, Location loc, Value rankedTensor) {
-  auto tensorTy = cast<RankedTensorType>(rankedTensor.getType());
-  SmallVector<OpFoldResult> dims;
-  for (const auto &en : llvm::enumerate(tensorTy.getShape())) {
-    if (ShapedType::isDynamic(en.value())) {
-      dims.push_back(
-          b.createOrFold<tensor::DimOp>(loc, rankedTensor, en.index()));
-    } else {
-      dims.push_back(b.getIndexAttr(en.value()));
-    }
-  }
-  return dims;
-}
-
 FailureOr<RankedTensorType>
 mlir::tensor::computeTransposedType(RankedTensorType rankedTensorType,
                                     ArrayRef<int64_t> transposeVector) {
   if (transposeVector.empty())
     return rankedTensorType;
 
   if (!isPermutationVector(transposeVector) ||
       transposeVector.size() != static_cast<size_t>(rankedTensorType.getRank()))