refactor: Implement a macro for pytorch type checking

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

core/conversion/converters/impl/shuffle.cpp

@@ -86,7 +86,7 @@ static auto shuffle_registrations TORCHTRT_UNUSED =
                    shape_tensor = tensor_to_const(ctx, torch::tensor(shape_vec).to(torch::kI32));
                  } else {
                    LOG_ERROR(
-                       "Invalid IValue type of " << args[1].ivalue_type()
+                       "Invalid IValue type of " << args[1].IValue()->type()
                                                  << " detected for shape tensor from node: " << *n);
                  }
                } else {

core/conversion/var/Var.cpp

@@ -21,28 +21,6 @@ Var::Var(nvinfer1::ITensor* p) : type_(Type::kITensor) {
   ptr_.tensor = p;
 }
 
-Var::IValueType Var::determineIValueType(torch::jit::IValue* p) {
-  if (p->isInt()) {
-    return IValueType::kInt;
-  } else if (p->isDouble()) {
-    return IValueType::kDouble;
-  } else if (p->isBool()) {
-    return IValueType::kBool;
-  } else if (p->isTensor()) {
-    return IValueType::kTensor;
-  } else if (p->isIntList()) {
-    return IValueType::kIntList;
-  } else if (p->isDoubleList()) {
-    return IValueType::kDoubleList;
-  } else if (p->isBoolList()) {
-    return IValueType::kBoolList;
-  } else if (p->isTensorList()) {
-    return IValueType::kTensorList;
-  } else if (p->isList()) {
-    return IValueType::kITensorList;
-  }
-}
-
 Var::Var(const Var& a) {
   switch (a.type_) {
     case Type::kITensor:
@@ -52,7 +30,6 @@ Var::Var(const Var& a) {
     case Type::kIValue:
       ptr_.ivalue = a.ptr_.ivalue;
       type_ = Type::kIValue;
-      ivalue_type_ = determineIValueType(ptr_.ivalue);
       break;
     case Type::kNone:
     default:
@@ -70,7 +47,6 @@ Var& Var::operator=(const Var& a) {
     case Type::kIValue:
       ptr_.ivalue = a.ptr_.ivalue;
       type_ = Type::kIValue;
-      ivalue_type_ = determineIValueType(ptr_.ivalue);
       break;
     case Type::kNone:
     default:
@@ -83,7 +59,6 @@ Var& Var::operator=(const Var& a) {
 Var& Var::operator=(torch::jit::IValue* in) {
   ptr_.ivalue = in;
   type_ = Type::kIValue;
-  ivalue_type_ = determineIValueType(ptr_.ivalue);
   return (*this);
 }
 
@@ -97,10 +72,6 @@ Var::Type Var::type() const {
   return type_;
 }
 
-Var::IValueType Var::ivalue_type() const {
-  return ivalue_type_;
-}
-
 std::string Var::type_name() const {
   switch (type_) {
     case Type::kITensor:
@@ -175,40 +146,8 @@ bool Var::isITensor() const {
   }
 }
 
-bool Var::isITensorList() const {
-  if (ivalue_type_ == IValueType::kITensorList) {
-    return true;
-  } else {
-    return false;
-  }
-}
-
-bool Var::isIntList() const {
-  if (ivalue_type_ == IValueType::kIntList) {
-    return true;
-  } else {
-    return false;
-  }
-}
-
-bool Var::isDoubleList() const {
-  if (ivalue_type_ == IValueType::kDoubleList) {
-    return true;
-  } else {
-    return false;
-  }
-}
-
-bool Var::isTensorList() const {
-  if (ivalue_type_ == IValueType::kTensorList) {
-    return true;
-  } else {
-    return false;
-  }
-}
-
-bool Var::isBoolList() const {
-  if (ivalue_type_ == IValueType::kBoolList) {
+bool Var::isITensorList() {
+  if (isList() && ptr_.ivalue->isCustomClass()) {
     return true;
   } else {
     return false;
@@ -218,10 +157,7 @@ bool Var::isBoolList() const {
 std::vector<nvinfer1::ITensor*> Var::unwrapToITensorList() {
   TORCHTRT_CHECK(
       isIValue(), "Requested unwrapping of arg assuming it was an IValue, however arg type is " << type_name());
-  TORCHTRT_CHECK(
-      isITensorList(),
-      "Expected IValue to be an ITensorList, however the type is "
-          << static_cast<std::underlying_type<IValueType>::type>(ivalue_type_));
+  TORCHTRT_CHECK(isITensorList(), "Expected IValue to be an ITensorList");
   auto ivalue_list = ptr_.ivalue->toList();
   std::vector<nvinfer1::ITensor*> outputs;
   for (int i = 0; i < ivalue_list.size(); i++) {

core/conversion/var/Var.h

@@ -15,7 +15,7 @@ namespace conversion {
 class Var : torch::CustomClassHolder {
  public:
   enum Type { kITensor, kIValue, kNone };
-  enum IValueType { kInt, kDouble, kBool, kTensor, kIntList, kDoubleList, kBoolList, kTensorList, kITensorList };
+
   Var();
   Var(torch::jit::IValue* p);
   Var(nvinfer1::ITensor* p);
@@ -59,16 +59,23 @@ class Var : torch::CustomClassHolder {
 
   bool isIValue() const;
   bool isITensor() const;
-  bool isITensorList() const;
-  bool isTensorList() const;
-  bool isDoubleList() const;
-  bool isIntList() const;
-  bool isBoolList() const;
   bool isNone() const;
+
+  bool isInt();
+  bool isDouble();
+  bool isBool();
+  bool isString();
+  bool isScalar();
+  bool isTensor();
+  bool isIntList();
+  bool isDoubleList();
+  bool isBoolList();
+  bool isTensorList();
+  bool isITensorList();
+  bool isList();
+
   Var::Type type() const;
-  Var::IValueType ivalue_type() const;
   std::string type_name() const;
-  Var::IValueType determineIValueType(torch::jit::IValue* p);
 
  private:
   union VarContainer {
@@ -79,7 +86,6 @@ class Var : torch::CustomClassHolder {
 
   VarContainer ptr_;
   Type type_;
-  IValueType ivalue_type_;
 };
 
 } // namespace conversion

core/conversion/var/Var_inl.h

@@ -4,6 +4,13 @@ namespace torch_tensorrt {
 namespace core {
 namespace conversion {
 
+#define DEFINE_IS_IVAL_TYPE(method_variant)                                                                        \
+  inline bool Var::is##method_variant() {                                                                          \
+    TORCHTRT_CHECK(                                                                                                \
+        isIValue(), "Requested unwrapping of arg assuming it was an IValue, however arg type is " << type_name()); \
+    return ptr_.ivalue->is##method_variant();                                                                      \
+  }
+
 #define DEFINE_UNWRAP_TO(ival_type, method_variant)                                                                \
   template <>                                                                                                      \
   inline ival_type Var::unwrapTo<ival_type>() {                                                                    \
@@ -34,6 +41,18 @@ namespace conversion {
     return this->unwrapTo<ival_type>();                                                                            \
   }
 
+DEFINE_IS_IVAL_TYPE(Int)
+DEFINE_IS_IVAL_TYPE(Double)
+DEFINE_IS_IVAL_TYPE(Bool)
+DEFINE_IS_IVAL_TYPE(String)
+DEFINE_IS_IVAL_TYPE(Scalar)
+DEFINE_IS_IVAL_TYPE(Tensor)
+DEFINE_IS_IVAL_TYPE(IntList)
+DEFINE_IS_IVAL_TYPE(DoubleList)
+DEFINE_IS_IVAL_TYPE(BoolList)
+DEFINE_IS_IVAL_TYPE(TensorList)
+DEFINE_IS_IVAL_TYPE(List)
+
 DEFINE_UNWRAP_TO(at::Tensor, Tensor)
 DEFINE_UNWRAP_TO(int64_t, Int)
 DEFINE_UNWRAP_TO(double, Double)