获取目标源码#
import numpy as np
import tvm
from tvm import relay
x = relay.var("x", shape=(1,), dtype="float32")
y = relay.var("y", shape=(1,), dtype="float32")
params = {"y": np.ones((1,), dtype="float32")}
mod = tvm.IRModule.from_expr(relay.Function([x, y], x + y))
mod.show()
def @main(%x: Tensor[(1), float32], %y: Tensor[(1), float32]) {
add(%x, %y)
}
获取 C 源码:
compiled_module = relay.build(mod, "c", params=params)
print(compiled_module.lib.get_source())
// tvm target: c -keys=cpu
#define TVM_EXPORTS
#include "tvm/runtime/c_runtime_api.h"
#include "tvm/runtime/c_backend_api.h"
#include <math.h>
#include <stdbool.h>
#ifdef __cplusplus
extern "C"
#endif
TVM_DLL int32_t tvmgen_default_fused_add(void* args, int32_t* arg_type_ids, int32_t num_args, void* out_ret_value, int32_t* out_ret_tcode, void* resource_handle) {
int32_t p0_code = arg_type_ids[0];
int32_t p1_code = arg_type_ids[1];
int32_t T_add_code = arg_type_ids[2];
void* p0 = (((TVMValue*)args)[0].v_handle);
void* p1 = (((TVMValue*)args)[1].v_handle);
void* T_add = (((TVMValue*)args)[2].v_handle);
void* p0_1 = (((DLTensor*)p0)[0].data);
void* tvmgen_default_fused_add_p0_shape = (((DLTensor*)p0)[0].shape);
void* tvmgen_default_fused_add_p0_strides = (((DLTensor*)p0)[0].strides);
int32_t dev_id = (((DLTensor*)p0)[0].device.device_id);
void* p1_1 = (((DLTensor*)p1)[0].data);
void* tvmgen_default_fused_add_p1_shape = (((DLTensor*)p1)[0].shape);
void* tvmgen_default_fused_add_p1_strides = (((DLTensor*)p1)[0].strides);
void* T_add_1 = (((DLTensor*)T_add)[0].data);
void* tvmgen_default_fused_add_T_add_shape = (((DLTensor*)T_add)[0].shape);
void* tvmgen_default_fused_add_T_add_strides = (((DLTensor*)T_add)[0].strides);
if (!(tvmgen_default_fused_add_p0_strides == NULL)) {
}
if (!(tvmgen_default_fused_add_p1_strides == NULL)) {
}
if (!(tvmgen_default_fused_add_T_add_strides == NULL)) {
}
((float*)T_add_1)[0] = (((float*)p0_1)[0] + ((float*)p1_1)[0]);
return 0;
}
// CodegenC: NOTE: Auto-generated entry function
#ifdef __cplusplus
extern "C"
#endif
TVM_DLL int32_t __tvm_main__(void* args, int* arg_type_ids, int num_args, void* out_ret_value, int* out_ret_tcode, void* resource_handle) {
return tvmgen_default_fused_add(args, arg_type_ids, num_args, out_ret_value, out_ret_tcode, resource_handle);
}