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,10 @@ ArrayRef low, ArrayRef high, bool nofold, Location loc, OpBuilder &builder); +// Create dim ops for the dynamic dimensions of a tensor value. +SmallVector createDynamicDimOps(OpBuilder &b, Location loc, + Value tensor); + } // 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,14 @@ } return createPadScalarOp(type, source, pad, low, high, nofold, loc, b); } + +SmallVector mlir::tensor::createDynamicDimOps(OpBuilder &b, Location loc, + Value tensor) { + auto tensorTy = tensor.getType().cast(); + SmallVector dynamicDims; + for (const auto &it : llvm::enumerate(tensorTy.getShape())) { + if (it.value() == ShapedType::kDynamicSize) + dynamicDims.push_back(b.create(loc, tensor, it.index())); + } + return dynamicDims; +}