From c13276f8a603d6005e8d06fdcb800ced13c3e567 Mon Sep 17 00:00:00 2001 From: shen_jingxing Date: Mon, 10 May 2021 10:17:50 +0800 Subject: [PATCH] sub --- .../core/abstract/primitive_infer_map.cc | 1 + mindspore/core/base/core_ops.h | 3 +- mindspore/core/ops/op_utils.cc | 33 +++++++++++++++---- mindspore/core/ops/sub.cc | 10 ++++-- mindspore/core/ops/sub.h | 4 +-- 5 files changed, 39 insertions(+), 12 deletions(-) diff --git a/mindspore/core/abstract/primitive_infer_map.cc b/mindspore/core/abstract/primitive_infer_map.cc index 2728f855e9..30f6780c51 100644 --- a/mindspore/core/abstract/primitive_infer_map.cc +++ b/mindspore/core/abstract/primitive_infer_map.cc @@ -195,6 +195,7 @@ PrimitiveEvalImplMap &GetPrimitiveToBackendEvalImplMap() { {prim::kPrimDivNoNan, {InferImplDivNoNan, nullptr, true}}, {prim::kPrimLinSpace, {InferImplLinSpace, nullptr, true}}, {prim::kPrimAddN, {InferImplAddN, nullptr, true}}, + {prim::kPrimSub, {InferImplSub, nullptr, false}}, {prim::kPrimLess, {InferImplLess, nullptr, true}}, {prim::kPrimStack, {InferImplStack, nullptr, true}}, diff --git a/mindspore/core/base/core_ops.h b/mindspore/core/base/core_ops.h index ee02232295..637696e7d5 100644 --- a/mindspore/core/base/core_ops.h +++ b/mindspore/core/base/core_ops.h @@ -45,6 +45,7 @@ constexpr auto kScalarTrunc = "ScalarTrunc"; constexpr auto kScalarFloor = "ScalarFloor"; constexpr auto kScalarUadd = "ScalarUadd"; constexpr auto kScalarUsub = "ScalarUsub"; +constexpr auto kSub = "Sub"; // Arrays constexpr auto kStack = "Stack"; @@ -371,7 +372,7 @@ inline const PrimitivePtr kPrimCentralization = std::make_shared("Cen inline const PrimitivePtr kPrimNeg = std::make_shared("Neg"); inline const PrimitivePtr kPrimSin = std::make_shared("Sin"); inline const PrimitivePtr kPrimCos = std::make_shared("Cos"); -inline const PrimitivePtr kPrimSub = std::make_shared("Sub"); +inline const PrimitivePtr kPrimSub = std::make_shared(kSub); inline const PrimitivePtr kPrimMul = std::make_shared("Mul"); inline const PrimitivePtr kPrimDiv = std::make_shared("Div"); inline const PrimitivePtr kPrimMod = std::make_shared("Mod"); diff --git a/mindspore/core/ops/op_utils.cc b/mindspore/core/ops/op_utils.cc index 64e61e0a8b..ae886934c0 100644 --- a/mindspore/core/ops/op_utils.cc +++ b/mindspore/core/ops/op_utils.cc @@ -25,14 +25,11 @@ namespace mindspore { namespace ops { -abstract::ShapePtr BroadCastInferShape(const std::string &op_name, const std::vector &input_args) { - MS_LOG(INFO) << "Do infer shape for op " << op_name; - auto x_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[0]->GetShapeTrack())[kShape]; - auto y_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[1]->GetShapeTrack())[kShape]; +std::vector CalBroadCastShape(std::vector x_shape, std::vector y_shape, + const std::string &op_name) { if (x_shape == y_shape) { - return std::make_shared(x_shape); + return x_shape; } - auto x_length = x_shape.size(); auto y_length = y_shape.size(); auto length = x_length < y_length ? x_length : y_length; @@ -53,7 +50,29 @@ abstract::ShapePtr BroadCastInferShape(const std::string &op_name, const std::ve MS_EXCEPTION(ValueError) << "For op " << op_name << ", the two input can not broadcast"; } } - return std::make_shared(broadcast_shape); + return broadcast_shape; +} +abstract::ShapePtr BroadCastInferShape(const std::string &op_name, const std::vector &input_args) { + MS_LOG(INFO) << "Do infer shape for op " << op_name; + CheckAndConvertUtils::CheckInteger("input numbers", input_args.size(), kGreaterEqual, 2, prim_name); + for (const auto &item : input_args) { + MS_EXCEPTION_IF_NULL(item); + } + auto x_shape_map = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[0]->GetShapeTrack()); + auto y_shape_map = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[1]->GetShapeTrack()); + auto x_shape = x_shape_map[kShape]; + auto y_shape = y_shape_map[kShape]; + auto x_min_shape = x_shape_map[kMinShape]; + auto x_max_shape = x_shape_map[kMaxShape]; + auto y_min_shape = y_shape_map[kMinShape]; + auto y_max_shape = y_shape_map[kMaxShape]; + if (x_shape == y_shape) { + return std::make_shared(x_shape, x_min_shape, x_max_shape); + } + auto broadcast_shape = CalBroadCastShape(x_shape, y_shape, op_name); + auto min_broadcast_shape = CalBroadCastShape(x_min_shape, y_min_shape, op_name); + auto max_broadcast_shape = CalBroadCastShape(x_max_shape, y_max_shape, op_name); + return std::make_shared(broadcast_shape, min_broadcast_shape, max_broadcast_shape); } } // namespace ops } // namespace mindspore diff --git a/mindspore/core/ops/sub.cc b/mindspore/core/ops/sub.cc index c45b15d440..22c0df31b1 100644 --- a/mindspore/core/ops/sub.cc +++ b/mindspore/core/ops/sub.cc @@ -1,5 +1,5 @@ /** - * Copyright 2020 Huawei Technologies Co., Ltd + * Copyright 2020-2021 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. @@ -14,12 +14,12 @@ * limitations under the License. */ +#include "ops/sub.h" #include #include #include #include -#include "ops/sub.h" #include "ops/op_utils.h" #include "utils/check_convert_utils.h" #include "abstract/primitive_infer_map.h" @@ -30,6 +30,10 @@ namespace { abstract::ShapePtr InferShape(const PrimitivePtr &primitive, const std::vector &input_args) { MS_EXCEPTION_IF_NULL(primitive); auto prim_name = primitive->name(); + CheckAndConvertUtils::CheckInteger("input numbers", input_args.size(), kGreaterEqual, 2, prim_name); + for (const auto &item : input_args) { + MS_EXCEPTION_IF_NULL(item); + } return BroadCastInferShape(prim_name, input_args); } @@ -37,6 +41,8 @@ TypePtr InferType(const PrimitivePtr &prim, const std::vector & for (const auto &item : input_args) { MS_EXCEPTION_IF_NULL(item); } + auto op_name = prim->name(); + CheckAndConvertUtils::CheckInteger("Sub infer", input_args.size(), kGreaterEqual, 2, op_name); std::map types; types.emplace("x", input_args[0]->BuildType()); types.emplace("y", input_args[1]->BuildType()); diff --git a/mindspore/core/ops/sub.h b/mindspore/core/ops/sub.h index 12cd95d6a3..89c9653837 100644 --- a/mindspore/core/ops/sub.h +++ b/mindspore/core/ops/sub.h @@ -1,5 +1,5 @@ /** - * Copyright 2020 Huawei Technologies Co., Ltd + * Copyright 2020-2021 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. @@ -26,7 +26,7 @@ namespace mindspore { namespace ops { -constexpr auto kNameSub = "Sub"; +constexpr auto kNameSub = prim::kSub; class Sub : public PrimitiveC { public: Sub() : PrimitiveC(kNameSub) { InitIOName({"x", "y"}, {"output"}); }