• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /**
2  * Copyright 2019 Huawei Technologies Co., Ltd
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #ifndef MINDSPORE_CCSRC_PIPELINE_JIT_RESOURCE_H_
18 #define MINDSPORE_CCSRC_PIPELINE_JIT_RESOURCE_H_
19 
20 #include <iostream>
21 #include <vector>
22 #include <string>
23 #include <unordered_map>
24 #include <memory>
25 #include <unordered_set>
26 
27 #include "pybind11/pybind11.h"
28 #include "pybind11/stl.h"
29 
30 #include "utils/any.h"
31 #include "utils/profile.h"
32 #include "ir/manager.h"
33 
34 #include "pipeline/jit/resource_base.h"
35 #include "pipeline/jit/static_analysis/prim.h"
36 #include "pipeline/jit/static_analysis/static_analysis.h"
37 
38 namespace mindspore {
39 namespace pipeline {
40 
41 namespace py = pybind11;
42 
43 const char kBackend[] = "backend";
44 const char kStepParallelGraph[] = "step_parallel";
45 const char kOutput[] = "output";
46 const char kPynativeGraphId[] = "graph_id";
47 
48 class InferenceResource;
49 
50 using BuiltInTypeMap = std::unordered_map<int64_t, std::unordered_map<std::string, Any>>;
51 
52 BuiltInTypeMap &GetMethodMap();
53 
54 BuiltInTypeMap &GetAttrMap();
55 
56 class Resource : public ResourceBase {
57  public:
58   explicit Resource(const py::object &obj = py::none());
59 
60   ~Resource() override;
61 
engine()62   abstract::AnalysisEnginePtr engine() { return engine_; }
63 
64   static bool IsTypeInBuiltInMap(const TypeId &type);
65 
66   static Any GetMethodPtr(const TypeId &type, const std::string &name);
67 
68   static Any GetAttrPtr(const TypeId &type, const std::string &name);
69 
source_input()70   const py::object &source_input() const { return source_input_; }
71 
func_graph()72   FuncGraphPtr func_graph() const { return func_graph_; }
set_func_graph(const FuncGraphPtr & func_graph)73   void set_func_graph(const FuncGraphPtr &func_graph) { func_graph_ = func_graph; }
74 
args_spec()75   const abstract::AbstractBasePtrList &args_spec() const { return args_spec_; }
set_args_spec(const abstract::AbstractBasePtrList & args_spec)76   void set_args_spec(const abstract::AbstractBasePtrList &args_spec) { args_spec_ = args_spec; }
77 
set_vm_loop(const bool & flag,const int64_t size)78   void set_vm_loop(const bool &flag, const int64_t size) {
79     vm_loop_flag_ = flag;
80     loop_size_ = size;
81   }
set_is_load(bool flag)82   void set_is_load(bool flag) { is_load_ = flag; }
is_load()83   bool is_load() { return is_load_; }
vm_loop_flag()84   bool vm_loop_flag() { return vm_loop_flag_; }
loop_size()85   int64_t loop_size() { return loop_size_; }
86   // Reclaim resource and clear the cache.
87   // GraphExecutorPy::Compile() can be called multiple times, so cache
88   // should be cleared.
89   void Clean();
90 
91  private:
92   abstract::AnalysisEnginePtr engine_;
93   FuncGraphPtr func_graph_;
94   abstract::AbstractBasePtrList args_spec_;
95   // The source obj to compile, usually a `Cell` or `ms_function` decorated function.
96   py::object source_input_;
97   bool is_cleaned_;
98   // The func_graph_ is loaded from mindir
99   bool is_load_{false};
100   bool vm_loop_flag_{false};
101   int64_t loop_size_{1};
102 };
103 
104 using ResourcePtr = std::shared_ptr<pipeline::Resource>;
105 
106 }  // namespace pipeline
107 }  // namespace mindspore
108 
109 #endif  // MINDSPORE_CCSRC_PIPELINE_JIT_RESOURCE_H_
110