forked from huawei/mindspore2022
120 lines
5.5 KiB
C++
120 lines
5.5 KiB
C++
/**
|
|
* Copyright 2021 Huawei Technologies Co., Ltd
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#ifndef MINDSPORE_CCSRC_RUNTIME_FRAMEWORK_GRAPH_COMPILER_H_
|
|
#define MINDSPORE_CCSRC_RUNTIME_FRAMEWORK_GRAPH_COMPILER_H_
|
|
|
|
#include <vector>
|
|
#include <memory>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
#include <map>
|
|
#include "runtime/hardware/device_context.h"
|
|
#include "backend/session/session_basic.h"
|
|
#include "ir/tensor.h"
|
|
|
|
namespace mindspore {
|
|
using device::DeviceContext;
|
|
using mindspore::tensor::TensorPtr;
|
|
using session::InputTensorInfo;
|
|
using session::KernelWithIndex;
|
|
using session::OpRunInfo;
|
|
|
|
namespace runtime {
|
|
class GraphCompiler {
|
|
public:
|
|
static GraphCompiler &GetInstance() {
|
|
static GraphCompiler instance;
|
|
return instance;
|
|
}
|
|
|
|
// Set device context which is initialized, the function must be called
|
|
// before using GraphCompiler and after changing device type or device id.
|
|
void set_device_context(DeviceContext *device_context);
|
|
|
|
// Construct kernel graph from anf nodes list and compile kernel graph in Graph mode,
|
|
// the detailed implementation of compiling graph is in 'CompileGraphImpl'.
|
|
GraphId CompileGraph(const AnfNodePtrList &nodes, const AnfNodePtrList &outputs);
|
|
|
|
// Construct single op kernel graph and compile the kernel graph in PyNative mode.
|
|
GraphId CompileGraph(const session::OpRunInfo &op_run_info, const GraphInfo &graph_info,
|
|
const std::vector<int64_t> *tensors_mask, std::vector<TensorPtr> *input_tensors,
|
|
bool *single_op_cache_hit);
|
|
|
|
// Get graph by graph id, if not exist return nullptr, used in Graph mode.
|
|
KernelGraphPtr Fetch(GraphId graph_id) const;
|
|
|
|
// Get graph by graph info, if not exist return nullptr, used in PyNative mode.
|
|
KernelGraphPtr Fetch(const GraphInfo &graph_info) const;
|
|
|
|
// The following four methods used in PyNative back propagation to split complete kernel graph to single
|
|
// op graph, and these methods will be removed to class MindRTBackend after deleting session module.
|
|
|
|
// Cache index for all parameter and output nodes of kernel graph, used to get parameter of single op and
|
|
// recover output of original complete back propagation kernel graph.
|
|
void GetParamAndOutputIndex(const KernelGraphPtr &graph, const std::vector<TensorPtr> &inputs, VectorRef *outputs,
|
|
std::map<AnfNodePtr, size_t> *parameter_index,
|
|
std::map<KernelWithIndex, std::vector<std::vector<size_t>>> *output_indexes);
|
|
|
|
// Get input tensors for single op compile and run, input tensors may convert from value node and parameter in graph
|
|
// and prev kernel node's output.
|
|
void GetSingleOpInputTensors(const CNodePtr &kernel, const std::map<KernelWithIndex, TensorPtr> &op_output,
|
|
const std::map<AnfNodePtr, size_t> ¶meter_index,
|
|
const std::vector<TensorPtr> &graph_inputs, InputTensorInfo *input_tensor_info);
|
|
|
|
// Get OpRunInfo and GraphInfo for single op compile and run.
|
|
void GetSingleOpRunInfoAndGraphInfo(const CNodePtr &kernel, const std::vector<TensorPtr> &input_tensors,
|
|
OpRunInfo *run_info, GraphInfo *graph_info);
|
|
|
|
// Handle single op output tensor and recover output of original complete kernel graph.
|
|
void RecoverGraphOutput(const AnfNodePtr &kernel, const VectorRef &op_outputs,
|
|
const std::map<KernelWithIndex, std::vector<std::vector<size_t>>> &output_indexes,
|
|
std::map<KernelWithIndex, TensorPtr> *op_output_map, VectorRef *outputs,
|
|
std::vector<TensorPtr> *runop_output_tensors);
|
|
|
|
// Collect output tensors of back propagation graph for allreduce operators to average gradient,
|
|
// used in PyNative distributed training mode.
|
|
void AddGradAddrToBucket(const GraphId &graph_id, const std::vector<tensor::TensorPtr> &grad_tensor);
|
|
|
|
// Clear resource in bucket, such as useless tensors and device memory of all communication operators,
|
|
// Bucket is used in PyNative distributed training mode, one bucket handles all resource to launch and sync allreduce
|
|
// operator.
|
|
void ClearAllBucket(const GraphId &graph_id);
|
|
|
|
private:
|
|
GraphCompiler() = default;
|
|
~GraphCompiler() = default;
|
|
DISABLE_COPY_AND_ASSIGN(GraphCompiler);
|
|
|
|
// The implementation of compiling graph in Graph Mode, including optimizing graph,
|
|
// setting operator info, creating kernel and transforming kernel graph to ActorSet.
|
|
GraphId CompileGraphImpl(const KernelGraphPtr &graph) const;
|
|
|
|
// Create device address for all anf nodes of graph.
|
|
void CreateDeviceAddress(const KernelGraphPtr &graph) const;
|
|
|
|
DeviceContext *device_context_{nullptr};
|
|
|
|
// Single op kernel graph cache for PyNative mode.
|
|
std::unordered_map<GraphInfo, KernelGraphPtr> run_op_graphs_;
|
|
|
|
// The member variable 'session_' will be removed after removing session module.
|
|
session::SessionPtr session_{nullptr};
|
|
};
|
|
} // namespace runtime
|
|
} // namespace mindspore
|
|
#endif // MINDSPORE_CCSRC_RUNTIME_FRAMEWORK_GRAPH_COMPILER_H_
|