Allow implicit conversion in TVM FFI to tvm::Bool (apache#5907)

trevor-m · Jun 30, 2020 · 257e5f3 · 257e5f3
1 parent a0a57fa
commit 257e5f3
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 21 deletions.
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
@@ -483,9 +483,9 @@ inline const TTypeNode* RelayExprNode::type_as() const {
 
 namespace tvm {
 namespace runtime {
+// common rule for RetValue and ArgValue
 template <>
 struct PackedFuncValueConverter<PrimExpr> {
-  // common rule for both RetValue and ArgValue.
   static PrimExpr From(const TVMPODValue_& val) {
     if (val.type_code() == kTVMNullptr) {
       return PrimExpr(ObjectPtr<Object>(nullptr));
@@ -500,6 +500,35 @@ struct PackedFuncValueConverter<PrimExpr> {
     return PrimExpr::FromObject_(val.AsObjectRef<ObjectRef>());
   }
 };
+
+template <>
+struct PackedFuncValueConverter<tvm::Integer> {
+  static tvm::Integer From(const TVMPODValue_& val) {
+    if (val.type_code() == kTVMNullptr) {
+      return Integer(ObjectPtr<Object>(nullptr));
+    }
+    if (val.type_code() == kTVMArgInt) {
+      return Integer(val.operator int());
+    }
+    return val.AsObjectRef<tvm::Integer>();
+  }
+};
+
+template <>
+struct PackedFuncValueConverter<tvm::Bool> {
+  static tvm::Bool From(const TVMPODValue_& val) {
+    if (val.type_code() == kTVMNullptr) {
+      return Bool(ObjectPtr<Object>(nullptr));
+    }
+    if (val.type_code() == kTVMArgInt) {
+      int v = val.operator int();
+      CHECK(v == 0 || v == 1) << "ValueError: boolean value can only be 0 or 1, but get " << v;
+      return Bool(static_cast<bool>(v));
+    }
+    return val.AsObjectRef<tvm::Bool>();
+  }
+};
+
 }  // namespace runtime
 }  // namespace tvm
 #endif  // TVM_IR_EXPR_H_
diff --git a/include/tvm/tir/expr.h b/include/tvm/tir/expr.h
@@ -1147,26 +1147,6 @@ inline std::unordered_map<K, V> as_unordered_map(const Map<K, V>& dmap) {
 }  // namespace tir
 }  // namespace tvm
 
-namespace tvm {
-namespace runtime {
-// Additional implementattion overloads for PackedFunc.
-
-template <>
-struct PackedFuncValueConverter<tvm::Integer> {
-  // common rule for RetValue and ArgValue
-  static tvm::Integer From(const TVMPODValue_& val) {
-    if (val.type_code() == kTVMNullptr) {
-      return Integer(ObjectPtr<Object>(nullptr));
-    }
-    if (val.type_code() == kDLInt) {
-      return Integer(val.operator int());
-    }
-    return val.AsObjectRef<tvm::Integer>();
-  }
-};
-}  // namespace runtime
-}  // namespace tvm
-
 namespace std {
 template <>
 struct hash<::tvm::tir::IterVar> : public ::tvm::ObjectPtrHash {};