[Bug] Using AOT with UMA and TFLITE

[MichaelJKlaiber]

Hi everyone,

found a possible issue.
When executing a testcase using UMA and generating the test .c files via aot.py, it seems that the pointers of the offloaded functions are not in the right order.
Can anyone help?

Thanks,
Michael

CC: @cgerum @PaulPalomeroBernardo @manupa-arm @areusch @lhutton1 @d-smirnov @mehrdadh @SebastianBoblest @sezgin1947

Expected behavior

Correct execution of provided NN imported from TFLITE file.

Actual behavior

Segmentation fault, due to mixed up pointers in generated file.
In particular: in default_lib2.c the pointers for placeholder_1, placeholder are mixed up

Environment

Ubuntu 18.04
TVM c3c7c4c

Steps to reproduce

Run tests/python/contrib/test_uma/test_uma_pipeline.py from this branch
https://github.com/MichaelJKlaiber/tvm/tree/issue/uma-tflite

Generated file looks like this:

#include "/home/michael/Documents/code/tvm/apps/uma/_template/conv2dnchw.cc"// tvm target: vanilla_accelerator
#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 tvmgen_default_vanilla_accelerator_main_0(float* placeholder, float* placeholder_1, float* T_matmul_NT, uint8_t* global_const_workspace_4_var, uint8_t* global_workspace_5_var) {
  for (int32_t i1 = 0; i1 < 10; ++i1) {
    for (int32_t i2 = 0; i2 < 784; ++i2) {
      if (i2 == 0) {
        T_matmul_NT[i1] = 0.000000e+00f;
      }
      T_matmul_NT[i1] = (T_matmul_NT[i1] + (placeholder_1[i2] * placeholder[((i1 * 784) + i2)]));
    }
  }
  return 0;
}

when changing the function signature manually to
TVM_DLL int32_t tvmgen_default_vanilla_accelerator_main_0(float* placeholder_1, float* placeholder, float* T_matmul_NT, uint8_t* global_const_workspace_4_var, uint8_t* global_workspace_5_var)
the test is successful

[MichaelJKlaiber]

CC: @cgerum @PaulPalomeroBernardo @manupa-arm @areusch @lhutton1 @d-smirnov @mehrdadh @SebastianBoblest @sezgin1947

[areusch]

this looks odd--it's late for today for me to investigate but hopefully we can look into it next week.

[MichaelJKlaiber]

Seems the bug is in the lowering step from relay_prim_func to tir_prim_func.

When collecting the input arguments, the order gets messed up. By replacing the function _get_tensors here
https://github.com/MichaelJKlaiber/tvm/blob/issue/uma-tflite/python/tvm/relay/backend/contrib/uma/api/lower.py#L83

by

            return list(te_cached_func.inputs) + list(te_cached_func.outputs)

everything works perfectly fine for all the test case we provided including the tflite import

Does anyone know if it is sufficient to rely on the .inputs parameters from te_cached_func ?

[areusch]

@MichaelJKlaiber is this one fixed?

[MichaelJKlaiber]

yes, the bugfix is in main