1 /**
2 * Copyright 2021 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 #include "backend/session/single_kernel_graph.h"
18 #include "backend/session/anf_runtime_algorithm.h"
19
20 namespace mindspore {
21 namespace session {
ConstructKernelGraphBasedOnSingleOp(const std::string & op_name,const std::vector<TypeId> & input_dtypes,const std::vector<ShapeVector> & input_shapes,const std::vector<TypeId> & output_dtypes,const std::vector<std::vector<size_t>> & output_shapes)22 std::shared_ptr<session::KernelGraph> SingleKernelGraph::ConstructKernelGraphBasedOnSingleOp(
23 const std::string &op_name, const std::vector<TypeId> &input_dtypes, const std::vector<ShapeVector> &input_shapes,
24 const std::vector<TypeId> &output_dtypes, const std::vector<std::vector<size_t>> &output_shapes) {
25 auto graph = std::make_shared<session::KernelGraph>();
26 MS_EXCEPTION_IF_NULL(graph);
27 std::vector<AnfNodePtr> inputs;
28 // set input[0]
29 PrimitivePtr op_prim = std::make_shared<Primitive>(op_name);
30 MS_EXCEPTION_IF_NULL(op_prim);
31 inputs.push_back(std::make_shared<ValueNode>(op_prim));
32 // construct real input
33 if (input_dtypes.size() != input_shapes.size()) {
34 MS_LOG(EXCEPTION) << " input_dtypes size should equal to input_shapes size, the op name is: " << op_name;
35 }
36 auto input_num = input_dtypes.size();
37 for (size_t i = 0; i < input_num; ++i) {
38 auto tensor = std::make_shared<tensor::Tensor>(input_dtypes[i], input_shapes[i]);
39 auto value_node = graph->NewValueNode(tensor);
40 inputs.push_back(value_node);
41 }
42 // obtain cnode
43 auto cnode = graph->NewCNode(inputs);
44 MS_EXCEPTION_IF_NULL(cnode);
45 // get output dynamic shape info
46 AnfAlgo::SetNodeAttr(kAttrOutputIsDynamicShape, MakeValue(false), cnode);
47 if (output_dtypes.size() != output_shapes.size()) {
48 MS_LOG(EXCEPTION) << " output_dtypes size should equal to output_shapes size, the op name is: " << op_name;
49 }
50 AnfAlgo::SetOutputInferTypeAndShape(output_dtypes, output_shapes, cnode.get());
51 // set execution order
52 std::vector<CNodePtr> exe_order = {cnode};
53 graph->set_execution_order(exe_order);
54 // set graph output
55 graph->set_output(cnode);
56 graph->SetInputNodes();
57 return graph;
58 }
59 } // namespace session
60 } // namespace mindspore
61