Heterogeneous cpu performance in GPU

This commit is contained in:
zhangyihui 2021-09-10 20:11:42 +08:00
parent 23d03a0368
commit e121fb2e8c
1 changed files with 7 additions and 0 deletions

View File

@ -31,6 +31,12 @@ void CpuDataSaver::WriteFile(const std::string out_path_dir) {
MS_LOG(INFO) << "No cpu operation detail infos to write.";
return;
}
#if ENABLE_GPU
auto context_ptr = MsContext::GetInstance();
MS_EXCEPTION_IF_NULL(context_ptr);
auto device_id = context_ptr->get_param<uint32_t>(MS_CTX_DEVICE_ID);
device_id_ = std::to_string(device_id);
#else
auto rank_id = common::GetEnv("RANK_ID");
// If RANK_ID is not set, default value is 0.
if (rank_id.empty()) {
@ -45,6 +51,7 @@ void CpuDataSaver::WriteFile(const std::string out_path_dir) {
}
}
device_id_ = rank_id;
#endif
op_side_ = "cpu";
WriteOpDetail(out_path_dir);
WriteOpType(out_path_dir);