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 #include "frontend/operator/prim_to_function.h"
18 #include "base/core_ops.h"
19
20 namespace mindspore {
21 // namespace to support prim related definition
22 namespace prim {
23
PrimToFunction()24 PrimToFunction::PrimToFunction()
25 : prim_func_type_map_(
26 {{"bool_not", kPrimTypeOneArg}, {"scalar_cos", kPrimTypeOneArg}, {"scalar_exp", kPrimTypeOneArg},
27 {kScalarFloor, kPrimTypeOneArg}, {"scalar_log", kPrimTypeOneArg}, {"scalar_sin", kPrimTypeOneArg},
28 {"scalar_tan", kPrimTypeOneArg}, {kScalarTrunc, kPrimTypeOneArg}, {"typeof", kPrimTypeOneArg},
29 {kScalarUadd, kPrimTypeOneArg}, {kScalarUsub, kPrimTypeOneArg}, {kScalarAdd, kPrimTypeTwoArgs},
30 {"bool_and", kPrimTypeTwoArgs}, {"bool_eq", kPrimTypeTwoArgs}, {"bool_or", kPrimTypeTwoArgs},
31 {kScalarDiv, kPrimTypeTwoArgs}, {"scalar_eq", kPrimTypeTwoArgs}, {"scalar_ge", kPrimTypeTwoArgs},
32 {"scalar_gt", kPrimTypeTwoArgs}, {"scalar_le", kPrimTypeTwoArgs}, {"scalar_lt", kPrimTypeTwoArgs},
33 {"scalar_ne", kPrimTypeTwoArgs}, {kScalarMod, kPrimTypeTwoArgs}, {kScalarMul, kPrimTypeTwoArgs},
34 {kScalarPow, kPrimTypeTwoArgs}, {kScalarSub, kPrimTypeTwoArgs}, {kScalarFloordiv, kPrimTypeTwoArgs}}) {}
35
GetFunction(const PrimitivePtr & prim,FunctionPtr * const func) const36 bool PrimToFunction::GetFunction(const PrimitivePtr &prim, FunctionPtr *const func) const {
37 bool result = false;
38
39 if (func != nullptr) {
40 int64_t args_num = GetPrimType(prim);
41 std::vector<TypePtr> one_arg{std::make_shared<Number>()};
42 std::vector<TypePtr> two_args{std::make_shared<Number>(), std::make_shared<Number>()};
43 TypePtr retval = std::make_shared<Number>();
44 result = true;
45 switch (args_num) {
46 case kPrimTypeOneArg:
47 *func = Function(one_arg, retval).DeepCopy()->cast<FunctionPtr>();
48 break;
49 case kPrimTypeTwoArgs:
50 *func = Function(two_args, retval).DeepCopy()->cast<FunctionPtr>();
51 break;
52 default:
53 result = false;
54 break;
55 }
56 }
57
58 return result;
59 }
60
GetPrimType(const PrimitivePtr & prim) const61 int64_t PrimToFunction::GetPrimType(const PrimitivePtr &prim) const {
62 MS_EXCEPTION_IF_NULL(prim);
63 int64_t prim_type = static_cast<int64_t>(kPrimTypeUnknown);
64
65 auto value = prim_func_type_map_.find(prim->name());
66 if (value != prim_func_type_map_.end()) {
67 prim_type = value->second;
68 }
69 return prim_type;
70 }
71 } // namespace prim
72 } // namespace mindspore
73