[Mlir-commits] [mlir] 2c3ca3b - [MLIR] Add utility function to create values for all dimensions of a tensor value

[Frederik Gossen]

Author: Frederik Gossen
Date: 2022-08-12T14:42:27-04:00
New Revision: 2c3ca3b684bb2b188d977d47548e79dc559fb8ad

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

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

This is a variant of the already provided `createDynamicDimValues` helper.

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

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 8426c63d191cc..c7547c050c90b 100644
--- a/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
@@ -28,11 +28,16 @@ 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
 

diff  --git a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
index 1d9ac83d1abf8..e3d008b6eed2a 100644
--- a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
@@ -68,3 +68,14 @@ SmallVector<Value> mlir::tensor::createDynamicDimValues(OpBuilder &b,
   }
   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;
+}