diff --git a/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h b/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h --- a/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h +++ b/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h @@ -28,6 +28,11 @@ ArrayRef low, ArrayRef 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 createDynamicDimValues(OpBuilder &b, Location loc, + Value rankedTensor); + } // namespace tensor } // namespace mlir diff --git a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp --- a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp @@ -55,3 +55,16 @@ } return createPadScalarOp(type, source, pad, low, high, nofold, loc, b); } + +SmallVector mlir::tensor::createDynamicDimValues(OpBuilder &b, + Location loc, + Value rankedTensor) { + auto tensorTy = rankedTensor.getType().cast(); + SmallVector dynamicDims; + for (const auto &en : llvm::enumerate(tensorTy.getShape())) { + if (en.value() == ShapedType::kDynamicSize) + dynamicDims.push_back( + b.create(loc, rankedTensor, en.index())); + } + return dynamicDims; +}