convert the reduce axis attr when the reduce node selected the special format

mindspore/ccsrc/kernel/common_utils.cc

@@ -32,6 +32,8 @@
 
 namespace mindspore {
 namespace kernel {
+constexpr char kAxis[] = "axis";
+constexpr char kTypeInt32[] = "Int32";
 const std::unordered_map<std::string, TypeId> type_id_maps = {
   {"float", TypeId::kNumberTypeFloat32},   {"float16", TypeId::kNumberTypeFloat16},
   {"float32", TypeId::kNumberTypeFloat32}, {"float64", TypeId::kNumberTypeFloat64},
@@ -989,5 +991,39 @@ void MultiThreadCompute(const MultiThreadComputeFunc &func, MultiThreadComputePa
     threads[i].join();
   }
 }
+
+std::vector<int> GetReduceAttrAxis(const CNodePtr &cnode) {
+  if (AnfAlgo::GetInputTensorNum(cnode) != AnfAlgo::GetOutputTensorNum(cnode) &&
+      AnfAlgo::GetInputTensorNum(cnode) != 1) {
+    MS_LOG(EXCEPTION) << "the kind of reduce node [" << cnode->DebugString()
+                      << "] is not single input or single output ";
+  }
+  std::vector<int> axis;
+  auto input_shape = AnfAlgo::GetPrevNodeOutputInferShape(cnode, 0);
+  auto primitive = AnfAlgo::GetCNodePrimitive(cnode);
+  MS_EXCEPTION_IF_NULL(primitive);
+  auto axis_attr = primitive->GetAttr(kAxis);
+  if (axis_attr == nullptr) {
+    MS_LOG(ERROR) << "This node does't have axie attr.";
+    return std::vector<int>();
+  }
+  auto type = axis_attr->type();
+  MS_EXCEPTION_IF_NULL(type);
+  std::vector<int> axis_list;
+  if (type->ToString() == kTypeInt32) {
+    axis_list.emplace_back(GetValue<int>(axis_attr));
+  } else {
+    axis_list = GetValue<std::vector<int>>(axis_attr);
+  }
+  for (const auto &elem : axis_list) {
+    if (elem < 0) {
+      axis.emplace_back(input_shape.size() + elem);
+    } else {
+      axis.emplace_back(elem);
+    }
+  }
+  AnfAlgo::SetNodeAttr(kAttrAxis, MakeValue(axis), cnode);
+  return axis;
+}
 }  // namespace kernel
 }  // namespace mindspore

mindspore/ccsrc/kernel/common_utils.h

@@ -138,6 +138,7 @@ void ReduceMultiSparseGradient(const std::vector<std::shared_ptr<SparseGradient>
                                size_t outer_dim);
 void TwoLevelReduceSparseGradient(const SparseGradient &origin_sparse_grad, SparseGradient *tmp_grad,
                                   SparseGradient *unique_grad, size_t first_dim, size_t outer_dim);
+std::vector<int> GetReduceAttrAxis(const CNodePtr &cnode);
 }  // namespace kernel
 }  // namespace mindspore
 

mindspore/ccsrc/kernel/tbe/tbe_kernel_select/tbe_kernel_reduce_selecter.cc

@@ -20,11 +20,10 @@
 #include "utils/utils.h"
 #include "session/anf_runtime_algorithm.h"
 #include "kernel/tbe/tbe_kernel_select/common_utils.h"
