Merge pull request #452 from Xilinx/jrickert.dynamic_shapes

Fix crashes realted to dynamic shapes for GridSample and Cast

Author: jorickert

src/Dialect/ONNX/ONNXOps/Canonicalize.cpp

@@ -126,10 +126,8 @@ SmallVector<Value, 4> castVariadicInput(PatternRewriter &rewriter, Location loc,
   SmallVector<Value, 4> castInputs;
   for (Value inp : inputs) {
     ShapedType inpType = mlir::cast<ShapedType>(inp.getType());
-    assert(inpType && "Type is not ShapedType");
-    ONNXCastOp castOp = rewriter.create<ONNXCastOp>(loc,
-        UnrankedTensorType::get(inpType.getElementType()), inp, saturate, to);
-    static_cast<void>(castOp.inferShapes([](Region &region) {}));
+    ONNXCastOp castOp = rewriter.create<ONNXCastOp>(
+        loc, inpType.clone(to.getValue()), inp, saturate, to);
     castInputs.emplace_back(castOp.getResult());
   }
   return castInputs;

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_canonicalization.mlir

@@ -85,6 +85,22 @@ func.func @cast_concat_swap(%arg0: tensor<1xi32>, %arg1: tensor<1xi32>) -> tenso
 
 // -----
 
+func.func @cast_concat_swap_dynamic(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>) -> tensor<*xi64> {
+  %0 = "onnx.Concat"(%arg0, %arg1) {axis = 0 : si64} : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
+  %1 = "onnx.Cast"(%0) {to = i64} : (tensor<*xi32>) -> tensor<*xi64>
+  onnx.Return %1 : tensor<*xi64>
+
+// CHECK-LABEL:  func.func @cast_concat_swap_dynamic
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<*xi32>, [[PARAM_1_:%.+]]: tensor<*xi32>) -> tensor<*xi64> {
+// CHECK-DAG:       [[VAR_0_:%.+]] = "onnx.Cast"([[PARAM_0_]]) {saturate = 1 : si64, to = i64} : (tensor<*xi32>) -> tensor<*xi64>
+// CHECK-DAG:       [[VAR_1_:%.+]] = "onnx.Cast"([[PARAM_1_]]) {saturate = 1 : si64, to = i64} : (tensor<*xi32>) -> tensor<*xi64>
+// CHECK:           [[VAR_2_:%.+]] = "onnx.Concat"([[VAR_0_]], [[VAR_1_]]) {axis = 0 : si64} : (tensor<*xi64>, tensor<*xi64>) -> tensor<*xi64>
+// CHECK:           onnx.Return [[VAR_2_]] : tensor<*xi64>
+// CHECK:         }
+}
+
+// -----
+
 func.func @cast_slice_swap(%arg0: tensor<3xi32>, %arg1: tensor<1xi64>, %arg2: tensor<1xi64>, %arg3: tensor<1xi64>, %arg4: tensor<1xi64>) -> tensor<1xi64> {
   %0 = "onnx.Slice"(%arg0, %arg1, %arg2, %arg3, %arg4) {axis = 0 : si64} : (tensor<3xi32>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> tensor<1xi32>
   %1 = "onnx.Cast"(%0) {to = i64} : (tensor<1xi32>) -> tensor<1xi64>

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>