Merge from upstream

.jenkins/pytorch/build.sh

@@ -42,11 +42,11 @@ if [[ "$BUILD_ENVIRONMENT" == *rocm* ]]; then
   # This environment variable enabled HCC Optimizations that speed up the linking stage.
   # https://github.com/RadeonOpenCompute/hcc#hcc-with-thinlto-linking
   export KMTHINLTO=1
-  
+
   # Need the libc++1 and libc++abi1 libraries to allow torch._C to load at runtime
   sudo apt-get install libc++1
   sudo apt-get install libc++abi1
-  
+
   python tools/amd_build/build_pytorch_amd.py
   USE_ROCM=1 python setup.py install --user
   exit 0
@@ -118,5 +118,9 @@ if [[ "$BUILD_TEST_LIBTORCH" == "1" ]]; then
   echo "Building libtorch"
   # NB: Install outside of source directory (at the same level as the root
   # pytorch folder) so that it doesn't get cleaned away prior to docker push.
-  WERROR=1 VERBOSE=1 tools/cpp_build/build_caffe2.sh "$PWD/../cpp-build"
+  BUILD_LIBTORCH_PY=$PWD/tools/build_libtorch.py
+  mkdir -p ../cpp-build/caffe2
+  pushd ../cpp-build/caffe2
+  WERROR=1 VERBOSE=1 DEBUG=1 python $BUILD_LIBTORCH_PY
+  popd
 fi

.jenkins/pytorch/macos-test.sh

