* \brief This demonstration external codegen target emits C/C++ for compilation by the native c

 * compiler on CPU.

 *  - Patterns: None, functions must be explicitly marked as "Primitive" and "Compiler=ccompiler".

 *  - Custom compiler: relay/backend/contrib/codegen_c/codegen.cc

 */

TVM_REGISTER_TARGET_KIND("ccompiler", kDLCPU)

    .set_attr<Bool>(tvm::attr::kIsExternalCodegen, Bool(true))

    .set_attr<relay::transform::FTVMRelayToTIR>(tvm::attr::kRelayToTIR, CCompilerPass())

    .add_attr_option<String>("header", String(""));

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

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

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

    mod = tvm.IRModule.from_expr(call)

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

    print(mod.script())

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

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

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

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

    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)  

    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)

def @main(%x: Tensor[(5, 5), float32], %y: Tensor[(5, 5), float32]) {

  %2 = fn (%x0: Tensor[(5, 5), float32], %y0: Tensor[(5, 5), float32], Primitive=1, Compiler="ccompiler", global_symbol="ccompiler_01") {

    %0 = add(%x0, %y0);

    %1 = multiply(%x0, %0);

    subtract(%1, %y0)

  };

  %2(%x, %y)

}

************** result **************

 [[-0.2741713   0.08935148 -0.04881942 -0.02653596  0.6241337 ]

 [ 0.6122713  -0.62298256  0.68201226  0.37773138 -0.4962681 ]

 [-0.5679778   0.28375486  0.38151044 -0.41397965 -0.3218666 ]

 [-0.24529386 -0.4346298  -0.10323387 -0.6108544  -0.22057068]

 [-0.44894195  0.16438246 -0.47126824  0.546167   -0.8994445 ]]

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

      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];                                    \

        }                                                              \

      }

      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;

}

extern "C" {

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;

}

}

std::vector<Output> VisitExpr_(const CallNode* call) override {

    std::ostringstream macro_stream;

    std::ostringstream decl_stream;

    std::ostringstream buf_stream;

    std::string func_name = ext_func_id_ + "_" + std::to_string(func_idx++);

    macro_stream << "CSOURCE_BINARY_OP_" << call->args.size() << "D(" << func_name << ", ";  

    if (backend::IsOp(call, "add")) {

      macro_stream << "+";

    } else if (backend::IsOp(call, "subtract")) {

      macro_stream << "-";

    } else if (backend::IsOp(call, "multiply")) {

      macro_stream << "*";

    } else {

      LOG(FATAL) << "Unrecognized op";

    }

    auto in_shape = backend::GetShape(call->args[0]->checked_type());

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

      macro_stream << ", " << in_shape[i];

    }

    const auto* type_node = call->checked_type().as<TensorTypeNode>();

    ICHECK(type_node);

    const auto& dtype = GetDtypeString(type_node);

    macro_stream << ", " << dtype;

    macro_stream << ");";

    func_decl_.push_back(macro_stream.str());

    bool first = true;    

    decl_stream << func_name << "(";

    for (size_t i = 0; i < call->args.size(); ++i) {

      auto res = VisitExpr(call->args[i]);

      for (auto out : res) {

        if (!first) {

          decl_stream << ", ";

        }

        first = false;

        decl_stream << out.name;

      }

    }

    std::string out = "buf_" + std::to_string(buf_idx_++);  

    auto out_shape = backend::GetShape(call->checked_type());

    int out_size = 1;

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

      out_size *= out_shape[i];

    }

    buf_stream << dtype << "* " << out << " = (" << dtype << "*)malloc(4 * " << out_size << ");";

    buf_decl_.push_back(buf_stream.str());

    decl_stream << ", " << out << ");";

    ext_func_body_.push_back(decl_stream.str());

    Output output;

    output.name = out;

    output.dtype = dtype;

    output.need_copy = true;

    output.size = out_size;

    return {output};

  }

['/usr/bin/g++', '-shared', '-fPIC', '-o', '/home/user/tempworkspace/test_ccompiler.so', '/home/user/tempworkspace/lib0.c', '/home/user/tempworkspace/lib1.o', '/home/user/tempworkspace/devc.o', '-I/home/user/TVM/include', '-I/home/user/TVM/3rdparty/dlpack/include', '-I/home/user/TVM/3rdparty/dmlc-core/include']