Create a union type GlowIValue for loading values from PyTorch (pytorch#3479)

Summary:
* Add a type which maps closely to torch::jit::IValue but holds Glow tensors. This is a departure from the previous approach which is to just cast everything as Glow Constants.
* Separate values loaded from PyTorch into two categories, Tensors for which only shapes are known at load time, and all other values (ints, lists, etc) for which values can be reasoned about at load time. For each of these categories, the information known at load-time will eventually be guarded so that, for example if a tensor's shape changes or any other values's contents change, the graph will not be valid and reloading will happen
* This also adds most of the necessary ingredients to make loading quantized graphs more simple

Documentation:
doxygen
Will add some information to pytorch.md about how this works once it's settled
Pull Request resolved: pytorch#3479

Test Plan: CI

Differential Revision: D17184105

Pulled By: jackm321

fbshipit-source-id: fad7fbd67e904d5f4ce40347fc7dfd377764acf8

Author: jackm321

include/glow/Graph/Nodes.h

@@ -91,14 +91,20 @@ class Constant : public Storage {
     return k->getKind() == Kinded::Kind::ConstantKind;
   }
 
-  /// \returns a mutable reference to the payload tensor. If the payload tensor
-  /// is unowned then it will be converted to an owned copy before returning.
-  Tensor &getPayloadMutable() {
-    // If payload is unowned, make an owned copy of the payload for
-    // modification.
+  /// If payload is unowned, make an owned copy of the payload for
+  /// modification.
+  void ensureIsOwned() {
     if (payload_.isUnowned()) {
       payload_ = payload_.clone();
     }
+  }
+
+  /// \returns a mutable reference to the payload tensor. If the payload tensor
+  /// is unowned then it will be converted to an owned copy before returning.
+  Tensor &getPayloadMutable() {
+    /// Make sure the payload is owned before handing out a mutable reference.
+    ensureIsOwned();
+
     assert(!payload_.isUnowned() &&
            "Can only modify Constants with owned payloads");
     return payload_;

torch_glow/setup.cfg

@@ -3,4 +3,4 @@ test=pytest
 
 [tool:pytest]
 testpaths = tests
-addopts = --verbose
+addopts = --verbose

torch_glow/src/CMakeLists.txt

@@ -14,6 +14,7 @@ link_directories(${PYTORCH_DIR}/lib)
 add_library(PyTorchModelLoader
                      CachingGraphRunner.cpp
                      GlowFuser.cpp
+                     GlowIValue.cpp
                      PyTorchCommon.cpp
                      FuseLinear.cpp
                      PyTorchModelLoader.cpp)

torch_glow/src/CachingGraphRunner.cpp

@@ -104,9 +104,9 @@ llvm::Error CachingGraphRunner::runImpl(const PerGlowGraphInfo &info,
     for (auto size : ph->dims()) {
       sizes.push_back(static_cast<int64_t>(size));
     }
-    auto ptT = at::empty(
-        sizes, at::TensorOptions().dtype(
-                   PyTorchModelLoader::convertGlowType(ph->getType())));
+
+    auto ptT = glowTypeToEmptyPTTensor(*ph->getType());
+
     glow::Tensor t(ptT.data_ptr(), ph->getType());
 
     outputs.push_back(std::move(ptT));

torch_glow/src/GlowIValue.cpp

@@ -0,0 +1,277 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "GlowIValue.h"
+
+#include <ATen/core/ivalue.h>
+
+#include "PyTorchCommon.h"
+
+#include "glow/Support/Error.h"
+#include "glow/Support/Support.h"
+
+namespace glow {
+
+// static
+const char *GlowIValue::tagToStr(GlowIValue::Tag tag) {
+  switch (tag) {
+  case GlowIValue::Tag::None:
+    return "None";
+  case GlowIValue::Tag::Tensor:
+    return "Tensor";
+  case GlowIValue::Tag::Double:
+    return "Double";
+  case GlowIValue::Tag::Int:
+    return "Int";
+  case GlowIValue::Tag::Bool:
+    return "Bool";
+  case GlowIValue::Tag::IntList:
+    return "IntList";
+  case GlowIValue::Tag::DoubleList:
+    return "DoubleList";
+  case GlowIValue::Tag::BoolList:
+    return "BoolList";
+  case GlowIValue::Tag::Tuple:
+    return "Tuple";
+  }
+}
+
+void GlowIValue::reset() {
+  switch (tag_) {
+  case Tag::Tensor:
+    delete payload_.asTensor;
+    break;
+  case Tag::IntList:
+    delete payload_.asIntList;
+    break;
+  case Tag::DoubleList:
+    delete payload_.asDoubleList;
+    break;
+  case Tag::BoolList:
+    delete payload_.asBoolList;
+    break;
+  case Tag::Tuple:
+    delete payload_.asTuple;
+    break;
+  case Tag::None:
+  case Tag::Double:
+  case Tag::Int:
+  case Tag::Bool:
+    // Nothing to free.
+    break;
+  }
+  tag_ = Tag::None;
+}
+
+GlowIValue::~GlowIValue() { reset(); }
+
+GlowIValue::GlowIValue(GlowIValue &&other) {
+  std::swap(tag_, other.tag_);
+  std::swap(payload_, other.payload_);
+}
+
+GlowIValue &GlowIValue::operator=(GlowIValue &&other) {
+  reset();
+  std::swap(tag_, other.tag_);
+  std::swap(payload_, other.payload_);
+  return *this;
+}
+
+GlowIValue::Tag GlowIValue::getTag() const { return tag_; }
+
+const char *GlowIValue::getTagString() const { return tagToStr(tag_); }
+
+bool GlowIValue::isNone() const { return Tag::None == tag_; }
+bool GlowIValue::isTensor() const { return Tag::Tensor == tag_; }
+bool GlowIValue::isDouble() const { return Tag::Double == tag_; }
+bool GlowIValue::isInt() const { return Tag::Int == tag_; }
+bool GlowIValue::isBool() const { return Tag::Bool == tag_; }
+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_; }
+
+#define ExpectTag(EXPECTED_TAG)                                                \
+  RETURN_ERR_IF_NOT(tag_ == (EXPECTED_TAG),                                    \
+                    strFormat("Expected GlowIValue with tag %s but found %s",  \
+                              tagToStr((EXPECTED_TAG)), tagToStr(tag_)))
+
+llvm::Expected<Tensor *> GlowIValue::toTensor() {
+  ExpectTag(Tag::Tensor);
+  return payload_.asTensor;
+}
+
+llvm::Expected<const Tensor *> GlowIValue::toTensor() const {
+  ExpectTag(Tag::Tensor);
+  return payload_.asTensor;
+}
+
+llvm::Expected<double> GlowIValue::toDouble() const {
+  ExpectTag(Tag::Double);
+  return payload_.asDouble;
+}
+
+llvm::Expected<int64_t> GlowIValue::toInt() const {
+  ExpectTag(Tag::Int);
+  return payload_.asInt;
+}
+
+llvm::Expected<bool> GlowIValue::toBool() const {
+  ExpectTag(Tag::Bool);
+  return payload_.asBool;
+}
+
+llvm::Expected<std::vector<int64_t> *> GlowIValue::toIntList() {
+  ExpectTag(Tag::IntList);
+  return payload_.asIntList;
+}
+
+llvm::Expected<const std::vector<int64_t> *> GlowIValue::toIntList() const {
+  ExpectTag(Tag::IntList);
+  return payload_.asIntList;
+}
+
+llvm::Expected<std::vector<double> *> GlowIValue::toDoubleList() {
+  ExpectTag(Tag::DoubleList);
+  return payload_.asDoubleList;
+}
+
+llvm::Expected<const std::vector<double> *> GlowIValue::toDoubleList() const {
+  ExpectTag(Tag::DoubleList);
+  return payload_.asDoubleList;
+}
+
+llvm::Expected<std::vector<bool> *> GlowIValue::toBoolList() {
+  ExpectTag(Tag::BoolList);
+  return payload_.asBoolList;
+}
+
+llvm::Expected<const std::vector<bool> *> GlowIValue::toBoolList() const {
+  ExpectTag(Tag::BoolList);
+  return payload_.asBoolList;
+}
+
+llvm::Expected<std::vector<GlowIValue> *> GlowIValue::toTuple() {
+  ExpectTag(Tag::Tuple);
+  return payload_.asTuple;
+}
+
+llvm::Expected<const std::vector<GlowIValue> *> GlowIValue::toTuple() const {
+  ExpectTag(Tag::Tuple);
+  return payload_.asTuple;
+}
+
+#undef ExpectTag
+
+void GlowIValue::fromNone() {
+  reset();
+  tag_ = Tag::None;
+}
+
+void GlowIValue::fromTensor(Tensor tensor) {
+  reset();
+  tag_ = Tag::Tensor;
+  payload_.asTensor = new glow::Tensor(std::move(tensor));
+}
+
+void GlowIValue::fromDouble(double d) {
+  reset();
+  tag_ = Tag::Double;
+  payload_.asDouble = d;
+}
+
+void GlowIValue::fromInt(int64_t i) {
+  reset();
+  tag_ = Tag::Int;
+  payload_.asInt = i;
+}
+
+void GlowIValue::fromBool(bool b) {
+  reset();
+  tag_ = Tag::Bool;
+  payload_.asBool = b;
+}
+
+void GlowIValue::fromIntList(std::vector<int64_t> intList) {
+  reset();
+  tag_ = Tag::IntList;
+  payload_.asIntList = new std::vector<int64_t>;
+  std::swap(intList, *payload_.asIntList);
+}
+
+void GlowIValue::fromDoubleList(std::vector<double> doubleList) {
+  reset();
+  tag_ = Tag::DoubleList;
+  payload_.asDoubleList = new std::vector<double>;
+  std::swap(doubleList, *payload_.asDoubleList);
+}
+
+void GlowIValue::fromBoolList(std::vector<bool> boolList) {
+  reset();
+  tag_ = Tag::BoolList;
+  payload_.asBoolList = new std::vector<bool>;
+  std::swap(boolList, *payload_.asBoolList);
+}
+
+void GlowIValue::fromTuple(std::vector<GlowIValue> glowIValList) {
+  reset();
+  tag_ = Tag::Tuple;
+  payload_.asTuple = new std::vector<GlowIValue>;
+  std::swap(glowIValList, *payload_.asTuple);
+}
+
+llvm::Error GlowIValue::fromIValue(const at::IValue &ival) {
+  reset();
+  if (ival.isNone()) {
+    fromNone();
+  } else if (ival.isTensor()) {
+    glow::Tensor t = ptTensorToGlowTensor(ival.toTensor());
+    fromTensor(std::move(t));
+  } else if (ival.isDouble()) {
+    fromDouble(ival.toDouble());
+  } else if (ival.isInt()) {
+    fromInt(ival.toInt());
+  } else if (ival.isBool()) {
+    fromBool(ival.toBool());
+  } else if (ival.isDoubleList()) {
+    const auto ivalDoubles = ival.toDoubleList();
+    std::vector<double> doubles(ivalDoubles.begin(), ivalDoubles.end());
+    fromDoubleList(std::move(doubles));
+  } else if (ival.isIntList()) {
+    const auto ivalInts = ival.toIntList();
+    std::vector<int64_t> ints(ivalInts.begin(), ivalInts.end());
+    fromIntList(std::move(ints));
+  } else if (ival.isBoolList()) {
+    const auto ivalBools = ival.toBoolList();
+    std::vector<bool> bools(ivalBools.begin(), ivalBools.end());
+    fromBoolList(std::move(bools));
+  } else if (ival.isTuple()) {
+    const auto ivalTuple = ival.toTuple();
+    const auto &elems = ivalTuple->elements();
+    std::vector<GlowIValue> tuple;
+    for (const auto &elem : elems) {
+      GlowIValue glowIVal;
+      RETURN_IF_ERR(glowIVal.fromIValue(elem));
+      tuple.push_back(std::move(glowIVal));
+    }
+    fromTuple(std::move(tuple));
+  } else {
+    RETURN_ERR("Encountered unhandled IValue type");
+  }
+  return llvm::Error::success();
+}
+
+} // namespace glow