import onnx

from PIL import Image

import numpy as np

import tvm.relay as relay

import tvm

from tvm.contrib import graph_executor

model_path = "/home/xianmu/module/resnet18.onnx"

save_path = "/home/xianmu/module/pythonSave/"

onnx_model = onnx.load(model_path)

img_path = "/home/xianmu/.tvm_test_data/data/imagenet_cat.png"

resized_image = Image.open(img_path).resize((224, 224))

img_data = np.asarray(resized_image).astype("float32")

img_data = np.transpose(img_data, (2, 0, 1))

imagenet_mean = np.array([0.485, 0.456, 0.406]).reshape((3, 1, 1))

imagenet_stddev = np.array([0.229, 0.224, 0.225]).reshape((3, 1, 1))

norm_img_data = (img_data / 255 - imagenet_mean) / imagenet_stddev

img_data = np.expand_dims(norm_img_data, axis=0)

input_name = "data"

target = tvm.target.Target(target="llvm", host="llvm")

shape_dict = {input_name: img_data.shape}

mod, params = relay.frontend.from_onnx(onnx_model, shape_dict, export_node_renamed_model_path=save_path)  

with tvm.transform.PassContext(opt_level=3):

    lib = relay.build(mod, target=target, params=params)   

lib_fname = save_path + "mod.so"

lib.get_lib().export_library(lib_fname)

params_fname = save_path + "mod.params"

with open(params_fname, "wb") as param_file:

    param_file.write(relay.save_param_dict(lib.get_params()))

json_fname = save_path + "mod.json"

with open(json_fname, "w") as json_file:

    json_file.write(lib.get_executor_config())

dev = tvm.device(str(target), 0)

module = graph_executor.GraphModule(lib["default"](dev))    

module.set_input(input_name, img_data)

module.run()

output_shape = (1, 1000)

tvm_output = module.get_output(0, tvm.nd.empty(output_shape)).numpy()

print(tvm_output)

from tvm.relay import transform

print("===========================================")

print(mod)  

mod = transform.FuseOps()(mod)

print("===========================================")

print(mod)  

def @main(%data: Tensor[(1, 3, 224, 224), float32]) -> Tensor[(1, 10), float32] {

  %0 = nn.conv2d(%data, meta[relay.Constant][0], padding=[0, 0, 0, 0], channels=4, kernel_size=[3, 3]) ;

  %1 = nn.bias_add(%0, meta[relay.Constant][1] );

  %2 = nn.relu(%1) ;

  %3 = nn.max_pool2d(%2, pool_size=[3, 3], strides=[2, 2], padding=[1, 1, 1, 1]) ;

  %4 = nn.conv2d(%3, meta[relay.Constant][2], padding=[0, 0, 0, 0], channels=5, kernel_size=[3, 3]) ;

  %5 = nn.bias_add(%4, meta[relay.Constant][3] ) ;

  ...

  }

def @main(%data: Tensor[(1, 3, 224, 224), float32]) -> Tensor[(1, 10), float32] {

  %5 = fn (%p05: Tensor[(1, 3, 224, 224), float32] , %p12: Tensor[(4, 3, 3, 3), float32] , %p22: Tensor[(4), float32], Primitive=1) -> Tensor[(1, 4, 222, 222), float32] { 

    %3 = nn.conv2d(%p05, %p12, padding=[0, 0, 0, 0], channels=4, kernel_size=[3, 3]);

    %4 = nn.bias_add(%3, %p22) ;

    nn.relu(%4) 

  } ;

  %6 = %5(%data, meta[relay.Constant][0], meta[relay.Constant][1] ) ;

  ...

  }

Pass FuseOps(int fuse_opt_level) {

 runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func = 

     [=](Function f, IRModule m, PassContext pc) {

       bool link_params = false;

       Executor executor =

           m->GetAttr<Executor>(tvm::attr::kExecutor).value_or(NullValue<Executor>());

       link_params = executor.defined()

                         ? executor->attrs.GetAttr<Bool>("link-params").value_or(Bool(link_params))

                         : link_params;

       link_params = pc->GetConfig("relay.FuseOps.link_params", Bool(link_params)).value();

       int opt_level = fuse_opt_level == -1 ? pc->opt_level : fuse_opt_level;

       auto max_fuse_depth = pc->GetConfig("relay.FuseOps.max_depth", Integer(kMaxFusedOps));

       auto target = Target::Current();

       size_t max_function_args =

           (target.defined())

               ? target->GetAttr<Integer>("max_function_args", Integer(0)).value().IntValue()

               : 0;

       return Downcast<Function>(FuseOps(f, opt_level, max_fuse_depth.value().IntValue(),

                                         max_function_args, link_params, m));

     };

 return CreateFunctionPass(pass_func, 0, "FuseOps", {"InferType"});  

}

