[RELAY][PASS] FoldScaleAxis Backward (#2024)

include/tvm/relay/expr_functor.h

@@ -135,9 +135,9 @@ class ExprVisitor
   void VisitExpr_(const TupleGetItemNode* op) override;
   virtual void VisitType(const Type& t);
 
- private:
-  // internal visited flag.
-  std::unordered_set<const Node*> visited_;
+ protected:
+  // Internal visiting counter
+  std::unordered_map<const Node*, size_t> visit_counter_;
 };
 
 /*!

python/tvm/relay/ir_pass.py

@@ -31,6 +31,29 @@ def infer_type(expr, env=None):
     return _ir_pass.infer_type(expr, env)
 
 
+def backward_fold_scale_axis(expr):
+    """Backward fold axis scaling into weights of conv2d/dense.
+
+    Parameters
+    ----------
+    expr : tvm.relay.Expr
+        The input expression, we expect that expr's types
+        should be fully inferred by infer_type.
+
+    Returns
+    -------
+    folded_expr : tvm.relay.Expr
+        The folded expression after transformation.
+
+    Note
+    ----
+    It is recommended to call backward_fold_scale_axis
+    before using forward_fold_scale_axis.
+    As backward folding targets common conv-bn pattern.
+    """
+    return _ir_pass.backward_fold_scale_axis(expr)
+
+
 def forward_fold_scale_axis(expr):
     """Fold the scaling of axis into weights of conv2d/dense.
 
@@ -44,6 +67,12 @@ def forward_fold_scale_axis(expr):
     -------
     folded_expr : tvm.relay.Expr
         The folded expression after transformation.
+
+    Note
+    ----
+    It is recommended to call backward_fold_scale_axis
+    before using forward_fold_scale_axis.
+    As backward folding targets common conv-bn pattern.
     """
     return _ir_pass.forward_fold_scale_axis(expr)
 

src/relay/ir/expr_functor.cc

@@ -160,10 +160,14 @@ Expr ExprMutator::VisitExpr_(const TupleGetItemNode* g) {
 Type ExprMutator::VisitType(const Type& t) { return t; }
 
 void ExprVisitor::VisitExpr(const Expr& expr) {
-  if (visited_.count(expr.get())) return;
-  using TParent = ExprFunctor<void(const Expr&)>;
-  TParent::VisitExpr(expr);
-  visited_.insert(expr.get());
+  auto it = visit_counter_.find(expr.get());
+  if (it != visit_counter_.end()) {
+    ++it->second;
+  } else {
+    using TParent = ExprFunctor<void(const Expr&)>;
+    TParent::VisitExpr(expr);
+    visit_counter_.insert({expr.get(), 1});
+  }
 }
 
 void ExprVisitor::ExprVisitor::VisitExpr_(const VarNode* op) {

src/relay/pass/fold_scale_axis.cc

@@ -24,9 +24,9 @@ namespace fold_scale_axis {
 using runtime::TypedPackedFunc;
 
 
-// FoldScaleAxisFoward algorithm:
+// FoldScaleAxis algorithm:
 //
-// The general idea is that we transform Expr to tuple of
+// The general idea is to transform Expr to tuple of
 // (value, axes, scale), where the final result satiesfies:
 //
 // result = value
@@ -41,9 +41,14 @@ using runtime::TypedPackedFunc;
 // we run a backward "preparation phase", which propagates the demand
 // of the potential axes scaling back to its input.
 //
-// The folding process is done in two steps:
+// Forward folding process is done in two steps:
 // - Prepare phase: backward propagation of demand.
 // - Transform phase: forward transformation,
+//
+// Similarly, backward folding process is done in two steps:
+// - Prepare phase: forward propagation of demand.
+// - Transform phase: transformation by push down the axes scale signal to inputs.
+//
 
 /*!
  * \brief sorted array axis, can also be nullptr.
@@ -99,7 +104,7 @@ ValueType GetFunc(const OpMap<ValueType>& op_map,
 }
 
 /*!
- * \brief Preparation function for for pass scale forward.
+ * \brief Preparation function for pass scale forward.
  * \param call The call node.
  * \param out_scale_axes Possible scaling on axes of the output.
  * \return The result scaling on axes of the input.
@@ -144,7 +149,7 @@ using FForwardTransform = TypedPackedFunc<
 //----------------------------------------------
 // Generic Visitors for FScaleAxisForward
 //----------------------------------------------
-class FScaleAxisForwardPrep : private ExprVisitor {
+class ForwardPrep : private ExprVisitor {
  public:
   std::unordered_map<const Node*, AxesSet>
   Prepare(const Expr& body) {
@@ -255,12 +260,12 @@ class FScaleAxisForwardPrep : private ExprVisitor {
   }
 };
 
-class FScaleAxisForwardTransform : private ExprMutator {
+class ForwardTransformer : private ExprMutator {
  public:
   // Transform expression.
-  Expr Transform(Expr expr) {
+  Expr Fold(Expr expr) {
     expected_scale_axes_ =
-        FScaleAxisForwardPrep().Prepare(expr);
+        ForwardPrep().Prepare(expr);
     return this->Mutate(expr);
   }
 
@@ -346,13 +351,13 @@ Array<AxesSet> ReluForwardPrep(const Call& call, AxesSet out) {
 }
 
 STuple ReluForwardTransform(const Call& ref_call,
-                              const AxesSet& expected_axes,
-                              const Array<STuple>& sargs) {
+                            const AxesSet& expected_axes,
+                            const Array<STuple>& sargs) {
   if (!sargs[0]->axes.defined()) return STuple();
   // return transformed conv2d
   auto rnode = make_node<STupleNode>();
   rnode->value = CallNode::make(
-      ref_call->op, {sargs[0]->value}, ref_call->attrs, {});
+      ref_call->op, {sargs[0]->value}, ref_call->attrs, ref_call->type_args);
   rnode->scale = sargs[0]->scale;
   rnode->axes = sargs[0]->axes;
   return STuple(rnode);
@@ -474,8 +479,6 @@ Array<AxesSet> Conv2DForwardPrep(const Call& call, AxesSet out) {
   Layout weight_layout(param->weight_layout);
   int c_big_axis = data_layout.indexof('C');
   int c_small_axis = data_layout.indexof('c');
-  const auto* tdata = call->args[0]->type_as<TensorTypeNode>();
-  CHECK(tdata) << "require checked type";
 
   CHECK_GE(c_big_axis, 0);
   AxesSet data_axes = NullValue<AxesSet>();
@@ -486,8 +489,7 @@ Array<AxesSet> Conv2DForwardPrep(const Call& call, AxesSet out) {
   //
   // only handle depthwise or full conv2d.
   // TODO(tvm-team) handle grouped conv by reshape + bcast
-  bool is_depthwise_conv2d =
-      is_const_int(tdata->shape[c_big_axis], param->groups);
+  bool is_depthwise_conv2d = IsDepthwiseConv2D(call, param, weight_layout);
   if (weight_layout.indexof('i') < 0 &&
       c_small_axis < 0 &&
       (param->groups == 1 || is_depthwise_conv2d)) {
@@ -515,18 +517,24 @@ STuple Conv2DForwardTransform(const Call& ref_call,
   CHECK_EQ(weight_layout.indexof('i'), -1);
   CHECK(sdata->axes.size() == 1 &&
         c_big_axis == sdata->axes[0]->value);
+  int big_oc_axis = weight_layout.indexof('O');
   int big_ic_axis = weight_layout.indexof('I');
 
-  const auto* tdata = ref_call->args[0]->type_as<TensorTypeNode>();
   // Check it must be depthwise or full conv2d.
-  bool is_depthwise_conv2d =
-      is_const_int(tdata->shape[c_big_axis], param->groups);
+  bool is_depthwise_conv2d = IsDepthwiseConv2D(ref_call, param, weight_layout);
   CHECK(param->groups == 1 || is_depthwise_conv2d);
+  Expr weight = sargs[1]->value;
 
   // match the ic_axis
-  Expr scale = ExpandBiasToMatchAxis(
-      sdata->scale, weight_layout.ndim(), {big_ic_axis});
-  Expr weight = Multiply(sargs[1]->value, scale);
+  if (is_depthwise_conv2d) {
+    Expr scale = ExpandBiasToMatchAxis(
+        sdata->scale, weight_layout.ndim(), {big_oc_axis});
+    weight = Multiply(weight, scale);
+  } else {
+    Expr scale = ExpandBiasToMatchAxis(
+        sdata->scale, weight_layout.ndim(), {big_ic_axis});
+    weight = Multiply(weight, scale);
+  }
   // return transformed conv2d
   auto rnode = make_node<STupleNode>();
   rnode->value = CallNode::make(
@@ -542,13 +550,416 @@ RELAY_REGISTER_OP("nn.conv2d")
 
 
 Expr ForwardFoldScaleAxis(Expr data) {
-  return FScaleAxisForwardTransform().Transform(data);
+  return ForwardTransformer().Fold(data);
 }
 
 // Expose the FoldScaleAxisFoward
 TVM_REGISTER_API("relay._ir_pass.forward_fold_scale_axis")
 .set_body_typed<Expr(Expr)>(ForwardFoldScaleAxis);
 
+//----------------------------------------
+// Implement backward transformations.
+//----------------------------------------
+class BackwardTransformer;
+
+/*!
+ * \brief Preparation function for for pass scale backward.
+ * \param call The call node.
+ * \param in_scale_axes Allowed input scaling.
+ * \return The result scaling on axes of the input.
+ */
+using FBackwardPrep = TypedPackedFunc<
+  AxesSet(const Call& call, const Array<AxesSet>& in_scale_axes)>;
+
+using FBackwardTransform = TypedPackedFunc<
+  Expr(const Call& call,
+       const AxesSet& axes,
+       const Expr& scale,
+       const BackwardTransformer& transformer)>;
+
+//----------------------------------------------
+// Generic Visitors for FScaleAxisBackward
+//----------------------------------------------
+/*!
+ * \brief Get reference counter of each internal ExprNode in body.
+ * \param body The body expression.
+ * \return The reference count mapping.
+ */
+std::unordered_map<const Node*, size_t>
+GetExprRefCount(const Expr& body) {
+  class ExprRefCounter : private ExprVisitor {
+   public:
+    std::unordered_map<const Node*, size_t>
+    Get(const Expr& body) {
+      this->VisitExpr(body);
+      return std::move(this->visit_counter_);
+    }
+  };
+  return ExprRefCounter().Get(body);
+}
+
+class BackwardPrep : private ExprVisitor {
+ public:
+  // The message on each node.
+  std::unordered_map<const Node*, AxesSet>
+  Prepare(const Expr& body) {
+    ref_counter_ = GetExprRefCount(body);
+    this->VisitExpr(body);
+    return std::move(message_);
+  }
+
+ private:
+  // The message on each node.
+  std::unordered_map<const Node*, AxesSet> message_;
+  // reference counter of an internal expr
+  std::unordered_map<const Node*, size_t> ref_counter_;
+  // Visit the expression.
+  void VisitExpr_(const CallNode* call) {
+    ExprVisitor::VisitExpr_(call);
+    static const auto& fprep =
+        Op::GetAttr<FBackwardPrep>("FScaleAxisBackwardPrep");
+    auto f = GetFunc(fprep, call->op);
+    if (f == nullptr) return;
+    auto rit = ref_counter_.find(call);
+    CHECK(rit != ref_counter_.end());
+    // We only allow propagation of scale backward
+    // if the expression is only referred by a single parent.
+    if (rit->second != 1) return;
+    Array<AxesSet> in_axes;
+    for (Expr arg : call->args) {
+      auto it = message_.find(arg.get());
+      if (it != message_.end()) {
+        in_axes.push_back(it->second);
+      } else {
+        in_axes.push_back(NullValue<AxesSet>());
+      }
+    }
+    AxesSet out_axes = f(GetRef<Call>(call), in_axes);
+    if (out_axes.defined()) {
+      message_[call] = out_axes;
+    }
+  }
+};
+
+class BackwardTransformerNode :
+      public Node,
+      private ExprMutator {
+ public:
+  // Run forward transform.
+  Expr Fold(Expr expr) {
+    expected_scale_axes_ = BackwardPrep().Prepare(expr);
+    return this->Mutate(expr);
+  }
+  /*!
+   * \brief Transform the expr to consider the scaling.
+   *
+   * \param expr The input expression.
+   * \param axes The axes to scale.
+   * \param scale The scale applied to the axes.
+   * \return The result of transformation.
+   */
+  Expr Transform(const Expr& expr, AxesSet axes, Expr scale) {
+    // NOTE: the result of Transform is not memoized.
+    // However, in the current rule, Transform will
+    // only be called to expr that is referred once.
+    if (const CallNode* call_node = expr.as<CallNode>()) {
+      return Transform(call_node, axes, scale);
+    } else {
+      CHECK(!axes.defined()) << "outstanding scale";
+      return ExprMutator::VisitExpr(expr);
+    }
+  }
+  /*!
+   * \brief Normal way of mutating call node.
+   * \param call_node The call node to be mutated.
+   * \return the result of the call Mutation.
+   */
+  Expr NormalCallTransform(const CallNode* call_node) {
+    return ExprMutator::VisitExpr_(call_node);
+  }
+  /*!
+   * \brief Get the expected axes on expr.
+   * \param expr The expresison.
+   * \return The expected axes.
+   */
+  AxesSet GetExpectedAxes(const Expr& expr) const {
+    auto it = expected_scale_axes_.find(expr.get());
+    if (it != expected_scale_axes_.end()) return it->second;
+    return NullValue<AxesSet>();
+  }
+
+  // solver is not serializable.
+  void VisitAttrs(tvm::AttrVisitor* v) final {}
+
+  static constexpr const char* _type_key = "relay.fold_scale_axis.FBackwardTransformer";
+  TVM_DECLARE_NODE_TYPE_INFO(BackwardTransformerNode, Node);
+
+ private:
+  // Valid axes on each node.
+  std::unordered_map<const Node*, AxesSet> expected_scale_axes_;
+  // Override mutation of call.
+  Expr VisitExpr_(const CallNode* call_node) final {
+    return Transform(call_node, NullValue<AxesSet>(), NullValue<Expr>());
+  }
+  // Transform of CallNode.
+  Expr Transform(const CallNode* call_node, AxesSet axes, Expr scale);
+};
+
+class BackwardTransformer : public NodeRef {
+ public:
+  BackwardTransformer() {}
+  explicit BackwardTransformer(
+      ::tvm::NodePtr<::tvm::Node> n) : NodeRef(n) {
+  }
+  BackwardTransformerNode* operator->() const {
+    return static_cast<BackwardTransformerNode*>(node_.get());
+  }
+  using ContainerType = BackwardTransformerNode;
+};
+
+Expr BackwardTransformerNode::Transform(
+    const CallNode* call_node, AxesSet axes, Expr scale) {
+  static const auto& ftransform =
+      Op::GetAttr<FBackwardTransform>("FScaleAxisBackwardTransform");
+  auto f = GetFunc(ftransform, call_node->op);
+  if (f != nullptr) {
+    return f(GetRef<Call>(call_node),
+             axes,
+             scale,
+             GetRef<BackwardTransformer>(this));
+  } else {
+    CHECK(!axes.defined()) << "outstanding scale";
+    return NormalCallTransform(call_node);
+  }
+}
+
+
+//----------------------------------------------
+// Per operator defs for FScaleAxisForward
+//----------------------------------------------
+
+// Intermediate operators
+AxesSet ReluBackwardPrep(const Call& call, const Array<AxesSet>& in_axes) {
+  return in_axes[0];
+}
+
+Expr ReluBackwardTransform(const Call& call,
+                           const AxesSet& axes,
+                           const Expr& scale,
+                           const BackwardTransformer& transformer) {
+  if (!axes.defined()) {
+    return transformer->NormalCallTransform(call.operator->());
+  }
+  Expr input = transformer->Transform(
+      call->args[0], axes, scale);
+  return CallNode::make(call->op, {input}, call->attrs, call->type_args);
+}
+
+RELAY_REGISTER_OP("nn.relu")
+.set_attr<FBackwardPrep>("FScaleAxisBackwardPrep", ReluBackwardPrep);
+
+RELAY_REGISTER_OP("nn.relu")
+.set_attr<FBackwardTransform>("FScaleAxisBackwardTransform", ReluBackwardTransform);
+
+RELAY_REGISTER_OP("nn.leaky_relu")
+.set_attr<FBackwardPrep>("FScaleAxisBackwardPrep", ReluBackwardPrep);
+
+RELAY_REGISTER_OP("nn.leaky_relu")
+.set_attr<FBackwardTransform>("FScaleAxisBackwardTransform", ReluBackwardTransform);
+
+// AddSub
+AxesSet AddSubBackwardPrep(const Call& call, const Array<AxesSet>& in_axes) {
+  const auto* tlhs = call->args[0]->type_as<TensorTypeNode>();
+  const auto* trhs = call->args[1]->type_as<TensorTypeNode>();
+  AttrsEqual equal;
+  if (in_axes[0].defined() &&
+      MatchBroadcastToLeftAxes(tlhs, trhs, in_axes[0])) {
+    return in_axes[0];
+  } else if (in_axes[1].defined() &&
+             MatchBroadcastToLeftAxes(trhs, tlhs, in_axes[1])) {
+    return in_axes[1];
+  } else if (in_axes[0].defined() &&
+             in_axes[1].defined() &&
+             equal(in_axes[0], in_axes[1]) &&
+             equal(tlhs->shape, trhs->shape)) {
+    // add of two elements.
+    return in_axes[0];
+  } else {
+    return NullValue<AxesSet>();
+  }
+}
+
+Expr AddSubBackwardTransform(const Call& call,
+                             const AxesSet& axes,
+                             const Expr& scale,
+                             const BackwardTransformer& transformer) {
+  const auto* tlhs = call->args[0]->type_as<TensorTypeNode>();
+  const auto* trhs = call->args[1]->type_as<TensorTypeNode>();
+  if (!axes.defined()) {
+    return transformer->NormalCallTransform(call.operator->());
+  }
+  AxesSet lhs_axes = transformer->GetExpectedAxes(call->args[0]);
+  AxesSet rhs_axes = transformer->GetExpectedAxes(call->args[1]);
+  AttrsEqual equal;
+
+  if (lhs_axes.defined() && rhs_axes.defined()) {
+    CHECK(equal(lhs_axes, rhs_axes));
+    CHECK(equal(axes, lhs_axes));
+    Expr lhs = transformer->Transform(call->args[0], axes, scale);
+    Expr rhs = transformer->Transform(call->args[1], axes, scale);
+    return CallNode::make(call->op, {lhs, rhs}, call->attrs, call->type_args);
+  } else if (lhs_axes.defined()) {
+    CHECK(equal(axes, lhs_axes));
+    Expr lhs = transformer->Transform(call->args[0], axes, scale);
+    Expr rhs = transformer->Transform(
+        call->args[1], NullValue<AxesSet>(), NullValue<Expr>());
+    Expr rhs_scale = ExpandBiasToMatchAxis(
+        scale, tlhs->shape.size(), axes);
+    rhs = Multiply(rhs, rhs_scale);
+    return CallNode::make(call->op, {lhs, rhs}, call->attrs, call->type_args);
+  } else if (rhs_axes.defined()) {
+    CHECK(equal(axes, rhs_axes));
+    Expr lhs = transformer->Transform(
+        call->args[0], NullValue<AxesSet>(), NullValue<Expr>());
+    Expr rhs = transformer->Transform(call->args[1], axes, scale);
+    Expr lhs_scale = ExpandBiasToMatchAxis(
+        scale, trhs->shape.size(), axes);
+    lhs = Multiply(lhs, lhs_scale);
+    return CallNode::make(call->op, {lhs, rhs}, call->attrs, call->type_args);
+  } else {
+    LOG(FATAL) << "outstanding scale";
+    return Expr();
+  }
+}
+
+RELAY_REGISTER_OP("add")
+.set_attr<FBackwardPrep>("FScaleAxisBackwardPrep", AddSubBackwardPrep);
+
+RELAY_REGISTER_OP("add")
+.set_attr<FBackwardTransform>("FScaleAxisBackwardTransform", AddSubBackwardTransform);
+
+RELAY_REGISTER_OP("subtract")
+.set_attr<FBackwardPrep>("FScaleAxisBackwardPrep", AddSubBackwardPrep);
+
+RELAY_REGISTER_OP("subtract")
+.set_attr<FBackwardTransform>("FScaleAxisBackwardTransform", AddSubBackwardTransform);
+
+// Producer operators
+// Multiply produces the scale-axis pair.
+Expr MultiplyBackwardTransform(const Call& call,
+                               const AxesSet& axes,
+                               const Expr& scale,
+                               const BackwardTransformer& transformer) {
+  CHECK(!axes.defined()) << "outstanding scale";
+  const auto* tlhs = call->args[0]->type_as<TensorTypeNode>();
+  const auto* trhs = call->args[1]->type_as<TensorTypeNode>();
+  AxesSet lhs_axes = transformer->GetExpectedAxes(call->args[0]);
+  AxesSet rhs_axes = transformer->GetExpectedAxes(call->args[1]);
+  if (lhs_axes.defined()) {
+    // NOTE we won't recursively call mutating on scale part.
+    // since there  won't be scale chance within scale part.
+    Expr rhs = call->args[1];
+    if (MatchBroadcastToLeftAxes(tlhs, trhs, lhs_axes, &rhs)) {
+      return transformer->Transform(call->args[0], lhs_axes, rhs);
+    }
+  } else if (rhs_axes.defined()) {
+    Expr lhs = call->args[0];
+    if (MatchBroadcastToLeftAxes(trhs, tlhs, rhs_axes, &lhs)) {
+      return transformer->Transform(call->args[1], rhs_axes, lhs);
+    }
+  }
+  return transformer->NormalCallTransform(call.operator->());
+}
+
+RELAY_REGISTER_OP("multiply")
+.set_attr<FBackwardTransform>("FScaleAxisBackwardTransform", MultiplyBackwardTransform);
+
+// Consumer operators
+// Conv2D send out requirement of axis folding.
+AxesSet Conv2DBackwardPrep(const Call& call, const Array<AxesSet>& in_axes) {
+  const auto* param = call->attrs.as<Conv2DAttrs>();
+  CHECK(param != nullptr);
+  Layout out_layout(param->out_layout);
+  if (!out_layout.defined()) {
+    out_layout = Layout(param->data_layout);
+  }
+  Layout weight_layout(param->weight_layout);
+  int c_big_axis = out_layout.indexof('C');
+  int c_small_axis = out_layout.indexof('c');
+
+  CHECK_GE(c_big_axis, 0);
+  // For now, we only support simple pattern (no folded weight/data)
+  // More general layout can be supported under the current framework.
+  // By using a unified layout transformation.
+  // We only need to change the Prep and Mutate function.
+  //
+  // only handle depthwise or full conv2d.
+  // TODO(tvm-team) handle grouped conv by reshape + bcast
+  bool is_depthwise_conv2d = IsDepthwiseConv2D(call, param, weight_layout);
+  if (weight_layout.indexof('o') < 0 &&
+      weight_layout.indexof('i') < 0 &&
+      c_small_axis < 0 &&
+      (param->groups == 1 || is_depthwise_conv2d)) {
+    return {c_big_axis};
+  } else {
+    return NullValue<AxesSet>();
+  }
+}
+
+// Conv2D consumes the scale axis during transformation.
+Expr Conv2DBackwardTransform(const Call& call,
+                             const AxesSet& axes,
+                             const Expr& scale,
+                             const BackwardTransformer& transformer) {
+  if (!axes.defined()) {
+    return transformer->NormalCallTransform(call.operator->());
+  }
+  const auto* param = call->attrs.as<Conv2DAttrs>();
+  CHECK(param != nullptr);
+  Layout out_layout(param->out_layout);
+  if (!out_layout.defined()) {
+    out_layout = Layout(param->data_layout);
+  }
+  Layout weight_layout(param->weight_layout);
+  int c_big_axis = out_layout.indexof('C');
+  CHECK_GE(c_big_axis, 0);
+  // For now, we only support simple pattern (no folded weight/data)
+  // TODO(tvm-team) support general data layout
+  CHECK_EQ(weight_layout.indexof('o'), -1);
+  CHECK_EQ(weight_layout.indexof('i'), -1);
+  CHECK(axes.size() == 1 &&
+        c_big_axis == axes[0]->value);
+
+  int big_oc_axis = weight_layout.indexof('O');
+  // Check it must be depthwise or full conv2d.
+  bool is_depthwise_conv2d = IsDepthwiseConv2D(call, param, weight_layout);
+  CHECK(param->groups == 1 || is_depthwise_conv2d);
+
+  Expr data = transformer->Transform(
+      call->args[0], NullValue<AxesSet>(), NullValue<Expr>());
+  Expr weight = transformer->Transform(
+      call->args[1], NullValue<AxesSet>(), NullValue<Expr>());
+  // scale on input for deptwise.
+  Expr wscale = ExpandBiasToMatchAxis(
+      scale, weight_layout.ndim(), {big_oc_axis});
+  weight = Multiply(weight, wscale);
+  return CallNode::make(
+      call->op, {data, weight}, call->attrs, call->type_args);
+}
+
+RELAY_REGISTER_OP("nn.conv2d")
+.set_attr<FBackwardPrep>("FScaleAxisBackwardPrep", Conv2DBackwardPrep);
+
+RELAY_REGISTER_OP("nn.conv2d")
+.set_attr<FBackwardTransform>("FScaleAxisBackwardTransform", Conv2DBackwardTransform);
+
+Expr BackwardFoldScaleAxis(Expr data) {
+  return make_node<BackwardTransformerNode>()->Fold(data);
+}
+
+TVM_REGISTER_API("relay._ir_pass.backward_fold_scale_axis")
+.set_body_typed<Expr(Expr)>(BackwardFoldScaleAxis);
+
 }  // namespace fold_scale_axis
 }  // namespace relay
 }  // namespace tvm

