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 61 : inline int32_t GetPrimaryFormat(int32_t format) { return static_cast<int32_t>(static_cast<uint32_t>(format) & 0xffU); }
112 :
113 67 : inline int32_t GetSubFormat(int32_t format)
114 : {
115 67 : return static_cast<int32_t>((static_cast<uint32_t>(format) & 0xffff00U) >> 8);
116 : }
117 :
118 : class DumpResourceSafeMap {
119 : public:
120 897 : static DumpResourceSafeMap& Instance()
121 : {
122 907 : static DumpResourceSafeMap instance;
123 897 : return instance;
124 : }
125 :
126 : static void WaitInterval(uint32_t intervalSec);
127 :
128 93 : void insert(const std::string key, std::shared_ptr<DumpStreamInfo> ptr)
129 : {
130 93 : std::lock_guard<std::mutex> lock(mtx_);
131 93 : resourceMap_[key] = std::move(ptr);
132 93 : }
133 :
134 48 : void remove(const std::string key)
135 : {
136 48 : std::lock_guard<std::mutex> lock(mtx_);
137 48 : resourceMap_.erase(key);
138 48 : }
139 :
140 60 : std::shared_ptr<DumpStreamInfo> get(const std::string key)
141 : {
142 60 : std::lock_guard<std::mutex> lock(mtx_);
143 60 : auto it = resourceMap_.find(key);
144 60 : if (it != resourceMap_.end()) {
145 24 : return it->second;
146 : }
147 36 : return nullptr;
148 60 : }
149 :
150 51 : int32_t size()
151 : {
152 51 : std::lock_guard<std::mutex> lock(mtx_);
153 102 : return resourceMap_.size();
154 51 : }
155 :
156 408 : void clear()
157 : {
158 408 : std::lock_guard<std::mutex> lock(mtx_);
159 408 : resourceMap_.clear();
160 408 : }
161 :
162 216 : void waitAndClear()
163 : {
164 : // 等待 resourceMap_ 中的数据被清理
165 : // 检查是否有待处理的 dump 任务，如果有则等待 5 秒让任务完成
166 216 : constexpr uint32_t WAIT_INTERVAL_SEC = 5U;
167 216 : constexpr uint32_t MAX_WAIT_COUNT = 12U; // 最多等待 60 秒 (5s * 12)
168 216 : uint32_t waitCount = 0U;
169 :
170 375 : while (waitCount < MAX_WAIT_COUNT) {
171 : {
172 363 : std::lock_guard<std::mutex> lock(mtx_);
173 363 : if (resourceMap_.empty()) {
174 204 : IDE_LOGI("resourceMap_ is empty, proceed to cleanup");
175 204 : break;
176 : }
177 159 : IDE_LOGI(
178 : "resourceMap_ has %zu items, wait %us for dump tasks to complete, count: %u/%u",
179 : resourceMap_.size(), WAIT_INTERVAL_SEC, waitCount + 1, MAX_WAIT_COUNT);
180 363 : }
181 :
182 : // map 不为空，等待 5 秒让任务完成
183 159 : WaitInterval(WAIT_INTERVAL_SEC);
184 159 : waitCount++;
185 : }
186 :
187 216 : if (waitCount >= MAX_WAIT_COUNT) {
188 12 : std::lock_guard<std::mutex> lock(mtx_);
189 12 : IDE_LOGW("Wait timeout, resourceMap_ still has %zu items", resourceMap_.size());
190 12 : }
191 :
192 : // 停止 cleanup 线程并清理资源
193 216 : StopCleanupThread();
194 : {
195 216 : std::lock_guard<std::mutex> lock(mtx_);
196 216 : resourceMap_.clear();
197 216 : }
198 216 : IDE_LOGI("All dump operations completed and resources cleared");
199 216 : }
200 :
201 : void EnqueueCleanup(const std::string key);
202 : bool IsCleanupThreadActive();
203 :
204 : private:
205 : void StartCleanupThread();
206 : void StopCleanupThread();
207 : void CleanupThreadLoop();
208 :
209 3 : DumpResourceSafeMap() : cleanupThreadActive_(false) {}
210 3 : ~DumpResourceSafeMap()
211 : {
212 3 : StopCleanupThread();
213 3 : std::lock_guard<std::mutex> lock(mtx_);
214 3 : resourceMap_.clear();
215 3 : }
216 : DumpResourceSafeMap(const DumpResourceSafeMap&) = delete;
217 : DumpResourceSafeMap& operator=(const DumpResourceSafeMap&) = delete;
218 :
219 : std::map<const std::string, std::shared_ptr<DumpStreamInfo>> resourceMap_;
220 : std::mutex mtx_;
221 :
222 : std::thread cleanupThread_;
223 : std::queue<std::string> cleanupQueue_;
224 : std::mutex cleanupMtx_;
225 : std::condition_variable cleanupCv_;
226 : std::atomic<bool> cleanupThreadActive_;
227 : };
228 : } // namespace Adx
229 : #endif
|
