!32495 add conv2dbackprobfilter + mul pass

Merge pull request !32495 from jjfeing/master

mindspore/ccsrc/plugin/device/ascend/optimizer/ascend_backend_optimization.cc

@@ -59,6 +59,7 @@
 #include "plugin/device/ascend/optimizer/ir_fusion/transpose_transdata_fusion.h"
 #include "plugin/device/ascend/optimizer/ir_fission/transdata_split.h"
 #include "plugin/device/ascend/optimizer/ir_fission/topk_split.h"
+#include "plugin/device/ascend/optimizer/ir_fission/conv2d_backprop_filter_mul_fission.h"
 #include "plugin/device/ascend/optimizer/ir_fission/lin_space_fission.h"
 #include "plugin/device/ascend/optimizer/ir_fission/space_to_depth_split.h"
 #include "plugin/device/ascend/optimizer/ir_fission/diag_fission.h"
@@ -210,6 +211,7 @@ void AddAscendIRFusionPass(PassManager *ir_fusion_pm) {
   ir_fusion_pm->AddPass(std::make_shared<ConfusionSoftmaxGradRule>());
   ir_fusion_pm->AddPass(std::make_shared<ReshapeTransposeFusion>());
   ir_fusion_pm->AddPass(std::make_shared<TransposeReshapeFusion>());
+  ir_fusion_pm->AddPass(std::make_shared<Conv2dBackpropFilterMul>());
   ir_fusion_pm->AddPass(std::make_shared<TopKSplit>());
   ir_fusion_pm->AddPass(std::make_shared<LinSpaceFission>());
   ir_fusion_pm->AddPass(std::make_shared<DiagFission>());
@@ -385,6 +387,7 @@ void RunOpAscendBackendIRFusionOptimization(const std::shared_ptr<session::Kerne
   ir_fusion_pm->AddPass(std::make_shared<BnSplit>());
   ir_fusion_pm->AddPass(std::make_shared<BnGradSplit>());
   ir_fusion_pm->AddPass(std::make_shared<LayerNormGradSplit>());
+  ir_fusion_pm->AddPass(std::make_shared<Conv2dBackpropFilterMul>());
   ir_fusion_pm->AddPass(std::make_shared<TopKSplit>());
   ir_fusion_pm->AddPass(std::make_shared<LinSpaceFission>());
   ir_fusion_pm->AddPass(std::make_shared<SpaceToDepthSplit>());

mindspore/ccsrc/plugin/device/ascend/optimizer/format_type/set_fracz_group_attr.cc

@@ -44,7 +44,7 @@ AnfNodePtr GetOutputItem(const FuncGraphManagerPtr &manager, const CNodePtr &cno
   while (!depend_nodes.empty()) {
     auto node = depend_nodes.back();
     depend_nodes.pop_back();
-    for (auto node_index : manager->node_users()[node]) {
+    for (const auto &node_index : manager->node_users()[node]) {
       if (common::AnfAlgo::CheckPrimitiveType(node_index.first, prim::kPrimDepend) && node_index.second == 1) {
         (void)depend_nodes.emplace_back(node_index.first);
       } else if (common::AnfAlgo::CheckPrimitiveType(node_index.first, prim::kPrimTupleGetItem)) {
@@ -71,6 +71,14 @@ bool HasFraczGroupAttrAndSet(const AnfNodePtr &node, size_t index, int64_t group
     param->set_fracz_group(groups);
     return false;
   }
+  if (node->isa<ValueNode>()) {
+    auto value_node = node->cast<ValueNodePtr>();
+    if (value_node->fracz_group() != 1) {
+      return true;
+    }
+    value_node->set_fracz_group(groups);
+    return false;
+  }
   if (node->isa<CNode>()) {
     auto cnode = node->cast<CNodePtr>();
     auto node_name = common::AnfAlgo::GetCNodeName(cnode);
@@ -147,7 +155,7 @@ std::vector<KernelWithIndex> GetCNodeNeighborFraczNodes(const FuncGraphManagerPt
       auto output = GetOutputItem(manager, cnode, groups, i);
       if (output != nullptr) {
         (void)std::transform(node_user[output].begin(), node_user[output].end(), std::back_inserter(ret),
-                             [](KernelWithIndex node_index) {
+                             [](const KernelWithIndex &node_index) {
                                return KernelWithIndex{node_index.first, node_index.second - 1};
                              });
       }
@@ -160,9 +168,9 @@ std::vector<KernelWithIndex> GetNeighborFraczNodes(const FuncGraphManagerPtr &ma
                                                    size_t index, int64_t groups) {
   std::vector<KernelWithIndex> ret;
   auto node_user = manager->node_users();
-  if (node->isa<Parameter>()) {
+  if (node->isa<Parameter>() || node->isa<ValueNode>()) {
     std::transform(node_user[node].begin(), node_user[node].end(), std::back_inserter(ret),
-                   [](KernelWithIndex node_index) {
+                   [](const KernelWithIndex &node_index) {
                      return KernelWithIndex{node_index.first, node_index.second - 1};
                    });
   }
@@ -178,7 +186,7 @@ std::vector<KernelWithIndex> GetNeighborFraczNodes(const FuncGraphManagerPtr &ma
     auto output = GetOutputItem(manager, cnode, groups, index);
     if (output != nullptr) {
       (void)std::transform(node_user[output].begin(), node_user[output].end(), std::back_inserter(ret),
-                           [](KernelWithIndex node_index) {
+                           [](const KernelWithIndex &node_index) {
                              return KernelWithIndex{node_index.first, node_index.second - 1};
                            });
     }
@@ -210,7 +218,8 @@ bool SetAttrFraczGroup(const FuncGraphPtr &func_graph, const CNodePtr &cnode) {
   return true;
 }
 
-bool SetAttrFraczGroup(const FuncGraphPtr &func_graph, const ParameterPtr &param) {
+template <typename T>
+bool SetAttrFraczGroup(const FuncGraphPtr &func_graph, const T &param) {
   MS_EXCEPTION_IF_NULL(func_graph);
   MS_EXCEPTION_IF_NULL(param);
   auto groups = param->fracz_group();
@@ -253,8 +262,15 @@ bool SetFraczGroupAttr::Run(const FuncGraphPtr &func_graph) {
     if (node->isa<Parameter>()) {
       // transmit fracz_group attr through multi graph by parameter
       auto param = node->cast<ParameterPtr>();
+      MS_EXCEPTION_IF_NULL(param);
       changed = SetAttrFraczGroup(func_graph, param) || changed;
     }
+    if (node->isa<ValueNode>()) {
+      // transmit fracz_group attr through multi graph by value node
+      auto value_node = node->cast<ValueNodePtr>();
+      MS_EXCEPTION_IF_NULL(value_node);
+      changed = SetAttrFraczGroup(func_graph, value_node) || changed;
+    }
     if (node->isa<CNode>()) {
       auto cnode = node->cast<CNodePtr>();
       if (cnode == nullptr) {

mindspore/ccsrc/plugin/device/ascend/optimizer/ir_fission/conv2d_backprop_filter_mul_fission.cc

@@ -0,0 +1,130 @@
+/**
+ * Copyright 2022 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 "plugin/device/ascend/optimizer/ir_fission/conv2d_backprop_filter_mul_fission.h"
+
+#include <algorithm>
+#include <string>
+#include <vector>
+#include <memory>
+#include "backend/common/optimizer/const_input_to_attr.h"
+#include "kernel/kernel_build_info.h"
+#include "include/common/utils/utils.h"
+#include "backend/common/session/kernel_graph.h"
+#include "backend/common/session/anf_runtime_algorithm.h"
+#include "include/common/utils/anfalgo.h"
+#include "runtime/device/kernel_info.h"
+#include "utils/ms_context.h"
+
+namespace mindspore::opt {
+namespace {
+constexpr int64_t kGroupsDefaultValue = 1;
+
+template <typename T>
+void SetAssistTensorData(void *data, const T &value, size_t dims_size) {
+  MS_EXCEPTION_IF_NULL(data);
+  auto tensor_data = static_cast<T *>(data);
+  for (size_t i = 0; i < dims_size; ++i) {
+    tensor_data[i] = value;
+  }
+}
+
+ValueNodePtr CreateAssistNode(const FuncGraphPtr &func_graph, const AnfNodePtr &node, const std::vector<size_t> &shape,
+                              size_t matrix_size) {
+  MS_EXCEPTION_IF_NULL(func_graph);
+  MS_EXCEPTION_IF_NULL(node);
+  auto type = common::AnfAlgo::GetOutputInferDataType(node, 0);
+  std::vector<int64_t> assist_shape;
+  std::transform(shape.begin(), shape.end(), std::back_inserter(assist_shape), SizeToLong);
+  tensor::TensorPtr tensor = std::make_shared<tensor::Tensor>(type, assist_shape);
+  AbstractBasePtr x_abstract;
+  if (type == kNumberTypeInt32) {
+    SetAssistTensorData<int32_t>(tensor->data_c(), 1, matrix_size);
+    x_abstract = std::make_shared<abstract::AbstractTensor>(kInt32, assist_shape);
+  } else if (type == kNumberTypeFloat16) {
+    SetAssistTensorData<float16>(tensor->data_c(), float16(static_cast<float>(1)), matrix_size);
+    x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat16, assist_shape);
+  } else if (type == kNumberTypeFloat32) {
+    SetAssistTensorData<float>(tensor->data_c(), static_cast<float>(1), matrix_size);
+    x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat, assist_shape);
+  } else {
+    MS_EXCEPTION(TypeError) << "The type of node [" << node->DebugString()
+                            << "] should be int32, float16 or float32, but got" << node->Type()->ToString();
+  }
+  auto kernel_graph = func_graph->cast<KernelGraphPtr>();
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  auto assist_value_node = kernel_graph->NewValueNode(x_abstract, tensor);
+  kernel_graph->AddValueNodeToGraph(assist_value_node);
+  common::AnfAlgo::SetOutputInferTypeAndShape({type}, {shape}, assist_value_node.get());
+  return assist_value_node;
+}
+}  // namespace
+
+const BaseRef Conv2dBackpropFilterMul::DefinePattern() const {
+  VarPtr X1 = std::make_shared<Var>();
+  VarPtr X2 = std::make_shared<Var>();
+  auto prim = std::make_shared<Primitive>(kConv2DBackpropFilterOpName);
+  return VectorRef({prim, X1, X2});
+}
+
+const AnfNodePtr Conv2dBackpropFilterMul::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                                  const EquivPtr &) const {
+  MS_EXCEPTION_IF_NULL(func_graph);
+  MS_EXCEPTION_IF_NULL(node);
+  if (common::AnfAlgo::IsDynamicShape(node)) {
+    return nullptr;
+  }
+  if (GetBoolAttr(node, kAttrVisited)) {
+    return nullptr;
+  }
+  auto cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(cnode);
+  if (!common::AnfAlgo::HasNodeAttr(kAttrGroup, cnode)) {
+    MS_LOG(EXCEPTION) << "Get Conv2DBackpropFilter attr(groups) failed, node: " << node->DebugString();
+  }
+  auto groups = common::AnfAlgo::GetNodeAttr<int64_t>(node, kAttrGroup);
+  // if groups not > 1, skip process
+  if (groups <= kGroupsDefaultValue) {
+    return nullptr;
+  }
+  auto shape = common::AnfAlgo::GetOutputInferShape(node, 0);
+  if (shape.size() != kDim4) {
+    MS_LOG(ERROR) << "Conv2DBackpropFilter node output ori shape is: " << shape.size();
+    return nullptr;
+  }
+  auto filter_n = shape[kIndex0];
+  auto filter_c = shape[kIndex1];
+  auto filter_h = shape[kIndex2];
+  auto filter_w = shape[kIndex3];
+  auto matrix_size = filter_n * filter_c * filter_h * filter_w;
+  if (matrix_size <= 0 || filter_n % groups != 0) {
+    MS_LOG(ERROR) << "Conv2DBackpropFilter node shape value is error, matrix_size: " << matrix_size
+                  << ", shape: " << shape << ", groups: " << groups;
+    return nullptr;
+  }
+  // CreateAssitValueNode
+  auto value_node = CreateAssistNode(func_graph, node, shape, matrix_size);
+  MS_LOG(INFO) << "Create assist value node success.";
+  // CreateMulNode
+  std::vector<AnfNodePtr> mul_inputs{NewValueNode(std::make_shared<Primitive>(kMulOpName)), node, value_node};
+  CNodePtr mul_node = NewCNode(mul_inputs, func_graph);
+  MS_EXCEPTION_IF_NULL(mul_node);
+  mul_node->set_abstract(cnode->abstract());
+  mul_node->set_scope(cnode->scope());
+  MS_LOG(INFO) << "Create mul node success.";
+  common::AnfAlgo::SetNodeAttr(kAttrVisited, MakeValue(true), node);
+  return mul_node;
+}
+}  // namespace mindspore::opt

mindspore/ccsrc/plugin/device/ascend/optimizer/ir_fission/conv2d_backprop_filter_mul_fission.h

@@ -0,0 +1,35 @@
+/**
+ * Copyright 2022 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_OPTIMIZER_ASCEND_IR_CONV2D_BACKPROP_FILTER_MUL_H_
+#define MINDSPORE_CCSRC_OPTIMIZER_ASCEND_IR_CONV2D_BACKPROP_FILTER_MUL_H_
+
+#include <memory>
+#include "backend/common/optimizer/optimizer.h"
+#include "plugin/device/ascend/optimizer/ascend_helper.h"
+
+namespace mindspore {
+namespace opt {
+class Conv2dBackpropFilterMul : public PatternProcessPass {
+ public:
+  explicit Conv2dBackpropFilterMul(bool multigraph = true)
+      : PatternProcessPass("conv2d_backprop_filter_mul", multigraph) {}
+  ~Conv2dBackpropFilterMul() 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_OPTIMIZER_ASCEND_IR_CONV2D_BACKPROP_FILTER_MUL_H_

mindspore/ccsrc/utils/anfalgo.cc

@@ -1541,11 +1541,17 @@ int64_t AnfAlgo::GetAttrGroups(const AnfNodePtr &node, size_t index) {
       }
       return AnfAlgo::GetNodeAttr<int64_t>(cnode, kAttrFracZGroup);
     }
-  } else if (node->isa<Parameter>()) {
+  }
+  if (node->isa<Parameter>()) {
     auto param = node->cast<ParameterPtr>();
     MS_EXCEPTION_IF_NULL(param);
     return param->fracz_group();
   }
+  if (node->isa<ValueNode>()) {
+    auto value_node = node->cast<ValueNodePtr>();
+    MS_EXCEPTION_IF_NULL(value_node);
+    return value_node->fracz_group();
+  }
   return 1;
 }
 

mindspore/core/ir/anf.h

@@ -987,6 +987,16 @@ class MS_CORE_API ValueNode final : public ANode {
   /// \return The count of graphs using this ValueNode.
   size_t used_graph_count() const { return used_graph_count_; }
 
+  /// \brief Set the count of groups using this ValueNode.
+  ///
+  /// \param[in] group The count of groups using this ValueNode.
+  void set_fracz_group(int64_t group) { format_attr_.fracz_group = group; }
+
+  /// \brief Get groups attr in FracZ format.
+  ///
+  /// \return Groups attr in FracZ format.
+  int64_t fracz_group() const { return format_attr_.fracz_group; }
+
   /// \brief Set the count of graphs using this ValueNode.
   ///
   /// \param[in] used_graph_count The count of graphs using this ValueNode.
@@ -1010,6 +1020,12 @@ class MS_CORE_API ValueNode final : public ANode {
   }
 
  private:
+  struct FormatAttr {
+    int64_t fracz_group = 1;
+    int64_t input_size = 0;
+    int64_t hidden_size = 0;
+  };
+  FormatAttr format_attr_;
   ValuePtr value_;
   size_t used_graph_count_{0};
   bool has_new_value_ = false;