Add Enum TorchScript serialization and deserialization support (pytorch#42963)

Summary:
Pull Request resolved: pytorch#42963

* Adds code printing for enum type
* Enhance enum type to include all contained enum names and values
* Adds code parsing for enum type in deserialization
* Enabled serialization/deserialization test in most TestCases. (With a few dangling issues to be addressed in later PRs to avoid this PR grows too large)

Test Plan: Imported from OSS

Reviewed By: SplitInfinity

Differential Revision: D23223281

Pulled By: gmagogsfm

fbshipit-source-id: 716d1866b7770dfb7bd8515548cfe7dc4c4585f7

Author: gmagogsfm

aten/src/ATen/core/jit_type.h

@@ -1131,18 +1131,19 @@ struct CAFFE2_API TupleType : public NamedType {
 
 struct EnumType;
 using EnumTypePtr = std::shared_ptr<EnumType>;
+using EnumNameValue = std::pair<std::string, IValue>;
 struct CAFFE2_API EnumType : public NamedType {
   friend struct Type;
   static const TypeKind Kind = TypeKind::EnumType;
 
   static EnumTypePtr create(
       const c10::QualifiedName& qualified_class_name,
-      TypePtr value, std::weak_ptr<::torch::jit::CompilationUnit> cu) {
+      TypePtr value, std::vector<EnumNameValue> enum_names_values, std::weak_ptr<::torch::jit::CompilationUnit> cu) {
     switch (value->kind()) {
       case TypeKind::IntType:
       case TypeKind::FloatType:
       case TypeKind::StringType:
-        return EnumTypePtr(new EnumType(qualified_class_name, std::move(value), std::move(cu)));
+        return EnumTypePtr(new EnumType(qualified_class_name, std::move(value), std::move(enum_names_values), std::move(cu)));
       default:
         AT_ERROR(
             "Cannot create Enum with value type '",
@@ -1183,17 +1184,32 @@ struct CAFFE2_API EnumType : public NamedType {
     return name().value();
   }
 
+  at::ArrayRef<TypePtr> containedTypes() const override {
+    return value_type_;
+  }
+
+  const at::ArrayRef<EnumNameValue> enumNamesValues() const {
+    return enum_names_values_;
+  }
+
  private:
-  EnumType(c10::QualifiedName qualified_class_name, TypePtr value_type, std::weak_ptr<torch::jit::CompilationUnit> cu)
+  EnumType(
+      c10::QualifiedName qualified_class_name,
+      TypePtr value_type,
+      std::vector<EnumNameValue> enum_names_values,
+      std::weak_ptr<torch::jit::CompilationUnit> cu)
       : NamedType(TypeKind::EnumType, std::move(qualified_class_name)),
-        value_type_(std::move(value_type)), cu_(cu) {}
+        value_type_(std::move(value_type)),
+        enum_names_values_(std::move(enum_names_values)),
+        cu_(cu) {}
 
   std::string annotation_str_impl(TypePrinter printer = nullptr) const override {
     const auto& n = name().value();
     return n.qualifiedName();
   }
 
   TypePtr value_type_;
+  std::vector<EnumNameValue> enum_names_values_;
   std::weak_ptr<::torch::jit::CompilationUnit> cu_;
 };
 

aten/src/ATen/test/ivalue_test.cpp

@@ -261,36 +261,48 @@ TEST(IValueTest, ListNestedEquality) {
 
 TEST(IValueTest, EnumEquality) {
   auto cu = std::make_shared<CompilationUnit>();
-  auto int_enum_type1 = EnumType::create("enum_class_1", IntType::get(), cu);
-  auto int_enum_type2 = EnumType::create("enum_class_2", IntType::get(), cu);
-  auto string_enum_type = EnumType::create("enum_class_3", StringType::get(), cu);
+  IValue int_ivalue_1(1);
+  IValue int_ivalue_2(2);
+  IValue str_ivalue_1("1");
+  auto int_enum_type1 = EnumType::create(
+      "enum_class_1",
+      IntType::get(),
+      {{"enum_name_1", int_ivalue_1}, {"enum_name_2", int_ivalue_2}},
+      cu);
+  auto int_enum_type2 = EnumType::create(
+      "enum_class_2",
+      IntType::get(),
+      {{"enum_name_1", int_ivalue_1}, {"enum_name_2", int_ivalue_2}},
+      cu);
+  auto string_enum_type = EnumType::create(
+      "enum_class_3", StringType::get(), {{"enum_name_1", str_ivalue_1}}, cu);
 
   EXPECT_EQ(
       IValue(c10::make_intrusive<ivalue::EnumHolder>(
-          int_enum_type1, "enum_name_1", IValue(1))),
+          int_enum_type1, "enum_name_1", int_ivalue_1)),
       IValue(c10::make_intrusive<ivalue::EnumHolder>(
-          int_enum_type1, "enum_name_1", IValue(1)))
+          int_enum_type1, "enum_name_1", int_ivalue_1))
   );
 
   EXPECT_NE(
       IValue(c10::make_intrusive<ivalue::EnumHolder>(
-          int_enum_type1, "enum_name_1", IValue(1))),
+          int_enum_type1, "enum_name_1", int_ivalue_1)),
       IValue(c10::make_intrusive<ivalue::EnumHolder>(
-          int_enum_type2, "enum_name_1", IValue(1)))
+          int_enum_type2, "enum_name_1", int_ivalue_1))
   );
 
   EXPECT_NE(
       IValue(c10::make_intrusive<ivalue::EnumHolder>(
-          int_enum_type1, "enum_name_1", IValue(1))),
+          int_enum_type1, "enum_name_1", int_ivalue_1)),
       IValue(c10::make_intrusive<ivalue::EnumHolder>(
-          int_enum_type1, "enum_name_2", IValue(1)))
+          int_enum_type1, "enum_name_2", int_ivalue_2))
   );
 
   EXPECT_NE(
       IValue(c10::make_intrusive<ivalue::EnumHolder>(
-          int_enum_type1, "enum_name_1", IValue(1))),
+          int_enum_type1, "enum_name_1", int_ivalue_1)),
       IValue(c10::make_intrusive<ivalue::EnumHolder>(
-          string_enum_type, "enum_name_1", IValue("1")))
+          string_enum_type, "enum_name_1", str_ivalue_1))
   );
 }