Hi
I am trying to extract C source code from the Relay IR after I import the model from onnx.
For starters I tried to create a simple tvm expression and generated a schedule for it using the following method:
dylib_path = os.path.join(base_path, "test_relay_add.so")
compiled_lib.export_library(dylib_path)
def saxpy_sch(curr_path,n): “”" Returns the computing expresioon of saxpy using tvm schedule X: n size vector Y: n size vector Z: n size vector “”"
X = te.placeholder((n,), name = 'X')
Y = te.placeholder((n,), name = 'Y')
Z = te.compute(X.shape, lambda i: n*X[i] + Y[i], name = 'Z')
sch = te.create_schedule(Z.op)
target ="llvm -mtriple=aarch64-linux-gnu -mattr=+sve"
saxpy_dylib = tvm.build(sch, [X,Y, Z] , target='c')
c_source = str(saxpy_dylib.get_source())
print(c_source)
The above snippet is able to generate C code from tvm build schedule.
The C code is:
// tvm target: c -keys=cpu -link-params=0
#define TVM_EXPORTS
#include "tvm/runtime/c_runtime_api.h"
#include "tvm/runtime/c_backend_api.h"
#include <math.h>
#ifdef __cplusplus
extern "C"
#endif
TVM_DLL int32_t default_function(void* args, void* arg_type_ids, int32_t num_args, void* out_ret_value, void* out_ret_tcode, void* resource_handle) {
void* arg0 = (((TVMValue*)args)[0].v_handle);
int32_t arg0_code = ((int32_t*)arg_type_ids)[(0)];
void* arg1 = (((TVMValue*)args)[1].v_handle);
int32_t arg1_code = ((int32_t*)arg_type_ids)[(1)];
void* arg2 = (((TVMValue*)args)[2].v_handle);
int32_t arg2_code = ((int32_t*)arg_type_ids)[(2)];
void* X = (((DLTensor*)arg0)[0].data);
void* arg0_shape = (((DLTensor*)arg0)[0].shape);
void* arg0_strides = (((DLTensor*)arg0)[0].strides);
int32_t dev_id = (((DLTensor*)arg0)[0].device.device_id);
void* Y = (((DLTensor*)arg1)[0].data);
void* arg1_shape = (((DLTensor*)arg1)[0].shape);
void* arg1_strides = (((DLTensor*)arg1)[0].strides);
void* Z = (((DLTensor*)arg2)[0].data);
void* arg2_shape = (((DLTensor*)arg2)[0].shape);
void* arg2_strides = (((DLTensor*)arg2)[0].strides);
if (!(arg0_strides == NULL)) {
}
if (!(arg1_strides == NULL)) {
}
if (!(arg2_strides == NULL)) {
}
for (int32_t i = 0; i < 1024; ++i) {
((float*)Z)[(i)] = ((1.024000e+03f * ((float*)X)[(i)]) + ((float*)Y)[(i)]);
}
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 default_function(args, arg_type_ids, num_args, out_ret_value, out_ret_tcode, resource_handle);
}
But I cannot do the same for models. Since I will need to use tvm.relay.build instead of tvm.build (as above) for building models from onnx, I don’t know how we can extract the c source code in a similar way.
I see that the build generated from relay.build doesn’t contain a ‘get_source’ function as above for extracting c source code.
for now I was able to print the realy IR using str(mod), where mod is the module returned after importing onnx models using relay:
def @main(%Input3: Tensor[(1, 1, 28, 28), float32], %Parameter5: Tensor[(8, 1, 5, 5), float32], %Parameter6: Tensor[(8, 1, 1), float32], %Parameter87: Tensor[(16, 8, 5, 5), float32], %Parameter88: Tensor[(16, 1, 1), float32], %Pooling160_Output_0_reshape0_shape: Tensor[(2), int64], %Parameter193: Tensor[(16, 4, 4, 10), float32], %Parameter193_reshape1_shape: Tensor[(2), int64], %Parameter194: Tensor[(1, 10), float32]) {
%0 = nn.pad(%Input3, 0f, pad_width=[[0, 0], [0, 0], [2, 2], [2, 2]]);
%1 = nn.conv2d(%0, %Parameter5, padding=[0, 0, 0, 0], channels=8, kernel_size=[5, 5]);
%2 = add(%1, %Parameter6);
%3 = nn.relu(%2);
%4 = nn.max_pool2d(%3, pool_size=[2, 2], strides=[2, 2], padding=[0, 0, 0, 0]);
%5 = nn.pad(%4, 0f, pad_width=[[0, 0], [0, 0], [2, 2], [2, 2]]);
%6 = nn.conv2d(%5, %Parameter87, padding=[0, 0, 0, 0], channels=16, kernel_size=[5, 5]);
%7 = add(%6, %Parameter88);
%8 = nn.relu(%7);
%9 = nn.max_pool2d(%8, pool_size=[3, 3], strides=[3, 3], padding=[0, 0, 0, 0]);
%10 = reshape(%Parameter193, newshape=[256, 10]);
%11 = reshape(%9, newshape=[1, 256]);
%12 = transpose(%10, axes=[1, 0]);
%13 = nn.dense(%11, %12, units=None);
add(%13, %Parameter194)
}
So i wanted to know if there is any other way to get relay to print the ‘C’ source code used to export models as libraries rather than printing relay IR.
Any help will be great!
Thanks