saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Do not inline ExternOp's scalar tensor inputs (#3426)

Browse source 
An ExternOp's input needs buffers, so we cannot add compute_inline
schedule on it even if it's a scalar tensor. Instead, we need to
schedule it as compute_root.
This commit is contained in:
Yang Chen 2020-04-05 18:35:09 -07:00 committed by GitHub
parent 517693a507
commit d361121d98
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 31 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
onnxruntime/core/providers/nuphar/compiler/nuphar_schedule_builder.cc
View file

@ -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<tvm::ExternOpNode>() != nullptr) {
InsertRootSchedule(t, ctx_schedule);
} else {
TryInlineSchedule(t, ctx_schedule);
}
}
}
}

26
onnxruntime/test/providers/cpu/tensor/scatter_op_test.cc
View file

@ -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<float> input;
input.resize(3 * 3);
std::fill(input.begin(), input.end(), .0f);
test.AddInput<float>("data", {3, 3}, input);
test.AddInput<int64_t>("indices", {1, 2},
{1, 1}, /*is_initializer*/ true);
test.AddInput<float>("updates", {1, 2},
{2.0f, 2.0f});
test.AddOutput<float>("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