Relax dynamic restriction on ScatterND stablehlo conversion (#2772)

Signed-off-by: Sam <[email protected]>
Co-authored-by: Tung D. Le <[email protected]>

Author: srcarroll

src/Conversion/ONNXToStablehlo/Tensor/ScatterND.cpp

@@ -40,8 +40,9 @@ struct ONNXScatterNDOpLoweringToStablehlo
     auto indicesType = indices.getType().cast<ShapedType>();
     int64_t dataRank = dataType.getRank();
     int64_t indicesRank = indicesType.getRank();
-    assert(indicesType.hasStaticShape() &&
-           "only support indices with static shape");
+    if (indicesType.isDynamicDim(indicesRank - 1))
+      return rewriter.notifyMatchFailure(
+          op, "only support indices with static last dim");
     int64_t partialIdxDim = indicesType.getDimSize(indicesRank - 1);
 
     assert(dataRank >= 1 && "The rank of 'data' must be >= 1");

test/mlir/conversion/onnx_to_stablehlo/Tensor/ScatterND.mlir

@@ -27,3 +27,18 @@ func.func @test_scatternd_2(%arg0 : tensor<4x4x4xi32>, %arg1 : tensor<2x1xi64>,
 // CHECK:           return [[VAR_0_]] : tensor<4x4x4xi32>
 // CHECK:         }
 }
+
+// -----
+
+func.func @test_scatternd_dynamic(%arg0 : tensor<1x?x32x128xf32>, %arg1 : tensor<?x?x32x64x4xi64>, %arg2 : tensor<?x?x?x?xf32>) -> tensor<1x?x32x128xf32> {
+  %0 = "onnx.ScatterND"(%arg0, %arg1, %arg2) : (tensor<1x?x32x128xf32>, tensor<?x?x32x64x4xi64>, tensor<?x?x?x?xf32>) -> tensor<1x?x32x128xf32>
+  return %0 : tensor<1x?x32x128xf32>
+// CHECK-LABEL:  func.func @test_scatternd_dynamic
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<1x?x32x128xf32>, [[PARAM_1_:%.+]]: tensor<?x?x32x64x4xi64>, [[PARAM_2_:%.+]]: tensor<?x?x?x?xf32>) -> tensor<1x?x32x128xf32> {
+// CHECK:           [[VAR_0_:%.+]] = "stablehlo.scatter"([[PARAM_0_]], [[PARAM_1_]], [[PARAM_2_]]) ({
+// CHECK:           ^bb0([[arg3_:%.+]]: tensor<f32>, [[arg4_:%.+]]: tensor<f32>):
+// CHECK:             stablehlo.return [[arg4_]] : tensor<f32>
+// CHECK:           }) {indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0, 1, 2, 3], scatter_dims_to_operand_dims = [0, 1, 2, 3], index_vector_dim = 4>, unique_indices = false} : (tensor<1x?x32x128xf32>, tensor<?x?x32x64x4xi64>, tensor<?x?x?x?xf32>) -> tensor<1x?x32x128xf32>
+// CHECK:           return [[VAR_0_]] : tensor<1x?x32x128xf32>
+// CHECK:         }
+}