Add createDynamicDimOps to tensor dialect utils

[pifon2a]

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

32

nit: I think the name is a big misleading. What you are doing here is you compute values for every dynamic dimensions of the shape. Maybe call it getDynamicDimValues?

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

63

nit: usually this is called en.

[akuegel]

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

31

nit: maybe "ranked tensor value" instead of just "tensor value"?

[frgossen]

Address comments

[frgossen]

Thanks

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

 
 // Return a PadOp that pads `source to `type` size with `pad` value.
 // I.e., a block will be created and the `pad` value will be yielded
 // directly. If the type passed is nullptr, it is inferred.
 PadOp createPadScalarOp(Type type, Value source, Value pad,
                         ArrayRef<OpFoldResult> low, ArrayRef<OpFoldResult> high,
                         bool nofold, Location loc, OpBuilder &builder);
 
+// Creates dim ops for each dynamic dimension of the raked tensor argument and
+// returns these as values.
+SmallVector<Value> createDynamicDimValues(OpBuilder &b, Location loc,
+                                          Value rankedTensor);
+
 } // namespace tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_UTILS_UTILS_H_

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

for (const auto &en : enumerate(type.getShape())) {
     AffineExpr d0;
     bindDims(b.getContext(), d0);
     auto dimOp = b.createOrFold<tensor::DimOp>(loc, source, en.index());
     high[en.index()] =
         makeComposedAffineApply(b, loc, en.value() - d0, {dimOp}).getResult();
   }
   return createPadScalarOp(type, source, pad, low, high, nofold, loc, b);
 }
+
+SmallVector<Value> mlir::tensor::createDynamicDimValues(OpBuilder &b,
+                                                        Location loc,
+                                                        Value rankedTensor) {
+  auto tensorTy = rankedTensor.getType().cast<RankedTensorType>();
+  SmallVector<Value> dynamicDims;
+  for (const auto &en : llvm::enumerate(tensorTy.getShape())) {
+    if (en.value() == ShapedType::kDynamicSize)
+      dynamicDims.push_back(
+          b.create<tensor::DimOp>(loc, rankedTensor, en.index()));
+  }
+  return dynamicDims;
+}