Move QScheme ops to c10 (pytorch#30134)

Summary:
Pull Request resolved: pytorch#30134

ghstack-source-id: 95055387

Test Plan: buck build mode/dev caffe2:generate-code

Differential Revision: D18609716

fbshipit-source-id: fec39359e0b97387a9b13f8179d72a731cc61808

Author: xcheng16

aten/src/ATen/core/boxing/kernel_functor.h

@@ -39,7 +39,8 @@ namespace detail {
     bool,
     std::string,
     at::Tensor,
-    at::Scalar
+    at::Scalar,
+    c10::QScheme
   >;
 
   template<class T, bool AllowDeprecatedTypes, class Enable = void> struct assert_is_valid_input_type {

aten/src/ATen/core/jit_type.h

@@ -49,7 +49,8 @@ using OptNameList = c10::optional<std::vector<std::string>>;
   _(FunctionType)           \
   _(ClassType)              \
   _(CapsuleType)            \
-  _(InterfaceType)
+  _(InterfaceType)          \
+  _(QSchemeType)
 
 enum class TypeKind {
 #define DEFINE_TYPE(T) T,
@@ -1090,6 +1091,28 @@ struct CAFFE2_API GeneratorType : public Type {
   GeneratorType() : Type(TypeKind::GeneratorType) {}
 };
 
+struct QSchemeType;
+using QSchemeTypePtr = std::shared_ptr<QSchemeType>;
+// This type represents a QScheme
+struct CAFFE2_API QSchemeType : public Type {
+  static QSchemeTypePtr create() {
+    return QSchemeTypePtr(
+        new QSchemeType()); // NOLINT(modernize-make-shared)
+  }
+  bool operator==(const Type& rhs) const override {
+    return rhs.kind() == kind();
+  }
+  std::string str() const override {
+    return "QScheme";
+  }
+  static const TypeKind Kind = TypeKind::QSchemeType;
+  // global singleton
+  static QSchemeTypePtr get();
+
+ private:
+  QSchemeType() : Type(TypeKind::QSchemeType) {}
+};
+
 struct DeviceObjType;
 using DeviceObjTypePtr = std::shared_ptr<DeviceObjType>;
 // This type represents a Generator
@@ -1261,6 +1284,12 @@ struct getTypePtr_<at::Scalar> final {
   }
 };
 template <>
+struct getTypePtr_<c10::QScheme> final {
+  static TypePtr call() {
+    return QSchemeType::get();
+  }
+};
+template <>
 struct getTypePtr_<at::Generator*> final {
   static TypePtr call() {
     return OptionalType::create(GeneratorType::get());
@@ -1392,8 +1421,8 @@ struct CAFFE2_API ClassType : public NamedType {
   }
 
   std::string str() const override {
-    return python_str();
-  }
+     return python_str();
+   }
 
   std::string python_str() const override {
     const auto& n = name().value();

aten/src/ATen/core/type.cpp

@@ -104,6 +104,10 @@ GeneratorTypePtr GeneratorType::get() {
   static auto value = GeneratorType::create();
   return value;
 }
+QSchemeTypePtr QSchemeType::get() {
+  static auto value = QSchemeType::create();
+  return value;
+}
 StringTypePtr StringType::get() {
   static auto value = StringType::create();
   return value;

aten/src/ATen/native/native_functions.yaml

@@ -3533,6 +3533,7 @@
     CPU: make_per_channel_quantized_tensor_cpu
 
 - func: qscheme(Tensor self) -> QScheme
+  use_c10_dispatcher: full
   variants: method
   dispatch:
     QuantizedCPU: qscheme_quant

torch/csrc/jit/pybind_utils.h

@@ -541,6 +541,7 @@ inline IValue toIValue(
     case TypeKind::GeneratorType:
     case TypeKind::VarType:
     case TypeKind::FutureType:
+    case TypeKind::QSchemeType:
       break;
     case TypeKind::FunctionType:
       AT_ERROR("Function Values aren't yet supported");

torch/csrc/jit/script/schema_type_parser.cpp

@@ -22,6 +22,7 @@ using c10::NumberType;
 using c10::OptionalType;
 using c10::StringType;
 using c10::Symbol;
+using c10::QSchemeType;
 using c10::TensorType;
 using c10::TupleType;
 using c10::VarType;
@@ -38,7 +39,7 @@ TypeAndAlias SchemaTypeParser::parseBaseType() {
       {"Layout", IntType::get()},
       {"MemoryFormat", IntType::get()},
       {"Storage", IntType::get()},
-      {"QScheme", IntType::get()},
+      {"QScheme", QSchemeType::get()},
       {"ConstQuantizerPtr", IntType::get()},  // TODO This type should be removed from the schema parser, it should use the custom class mechanism instead. @jerryzh
       {"Device", DeviceObjType::get()},
       {"Scalar", NumberType::get()},

torch/csrc/jit/unpickler.cpp

@@ -83,6 +83,7 @@ void restoreAccurateTypeTags(const IValue& root, const TypePtr& type_tag) {
       case StringType::Kind:
       case FunctionType::Kind:
       case DeviceObjType::Kind:
+      case QSchemeType::Kind:
         // no op, there is nothing to tag
         break;
       case AnyType::Kind: