import tvm

import numpy as np

from tvm import relay

import os

from tvm.contrib import graph_executor

from collections import OrderedDict

def update_lib(lib):  # lib : llvmModuleNode   ret_.mod

    kwargs = {}

    kwargs["options"] = ["-O2", "-std=c++17"]

    tmp_path = "/home/xianmu/temp/repare_test/"

    lib_name = "lib.so"

    lib_path = os.path.join(tmp_path, lib_name)

    lib.export_library(lib_path, fcompile=False, workspace_dir='/home/user/tempworkspace', **kwargs)

    lib = tvm.runtime.load_module(lib_path)

    return lib

def set_external_func_attr(func, compiler, ext_symbol):

    func = func.with_attr("Primitive", tvm.tir.IntImm("int32", 1))

    func = func.with_attr("Compiler", compiler)

    func = func.with_attr("global_symbol", ext_symbol)

    return func

def prepare_graph_lib(base_path):

    x = relay.var("x", shape=(5, 5), dtype="float32")

    y = relay.var("y", shape=(5, 5), dtype="float32")

    x0 = relay.var("x0", shape=(5, 5), dtype="float32")

    y0 = relay.var("y0", shape=(5, 5), dtype="float32")

    params = {"y0": np.ones((5, 5), dtype="float32")}

    f = relay.Function([x0, y0], x0 * (x0 + y0) - y0)

    f = set_external_func_attr(f, "ccompiler", "ccompiler_01")  # 添加 `Primitive` 和 `Compiler=ccompiler` 标记

    call = relay.Call(f, [x, y])

    mod = tvm.IRModule.from_expr(call)

    print("----------------------------------------------")

    print(mod.script())

    print("----------------------------------------------")

    # 编译 Relay->TIR

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

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

        compiled_lib = relay.build(mod, target, params=params)  # const_loader

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

    device = tvm.cpu()

    lib = update_lib(compiled_lib.get_lib()) 

    rt_mod = tvm.contrib.graph_executor.create(compiled_lib.get_executor_config(), lib, device)  # runtime_mod  图执行器

    # 模型执行

    x_data = np.random.rand(5, 5).astype("float32")

    y_data = np.random.rand(5, 5).astype("float32")

    inputs = OrderedDict([

            ("y", y_data),

            ("x", x_data),

    ])

    for name, data in inputs.items():

        rt_mod.set_input(name, data)

    rt_mod.run()

    out = tvm.nd.empty((5, 5), device=device)

    out = rt_mod.get_output(0, out)

    # 保存模型

    dylib_path = os.path.join(base_path, "test_ccompiler.so")

    compiled_lib.get_lib().export_library(dylib_path)

if __name__ == "__main__":

    curr_path = "/home/xianmu/temp/repare_test/"

    os.environ["TVM_HOME"] = "/home/xianmu/TVM/tvm/"

    prepare_graph_lib(curr_path)

lib.export_library(lib_path, fcompile=False, workspace_dir='/home/user/tempworkspace', **kwargs)

#include <stdio.h>  // 从 1 到 23 行是固定输出代码（头文件和宏定义）

#include <stdlib.h>

#include  <string.h>

#include <tvm/runtime/c_runtime_api.h>

#include <tvm/runtime/c_backend_api.h>

    #define CSOURCE_BINARY_OP_1D(p_ID_, p_OP_, p_DIM1_, p_DTYPE)       \

      void p_ID_(p_DTYPE* a, p_DTYPE* b, p_DTYPE* out) {    \

        for (int64_t i = 0; i < p_DIM1_; ++i) {                        \

          out[i] = a[i] p_OP_ b[i];                                    \

        }                                                              \

      }

    #define CSOURCE_BINARY_OP_2D(p_ID_, p_OP_, p_DIM1_, p_DIM2_, p_DTYPE)  \

      void p_ID_(p_DTYPE* a, p_DTYPE* b, p_DTYPE* out) {        \

        for (int64_t i = 0; i < p_DIM1_; ++i) {                            \

          for (int64_t j = 0; j < p_DIM2_; ++j) {                          \

            int64_t k = i * p_DIM2_ + j;                                   \

            out[k] = a[k] p_OP_ b[k];                                      \

          }                                                                \

        }                                                                  \

      }

CSOURCE_BINARY_OP_2D(ccompiler_01_0, -, 5, 5, float); // 根据AST生成所有Call调用的函数

CSOURCE_BINARY_OP_2D(ccompiler_01_1, *, 5, 5, float);

CSOURCE_BINARY_OP_2D(ccompiler_01_2, +, 5, 5, float);

void ccompiler_01_(float* x0, float* y0, float* out0) {  // 执行过程

  float* buf_0 = (float*)malloc(4 * 25);

  float* buf_1 = (float*)malloc(4 * 25);

  float* buf_2 = (float*)malloc(4 * 25);

  ccompiler_01_2(x0, y0, buf_0);  // 依次调用函数

  ccompiler_01_1(x0, buf_0, buf_1);

  ccompiler_01_0(buf_1, y0, buf_2);

  memcpy(out0, buf_2, 4 * 25);

  free(buf_0);

  free(buf_1);

  free(buf_2);

}

int ccompiler_01_wrapper_(DLTensor* arg0,  // 函数封装

	DLTensor* arg1,

	DLTensor* out0) {

  ccompiler_01_((float*)(arg0->data),

  (float*)(arg1->data),

  (float*)(out0->data));

  return 0;

}

#ifdef __cplusplus

extern "C" {

#endif

TVM_DLL int32_t ccompiler_01(TVMValue* args, int* type_code, int num_args, TVMValue* out_value, int* out_type_code) {

  DLTensor* arg0 = (DLTensor*)(((TVMValue*)args)[0].v_handle);

  DLTensor* arg1 = (DLTensor*)(((TVMValue*)args)[1].v_handle);

  DLTensor* ret2 = (DLTensor*)(((TVMValue*)args)[2].v_handle);

  ccompiler_01_wrapper_(arg0,arg1,ret2);

  return 0;

}

#ifdef __cplusplus

}

#endif