m53297601
/
mindspore
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

!3872 add internal output tensor

Browse Source 
Merge pull request !3872 from kisnwang/cache-internal-tensor
pull/3872/MERGE
mindspore-ci-bot 5 years ago committed by Gitee
parent def45cc69a 3ce0f33b27
commit a11e0e35f4
5 changed files with 212 additions and 135 deletions
Show Stats Download Patch File Download Diff File

68
mindspore/ccsrc/backend/session/kernel_graph.cc
Unescape View File

@ -961,18 +961,40 @@ void KernelGraph::PrintGraphExecuteOrder() const {
}
}
void KernelGraph::AddInternalOutput(const AnfNodePtr &front_node, const AnfNodePtr &node) {
void KernelGraph::AddInternalOutput(const AnfNodePtr &front_node, const AnfNodePtr &node, int output_idx,
bool unique_target) {
if (front_node == nullptr || node == nullptr) {
MS_LOG(INFO) << "Front node or node is nullptr";
return;
}
MS_LOG(INFO) << "Add internal node " << node->DebugString() << " with front node " << front_node->DebugString();
front_to_internal_outputs_map_[front_node] = node;
int output_idx = 0;
if (AnfAlgo::CheckPrimitiveType(front_node, prim::kPrimTupleGetItem)) {
output_idx = AnfAlgo::GetTupleGetItemOutIndex(front_node->cast<CNodePtr>());
}
internal_outputs_to_front_map_[node][output_idx] = front_node;
internal_outputs_to_front_map_[node][output_idx] = std::pair<AnfNodePtr, bool>(front_node, unique_target);
}
void KernelGraph::AddInternalOutputTensor(const AnfNodePtr &node, int output_idx, const tensor::TensorPtr &tensor) {
if (node == nullptr) {
return;
}
internal_outputs_tensor_map_[node][output_idx] = tensor;
}
tensor::TensorPtr KernelGraph::GetInternalOutputTensor(const AnfNodePtr &node, int output_idx) {
if (node == nullptr) {
return nullptr;
}
auto iter = internal_outputs_tensor_map_.find(node);
if (iter == internal_outputs_tensor_map_.end()) {
return nullptr;
}
auto idx_iter = iter->second.find(output_idx);
if (idx_iter == iter->second.end()) {
return nullptr;
}
return idx_iter->second;
}
void KernelGraph::ReplaceInternalOutput(const AnfNodePtr &node, const AnfNodePtr &new_node, int src_output_idx,
@ -996,7 +1018,7 @@ void KernelGraph::ReplaceInternalOutput(const AnfNodePtr &node, const AnfNodePtr
if (src_output_idx == -1) {
internal_outputs_to_front_map_[new_node] = front_nodes;
for (const auto &front_node_iter : front_nodes) {
front_to_internal_outputs_map_[front_node_iter.second] = new_node;
front_to_internal_outputs_map_[front_node_iter.second.first] = new_node;
}
internal_outputs_to_front_map_.erase(iter);
return;
@ -1008,9 +1030,9 @@ void KernelGraph::ReplaceInternalOutput(const AnfNodePtr &node, const AnfNodePtr
MS_LOG(INFO) << "The output " << src_output_idx << " of node " << node->DebugString() << " is not an internal node";
return;
}
auto front_node = front_node_iter->second;
internal_outputs_to_front_map_[new_node][dst_output_idx] = front_node;
front_to_internal_outputs_map_[front_node] = new_node;
auto front_node_pair = front_node_iter->second;
internal_outputs_to_front_map_[new_node][dst_output_idx] = front_node_pair;
front_to_internal_outputs_map_[front_node_pair.first] = new_node;
front_nodes.erase(index);
if (front_nodes.empty()) {
internal_outputs_to_front_map_.erase(iter);
@ -1027,16 +1049,30 @@ AnfNodePtr KernelGraph::GetInternalOutputByFrontNode(const AnfNodePtr &front_nod
bool KernelGraph::IsInternalOutput(const AnfNodePtr &node, int output_idx) const {
auto front_nodes_iter = internal_outputs_to_front_map_.find(node);
if (front_nodes_iter != internal_outputs_to_front_map_.end()) {
if (output_idx == -1) {
return true;
}
auto &front_nodes = front_nodes_iter->second;
if (front_nodes.find(output_idx) != front_nodes.end()) {
return true;
}
if (front_nodes_iter == internal_outputs_to_front_map_.end()) {
return false;
}
return false;
if (output_idx == -1) {
return true;
}
auto &front_nodes = front_nodes_iter->second;
if (front_nodes.find(output_idx) == front_nodes.end()) {
return false;
}
return true;
}
bool KernelGraph::IsUniqueTargetInternalOutput(const AnfNodePtr &node, int output_idx) const {
auto front_nodes_iter = internal_outputs_to_front_map_.find(node);
if (front_nodes_iter == internal_outputs_to_front_map_.end()) {
return false;
}
auto &front_nodes = front_nodes_iter->second;
auto idx_iter = front_nodes.find(output_idx);
if (idx_iter == front_nodes.end()) {
return false;
}
return idx_iter->second.second;
}
void KernelGraph::UpdateChildGraphOrder() {

10
mindspore/ccsrc/backend/session/kernel_graph.h
Unescape View File

@ -143,11 +143,16 @@ class KernelGraph : public FuncGraph {
void PrintGraphExecuteOrder() const;
const std::map<std::string, std::pair<AnfNodePtr, int>> &summary_nodes() const { return summary_nodes_; }
void set_summary_nodes(const std::map<std::string, std::pair<AnfNodePtr, int>> &nodes) { summary_nodes_ = nodes; }
void AddInternalOutput(const AnfNodePtr &front_node, const AnfNodePtr &node);
void AddInternalOutput(const AnfNodePtr &front_node, const AnfNodePtr &node, int output_idx = 0,
bool unique_target = false);
void ReplaceInternalOutput(const AnfNodePtr &node, const AnfNodePtr &new_node, int src_output_idx = -1,
int dst_output_idx = -1);
AnfNodePtr GetInternalOutputByFrontNode(const AnfNodePtr &front_node) const;
bool IsInternalOutput(const AnfNodePtr &node, int output_idx = -1) const;
bool IsUniqueTargetInternalOutput(const AnfNodePtr &node, int output_idx) const;
void AddInternalOutputTensor(const AnfNodePtr &node, int output_idx, const tensor::TensorPtr &tensor);
tensor::TensorPtr GetInternalOutputTensor(const AnfNodePtr &node, int output_idx);
uint32_t current_epoch() const { return current_epoch_; }
void set_current_epoch(uint32_t epoch) { current_epoch_ = epoch; }
void UpdateChildGraphOrder();
@ -217,7 +222,8 @@ class KernelGraph : public FuncGraph {
CNodePtr end_goto_;
bool null_output_;
std::unordered_map<AnfNodePtr, AnfNodePtr> front_to_internal_outputs_map_;
std::unordered_map<AnfNodePtr, std::unordered_map<int, AnfNodePtr>> internal_outputs_to_front_map_;
std::unordered_map<AnfNodePtr, std::unordered_map<int, std::pair<AnfNodePtr, bool>>> internal_outputs_to_front_map_;
std::unordered_map<AnfNodePtr, std::unordered_map<int, tensor::TensorPtr>> internal_outputs_tensor_map_;
uint32_t current_epoch_;
};
} // namespace session

200
mindspore/ccsrc/backend/session/session_basic.cc
View File

File diff suppressed because it is too large Load Diff

56
mindspore/ccsrc/runtime/device/cpu/cpu_kernel_runtime.cc
Unescape View File

@ -20,7 +20,7 @@
#include <numeric>
#include <utility>
#include <functional>
#include <unordered_map>
#include <map>
#include <set>
#include "backend/kernel_compiler/kernel.h"
#include "runtime/device/cpu/cpu_device_address.h"
@ -124,11 +124,10 @@ DeviceAddressPtr CPUKernelRuntime::CreateDeviceAddress(void *device_ptr, size_t
return std::make_shared<CPUDeviceAddress>(device_ptr, device_size, format, type_id);
}
tensor::TensorPtr CPUKernelRuntime::CreatTensorForOutput(const CNodePtr &node, size_t index,
std::set<DeviceAddressPtr> *bound_addresses,
tensor::TensorPtr CPUKernelRuntime::CreatTensorForOutput(session::KernelGraph *kernel_graph, const CNodePtr &node,
size_t index,
std::vector<tensor::TensorPtr> *need_sync_outputs) {
MS_EXCEPTION_IF_NULL(node);
MS_EXCEPTION_IF_NULL(bound_addresses);
MS_EXCEPTION_IF_NULL(need_sync_outputs);
size_t output_size = AnfAlgo::GetOutputTensorNum(node);
if (index >= output_size) {
@ -136,14 +135,21 @@ tensor::TensorPtr CPUKernelRuntime::CreatTensorForOutput(const CNodePtr &node, s
}
auto address = AnfAlgo::GetMutableOutputAddr(node, index);
MS_EXCEPTION_IF_NULL(address);
auto shape = AnfAlgo::GetOutputInferShape(node, index);
std::vector<int> temp_shape;
(void)temp_shape.insert(temp_shape.end(), shape.begin(), shape.end());
TypeId infer_type_id = AnfAlgo::GetOutputInferDataType(node, index);
TypeId device_type_id = AnfAlgo::GetOutputDeviceDataType(node, index);
tensor::TensorPtr tensor = std::make_shared<tensor::Tensor>(infer_type_id, temp_shape);
MS_EXCEPTION_IF_NULL(tensor);
if (bound_addresses->find(address) != bound_addresses->end()) {
tensor::TensorPtr tensor = kernel_graph->GetInternalOutputTensor(node, index);
if (tensor == nullptr) {
auto shape = AnfAlgo::GetOutputInferShape(node, index);
std::vector<int> temp_shape;
(void)temp_shape.insert(temp_shape.end(), shape.begin(), shape.end());
tensor = std::make_shared<tensor::Tensor>(infer_type_id, temp_shape);
bool is_internal_output = kernel_graph->IsInternalOutput(node, index);
if (is_internal_output) {
kernel_graph->AddInternalOutputTensor(node, index, tensor);
}
}
if (bound_addresses_.find(address) != bound_addresses_.end()) {
tensor->set_device_address(address);
need_sync_outputs->emplace_back(tensor);
} else {
@ -159,15 +165,14 @@ tensor::TensorPtr CPUKernelRuntime::CreatTensorForOutput(const CNodePtr &node, s
address->ptr_ = tensor->data_c();
}
address->ref_count_ = INIT_NODE_REF;
(void)bound_addresses->insert(address);
(void)bound_addresses_.insert(address);
}
tensor->set_dirty(false);
return tensor;
}
BaseRef CPUKernelRuntime::CreatTensorForOutput(const session::KernelWithIndex &kernel_with_index,
const std::unordered_map<AnfNode *, tensor::TensorPtr> &input_map,
std::set<DeviceAddressPtr> *bound_addresses,
BaseRef CPUKernelRuntime::CreatTensorForOutput(session::KernelGraph *kernel_graph,
const session::KernelWithIndex &kernel_with_index,
std::vector<tensor::TensorPtr> *need_sync_outputs) {
auto &input_node = kernel_with_index.first;
auto index = kernel_with_index.second;
@ -179,15 +184,15 @@ BaseRef CPUKernelRuntime::CreatTensorForOutput(const session::KernelWithIndex &k
VectorRef ret;
for (size_t i = 1; i < node->inputs().size(); i++) {
auto item_with_index = AnfAlgo::VisitKernelWithReturnType(node->input(i), 0);
auto out = CreatTensorForOutput(item_with_index, input_map, bound_addresses, need_sync_outputs);
auto out = CreatTensorForOutput(kernel_graph, item_with_index, need_sync_outputs);
ret.push_back(out);
}
return ret;
}
return CreatTensorForOutput(node, index, bound_addresses, need_sync_outputs);
return CreatTensorForOutput(kernel_graph, node, index, need_sync_outputs);
} else if (input_node->isa<Parameter>()) {
auto iter = input_map.find(input_node.get());
if (iter != input_map.end()) {
auto iter = input_param_tensor_map_.find(input_node);
if (iter != input_param_tensor_map_.end()) {
return iter->second;
}
} else if (input_node->isa<ValueNode>()) {
@ -197,10 +202,8 @@ BaseRef CPUKernelRuntime::CreatTensorForOutput(const session::KernelWithIndex &k
}
return BaseRef();
}
void CPUKernelRuntime::BindInputOutput(const session::KernelGraph *kernel_graph,
const std::vector<tensor::TensorPtr> &inputs, VectorRef *outputs,
std::vector<tensor::TensorPtr> *need_sync_outputs) {
void CPUKernelRuntime::BindInputOutput(session::KernelGraph *kernel_graph, const std::vector<tensor::TensorPtr> &inputs,
VectorRef *outputs, std::vector<tensor::TensorPtr> *need_sync_outputs) {
MS_EXCEPTION_IF_NULL(kernel_graph);
MS_EXCEPTION_IF_NULL(outputs);
// bind input ptr
@ -208,11 +211,11 @@ void CPUKernelRuntime::BindInputOutput(const session::KernelGraph *kernel_graph,
if (input_nodes.size() != inputs.size()) {
MS_LOG(EXCEPTION) << "Input size not equal to input node size!";
}
std::unordered_map<AnfNode *, tensor::TensorPtr> input_map;
input_param_tensor_map_.clear();
size_t input_idx = 0;
for (auto &item : input_nodes) {
MS_EXCEPTION_IF_NULL(item);
input_map[item.get()] = inputs[input_idx];
input_param_tensor_map_[item] = inputs[input_idx];
if (item->isa<Parameter>()) {
auto address = AnfAlgo::GetMutableOutputAddr(item, 0);
auto tensor = inputs[input_idx];
@ -222,7 +225,6 @@ void CPUKernelRuntime::BindInputOutput(const session::KernelGraph *kernel_graph,
if (tensor_address != nullptr && tensor_address != address) {
(void)tensor->data_sync();
}
if (tensor->data_type() == address->type_id_ || tensor->data_type() == kNumberTypeFloat32 ||
tensor->data_type() == kNumberTypeInt32) {
address->ptr_ = tensor->data_c();
@ -243,11 +245,11 @@ void CPUKernelRuntime::BindInputOutput(const session::KernelGraph *kernel_graph,
input_idx++;
}
// new output and bind ptr
std::set<DeviceAddressPtr> bound_addresses;
bound_addresses_.clear();
auto output_nodes = kernel_graph->outputs();
for (const auto &item : output_nodes) {
auto item_with_index = AnfAlgo::VisitKernelWithReturnType(item, 0, true);
auto out = CreatTensorForOutput(item_with_index, input_map, &bound_addresses, need_sync_outputs);
auto out = CreatTensorForOutput(kernel_graph, item_with_index, need_sync_outputs);
outputs->push_back(std::move(out));
}
}

13
mindspore/ccsrc/runtime/device/cpu/cpu_kernel_runtime.h
Unescape View File

@ -19,7 +19,7 @@
#include <memory>
#include <vector>
#include <string>
#include <unordered_map>
#include <map>
#include <set>
#include "runtime/device/kernel_runtime.h"
#include "backend/session/kernel_graph.h"
@ -38,7 +38,7 @@ class CPUKernelRuntime : public KernelRuntime {
bool Init() override { return true; }
bool Run(session::KernelGraph *graph, Debugger *debugger = nullptr) override;
void AssignKernelAddress(session::KernelGraph *kernel_graph);
void BindInputOutput(const session::KernelGraph *kernel_graph, const std::vector<tensor::TensorPtr> &inputs,
void BindInputOutput(session::KernelGraph *kernel_graph, const std::vector<tensor::TensorPtr> &inputs,
VectorRef *outputs, std::vector<tensor::TensorPtr> *need_sync_outputs);
void IncreaseSummaryRefCount(const session::NamedSummaryOutputs &summary_outputs);
void DecreaseSummaryRefCount(const session::NamedSummaryOutputs &summary_outputs);
@ -49,19 +49,18 @@ class CPUKernelRuntime : public KernelRuntime {
TypeId type_id) override;
private:
tensor::TensorPtr CreatTensorForOutput(const CNodePtr &node, size_t index,
std::set<DeviceAddressPtr> *bound_addresses,
tensor::TensorPtr CreatTensorForOutput(session::KernelGraph *kernel_graph, const CNodePtr &node, size_t index,
std::vector<tensor::TensorPtr> *need_sync_outputs);
BaseRef CreatTensorForOutput(const session::KernelWithIndex &kernel_with_index,
const std::unordered_map<AnfNode *, tensor::TensorPtr> &input_map,
std::set<DeviceAddressPtr> *bound_addresses,
BaseRef CreatTensorForOutput(session::KernelGraph *kernel_graph, const session::KernelWithIndex &kernel_with_index,
std::vector<tensor::TensorPtr> *need_sync_outputs);
void AssignValueNodeAddress(session::KernelGraph *kernel_graph);
void AssignInputNodeAddress(const session::KernelGraph *kernel_graph);
void AssignKernelOutputAddress(const session::KernelGraph *kernel_graph);
void AddRuntimeAddress(DeviceAddress *address, std::vector<kernel::AddressPtr> *input_list);
CPUResourceManager resource_manager_;
std::set<DeviceAddressPtr> bound_addresses_;
std::map<AnfNodePtr, tensor::TensorPtr> input_param_tensor_map_;
};
} // namespace cpu
} // namespace device

Write Preview
Loading…
Cancel
Save