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)

void BuildRelay(IRModule relay_module, const String& mod_name) {

    IRModule module = WithAttrs(  

        relay_module, <!--swig0-->);

    relay_module = OptimizeImpl(std::move(module));  

    Function func = Downcast<Function>(relay_module->Lookup("main"));  

    IRModule func_module = WithAttrs(IRModule::FromExpr(func),  

                                     <!--swig1-->);

    executor_codegen_ = MakeExecutorCodegen(executor_->name);   

    executor_codegen_->Init(nullptr, config_->primitive_targets);  

    executor_codegen_->Codegen(func_module, func, mod_name);   

    executor_codegen_->UpdateOutput(&ret_);  

    ret_.params = executor_codegen_->GetParams();   

    auto lowered_funcs = executor_codegen_->GetIRModule(); 

    Target ext_dev("ext_dev");

    if (lowered_funcs.find(ext_dev) != lowered_funcs.end()) {

      lowered_funcs.Set(ext_dev, IRModule());

    }

    const Target& host_target = config_->host_virtual_device->target;

    const runtime::PackedFunc* pf = runtime::Registry::Get("codegen.LLVMModuleCreate");

    if (lowered_funcs.size() == 0) {

      if (host_target->kind->name == "llvm") {

        CHECK(pf != nullptr) << "Unable to create empty module for llvm without llvm codegen.";

        ret_.mod = (*pf)(host_target->str(), "empty_module");

      } else {

        ret_.mod = tvm::codegen::CSourceModuleCreate(";", "", Array<String>{});

      }

    } else {

      ret_.mod = tvm::TIRToRuntime(lowered_funcs, host_target);  

    }

    auto ext_mods = executor_codegen_->GetExternalModules();

    ret_.mod = tvm::codegen::CreateMetadataModule(ret_.params, ret_.mod, ext_mods, host_target,

                                                  runtime_, executor_,

                                                  executor_codegen_->GetExecutorCodegenMetadata());

    for (tvm::runtime::Module mod : ext_mods) {

      auto pf_var = mod.GetFunction("get_const_vars");

      if (pf_var != nullptr) {

        Array<String> variables = pf_var();

        for (size_t i = 0; i < variables.size(); i++) {

          auto it = ret_.params.find(variables[i].operator std::string());

          if (it != ret_.params.end()) {

            VLOG(1) << "constant '" << variables[i] << "' has been captured in external module";

            ret_.params.erase(it);

          }

        }

      }

    }

  }

 * \brief GraphCodegen module wrapper

 */

struct GraphCodegen : ExecutorCodegen {  

  GraphCodegen() {

    auto pf = GetPackedFunc("relay.build_module._GraphExecutorCodegen");  

    mod = (*pf)();  

  }

  void UpdateOutput(BuildOutput* ret) override { ret->graph_json = GetGraphJSON(); }

  std::string GetGraphJSON() { return CallFunc<std::string>("get_graph_json", nullptr); }

  ~GraphCodegen() {}

};

TVM_REGISTER_GLOBAL("relay.build_module._GraphExecutorCodegen")

    .set_body([](TVMArgs args, TVMRetValue* rv) { *rv = CreateGraphCodegenMod(); });

     }

runtime::Module CreateGraphCodegenMod() {

  auto ptr = make_object<GraphExecutorCodegenModule>();  

  return runtime::Module(ptr);  

}

virtual PackedFunc GetFunction(const String& name, const ObjectPtr<Object>& sptr_to_self) {

    if (name == "init") {

      return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {

        ICHECK_EQ(args.num_args, 2) << "The expected of arguments are: "

                                    << "runtime::Module mod and Array<Target> targets";

        void* mod = args[0];

        Array<Target> targets = args[1];

        codegen_ = std::make_shared<GraphExecutorCodegen>(reinterpret_cast<runtime::Module*>(mod),

                                                          std::move(targets));

      });

    }

else if (name == "codegen") {

      return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {

        IRModule mod = args[0];

        Function func = args[1];

        String mod_name = args[2];

        this->output_ = this->codegen_->Codegen(mod, func, mod_name); 

      });

