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 35 kB

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
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
5 years ago
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873
  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 "common/math/math_util.h"
  25. #include "framework/common/debug/log.h"
  26. #include "register/op_tiling.h"
  27. #include "runtime/rt.h"
  28. #include "build_task_utils.h"
  29. namespace ge {
  30. namespace {
  31. constexpr int kLaunchRetryTimes = 1000;
  32. constexpr int kSleepTime = 10;
  33. constexpr uint64_t kReleaseFlag = 1;
  34. constexpr int kCopyNum = 2;
  35. void FreeHbm(void *var) {
  36. if (var) {
  37. (void)rtFree(var);
  38. }
  39. }
  40. } // namespace
  41. Status OpTask::OpenDump(rtStream_t stream) {
  42. if (DumpManager::GetInstance().GetDumpProperties().IsSingleOpNeedDump()) {
  43. GELOGI("Dump is open in single op,start to set dump info");
  44. std::vector<uint64_t> input_addrs;
  45. std::vector<uint64_t> output_adds;
  46. auto input_size = op_desc_->GetInputsSize();
  47. auto output_size = op_desc_->GetOutputsSize();
  48. uintptr_t *arg_base = nullptr;
  49. size_t arg_num = 0;
  50. GetIoAddr(arg_base, arg_num);
  51. if (arg_num < input_size + output_size) {
  52. GELOGE(FAILED, "io_addrs_for_dump_ size %zu is not equal input and output size %zu",
  53. arg_num,
  54. input_size + output_size);
  55. return FAILED;
  56. }
  57. for (size_t i = 0; i < input_size; i++) {
  58. uint64_t input_addr = arg_base[i];
  59. input_addrs.emplace_back(input_addr);
  60. }
  61. for (size_t j = 0; j < output_size; j++) {
  62. uint64_t output_addr = arg_base[input_size + j];
  63. output_adds.emplace_back(output_addr);
  64. }
  65. dump_op_.SetDumpInfo(DumpManager::GetInstance().GetDumpProperties(), op_desc_, input_addrs, output_adds, stream);
  66. auto status = dump_op_.LaunchDumpOp();
  67. if (status != SUCCESS) {
  68. GELOGE(status, "Launch dump op failed in single op");
  69. return status;
  70. }
  71. return SUCCESS;
  72. }
  73. GELOGI("Dump is not open in single op");
  74. return SUCCESS;
  75. }
  76. void TbeOpTask::SetStubFunc(const std::string &name, const void *stub_func) {
  77. this->stub_name_ = name;
  78. this->stub_func_ = stub_func;
  79. }
  80. void TbeOpTask::SetKernelArgs(std::unique_ptr<uint8_t[]> &&args, size_t arg_size, uint32_t block_dim,
  81. const OpDescPtr &op_desc) {
  82. args_ = std::move(args);
  83. arg_size_ = arg_size;
  84. block_dim_ = block_dim;
  85. op_desc_ = op_desc;
  86. }
  87. void TbeOpTask::SetSmDesc(void *sm_desc) { sm_desc_ = sm_desc; }
  88. void OpTask::SetModelArgs(std::string model_name, uint32_t model_id) {
  89. model_name_ = model_name;
  90. model_id_ = model_id;
  91. }
  92. Status OpTask::GetProfilingArgs(std::string &model_name, std::string &op_name, uint32_t &model_id,
  93. uint32_t &block_dim) {
  94. model_name = model_name_;
  95. model_id = model_id_;
  96. block_dim = block_dim_;
  97. GE_CHECK_NOTNULL(op_desc_);
  98. op_name = op_desc_->GetName();
  99. return SUCCESS;
  100. }
  101. Status OpTask::UpdateRunInfo(const vector<GeTensorDesc> &input_desc, const vector<GeTensorDesc> &output_desc) {
  102. return UNSUPPORTED;
  103. }
  104. Status OpTask::DoUpdateArgTable(const SingleOpModelParam &param, bool keep_workspace) {
  105. auto addresses = BuildTaskUtils::GetAddresses(op_desc_, param, keep_workspace);
  106. auto all_addresses = BuildTaskUtils::JoinAddresses(addresses);
  107. uintptr_t *arg_base = nullptr;
  108. size_t arg_num = 0;
  109. GetIoAddr(arg_base, arg_num);
  110. if (arg_num < all_addresses.size()) {
  111. GELOGE(INTERNAL_ERROR, "[%s] arg number mismatches, expect at least = %zu, but got = %zu",
  112. op_desc_->GetName().c_str(),
  113. all_addresses.size(),
  114. arg_num);
  115. return INTERNAL_ERROR;
  116. }
  117. for (void *addr : all_addresses) {
  118. *arg_base++ = reinterpret_cast<uintptr_t >(addr);
  119. }
  120. return SUCCESS;
  121. }
  122. Status OpTask::UpdateArgTable(const SingleOpModelParam &param) {
  123. return DoUpdateArgTable(param, true);
  124. }
  125. Status OpTask::LaunchKernel(const vector<GeTensorDesc> &input_desc,
  126. const vector<DataBuffer> &input_buffers,
  127. vector<GeTensorDesc> &output_desc,
  128. vector<DataBuffer> &output_buffers,
  129. rtStream_t stream) {
  130. return UNSUPPORTED;
  131. }
  132. TbeOpTask::~TbeOpTask() {
  133. if (sm_desc_ != nullptr) {
  134. (void)rtMemFreeManaged(sm_desc_);
  135. }
  136. if (tiling_buffer_ != nullptr) {
  137. (void)rtFree(tiling_buffer_);
  138. }
  139. }
  140. const void *TbeOpTask::GetArgs() const { return args_.get(); }
  141. size_t TbeOpTask::GetArgSize() const { return arg_size_; }
  142. const std::string &TbeOpTask::GetStubName() const { return stub_name_; }
  143. Status TbeOpTask::LaunchKernel(rtStream_t stream) {
  144. GELOGD("To invoke rtKernelLaunch. task = %s, block_dim = %u", this->stub_name_.c_str(), block_dim_);
  145. auto *sm_desc = reinterpret_cast<rtSmDesc_t *>(sm_desc_);
  146. auto ret = rtKernelLaunch(stub_func_, block_dim_, args_.get(), static_cast<uint32_t>(arg_size_), sm_desc, stream);
  147. int retry_times = 0;
  148. while (ret != RT_ERROR_NONE && retry_times < kLaunchRetryTimes) {
  149. retry_times++;
  150. GELOGW("Retry after %d ms, retry_times: %d", kSleepTime, retry_times);
  151. std::this_thread::sleep_for(std::chrono::milliseconds(kSleepTime));
  152. ret = rtKernelLaunch(stub_func_, block_dim_, args_.get(), arg_size_, sm_desc, stream);
  153. }
  154. if (ret != RT_ERROR_NONE) {
  155. GELOGE(RT_FAILED, "Invoke rtKernelLaunch failed. ret = %d, task = %s", ret, this->stub_name_.c_str());
  156. return RT_FAILED;
  157. }
  158. GELOGI("[TASK_INFO] %s", this->stub_name_.c_str());
  159. auto status = OpenDump(stream);
  160. if (status != SUCCESS) {
  161. GELOGE(status, "Open dump failed in the tbe single op %s", this->stub_name_.c_str());
  162. return status;
  163. }
  164. return SUCCESS;
  165. }
  166. Status TbeOpTask::UpdateRunInfo(const vector<GeTensorDesc> &input_desc, const vector<GeTensorDesc> &output_desc) {
  167. GE_CHK_STATUS_RET_NOLOG(UpdateNodeByShape(input_desc, output_desc));
  168. // invoke OpParaCalculate
  169. GELOGD("Start to invoke OpParaCalculate.");
  170. optiling::OpRunInfo run_info;
  171. run_info.block_dim = 0;
  172. auto ret = optiling::OpParaCalculate(*node_, run_info);
  173. if (ret != GRAPH_SUCCESS) {
  174. GELOGE(FAILED, "Failed to invoke OpParaCalculate. ret = %u", ret);
  175. return FAILED;
  176. }
  177. block_dim_ = run_info.block_dim;
  178. tiling_data_ = run_info.tiling_data.str();
  179. GELOGD("Done invoking OpParaCalculate successfully. block_dim = %u, tiling size = %zu", block_dim_,
  180. tiling_data_.size());
  181. GE_CHK_STATUS_RET(AllocateWorkspaces(run_info.workspaces), "Failed to allocate workspaces");
  182. return SUCCESS;
  183. }
  184. Status TbeOpTask::UpdateTensorDesc(const GeTensorDesc &src_tensor, GeTensorDesc &dst_tensor) {
  185. int64_t storage_format_val = static_cast<Format>(FORMAT_RESERVED);
  186. (void)AttrUtils::GetInt(src_tensor, ge::ATTR_NAME_STORAGE_FORMAT, storage_format_val);
  187. auto storage_format = static_cast<Format>(storage_format_val);
  188. if (storage_format == FORMAT_RESERVED) {
  189. GELOGD("Storage format not set. update shape to [%s], and original shape to [%s]",
  190. src_tensor.GetShape().ToString().c_str(), src_tensor.GetOriginShape().ToString().c_str());
  191. dst_tensor.SetShape(src_tensor.GetShape());
  192. dst_tensor.SetOriginShape(src_tensor.GetOriginShape());
  193. } else {
  194. std::vector<int64_t> storage_shape;
  195. if (!AttrUtils::GetListInt(src_tensor, ge::ATTR_NAME_STORAGE_SHAPE, storage_shape)) {
  196. GELOGE(PARAM_INVALID, "Failed to get storage_shape while storage_format was set");
  197. return PARAM_INVALID;
  198. }
  199. GELOGD("Storage format set. update shape to [%s], and original shape to [%s]",
  200. GeShape(storage_shape).ToString().c_str(), src_tensor.GetShape().ToString().c_str());
  201. dst_tensor.SetShape(GeShape(std::move(storage_shape)));
  202. dst_tensor.SetOriginShape(src_tensor.GetShape());
  203. }
  204. return SUCCESS;
  205. }
  206. Status TbeOpTask::UpdateNodeByShape(const vector<GeTensorDesc> &input_desc, const vector<GeTensorDesc> &output_desc) {
  207. auto op_desc = node_->GetOpDesc();
  208. GE_CHECK_NOTNULL(op_desc);
  209. // Set runtime shape to node
  210. for (size_t i = 0; i < input_desc.size(); ++i) {
  211. auto tensor_desc = op_desc->MutableInputDesc(i);
  212. auto &runtime_tensor_desc = input_desc[i];
  213. GE_CHECK_NOTNULL(tensor_desc);
  214. GE_CHK_STATUS_RET(UpdateTensorDesc(runtime_tensor_desc, *tensor_desc));
  215. }
  216. for (size_t i = 0; i < output_desc.size(); ++i) {
  217. auto tensor_desc = op_desc->MutableOutputDesc(i);
  218. auto &runtime_tensor_desc = output_desc[i];
  219. GE_CHECK_NOTNULL(tensor_desc);
  220. GE_CHK_STATUS_RET(UpdateTensorDesc(runtime_tensor_desc, *tensor_desc));
  221. }
  222. return SUCCESS;
  223. }
  224. void TbeOpTask::EnableDynamicSupport(const NodePtr &node, void *tiling_buffer, size_t max_tiling_size) {
  225. node_ = node;
  226. tiling_buffer_ = tiling_buffer;
  227. max_tiling_size_ = max_tiling_size;
  228. }
  229. Status TbeOpTask::AllocateWorkspaces(const vector<int64_t> &workspace_sizes) {
  230. static const std::string kPurpose("malloc workspace memory for dynamic op.");
  231. if (workspace_sizes.empty()) {
  232. GELOGD("No need to allocate workspace.");
  233. return SUCCESS;
  234. }
  235. int64_t total_size = 0;
  236. std::vector<int64_t> ws_offsets;
  237. for (auto ws_size : workspace_sizes) {
  238. // alignment and padding should be done in OpParaCalculate
  239. GE_CHK_STATUS_RET_NOLOG(CheckInt64AddOverflow(total_size, ws_size));
  240. ws_offsets.emplace_back(total_size);
  241. total_size += ws_size;
  242. }
  243. GELOGD("Total workspace size is %ld", total_size);
  244. GE_CHECK_NOTNULL(stream_resource_);
  245. auto ws_base = stream_resource_->MallocMemory(kPurpose, static_cast<size_t>(total_size));
  246. if (ws_base == nullptr) {
  247. GELOGE(ACL_ERROR_GE_MEMORY_ALLOCATION, "Failed to allocate memory of size: %ld", total_size);
  248. return ACL_ERROR_GE_MEMORY_ALLOCATION;
  249. }
  250. GELOGD("Done allocating workspace memory successfully.");
  251. for (auto ws_offset : ws_offsets) {
  252. workspaces_.emplace_back(ws_base + ws_offset);
  253. }
  254. return SUCCESS;
  255. }
  256. Status TbeOpTask::LaunchKernel(const vector<GeTensorDesc> &input_desc,
  257. const vector<DataBuffer> &input_buffers,
  258. vector<GeTensorDesc> &output_desc,
  259. vector<DataBuffer> &output_buffers,
  260. rtStream_t stream) {
  261. GE_CHK_STATUS_RET_NOLOG(UpdateRunInfo(input_desc, output_desc));
  262. GELOGD("[%s] Start to launch kernel", node_->GetName().c_str());
  263. std::vector<void *> args;
  264. for (auto &buffer : input_buffers) {
  265. args.emplace_back(buffer.data);
  266. }
  267. for (auto &buffer : output_buffers) {
  268. args.emplace_back(buffer.data);
  269. }
  270. for (auto &buffer : workspaces_) {
  271. args.emplace_back(buffer);
  272. }
  273. if (tiling_buffer_ != nullptr) {
  274. GELOGD("[%s] Start to copy tiling info. size = %zu", node_->GetName().c_str(), tiling_data_.size());
  275. GE_CHK_RT_RET(rtMemcpyAsync(tiling_buffer_, max_tiling_size_, tiling_data_.data(), tiling_data_.size(),
  276. RT_MEMCPY_HOST_TO_DEVICE_EX, stream));
  277. args.emplace_back(tiling_buffer_);
  278. }
  279. if (memcpy_s(args_.get(), arg_size_, args.data(), args.size() * sizeof(void *)) != EOK) {
  280. GELOGE(INTERNAL_ERROR, "[%s] Failed to update kernel args.", node_->GetName().c_str());
  281. return INTERNAL_ERROR;
  282. }
  283. GELOGD("[%s] Start to invoke rtKernelLaunch", node_->GetName().c_str());
  284. GE_CHK_RT_RET(rtKernelLaunch(stub_func_, block_dim_, args_.get(), arg_size_, nullptr, stream));
  285. GELOGD("[%s] Done invoking rtKernelLaunch successfully", node_->GetName().c_str());
  286. return SUCCESS;
  287. }
  288. void TbeOpTask::GetIoAddr(uintptr_t *&arg_base, size_t &arg_count) {
  289. arg_base = reinterpret_cast<uintptr_t *>(args_.get());
  290. arg_count = arg_size_ / sizeof(void *);
  291. if (tiling_buffer_ != nullptr) {
  292. --arg_count;
  293. }
  294. }
  295. AiCpuBaseTask::~AiCpuBaseTask() {
  296. if (ext_info_addr_dev_ != nullptr) {
  297. (void)rtFree(ext_info_addr_dev_);
  298. }
  299. }
  300. Status AiCpuBaseTask::SetExtInfoAndType(const std::string &kernel_ext_info, uint64_t kernel_id) {
  301. if (kernel_ext_info.empty()) {
  302. GELOGI("Kernel_ext_info is empty, no need copy to device.");
  303. return SUCCESS;
  304. }
  305. int32_t unknown_shape_type_val = 0;
  306. (void) AttrUtils::GetInt(op_desc_, ::ge::ATTR_NAME_UNKNOWN_SHAPE_TYPE, unknown_shape_type_val);
  307. GELOGD("Get unknown_type is %d.", unknown_shape_type_val);
  308. unknown_type_ = static_cast<UnknowShapeOpType>(unknown_shape_type_val);
  309. aicpu_ext_handle_.reset(new(std::nothrow) ::ge::hybrid::AicpuExtInfoHandler(op_desc_->GetName(),
  310. num_inputs_,
  311. num_outputs_,
  312. unknown_type_));
  313. GE_CHK_BOOL_RET_STATUS(aicpu_ext_handle_ != nullptr, FAILED, "Malloc aicpu_ext_handle mem failed!");
  314. Status ret = aicpu_ext_handle_->Parse(kernel_ext_info);
  315. if (ret != SUCCESS) {
  316. GELOGE(ret, "Parse kernel ext info failed, kernel_ext_info_size=%zu.", kernel_ext_info.size());
  317. return ret;
  318. }
  319. GE_CHK_STATUS_RET(aicpu_ext_handle_->UpdateSessionInfo(ULLONG_MAX, kernel_id, false),
  320. "UpdateSessionInfo failed.");
  321. GE_CHK_RT_RET(rtMalloc(&ext_info_addr_dev_, aicpu_ext_handle_->GetExtInfoLen(), RT_MEMORY_HBM));
  322. GE_CHK_RT_RET(rtMemcpy(ext_info_addr_dev_, aicpu_ext_handle_->GetExtInfoLen(),
  323. aicpu_ext_handle_->GetExtInfo(), aicpu_ext_handle_->GetExtInfoLen(),
  324. RT_MEMCPY_HOST_TO_DEVICE));
  325. return SUCCESS;
  326. }
  327. Status AiCpuBaseTask::SetInputConst() {
  328. input_is_const_.clear();
  329. const vector<bool> v_is_input_const = op_desc_->GetIsInputConst();
  330. for (size_t i = 0; i < op_desc_->GetAllInputsSize(); ++i) {
  331. const GeTensorDescPtr tensor_desc = op_desc_->MutableInputDesc(static_cast<uint32_t>(i));
  332. if (tensor_desc == nullptr) {
  333. GELOGD("SingleOp: %s, Index: %zu, has no input", op_desc_->GetName().c_str(), i);
  334. continue;
  335. }
  336. if (i < v_is_input_const.size() && v_is_input_const[i]) {
  337. GELOGD("SingleOp: %s, Index: %zu, input is const", op_desc_->GetName().c_str(), i);
  338. input_is_const_.push_back(true);
  339. continue;
  340. }
  341. input_is_const_.push_back(false);
  342. }
  343. return SUCCESS;
  344. }
  345. Status AiCpuBaseTask::UpdateExtInfo(const std::vector<GeTensorDesc> &input_desc,
  346. std::vector<GeTensorDesc> &output_desc,
  347. rtStream_t stream) {
  348. GELOGI("Update ext info begin, unknown_type=%d.", unknown_type_);
  349. if (num_inputs_ == 0 && num_outputs_ == 0) {
  350. GELOGI("No input and output, no need update ext info.");
  351. return SUCCESS;
  352. }
  353. GE_CHECK_NOTNULL(aicpu_ext_handle_);
  354. size_t non_const_index = 0;
  355. for (size_t input_index = 0; input_index < num_inputs_; input_index++) {
  356. if (input_index < input_is_const_.size() && input_is_const_[input_index]) {
  357. // get input_desc from op_desc if const input, num_inputs_ is op_desc_ input_size
  358. auto const_input_desc = op_desc_->MutableInputDesc(static_cast<uint32_t>(input_index));
  359. GE_CHECK_NOTNULL(const_input_desc);
  360. GE_CHK_STATUS_RET(aicpu_ext_handle_->UpdateInputShapeAndType(input_index, *const_input_desc),
  361. "Input[%zu] update input shape failed.", input_index);
  362. continue;
  363. }
  364. GE_CHK_BOOL_RET_STATUS(non_const_index < input_desc.size(), PARAM_INVALID,
  365. "Input_desc size is %zu, but get non_const_index is %zu",
  366. input_desc.size(), non_const_index);
  367. GE_CHK_STATUS_RET(aicpu_ext_handle_->UpdateInputShapeAndType(input_index, input_desc[non_const_index]),
  368. "Input[%zu] update input shape failed.", input_index);
  369. non_const_index++;
  370. }
  371. if (unknown_type_ != DEPEND_COMPUTE) {
  372. for (size_t j = 0; j < num_outputs_; ++j) {
  373. GE_CHK_STATUS_RET(aicpu_ext_handle_->UpdateOutputShapeAndType(j, output_desc[j]),
  374. "Output[%zu] UpdateOutputShapeAndType failed.", j);
  375. }
  376. }
  377. GE_CHK_RT_RET(rtMemcpyAsync(ext_info_addr_dev_,
  378. aicpu_ext_handle_->GetExtInfoLen(), // check size
  379. aicpu_ext_handle_->GetExtInfo(),
  380. aicpu_ext_handle_->GetExtInfoLen(),
  381. RT_MEMCPY_HOST_TO_DEVICE_EX,
  382. stream));
  383. GELOGI("Update ext info end.");
  384. return SUCCESS;
  385. }
  386. Status AiCpuBaseTask::UpdateOutputShape(vector<GeTensorDesc> &output_desc) {
  387. if (num_outputs_ == 0) {
  388. GELOGD("AiCpuBaseTask output_num is 0, no need update output shape.");
  389. return SUCCESS;
  390. }
  391. GELOGD("Start to update DEPEND_SHAPE_RANGE AiCpuBaseTask outputshape.");
  392. GE_CHK_RT_RET(rtMemcpy(aicpu_ext_handle_->GetExtInfo(),
  393. aicpu_ext_handle_->GetExtInfoLen(),
  394. ext_info_addr_dev_,
  395. aicpu_ext_handle_->GetExtInfoLen(),
  396. RT_MEMCPY_DEVICE_TO_HOST));
  397. for (size_t i = 0; i < num_outputs_; ++i) {
  398. GeShape shape;
  399. DataType data_type;
  400. aicpu_ext_handle_->GetOutputShapeAndType(i, shape, data_type);
  401. GE_CHK_STATUS_RET(UpdateShapeToOutputDesc(shape, output_desc[i]),
  402. "AiCpuCCTask Update [%zu]th output shape failed.", i);
  403. }
  404. GELOGD("Update DEPEND_SHAPE_RANGE AiCpuBaseTask outputshape finished.");
  405. return SUCCESS;
  406. }
  407. Status AiCpuBaseTask::UpdateShapeToOutputDesc(const GeShape &shape_new, GeTensorDesc &output_desc) {
  408. auto shape_old = output_desc.GetShape();
  409. output_desc.SetShape(shape_new);
  410. GELOGD("Update AiCpuBaseTask shape from %s to %s", shape_old.ToString().c_str(), shape_new.ToString().c_str());
  411. auto origin_shape_old = output_desc.GetOriginShape();
  412. auto origin_format = output_desc.GetOriginFormat();
  413. auto format = output_desc.GetFormat();
  414. if (origin_format == format) {
  415. output_desc.SetOriginShape(shape_new);
  416. return SUCCESS;
  417. }
  418. std::vector<int64_t> origin_dims_new;
  419. auto trans_ret = formats::TransShape(format, shape_new.GetDims(),
  420. output_desc.GetDataType(), origin_format, origin_dims_new);
  421. GE_CHK_STATUS_RET(trans_ret,
  422. "AiCpuTask originFormat[%d] is not same as format[%d], but TransShape failed, shape=%s.",
  423. origin_format, format, shape_new.ToString().c_str());
  424. auto origin_shape_new = GeShape(origin_dims_new);
  425. output_desc.SetOriginShape(origin_shape_new);
  426. GELOGD("AiCpuTask originFormat[%d] is not same as format[%d], need update from %s ro %s.",
  427. origin_format, format, origin_shape_old.ToString().c_str(), origin_shape_new.ToString().c_str());
  428. return SUCCESS;
  429. }
  430. Status AiCpuBaseTask::UpdateIoAddr(const vector<DataBuffer> &inputs, const vector<DataBuffer> &outputs) {
  431. uintptr_t *arg_base = nullptr;
  432. size_t arg_num = 0;
  433. GetIoAddr(arg_base, arg_num);
  434. // input number and output number was check in ValidateParams
  435. size_t non_const_index = 0;
  436. for (size_t input_index = 0; input_index < num_inputs_; input_index++) {
  437. if (input_index < input_is_const_.size() && input_is_const_[input_index]) {
  438. // const input no need update addr
  439. GE_CHECK_NOTNULL(arg_base);
  440. GELOGD("AICpuTask input[%zu] addr = %u", input_index, *arg_base);
  441. arg_base++;
  442. continue;
  443. }
  444. GE_CHK_BOOL_RET_STATUS(non_const_index < inputs.size(), PARAM_INVALID,
  445. "Input size is %zu, but get non_const_index is %zu",
  446. inputs.size(), non_const_index);
  447. auto addr = inputs[non_const_index].data;
  448. GE_CHECK_NOTNULL(addr);
  449. GELOGD("AICpuTask input[%zu] addr = %p", input_index, addr);
  450. *arg_base++ = reinterpret_cast<uintptr_t>(addr);
  451. non_const_index++;
  452. }
  453. for (size_t i = 0; i < outputs.size(); ++i) {
  454. auto addr = outputs[i].data;
  455. GE_CHECK_NOTNULL(addr);
  456. GELOGD("AICpuTask output[%zu] addr = %p", i, addr);
  457. *arg_base++ = reinterpret_cast<uintptr_t>(addr);
  458. }
  459. return SUCCESS;
  460. }
  461. AiCpuTask::~AiCpuTask() {
  462. FreeHbm(args_);
  463. FreeHbm(io_addr_);
  464. if (dynamic_flag_) {
  465. FreeHbm(workspace_addr_);
  466. }
  467. FreeHbm(copy_workspace_buf_);
  468. FreeHbm(copy_ioaddr_dev_);
  469. FreeHbm(copy_input_release_flag_dev_);
  470. FreeHbm(copy_input_data_size_dev_);
  471. FreeHbm(copy_input_src_dev_);
  472. FreeHbm(copy_input_dst_dev_);
  473. FreeHbm(copy_task_args_buf_);
  474. for (auto summary : output_summary_) {
  475. FreeHbm(summary);
  476. }
  477. for (auto out_shape : out_shape_hbm_) {
  478. FreeHbm(out_shape);
  479. }
  480. }
  481. Status AiCpuTask::LaunchKernel(rtStream_t stream) {
  482. GELOGD("Start to launch kernel. task = %s", this->op_type_.c_str());
  483. auto ret = rtMemcpyAsync(io_addr_,
  484. io_addr_size_,
  485. io_addr_host_.data(),
  486. io_addr_host_.size() * sizeof(void *),
  487. RT_MEMCPY_HOST_TO_DEVICE_EX,
  488. stream);
  489. if (ret != RT_ERROR_NONE) {
  490. GELOGE(RT_FAILED, "rtMemcpyAsync workspace data failed. ret = %d, task = %s", ret, this->op_type_.c_str());
  491. return RT_FAILED;
  492. }
  493. GELOGI("To invoke rtKernelLaunchEx. task = %s", this->op_type_.c_str());
  494. ret = rtKernelLaunchEx(args_, arg_size_, 0, stream);
  495. if (ret != RT_ERROR_NONE) {
  496. GELOGE(RT_FAILED, "Invoke rtKernelLaunch failed. ret = %d, task = %s", ret, this->op_type_.c_str());
  497. return RT_FAILED;
  498. }
  499. GELOGI("[TASK_INFO] %s/%s", std::to_string(kernel_id_).c_str(), op_type_.c_str());
  500. auto status = OpenDump(stream);
  501. if (status != SUCCESS) {
  502. GELOGE(status, "Open dump failed in aicpu single op %s", this->op_type_.c_str());
  503. return status;
  504. }
  505. GELOGD("Done launch kernel successfully. task = %s", this->op_type_.c_str());
  506. return SUCCESS;
  507. }
  508. Status AiCpuTask::PrepareCopyInputs(vector<DataBuffer> &outputs) {
  509. std::vector<uint64_t> copy_input_release_flag;
  510. std::vector<uint64_t> copy_input_data_size;
  511. std::vector<uint64_t> copy_input_src;
  512. std::vector<uint64_t> copy_input_dst;
  513. for (size_t i = 0; i < num_outputs_; ++i) {
  514. const auto &summary = output_summary_host_[i];
  515. GELOGI("Node out[%zu] summary, shape data=0x%lx, shape data size=%lu, raw data=0x%lx, raw data size=%lu.",
  516. i, summary.shape_data_ptr, summary.shape_data_size,
  517. summary.raw_data_ptr, summary.raw_data_size);
  518. auto output = outputs[i];
  519. copy_input_release_flag.emplace_back(kReleaseFlag);
  520. if (summary.raw_data_size > 0) {
  521. copy_input_data_size.emplace_back(output.length);
  522. } else {
  523. copy_input_data_size.emplace_back(summary.raw_data_size);
  524. }
  525. copy_input_src.emplace_back(summary.raw_data_ptr);
  526. copy_input_dst.emplace_back(reinterpret_cast<uintptr_t>(output.data));
  527. const auto &shape_buffer = out_shape_hbm_[i];
  528. copy_input_release_flag.emplace_back(kReleaseFlag);
  529. copy_input_data_size.emplace_back(summary.shape_data_size);
  530. copy_input_src.emplace_back(summary.shape_data_ptr);
  531. copy_input_dst.emplace_back(reinterpret_cast<uintptr_t>(shape_buffer));
  532. }
  533. const size_t copy_input_buf_len = num_outputs_ * kCopyNum * sizeof(uint64_t);
  534. GE_CHK_RT_RET(rtMemcpy(copy_input_release_flag_dev_, copy_input_buf_len,
  535. copy_input_release_flag.data(), copy_input_buf_len, RT_MEMCPY_HOST_TO_DEVICE));
  536. GE_CHK_RT_RET(rtMemcpy(copy_input_data_size_dev_, copy_input_buf_len,
  537. copy_input_data_size.data(), copy_input_buf_len, RT_MEMCPY_HOST_TO_DEVICE));
  538. GE_CHK_RT_RET(rtMemcpy(copy_input_src_dev_, copy_input_buf_len,
  539. copy_input_src.data(), copy_input_buf_len, RT_MEMCPY_HOST_TO_DEVICE));
  540. GE_CHK_RT_RET(rtMemcpy(copy_input_dst_dev_, copy_input_buf_len,
  541. copy_input_dst.data(), copy_input_buf_len, RT_MEMCPY_HOST_TO_DEVICE));
  542. return SUCCESS;
  543. }
  544. Status AiCpuTask::ReadResultSummaryAndPrepareMemory() {
  545. for (size_t i = 0; i < num_outputs_; ++i) {
  546. auto &result_summary = output_summary_host_[i];
  547. GE_CHK_RT_RET(rtMemcpy(&result_summary, sizeof(aicpu::FWKAdapter::ResultSummary),
  548. output_summary_[i], sizeof(aicpu::FWKAdapter::ResultSummary),
  549. RT_MEMCPY_DEVICE_TO_HOST));
  550. auto shape_data_size = result_summary.shape_data_size;
  551. void *shape_buffer = nullptr;
  552. if (shape_data_size > 0) {
  553. GE_CHK_RT_RET(rtMalloc(&shape_buffer, shape_data_size, RT_MEMORY_HBM));
  554. }
  555. out_shape_hbm_.emplace_back(shape_buffer);
  556. }
  557. return SUCCESS;
  558. }
  559. Status AiCpuTask::CopyDataToHbm(vector<DataBuffer> &outputs,
  560. rtStream_t stream) {
  561. GE_CHK_STATUS_RET_NOLOG(PrepareCopyInputs(outputs));
  562. GE_CHK_RT_RET(rtKernelLaunchEx(copy_task_args_buf_, sizeof(STR_FWK_OP_KERNEL),
  563. RT_KERNEL_DEFAULT, stream));
  564. GE_CHK_RT_RET(rtStreamSynchronize(stream));
  565. return SUCCESS;
  566. }
  567. Status AiCpuTask::UpdateShapeByHbmBuffer(vector<GeTensorDesc> &output_desc) {
  568. for (size_t i = 0; i < num_outputs_; ++i) {
  569. const auto &result_summary = output_summary_host_[i];
  570. std::vector<int64_t> shape_dims;
  571. if (result_summary.shape_data_size > 0) {
  572. const auto &shape_hbm = out_shape_hbm_[i];
  573. uint32_t dim_num = result_summary.shape_data_size / sizeof(int64_t);
  574. std::unique_ptr<int64_t[]> shape_addr(new(std::nothrow) int64_t[dim_num]());
  575. GE_CHECK_NOTNULL(shape_addr);
  576. GE_CHK_RT_RET(rtMemcpy(shape_addr.get(), result_summary.shape_data_size,
  577. shape_hbm, result_summary.shape_data_size, RT_MEMCPY_DEVICE_TO_HOST));
  578. for (uint32_t dim_idx = 0; dim_idx < dim_num; ++dim_idx) {
  579. shape_dims.emplace_back(shape_addr[dim_idx]);
  580. GELOGD("Node [%zu]th output dim[%u]=%ld.", i, dim_idx, shape_addr[dim_idx]);
  581. }
  582. }
  583. GE_CHK_STATUS_RET(UpdateShapeToOutputDesc(GeShape(shape_dims), output_desc[i]),
  584. "AiCpuTask update [%zu]th output shape failed.", i);
  585. }
  586. return SUCCESS;
  587. }
  588. Status AiCpuTask::UpdateShapeAndDataByResultSummary(vector<GeTensorDesc> &output_desc,
  589. vector<DataBuffer> &outputs,
  590. rtStream_t stream) {
  591. if (num_outputs_ == 0) {
  592. GELOGI("Output num is 0, there is no need to update the output and size.");
  593. return SUCCESS;
  594. }
  595. GELOGI("Update shape and data by result summary begin.");
  596. for (auto out_shape : out_shape_hbm_) {
  597. FreeHbm(out_shape);
  598. }
  599. out_shape_hbm_.clear();
  600. GE_CHK_STATUS_RET(ReadResultSummaryAndPrepareMemory(),
  601. "Read ResultSummary and update output shape failed.");
  602. GE_CHK_STATUS_RET(CopyDataToHbm(outputs, stream),
  603. "Copy data to output failed.");
  604. GE_CHK_STATUS_RET(UpdateShapeByHbmBuffer(output_desc),
  605. "Update shape by hbm buffer failed.");
  606. for (auto out_shape : out_shape_hbm_) {
  607. FreeHbm(out_shape);
  608. }
  609. out_shape_hbm_.clear();
  610. GELOGI("Update shape and data by result summary end.");
  611. return SUCCESS;
  612. }
  613. Status AiCpuTask::InitForSummaryAndCopy() {
  614. if (unknown_type_ != DEPEND_COMPUTE || num_outputs_ == 0) {
  615. GELOGI("Unknown_type is %d, output num is %d.", unknown_type_, num_outputs_);
  616. return SUCCESS;
  617. }
  618. output_summary_.resize(num_outputs_);
  619. constexpr auto result_summary_size = sizeof(aicpu::FWKAdapter::ResultSummary);
  620. for (size_t i = 0; i < num_outputs_; ++i) {
  621. GE_CHK_RT_RET(rtMalloc(&output_summary_[i], result_summary_size, RT_MEMORY_HBM));
  622. }
  623. output_summary_host_.resize(num_outputs_);
  624. const size_t copy_input_buf_len = num_outputs_ * kCopyNum * sizeof(uint64_t);
  625. GE_CHK_RT_RET(rtMalloc(&copy_input_release_flag_dev_, copy_input_buf_len, RT_MEMORY_HBM));
  626. GE_CHK_RT_RET(rtMalloc(&copy_input_data_size_dev_, copy_input_buf_len, RT_MEMORY_HBM));
  627. GE_CHK_RT_RET(rtMalloc(&copy_input_src_dev_, copy_input_buf_len, RT_MEMORY_HBM));
  628. GE_CHK_RT_RET(rtMalloc(&copy_input_dst_dev_, copy_input_buf_len, RT_MEMORY_HBM));
  629. GE_CHK_RT_RET(rtMalloc(&copy_task_args_buf_, sizeof(STR_FWK_OP_KERNEL), RT_MEMORY_HBM));
  630. std::vector<uint64_t> copy_io_addr;
  631. copy_io_addr.emplace_back(reinterpret_cast<uintptr_t>(copy_input_release_flag_dev_));
  632. copy_io_addr.emplace_back(reinterpret_cast<uintptr_t>(copy_input_data_size_dev_));
  633. copy_io_addr.emplace_back(reinterpret_cast<uintptr_t>(copy_input_src_dev_));
  634. copy_io_addr.emplace_back(reinterpret_cast<uintptr_t>(copy_input_dst_dev_));
  635. const auto copy_io_addr_size = sizeof(uint64_t) * copy_io_addr.size();
  636. GE_CHK_RT_RET(rtMalloc(&copy_ioaddr_dev_, copy_io_addr_size, RT_MEMORY_HBM));
  637. GE_CHK_RT_RET(rtMemcpy(copy_ioaddr_dev_, copy_io_addr_size,
  638. copy_io_addr.data(), copy_io_addr_size, RT_MEMCPY_HOST_TO_DEVICE));
  639. return SUCCESS;
  640. }
  641. Status AiCpuTask::SetMemCopyTask(const domi::KernelExDef &kernel_def) {
  642. if (kernel_def.args_size() > sizeof(STR_FWK_OP_KERNEL)) {
  643. GELOGE(PARAM_INVALID, "sizeof STR_FWK_OP_KERNEL is: %lu, but args_size is: %d",
  644. sizeof(STR_FWK_OP_KERNEL), kernel_def.args_size());
  645. return PARAM_INVALID;
  646. }
  647. GE_CHK_RT_RET(rtMalloc(&copy_workspace_buf_, kernel_def.task_info_size(), RT_MEMORY_HBM));
  648. GE_CHK_RT_RET(rtMemcpy(copy_workspace_buf_, kernel_def.task_info_size(),
  649. kernel_def.task_info().data(), kernel_def.task_info_size(), RT_MEMCPY_HOST_TO_DEVICE));
  650. STR_FWK_OP_KERNEL aicpu_task = {0};
  651. auto sec_ret = memcpy_s(&aicpu_task, sizeof(STR_FWK_OP_KERNEL),
  652. kernel_def.args().data(), kernel_def.args().size());
  653. if (sec_ret != EOK) {
  654. GELOGE(FAILED, "memcpy failed, ret: %d", sec_ret);
  655. return FAILED;
  656. }
  657. aicpu_task.fwkKernelBase.fwk_kernel.inputOutputAddr = reinterpret_cast<uintptr_t>(copy_ioaddr_dev_);
  658. aicpu_task.fwkKernelBase.fwk_kernel.workspaceBaseAddr = reinterpret_cast<uintptr_t>(copy_workspace_buf_);
  659. aicpu_task.fwkKernelBase.fwk_kernel.extInfoAddr = 0;
  660. aicpu_task.fwkKernelBase.fwk_kernel.extInfoLen = 0;
  661. GE_CHK_RT_RET(rtMemcpy(copy_task_args_buf_, sizeof(STR_FWK_OP_KERNEL),
  662. &aicpu_task, sizeof(STR_FWK_OP_KERNEL), RT_MEMCPY_HOST_TO_DEVICE));
  663. return SUCCESS;
  664. }
  665. Status AiCpuTask::LaunchKernel(const std::vector<GeTensorDesc> &input_desc,
  666. const std::vector<DataBuffer> &input_buffers,
  667. std::vector<GeTensorDesc> &output_desc,
  668. std::vector<DataBuffer> &output_buffers,
  669. rtStream_t stream) {
  670. GE_CHK_STATUS_RET_NOLOG(UpdateExtInfo(input_desc, output_desc, stream));
  671. if (unknown_type_ == DEPEND_COMPUTE) {
  672. std::vector<DataBuffer> summary_buffers;
  673. for (size_t i = 0; i < num_outputs_; ++i) {
  674. summary_buffers.emplace_back(output_summary_[i], sizeof(aicpu::FWKAdapter::ResultSummary), false);
  675. }
  676. GE_CHK_STATUS_RET_NOLOG(UpdateIoAddr(input_buffers, summary_buffers));
  677. } else {
  678. GE_CHK_STATUS_RET_NOLOG(UpdateIoAddr(input_buffers, output_buffers));
  679. }
  680. GE_CHK_STATUS_RET_NOLOG(LaunchKernel(stream));
  681. if (unknown_type_ == DEPEND_SHAPE_RANGE) {
  682. GE_CHK_RT_RET(rtStreamSynchronize(stream));
  683. GE_CHK_STATUS_RET_NOLOG(UpdateOutputShape(output_desc));
  684. } else if (unknown_type_ == DEPEND_COMPUTE) {
  685. GE_CHK_RT_RET(rtStreamSynchronize(stream));
  686. GE_CHK_STATUS_RET_NOLOG(UpdateShapeAndDataByResultSummary(output_desc, output_buffers, stream));
  687. }
  688. return SUCCESS;
  689. }
  690. Status AiCpuBaseTask::UpdateArgTable(const SingleOpModelParam &param) {
  691. // aicpu do not have workspace, for now
  692. return DoUpdateArgTable(param, false);
  693. }
  694. void AiCpuTask::GetIoAddr(uintptr_t *&arg_base, size_t &arg_count) {
  695. arg_base = reinterpret_cast<uintptr_t *>(io_addr_host_.data());
  696. arg_count = io_addr_host_.size();
  697. }
  698. void AiCpuCCTask::SetKernelArgs(std::unique_ptr<uint8_t[]> args, size_t arg_size) {
  699. args_ = std::move(args);
  700. arg_size_ = arg_size;
  701. // The blockdim value is defult "1" for rtCpuKernelLaunch
  702. block_dim_ = 1;
  703. }
  704. void AiCpuCCTask::SetSoName(const std::string &so_name) { so_name_ = so_name; }
  705. void AiCpuCCTask::SetkernelName(const std::string &kernel_Name) { kernel_name_ = kernel_Name; }
  706. void AiCpuCCTask::SetIoAddr(uintptr_t *io_addr) { io_addr_ = io_addr; }
  707. const void *AiCpuCCTask::GetArgs() const { return args_.get(); }
  708. size_t AiCpuCCTask::GetArgSize() const { return arg_size_; }
  709. AiCpuCCTask::~AiCpuCCTask() {
  710. }
  711. Status AiCpuCCTask::LaunchKernel(rtStream_t stream) {
  712. GELOGI("To invoke rtCpuKernelLaunch. block_dim = %u, so_name is %s, kernel_name is %s", block_dim_, so_name_.data(),
  713. kernel_name_.data());
  714. // sm_desc is nullptr, because l2 buffer does not support
  715. auto *sm_desc = reinterpret_cast<rtSmDesc_t *>(sm_desc_);
  716. auto ret = rtCpuKernelLaunchWithFlag(static_cast<const void *>(so_name_.data()),
  717. static_cast<const void *>(kernel_name_.data()),
  718. block_dim_, args_.get(), static_cast<uint32_t>(arg_size_),
  719. sm_desc, stream, dump_flag_);
  720. if (ret != RT_ERROR_NONE) {
  721. GELOGE(ret, "Invoke rtCpuKernelLaunch failed. ret = %d", ret);
  722. return ret;
  723. }
  724. GELOGD("Invoke rtCpuKernelLaunch succeeded");
  725. auto status = OpenDump(stream);
  726. if (status != SUCCESS) {
  727. GELOGE(status, "Open dump failed in the aicpucc single op %s", this->kernel_name_.c_str());
  728. return status;
  729. }
  730. return SUCCESS;
  731. }
  732. Status AiCpuCCTask::LaunchKernel(const std::vector<GeTensorDesc> &input_desc,
  733. const std::vector<DataBuffer> &input_buffers,
  734. std::vector<GeTensorDesc> &output_desc,
  735. std::vector<DataBuffer> &output_buffers,
  736. rtStream_t stream) {
  737. GE_CHK_STATUS_RET_NOLOG(UpdateExtInfo(input_desc, output_desc, stream));
  738. GE_CHK_STATUS_RET_NOLOG(UpdateIoAddr(input_buffers, output_buffers));
  739. GE_CHK_STATUS_RET_NOLOG(LaunchKernel(stream));
  740. if (unknown_type_ == DEPEND_SHAPE_RANGE) {
  741. GE_CHK_RT_RET(rtStreamSynchronize(stream));
  742. GE_CHK_STATUS_RET_NOLOG(UpdateOutputShape(output_desc));
  743. }
  744. return SUCCESS;
  745. }
  746. void AiCpuCCTask::GetIoAddr(uintptr_t *&arg_base, size_t &arg_count) {
  747. arg_base = io_addr_;
  748. arg_count = io_addr_num_;
  749. }
  750. } // namespace ge

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