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.
llvm · Feb 8, 2022 · 64b9188 · 64b9188
1 parent 99d9502
commit 64b9188
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git 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
@@ -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