TVM C codegen clarification

I am dumping C code for super_resolution_0.2.onnx using relay.build(…) and want to understand how the loop bounds(224) are computed? Is it from the tensor shapes and compiler doing constant folding?

target = "c"
dev=tvm.cpu()
mod, params = relay.frontend.from_onnx(onnx_model,shape_dict,freeze_params=True)
with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
    graph , lib , params = relay.build(mod, target=target, params=None)
print(lib.get_source())

#ifdef __cplusplus
extern "C"
#endif
TVM_DLL int32_t tvmgen_default_fused_layout_transform(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* placeholder = (((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* T_layout_trans = (((DLTensor*)arg1)[0].data);
  void* arg1_shape = (((DLTensor*)arg1)[0].shape);
  void* arg1_strides = (((DLTensor*)arg1)[0].strides);
  if (!(arg0_strides == NULL)) {
  }
  if (!(arg1_strides == NULL)) {
  }
  for (int32_t ax0_ax1_fused_ax2_fused = 0; ax0_ax1_fused_ax2_fused < **224**; ++ax0_ax1_fused_ax2_fused) {
    for (int32_t ax3 = 0; ax3 < **224**; ++ax3) {
      ((float*)T_layout_trans)[(((ax0_ax1_fused_ax2_fused * 224) + ax3))] = ((float*)placeholder)[(((ax0_ax1_fused_ax2_fused * 224) + ax3))];
    }
  }
  return 0;
}

most often, yes. it depends a bit on any layout transformations chosen and which schedules were used to implement your model.

1 Like

Thanks for the reply, Is it possible to disable infraBoundpass and generate code with bounds as non-const expressions?

@harishch4 is this for dynamic shapes? I believe Relax is bringing support for this.

1 Like

Yes, support for dynamic shapes is what I’m looking for. Thanks for pointing me to Relax. I found an example that uses Relax here.

Is it possible to dump the C source from relax.vm.build()'s output? like we do from relay.build()'s lib.get_source().

cc @yuchenj who knows more here :slight_smile:

Thanks for pointing this to me! I replied in another thread as well:

Hi @harishch4, you can get the C source for the compiled kernel library by:

exec = relax.vm.build(relax_mod: IRModule, target)
c_str = exec.mod.imported_modules[0].get_source()

Thanks @yuchenj @areusch :slight_smile: It worked!

Just in case anyone is trying to convert the relay model to relax model, it can be done as below: relax_mod = relay_translator.from_relay(relay_mod["main"], target)

Glad it worked! :slight_smile:

Yep, an example to convert Relay ResNet model to Relax is here: https://github.com/tlc-pack/relax/blob/relax/apps/relax_examples/resnet.py.

1 Like