mindspore2022/mindspore/ccsrc/vm/transform.h

141 lines
4.3 KiB
C++

/**
* This is the C++ adaptation and derivative work of Myia (https://github.com/mila-iqia/myia/).
*
* Copyright 2019 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_VM_TRANSFORM_H_
#define MINDSPORE_CCSRC_VM_TRANSFORM_H_
#include <string>
#include <memory>
#include <functional>
#include <utility>
#include <unordered_map>
#include <vector>
#include "vm/vm.h"
#include "ir/anf.h"
#include "operator/ops.h"
#include "vm/segment_runner.h"
#include "vm/backend.h"
// mindspore namespace is the top level namespace of Mindsporeession project.
// Other namespace should be a sub namespace of mindspore namespace in the ME project.
namespace mindspore {
extern const char kMsVm[];
extern const char kGeVm[];
// compile namespace
// A sub namespace in ME to support compile related definition.
namespace compile {
extern std::vector<PrimitivePtr> nonlinear_ops;
const std::vector<PrimitivePtr> &GetMsNonlinearOps();
using VmEvalFunc = std::function<BaseRef(const VectorRef &)>;
using VmEvalFuncPtr = std::shared_ptr<std::function<BaseRef(const VectorRef &)>>;
class CompileGraph {
public:
explicit CompileGraph(const BackendPtr &backend, const std::vector<PrimitivePtr> &cut_list = nonlinear_ops);
~CompileGraph() = default;
InstSet Run(const FuncGraphPtr &func_graph);
InstSet GenMultiGraphsSinkInst(const FuncGraphPtr &graph);
bool IsCut(const AnfNodePtr &node);
void Push(const AnfNodePtr &node);
void Tie(const AnfNodePtr &n1, const AnfNodePtr &n2) { slots_[n2] = slots_[n1]; }
void Ret(int nargs);
void GenMultiGraphsRun(const FuncGraphPtr &graph);
int Ref(const AnfNodePtr &node);
VectorRef SplitNodes(const FuncGraphPtr &func_graph);
void set_height(int h) {
height_ = h;
if (height_ > max_height_) {
max_height_ = height_;
}
}
void Reset() {
height_ = 0;
max_height_ = 0;
slots_.clear();
inst_.clear();
}
private:
void PushParameters(const FuncGraphPtr &func_graph);
bool SplitGraph(const FuncGraphPtr &func_graph);
int LinConvert(const FuncGraphPtr &func_graph, const AnfNodePtrList &node_list);
int InterpretNode(const FuncGraphPtr &func_graph, const CNodePtr &node);
int AddCall(const FuncGraphPtr &graph, const CNodePtr &node);
void AddSinkSwitch(const CNodePtr &node);
void AddPadStack(int param_height);
void AddTailCall(const AnfNodePtr &fn, size_t size);
void AddPartial(const CNodePtr &node);
void AddMakeTuple(const CNodePtr &node);
void AddSwitch(const CNodePtr &node);
void AddReturn(const CNodePtr &node);
void AddPrimitive(const CNodePtr &node, const PrimitivePtr &prim);
void AddInput(const AnfNodePtr &node);
void AddExternal(const LinConvertResult &result);
void AddInst(const Instruction &inst, const int &arg);
void AddInst(const Instruction &inst, const ValuePtr &arg);
void AddInst(const Instruction &inst, const VectorRef &args);
BackendPtr backend_;
LinkFuncType lin_convert_;
bool is_gevm_convert_;
int height_{0};
int max_height_{0};
std::vector<PrimitivePtr> cut_list_;
std::unordered_map<AnfNodePtr, int> slots_;
InstSet inst_;
};
using CompileGraphPtr = std::shared_ptr<CompileGraph>;
// CompileGraphs is used to Convert a graph cluster into instruction lists.
class CompileGraphs {
public:
explicit CompileGraphs(const BackendPtr &backend, const std::vector<PrimitivePtr> &cut_list = nonlinear_ops);
~CompileGraphs() = default;
void Reset() {
insts_.clear();
mapping_.clear();
}
void Compile(const FuncGraphPtr &func_graph);
FinalVMPtr Link(const FuncGraphPtr &func_graph);
FinalVMPtr CompileAndLink(const FuncGraphPtr &func_graph);
private:
InstSet insts_;
std::unordered_map<FuncGraphPtr, int> mapping_;
CompileGraphPtr transform_;
BackendPtr backend_;
};
BackendPtr CreateBackend();
} // namespace compile
} // namespace mindspore
#endif // MINDSPORE_CCSRC_VM_TRANSFORM_H_