@@ -56,8 +56,12 @@ test_cpp_api() {
   #
   CPP_BUILD="$PWD/../cpp-build"
   rm -rf $CPP_BUILD
-  mkdir -p $CPP_BUILD
-  WERROR=1 VERBOSE=1 tools/cpp_build/build_caffe2.sh "$CPP_BUILD"
+  mkdir -p $CPP_BUILD/caffe2
+
+  BUILD_LIBTORCH_PY=$PWD/tools/build_libtorch.py
+  pushd $CPP_BUILD/caffe2
+  WERROR=1 VERBOSE=1 DEBUG=1 python $BUILD_LIBTORCH_PY
+  popd
 
   python tools/download_mnist.py --quiet -d test/cpp/api/mnist
 

CMakeLists.txt

@@ -230,7 +230,7 @@ if(NOT MSVC)
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Qunused-arguments")
   endif()
   if ((APPLE AND (NOT ("${CLANG_VERSION_STRING}" VERSION_LESS "9.0")))
-    OR (CMAKE_COMPILER_IS_GNUCXX 
+    OR (CMAKE_COMPILER_IS_GNUCXX
     AND (CMAKE_CXX_COMPILER_VERSION VERSION_GREATER 7.0 AND NOT APPLE)))
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -faligned-new")
   endif()

aten/src/ATen/Declarations.cwrap

@@ -2394,7 +2394,7 @@
     if (self_->dim() == 0) {
       throw std::runtime_error("Input must be 1-d or 2-d");
     }
-    ${THTensor}_diag(${state,}result_->tensor, self_->tensor, diagonal);
+    ${THTensor}_diag(${state,}result_, self_, diagonal);
     result_->maybe_zero_dim(self_->dim() == 0);
 ]]
 [[
@@ -2986,7 +2986,7 @@
     - arg: real tol
       default: -1
   aten_custom_call: |
-    ${THTensor}_pstrf(res1_->tensor, res2_->tensor, self_->tensor, (upper) ? "U" : "L", tol_);
+    ${THTensor}_pstrf(res1_, res2_, self_, (upper) ? "U" : "L", tol_);
     res2 -= 1;  // LAPACK returns 1-indexed pivots
 ]]
 [[

aten/src/ATen/Registry.h

@@ -114,7 +114,7 @@ class AT_API Registry {
 };
 
 template <class SrcType, class ObjectPtrType, class... Args>
-class Registerer {
+class AT_API Registerer {
  public:
   Registerer(
       const SrcType& key,

aten/src/ATen/SparseTensorImpl.cpp

@@ -31,7 +31,7 @@ namespace {
 // we don't currently support zero-size dimensions, so we can't actually
 // do this; so we just allocate zero-size tensors for everything.
 SparseTensorImpl::SparseTensorImpl(at::TensorTypeId type_id, at::ScalarType scalar_type)
-    : TensorImpl(type_id, scalar_type, nullptr, false)
+    : TensorImpl(type_id, scalar_type, false)
     , size_{0}
     , sparseDims_(1)
     , denseDims_(0)
@@ -44,6 +44,14 @@ IntList SparseTensorImpl::sizes() const {
 IntList SparseTensorImpl::strides() const {
   AT_ERROR("sparse tensors do not have strides");
 }
+int64_t SparseTensorImpl::size(int64_t d) const {
+  d = at::maybe_wrap_dim(d, dim(), false);
+  return size_[d];
+}
+int64_t SparseTensorImpl::stride(int64_t d) const {
+  AT_ERROR("sparse tensors do not have strides");
+}
+
 int64_t SparseTensorImpl::dim() const {
   return sparseDims_ + denseDims_;
 }
@@ -54,13 +62,15 @@ TensorImpl* SparseTensorImpl::maybe_zero_dim(bool condition_when_zero_dim) {
            " changing dimensionality via maybe_zero_dim");
   return this;
 }
-void * SparseTensorImpl::unsafeGetTH(bool retain) {
-  AT_ERROR("unsafeGetTH not supported for new style TensorImpl");
-}
 std::unique_ptr<Storage> SparseTensorImpl::storage() {
   AT_ERROR("sparse tensors do not have storage");
 }
-
+at::StorageImpl* SparseTensorImpl::storageImpl() const {
+  AT_ERROR("sparse tensors do not have storage");
+}
+int64_t SparseTensorImpl::storage_offset() const {
+  AT_ERROR("sparse tensors do not have storage");
+}
 void SparseTensorImpl::set_indices_and_values(const Tensor& indices, const Tensor& values) {
   // TODO: Explicit empty test is needed because we don't handle size zero
   // dimensions at the moment

aten/src/ATen/SparseTensorImpl.h

@@ -59,10 +59,14 @@ struct AT_API SparseTensorImpl : public TensorImpl {
 
   IntList sizes() const override;
   IntList strides() const override;
+  int64_t size(int64_t d) const override;
+  int64_t stride(int64_t d) const override;
+
   int64_t dim() const override;
   TensorImpl* maybe_zero_dim(bool condition_when_zero_dim) override;
-  void * unsafeGetTH(bool retain) override;
   std::unique_ptr<Storage> storage() override;
+  at::StorageImpl* storageImpl() const override;
+  int64_t storage_offset() const override;
 
   // Some ops do some manual size fiddling.
   // TODO: Figure out a more safe way to provide this functionality

aten/src/ATen/StorageImpl.h

@@ -1,7 +1,5 @@
 #pragma once
 
-#include <ATen/Scalar.h>
-
 #include <ATen/Allocator.h>
 #include <ATen/ScalarType.h>
 #include <ATen/ScalarTypeUtils.h>

aten/src/ATen/TensorImpl.cpp

@@ -59,59 +59,77 @@ void Tensor::backward(
   pImpl->backward(std::move(gradient), keep_graph, create_graph);
 }
 
-TensorImpl::TensorImpl(TensorTypeId type_id, ScalarType scalar_type)
-    : type_id_(type_id), scalar_type_(scalar_type) {
-  auto type = &globalContext().getType(tensorTypeIdToBackend(type_id), scalar_type);
-  auto storage = type->storage(true);
-  StorageImpl* storage_impl = storage->pImpl();
-  storage_impl->retain();
-  tensor = new THTensor(storage_impl);
+TensorImpl::TensorImpl(TensorTypeId type_id, ScalarType scalar_type, bool is_variable)
+    : TensorImpl(nullptr, type_id, scalar_type, is_variable) {
+  // UndefinedTensors and SparseTensors don't have storages.
+  if (type_id != UndefinedTensorId() && scalar_type != ScalarType::Undefined
+      && type_id != SparseCPUTensorId() && type_id != SparseCUDATensorId()) {
+    auto type = &globalContext().getType(tensorTypeIdToBackend(type_id), scalar_type);
+    auto storage = type->storage(true);
+    storage_ = storage->pImpl();
+    storage_->retain();
+  }
 }
 
+TensorImpl::TensorImpl(StorageImpl* storage, TensorTypeId type_id, bool is_variable)
+    : TensorImpl(storage, type_id, storage->scalar_type(), is_variable) {}
+
+TensorImpl::TensorImpl(StorageImpl* storage, TensorTypeId type_id, ScalarType scalar_type, bool is_variable)
+    : storage_(storage),
+      storage_offset_(0),
+      sizes_{0},
+      strides_{1},
+      type_id_(type_id),
+      scalar_type_(scalar_type),
+      is_variable_(is_variable) {}
+
 TensorImpl::~TensorImpl() {
-  if (tensor) tensor->release();
+  if (storage_) {
+    storage_->release();
+    storage_ = nullptr;
+  }
 }
 
 IntList TensorImpl::sizes() const {
-  // NB: dim in tensor is not synchronized with THTensor, so it's
-  // important to apply dim here
-  return IntList(THTensor_getSizePtr(tensor), dim());
+  return sizes_;
 }
 
 IntList TensorImpl::strides() const {
-  // NB: dim in tensor is not synchronized with THTensor, so it's
-  // important to apply dim here
-  return IntList(THTensor_getStridePtr(tensor), dim());
+  return strides_;
 }
 
 void TensorImpl::release_resources() {
-  if (tensor) {
-      tensor->release();
-      tensor = nullptr;
+  if (storage_) {
+    storage_->release();
+    storage_ = nullptr;
   }
 }
 
 int64_t TensorImpl::dim() const {
-  return tensor->dim();
+  return sizes_.size();
 }
 
-TensorImpl* TensorImpl::maybe_zero_dim(bool condition_when_zero_dim) {
-  AT_CHECK(tensor, "TensorImpl without THTensor in maybe_zero_dim");
-  THTensor_maybe_zero_dim(tensor, condition_when_zero_dim);
-  return this;
+int64_t TensorImpl::size(int64_t d) const {
+  d = at::maybe_wrap_dim(d, dim(), false);
+  return sizes_[d];
+}
+
+int64_t TensorImpl::stride(int64_t d) const {
+  d = at::maybe_wrap_dim(d, dim(), false);
+  return strides_[d];
 }
 
-void * TensorImpl::unsafeGetTH(bool retain) {
-  if (retain) {
-    tensor->retain();
+TensorImpl* TensorImpl::maybe_zero_dim(bool condition_when_zero_dim) {
+  bool set_zero_dim = condition_when_zero_dim && this->sizes().size() == 1 && this->size(0) == 1;
+  if (set_zero_dim) {
+    THTensor_resizeDim(this, 0);
   }
-  return tensor;
+  return this;
 }
 
 std::unique_ptr<Storage> TensorImpl::storage() {
-  StorageImpl* storage = tensor->storage_;
-  storage->retain();
-  return std::unique_ptr<Storage>(new Storage(storage));
+  storage_->retain();
+  return std::unique_ptr<Storage>(new Storage(storage_));
 }
 
 } // namespace at

aten/src/ATen/TensorImpl.h

@@ -3,8 +3,7 @@
 #include <atomic>
 #include <memory>
 
-#include "ATen/Retainable.h"
-#include "ATen/ScalarType.h"
+#include "ATen/StorageImpl.h"
 #include "ATen/core/optional.h"
 #include "ATen/core/TensorTypeId.h"
 #include "ATen/core/TensorTypeIdRegistration.h"
@@ -20,9 +19,8 @@ struct Tensor;
 
 namespace at {
 struct AT_API TensorImpl : public Retainable {
-  explicit TensorImpl(TensorTypeId type_id, ScalarType scalar_type, THTensor * tensor, bool is_variable)
-  : type_id_(type_id), scalar_type_(scalar_type), is_variable_(is_variable), tensor(tensor) {}
-  TensorImpl(TensorTypeId type_id, ScalarType scalar_type);
+  TensorImpl(TensorTypeId type_id, ScalarType scalar_type, bool is_variable);
+  TensorImpl(StorageImpl* storage, TensorTypeId type_id, bool is_variable);
 
   virtual ~TensorImpl();
 
@@ -37,7 +35,6 @@ struct AT_API TensorImpl : public Retainable {
   virtual IntList sizes() const;
   virtual IntList strides() const;
   virtual int64_t dim() const;
-  virtual void * unsafeGetTH(bool retain);
   virtual std::unique_ptr<Storage> storage();
   friend struct Type;
 
@@ -95,14 +92,58 @@ struct AT_API TensorImpl : public Retainable {
 
   virtual void set_data(Tensor new_data);
 
+  // TODO: make these protected
+  // Note: storage->size() may be greater than the recorded size
+  // of a tensor
+  at::StorageImpl* storage_;
+  int64_t storage_offset_;
+
+  std::vector<int64_t> sizes_;
+  std::vector<int64_t> strides_;
+
+  template <typename T>
+  inline T * data() const {
+    return storageImpl()->data<T>() + storage_offset_;
+  }
+
+  template <typename T>
+  inline T * unsafe_data() const {
+    return storageImpl()->unsafe_data<T>() + storage_offset_;
+  }
+
+  inline at::ScalarType scalar_type() const {
+    return scalar_type_;
+  }
+
+  virtual int64_t storage_offset() const {
+    return storage_offset_;
+  }
+
+  // represents that numel() == 0.
+  inline bool is_empty() const {
+    for (int64_t i = 0; i < dim(); ++i) {
+      if (sizes()[i] == 0) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  virtual int64_t size(int64_t d) const;
+  virtual int64_t stride(int64_t d) const;
+
+  // TODO: get rid of this.
+  virtual at::StorageImpl* storageImpl() const { return storage_; }
+
 protected:
   TensorTypeId type_id_;
   // INVARIANT: When storage is non-null, this scalar type must
   // agree with the scalar type in storage
   ScalarType scalar_type_;
   bool is_variable_ = false;
   bool is_wrapped_number_ = false;
-public:
-  THTensor * tensor;
+
+private:
+  TensorImpl(StorageImpl* storage, TensorTypeId type_id, ScalarType scalar_type, bool is_variable);
 };
 } // namespace at

aten/src/ATen/UndefinedTensor.cpp

@@ -6,24 +6,37 @@ namespace at {
 
 // should this use the globalContext?  Can it get a context passed in somehow?
 UndefinedTensor::UndefinedTensor()
-: TensorImpl(UndefinedTensorId(), ScalarType::Undefined, nullptr, /* is variable */ false) {
+: TensorImpl(UndefinedTensorId(), ScalarType::Undefined,  /* is variable */ false) {
 }
 
 IntList UndefinedTensor::sizes() const {
   AT_ERROR("sizes() called on undefined Tensor");
 }
 
+int64_t UndefinedTensor::size(int64_t d) const {
+  AT_ERROR("size(dim) called on an undefined Tensor");
+}
+
+int64_t UndefinedTensor::stride(int64_t d) const {
+  AT_ERROR("stride(dim) called on an undefined Tensor");
+}
+
 int64_t UndefinedTensor::dim() const {
   AT_ERROR("dim() called on undefined Tensor");
 }
 
-void * UndefinedTensor::unsafeGetTH(bool retain) {
-  AT_ERROR("unsafeGetTH(bool retain) called on undefined Tensor");
-}
 std::unique_ptr<Storage> UndefinedTensor::storage() {
   AT_ERROR("storage() called on undefined Tensor");
 }
 
+at::StorageImpl* UndefinedTensor::storageImpl() const {
+  AT_ERROR("storageImpl() called on an undefined Tensor");
+}
+
+int64_t UndefinedTensor::storage_offset() const {
+  AT_ERROR("storage_offset() called on an undefined Tensor");
+}
+
 IntList UndefinedTensor::strides() const {
   AT_ERROR("strides() called on undefined Tensor");
 }

aten/src/ATen/UndefinedTensor.h

@@ -11,9 +11,12 @@ struct AT_API UndefinedTensor final : public TensorImpl {
   }
   IntList sizes() const override;
   IntList strides() const override;
+  int64_t size(int64_t d) const override;
+  int64_t stride(int64_t d) const override;
   int64_t dim() const override;
-  void * unsafeGetTH(bool retain) override;
   std::unique_ptr<Storage> storage() override;
+  at::StorageImpl* storageImpl() const override;
+  int64_t storage_offset() const override;
 private:
   UndefinedTensor();
   static UndefinedTensor _singleton;