Batchnorm op Fusion in TVM

mbrookhart · April 1, 2022, 8:24pm

If you add CanonicalizeOps to your pass list before Simplify Expr, it will turn the bias_add to an add

apache/tvm/blob/main/src/relay/transforms/canonicalize_ops.cc#L36-L69


class BiasAddSimplifier : public ExprRewriter {
 public:
  BiasAddSimplifier() : bias_add_op_(Op::Get("nn.bias_add")) {}


  Expr Rewrite_(const CallNode* n, const Expr& post) override {
    auto new_n = post;
    if (n->op == bias_add_op_) {
      Call call = Downcast<Call>(new_n);
      ICHECK_EQ(call->args.size(), 2);
      const BiasAddAttrs* param = call->attrs.as<BiasAddAttrs>();


      auto ttype = n->args[0]->type_as<TensorTypeNode>();
      size_t n_dim = ttype->shape.size();
      int axis = param->axis;
      if (axis < 0) {
        axis += n_dim;
      }
      Expr expanded_bias = ExpandBiasToMatchAxis(call->args[1], n_dim, {axis});
      Expr ret = Add(call->args[0], expanded_bias);
      ret->checked_type_ = n->checked_type_;

This file has been truncated. show original