[RUNTIME] Quick fix PackedFunc String passing (apache#5266)

dpankratz · Apr 24, 2020 · 3761539 · 3761539
1 parent df87292
commit 3761539
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/include/tvm/runtime/packed_func.h b/include/tvm/runtime/packed_func.h
@@ -513,8 +513,11 @@ class TVMArgValue : public TVMPODValue_ {
     }
   }
   operator tvm::runtime::String() const {
-    // directly use the std::string constructor for now.
-    return tvm::runtime::String(operator std::string());
+    if (IsObjectRef<tvm::runtime::String>()) {
+      return AsObjectRef<tvm::runtime::String>();
+    } else {
+      return tvm::runtime::String(operator std::string());
+    }
   }
   operator DLDataType() const {
     if (type_code_ == kTVMStr) {
@@ -605,8 +608,11 @@ class TVMRetValue : public TVMPODValue_ {
     return *ptr<std::string>();
   }
   operator tvm::runtime::String() const {
-    // directly use the std::string constructor for now.
-    return tvm::runtime::String(operator std::string());
+    if (IsObjectRef<tvm::runtime::String>()) {
+      return AsObjectRef<tvm::runtime::String>();
+    } else {
+      return tvm::runtime::String(operator std::string());
+    }
   }
   operator DLDataType() const {
     if (type_code_ == kTVMStr) {

diff --git a/tests/cpp/packed_func_test.cc b/tests/cpp/packed_func_test.cc
@@ -95,6 +95,12 @@ TEST(PackedFunc, str) {
       CHECK(y == "hello");
       *rv = x;
     })("hello");
+
+  PackedFunc([&](TVMArgs args, TVMRetValue* rv) {
+      CHECK(args.num_args == 1);
+      runtime::String s = args[0];
+      CHECK(s == "hello");
+  })(runtime::String("hello"));
 }