+#include "kernel/common_utils.h"
 
 namespace mindspore {
 namespace kernel {
-constexpr char kAxis[] = "axis";
-constexpr char kTypeInt32[] = "Int32";
 constexpr size_t kInputIndex_0 = 0;
 constexpr size_t kOutputIndex_0 = 0;
 constexpr size_t kChannelN = 0;
@@ -50,7 +49,7 @@ bool TbeKernelReduceSelecter::GetShapeInfo(SupportFormat *support_format) {
   // get keep dim attr
   GetReduceAttrKeepDim();
   // get axis attr
-  GetReduceAttrAxis();
+  axis_ = GetReduceAttrAxis(cnode_ptr_);
   AssignSupportFormat(kOpFormat_DEFAULT, support_format);
   return true;
 }
@@ -121,31 +120,6 @@ bool TbeKernelReduceSelecter::IsFracZAndC1HWNCoC0Common(const std::string &forma
   return true;
 }
 
-void TbeKernelReduceSelecter::GetReduceAttrAxis() {
-  auto primitive = AnfAlgo::GetCNodePrimitive(cnode_ptr_);
-  MS_EXCEPTION_IF_NULL(primitive);
-  auto axis = primitive->GetAttr(kAxis);
-  if (axis == nullptr) {
-    MS_LOG(INFO) << "This node does't have axie attr.";
-    return;
-  }
-  auto type = axis->type();
-  MS_EXCEPTION_IF_NULL(type);
-  std::vector<int> axis_list;
-  if (type->ToString() == kTypeInt32) {
-    axis_list.emplace_back(GetValue<int>(axis));
-  } else {
-    axis_list = GetValue<std::vector<int>>(axis);
-  }
-  for (const auto &elem : axis_list) {
-    if (elem < 0) {
-      axis_.emplace_back(input_shape_.size() + elem);
-    } else {
-      axis_.emplace_back(IntToSize(elem));
-    }
-  }
-}
-
 void TbeKernelReduceSelecter::GetReduceAttrKeepDim() {
   if (!AnfAlgo::HasNodeAttr(kAttrKeepDims, cnode_ptr_)) {
     MS_LOG(INFO) << "This node does't have keep_attr.";

mindspore/ccsrc/kernel/tbe/tbe_kernel_select/tbe_kernel_reduce_selecter.h

@@ -36,7 +36,6 @@ class TbeKernelReduceSelecter {
 
  private:
   bool IsFracZAndC1HWNCoC0Common(const std::string &format, SupportFormat *support_format) const;
-  void GetReduceAttrAxis();
   void GetReduceAttrKeepDim();
   void AssignSupportFormat(const std::string &support_format_str, SupportFormat *support_format) const;
   bool Is4DShape(const std::vector<size_t> &shape) const;
@@ -44,7 +43,7 @@ class TbeKernelReduceSelecter {
   CNodePtr cnode_ptr_;
   std::vector<size_t> input_shape_{};
   std::vector<size_t> output_shape_{};
-  std::vector<size_t> axis_{};
+  std::vector<int> axis_{};
   bool keep_dims_ = false;
 };
 }  // namespace kernel

mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc

@@ -57,6 +57,7 @@
 #include "pre_activate/ascend/ir_fusion/softmax_grad_ext_fusion.h"
 #include "pre_activate/ascend/format_type/insert_trans_op.h"
 #include "pre_activate/ascend/format_type/rectify_do_mask_kernel_info.h"
+#include "pre_activate/ascend/format_type/chang_axis_of_reduce_kernel.h"
 #include "pre_activate/pass/getitem_tuple.h"
 #include "pre_activate/pass/optimize_dependence.h"
 #include "pre_activate/pass/erase_visit_attr.h"
@@ -157,6 +158,7 @@ void RunOpAscendDataLayout(const std::shared_ptr<session::KernelGraph> &kernel_g
   MS_EXCEPTION_IF_NULL(kernel_graph);
   auto optimizer = std::make_shared<GraphOptimizer>();
   auto data_layout_pm = std::make_shared<PassManager>("pynative_transop_pm");
+  data_layout_pm->AddPass(std::make_shared<ChangeAxisOfReduceKernel>());
   data_layout_pm->AddPass(std::make_shared<RectifyDoMaskKernelInfo>());
   data_layout_pm->AddPass(std::make_shared<RunOpInsertTransData>());
   data_layout_pm->AddPass(std::make_shared<GetitemTuple>());

mindspore/ccsrc/pre_activate/ascend/format_type/chang_axis_of_reduce_kernel.cc

@@ -0,0 +1,103 @@
+/**
+ * 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 "pre_activate/ascend/format_type/chang_axis_of_reduce_kernel.h"
+
+#include <string>
+#include <memory>
+#include <vector>
+#include <map>
+
+#include "utils/utils.h"
+#include "session/anf_runtime_algorithm.h"
+#include "common/utils.h"
+#include "kernel/common_utils.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+using ConvertFunction = std::function<void(const CNodePtr &)>;
+
+void ConvertReduceAttrFraczAnd6HD(const CNodePtr &cnode);
+const size_t kAxis_H = 2;
+const size_t kAxis_W = 3;
+const size_t kAxis_6HD_H = 1;
+const size_t kAxis_6HD_W = 2;
+const std::map<std::string, ConvertFunction> kReduceConvertMap = {{kOpFormat_FRAC_Z, ConvertReduceAttrFraczAnd6HD},
+                                                                  {kOpFormat_C1HWNCoC0, ConvertReduceAttrFraczAnd6HD}};
+void SafeCheckFunction(const CNodePtr &cnode, const std::vector<int> &reduce_axis) {
+  if (reduce_axis.empty()) {
+    MS_LOG(EXCEPTION) << "The node " << cnode->DebugString() << "'s reduce axis got a empty vector";
+  }
+  if (AnfAlgo::GetInputTensorNum(cnode) != AnfAlgo::GetOutputTensorNum(cnode) &&
+      AnfAlgo::GetInputTensorNum(cnode) != 1) {
+    MS_LOG(EXCEPTION) << "the kind of reduce node [" << cnode->DebugString()
+                      << "] is not single input or single output ";
+  }
+  for (auto elem : reduce_axis) {
+    if (elem > 4) {
+      MS_LOG(INFO) << "reduce axis is larger than 4 dims reduce axis : [" << elem << "]";
+    }
+  }
+}
+
+void ConvertReduceAttrFraczAnd6HD(const CNodePtr &cnode) {
+  auto axis = kernel::GetReduceAttrAxis(cnode);
+  std::vector<int> convert_axis;
+  SafeCheckFunction(cnode, axis);
+  auto format = AnfAlgo::GetInputFormat(cnode, 0);
+  if (format != kOpFormat_FRAC_Z || format != kOpFormat_C1HWNCoC0) {
+    MS_LOG(EXCEPTION) << "The node [" << cnode->DebugString() << "] format " << format << " is not 5hd";
+  }
+  for (auto elem : axis) {
+    switch (elem) {
+      case kAxis_H:
+        convert_axis.emplace_back(kAxis_6HD_H);
+        break;
+      case kAxis_W:
+        convert_axis.emplace_back(kAxis_6HD_W);
+        break;
+      default:
+        MS_LOG(INFO) << "reduce axis is axis : [" << elem << "]"
+                     << " but the format is not supported this reduce axis";
+    }
+  }
+  AnfAlgo::SetNodeAttr(kAttrAxis, MakeValue(convert_axis), cnode);
+}
+}  // namespace
+
+const BaseRef ChangeAxisOfReduceKernel::DefinePattern() const {
+  VarPtr X = std::make_shared<Var>();
+  VarPtr Xs = std::make_shared<SeqVar>();
+  return VectorRef({X, Xs});
+}
+
+const AnfNodePtr ChangeAxisOfReduceKernel::Process(const FuncGraphPtr &, const AnfNodePtr &node,
+                                                   const EquivPtr &) const {
+  if (node == nullptr || !node->isa<CNode>() || !AnfAlgo::IsRealKernel(node)) {
+    return nullptr;
+  }
+  if (AnfAlgo::GetOpPattern(node) != kernel::kReducePattern) {
+    return nullptr;
+  }
+  auto convert_map = kReduceConvertMap.find(AnfAlgo::GetInputFormat(node, 0));
+  if (convert_map == kReduceConvertMap.end()) {
+    return nullptr;
+  }
+  convert_map->second(node->cast<CNodePtr>());
+  return nullptr;
+}
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/pre_activate/ascend/format_type/chang_axis_of_reduce_kernel.h

@@ -0,0 +1,33 @@
+/**
+ * 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.
+ */
+#ifndef MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_FORMAT_TYPE_CHANGE_AXIS_OF_REDUCE_KENRNEL_H_
+#define MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_FORMAT_TYPE_CHANGE_AXIS_OF_REDUCE_KENRNEL_H_
+
+#include "pre_activate/common/optimizer.h"
+
+namespace mindspore {
+namespace opt {
+class ChangeAxisOfReduceKernel : public PatternProcessPass {
+ public:
+  explicit ChangeAxisOfReduceKernel(bool multigraph = true)
+      : PatternProcessPass("change_axis_of_reduce_kernel", multigraph) {}
+  ~ChangeAxisOfReduceKernel() override = default;
+  const BaseRef DefinePattern() const override;
+  const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override;
+};
+}  // namespace opt
+}  // namespace mindspore
+#endif  // MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_FORMAT_TYPE_CHANGE_AXIS_OF_REDUCE_KENRNEL_H_