未验证 提交 8e4c1936 编写于 作者: L Leo Chen 提交者: GitHub

[NPU] register npu finalize on exit (#32390)

* [NPU] register finalize on exit

* fix
上级 ab6f8745
......@@ -167,7 +167,7 @@ endif()
if (WITH_ASCEND_CL)
cc_test(range_op_npu_test SRCS range_op_npu_test.cc DEPS op_registry range_op scope device_context enforce executor)
cc_test(expand_op_npu_test SRCS expand_op_npu_test.cc DEPS op_registry expand_op scope device_context enforce executor compare_op)
cc_test(expand_op_npu_test SRCS expand_op_npu_test.cc DEPS op_registry expand_op eigen_cc_function scope device_context enforce executor compare_op)
endif()
set(GLOB_OP_LIB ${OP_LIBRARY} CACHE INTERNAL "Global OP library")
......
......@@ -190,6 +190,8 @@ void NPUMemcpySync(void *dst, const void *src, size_t count,
enum aclrtMemcpyKind kind, size_t dst_max_count) {
// NOTE(zhiqiu): The default max_count is count
dst_max_count = dst_max_count ? dst_max_count : count;
VLOG(4) << dst << " " << dst_max_count << " " << src << " " << count << " "
<< kind;
PADDLE_ENFORCE_NPU_SUCCESS(aclrtMemcpy(dst, dst_max_count, src, count, kind));
}
......
......@@ -2178,9 +2178,7 @@ All parameter, weight, gradient are variables in Paddle.
#ifdef PADDLE_WITH_ASCEND_CL
m.def("get_npu_device_count", platform::GetNPUDeviceCount);
m.def("_npu_finalize", []() {
platform::AclInstance::Instance().Finalize();
}); // private interface
m.def("npu_finalize", []() { platform::AclInstance::Instance().Finalize(); });
py::class_<platform::NPUProfConfigWrapper>(m, "NPUProfConfigWrapper");
......
......@@ -15,6 +15,7 @@
from __future__ import print_function
import os
import sys
import atexit
# The legacy core need to be removed before "import core",
# in case of users installing paddlepadde without -U option
......@@ -255,3 +256,8 @@ def __bootstrap__():
monkey_patch_variable()
__bootstrap__()
monkey_patch_varbase()
# NOTE(zhiqiu): register npu_finalize on the exit of Python,
# do some clean up manually.
if core.is_compiled_with_npu():
atexit.register(core.npu_finalize)
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册