[Mlir-commits] [mlir] ff6ce9e - Add `createDynamicDimValues` to tensor dialect utils

[Frederik Gossen]

Author: Frederik Gossen
Date: 2022-06-13T08:26:28-04:00
New Revision: ff6ce9e8fc7c8b4e5bd7275583d571354454162a

URL: https://github.com/llvm/llvm-project/commit/ff6ce9e8fc7c8b4e5bd7275583d571354454162a
DIFF: https://github.com/llvm/llvm-project/commit/ff6ce9e8fc7c8b4e5bd7275583d571354454162a.diff

LOG: Add `createDynamicDimValues` to tensor dialect utils

The function creates dim ops for each dynamic dimension of the raked tensor
argument and returns these as values.

Differential Revision: https://reviews.llvm.org/D127533

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
    mlir/lib/Dialect/Tensor/Utils/Utils.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h b/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
index f8c8d3dd357b4..8426c63d191cc 100644
--- a/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
@@ -28,6 +28,11 @@ 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
 

diff  --git a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
index b35f554ca1a52..5eb40a2cb5dab 100644
--- a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
@@ -55,3 +55,16 @@ PadOp mlir::tensor::createPadHighOp(RankedTensorType type, Value source,
   }
   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;
+}