Expr FuseOps(const Expr& expr, int fuse_opt_level, size_t max_fuse_depth, size_t max_function_args,

            bool link_params, const IRModule& module) {

 return FuseMutator(fuse_opt_level, max_fuse_depth, max_function_args, link_params)

     .Transform(expr);

}

 Expr Transform(const Expr& body, int fuse_opt_level, size_t max_fuse_depth, bool link_params) {

   auto graph = IndexedForwardGraphCreator::Create(&arena_, body);

   auto groups = GraphPartitioner(&arena_, fuse_opt_level, max_fuse_depth, max_function_args_)

                     .Partition(graph);

   for (size_t nid = 0; nid < graph.post_dfs_order.size(); ++nid) {

     ICHECK(graph.post_dfs_order[nid]->ref != nullptr);

     gmap_[graph.post_dfs_order[nid]->ref] = groups[nid];

   }

   return this->Mutate(body);

 }

class IndexedForwardGraphCreator : private ExprVisitor {

 public:

  static IndexedForwardGraph Create(support::Arena* arena, const Expr& body) {

    IndexedForwardGraphCreator creator(arena);

    return creator.Prepare(body);

  }

 private:

  explicit IndexedForwardGraphCreator(support::Arena* arena) : arena_(arena) {}

  IndexedForwardGraph Prepare(const Expr& body) { 

    this->Update(body, nullptr, kOpaque);  

    this->VisitExpr(body);  

    return std::move(graph_);

  }

   private:

  support::Arena* arena_;

  IndexedForwardGraph graph_;  

class IndexedForwardGraph {

 public:

  struct Edge {

    Node* node{nullptr};

    OpPatternKind pattern{kOpaque};

  };

  struct Node {  

    const tvm::Object* ref{nullptr};

    size_t index{0};

    bool extern_ref{false};

    OpPatternKind pattern{kOpaque};

    LinkedList<Edge> outputs;

  };

  std::unordered_map<const tvm::Object*, Node*> node_map;

  std::vector<Node*> post_dfs_order;

}

std::vector<GraphPartitioner::Group*> GraphPartitioner::Partition(

    const IndexedForwardGraph& graph) {

  this->InitGroups(graph);  

  if (opt_level_ == 0) return std::move(groups_);

  auto post_dom_tree = DominatorTree::PostDom(arena_, graph);  

  for (int phase = 0; phase < 3; ++phase) {  

    this->RunFuse(graph, post_dom_tree, phase);

  }

  return std::move(groups_);

}

void GraphPartitioner::MergeFromTo(Group* child, Group* parent) {

  child = child->FindRoot();

  parent = parent->FindRoot();

  if (child == parent) return;

  parent->num_nodes += child->num_nodes;

  parent->args_num += child->args_num;

  child->parent = parent;    

  if (child->anchor_ref != nullptr) {

    ICHECK(parent->anchor_ref == nullptr);

    parent->anchor_ref = child->anchor_ref;

    parent->pattern = CombinePattern(child->pattern, parent->pattern);

  }

}

  Expr Rewrite_(const CallNode* call, const Expr& post) {

    if (call->op.as<OpNode>()) {

      static auto fnoncomputational = Op::GetAttrMap<TNonComputational>("TNonComputational");

      static auto fqnncanonicalize = Op::GetAttrMap<FTVMLegalize>("FTVMQnnCanonicalize");

      Op op = Downcast<Op>(call->op);

      if (fnoncomputational.get(op, false) && !fqnncanonicalize.count(op)) {

        return ExprMutator::VisitExpr_(call);

      }

      ICHECK(gmap_.count(call));

      if (call->op == stop_fusion_op) {

        return ExprMutator::VisitExpr(call->args[0]);

      }

      auto* ret_group = gmap_.at(call)->FindRoot();

      Array<Expr> new_args = GetNewArguments(call->args, ret_group);

      auto new_call = Call(call->op, new_args, call->attrs, call->type_args, call->span);

      if (ret_group->root_ref == call) {  

        return MakeNewFunction(ret_group, call->checked_type(), new_call); 

      } else {

        return std::move(new_call);

      }

    } else {

      return ExprMutator::VisitExpr_(call);

    }

  }