LoweredOutput Codegen(IRModule mod, relay::Function func, String mod_name) {

    mod_name_ = mod_name;

    VLOG_CONTEXT << "GraphExecutorCodegen";

    VLOG(1) << "compiling:" << std::endl << PrettyPrint(func);

    memory_plan_ = GraphPlanMemory(func);  

    backend::FunctionInfo func_info;

    if (memory_plan_.defined()) {

      func_info =

          relay::tec::UpdateMainWorkspaceSize(mod, config_, memory_plan_->expr_to_storage_info);

      mod = WithAttr(mod, "main_func_info", func_info);

    }

    IRModule lowered_mod = tec::LowerTE(mod_name_, config_, [this](BaseFunc func) {  

      if (func->GetAttr<String>(attr::kCompiler).defined()) {

        UpdateConstants(func, &params_);

      }

      tec::UpdateFunctionMetadata(func, this->function_metadata_);

    })(mod);  

    Optional<backend::FunctionInfo> main_func_info =

        lowered_mod->GetAttr<backend::FunctionInfo>("main_func_info");  

    function_metadata_.Set(runtime::symbol::tvm_module_main, main_func_info.value());

    Function lowered_main_func = Downcast<Function>(lowered_mod->Lookup("main"));

    memory_plan_ = GraphPlanMemory(lowered_main_func);  

    for (auto param : lowered_main_func->params) {

      auto node_ptr = GraphInputNode::make_node_ptr(param->name_hint(), GraphAttrs());

      var_map_[param.get()] = AddNode(node_ptr, param);

    }

    heads_ = VisitExpr(lowered_main_func->body);

    std::ostringstream os;

    dmlc::JSONWriter writer(&os);

    GetJSON(&writer);    

    LoweredOutput ret;  

    ret.graph_json = os.str();  

    ret.external_mods =  

        lowered_mod->GetAttr<Array<runtime::Module>>(tvm::attr::kExternalMods).value_or({});

    ret.params = std::unordered_map<std::string, tvm::runtime::NDArray>();

    Map<String, runtime::NDArray> const_name_to_constant =

        lowered_mod->GetAttr<Map<String, runtime::NDArray>>(tvm::attr::kConstNameToConstant)

            .value_or({});

    for (const auto& kv : const_name_to_constant) {

      VLOG(1) << "constant '" << kv.first << "' contributed by external codegen";

      ICHECK(ret.params.emplace(kv.first, kv.second).second);

    }

    for (const auto& kv : params_) {

      VLOG(1) << "constant '" << kv.first << "' contributed by TECompiler";

      ICHECK(ret.params.emplace(kv.first, kv.second).second);

    }

    ret.function_metadata = std::move(function_metadata_);  

    ret.lowered_funcs = tec::GetPerTargetModules(lowered_mod);

    ret.metadata =  

        ExecutorCodegenMetadata({} , {} , {} ,

                                {} , {} , {} ,

                                runtime::kTvmExecutorGraph , mod_name_ ,

                                "packed" , Bool(false) );

    return ret;  

  }

IRModule lowered_mod = tec::LowerTE(mod_name_, config_, [this](BaseFunc func) {  

      if (func->GetAttr<String>(attr::kCompiler).defined()) {

        UpdateConstants(func, &params_);

      }

      tec::UpdateFunctionMetadata(func, this->function_metadata_);

    })(mod);  

Pass LowerTE(String module_name, CompilationConfig complilation_config, ProcessFn process_fn) {

  runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func = [=](IRModule module,  

                                                                            PassContext ctx) {

    return LowerTE(module, module_name, process_fn, complilation_config);

  };

  return tvm::transform::Sequential(

      {tvm::relay::transform::RelayToTIRTargetHook(complilation_config),

       tvm::transform::CreateModulePass(pass_func, 0, "LowerTE", {"InferType"}), InferType(),

       tvm::tir::transform::ExtractPrimFuncConstants()});

}