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)

else if (name == "build") {

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

        ICHECK_EQ(args.num_args, 8);

        this->Build(args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7]);

      });

void Build(IRModule mod, const Array<Target>& raw_targets, const tvm::Target& target_host,

             const Executor& executor, const Runtime& runtime,

             const WorkspaceMemoryPools& workspace_memory_pools,

             const ConstantMemoryPools& constant_memory_pools, const String mod_name) {

    VLOG_CONTEXT << "Build";

    executor_ = executor;

    runtime_ = runtime;

    workspace_memory_pools_ = workspace_memory_pools;

    constant_memory_pools_ = constant_memory_pools;

    config_ = CompilationConfig(PassContext::Current(), raw_targets);

    VLOG(1) << "Using compilation config:" << std::endl << config_;

    BuildRelay(std::move(mod), mod_name);

  }

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);

          }

        }

      }

    }

  }

IRModule OptimizeImpl(IRModule relay_module) {

    ICHECK(relay_module.defined()) << "The IRModule must be defined for the Relay compiler.";

    backend::BindParamsInModule(relay_module, params_);  

    Array<Pass> pass_seqs =  

        GetPassPrefix(config_->primitive_targets.size() == 1, false);

    transform::PassContext pass_ctx = PassContext::Current();

    if (config_->optional_homogeneous_target.defined()) {

      pass_seqs.push_back(transform::SplitArgs(

          config_->optional_homogeneous_target->GetAttr<Integer>("max_function_args", 0)

              .value()

              .IntValue()));

    }

    pass_seqs.push_back(transform::PlanDevices(config_));

    pass_seqs.push_back(transform::FuseOps());  

    transform::Pass seq = transform::Sequential(pass_seqs);  

    if (config_->optional_homogeneous_target.defined()) {

      With<Target> tctx(config_->optional_homogeneous_target);

      relay_module = seq(relay_module);

    } else {

      relay_module = seq(relay_module);  

    }

    if (backend::IsAutoSchedulerEnabled() && config_->optional_homogeneous_target.defined()) {  

      Pass major_pass = transform::AutoSchedulerLayoutRewrite();  

      bool enable_layout_rewrite_targets =

          config_->optional_homogeneous_target->GetTargetDeviceType() == kDLCPU ||

          config_->optional_homogeneous_target->GetAttr<String>("device", "") == "mali";

      if (enable_layout_rewrite_targets && pass_ctx.PassEnabled(major_pass->Info())) {

        With<Target> tctx(config_->optional_homogeneous_target);

        relay_module = major_pass(relay_module);

        relay_module = transform::DefuseOps()(relay_module);  

        relay_module = transform::FoldConstant()(relay_module);

        relay_module = transform::FuseOps()(relay_module);

      }

    }

    if (backend::IsMetaScheduleEnabled() && config_->optional_homogeneous_target.defined()) { 

      Pass major_pass = transform::MetaScheduleLayoutRewrite();

      bool enable_layout_rewrite_targets =

          config_->optional_homogeneous_target->GetTargetDeviceType() == kDLCPU ||

          config_->optional_homogeneous_target->GetAttr<String>("device", "") == "mali";

      if (enable_layout_rewrite_targets && pass_ctx.PassEnabled(major_pass->Info())) {

        With<Target> tctx(config_->optional_homogeneous_target);

        relay_module = major_pass(relay_module);

        relay_module = transform::DefuseOps()(relay_module);

        relay_module = transform::FoldConstant()(relay_module);  

        relay_module = transform::FuseOps()(relay_module);

      }

    }

    relay_module = transform::InferType()(relay_module);

    relay_module = transform::Inline()(relay_module);

    relay_module = transform::InferType()(relay_module);

    relay_module = transform::LabelOps()(relay_module);

    relay_module = transform::AnnotateMemoryScope()(relay_module);

    return relay_module;

  }

Array<Pass> GetPassPrefix(bool is_homogeneous, bool is_vm) {

  Array<Pass> pass_seqs;

  Array<runtime::String> entry_functions{"main"};

  pass_seqs.push_back(transform::RemoveUnusedFunctions(entry_functions));

  pass_seqs.push_back(transform::ToBasicBlockNormalForm());

  pass_seqs.push_back(relay::qnn::transform::Legalize());

  if (is_homogeneous) {

    pass_seqs.push_back(transform::Legalize());

  }

  pass_seqs.push_back(transform::SimplifyInference());

  if (is_vm) {

    pass_seqs.push_back(transform::EtaExpand(

         true,  false));

  }

  PackedFunc fskip = PackedFunc([](TVMArgs args, TVMRetValue* rv) {

    Expr expr = args[0];

    if (auto* call_node = expr.as<CallNode>()) {

      auto op_node = call_node->op.as<OpNode>();

      if (op_node->name == "cast") {

        auto attrs = call_node->attrs.as<CastAttrs>();

        if (attrs->dtype == DataType::Int(32)) {

          *rv = true;

        }

      }

    }

    *rv = false;

  });

  pass_seqs.push_back(transform::EliminateCommonSubexpr(fskip));

  pass_seqs.push_back(transform::CombineParallelConv2D(3));

  pass_seqs.push_back(transform::CombineParallelDense(3));

  pass_seqs.push_back(transform::CombineParallelBatchMatmul(3));

  pass_seqs.push_back(transform::FoldConstant());

  pass_seqs.push_back(transform::FoldScaleAxis());

  pass_seqs.push_back(transform::SimplifyExpr());

  pass_seqs.push_back(transform::CanonicalizeCast());

  pass_seqs.push_back(transform::CanonicalizeOps());

  pass_seqs.push_back(transform::FlattenAtrousConv());

  if (is_homogeneous) {

    if (!is_vm) {

      pass_seqs.push_back(transform::InferType());

    }

    pass_seqs.push_back(transform::AlterOpLayout());

    pass_seqs.push_back(transform::SimplifyExprPostAlterOp());

  }

  pass_seqs.push_back(transform::FastMath());

  pass_seqs.push_back(transform::FoldConstant());

  return pass_seqs;

}