- preserve legacy print behavior for unrecognised dtypes.

Author: mbs-octoml

src/printer/relay_text_printer.cc

@@ -348,14 +348,15 @@ Doc RelayTextPrinter::VisitExpr_(const VarNode* op) { return AllocVar(GetRef<Var
 
 Doc RelayTextPrinter::VisitExpr_(const ConstantNode* op) {
   // Print out simple scalars directly.
-  if (op->is_scalar()) {
+  if (support::IsSimpleScalar(op)) {
     return Doc::Text(support::NDArrayScalarToString(op->data));
   }
-  // default fall-back, record it as meta node.
+  // Fallbock: record it as a meta node.
   Doc doc;
   // Don't append optional_info. Because the entry function is Print,
   // and it will append the optional_info afterwards.
-  return doc << PrintExpr(GetRef<Expr>(op), true, false, false);
+  return doc << PrintExpr(GetRef<Expr>(op), /*meta=*/true, /*try_inline=*/false,
+                          /*optional_info=*/false);
 }
 
 Doc RelayTextPrinter::VisitExpr_(const TupleNode* op) {
@@ -772,11 +773,21 @@ Doc RelayTextPrinter::VisitAttr_(const ArrayNode* op) {
 }
 
 Doc RelayTextPrinter::VisitAttr_(const tir::IntImmNode* op) {
-  return Doc::Text(support::IntImmToString(GetRef<IntImm>(op)));
+  if (support::IsSimpleScalarDtype(op->dtype)) {
+    return Doc::Text(support::IntImmToString(GetRef<IntImm>(op)));
+  } else {
+    // Fallback: Print int64_t without width suffix.
+    return Doc::Text(std::to_string(op->value));
+  }
 }
 
 Doc RelayTextPrinter::VisitAttr_(const tir::FloatImmNode* op) {
-  return Doc::Text(support::FloatImmToString(GetRef<FloatImm>(op)));
+  if (support::IsSimpleScalarDtype(op->dtype)) {
+    return Doc::Text(support::FloatImmToString(GetRef<FloatImm>(op)));
+  } else {
+    // Fallbock: Print double without width suffix.
+    return Doc::Text(std::to_string(op->value));
+  }
 }
 
 Doc RelayTextPrinter::VisitAttr_(const tir::StringImmNode* op) {

src/support/scalars.cc

@@ -38,27 +38,27 @@ static const DataType kFloat32 = DataType::Float(32);
 static const DataType kFloat64 = DataType::Float(64);
 static const DataType kBool = DataType::Bool();
 
-bool IsSimpleScalar(const relay::ConstantNode* constant_node) {
-  if (!constant_node->is_scalar()) {
-    return false;
-  }
-  DataType dtype(constant_node->data->dtype);
+bool IsSimpleScalarDtype(DataType dtype) {
   return dtype == kInt16 || dtype == kInt32 || dtype == kInt64 || dtype == kFloat16 ||
          dtype == kFloat32 || dtype == kFloat64 || dtype == kBool;
 }
 
+bool IsSimpleScalar(const relay::ConstantNode* constant_node) {
+  return constant_node->is_scalar() && IsSimpleScalarDtype(DataType(constant_node->data->dtype));
+}
+
 runtime::NDArray IntImmToNDArray(const IntImm& int_imm) {
   DLDevice dev = {DLDeviceType::kDLCPU, 0};
   auto data = runtime::NDArray::Empty({}, int_imm->dtype, dev);
-  if (int_imm.dtype() == kInt64) {
-    auto array = reinterpret_cast<int64_t*>(data->data);
-    array[0] = int_imm->value;
+  if (int_imm.dtype() == kInt16) {
+    auto array = reinterpret_cast<int16_t*>(data->data);
+    array[0] = static_cast<int16_t>(int_imm->value);
   } else if (int_imm.dtype() == kInt32) {
     auto array = reinterpret_cast<int32_t*>(data->data);
     array[0] = static_cast<int32_t>(int_imm->value);
-  } else if (int_imm.dtype() == kInt16) {
-    auto array = reinterpret_cast<int16_t*>(data->data);
-    array[0] = static_cast<int16_t>(int_imm->value);
+  } else if (int_imm.dtype() == kInt64) {
+    auto array = reinterpret_cast<int64_t*>(data->data);
+    array[0] = int_imm->value;
   } else {
     LOG(FATAL) << "Unrecognized numeric literal dtype: " << DLDataType2String(int_imm.dtype());
   }
@@ -68,15 +68,15 @@ runtime::NDArray IntImmToNDArray(const IntImm& int_imm) {
 runtime::NDArray FloatImmToNDArray(const FloatImm& float_imm) {
   DLDevice dev = {DLDeviceType::kDLCPU, 0};
   auto data = runtime::NDArray::Empty({}, float_imm->dtype, dev);
-  if (float_imm.dtype() == kFloat64) {
-    auto array = reinterpret_cast<double*>(data->data);
-    array[0] = float_imm->value;
+  if (float_imm.dtype() == kFloat16) {
+    auto array = reinterpret_cast<uint16_t*>(data->data);
+    array[0] = __gnu_f2h_ieee(static_cast<float>(float_imm->value));
   } else if (float_imm.dtype() == kFloat32) {
     auto array = reinterpret_cast<float*>(data->data);
     array[0] = static_cast<float>(float_imm->value);
-  } else if (float_imm.dtype() == kFloat16) {
-    auto array = reinterpret_cast<uint16_t*>(data->data);
-    array[0] = __gnu_f2h_ieee(static_cast<float>(float_imm->value));
+  } else if (float_imm.dtype() == kFloat64) {
+    auto array = reinterpret_cast<double*>(data->data);
+    array[0] = float_imm->value;
   } else {
     LOG(FATAL) << "Unrecognized numeric literal dtype: " << DLDataType2String(float_imm.dtype());
   }

src/support/scalars.h

@@ -35,6 +35,9 @@
 namespace tvm {
 namespace support {
 
+/*! \brief Returns true if a tensor of empty shape and given dtype is considered a Relay scalar. */
+bool IsSimpleScalarDtype(DataType dtype);
+
 /*! \brief Returns true if \p constant_node is a float/int/bool scalar. */
 bool IsSimpleScalar(const relay::ConstantNode* constant_node);
 

tests/python/relay/test_target_hooks.py

@@ -73,4 +73,5 @@ def test_runtime_module_generation(check_result):
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    # sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    test_runtime_module_generation(check_aot_executor_result)