[PyTorch] Add quantized packing weights support for fc & conv

.circleci/build.sh

@@ -135,7 +135,7 @@ elif [[ "$CIRCLE_JOB" == "PYTORCH" ]]; then
     git clone https://github.com/pytorch/pytorch.git --recursive --depth 1
     cd pytorch
     pip install -r requirements.txt
-    BUILD_BINARY=OFF BUILD_TEST=0 BUILD_CAFFE2_OPS=0 python setup.py install
+    BUILD_BINARY=OFF BUILD_TEST=0 BUILD_CAFFE2_OPS=0 USE_FBGEMM=ON python setup.py install
     cd ${GLOW_DIR}
     cd build
 elif [[ "$CIRCLE_JOB" == "OPENCL" ]]; then

torch_glow/src/GlowIValue.cpp

@@ -46,6 +46,8 @@ const char *GlowIValue::tagToStr(GlowIValue::Tag tag) {
     return "BoolList";
   case GlowIValue::Tag::Tuple:
     return "Tuple";
+  case GlowIValue::Tag::PTTensor:
+    return "PyTorch Tensor";
   }
   LOG(DFATAL) << "Cannot reach here.";
 }
@@ -67,6 +69,9 @@ void GlowIValue::reset() {
   case Tag::Tuple:
     delete payload_.asTuple;
     break;
+  case Tag::PTTensor:
+    delete payload_.asPTTensor;
+    break;
   case Tag::None:
   case Tag::Double:
   case Tag::Int:
@@ -104,6 +109,7 @@ bool GlowIValue::isIntList() const { return Tag::IntList == tag_; }
 bool GlowIValue::isDoubleList() const { return Tag::DoubleList == tag_; }
 bool GlowIValue::isBoolList() const { return Tag::BoolList == tag_; }
 bool GlowIValue::isTuple() const { return Tag::Tuple == tag_; }
+bool GlowIValue::isPTTensor() const { return Tag::PTTensor == tag_; }
 
 #define ExpectTag(EXPECTED_TAG)                                                \
   RETURN_ERR_IF_NOT(tag_ == (EXPECTED_TAG),                                    \
@@ -175,6 +181,16 @@ Expected<const std::vector<GlowIValue> *> GlowIValue::toTuple() const {
   return payload_.asTuple;
 }
 
+Expected<at::Tensor *> GlowIValue::toPTTensor() {
+  ExpectTag(Tag::PTTensor);
+  return payload_.asPTTensor;
+}
+
+Expected<const at::Tensor *> GlowIValue::toPTTensor() const {
+  ExpectTag(Tag::PTTensor);
+  return payload_.asPTTensor;
+}
+
 #undef ExpectTag
 
 void GlowIValue::fromNone() {
@@ -234,6 +250,12 @@ void GlowIValue::fromTuple(std::vector<GlowIValue> glowIValList) {
   std::swap(glowIValList, *payload_.asTuple);
 }
 
+void GlowIValue::fromPTTensor(at::Tensor tensor) {
+  reset();
+  tag_ = Tag::PTTensor;
+  payload_.asPTTensor = new at::Tensor(tensor);
+}
+
 Error GlowIValue::fromIValue(const at::IValue &ival) {
   reset();
   if (ival.isNone()) {

torch_glow/src/GlowIValue.h

@@ -41,6 +41,7 @@ class GlowIValue {
     DoubleList,
     BoolList,
     Tuple,
+    PTTensor,
   };
 
 private:
@@ -56,6 +57,7 @@ class GlowIValue {
     std::vector<double> *asDoubleList;
     std::vector<bool> *asBoolList;
     std::vector<GlowIValue> *asTuple;
+    at::Tensor *asPTTensor;
   };
 
   Tag tag_ = Tag::None;
@@ -96,6 +98,7 @@ class GlowIValue {
   bool isDoubleList() const;
   bool isBoolList() const;
   bool isTuple() const;
+  bool isPTTensor() const;
 
   /// \returns Payload a glow Tensor or error if the tag is not Tensor.
   Expected<Tensor *> toTensor();
@@ -138,12 +141,23 @@ class GlowIValue {
   /// \returns Payload a vector of GlowIValues or error if the tag is not Tuple.
   Expected<const std::vector<GlowIValue> *> toTuple() const;
 
+  /// \returns Payload a PyTorch Tensor* or error if the tag is not a PyTorch
+  /// Tensor.
+  Expected<at::Tensor *> toPTTensor();
+
+  /// \returns Payload a const Pytorch Tensor* or error if the tag is not
+  /// Tensor.
+  Expected<const at::Tensor *> toPTTensor() const;
+
   /// Set the tag to None.
   void fromNone();
 
   /// Set the tag to Tensor.
   void fromTensor(Tensor tensor);
 
+  /// Set the tag to PyTorch Tensor.
+  void fromPTTensor(at::Tensor tensor);
+
   /// Set the tag to Double.
   void fromDouble(double d);
 

torch_glow/src/PyTorchCommon.cpp

@@ -132,6 +132,14 @@ glow::ElemKind scalarTypeToElemKind(c10::ScalarType ty) {
     return ElemKind::Int64ITy;
   } else if (ty == at::kBool) {
     return ElemKind::BoolTy;
+  } else if (ty == at::kByte) {
+    // We should have an 8-byte non-quantized integer type eventually
+    // Currently usage of Bool is fine
+    return ElemKind::BoolTy;
+  } else if (ty == at::kQInt8) {
+    return ElemKind::Int8QTy;
+  } else if (ty == at::kQUInt8) {
+    return ElemKind::UInt8QTy;
   } else {
     LOG(DFATAL) << "ScalarType " << static_cast<int>(ty)
                 << " not supported yet.";
@@ -177,6 +185,20 @@ glow::Type ptTypeToGlowType(const c10::TensorType &ptType) {
   return glow::Type(scalarTypeToElemKind(scalarType), dims);
 }
 
+glow::Type ptTypeToGlowType(const c10::TensorType &ptType, float scale,
+                            int32_t zero_point) {
+  DCHECK(ptType.scalarType().has_value())
+      << "TensorType has no associated scalar type.";
+  const auto concreteSizes = ptType.sizes().concrete_sizes().value();
+  std::vector<glow::dim_t> dims;
+  for (const auto &size : concreteSizes) {
+    dims.push_back(static_cast<glow::dim_t>(size));
+  }
+
+  auto scalarType = ptType.scalarType().value();
+  return glow::Type(scalarTypeToElemKind(scalarType), dims, scale, zero_point);
+}
+
 at::Tensor glowTypeToEmptyPTTensor(const glow::Type &glowType) {
   std::vector<int64_t> sizes;
   for (const auto dim : glowType.dims()) {
@@ -188,7 +210,41 @@ at::Tensor glowTypeToEmptyPTTensor(const glow::Type &glowType) {
 }
 
 glow::Tensor ptTensorToGlowTensor(const at::Tensor &ptTensor) {
-  auto glowType = ptTypeToGlowType(*c10::TensorType::create(ptTensor));
-  return glow::Tensor(ptTensor.data_ptr(), &glowType);
+  if (ptTensor.is_quantized()) {
+    float scale = 1.0;
+    int32_t offset = 0;
+    if (ptTensor.qscheme() == at::kPerChannelAffine) {
+      // If it is channel wise quantized, which means
+      // this tensor is the weight of quantized linear or conv
+      // Then we dont deal with the qparams here,
+      // and only set up soome dummy scale & offset by using the first
+      // elements's scale & offset.
+      scale = ptTensor.q_per_channel_scales()[0].item<float>();
+      offset = ptTensor.q_per_channel_zero_points()[0].item<int32_t>();
+    } else if (ptTensor.qscheme() == at::kPerTensorAffine) {
+      scale = static_cast<float>(ptTensor.q_scale());
+      offset = static_cast<int32_t>(ptTensor.q_zero_point());
+    } else {
+      LOG(DFATAL)
+          << "PyTorch tensor with unsupported quantization scheme detected.";
+    }
+    auto glowType =
+        ptTypeToGlowType(*c10::TensorType::create(ptTensor), scale, offset);
+    return glow::Tensor(ptTensor.data_ptr(), &glowType);
+  } else {
+    auto glowType = ptTypeToGlowType(*c10::TensorType::create(ptTensor));
+    return glow::Tensor(ptTensor.data_ptr(), &glowType);
+  }
+}
+
+at::Tensor glowTensorToPTTensor(const glow::Tensor &glowTensor,
+                                const at::ScalarType &torch_type) {
+  std::vector<int64_t> sizes;
+  for (const auto dim : glowTensor.dims()) {
+    sizes.push_back(dim);
+  }
+  return at::from_blob(glowTensor.getUnsafePtr(), sizes,
+                       at::device(at::kCPU).dtype(torch_type));
 }
+
 } // namespace glow

torch_glow/src/PyTorchCommon.h

@@ -26,6 +26,11 @@
 
 namespace glow {
 
+/// For Glow: -128 <= orig_fp32/scale_1 + offset_1 <= 127
+/// For PyTorch: 0 <= orig_fp32/scale_2 + offset_2 <= 255
+/// Therefore, we can make scale_1 == scale_2, and offset_1 = offset2 - 128
+const int32_t OFFSETSHIFT = 128;
+
 extern bool GlowCompilePyTorchModule;
 /// Various settings to be used by code that loads PyTorch models. There should
 /// only be one of these and it should be obtained by calling
@@ -88,6 +93,11 @@ glow::Tensor ptTensorToGlowTensor(const at::Tensor &ptTensor);
 /// matching type.
 at::Tensor glowTypeToEmptyPTTensor(const glow::Type &glowType);
 
+/// Given a Glow Tensor \p glowTensor, \returns a PyTorch Tensor with the same
+/// type, shape and content.
+at::Tensor glowTensorToPTTensor(const glow::Tensor &glowTensor,
+                                const at::ScalarType &torch_type);
+
 } // namespace glow
 
 #endif // GLOW_TORCH_GLOW_SRC_COMMON_H