|
|
@@ -161,6 +161,8 @@ class AiCpuBaseTask : public OpTask { |
|
|
|
Status PrepareCopyInputs(vector<DataBuffer> &outputs); |
|
|
|
|
|
|
|
Status UpdateShapeByHbmBuffer(vector<GeTensorDesc> &output_desc); |
|
|
|
|
|
|
|
virtual Status CopyDataToHbm(vector<DataBuffer> &outputs, rtStream_t stream) = 0; |
|
|
|
protected: |
|
|
|
size_t num_inputs_ = 0; |
|
|
|
size_t num_outputs_ = 0; |
|
|
@@ -201,7 +203,7 @@ class AiCpuTask : public AiCpuBaseTask { |
|
|
|
// for copy task. |
|
|
|
Status InitForSummaryAndCopy(); |
|
|
|
|
|
|
|
Status CopyDataToHbm(vector<DataBuffer> &outputs, rtStream_t stream); |
|
|
|
Status CopyDataToHbm(vector<DataBuffer> &outputs, rtStream_t stream) override; |
|
|
|
private: |
|
|
|
friend class AiCpuTaskBuilder; |
|
|
|
void *workspace_addr_ = nullptr; |
|
|
@@ -247,7 +249,7 @@ class AiCpuCCTask : public AiCpuBaseTask { |
|
|
|
private: |
|
|
|
Status InitForSummaryAndCopy(); |
|
|
|
|
|
|
|
Status CopyDataToHbm(vector<DataBuffer> &outputs, rtStream_t stream); |
|
|
|
Status CopyDataToHbm(vector<DataBuffer> &outputs, rtStream_t stream) override; |
|
|
|
private: |
|
|
|
friend class AiCpuCCTaskBuilder; |
|
|
|
std::string so_name_; |
|
|
@@ -266,7 +268,7 @@ private: |
|
|
|
std::string memcpy_so_name_; |
|
|
|
std::string memcpy_kernel_name_; |
|
|
|
// args size |
|
|
|
size_t memcpy_args_size_ = 0; |
|
|
|
size_t memcpy_args_size_ = 0; |
|
|
|
}; |
|
|
|
|
|
|
|
class MemcpyAsyncTask : public OpTask { |
|
|
|