[Mlir-commits] [mlir] 64b9188 - Remove restriction on static dimensions in Shape method

[Tres Popp]

Author: Tres Popp
Date: 2022-02-08T11:20:01+01:00
New Revision: 64b918852c09db6461d4c4c97df288c7b7fee865

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

LOG: Remove restriction on static dimensions in Shape method

mlir::shape::ToExtentTensorOp::areCastCompatible didn't allow the input
to have a static dimension, but that is allowed.

Added: 
    

Modified: 
    mlir/lib/Dialect/Shape/IR/Shape.cpp
    mlir/test/Dialect/Shape/ops.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index 1ae9e68d2024f..89a976686492c 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -1767,7 +1767,7 @@ bool ToExtentTensorOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
     return false;
   if (auto inputTensor = inputs[0].dyn_cast<RankedTensorType>()) {
     if (!inputTensor.getElementType().isa<IndexType>() ||
-        inputTensor.getRank() != 1 || !inputTensor.isDynamicDim(0))
+        inputTensor.getRank() != 1)
       return false;
   } else if (!inputs[0].isa<ShapeType>()) {
     return false;

diff  --git a/mlir/test/Dialect/Shape/ops.mlir b/mlir/test/Dialect/Shape/ops.mlir
index e4f22a2edf502..8c21878c1ef85 100644
--- a/mlir/test/Dialect/Shape/ops.mlir
+++ b/mlir/test/Dialect/Shape/ops.mlir
@@ -168,6 +168,11 @@ func @test_to_extent_tensor(%arg: !shape.shape) -> tensor<3xindex> {
   return %0 : tensor<3xindex>
 }
 
+func @test_identity_to_extent_tensor(%arg: tensor<3xindex>) -> tensor<3xindex> {
+  %0 = shape.to_extent_tensor %arg : tensor<3xindex> -> tensor<3xindex>
+  return %0 : tensor<3xindex>
+}
+
 func @test_from_extent_tensor(%arg: tensor<?xindex>) -> !shape.shape {
   %0 = shape.from_extent_tensor %arg : tensor<?xindex>
   return %0 : !shape.shape