diff --git a/ge/single_op/task/op_task.cc b/ge/single_op/task/op_task.cc index 5345f0f7..7117368e 100755 --- a/ge/single_op/task/op_task.cc +++ b/ge/single_op/task/op_task.cc @@ -761,7 +761,7 @@ Status AiCpuTask::UpdateShapeByHbmBuffer(vector &output_desc) { } -Status AiCpuTask::UpdateShapeAndDataByResultSummary(vector &output_desc, +Status AiCpuBaseTask::UpdateShapeAndDataByResultSummary(vector &output_desc, vector &outputs, rtStream_t stream) { if (num_outputs_ == 0) { @@ -951,6 +951,9 @@ Status AiCpuCCTask::LaunchKernel(const std::vector &input_desc, if (unknown_type_ == DEPEND_SHAPE_RANGE) { GE_CHK_RT_RET(rtStreamSynchronize(stream)); GE_CHK_STATUS_RET_NOLOG(UpdateOutputShape(output_desc)); + } else if (unknown_type_ == DEPEND_COMPUTE) { + GE_CHK_RT_RET(rtStreamSynchronize(stream)); + GE_CHK_STATUS_RET_NOLOG(UpdateShapeAndDataByResultSummary(output_desc, output_buffers, stream)); } return SUCCESS; diff --git a/ge/single_op/task/op_task.h b/ge/single_op/task/op_task.h index 84eb9d00..546b1681 100644 --- a/ge/single_op/task/op_task.h +++ b/ge/single_op/task/op_task.h @@ -140,7 +140,16 @@ class AiCpuBaseTask : public OpTask { rtStream_t stream); Status UpdateOutputShape(vector &output_desc); Status UpdateShapeToOutputDesc(const GeShape &shape_new, GeTensorDesc &output_desc); + + Status UpdateShapeAndDataByResultSummary(vector &output_desc, + vector &outputs, + rtStream_t stream); + Status ReadResultSummaryAndPrepareMemory(); + Status CopyDataToHbm(vector &outputs, rtStream_t stream); + Status PrepareCopyInputs(vector &outputs); + + Status UpdateShapeByHbmBuffer(vector &output_desc); protected: size_t num_inputs_ = 0; size_t num_outputs_ = 0; @@ -180,15 +189,6 @@ class AiCpuTask : public AiCpuBaseTask { private: // for copy task. Status InitForSummaryAndCopy(); - Status UpdateShapeAndDataByResultSummary(vector &output_desc, - vector &outputs, - rtStream_t stream); - Status ReadResultSummaryAndPrepareMemory(); - - Status CopyDataToHbm(vector &outputs, rtStream_t stream); - Status PrepareCopyInputs(vector &outputs); - - Status UpdateShapeByHbmBuffer(vector &output_desc); friend class AiCpuTaskBuilder; void *workspace_addr_ = nullptr;