From 52c5f9befe23a25e594c0ad0e7cc63ba5b39d17c Mon Sep 17 00:00:00 2001 From: simson Date: Thu, 5 Aug 2021 11:13:02 +0800 Subject: [PATCH] modify the limit of axis of reduce ops --- .../format_type/change_axis_of_reduce_kernel.cc | 15 +++++++++++++++ mindspore/ccsrc/backend/session/kernel_graph.cc | 4 ++-- mindspore/ccsrc/backend/session/kernel_graph.h | 2 +- mindspore/ccsrc/utils/utils.h | 1 + mindspore/ops/operations/math_ops.py | 4 ++-- 5 files changed, 21 insertions(+), 5 deletions(-) diff --git a/mindspore/ccsrc/backend/optimizer/ascend/format_type/change_axis_of_reduce_kernel.cc b/mindspore/ccsrc/backend/optimizer/ascend/format_type/change_axis_of_reduce_kernel.cc index 30d899f3916..8e240b308fd 100644 --- a/mindspore/ccsrc/backend/optimizer/ascend/format_type/change_axis_of_reduce_kernel.cc +++ b/mindspore/ccsrc/backend/optimizer/ascend/format_type/change_axis_of_reduce_kernel.cc @@ -53,6 +53,15 @@ void SafeCheckFunction(const CNodePtr &cnode, const std::vector &reduce } } +void DynamicAttrUpdate(const AnfNodePtr &node) { + MS_EXCEPTION_IF_NULL(node); + auto primitive = AnfAlgo::GetCNodePrimitive(node); + MS_EXCEPTION_IF_NULL(primitive); + auto axis_attr = primitive->GetAttr(kAttrAxis); + AnfAlgo::SetNodeAttr(kAttrAxes, axis_attr, node); + AnfAlgo::EraseNodeAttr(kAttrAxis, node); +} + void ConvertReduceAttrFraczAnd6HD(const CNodePtr &cnode) { auto axis = kernel::GetReduceAttrAxis(cnode); std::vector convert_axis; @@ -95,9 +104,15 @@ const AnfNodePtr ChangeAxisOfReduceKernel::Process(const FuncGraphPtr &, const A } auto convert_map = kReduceConvertMap.find(AnfAlgo::GetInputFormat(node, 0)); if (convert_map == kReduceConvertMap.end()) { + if (AnfAlgo::IsDynamicShape(node)) { + DynamicAttrUpdate(node); + } return nullptr; } convert_map->second(node->cast()); + if (AnfAlgo::IsDynamicShape(node)) { + DynamicAttrUpdate(node); + } return nullptr; } } // namespace opt diff --git a/mindspore/ccsrc/backend/session/kernel_graph.cc b/mindspore/ccsrc/backend/session/kernel_graph.cc index 02e8d6baf48..e84c58c383b 100644 --- a/mindspore/ccsrc/backend/session/kernel_graph.cc +++ b/mindspore/ccsrc/backend/session/kernel_graph.cc @@ -470,7 +470,7 @@ void KernelGraph::CreateKernelInfoFromNewParameter(const CNodePtr &cnode) { } } -void KernelGraph::ResetAssignInputFeaatureMapFlag(const CNodePtr &cnode) const { +void KernelGraph::ResetAssignInputFeatureMapFlag(const CNodePtr &cnode) const { if (kOpAssignKernelNameList.find(AnfAlgo::GetCNodeName(cnode)) == kOpAssignKernelNameList.end()) { MS_LOG(EXCEPTION) << "Only supported to change the node [Assign , AssignSub, AssignAdd] node's input feature map " "flag but got the node :" @@ -493,7 +493,7 @@ void KernelGraph::SetKernelInfoForNode(const AnfNodePtr &node) const { node->set_kernel_info(kernel_info); if (node->isa()) { if (kOpAssignKernelNameList.find(AnfAlgo::GetCNodeName(node)) != kOpAssignKernelNameList.end()) { - ResetAssignInputFeaatureMapFlag(node->cast()); + ResetAssignInputFeatureMapFlag(node->cast()); } #if defined(__APPLE__) std::vector feature_map_input_indexs; diff --git a/mindspore/ccsrc/backend/session/kernel_graph.h b/mindspore/ccsrc/backend/session/kernel_graph.h index e303e0863dd..bc9e2c4de0d 100644 --- a/mindspore/ccsrc/backend/session/kernel_graph.h +++ b/mindspore/ccsrc/backend/session/kernel_graph.h @@ -111,7 +111,7 @@ class KernelGraph : public FuncGraph { CNodePtr NewCNodeWithInfos(const std::vector &inputs, const CNodePtr &ori_cnode = nullptr); void CreateKernelInfoFromNewParameter(const CNodePtr &cnode); CNodePtr NewCNode(const CNodePtr &cnode); - void ResetAssignInputFeaatureMapFlag(const CNodePtr &cnode) const; + void ResetAssignInputFeatureMapFlag(const CNodePtr &cnode) const; ParameterPtr NewParameter(const ParameterPtr ¶meter = nullptr); ParameterPtr NewParameter(const abstract::AbstractBasePtr &abstract); ValueNodePtr NewValueNode(const AbstractBasePtr &abstract, const ValuePtr &value); diff --git a/mindspore/ccsrc/utils/utils.h b/mindspore/ccsrc/utils/utils.h index 57323d48e69..a4983b310a4 100644 --- a/mindspore/ccsrc/utils/utils.h +++ b/mindspore/ccsrc/utils/utils.h @@ -335,6 +335,7 @@ constexpr auto kAttrDataShape = "data_shape"; constexpr auto kAttrFormat = "format"; constexpr auto kAttrReshapeType = "reshape_type"; constexpr auto kAttrAxis = "axis"; +constexpr auto kAttrAxes = "axes"; constexpr auto kAttrKeepDims = "keep_dims"; constexpr auto kAttrShapeGamma = "shape_gamma"; constexpr auto kAttrPerm = "perm"; diff --git a/mindspore/ops/operations/math_ops.py b/mindspore/ops/operations/math_ops.py index e850496f612..9ef89e4d096 100644 --- a/mindspore/ops/operations/math_ops.py +++ b/mindspore/ops/operations/math_ops.py @@ -371,6 +371,8 @@ class _Reduce(PrimitiveWithInfer): input_shp = input_x['shape'] args = {'input_x': input_x['dtype']} validator.check_tensors_dtypes_same_and_valid(args, valid_dtype, self.name) + if not isinstance(axis, mstype.tensor_type) and axis_v is None: + raise ValueError(f"For {self.name}, axis must be const.") out_shape = _infer_shape_reduce(input_shp, axis_v, self.keep_dims, self.name) if -1 in input_shp: if axis_v is None: @@ -423,8 +425,6 @@ class _Reduce(PrimitiveWithInfer): value = np_reduce_func(value, axis_v, keepdims=self.keep_dims) value = np.array(value) value = Tensor(value) - if (-1 in input_shp or axis_v is None) and context.get_context("device_target") == "Ascend": - self.init_prim_io_names(inputs=['x', 'axes'], outputs=['y']) return {'shape': out_shape, 'min_shape': output_min_shape, 'max_shape': output_max_shape,