!10806 insert reformat op

From: @liubuyu
Reviewed-by: @zhoufeng54,@kisnwang
Signed-off-by: @kisnwang

mindspore/ccsrc/backend/optimizer/ascend/ascend_backend_optimization.cc

@@ -65,7 +65,7 @@
 #include "backend/optimizer/ascend/ir_fusion/confusion_mul_grad_fusion.h"
 #include "backend/optimizer/ascend/ir_fusion/softmax_grad_ext_fusion.h"
 #include "backend/optimizer/ascend/format_type/insert_trans_op.h"
-#include "backend/optimizer/ascend/format_type/add_attr_for_3d_graph.h"
+#include "backend/optimizer/ascend/format_type/add_reformat_op.h"
 #include "backend/optimizer/ascend/format_type/dynamic_rnn_grad_reformat.h"
 #include "backend/optimizer/ascend/format_type/insert_transpose_for_basiclstm_op.h"
 #include "backend/optimizer/ascend/format_type/insert_transpose_for_dyanmic_gru_v2.h"
@@ -117,6 +117,7 @@
 #include "backend/optimizer/ascend/enhancer/split_inputs_for_reduce_scatter.h"
 #include "backend/optimizer/ascend/enhancer/add_placeholder_for_dynamic_rnn.h"
 #include "backend/optimizer/ascend/enhancer/add_placeholder_for_dynamic_gru.h"
+#include "backend/optimizer/ascend/enhancer/add_attr_for_3d_graph.h"
 #include "utils/ms_context.h"
 #include "utils/config_manager.h"
 #include "debug/anf_ir_dump.h"
@@ -210,7 +211,6 @@ void AscendDataLayout(const std::shared_ptr<session::KernelGraph> &kernel_graph)
   data_layout_pm->AddPass(std::make_shared<RectifyDoMaskKernelInfo>());
   data_layout_pm->AddPass(std::make_shared<DynamicRNNGradReformat>());
   data_layout_pm->AddPass(std::make_shared<ChangeAxisOfReduceKernel>());
-  data_layout_pm->AddPass(std::make_shared<AddIoFormatAttrFor3DGraph>());
   auto ms_context = MsContext::GetInstance();
   MS_EXCEPTION_IF_NULL(ms_context);
   if (ms_context->get_param<int>(MS_CTX_EXECUTION_MODE) == kPynativeMode) {
@@ -222,6 +222,8 @@ void AscendDataLayout(const std::shared_ptr<session::KernelGraph> &kernel_graph)
     data_layout_pm->AddPass(std::make_shared<InsertTransOp>());
     data_layout_pm->AddPass(std::make_shared<GetitemTuple>());
   }
+  data_layout_pm->AddPass(std::make_shared<EraseVisitAttr>());
+  data_layout_pm->AddPass(std::make_shared<AddIoFormatAttrFor3DGraph>());
   data_layout_pm->AddPass(std::make_shared<CommonSubexpressionElimination>());
   data_layout_pm->AddPass(std::make_shared<RemoveReshapePair>());
   data_layout_pm->AddPass(std::make_shared<EliminateRedundantOp>());
@@ -381,6 +383,7 @@ void AscendBackendOptimization(const std::shared_ptr<session::KernelGraph> &kern
       ConfigManager::GetInstance().iter_num() > 1) {
     other2_pm->AddPass(std::make_shared<GetnextMemcpyElimination>());
   }
+  other2_pm->AddPass(std::make_shared<AddReFormatOp>());
   other2_pm->AddPass(std::make_shared<CheckConsistency>());
   optimizer2->AddPassManager(other2_pm);
   (void)optimizer2->Optimize(kernel_graph);

mindspore/ccsrc/backend/optimizer/ascend/ascend_helper.cc

