forked from huawei/mindspore2022
94 lines
3.6 KiB
C++
94 lines
3.6 KiB
C++
/**
|
|
* 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.
|
|
*/
|
|
|
|
#include "operator/prim_to_function.h"
|
|
#include <exception>
|
|
#include <iostream>
|
|
#include <string>
|
|
|
|
namespace mindspore {
|
|
// namespace to support prim related definition
|
|
namespace prim {
|
|
|
|
PrimToFunction::PrimToFunction()
|
|
: prim_func_type_map_({// ONE_ARG prim
|
|
{"bool_not", kPrimTypeOneArg},
|
|
{"scalar_cos", kPrimTypeOneArg},
|
|
{"scalar_exp", kPrimTypeOneArg},
|
|
{"scalar_floor", kPrimTypeOneArg},
|
|
{"scalar_log", kPrimTypeOneArg},
|
|
{"scalar_sin", kPrimTypeOneArg},
|
|
{"scalar_tan", kPrimTypeOneArg},
|
|
{"scalar_trunc", kPrimTypeOneArg},
|
|
{"typeof", kPrimTypeOneArg},
|
|
{"scalar_uadd", kPrimTypeOneArg},
|
|
{"scalar_usub", kPrimTypeOneArg},
|
|
// TWO_ARGS prim
|
|
{"scalar_add", kPrimTypeTwoArgs},
|
|
{"bool_and", kPrimTypeTwoArgs},
|
|
{"bool_eq", kPrimTypeTwoArgs},
|
|
{"bool_or", kPrimTypeTwoArgs},
|
|
{"scalar_div", kPrimTypeTwoArgs},
|
|
{"scalar_eq", kPrimTypeTwoArgs},
|
|
{"scalar_ge", kPrimTypeTwoArgs},
|
|
{"scalar_gt", kPrimTypeTwoArgs},
|
|
{"scalar_le", kPrimTypeTwoArgs},
|
|
{"scalar_lt", kPrimTypeTwoArgs},
|
|
{"scalar_ne", kPrimTypeTwoArgs},
|
|
{"scalar_mod", kPrimTypeTwoArgs},
|
|
{"scalar_mul", kPrimTypeTwoArgs},
|
|
{"scalar_pow", kPrimTypeTwoArgs},
|
|
{"scalar_sub", kPrimTypeTwoArgs},
|
|
{"scalar_floordiv", kPrimTypeTwoArgs}}) {}
|
|
|
|
bool PrimToFunction::GetFunction(const PrimitivePtr &prim, FunctionPtr *const func) const {
|
|
bool result = false;
|
|
|
|
if (func != nullptr) {
|
|
int args_num = GetPrimType(prim);
|
|
std::vector<TypePtr> one_arg{std::make_shared<Number>()};
|
|
std::vector<TypePtr> two_args{std::make_shared<Number>(), std::make_shared<Number>()};
|
|
TypePtr retval = std::make_shared<Number>();
|
|
result = true;
|
|
switch (args_num) {
|
|
case kPrimTypeOneArg:
|
|
*func = Function(one_arg, retval).DeepCopy()->cast<FunctionPtr>();
|
|
break;
|
|
case kPrimTypeTwoArgs:
|
|
*func = Function(two_args, retval).DeepCopy()->cast<FunctionPtr>();
|
|
break;
|
|
default:
|
|
result = false;
|
|
break;
|
|
}
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
int PrimToFunction::GetPrimType(const PrimitivePtr &prim) const {
|
|
MS_EXCEPTION_IF_NULL(prim);
|
|
int prim_type = static_cast<int>(kPrimTypeUnknown);
|
|
|
|
auto value = prim_func_type_map_.find(prim->name());
|
|
if (value != prim_func_type_map_.end()) {
|
|
prim_type = value->second;
|
|
}
|
|
return prim_type;
|
|
}
|
|
} // namespace prim
|
|
} // namespace mindspore
|