You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.

op_task.cc 29 kB

5 years ago
5 years ago
4 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
4 years ago
5 years ago
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713
  1. /**
  2. * Copyright 2019-2020 Huawei Technologies Co., Ltd
  3. *
  4. * Licensed under the Apache License, Version 2.0 (the "License");
  5. * you may not use this file except in compliance with the License.
  6. * You may obtain a copy of the License at
  7. *
  8. * http://www.apache.org/licenses/LICENSE-2.0
  9. *
  10. * Unless required by applicable law or agreed to in writing, software
  11. * distributed under the License is distributed on an "AS IS" BASIS,
  12. * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. * See the License for the specific language governing permissions and
  14. * limitations under the License.
  15. */
  16. #include "single_op/task/op_task.h"
  17. #include <google/protobuf/extension_set.h>
  18. #include <chrono>
  19. #include <thread>
  20. #include "aicpu/common/aicpu_task_struct.h"
  21. #include "common/dump/dump_manager.h"
  22. #include "common/dump/dump_op.h"
  23. #include "common/formats/formats.h"
  24. #include "framework/common/debug/log.h"
  25. #include "register/op_tiling.h"
  26. #include "runtime/rt.h"
  27. namespace ge {
  28. namespace {
  29. constexpr int kLaunchRetryTimes = 1000;
  30. constexpr int kSleepTime = 10;
  31. constexpr uint64_t kReleaseFlag = 1;
  32. constexpr int kCopyNum = 2;
  33. void FreeHbm(void *var) {
  34. if (var) {
  35. (void)rtFree(var);
  36. }
  37. }
  38. }
  39. Status OpTask::OpenDump(const std::vector<uintptr_t> &io_addr, rtStream_t stream) {
  40. if (DumpManager::GetInstance().GetDumpProperties().IsSingleOpNeedDump()) {
  41. GELOGI("Dump is open in single op,start to set dump info");
  42. std::vector<uint64_t> input_addrs;
  43. std::vector<uint64_t> output_adds;
  44. auto input_size = op_desc_->GetInputsSize();
  45. auto output_size = op_desc_->GetOutputsSize();
  46. auto all_size = io_addr.size();
  47. if (input_size + output_size != all_size) {
  48. GELOGE(FAILED, "io_addr size is not equal input and output size");
  49. return FAILED;
  50. }
  51. for (size_t i = 0; i < input_size; i++) {
  52. uint64_t input_addr = static_cast<uint64_t>(io_addr[i]);
  53. input_addrs.emplace_back(input_addr);
  54. }
  55. for (size_t j = 0; j < output_size; j++) {
  56. uint64_t output_addr = static_cast<uint64_t>(io_addr[input_size + j]);
  57. output_adds.emplace_back(output_addr);
  58. }
  59. dump_op_.SetDumpInfo(DumpManager::GetInstance().GetDumpProperties(), op_desc_, input_addrs, output_adds, stream);
  60. auto status = dump_op_.LaunchDumpOp();
  61. if (status != SUCCESS) {
  62. GELOGE(status, "Launch dump op failed in single op");
  63. return status;
  64. }
  65. return SUCCESS;
  66. }
  67. GELOGI("Dump is not open in single op");
  68. return SUCCESS;
  69. }
  70. void TbeOpTask::SetStubFunc(const std::string &name, const void *stub_func) {
  71. this->stub_name_ = name;
  72. this->stub_func_ = stub_func;
  73. }
  74. void TbeOpTask::SetKernelArgs(std::unique_ptr<uint8_t[]> &&args, size_t arg_size, uint32_t block_dim,
  75. const OpDescPtr &op_desc) {
  76. args_ = std::move(args);
  77. arg_size_ = arg_size;
  78. block_dim_ = block_dim;
  79. op_desc_ = op_desc;
  80. }
  81. void TbeOpTask::SetSmDesc(void *sm_desc) { sm_desc_ = sm_desc; }
  82. const vector<int64_t> &OpTask::GetWorkspaceSizes() const { return workspace_sizes_; }
  83. void OpTask::SetWorkspaceSizes(const vector<int64_t> &workspace_sizes) { workspace_sizes_ = workspace_sizes; }
  84. TbeOpTask::~TbeOpTask() {
  85. if (sm_desc_ != nullptr) {
  86. (void)rtMemFreeManaged(sm_desc_);
  87. }
  88. if (tiling_buffer_ != nullptr) {
  89. (void)rtFree(tiling_buffer_);
  90. }
  91. }
  92. const void *TbeOpTask::GetArgs() const { return args_.get(); }
  93. size_t TbeOpTask::GetArgSize() const { return arg_size_; }
  94. const std::string &TbeOpTask::GetStubName() const { return stub_name_; }
  95. Status TbeOpTask::LaunchKernel(rtStream_t stream) {
  96. GELOGD("To invoke rtKernelLaunch. task = %s, block_dim = %u", this->stub_name_.c_str(), block_dim_);
  97. auto *sm_desc = reinterpret_cast<rtSmDesc_t *>(sm_desc_);
  98. auto ret = rtKernelLaunch(stub_func_, block_dim_, args_.get(), static_cast<uint32_t>(arg_size_), sm_desc, stream);
  99. int retry_times = 0;
  100. while (ret != RT_ERROR_NONE && retry_times < kLaunchRetryTimes) {
  101. retry_times++;
  102. GELOGW("Retry after %d ms, retry_times: %d", kSleepTime, retry_times);
  103. std::this_thread::sleep_for(std::chrono::milliseconds(kSleepTime));
  104. ret = rtKernelLaunch(stub_func_, block_dim_, args_.get(), arg_size_, sm_desc, stream);
  105. }
  106. if (ret != RT_ERROR_NONE) {
  107. GELOGE(RT_FAILED, "Invoke rtKernelLaunch failed. ret = %d, task = %s", ret, this->stub_name_.c_str());
  108. return RT_FAILED;
  109. }
  110. GELOGI("[TASK_INFO] %s", this->stub_name_.c_str());
  111. return SUCCESS;
  112. }
  113. Status TbeOpTask::UpdateRunInfo(const vector<GeTensorDesc> &input_desc, const vector<GeTensorDesc> &output_desc) {
  114. GE_CHK_STATUS_RET_NOLOG(UpdateNodeByShape(input_desc, output_desc));
  115. // invoke OpParaCalculate
  116. GELOGD("Start to invoke OpParaCalculate.");
  117. optiling::OpRunInfo run_info;
  118. run_info.block_dim = 0;
  119. auto ret = optiling::OpParaCalculate(*node_, run_info);
  120. if (ret != GRAPH_SUCCESS) {
  121. GELOGE(FAILED, "Failed to invoke OpParaCalculate. ret = %u", ret);
  122. return FAILED;
  123. }
  124. SetWorkspaceSizes(run_info.workspaces);
  125. block_dim_ = run_info.block_dim;
  126. tiling_data_ = run_info.tiling_data.str();
  127. GELOGD("Done invoking OpParaCalculate successfully. block_dim = %u, tiling size = %zu", block_dim_,
  128. tiling_data_.size());
  129. return SUCCESS;
  130. }
  131. Status TbeOpTask::UpdateTensorDesc(const GeTensorDesc &src_tensor, GeTensorDesc &dst_tensor) {
  132. int64_t storage_format_val = static_cast<Format>(FORMAT_RESERVED);
  133. (void)AttrUtils::GetInt(src_tensor, ge::ATTR_NAME_STORAGE_FORMAT, storage_format_val);
  134. auto storage_format = static_cast<Format>(storage_format_val);
  135. if (storage_format == FORMAT_RESERVED) {
  136. GELOGD("Storage format not set. update shape to [%s], and original shape to [%s]",
  137. src_tensor.GetShape().ToString().c_str(), src_tensor.GetOriginShape().ToString().c_str());
  138. dst_tensor.SetShape(src_tensor.GetShape());
  139. dst_tensor.SetOriginShape(src_tensor.GetOriginShape());
  140. } else {
  141. std::vector<int64_t> storage_shape;
  142. if (!AttrUtils::GetListInt(src_tensor, ge::ATTR_NAME_STORAGE_SHAPE, storage_shape)) {
  143. GELOGE(PARAM_INVALID, "Failed to get storage_shape while storage_format was set");
  144. return PARAM_INVALID;
  145. }
  146. GELOGD("Storage format set. update shape to [%s], and original shape to [%s]",
  147. GeShape(storage_shape).ToString().c_str(), src_tensor.GetShape().ToString().c_str());
  148. dst_tensor.SetShape(GeShape(std::move(storage_shape)));
  149. dst_tensor.SetOriginShape(src_tensor.GetShape());
  150. }
  151. return SUCCESS;
  152. }
  153. Status TbeOpTask::UpdateNodeByShape(const vector<GeTensorDesc> &input_desc, const vector<GeTensorDesc> &output_desc) {
  154. auto op_desc = node_->GetOpDesc();
  155. GE_CHECK_NOTNULL(op_desc);
  156. // Set runtime shape to node
  157. for (size_t i = 0; i < input_desc.size(); ++i) {
  158. auto tensor_desc = op_desc->MutableInputDesc(i);
  159. auto &runtime_tensor_desc = input_desc[i];
  160. GE_CHECK_NOTNULL(tensor_desc);
  161. GE_CHK_STATUS_RET(UpdateTensorDesc(runtime_tensor_desc, *tensor_desc));
  162. }
  163. for (size_t i = 0; i < output_desc.size(); ++i) {
  164. auto tensor_desc = op_desc->MutableOutputDesc(i);
  165. auto &runtime_tensor_desc = output_desc[i];
  166. GE_CHECK_NOTNULL(tensor_desc);
  167. GE_CHK_STATUS_RET(UpdateTensorDesc(runtime_tensor_desc, *tensor_desc));
  168. }
  169. return SUCCESS;
  170. }
  171. void TbeOpTask::EnableDynamicSupport(const NodePtr &node, void *tiling_buffer, size_t max_tiling_size) {
  172. node_ = node;
  173. tiling_buffer_ = tiling_buffer;
  174. max_tiling_size_ = max_tiling_size;
  175. }
  176. Status TbeOpTask::LaunchKernel(const vector<void *> &inputs, const vector<void *> &outputs,
  177. const vector<void *> &workspaces, rtStream_t stream) {
  178. GELOGD("[%s] Start to launch kernel", node_->GetName().c_str());
  179. std::vector<void *> args;
  180. args.insert(args.end(), inputs.begin(), inputs.end());
  181. args.insert(args.end(), outputs.begin(), outputs.end());
  182. args.insert(args.end(), workspaces.begin(), workspaces.end());
  183. if (tiling_buffer_ != nullptr) {
  184. GELOGD("[%s] Start to copy tiling info. size = %zu", node_->GetName().c_str(), tiling_data_.size());
  185. GE_CHK_RT_RET(rtMemcpyAsync(tiling_buffer_, max_tiling_size_, tiling_data_.data(), tiling_data_.size(),
  186. RT_MEMCPY_HOST_TO_DEVICE_EX, stream));
  187. args.emplace_back(tiling_buffer_);
  188. }
  189. if (memcpy_s(args_.get(), arg_size_, args.data(), args.size() * sizeof(void *)) != EOK) {
  190. GELOGE(INTERNAL_ERROR, "[%s] Failed to update kernel args.", node_->GetName().c_str());
  191. return INTERNAL_ERROR;
  192. }
  193. GELOGD("[%s] Start to invoke rtKernelLaunch", node_->GetName().c_str());
  194. GE_CHK_RT_RET(rtKernelLaunch(stub_func_, block_dim_, args_.get(), arg_size_, nullptr, stream));
  195. GELOGD("[%s] Done invoking rtKernelLaunch successfully", node_->GetName().c_str());
  196. return SUCCESS;
  197. }
  198. AiCpuBaseTask::~AiCpuBaseTask() {
  199. if (ext_info_addr_dev_ != nullptr) {
  200. (void)rtFree(ext_info_addr_dev_);
  201. }
  202. }
  203. Status AiCpuBaseTask::SetExtInfoAndType(const std::string &kernel_ext_info) {
  204. if (kernel_ext_info.empty()) {
  205. GELOGI("Kernel_ext_info is empty, no need copy to device.");
  206. return SUCCESS;
  207. }
  208. int32_t unknown_shape_type_val = 0;
  209. (void) AttrUtils::GetInt(op_desc_, ::ge::ATTR_NAME_UNKNOWN_SHAPE_TYPE, unknown_shape_type_val);
  210. GELOGD("Get unknown_type is %d.", unknown_shape_type_val);
  211. unknown_type_ = static_cast<UnknowShapeOpType>(unknown_shape_type_val);
  212. aicpu_ext_handle_.reset(new(std::nothrow) ::ge::hybrid::AicpuExtInfoHandler(op_desc_->GetName(),
  213. num_inputs_,
  214. num_outputs_,
  215. unknown_type_));
  216. GE_CHK_BOOL_RET_STATUS(aicpu_ext_handle_ != nullptr, FAILED, "Malloc aicpu_ext_handle mem failed!");
  217. Status ret = aicpu_ext_handle_->Parse(kernel_ext_info);
  218. if (ret != SUCCESS) {
  219. GELOGE(ret, "Parse kernel ext info failed, kernel_ext_info_size=%zu.", kernel_ext_info.size());
  220. return ret;
  221. }
  222. GE_CHK_RT_RET(rtMalloc(&ext_info_addr_dev_, kernel_ext_info.size(), RT_MEMORY_HBM));
  223. GE_CHK_RT_RET(rtMemcpy(ext_info_addr_dev_, kernel_ext_info.size(),
  224. kernel_ext_info.data(), kernel_ext_info.size(), RT_MEMCPY_HOST_TO_DEVICE));
  225. return SUCCESS;
  226. }
  227. Status AiCpuBaseTask::UpdateExtInfo(const std::vector<GeTensorDesc> &input_desc,
  228. std::vector<GeTensorDesc> &output_desc) {
  229. GELOGI("Update ext info begin, unknown_type=%d.", unknown_type_);
  230. if (num_inputs_ == 0 && num_outputs_ == 0) {
  231. GELOGI("No input and output, no need update ext info.");
  232. return SUCCESS;
  233. }
  234. GE_CHECK_NOTNULL(aicpu_ext_handle_);
  235. for (size_t i = 0; i < num_inputs_; ++i) {
  236. GE_CHK_STATUS_RET(aicpu_ext_handle_->UpdateInputShapeAndType(i, input_desc[i]),
  237. "Input[%zu] update input shape failed.", i);
  238. }
  239. if (unknown_type_ != DEPEND_COMPUTE) {
  240. for (size_t j = 0; j < num_outputs_; ++j) {
  241. GE_CHK_STATUS_RET(aicpu_ext_handle_->UpdateOutputShapeAndType(j, output_desc[j]),
  242. "Output[%zu] UpdateOutputShapeAndType failed.", j);
  243. }
  244. }
  245. GE_CHK_RT_RET(rtMemcpy(ext_info_addr_dev_,
  246. aicpu_ext_handle_->GetExtInfoLen(), // check size
  247. aicpu_ext_handle_->GetExtInfo(),
  248. aicpu_ext_handle_->GetExtInfoLen(),
  249. RT_MEMCPY_HOST_TO_DEVICE));
  250. GELOGI("Update ext info end.");
  251. return SUCCESS;
  252. }
  253. Status AiCpuBaseTask::UpdateOutputShape(vector<GeTensorDesc> &output_desc) {
  254. if (num_outputs_ == 0) {
  255. GELOGD("AiCpuBaseTask output_num is 0, no need update output shape.");
  256. return SUCCESS;
  257. }
  258. GELOGD("Start to update DEPEND_SHAPE_RANGE AiCpuBaseTask outputshape.");
  259. GE_CHK_RT_RET(rtMemcpy(aicpu_ext_handle_->GetExtInfo(),
  260. aicpu_ext_handle_->GetExtInfoLen(),
  261. ext_info_addr_dev_,
  262. aicpu_ext_handle_->GetExtInfoLen(),
  263. RT_MEMCPY_DEVICE_TO_HOST));
  264. for (size_t i = 0; i < num_outputs_; ++i) {
  265. GeShape shape;
  266. DataType data_type;
  267. aicpu_ext_handle_->GetOutputShapeAndType(i, shape, data_type);
  268. GE_CHK_STATUS_RET(UpdateShapeToOutputDesc(shape, output_desc[i]),
  269. "AiCpuCCTask Update [%zu]th output shape failed.", i);
  270. }
  271. GELOGD("Update DEPEND_SHAPE_RANGE AiCpuBaseTask outputshape finished.");
  272. return SUCCESS;
  273. }
  274. Status AiCpuBaseTask::UpdateShapeToOutputDesc(const GeShape &shape_new, GeTensorDesc &output_desc) {
  275. auto shape_old = output_desc.GetShape();
  276. output_desc.SetShape(shape_new);
  277. GELOGD("Update AiCpuBaseTask shape from %s to %s", shape_old.ToString().c_str(), shape_new.ToString().c_str());
  278. auto origin_shape_old = output_desc.GetOriginShape();
  279. auto origin_format = output_desc.GetOriginFormat();
  280. auto format = output_desc.GetFormat();
  281. if (origin_format == format) {
  282. output_desc.SetOriginShape(shape_new);
  283. return SUCCESS;
  284. }
  285. std::vector<int64_t> origin_dims_new;
  286. auto trans_ret = formats::TransShape(format, shape_new.GetDims(),
  287. output_desc.GetDataType(), origin_format, origin_dims_new);
  288. GE_CHK_STATUS_RET(trans_ret,
  289. "AiCpuTask originFormat[%d] is not same as format[%d], but TransShape failed, shape=%s.",
  290. origin_format, format, shape_new.ToString().c_str());
  291. auto origin_shape_new = GeShape(origin_dims_new);
  292. output_desc.SetOriginShape(origin_shape_new);
  293. GELOGD("AiCpuTask originFormat[%d] is not same as format[%d], need update from %s ro %s.",
  294. origin_format, format, origin_shape_old.ToString().c_str(), origin_shape_new.ToString().c_str());
  295. return SUCCESS;
  296. }
  297. AiCpuTask::~AiCpuTask() {
  298. FreeHbm(args_);
  299. FreeHbm(io_addr_);
  300. if (dynamic_flag_) {
  301. FreeHbm(workspace_addr_);
  302. }
  303. FreeHbm(copy_workspace_buf_);
  304. FreeHbm(copy_ioaddr_dev_);
  305. FreeHbm(copy_input_release_flag_dev_);
  306. FreeHbm(copy_input_data_size_dev_);
  307. FreeHbm(copy_input_src_dev_);
  308. FreeHbm(copy_input_dst_dev_);
  309. FreeHbm(copy_task_args_buf_);
  310. for (auto summary : output_summary_) {
  311. FreeHbm(summary);
  312. }
  313. for (auto out_shape : out_shape_hbm_) {
  314. FreeHbm(out_shape);
  315. }
  316. }
  317. const void *AiCpuTask::GetIOAddr() const { return io_addr_; }
  318. Status AiCpuTask::LaunchKernel(rtStream_t stream) {
  319. GELOGD("Start to launch kernel. task = %s", this->op_type_.c_str());
  320. auto ret = rtMemcpyAsync(workspace_addr_, task_info_.size(), task_info_.data(), task_info_.size(),
  321. RT_MEMCPY_HOST_TO_DEVICE_EX, stream);
  322. if (ret != RT_ERROR_NONE) {
  323. GELOGE(RT_FAILED, "rtMemcpyAsync workspace data failed. ret = %d, task = %s", ret, this->op_type_.c_str());
  324. return RT_FAILED;
  325. }
  326. GELOGI("To invoke rtKernelLaunchEx. task = %s", this->op_type_.c_str());
  327. ret = rtKernelLaunchEx(args_, arg_size_, 0, stream);
  328. if (ret != RT_ERROR_NONE) {
  329. GELOGE(RT_FAILED, "Invoke rtKernelLaunch failed. ret = %d, task = %s", ret, this->op_type_.c_str());
  330. return RT_FAILED;
  331. }
  332. GELOGI("[TASK_INFO] is %s", this->task_info_.c_str());
  333. GELOGD("Done launch kernel successfully. task = %s", this->op_type_.c_str());
  334. return SUCCESS;
  335. }
  336. Status AiCpuTask::PrepareCopyInputs(vector<DataBuffer> &outputs) {
  337. std::vector<uint64_t> copy_input_release_flag;
  338. std::vector<uint64_t> copy_input_data_size;
  339. std::vector<uint64_t> copy_input_src;
  340. std::vector<uint64_t> copy_input_dst;
  341. for (size_t i = 0; i < num_outputs_; ++i) {
  342. const auto &summary = output_summary_host_[i];
  343. GELOGI("Node out[%zu] summary, shape data=0x%lx, shape data size=%lu, raw data=0x%lx, raw data size=%lu.",
  344. i, summary.shape_data_ptr, summary.shape_data_size,
  345. summary.raw_data_ptr, summary.raw_data_size);
  346. auto output = outputs[i];
  347. copy_input_release_flag.emplace_back(kReleaseFlag);
  348. if (summary.raw_data_size > 0) {
  349. copy_input_data_size.emplace_back(output.length);
  350. } else {
  351. copy_input_data_size.emplace_back(summary.raw_data_size);
  352. }
  353. copy_input_src.emplace_back(summary.raw_data_ptr);
  354. copy_input_dst.emplace_back(reinterpret_cast<uintptr_t>(output.data));
  355. const auto &shape_buffer = out_shape_hbm_[i];
  356. copy_input_release_flag.emplace_back(kReleaseFlag);
  357. copy_input_data_size.emplace_back(summary.shape_data_size);
  358. copy_input_src.emplace_back(summary.shape_data_ptr);
  359. copy_input_dst.emplace_back(reinterpret_cast<uintptr_t>(shape_buffer));
  360. }
  361. const size_t copy_input_buf_len = num_outputs_ * kCopyNum * sizeof(uint64_t);
  362. GE_CHK_RT_RET(rtMemcpy(copy_input_release_flag_dev_, copy_input_buf_len,
  363. copy_input_release_flag.data(), copy_input_buf_len, RT_MEMCPY_HOST_TO_DEVICE));
  364. GE_CHK_RT_RET(rtMemcpy(copy_input_data_size_dev_, copy_input_buf_len,
  365. copy_input_data_size.data(), copy_input_buf_len, RT_MEMCPY_HOST_TO_DEVICE));
  366. GE_CHK_RT_RET(rtMemcpy(copy_input_src_dev_, copy_input_buf_len,
  367. copy_input_src.data(), copy_input_buf_len, RT_MEMCPY_HOST_TO_DEVICE));
  368. GE_CHK_RT_RET(rtMemcpy(copy_input_dst_dev_, copy_input_buf_len,
  369. copy_input_dst.data(), copy_input_buf_len, RT_MEMCPY_HOST_TO_DEVICE));
  370. return SUCCESS;
  371. }
  372. Status AiCpuTask::ReadResultSummaryAndPrepareMemory() {
  373. for (size_t i = 0; i < num_outputs_; ++i) {
  374. auto &result_summary = output_summary_host_[i];
  375. GE_CHK_RT_RET(rtMemcpy(&result_summary, sizeof(aicpu::FWKAdapter::ResultSummary),
  376. output_summary_[i], sizeof(aicpu::FWKAdapter::ResultSummary),
  377. RT_MEMCPY_DEVICE_TO_HOST));
  378. auto shape_data_size = result_summary.shape_data_size;
  379. void *shape_buffer = nullptr;
  380. if (shape_data_size > 0) {
  381. GE_CHK_RT_RET(rtMalloc(&shape_buffer, shape_data_size, RT_MEMORY_HBM));
  382. }
  383. out_shape_hbm_.emplace_back(shape_buffer);
  384. }
  385. return SUCCESS;
  386. }
  387. Status AiCpuTask::CopyDataToHbm(vector<DataBuffer> &outputs,
  388. rtStream_t stream) {
  389. GE_CHK_STATUS_RET_NOLOG(PrepareCopyInputs(outputs));
  390. GE_CHK_RT_RET(rtKernelLaunchEx(copy_task_args_buf_, sizeof(STR_FWK_OP_KERNEL),
  391. RT_KERNEL_DEFAULT, stream));
  392. GE_CHK_RT_RET(rtStreamSynchronize(stream));
  393. return SUCCESS;
  394. }
  395. Status AiCpuTask::UpdateShapeByHbmBuffer(vector<GeTensorDesc> &output_desc) {
  396. for (size_t i = 0; i < num_outputs_; ++i) {
  397. const auto &result_summary = output_summary_host_[i];
  398. std::vector<int64_t> shape_dims;
  399. if (result_summary.shape_data_size > 0) {
  400. const auto &shape_hbm = out_shape_hbm_[i];
  401. uint32_t dim_num = result_summary.shape_data_size / sizeof(int64_t);
  402. std::unique_ptr<int64_t[]> shape_addr(new(std::nothrow) int64_t[dim_num]());
  403. GE_CHECK_NOTNULL(shape_addr);
  404. GE_CHK_RT_RET(rtMemcpy(shape_addr.get(), result_summary.shape_data_size,
  405. shape_hbm, result_summary.shape_data_size, RT_MEMCPY_DEVICE_TO_HOST));
  406. for (uint32_t dim_idx = 0; dim_idx < dim_num; ++dim_idx) {
  407. shape_dims.emplace_back(shape_addr[dim_idx]);
  408. GELOGD("Node [%zu]th output dim[%u]=%ld.", i, dim_idx, shape_addr[dim_idx]);
  409. }
  410. }
  411. GE_CHK_STATUS_RET(UpdateShapeToOutputDesc(GeShape(shape_dims), output_desc[i]),
  412. "AiCpuTask update [%zu]th output shape failed.", i);
  413. }
  414. return SUCCESS;
  415. }
  416. Status AiCpuTask::UpdateShapeAndDataByResultSummary(vector<GeTensorDesc> &output_desc,
  417. vector<DataBuffer> &outputs,
  418. rtStream_t stream) {
  419. if (num_outputs_ == 0) {
  420. GELOGI("Output num is 0, there is no need to update the output and size.");
  421. return SUCCESS;
  422. }
  423. GELOGI("Update shape and data by result summary begin.");
  424. for (auto out_shape : out_shape_hbm_) {
  425. FreeHbm(out_shape);
  426. }
  427. out_shape_hbm_.clear();
  428. GE_CHK_STATUS_RET(ReadResultSummaryAndPrepareMemory(),
  429. "Read ResultSummary and update output shape failed.");
  430. GE_CHK_STATUS_RET(CopyDataToHbm(outputs, stream),
  431. "Copy data to output failed.");
  432. GE_CHK_STATUS_RET(UpdateShapeByHbmBuffer(output_desc),
  433. "Update shape by hbm buffer failed.");
  434. for (auto out_shape : out_shape_hbm_) {
  435. FreeHbm(out_shape);
  436. }
  437. out_shape_hbm_.clear();
  438. GELOGI("Update shape and data by result summary end.");
  439. return SUCCESS;
  440. }
  441. Status AiCpuTask::SetIO(const vector<void *> &inputs, vector<void *> &outputs) {
  442. vector<uint64_t> io_addrs;
  443. io_addrs.reserve(num_inputs_ + num_outputs_);
  444. for (size_t i = 0; i < num_inputs_; ++i) {
  445. GE_CHECK_NOTNULL(inputs[i]);
  446. GELOGD("AiCpuTask input[%zu] addr = %p", i, inputs[i]);
  447. io_addrs.emplace_back(reinterpret_cast<uintptr_t>(inputs[i]));
  448. }
  449. if (unknown_type_ != DEPEND_COMPUTE) {
  450. for (size_t i = 0; i < num_outputs_; ++i) {
  451. GE_CHECK_NOTNULL(outputs[i]);
  452. GELOGD("AiCpuTask output[%zu] addr = %p", i, outputs[i]);
  453. io_addrs.emplace_back(reinterpret_cast<uintptr_t>(outputs[i]));
  454. }
  455. } else {
  456. for (size_t i = 0; i < num_outputs_; ++i) {
  457. void *summary_addr = output_summary_[i];
  458. io_addrs.emplace_back(reinterpret_cast<uintptr_t>(summary_addr));
  459. }
  460. }
  461. if (!io_addrs.empty()) {
  462. auto *dst_io_addr = const_cast<uintptr_t *>(reinterpret_cast<const uintptr_t *>(io_addr_));
  463. GE_CHK_RT_RET(rtMemcpy(dst_io_addr,
  464. sizeof(uint64_t) * io_addrs.size(),
  465. &io_addrs[0],
  466. sizeof(uint64_t) * io_addrs.size(),
  467. RT_MEMCPY_HOST_TO_DEVICE));
  468. GE_CHECK_NOTNULL(dst_io_addr);
  469. };
  470. return SUCCESS;
  471. }
  472. Status AiCpuTask::InitForSummaryAndCopy() {
  473. if (unknown_type_ != DEPEND_COMPUTE || num_outputs_ == 0) {
  474. GELOGI("Unknown_type is %d, output num is %d.", unknown_type_, num_outputs_);
  475. return SUCCESS;
  476. }
  477. output_summary_.resize(num_outputs_);
  478. constexpr auto result_summary_size = sizeof(aicpu::FWKAdapter::ResultSummary);
  479. for (size_t i = 0; i < num_outputs_; ++i) {
  480. GE_CHK_RT_RET(rtMalloc(&output_summary_[i], result_summary_size, RT_MEMORY_HBM));
  481. }
  482. output_summary_host_.resize(num_outputs_);
  483. const size_t copy_input_buf_len = num_outputs_ * kCopyNum * sizeof(uint64_t);
  484. GE_CHK_RT_RET(rtMalloc(&copy_input_release_flag_dev_, copy_input_buf_len, RT_MEMORY_HBM));
  485. GE_CHK_RT_RET(rtMalloc(&copy_input_data_size_dev_, copy_input_buf_len, RT_MEMORY_HBM));
  486. GE_CHK_RT_RET(rtMalloc(&copy_input_src_dev_, copy_input_buf_len, RT_MEMORY_HBM));
  487. GE_CHK_RT_RET(rtMalloc(&copy_input_dst_dev_, copy_input_buf_len, RT_MEMORY_HBM));
  488. GE_CHK_RT_RET(rtMalloc(&copy_task_args_buf_, sizeof(STR_FWK_OP_KERNEL), RT_MEMORY_HBM));
  489. std::vector<uint64_t> copy_io_addr;
  490. copy_io_addr.emplace_back(reinterpret_cast<uintptr_t>(copy_input_release_flag_dev_));
  491. copy_io_addr.emplace_back(reinterpret_cast<uintptr_t>(copy_input_data_size_dev_));
  492. copy_io_addr.emplace_back(reinterpret_cast<uintptr_t>(copy_input_src_dev_));
  493. copy_io_addr.emplace_back(reinterpret_cast<uintptr_t>(copy_input_dst_dev_));
  494. const auto copy_io_addr_size = sizeof(uint64_t) * copy_io_addr.size();
  495. GE_CHK_RT_RET(rtMalloc(&copy_ioaddr_dev_, copy_io_addr_size, RT_MEMORY_HBM));
  496. GE_CHK_RT_RET(rtMemcpy(copy_ioaddr_dev_, copy_io_addr_size,
  497. copy_io_addr.data(), copy_io_addr_size, RT_MEMCPY_HOST_TO_DEVICE));
  498. return SUCCESS;
  499. }
  500. Status AiCpuTask::SetMemCopyTask(const domi::KernelExDef &kernel_def) {
  501. if (kernel_def.args_size() > sizeof(STR_FWK_OP_KERNEL)) {
  502. GELOGE(PARAM_INVALID, "sizeof STR_FWK_OP_KERNEL is: %lu, but args_size is: %d",
  503. sizeof(STR_FWK_OP_KERNEL), kernel_def.args_size());
  504. return PARAM_INVALID;
  505. }
  506. GE_CHK_RT_RET(rtMalloc(&copy_workspace_buf_, kernel_def.task_info_size(), RT_MEMORY_HBM));
  507. GE_CHK_RT_RET(rtMemcpy(copy_workspace_buf_, kernel_def.task_info_size(),
  508. kernel_def.task_info().data(), kernel_def.task_info_size(), RT_MEMCPY_HOST_TO_DEVICE));
  509. STR_FWK_OP_KERNEL aicpu_task = {0};
  510. auto sec_ret = memcpy_s(&aicpu_task, sizeof(STR_FWK_OP_KERNEL),
  511. kernel_def.args().data(), kernel_def.args().size());
  512. if (sec_ret != EOK) {
  513. GELOGE(FAILED, "memcpy failed, ret: %d", sec_ret);
  514. return FAILED;
  515. }
  516. aicpu_task.fwkKernelBase.fwk_kernel.inputOutputAddr = reinterpret_cast<uintptr_t>(copy_ioaddr_dev_);
  517. aicpu_task.fwkKernelBase.fwk_kernel.workspaceBaseAddr = reinterpret_cast<uintptr_t>(copy_workspace_buf_);
  518. aicpu_task.fwkKernelBase.fwk_kernel.extInfoAddr = 0;
  519. aicpu_task.fwkKernelBase.fwk_kernel.extInfoLen = 0;
  520. GE_CHK_RT_RET(rtMemcpy(copy_task_args_buf_, sizeof(STR_FWK_OP_KERNEL),
  521. &aicpu_task, sizeof(STR_FWK_OP_KERNEL), RT_MEMCPY_HOST_TO_DEVICE));
  522. return SUCCESS;
  523. }
  524. Status AiCpuTask::LaunchKernel(const std::vector<GeTensorDesc> &input_desc,
  525. const std::vector<DataBuffer> &input_buffers,
  526. std::vector<GeTensorDesc> &output_desc,
  527. std::vector<DataBuffer> &output_buffers,
  528. rtStream_t stream) {
  529. GE_CHK_STATUS_RET_NOLOG(UpdateExtInfo(input_desc, output_desc));
  530. std::vector<void *> inputs;
  531. std::vector<void *> outputs;
  532. for (auto &buffer : input_buffers) {
  533. inputs.emplace_back(buffer.data);
  534. }
  535. for (auto &buffer : output_buffers) {
  536. outputs.emplace_back(buffer.data);
  537. }
  538. GE_CHK_STATUS_RET_NOLOG(SetIO(inputs, outputs));
  539. GE_CHK_STATUS_RET_NOLOG(LaunchKernel(stream));
  540. GE_CHK_RT_RET(rtStreamSynchronize(stream));
  541. if (unknown_type_ == DEPEND_SHAPE_RANGE) {
  542. GE_CHK_STATUS_RET_NOLOG(UpdateOutputShape(output_desc));
  543. } else if (unknown_type_ == DEPEND_COMPUTE) {
  544. GE_CHK_STATUS_RET_NOLOG(UpdateShapeAndDataByResultSummary(output_desc, output_buffers, stream));
  545. }
  546. return SUCCESS;
  547. }
  548. void AiCpuCCTask::SetKernelArgs(std::unique_ptr<uint8_t[]> args, size_t arg_size) {
  549. args_ = std::move(args);
  550. arg_size_ = arg_size;
  551. // The blockdim value is defult "1" for rtCpuKernelLaunch
  552. block_dim_ = 1;
  553. }
  554. void AiCpuCCTask::SetSoName(const std::string &so_name) { so_name_ = so_name; }
  555. void AiCpuCCTask::SetkernelName(const std::string &kernel_Name) { kernel_name_ = kernel_Name; }
  556. void AiCpuCCTask::SetIoAddr(void *io_addr) { io_addr_ = io_addr; }
  557. const void *AiCpuCCTask::GetIOAddr() const { return io_addr_; }
  558. const void *AiCpuCCTask::GetArgs() const { return args_.get(); }
  559. size_t AiCpuCCTask::GetArgSize() const { return arg_size_; }
  560. AiCpuCCTask::~AiCpuCCTask() {
  561. }
  562. Status AiCpuCCTask::LaunchKernel(rtStream_t stream) {
  563. GELOGI("To invoke rtCpuKernelLaunch. block_dim = %u, so_name is %s, kernel_name is %s", block_dim_, so_name_.data(),
  564. kernel_name_.data());
  565. // sm_desc is nullptr, because l2 buffer does not support
  566. auto *sm_desc = reinterpret_cast<rtSmDesc_t *>(sm_desc_);
  567. auto ret = rtCpuKernelLaunchWithFlag(static_cast<const void *>(so_name_.data()),
  568. static_cast<const void *>(kernel_name_.data()),
  569. block_dim_, args_.get(), static_cast<uint32_t>(arg_size_),
  570. sm_desc, stream, dump_flag_);
  571. if (ret != RT_ERROR_NONE) {
  572. GELOGE(RT_FAILED, "Invoke rtCpuKernelLaunch failed. ret = %d", ret);
  573. return RT_FAILED;
  574. }
  575. GELOGD("Invoke rtCpuKernelLaunch succeeded");
  576. return SUCCESS;
  577. }
  578. Status AiCpuCCTask::LaunchKernel(const std::vector<GeTensorDesc> &input_desc,
  579. const std::vector<DataBuffer> &input_buffers,
  580. std::vector<GeTensorDesc> &output_desc,
  581. std::vector<DataBuffer> &output_buffers,
  582. rtStream_t stream) {
  583. GE_CHK_BOOL_RET_STATUS(unknown_type_ != DEPEND_COMPUTE, FAILED,
  584. "AiCpuCCTask unknown type[%d] is depend compute, it's not supported now.",
  585. unknown_type_);
  586. GE_CHK_STATUS_RET_NOLOG(UpdateExtInfo(input_desc, output_desc));
  587. size_t arg_index = 0;
  588. auto *task_io_addr = reinterpret_cast<uintptr_t *>(io_addr_);
  589. GE_CHECK_NOTNULL(task_io_addr);
  590. for (auto &input : input_buffers) {
  591. task_io_addr[arg_index++] = reinterpret_cast<uintptr_t>(input.data);
  592. }
  593. for (auto &output : output_buffers) {
  594. task_io_addr[arg_index++] = reinterpret_cast<uintptr_t>(output.data);
  595. }
  596. GE_CHK_STATUS_RET_NOLOG(LaunchKernel(stream));
  597. GE_CHK_RT_RET(rtStreamSynchronize(stream));
  598. if (unknown_type_ == DEPEND_SHAPE_RANGE) {
  599. GE_CHK_STATUS_RET_NOLOG(UpdateOutputShape(output_desc));
  600. }
  601. return SUCCESS;
  602. }
  603. } // namespace ge

图引擎模块(GE)是MindSpore的一个子模块,其代码由C++实现,位于前端模块ME和底层硬件之间,起到承接作用。图引擎模块以ME下发的图作为输入,然后进行一系列的深度图优化操作,最后输出一张可以在底层硬件上高效运行的图。GE针对昇腾AI处理器的硬件结构特点,做了特定的优化工作,以此来充分发挥出昇腾AI处理器的强大算力。在进行模型训练/推理时,GE会被自动调用而用户并不感知。GE主要由GE API和GE Core两部分组成,详细的架构图如下所示