From e56280a8eb0676b1f7b10e919a064f6d7da12cdf Mon Sep 17 00:00:00 2001
From: TylunasLi <pwstudio@163.com>
Date: Mon, 29 Jul 2024 17:01:16 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E6=98=87=E8=85=BE=E5=8E=BB?=
 =?UTF-8?q?=E5=88=9D=E5=A7=8B=E5=8C=96=E6=93=8D=E4=BD=9C?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 example/benchmark/benchmark.cpp    | 1 +
 include/fastllm.h                  | 2 ++
 main.cpp                           | 2 +-
 src/devices/ascend/fastllm-acl.cpp | 4 +++-
 src/fastllm.cpp                    | 6 ++++++
 5 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/example/benchmark/benchmark.cpp b/example/benchmark/benchmark.cpp
index a09f2f7c..9077c180 100644
--- a/example/benchmark/benchmark.cpp
+++ b/example/benchmark/benchmark.cpp
@@ -198,5 +198,6 @@ int main(int argc, char **argv) {
     printf("prompt use %f s\n", promptSpend);
     printf("prompt speed = %f tokens / s\n", (float)promptTokenNum / promptSpend);
     printf("output %d tokens\nuse %f s\nspeed = %f tokens / s\n", tokens, spend, tokens / spend);
+    fastllm::Finalize();
     return 0;
 }
\ No newline at end of file
diff --git a/include/fastllm.h b/include/fastllm.h
index 02778ce5..8bac771d 100644
--- a/include/fastllm.h
+++ b/include/fastllm.h
@@ -508,6 +508,8 @@ namespace fastllm {
 
     void PrintProfiler();
 
+    void Finalize(); // 程序执行完，释放资源
+
     void ApplyDeviceMap(const std::map <std::string, int> &deviceMap, int current, int total); // 执行到了current, 一共total，使用deviceMap切换设备
 
     int LLMSamplingOnly(Data &logits, int outerOffset, const GenerationConfig &config);
diff --git a/main.cpp b/main.cpp
index 25c31718..7a8dc3cf 100644
--- a/main.cpp
+++ b/main.cpp
@@ -149,6 +149,6 @@ int main(int argc, char **argv) {
         }, generationConfig);
         messages.push_back(std::make_pair("assistant", ret));
     }
-
+    fastllm::Finalize();
     return 0;
 }
\ No newline at end of file
diff --git a/src/devices/ascend/fastllm-acl.cpp b/src/devices/ascend/fastllm-acl.cpp
index 2fe6b1e0..a849ed26 100644
--- a/src/devices/ascend/fastllm-acl.cpp
+++ b/src/devices/ascend/fastllm-acl.cpp
@@ -21,8 +21,10 @@ void showError(aclError state, char const* const message, const char* const file
            int const line) {
     if (ACL_SUCCESS != state) {
         printf("%s  AsecndCL error = %d at %s:%d\n", message, state, file, line);
-        if (state >= 500000)
+        if (state >= 500000) {
+            aclFinalize();
             exit(state);
+        }
     }  
 }
 
diff --git a/src/fastllm.cpp b/src/fastllm.cpp
index fdeb7139..b49aed04 100644
--- a/src/fastllm.cpp
+++ b/src/fastllm.cpp
@@ -3140,4 +3140,10 @@ namespace fastllm {
     std::map <std::string, int> GetDeviceMap() {
         return defaultDeviceMap;
     }
+
+    void Finalize() {
+#ifdef USE_ASCEND_NPU
+        npu::FastllmAclFinalize();
+#endif
+    }
 }