@@ -32,7 +32,6 @@ namespace mindspore {
 namespace opt {
 using KernelBuildInfoBuilder = kernel::KernelBuildInfo::KernelBuildInfoBuilder;
 namespace {
-const std::set<std::string> k3DFormatSet = {kOpFormat_NCDHW, kOpFormat_NDC1HWC0, kOpFormat_FRACTAL_Z_3D};
 AnfNodePtr CreateReshapeNode(const FuncGraphPtr &func_graph, const AnfNodePtr &input_node,
                              const KernelSelectPtr &kernel_select, const std::vector<size_t> &dst_shape) {
   std::vector<AnfNodePtr> trans_inputs;
@@ -66,19 +65,24 @@ void SetTransNodeAttr(const CNodePtr &trans_node) {
 
 std::string InitDefaultFormat(const AnfNodePtr &node) {
   MS_EXCEPTION_IF_NULL(node);
-  std::string default_format = kOpFormat_DEFAULT;
   if (node->isa<CNode>() && AnfAlgo::HasNodeAttr("io_format", node->cast<CNodePtr>())) {
     auto attr = AnfAlgo::GetNodeAttr<std::string>(node, "io_format");
     if (attr == kOpFormat_NCDHW) {
-      default_format = kOpFormat_NCDHW;
+      return kOpFormat_NCDHW;
     }
-  } else if (node->isa<ValueNode>() || node->isa<Parameter>()) {
-    auto out_format = AnfAlgo::GetOutputFormat(node, 0);
-    if (k3DFormatSet.find(out_format) != k3DFormatSet.end()) {
-      default_format = kOpFormat_NCDHW;
+  } else if (AnfAlgo::IsRealKernel(node)) {
+    auto formats = AnfAlgo::GetAllOutputFormats(node);
+    if (std::any_of(formats.begin(), formats.end(),
+                    [](const std::string &format) { return k3DFormatSet.find(format) != k3DFormatSet.end(); })) {
+      return kOpFormat_NCDHW;
+    }
+  } else {
+    auto format = AnfAlgo::GetOutputFormat(node, 0);
+    if (k3DFormatSet.find(format) != k3DFormatSet.end()) {
+      return kOpFormat_NCDHW;
     }
   }
-  return default_format;
+  return kOpFormat_DEFAULT;
 }
 
 void ReFreshInferShape(const AnfNodePtr &trans_node, const AnfNodePtr &node) {

mindspore/ccsrc/backend/optimizer/ascend/enhancer/add_attr_for_3d_graph.cc

@@ -0,0 +1,49 @@
+/**
+ * Copyright 2020 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 "backend/optimizer/ascend/enhancer/add_attr_for_3d_graph.h"
+#include <memory>
+#include "backend/session/anf_runtime_algorithm.h"
+#include "utils/utils.h"
+#include "base/core_ops.h"
+#include "runtime/device/kernel_info.h"
+#include "backend/optimizer/common/helper.h"
+
+namespace mindspore {
+namespace opt {
+const BaseRef AddIoFormatAttrFor3DGraph::DefinePattern() const {
+  std::shared_ptr<Var> V = std::make_shared<CondVar>(UnVisited);
+  std::shared_ptr<Var> Xs = std::make_shared<SeqVar>();
+  return VectorRef({V, Xs});
+}
+
+const AnfNodePtr AddIoFormatAttrFor3DGraph::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                                    const EquivPtr &) const {
+  MS_EXCEPTION_IF_NULL(node);
+  MS_EXCEPTION_IF_NULL(func_graph);
+  if (AnfAlgo::IsRealKernel(node)) {
+    AnfAlgo::SetNodeAttr(kAttrVisited, MakeValue(true), node);
+    auto formats = AnfAlgo::GetAllOutputFormats(node);
+    if (std::any_of(formats.begin(), formats.end(),
+                    [](const std::string &format) { return k3DFormatSet.find(format) != k3DFormatSet.end(); })) {
+      AnfAlgo::SetNodeAttr("io_format", MakeValue(kOpFormat_NCDHW), node);
+    }
+    return node;
+  }
+  return nullptr;
+}
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/optimizer/ascend/enhancer/add_attr_for_3d_graph.h

@@ -0,0 +1,40 @@
+/**
+ * Copyright 2020 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_ENHANCER_ADD_ATTR_FOR_3D_GRAPH_H
+#define MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_ENHANCER_ADD_ATTR_FOR_3D_GRAPH_H
+#include <vector>
+#include <string>
+#include <utility>
+#include <memory>
+#include "ir/anf.h"
+#include "backend/optimizer/common/optimizer.h"
+#include "backend/optimizer/ascend/ascend_helper.h"
+
+namespace mindspore {
+namespace opt {
+class AddIoFormatAttrFor3DGraph : public PatternProcessPass {
+ public:
+  explicit AddIoFormatAttrFor3DGraph(bool multigraph = true)
+      : PatternProcessPass("add_attr_for_3d_graph", multigraph) {}
+  ~AddIoFormatAttrFor3DGraph() 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_ENHANCER_ADD_ATTR_FOR_3D_GRAPH_H

mindspore/ccsrc/backend/optimizer/ascend/format_type/add_attr_for_3d_graph.cc

@@ -1,63 +0,0 @@
-/**
- * Copyright 2020 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 "backend/optimizer/ascend/format_type/add_attr_for_3d_graph.h"
-#include <memory>
-#include "backend/session/anf_runtime_algorithm.h"
-#include "utils/utils.h"
-#include "base/core_ops.h"
-#include "runtime/device/kernel_info.h"
-#include "backend/optimizer/common/helper.h"
-
-namespace mindspore {
-namespace opt {
-namespace {
-void AddAttrForAllCNode(const std::vector<AnfNodePtr> &node_list) {
-  for (auto node : node_list) {
-    if (node == nullptr || !node->isa<CNode>() || !AnfAlgo::IsRealKernel(node)) {
-      continue;
-    }
-    AnfAlgo::SetNodeAttr("io_format", MakeValue(kOpFormat_NCDHW), node);
-  }
-}
-
-bool NodeHasAttrIoFormat(const AnfNodePtr &node) {
-  MS_EXCEPTION_IF_NULL(node);
-  if (node->isa<CNode>()) {
-    auto cnode = node->cast<CNodePtr>();
-    MS_EXCEPTION_IF_NULL(cnode);
-    if (AnfAlgo::HasNodeAttr("io_format", cnode)) {
-      auto attr = AnfAlgo::GetNodeAttr<std::string>(cnode, "io_format");
-      return attr == kOpFormat_NCDHW;
-    }
-  }
-  return false;
-}
-}  // namespace
-
-bool AddIoFormatAttrFor3DGraph::Run(const FuncGraphPtr &func_graph) {
-  MS_EXCEPTION_IF_NULL(func_graph);
-  std::vector<AnfNodePtr> node_list = TopoSort(func_graph->get_return());
-  bool changed = false;
-  if (std::any_of(node_list.begin(), node_list.end(),
-                  [](const AnfNodePtr &node) { return NodeHasAttrIoFormat(node); })) {
-    AddAttrForAllCNode(node_list);
-    changed = true;
-  }
-  return changed;
-}
-}  // namespace opt
-}  // namespace mindspore

mindspore/ccsrc/backend/optimizer/ascend/format_type/add_reformat_op.cc

@@ -0,0 +1,136 @@
+/**
+ * Copyright 2020 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 "backend/optimizer/ascend/format_type/add_reformat_op.h"
+#include <memory>
+#include "backend/session/anf_runtime_algorithm.h"
+#include "utils/utils.h"
+#include "base/core_ops.h"
+#include "runtime/device/kernel_info.h"
+#include "backend/optimizer/common/helper.h"
+
+namespace mindspore {
+namespace opt {
+using KernelWithIndex = std::pair<AnfNodePtr, size_t>;
+namespace {
+AnfNodePtr InsertReFormatOp(const FuncGraphPtr &func_graph, const AnfNodePtr &node, const AnfNodePtr &in_node,
+                            size_t idx) {
+  MS_EXCEPTION_IF_NULL(node);
+  MS_EXCEPTION_IF_NULL(in_node);
+  MS_EXCEPTION_IF_NULL(func_graph);
+  std::vector<AnfNodePtr> reformat_inputs;
+  auto node_kernel_build_info = AnfAlgo::GetSelectKernelBuildInfo(node);
+  MS_EXCEPTION_IF_NULL(node_kernel_build_info);
+  auto reformat_prim = std::make_shared<Primitive>(prim::kPrimReformat->name());
+  reformat_inputs.push_back(NewValueNode(reformat_prim));
+  reformat_inputs.push_back(in_node);
+  auto reformat = func_graph->NewCNode(reformat_inputs);
+  auto reformat_builder = std::make_shared<kernel::KernelBuildInfo::KernelBuildInfoBuilder>();
+  reformat_builder->SetInputsFormat({AnfAlgo::GetPrevNodeOutputFormat(node, idx)});
+  reformat_builder->SetOutputsFormat({AnfAlgo::GetInputFormat(node, idx)});
+  reformat_builder->SetInputsDeviceType({AnfAlgo::GetPrevNodeOutputDeviceDataType(node, idx)});
+  reformat_builder->SetOutputsDeviceType({node_kernel_build_info->GetInputDeviceType(idx)});
+  AnfAlgo::SetSelectKernelBuildInfo(reformat_builder->Build(), reformat.get());
+
+  reformat->set_abstract(in_node->abstract());
+  AnfAlgo::SetNodeAttr("nop_op", MakeValue(true), reformat);
+  return reformat;
+}
+
+bool NeedInsert(const CNodePtr &cnode, const size_t input_index) {
+  KernelWithIndex kernel_with_index = AnfAlgo::GetPrevNodeOutput(cnode, input_index);
+  auto real_input_node = kernel_with_index.first;
+  auto idx = kernel_with_index.second;
+  auto input_format = AnfAlgo::GetInputFormat(cnode, input_index);
+  auto prev_format = AnfAlgo::GetOutputFormat(real_input_node, idx);
+  bool flag_format = (input_format != prev_format);
+  if (!flag_format) {
+    return false;
+  }
+
+  bool flag_shape = true;
+  auto input_origin_shape = AnfAlgo::GetOutputInferShape(real_input_node, idx);
+  if (prev_format == kOpFormat_DEFAULT || input_format == kOpFormat_DEFAULT) {
+    string checking_format = (prev_format == kOpFormat_DEFAULT) ? input_format : prev_format;
+    // when input shape size is 1D, default format and NC1HWC0 are compatible
+    if (input_origin_shape.size() == 1 && checking_format == kOpFormat_NC1HWC0) {
+      flag_shape = false;
+    }
+    if (kDefaultCompatibleFormat.find(checking_format) != kDefaultCompatibleFormat.end()) {
+      flag_shape = false;
+    }
+  }
+  if (input_origin_shape.size() == 0) {
+    flag_shape = false;
+  }
+  return flag_format && flag_shape;
+}
+
+AnfNodePtr NeedInSertReformatOp(const FuncGraphPtr &func_graph, const AnfNodePtr &node) {
+  MS_EXCEPTION_IF_NULL(node);
+  MS_EXCEPTION_IF_NULL(func_graph);
+  if (!node->isa<CNode>() || !AnfAlgo::IsRealKernel(node)) {
+    return nullptr;
+  }
+  auto cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(cnode);
+  auto in_nums = AnfAlgo::GetInputTensorNum(cnode);
+  bool need_insert = false;
+  std::vector<AnfNodePtr> new_inputs = {AnfAlgo::GetCNodePrimitiveNode(cnode)};
+  for (size_t i = 0; i < in_nums; i++) {
+    auto input_node = AnfAlgo::GetInputNode(cnode, i);
+    if (NeedInsert(cnode, i)) {
+      need_insert = true;
+      auto re_format = InsertReFormatOp(func_graph, cnode, input_node, i);
+      new_inputs.push_back(re_format);
+      continue;
+    }
+    new_inputs.push_back(input_node);
+  }
+  if (need_insert) {
+    auto kernel_graph = func_graph->cast<std::shared_ptr<session::KernelGraph>>();
+    CNodePtr new_node = nullptr;
+    if (kernel_graph == nullptr) {
+      new_node = std::make_shared<CNode>(*cnode);
+    } else {
+      new_node = kernel_graph->NewCNode(cnode);
+    }
+    MS_EXCEPTION_IF_NULL(new_node);
+    new_node->set_inputs(new_inputs);
+    AnfAlgo::CopyNodeAttrs(cnode, new_node);
+    return new_node;
+  }
+  return nullptr;
+}
+}  // namespace
+
+bool AddReFormatOp::Run(const FuncGraphPtr &func_graph) {
+  MS_EXCEPTION_IF_NULL(func_graph);
+  std::vector<AnfNodePtr> node_list = TopoSort(func_graph->get_return());
+  bool changed = false;
+  auto manager = func_graph->manager();
+  MS_EXCEPTION_IF_NULL(manager);
+  for (auto &node : node_list) {
+    auto new_node = NeedInSertReformatOp(func_graph, node);
+    if (new_node != nullptr) {
+      manager->Replace(node, new_node);
+      changed = true;
+    }
+  }
+  return changed;
+}
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/optimizer/ascend/format_type/add_attr_for_3d_graph.h → mindspore/ccsrc/backend/optimizer/ascend/format_type/add_reformat_op.h

@@ -25,10 +25,10 @@
 
 namespace mindspore {
 namespace opt {
-class AddIoFormatAttrFor3DGraph : public Pass {
+class AddReFormatOp : public Pass {
  public:
-  explicit AddIoFormatAttrFor3DGraph(size_t groups = 1) : Pass("add_attr_for_3d_graph"), groups_(groups) {}
-  ~AddIoFormatAttrFor3DGraph() override = default;
+  explicit AddReFormatOp(size_t groups = 1) : Pass("add_reformat_op"), groups_(groups) {}
+  ~AddReFormatOp() override = default;
   bool Run(const FuncGraphPtr &graph) override;
 
  private:

mindspore/ccsrc/backend/optimizer/common/helper.cc

@@ -281,7 +281,7 @@ bool IsNopNode(const AnfNodePtr &node) {
 
   static std::unordered_set<std::string> nop_nodes = {prim::kPrimReshape->name(), kExpandDimsOpName,
                                                       prim::kPrimSqueeze->name(), prim::kPrimFlatten->name(),
-                                                      kFlattenGradOpName};
+                                                      kFlattenGradOpName,         prim::kPrimReformat->name()};
   if (node == nullptr || !node->isa<CNode>()) {
     return false;
   }

mindspore/ccsrc/runtime/device/ascend/ascend_device_address.cc

@@ -93,9 +93,9 @@ namespace device {
 namespace ascend {
 const int FLOAT_LEN = sizeof(float);
 const int FLOAT16_LEN = 2;  // sizeof(float16);
-const std::set<std::string> kOpNeedTransFormat = {kOpFormat_NHWC,        kOpFormat_HWCN,         kOpFormat_NC1HWC0,
-                                                  kOpFormat_FRAC_Z,      kOpFormat_C1HWNCoC0,    kOpFormat_FRAC_NZ,
-                                                  kOpFormat_NC1HWC0_C04, kOpFormat_FRACTAL_Z_C04};
+const std::set<std::string> kOpNeedTransFormat = {kOpFormat_NHWC,        kOpFormat_HWCN,          kOpFormat_NC1HWC0,
+                                                  kOpFormat_FRAC_Z,      kOpFormat_C1HWNCoC0,     kOpFormat_FRAC_NZ,
+                                                  kOpFormat_NC1HWC0_C04, kOpFormat_FRACTAL_Z_C04, kOpFormat_NDC1HWC0};
 
 void SyncMemory(void *dst, const void *src, uint64_t size, rtMemcpyKind_t kind) {
   auto ms_context = MsContext::GetInstance();

mindspore/ccsrc/utils/utils.h

@@ -488,6 +488,8 @@ const std::set<TypeId> kFloatDataTypeSet = {kNumberTypeFloat16, kNumberTypeFloat
 const std::set<std::string> kComputeDepend = {kUniqueOpName, kComputeAccidentalHitsOpName, kSubAndFilterOpName,
                                               kPadAndShiftOpName};
 
+const std::set<std::string> k3DFormatSet = {kOpFormat_NCDHW, kOpFormat_NDC1HWC0, kOpFormat_FRACTAL_Z_3D};
+
 static inline void ChangeFileMode(const std::string &file_name, mode_t mode) {
   try {
     if (chmod(file_name.c_str(), mode) != 0) {

mindspore/core/base/core_ops.h

@@ -308,6 +308,7 @@ inline const PrimitivePtr kPrimErrorOnDynamicShapeInput = std::make_shared<Primi
 
 // Other miscellaneous
 inline const PrimitivePtr kPrimDepend = std::make_shared<Primitive>("Depend");
+inline const PrimitivePtr kPrimReformat = std::make_shared<Primitive>("Reformat");
 inline const PrimitivePtr kPrimPartial = std::make_shared<Primitive>("Partial");
 inline const PrimitivePtr kPrimIdentity = std::make_shared<Primitive>("identity");
 inline const PrimitivePtr kPrimHookBackward = std::make_shared<Primitive>("HookBackward");