[Bug] [TVMScript] Weird output of my script

[lightzhan-intellif]

Hi all,
I have a program like this:

import numpy as np
from tvm.script import tir as T
import tvm

@T.prim_func
def test_case(A: T.Buffer[(1, ), "uint8"], B: T.Buffer[(1, ), "float32"]):
  B[0] = A[0]


if __name__ == "__main__":
  compiled = tvm.build(test_case, target="llvm")
  in_data = np.empty((1,), dtype=np.uint8)
  in_data[0] = 100
  out_data = np.empty((1,), dtype=np.float32)
  A_tvm = tvm.nd.array(in_data, tvm.cpu(0))
  B_tvm = tvm.nd.array(out_data, tvm.cpu(0))
  compiled(A_tvm, B_tvm)
  print(B_tvm.numpy())

I just want to assign a value 100 typed uint8 to an address of type float. In python, we should get the result 100 in the B[0], but it give me a weird result [2.0000238] without any warning or error.

I change the target to c, and get the following c code:

// 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>
#ifdef __cplusplus
extern "C"
#endif
TVM_DLL int32_t default_function(void* args, int32_t* arg_type_ids, int32_t num_args, void* out_ret_value, int32_t* out_ret_tcode, void* resource_handle) {
  void* arg_A_handle = (((TVMValue*)args)[0].v_handle);
  int32_t arg_A_handle_code = arg_type_ids[0];
  void* arg_B_handle = (((TVMValue*)args)[1].v_handle);
  int32_t arg_B_handle_code = arg_type_ids[1];
  void* A = (((DLTensor*)arg_A_handle)[0].data);
  void* arg_A_handle_shape = (((DLTensor*)arg_A_handle)[0].shape);
  void* arg_A_handle_strides = (((DLTensor*)arg_A_handle)[0].strides);
  int32_t dev_id = (((DLTensor*)arg_A_handle)[0].device.device_id);
  void* B = (((DLTensor*)arg_B_handle)[0].data);
  void* arg_B_handle_shape = (((DLTensor*)arg_B_handle)[0].shape);
  void* arg_B_handle_strides = (((DLTensor*)arg_B_handle)[0].strides);
  if (!(arg_A_handle_strides == NULL)) {
  }
  if (!(arg_B_handle_strides == NULL)) {
  }
  *(uint8_t*)(((float*)B) + 0) = ((uint8_t*)A)[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 default_function(args, arg_type_ids, num_args, out_ret_value, out_ret_tcode, resource_handle);
}

From the code *(uint8_t*)(((float*)B) + 0) = ((uint8_t*)A)[0];, you can see that tvm does not do any casting to the assignment operation, which causes this strange behavior.

So, how can we fix this? Any idea is welcome.

@junrushao @wrongtest-intellif @Hzfengsy

[vinx13]

T.cast should be used to cast explicitly. In this case, without cast we should throw error in tvmscript parser instead

[lightzhan-intellif]

I am not very sure which way we should go for. But from the python view, if we can do auto casting from low precision datatype to the higher precision datatype and throw an error otherwise, it will be more pythonic.

[wrongtest-intellif]

In this case

For my curiosity, why we have divergent dtype value support in BufferStore? If this is the expected supporting behavior in backend, it seems that we may not just throw out, because we want to keep roundtrip property for all valid ast.

[vinx13]

We didn't add such check of data types between buffer and value to store yet in BufferStore. Previously we had some check in Store (now deprecated) that was disabled for some reasons. see https://github.com/apache/tvm/blob/main/src/tir/ir/stmt.cc#L274-L286
https://github.com/apache/tvm/blob/main/src/tir/ir/expr.cc#L651-L671

In general, storing value with inconsistent data type is not supported except a few cases like boolean vs uint8

[junrushao]

#13960