[mlir][tensor] Add getMixedSizes helper

[springerm]

Summary

This helper function computes the dimensions of a tensor value as OpFoldResults.

[springerm]

update more uses

[nicolasvasilache]

Thanks!

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 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,
                                                 Value tensor,
                                                 RankedTensorType targetType);
 
 /// Create a rank-reducing InsertSliceOp @[0 .. 0] with strides [1 .. 1] and

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

Operation *TensorDialect::materializeConstant(OpBuilder &builder,
   if (arith::ConstantOp::isBuildableWith(value, type))
     return builder.create<arith::ConstantOp>(loc, value, type);
   if (complex::ConstantOp::isBuildableWith(value, type))
     return builder.create<complex::ConstantOp>(loc, type,
                                                value.cast<ArrayAttr>());
   return nullptr;
 }
 
+SmallVector<OpFoldResult> tensor::getMixedSizes(OpBuilder &builder,
+                                                Location loc, Value value) {
+  auto tensorType = value.getType().cast<RankedTensorType>();
+  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)));
+    }
+  }
+  return result;
+}
+
 //===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//
 
 /// Returns true if `target` is a ranked tensor type that preserves static
 /// information available in the `source` ranked tensor type.
 bool mlir::tensor::preservesStaticInformation(Type source, Type target) {
   auto sourceType = source.dyn_cast<RankedTensorType>();
OpFoldResult ExtractSliceOp::fold(ArrayRef<Attribute> operands) {
 
   return OpFoldResult();
 }
 
 Value mlir::tensor::createCanonicalRankReducingExtractSliceOp(
     OpBuilder &b, Location loc, Value tensor, RankedTensorType targetType) {
   auto rankedTensorType = tensor.getType().cast<RankedTensorType>();
   unsigned rank = rankedTensorType.getRank();
-  auto shape = rankedTensorType.getShape();
   SmallVector<OpFoldResult> offsets(rank, b.getIndexAttr(0));
-  SmallVector<OpFoldResult> sizes;
-  for (unsigned i = 0, e = rank; i < e; ++i) {
-    OpFoldResult dim;
-    if (rankedTensorType.isDynamicDim(i))
-      dim = b.createOrFold<tensor::DimOp>(
-          loc, tensor, b.create<arith::ConstantIndexOp>(loc, i));
-    else
-      dim = b.getIndexAttr(shape[i]);
-    sizes.push_back(dim);
-  }
+  SmallVector<OpFoldResult> sizes = getMixedSizes(b, loc, tensor);
   SmallVector<OpFoldResult> strides(rank, b.getIndexAttr(1));
   return b.createOrFold<tensor::ExtractSliceOp>(loc, targetType, tensor,
                                                 offsets, sizes, strides);
 }
 
 //===----------------------------------------------------------------------===//
 // InsertSliceOp
 //===----------------------------------------------------------------------===//

 }
 
 Value mlir::tensor::createCanonicalRankReducingInsertSliceOp(OpBuilder &b,
                                                              Location loc,
                                                              Value tensor,
                                                              Value dest) {
   auto rankedTensorType = dest.getType().cast<RankedTensorType>();
   unsigned rank = rankedTensorType.getRank();
-  auto shape = rankedTensorType.getShape();
   SmallVector<OpFoldResult> offsets(rank, b.getIndexAttr(0));
-  SmallVector<OpFoldResult> sizes;
-  for (unsigned i = 0, e = rank; i < e; ++i) {
-    OpFoldResult dim;
-    if (rankedTensorType.isDynamicDim(i))
-      dim = b.createOrFold<tensor::DimOp>(
-          loc, dest, b.create<arith::ConstantIndexOp>(loc, i));
-    else
-      dim = b.getIndexAttr(shape[i]);
-    sizes.push_back(dim);
-  }
+  SmallVector<OpFoldResult> sizes = getMixedSizes(b, loc, dest);
   SmallVector<OpFoldResult> strides(rank, b.getIndexAttr(1));
   return b.createOrFold<tensor::InsertSliceOp>(loc, tensor, dest, offsets,
                                                sizes, strides);
 }
 
 //===----------------------------------------------------------------------===//
 // PadOp
 //===----------------------------------------------------------------------===//

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

if (padOp->hasAttr(BufferizationDialect::kEscapeAttrName))
       generateOp->setAttr(
           BufferizationDialect::kEscapeAttrName,
           padOp->getAttr(BufferizationDialect::kEscapeAttrName));
     // TODO: Memory space
     rewriter.inlineRegionBefore(padOp.getRegion(), generateOp.getBody(),
                                 generateOp.getBody().begin());
 
     // Create tensor::InsertSliceOp.
-    SmallVector<OpFoldResult> sliceSizes, sliceStrides;
-    for (int64_t i = 0; i < resultType.getRank(); ++i) {
-      sliceStrides.push_back(rewriter.getIndexAttr(1));
-      if (srcType.isDynamicDim(i)) {
-        Value size = rewriter.create<tensor::DimOp>(loc, padOp.getSource(), i);
-        sliceSizes.push_back(size);
-      } else {
-        sliceSizes.push_back(rewriter.getIndexAttr(srcType.getDimSize(i)));
-      }
-    }
+    SmallVector<OpFoldResult> sliceSizes =
+        getMixedSizes(rewriter, loc, padOp.getSource());
+    SmallVector<OpFoldResult> sliceStrides(srcType.getRank(),
+                                           rewriter.getIndexAttr(1));
     rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
         padOp, padOp.getSource(), generateOp.getResult(),
         /*offsets=*/padOp.getMixedLowPad(), sliceSizes, sliceStrides);
 
     return success();
   }
 };