diff --git a/onnxruntime/core/providers/nuphar/compiler/nuphar_schedule_builder.cc b/onnxruntime/core/providers/nuphar/compiler/nuphar_schedule_builder.cc index 47e7030551..3684035942 100644 --- a/onnxruntime/core/providers/nuphar/compiler/nuphar_schedule_builder.cc +++ b/onnxruntime/core/providers/nuphar/compiler/nuphar_schedule_builder.cc @@ -50,7 +50,11 @@ static void Traverse(const tvm::Tensor& tensor, auto current_node = ctx_codegen.FindNode(t); Traverse(t, current_node, ctx_codegen, ctx_schedule); } else if (ctx_codegen.CheckLiteral(t->op->name)) { - TryInlineSchedule(t, ctx_schedule); + if (tensor->op.as() != nullptr) { + InsertRootSchedule(t, ctx_schedule); + } else { + TryInlineSchedule(t, ctx_schedule); + } } } } diff --git a/onnxruntime/test/providers/cpu/tensor/scatter_op_test.cc b/onnxruntime/test/providers/cpu/tensor/scatter_op_test.cc index 5828b275a9..e5d611adff 100644 --- a/onnxruntime/test/providers/cpu/tensor/scatter_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/scatter_op_test.cc @@ -224,5 +224,31 @@ TEST(Scatter, BoolInputWithAxis) { scatter_bool_with_axis_tests("ScatterElements", 11); } +static void scatter_same_updates_tests(const char* op_name, int op_version) { + OpTester test(op_name, op_version); + + std::vector input; + input.resize(3 * 3); + std::fill(input.begin(), input.end(), .0f); + test.AddInput("data", {3, 3}, input); + + test.AddInput("indices", {1, 2}, + {1, 1}, /*is_initializer*/ true); + + test.AddInput("updates", {1, 2}, + {2.0f, 2.0f}); + + test.AddOutput("y", {3, 3}, + {0.0f, 0.0f, 0.0f, + 2.0f, 2.0f, 0.0f, + 0.0f, 0.0f, 0.0f}); + test.Run(); +} + +TEST(Scatter, SameUpdateWithoutAxis) { + scatter_same_updates_tests("Scatter", 9); + scatter_same_updates_tests("ScatterElements", 11); +} + } // namespace test } // namespace onnxruntime