Line data Source code
1 : /**
2 : * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3 : * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4 : * CANN Open Software License Agreement Version 2.0 (the "License").
5 : * Please refer to the License for details. You may not use this file except in compliance with the License.
6 : * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7 : * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8 : * See LICENSE in the root of the software repository for the full text of the License.
9 : */
10 : #ifndef DUMP_STREAM_INFO_H
11 : #define DUMP_STREAM_INFO_H
12 :
13 : #include <map>
14 : #include <mutex>
15 : #include <string>
16 : #include <iostream>
17 : #include <memory>
18 : #include <condition_variable>
19 : #include <chrono>
20 : #include <thread>
21 : #include <queue>
22 : #include <atomic>
23 : #include <sstream>
24 : #include <cstdint>
25 : #include <algorithm>
26 : #include <cstring>
27 : #include <vector>
28 : #include "acl/acl_base.h"
29 : #include "acl/acl_rt.h"
30 : #include "log/adx_log.h"
31 : #include "adump_pub.h"
32 : #include "dump_tensor.h"
33 : #include "proto/dump_task.pb.h"
34 : #include "dump_memory.h"
35 : #include "runtime/runtime/context.h"
36 : #include "runtime/runtime/event.h"
37 : #include "runtime/runtime/stream.h"
38 : #include "runtime/rts/rts_stream.h"
39 : #include "runtime/rts/rts_kernel.h"
40 :
41 : namespace Adx {
42 : constexpr uint32_t DUMP_SLICE_SIZE = 128 * 1024 * 1024; // 128MB
43 :
44 : struct ChunkContext {
45 : std::vector<char>& buffer;
46 : size_t& offset;
47 : const std::string& fileName;
48 : };
49 : typedef struct DumpStreamInfo {
50 : rtStream_t stm;
51 : rtEvent_t mainStmEvt;
52 : rtEvent_t dumpStmEvt;
53 : rtContext_t ctx{nullptr};
54 : std::vector<DumpTensor> inputTensors;
55 : std::vector<DumpTensor> outputTensors;
56 : std::string mainStreamKey;
57 : std::string opType;
58 : std::string opName;
59 : uint32_t dumpStmId;
60 : uint32_t dumpEvtId;
61 : uint32_t mainEvtId;
62 : uint32_t taskId;
63 : uint32_t streamId;
64 : uint32_t deviceId;
65 : uint32_t contextId;
66 : uint32_t threadId;
67 : uint64_t timestamp;
68 : uint64_t dumpNumber;
69 : std::string dumpPath;
70 : } DumpStreamInfo;
71 :
72 : typedef struct DumpInfoParams {
73 : std::string mainStreamKey;
74 : std::vector<DumpTensor> inputTensors;
75 : std::vector<DumpTensor> outputTensors;
76 : std::string opType;
77 : std::string opName;
78 : uint32_t streamId;
79 : uint32_t taskId;
80 : uint32_t deviceId;
81 : uint32_t contextId;
82 : uint32_t threadId;
83 : std::string dumpPath;
84 : } DumpInfoParams;
85 :
86 : int32_t DumpStreamCreate(DumpStreamInfo** ptr);
87 : void DumpStreamFree(DumpStreamInfo* ptr);
88 : uint64_t GetNextDumpNumber();
89 :
90 : std::string GenerateDumpFileName(const DumpStreamInfo* dumpInfoPtr);
91 : void FillTensorProtoInfo(const std::vector<DumpTensor>& tensors, toolkit::dump::DumpData& data, bool isInput);
92 : toolkit::dump::DumpData BuildDumpDataProto(const DumpStreamInfo* dumpInfoPtr);
93 : size_t CalculateTensorDataSize(const std::vector<DumpTensor>& tensors);
94 :
95 : int32_t DumpTensorPushToDumpQueue(
96 : void* dataBuf, uint32_t bufLen, const char* fileName, uint64_t offset, uint32_t isLastChunk);
97 : int32_t FlushCurrentChunk(ChunkContext& ctx, uint32_t isLastChunk);
98 : int32_t CopyTensorDataWithFlush(const DumpTensor& tensor, ChunkContext& ctx, bool isLastTensorForChunk);
99 : int32_t CopyTensorsWithChunking(const std::vector<DumpTensor>& tensors, ChunkContext& ctx, bool isLastTensorList);
100 : void DumpTensorToQueue(DumpStreamInfo* dumpInfoPtr);
101 :
102 : int32_t CollectStreamContextInfo(
103 : aclrtStream mainStream, const std::string& opName, const std::string& opType, uint32_t& streamId, uint32_t& taskId,
104 : uint32_t& deviceId, std::string& dumpPath);
105 : void DumpDataRecordInCaptureStream(void* fnArgs);
106 : int32_t SetupAsyncDump(
107 : std::shared_ptr<DumpStreamInfo> dumpInfoPtr, const std::string& opName, const std::string& opType,
108 : aclrtStream mainStream);
109 : int32_t GetDumpInfoFromMap(DumpInfoParams& params);
110 :
111 38 : inline int32_t GetPrimaryFormat(int32_t format)
112 : {
113 38 : return static_cast<int32_t>(static_cast<uint32_t>(format) & 0xffU);
114 : }
115 :
116 44 : inline int32_t GetSubFormat(int32_t format)
117 : {
118 44 : return static_cast<int32_t>((static_cast<uint32_t>(format) & 0xffff00U) >> 8);
119 : }
120 :
121 : class DumpResourceSafeMap {
122 : public:
123 897 : static DumpResourceSafeMap& Instance()
124 : {
125 907 : static DumpResourceSafeMap instance;
126 897 : return instance;
127 : }
128 :
129 : static void WaitInterval(uint32_t intervalSec);
130 :
131 93 : void insert(const std::string key, std::shared_ptr<DumpStreamInfo> ptr)
132 : {
133 93 : std::lock_guard<std::mutex> lock(mtx_);
134 93 : resourceMap_[key] = std::move(ptr);
135 93 : }
136 :
137 48 : void remove(const std::string key)
138 : {
139 48 : std::lock_guard<std::mutex> lock(mtx_);
140 48 : resourceMap_.erase(key);
141 48 : }
142 :
143 60 : std::shared_ptr<DumpStreamInfo> get(const std::string key)
144 : {
145 60 : std::lock_guard<std::mutex> lock(mtx_);
146 60 : auto it = resourceMap_.find(key);
147 60 : if (it != resourceMap_.end()) {
148 24 : return it->second;
149 : }
150 36 : return nullptr;
151 60 : }
152 :
153 51 : int32_t size()
154 : {
155 51 : std::lock_guard<std::mutex> lock(mtx_);
156 102 : return resourceMap_.size();
157 51 : }
158 :
159 408 : void clear()
160 : {
161 408 : std::lock_guard<std::mutex> lock(mtx_);
162 408 : resourceMap_.clear();
163 408 : }
164 :
165 216 : void waitAndClear()
166 : {
167 : // 等待 resourceMap_ 中的数据被清理
168 : // 检查是否有待处理的 dump 任务，如果有则等待 5 秒让任务完成
169 216 : constexpr uint32_t WAIT_INTERVAL_SEC = 5U;
170 216 : constexpr uint32_t MAX_WAIT_COUNT = 12U; // 最多等待 60 秒 (5s * 12)
171 216 : uint32_t waitCount = 0U;
172 :
173 372 : while (waitCount < MAX_WAIT_COUNT) {
174 : {
175 360 : std::lock_guard<std::mutex> lock(mtx_);
176 360 : if (resourceMap_.empty()) {
177 204 : IDE_LOGI("resourceMap_ is empty, proceed to cleanup");
178 204 : break;
179 : }
180 156 : IDE_LOGI(
181 : "resourceMap_ has %zu items, wait %us for dump tasks to complete, count: %u/%u",
182 : resourceMap_.size(), WAIT_INTERVAL_SEC, waitCount + 1, MAX_WAIT_COUNT);
183 360 : }
184 :
185 : // map 不为空，等待 5 秒让任务完成
186 156 : WaitInterval(WAIT_INTERVAL_SEC);
187 156 : waitCount++;
188 : }
189 :
190 216 : if (waitCount >= MAX_WAIT_COUNT) {
191 12 : std::lock_guard<std::mutex> lock(mtx_);
192 12 : IDE_LOGW("Wait timeout, resourceMap_ still has %zu items", resourceMap_.size());
193 12 : }
194 :
195 : // 停止 cleanup 线程并清理资源
196 216 : StopCleanupThread();
197 : {
198 216 : std::lock_guard<std::mutex> lock(mtx_);
199 216 : resourceMap_.clear();
200 216 : }
201 216 : IDE_LOGI("All dump operations completed and resources cleared");
202 216 : }
203 :
204 : void EnqueueCleanup(const std::string key);
205 : bool IsCleanupThreadActive();
206 :
207 : private:
208 : void StartCleanupThread();
209 : void StopCleanupThread();
210 : void CleanupThreadLoop();
211 :
212 3 : DumpResourceSafeMap() : cleanupThreadActive_(false)
213 3 : {}
214 3 : ~DumpResourceSafeMap()
215 : {
216 3 : StopCleanupThread();
217 3 : std::lock_guard<std::mutex> lock(mtx_);
218 3 : resourceMap_.clear();
219 3 : }
220 : DumpResourceSafeMap(const DumpResourceSafeMap&) = delete;
221 : DumpResourceSafeMap& operator=(const DumpResourceSafeMap&) = delete;
222 :
223 : std::map<const std::string, std::shared_ptr<DumpStreamInfo>> resourceMap_;
224 : std::mutex mtx_;
225 :
226 : std::thread cleanupThread_;
227 : std::queue<std::string> cleanupQueue_;
228 : std::mutex cleanupMtx_;
229 : std::condition_variable cleanupCv_;
230 : std::atomic<bool> cleanupThreadActive_;
231 : };
232 : } // namespace Adx
233 : #endif
|
