Browse Source

fix tbe handle not unregisted

pull/1640/head
lichun 4 years ago
parent
commit
6172a60eaa
1 changed files with 5 additions and 1 deletions
  1. +5
    -1
      tests/ut/ge/hybrid/ge_hybrid_unittest.cc

+ 5
- 1
tests/ut/ge/hybrid/ge_hybrid_unittest.cc View File

@@ -76,6 +76,10 @@ static ge::OpDescPtr CreateOpDesc(string name = "", string type = "") {
}

TEST_F(UtestGeHybrid, aicore_op_task_init_success) {
ComputeGraphPtr graph = std::make_shared<ComputeGraph>("test");
GeRootModelPtr ge_root_model = make_shared<GeRootModel>(graph);
HybridModel hybrid_model(ge_root_model);

// build aicore task
auto aicore_task = std::unique_ptr<hybrid::AiCoreOpTask>(new(std::nothrow)hybrid::AiCoreOpTask());
domi::TaskDef task_def;
@@ -99,7 +103,7 @@ TEST_F(UtestGeHybrid, aicore_op_task_init_success) {
op_desc->SetExtAttr(ge::OP_EXTATTR_NAME_TBE_KERNEL, tbe_kernel);
std::string kernel_name("kernel/Add");
AttrUtils::SetStr(op_desc, op_desc->GetName() + "_kernelname", kernel_name);
ASSERT_EQ(aicore_task->InitWithTaskDef(*op_desc.get(), task_def), SUCCESS);
ASSERT_EQ(aicore_task->InitWithTaskDef(hybrid_model, *op_desc.get(), task_def), SUCCESS);
rtStream_t stream = nullptr;
rtStreamCreate(&stream, 0);
ASSERT_EQ(aicore_task->LaunchKernel(stream), SUCCESS);


Loading…
Cancel
Save