[Mlir-commits] [mlir] 188b041 - [mlir][Tensor] Change `createDimValues` to return a list of `OpFoldResult`s.

[Mahesh Ravishankar]

Author: Mahesh Ravishankar
Date: 2022-10-27T03:13:56Z
New Revision: 188b041bf5d61fd0e6c911aeada1535cac58e1d8

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

LOG: [mlir][Tensor] Change `createDimValues` to return a list of `OpFoldResult`s.

Reviewed By: nicolasvasilache, hanchung, ThomasRaoux

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

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 66035efa7f36e..07e8bbf05d919 100644
--- a/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
@@ -28,8 +28,8 @@ SmallVector<Value> createDynamicDimValues(OpBuilder &b, Location loc,
 
 // 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);
+SmallVector<OpFoldResult> 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 06c344be0a06e..5f6d180a70999 100644
--- a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
@@ -51,13 +51,17 @@ SmallVector<Value> mlir::tensor::createDynamicDimValues(OpBuilder &b,
   return dynamicDims;
 }
 
-SmallVector<Value> mlir::tensor::createDimValues(OpBuilder &b, Location loc,
-                                                 Value rankedTensor) {
+SmallVector<OpFoldResult>
+mlir::tensor::createDimValues(OpBuilder &b, Location loc, Value rankedTensor) {
   auto tensorTy = rankedTensor.getType().cast<RankedTensorType>();
-  SmallVector<Value> dims;
+  SmallVector<OpFoldResult> dims;
   for (const auto &en : llvm::enumerate(tensorTy.getShape())) {
-    dims.push_back(
-        b.createOrFold<tensor::DimOp>(loc, rankedTensor, en.index()));
+    if (ShapedType::isDynamic(en.value())) {
+      dims.push_back(
+          b.createOrFold<tensor::DimOp>(loc, rankedTensor, en.index()));
+    } else {
+      dims.push_back(b.getIndexAttr(en.value()));
+    }
   }
   return dims;
 }