[MLIR] Add utility function to create values for all dimensions of a tensor value

[frgossen]

Summary

This is a variant of the already provided createDynamicDimValues helper.

[cota]

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

36

[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
+// 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);
 
+// Creates dim ops or constant ops for each dimension of the ranked tensor
+// argument and returns these as values.
+SmallVector<Value> createDimValues(OpBuilder &b, Location loc,
+                                   Value rankedTensor);
+
 } // namespace tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_UTILS_UTILS_H_

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

SmallVector<Value> mlir::tensor::createDynamicDimValues(OpBuilder &b,
   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;
 }
+
+SmallVector<Value> mlir::tensor::createDimValues(OpBuilder &b, Location loc,
+                                                 Value rankedTensor) {
+  auto tensorTy = rankedTensor.getType().cast<RankedTensorType>();
+  SmallVector<Value> dims;
+  for (const auto &en : llvm::enumerate(tensorTy.getShape())) {
+    dims.push_back(
+        b.createOrFold<tensor::DimOp>(loc, rankedTensor, en.index()));
+  }
+  return dims;
+}