From 4e3644c3a3a2dfd1a356dc01ef2b20b9b5cc7100 Mon Sep 17 00:00:00 2001
From: yaolun <yaolun2@hisilicon.com>
Date: Mon, 17 Jul 2023 22:04:45 +0800
Subject: [PATCH] =?UTF-8?q?=E5=88=AA=E9=99=A4=E5=86=97=E4=BD=99=E4=BB=A3?=
 =?UTF-8?q?=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 CMakeLists.txt                                    | 6 ++++++
 tf_adapter/kernels/aicpu/host_queue_dataset_op.cc | 4 ++--
 tf_adapter/tests/st/CMakeLists.txt                | 7 ++++++-
 tf_adapter/tests/ut/CMakeLists.txt                | 7 ++++++-
 4 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index d45b22369..1adfefbd7 100755
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -64,6 +64,12 @@ if (ENABLE_OPEN_SRC)
         ${CMAKE_CURRENT_LIST_DIR}/tf_adapter/optimizers/*.cc
         ${CMAKE_CURRENT_LIST_DIR}/tf_adapter/util/*.cc
     )
+    file(GLOB_RECURSE DISCARD_SOURCES
+        ${CMAKE_CURRENT_LIST_DIR}/tf_adapter/kernels/aicpu/map_dataset_op.cc
+        ${CMAKE_CURRENT_LIST_DIR}/tf_adapter/kernels/aicpu/map_and_batch_dataset_op.cc
+        ${CMAKE_CURRENT_LIST_DIR}/tf_adapter/kernels/aicpu/dataset_function.cc
+    )
+    list(REMOVE_ITEM SOURCES ${DISCARD_SOURCES})
     add_library(_tf_adapter SHARED ${SOURCES} ${CMAKE_BINARY_DIR}/dist/swig/ge_plugin_wrap.cxx)
 
     foreach (COMPILE_FLAG ${COMPILE_FLAGS})
diff --git a/tf_adapter/kernels/aicpu/host_queue_dataset_op.cc b/tf_adapter/kernels/aicpu/host_queue_dataset_op.cc
index 574abee93..d7f72079c 100644
--- a/tf_adapter/kernels/aicpu/host_queue_dataset_op.cc
+++ b/tf_adapter/kernels/aicpu/host_queue_dataset_op.cc
@@ -39,13 +39,13 @@
 #include "tf_adapter/util/util.h"
 #include "tf_adapter/util/npu_attrs.h"
 #include "tf_adapter/util/acl_channel.h"
+#include "tf_adapter/util/ge_plugin.h"
 #include "tf_adapter/util/host_queue.h"
 #include "tf_adapter/util/memory_pool.h"
 #include "tf_adapter/util/host_thread_pool.h"
 #include "tf_adapter/util/host_allocator.h"
 #include "tf_adapter/kernels/aicpu/data_item_deliver.h"
 #include "tf_adapter/kernels/aicpu/npu_tensor.h"
-#include "tf_adapter/kernels/aicpu/dataset_function.h"
 
 namespace tensorflow {
 namespace data {
@@ -1092,7 +1092,7 @@ class HostQueueDatasetOp : public DatasetOpKernel {
             TF_RETURN_IF_ERROR(data_deliver_->InitSocketServer());
           }
         }
-        TF_RETURN_IF_ERROR(DatasetFunction::RegisterNpuCancellation(
+        TF_RETURN_IF_ERROR(RegisterNpuCancellationCallback(
             [this]() {
               mutex_lock lck(mu_);
               finish_send_ = true;
diff --git a/tf_adapter/tests/st/CMakeLists.txt b/tf_adapter/tests/st/CMakeLists.txt
index 637cb6edc..16830c9a1 100644
--- a/tf_adapter/tests/st/CMakeLists.txt
+++ b/tf_adapter/tests/st/CMakeLists.txt
@@ -28,7 +28,12 @@ file(GLOB_RECURSE ST_SOURCES
     "kernels/testcase/*.cc"
     "util/testcase/*.cc"
 )
-
+file(GLOB_RECURSE DISCARD_ST_SOURCES
+    "kernels/testcase/dataset/map_dataset_st.cc"
+    "kernels/testcase/dataset/map_and_batch_dataset_st.cc"
+    "kernels/testcase/dataset/dataset_function_st.cc"
+)
+list(REMOVE_ITEM ST_SOURCES ${DISCARD_ST_SOURCES})
 add_executable(tfadapter_stest
                "main.cc"
                ${ST_SOURCES}
diff --git a/tf_adapter/tests/ut/CMakeLists.txt b/tf_adapter/tests/ut/CMakeLists.txt
index d40f64ae0..4caa091ee 100644
--- a/tf_adapter/tests/ut/CMakeLists.txt
+++ b/tf_adapter/tests/ut/CMakeLists.txt
@@ -29,7 +29,12 @@ file(GLOB_RECURSE UT_SOURCES
     "kernels/testcase/*.cc"
     "common/testcase/*.cc"
 )
-
+file(GLOB_RECURSE DISCARD_UT_SOURCES
+    "kernels/testcase/dataset/map_dataset_ut.cc"
+    "kernels/testcase/dataset/map_and_batch_dataset_ut.cc"
+    "kernels/testcase/dataset/dataset_function_ut.cc"
+)
+list(REMOVE_ITEM UT_SOURCES ${DISCARD_UT_SOURCES})
 add_executable(tfadapter_utest
                "main.cc"
                ${UT_SOURCES}
-- 
Gitee