[Mlir-commits] [mlir] [mlir][tosa] Guard pooling shape inference on unranked inputs (PR #177999)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-tosa

@llvm/pr-subscribers-mlir

Author: Yi-Chi Lee (yichi170)

<details>
<summary>Changes</summary>

poolingInferReturnTypes didn't properly guard the unknown rank, triggering an assertion in ShapeAdaptor::getDimSize.

---
Full diff: https://github.com/llvm/llvm-project/pull/177999.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Tosa/IR/TosaOps.cpp (+1-1) 
- (modified) mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir (+13) 


``````````diff
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index 6205161599899..c412d788c9b29 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -3409,7 +3409,7 @@ static LogicalResult poolingInferReturnTypes(
   outputShape.resize(4, ShapedType::kDynamic);
 
   // We only know the rank if the input type is unranked.
-  if (!inputShape) {
+  if (!inputShape.hasRank()) {
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
   }
diff --git a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
index 610fdb6d32ad4..5c73e37b47dbb 100644
--- a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
@@ -1743,3 +1743,16 @@ func.func @test_tconv2d_bias_broadcast(%input: tensor<2x6x7x3xf32>, %weight: ten
        : (tensor<2x6x7x3xf32>, tensor<?x3x3x3xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<?x?x?x?xf32>
     return
   }
+
+// -----
+
+// CHECK-LABEL: test_pool2d_unknown_rank
+func.func @test_pool2d_unknown_rank() {
+  %0 = gpu.block_id  x
+  %1 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+  %2 = tosa.reduce_sum %1 {axis = 0 : i32} : (tensor<1xi32>) -> tensor<1xi32>
+  %3 = tosa.bitwise_or %2, %2 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
+  %4 = tensor.cast %3 : tensor<1xi32> to tensor<*xi32>
+  %5 = tosa.avg_pool2d %4, %1, %1 {acc_type = i32, kernel = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<*xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xi32>
+  return
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/177999