src/relay/pass/pattern_util.h

@@ -11,6 +11,7 @@
 #include <tvm/relay/op.h>
 #include <tvm/relay/expr.h>
 #include <tvm/relay/attrs/transform.h>
+#include "../op/nn/layout.h"
 
 namespace tvm {
 namespace relay {
@@ -100,11 +101,31 @@ inline Expr ExpandBiasToMatchAxis(Expr bias,
   return bias;
 }
 
+/*!
+ * \brief Check if the call is depthwise conv2d.
+ *
+ * \param call The conv2d call.
+ * \param param The conv2d attributes.
+ * \return Whether it is depthwise_conv2d.
+ */
+inline bool IsDepthwiseConv2D(const Call& call,
+                              const Conv2DAttrs* param,
+                              const Layout& weight_layout) {
+  static const Layout kOIHW("OIHW");
+  auto wshape = ConvertLayout(
+      call->args[1]->type_as<TensorTypeNode>()->shape,
+      weight_layout, kOIHW);
+  return is_const_int(wshape[0], param->groups) &&
+      is_const_int(wshape[1], 1);
+}
+
+
 inline Expr Multiply(Expr lhs, Expr rhs) {
   static const Op& op = Op::Get("multiply");
   return CallNode::make(op, {lhs, rhs}, Attrs(), {});
 }
 
+
 inline Expr Divide(Expr lhs, Expr rhs) {
   static const Op& op = Op::Get("divide");
   return CallNode::make(op, {lhs, rhs}, Attrs(), {});
@@ -116,8 +137,6 @@ inline Expr ReshapeLike(Expr lhs, Expr rhs) {
   return CallNode::make(op, {lhs, rhs}, Attrs(), {});
 }
 
-
-
 }  // namespace relay
 }  // namespace tvm
 #endif  // TVM_RELAY_PASS_PATTERN_UTIL_H_

tests/python/relay/test_pass_fold_scale_axis.py

@@ -62,14 +62,14 @@ def test_fold_fwd_dual_path():
                              channels=channels,
                              kernel_size=(3, 3),
                              data_layout="NHWC",
-                             weight_layout="HWOI",
+                             weight_layout="HWIO",
                              groups=channels,
                              padding=(1, 1))
         y2 = relay.nn.conv2d(x, conv_weight,
                              channels=channels,
                              kernel_size=(3, 3),
                              data_layout="NHWC",
-                             weight_layout="HWOI",
+                             weight_layout="HWIO",
                              groups=channels,
                              padding=(1, 1))
         z = relay.add(y1, y2)
