Do not crash when verifiying/shape-infering GridSample if inputs have not static shape

Signed-off-by: Rickert, Jonas <[email protected]>

Author: jorickert

src/Dialect/ONNX/ONNXOps/Tensor/GridSample.cpp

@@ -29,6 +29,10 @@ LogicalResult ONNXGridSampleOpShapeHelper::computeShape() {
 
   // Read data and indices shapes as dim indices.
   ONNXGridSampleOpAdaptor operandAdaptor(operands);
+  if (!hasShapeAndRank(operandAdaptor.getX()) ||
+      !hasShapeAndRank(operandAdaptor.getGrid())) {
+    return failure();
+  }
   DimsExpr inputDims;
   DimsExpr gridDims;
   createIE->getShapeAsDims(operandAdaptor.getX(), inputDims);
@@ -78,11 +82,18 @@ LogicalResult ONNXGridSampleOp::verify() {
   if (!hasShapeAndRank(getOperation()))
     return success();
 
-  auto inputShape =
-      mlir::cast<ShapedType>(operandAdaptor.getX().getType()).getShape();
-  int64_t inputRank = inputShape.size();
-  auto gridShape =
-      mlir::cast<ShapedType>(operandAdaptor.getGrid().getType()).getShape();
+  auto inputType = mlir::cast<ShapedType>(operandAdaptor.getX().getType());
+  if (!inputType.hasStaticShape()) {
+    return success();
+  }
+  const auto inputShape = inputType.getShape();
+  const int64_t inputRank = inputShape.size();
+
+  auto gridType = mlir::cast<ShapedType>(operandAdaptor.getGrid().getType());
+  if (!gridType.hasStaticShape()) {
+    return success();
+  }
+  const auto gridShape = gridType.getShape();
 
   // Check whether the ranks of input and grid are valid and are equal.
   // Input's dimensions of rank r+2 should be in the form of (N,C,D1,D2,...,Dr)

test/mlir/onnx/onnx_shape_inference.mlir

@@ -4270,6 +4270,19 @@ func.func @test_grid_sample_same_dims(%arg0: tensor<1x3x1152x1344xf32>, %arg1: t
 // CHECK:         }
 }
 
+
+func.func @test_grid_sample_one_dynamic(%arg0: tensor<*xf32>, %arg1: tensor<1x1152x1344x2xf32>) -> tensor<*xf32> {
+  %0 = "onnx.GridSample"(%arg0, %arg1) {align_corners = 1 : si64, mode = "linear", onnx_node_name = "GridSample_181", padding_mode = "border"} : (tensor<*xf32>, tensor<1x1152x1344x2xf32>) -> tensor<*xf32>
+  return %0 : tensor<*xf32>
+
+// COM: Check that we do not crash
+// CHECK-LABEL:  func.func @test_grid_sample_one_dynamic
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<*xf32>, [[PARAM_1_:%.+]]: tensor<1x1152x1344x2xf32>) -> tensor<*xf32> {
+// CHECK:           [[VAR_0_:%.+]] = "onnx.GridSample"([[PARAM_0_]], [[PARAM_1_]]) {align_corners = 1 : si64, mode = "linear", onnx_node_name = "GridSample_181", padding_mode = "border"} : (tensor<*xf32>, tensor<1x1152x1344x2xf32>) -> tensor<*xf32>
+// CHECK:           return [[VAR_0_]] : tensor<*xf32>
+// CHECK:         }
+}
+
 func.func @test_grid_sample_diff_dims(%arg0: tensor<1x1x4x4xf32>, %arg1: tensor<1x6x6x2xf32>) -> tensor<*xf32> {
   %0 = "onnx.GridSample"(%arg0, %arg1) {align_corners = 1 : si64, mode = "linear", onnx_node_name = "GridSample_181", padding_mode = "border"} : (tensor<1x1x4x4xf32>, tensor<1x6x6x2xf32>) -> tensor<*xf32>
   return %0 : tensor<*xf32>