@@ -85,7 +85,7 @@ def test_fold_fwd_dual_path():
                              channels=channels,
                              kernel_size=(3, 3),
                              data_layout="NHWC",
-                             weight_layout="HWOI",
+                             weight_layout="HWIO",
                              groups=channels,
                              padding=(1, 1))
         y2 = relay.nn.conv2d(x,
@@ -93,7 +93,7 @@ def test_fold_fwd_dual_path():
                              channels=channels,
                              kernel_size=(3, 3),
                              data_layout="NHWC",
-                             weight_layout="HWOI",
+                             weight_layout="HWIO",
                              groups=channels,
                              padding=(1, 1))
         z = relay.add(y1, y2)
@@ -147,7 +147,176 @@ def test_fold_fwd_fail():
     check((2, 11, 10, 4), 4)
 
 
+def test_fold_bwd_simple():
+    """Simple testcase."""
+    def before(x, conv_weight, out_bias, out_scale, channels):
+        args = [x, conv_weight, out_bias, out_scale]
+        out_scale = relay.expand_dims(out_scale, axis=1, num_newaxis=2)
+        out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
+        y = relay.nn.conv2d(x, conv_weight,
+                            channels=channels,
+                            kernel_size=(3, 3),
+                            padding=(1, 1))
+        y = relay.add(y, out_bias)
+        y = relay.nn.relu(y)
+        y = relay.multiply(y, out_scale)
+        return relay.Function(args, y)
+
+    def expected(x, conv_weight, out_bias, out_scale, channels):
+        # use a fixed order of args so alpha equal check can pass
+        args = [x, conv_weight, out_bias, out_scale]
+        out_scale = relay.expand_dims(out_scale, axis=1, num_newaxis=2)
+        out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
+        squeezed_scale = relay.squeeze(out_scale, axis=[1,2])
+        conv_weight = relay.multiply(
+            conv_weight , relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
+
+        y = relay.nn.conv2d(x, conv_weight,
+                            channels=channels,
+                            kernel_size=(3, 3),
+                            padding=(1, 1))
+        out_bias = relay.multiply(out_bias,
+                                  relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2))
+        y = relay.add(y, out_bias)
+        y = relay.nn.relu(y)
+        return relay.Function(args, y)
+
+    def check(shape, channels):
+        x =  relay.var("x", shape=shape)
+        in_channels = shape[1]
+        weight = relay.var("weight")
+        out_bias = relay.var("out_bias", shape=(channels,))
+        out_scale = relay.var("out_scale", shape=(channels,))
+
+        y1 = before(x, weight, out_bias, out_scale, channels)
+        y1 = relay.ir_pass.infer_type(y1)
+        type_dict = {x.name_hint:x.checked_type for x in y1.params}
+        weight = relay.var("weight", type_dict["weight"])
+        y1_folded = relay.ir_pass.backward_fold_scale_axis(y1)
+        y1_expected = expected(x, weight, out_bias, out_scale, channels)
+        assert relay.ir_pass.alpha_equal(y1_folded, y1_expected)
+
+    check((2, 4, 10, 10), 8)
+
+
+def test_fold_bwd_dual_path():
+    """Dual path testcase."""
+    def before(x, conv_weight, out_bias, out_scale, channels):
+        args = [x, conv_weight, out_bias, out_scale]
+        out_scale = relay.expand_dims(out_scale, axis=1, num_newaxis=2)
+        out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
+        y1 = relay.nn.conv2d(x, conv_weight,
+                             channels=channels,
+                             kernel_size=(3, 3),
+                             padding=(1, 1))
+        y1 = relay.nn.relu(y1)
+        y2 = relay.nn.conv2d(x, conv_weight,
+                             channels=channels,
+                             kernel_size=(3, 3),
+                             padding=(1, 1))
+        y2 = relay.nn.relu(y2)
+        y = relay.add(y1, y2)
+        y = relay.multiply(y, out_scale)
+        return relay.Function(args, y)
+
+    def expected(x, conv_weight, out_bias, out_scale, channels):
+        # use a fixed order of args so alpha equal check can pass
+        args = [x, conv_weight, out_bias, out_scale]
+        out_scale = relay.expand_dims(out_scale, axis=1, num_newaxis=2)
+        out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
+        squeezed_scale = relay.squeeze(out_scale, axis=[1,2])
+        def fold_conv_weight():
+            return  relay.multiply(
+                conv_weight ,
+                relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
+        y1 = relay.nn.conv2d(x, fold_conv_weight(),
+                            channels=channels,
+                            kernel_size=(3, 3),
+                            padding=(1, 1))
+        y1 = relay.nn.relu(y1)
+        y2 = relay.nn.conv2d(x, fold_conv_weight(),
+                            channels=channels,
+                            kernel_size=(3, 3),
+                            padding=(1, 1))
+        y2 = relay.nn.relu(y2)
+        y = relay.add(y1, y2)
+        return relay.Function(args, y)
+
+    def check(shape, channels):
+        x =  relay.var("x", shape=shape)
+        in_channels = shape[1]
+        weight = relay.var("weight")
+        out_bias = relay.var("out_bias", shape=(channels,))
+        out_scale = relay.var("out_scale", shape=(channels,))
+
+        y1 = before(x, weight, out_bias, out_scale, channels)
+        y1 = relay.ir_pass.infer_type(y1)
+        type_dict = {x.name_hint:x.checked_type for x in y1.params}
+        weight = relay.var("weight", type_dict["weight"])
+        y1_folded = relay.ir_pass.backward_fold_scale_axis(y1)
+        y1_expected = expected(x, weight, out_bias, out_scale, channels)
+        assert relay.ir_pass.alpha_equal(y1_folded, y1_expected)
+
+    check((2, 4, 10, 10), 8)
+
+
+def test_fold_bwd_fail():
+    """Dual path testcase."""
+    def fail1(x, conv_weight, out_bias, out_scale, channels):
+        args = [x, conv_weight, out_bias, out_scale]
+        out_scale = relay.expand_dims(out_scale, axis=1, num_newaxis=2)
+        out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
+        y1 = relay.nn.conv2d(x, conv_weight,
+                             channels=channels,
+                             kernel_size=(3, 3),
+                             padding=(1, 1))
+        y1 = relay.nn.relu(y1)
+        y2 = relay.nn.conv2d(x, conv_weight,
+                             channels=channels,
+                             kernel_size=(3, 3),
+                             padding=(1, 1),
+                             out_layout="CNHW")
+        # fold will fail because the axis from two path
+        # differs from each other.
+        y2 = relay.nn.relu(y2)
+        y = relay.add(y1, y2)
+        y = relay.multiply(y, out_scale)
+        return relay.Function(args, y)
+
+    def fail2(x, conv_weight, out_bias, out_scale, channels):
+        args = [x, conv_weight, out_bias, out_scale]
+        out_scale = relay.expand_dims(out_scale, axis=1, num_newaxis=2)
+        out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
+        y1 = relay.nn.conv2d(x, conv_weight,
+                             channels=channels,
+                             kernel_size=(3, 3),
+                             padding=(1, 1))
+        y2 = relay.nn.relu(y1)
+        # fold will fail because y1 is referred also by y2
+        y1 = relay.multiply(y1, out_scale)
+        y = relay.add(y1, y2)
+        return relay.Function(args, y)
+
+
+    def check(shape, channels, fbefore):
+        x =  relay.var("x", shape=shape)
+        in_channels = shape[1]
+        weight = relay.var("weight")
+        out_bias = relay.var("out_bias", shape=(channels,))
+        out_scale = relay.var("out_scale", shape=(channels,))
+        y1 = fbefore(x, weight, out_bias, out_scale, channels)
+        y1 = relay.ir_pass.infer_type(y1)
+        y1_folded = relay.ir_pass.backward_fold_scale_axis(y1)
+        assert relay.ir_pass.alpha_equal(y1_folded, y1)
+
+    check((4, 4, 10, 10), 4, fail1)
+    check((4, 4, 10, 10), 4, fail2)
+
+
 if __name__ == "__main__":
     test_fold_fwd_simple()
     test_fold_fwd_dual_path()
     test_fold_fwd_fail()
+    test_fold_bwd_simple()
+    test_fold_bwd_dual_path()
+    test_fold_bwd_fail()